mrfluffy/bitBeam
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Update added download and user authentecation

Browse Source
This commit is contained in:
zastian@mrthoddata.com
2025-05-09 16:48:29 +01:00
parent 9b208edfde
commit 1c123f5a45
10 changed files with 456 additions and 68 deletions
Download Patch File Download Diff File Expand all files Collapse all files

381
src/api.rs
View File

@@ -1,27 +1,35 @@
use axum::{
body::Bytes,
extract::ConnectInfo,
extract::{ConnectInfo, Path},
http::{HeaderMap, StatusCode},
response::{IntoResponse, Response},
Extension, Json,
};
use chrono::Utc;
use log::{error, info, warn};
use rand::Rng;
use sqlx::AnyPool;
use std::path::Path;
use std::path::Path as PathBuf;
use tokio::fs;
use uuid::Uuid;
use log::{info, warn, error};
use std::net::SocketAddr;
use crate::data;
use std::net::SocketAddr;
use serde_json::json;
/// Handler to return all files as JSON
/// This function retrieves all files from the database
/// and returns them as a JSON response.
/// It also logs the IP address of the client making the request.
pub async fn all_files(Extension(pool): Extension<AnyPool>, ConnectInfo(addr): ConnectInfo<SocketAddr>) -> impl IntoResponse {
/// example request: curl -X GET http://localhost:3000/all_files
/// requires no parameters
/// returns a JSON array of files
/// TODO: add user authentication
pub async fn all_files(
Extension(pool): Extension<AnyPool>,
ConnectInfo(addr): ConnectInfo<SocketAddr>,
) -> impl IntoResponse {
//log the IP address of the client and the call
let ip = addr.ip().to_string();
info!("Received an all_files request from IP: {}", ip);
@@ -40,7 +48,7 @@ pub async fn all_files(Extension(pool): Extension<AnyPool>, ConnectInfo(addr): C
Ok(files) => {
info!("DB select all success");
(StatusCode::OK, Json(files)).into_response()
},
}
Err(e) => {
warn!("DB select all error: {}", e);
(
@@ -58,21 +66,67 @@ pub async fn all_files(Extension(pool): Extension<AnyPool>, ConnectInfo(addr): C
/// saves it to the server's file system,
/// and stores the file metadata in the database.
/// It also logs the IP address of the client making the request.
pub async fn upload(Extension(pool): Extension<AnyPool>,
ConnectInfo(addr): ConnectInfo<SocketAddr>,
Extension(config): Extension<data::Config> ,
headers: HeaderMap,
body: Bytes,
) -> Response {
/// example request: curl -X POST -H "key: <key>" -H "file_name: <file_name>" -H "content-type: <content_type>" -H "download_limit: <download_limit>" --data-binary @<file_path> http://localhost:3000/upload
/// requires the following headers:
/// - key: the key of the user (not optional)
/// - file_name: the name of the file (optional)
/// - content-type: the content type of the file (optional)
/// - download_limit: the download limit of the file (optional)
pub async fn upload(
Extension(pool): Extension<AnyPool>,
ConnectInfo(addr): ConnectInfo<SocketAddr>,
Extension(config): Extension<data::Config>,
headers: HeaderMap,
body: Bytes,
) -> Response {
//log the IP address of the client and the call
let ip = addr.ip().to_string();
info!("Received update from IP: {}", ip);
//get the key from the headers
let key = match headers.get("key") {
Some(hv) => hv.to_str().unwrap_or("unknown").to_string(),
None => {
return (
axum::http::StatusCode::BAD_REQUEST,
"Key header not supplied",
)
.into_response();
}
};
//check if the user exists
let owner = sqlx::query_as::<_, data::user>(
r#"
SELECT *
FROM users
WHERE key = ?
"#,
)
.bind(&key)
.fetch_one(&pool)
.await;
let owner = match owner {
Ok(user) => {
info!("User found in DB: {}", key);
user.username
}
Err(e) => {
error!("DB select error {}: {} Most likely because the Key is not valid", key, e);
return (
axum::http::StatusCode::INTERNAL_SERVER_ERROR,
"Your key is not valid",
)
.into_response();
}
};
// gets the content type from the headers
let content_type = headers
.get("content-type")
.and_then(|hv| hv.to_str().ok())
.unwrap_or("application/octet-stream")
.unwrap_or("unknown")
.to_string();
// gets the download limit from the headers
let download_limit = headers
@@ -80,6 +134,12 @@ pub async fn upload(Extension(pool): Extension<AnyPool>,
.and_then(|hv| hv.to_str().ok()) // Option<&str>
.and_then(|s| s.parse::<i32>().ok()) // Option<u32>
.unwrap_or(1); // u32
//get filename from the headers
let file_name = headers
.get("file_name")
.and_then(|hv| hv.to_str().ok())
.unwrap_or("unknown")
.to_string();
//generate a random UUID for the file ID
let id = {
// Fallback to random UUID if body is too small
@@ -87,9 +147,9 @@ pub async fn upload(Extension(pool): Extension<AnyPool>,
Uuid::from_u128(rng.random::<u128>()).to_string()
};
//create the directory if it doesn't exist
let dir = Path::new(&config.data_path);
let dir = PathBuf::new(&config.data_path);
if let Err(e) = fs::create_dir_all(dir).await {
warn!("could not make dir at {} error: {}", &config.data_path ,e);
warn!("could not make dir at {} error: {}", &config.data_path, e);
return (
axum::http::StatusCode::INTERNAL_SERVER_ERROR,
"Directory creation error",
@@ -98,20 +158,8 @@ pub async fn upload(Extension(pool): Extension<AnyPool>,
}
//create the file path
// the file path is the directory + the file ID + file type if file type is not application/x-executable
if content_type == "application/x-executable" {
info!("File type is application/x-executable");
} else {
info!("File type is {}", content_type);
}
let file_path = dir.join(
if content_type == "application/x-executable" {
format!("{}",id)
} else {
format!("{}.{}",
id,
content_type.split('/').last().unwrap_or("bin"))
},
);
info!("File type is {}", content_type);
let file_path = dir.join(&id);
if let Err(e) = fs::write(&file_path, &body).await {
warn!("write error {}: {}", id, e);
@@ -127,11 +175,17 @@ pub async fn upload(Extension(pool): Extension<AnyPool>,
let download_count = 0;
let download_url = match config.use_tls {
true => format!("https://{}/download/{}", config.base_url, id),
false => format!("http://{}/download/{}", config.base_url, id),
};
if let Err(e) = sqlx::query(
r#"
INSERT INTO files
(id, content_type, upload_time, download_limit, download_count, file_size)
VALUES (?, ?, ?, ?, ?, ?)
(id, content_type, upload_time, download_limit, download_count, file_size, download_url, file_name, owner)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
"#,
)
.bind(&id)
@@ -140,6 +194,9 @@ pub async fn upload(Extension(pool): Extension<AnyPool>,
.bind(download_limit)
.bind(download_count)
.bind(file_size as i64)
.bind(&download_url)
.bind(&file_name)
.bind(&owner)
.execute(&pool)
.await
{
@@ -151,13 +208,273 @@ pub async fn upload(Extension(pool): Extension<AnyPool>,
.into_response();
}
let uploaded_file = data::File {
id,
file_name,
content_type,
upload_time,
download_limit,
download_count,
file_size,
download_url,
owner,
};
Json(uploaded_file).into_response()
}
/// This is The file Download handler
/// This function handles the file download process.
/// It retrieves the file metadata from the database
/// and returns the file as a response.
/// It also logs the IP address of the client making the request.
/// example request: curl -X GET http://localhost:3000/download/<uuid>
/// requires the following path parameter:
/// - uuid: the UUID of the file (not optional)
pub async fn download_file(
Path(uuid): Path<String>, // Add this extractor
Extension(pool): Extension<AnyPool>,
ConnectInfo(addr): ConnectInfo<SocketAddr>,
Extension(config): Extension<data::Config>,
// Remove body: Bytes, // <-- GET handler shouldn't have a body
) -> Response {
// Get UUID directly from path
info!("Download request for UUID: {}", uuid);
// Log the IP address of the client and the call
let ip = addr.ip().to_string();
info!("Received download request for {} from IP: {}", uuid, ip);
// find file by uuid in the config.data_path
let file_path = PathBuf::new(&config.data_path).join(&uuid);
if !file_path.exists() {
error!("File not found: {}", file_path.display());
return (
axum::http::StatusCode::NOT_FOUND,
"File not found",
)
.into_response();
}
// Check if the file exists in the database
let file = sqlx::query_as::<_, data::File>(
r#"
SELECT *
FROM files
WHERE id = ?
"#,
)
.bind(&uuid)
.fetch_one(&pool)
.await;
let file = match file {
Ok(file) => {
info!("File found in DB: {}", uuid);
file
}
Err(e) => {
error!("DB select error {}: {}", uuid, e);
return (
axum::http::StatusCode::INTERNAL_SERVER_ERROR,
"Database select error",
)
.into_response();
}
};
//update download count
if let Err(e) = sqlx::query(
r#"
UPDATE files
SET download_count = download_count + 1
WHERE id = ?
"#,
)
.bind(&uuid)
.execute(&pool)
.await
{
error!("DB update error {}: {}", uuid, e);
return (
axum::http::StatusCode::INTERNAL_SERVER_ERROR,
"Database update error",
)
.into_response();
}
info!("Update Download Count Sucess for UUID: {}", uuid);
//rutn file to axum::body::Bytes
let file_bytes = match fs::read(&file_path).await {
Ok(file) => file,
Err(e) => {
error!("File read error {}: {}", uuid, e);
return (
axum::http::StatusCode::INTERNAL_SERVER_ERROR,
"File read error",
)
.into_response();
}
};
//if download count is greater or equal to download limit delete the file and remove it from the database
if (file.download_count) >= file.download_limit {
if let Err(e) = fs::remove_file(&file_path).await {
error!("File delete error {}: {}", uuid, e);
return (
axum::http::StatusCode::INTERNAL_SERVER_ERROR,
"File delete error",
)
.into_response();
}
if let Err(e) = sqlx::query(
r#"
DELETE FROM files
WHERE id = ?
"#,
)
.bind(&uuid)
.execute(&pool)
.await
{
error!("DB delete error {}: {}", uuid, e);
return (
axum::http::StatusCode::INTERNAL_SERVER_ERROR,
"Database delete error",
)
.into_response();
}
info!("File deleted from DB because max download limit was reached: {}", uuid);
}
// return the file as a response
return (
axum::http::StatusCode::OK,
axum::response::IntoResponse::into_response(
axum::response::Response::builder()
.header("Content-Disposition", format!("attachment; filename=\"{}\"", uuid))
.header("Content-Type", format!("{}", &file.content_type ))
.header("Content-Length", file.file_size)
.header("filename", file.file_name)
.body(axum::body::Body::from(file_bytes))
.unwrap(),
),
)
.into_response()
}
/// Handler to upload a file
/// This function registers a new user.
/// It receives the user data in the request headers,
/// saves it to the database,
/// and returns the user data as a JSON response.
/// It also logs the IP address of the client making the request.
/// example request: curl -X POST -H "username: <username>" -H "password: <password>" http://localhost:3000/register
/// requires the following headers:
/// - username: the username of the user (not optional)
/// - password: the password of the user (not optional)
pub async fn register_user (
Extension(pool): Extension<AnyPool>,
ConnectInfo(addr): ConnectInfo<SocketAddr>,
Extension(config): Extension<data::Config>,
headers: HeaderMap,
body: Bytes,
) -> Response {
//log the IP address of the client and the call
let ip = addr.ip().to_string();
info!("Received update from IP: {}", ip);
//check if registration is allowed
if !config.allow_register {
return (
axum::http::StatusCode::FORBIDDEN,
"Registration is not allowed",
)
.into_response();
}
// gets the content type from the headers return error if header is not suplyde
let username = match headers .get("username") {
Some(hv) => hv.to_str().unwrap_or("unknown").to_string(),
None => {
return (
axum::http::StatusCode::BAD_REQUEST,
"Username header not supplied",
)
.into_response();
}
};
let password = match headers .get("password") {
Some(hv) => hv.to_str().unwrap_or("unknown").to_string(),
None => {
return (
axum::http::StatusCode::BAD_REQUEST,
"Password header not supplied",
)
.into_response();
}
};
//generate a random UUID for the user key
let key = {
// Fallback to random UUID if body is too small
let mut rng = rand::rng();
Uuid::from_u128(rng.random::<u128>()).to_string()
};
// check if the user already exists
let user = sqlx::query_as::<_, data::user>(
r#"
SELECT *
FROM users
WHERE username = ?
"#,
)
.bind(&username)
.fetch_one(&pool)
.await;
match user {
Ok(_) => {
info!("User already exists: {}", username);
return (
axum::http::StatusCode::BAD_REQUEST,
"User already exists",
)
.into_response();
}
Err(e) => {
warn!("DB select error {}: {}", username, e);
}
}
//add the user to the database
if let Err(e) = sqlx::query(
r#"
INSERT INTO users
(key, username, password)
VALUES (?, ?, ?)
"#,
)
.bind(&key)
.bind(&username)
.bind(&password)
.execute(&pool)
.await
{
error!("DB insert error {}: {}", key, e);
return (
axum::http::StatusCode::INTERNAL_SERVER_ERROR,
"Database insert error",
)
.into_response();
}
info!("User registered: {}", username);
//return the user as a response
let registered_user = json!({
"key": key,
"username": username,
});
Json(registered_user)
.into_response()
}