From 136405d19b8a69477b1f63a6672030be0e0bbb21 Mon Sep 17 00:00:00 2001 From: Joey Hines Date: Sat, 10 Feb 2024 11:16:56 -0700 Subject: [PATCH] Refactored API code into its own module --- src/api/mod.rs | 240 ++++++++++++++++++++++++++++++++ src/api/models.rs | 118 ++++++++++++++++ src/main.rs | 335 +-------------------------------------------- src/model/album.rs | 2 +- 4 files changed, 360 insertions(+), 335 deletions(-) create mode 100644 src/api/models.rs diff --git a/src/api/mod.rs b/src/api/mod.rs index 8b13789..34a2e65 100644 --- a/src/api/mod.rs +++ b/src/api/mod.rs @@ -1 +1,241 @@ +use crate::api::models::{ + AddImage, AlbumQuery, CreateAlbum, ImageQuery, ImageSort, PicContext, PicOxError, Response, +}; +use crate::config::PicOxConfig; +use crate::model::album::Album; +use crate::model::api_key::ApiKey; +use crate::model::image::{Image, ImageData}; +use crate::state::Context; +use crate::storage_manager::StorageManager; +use axum::body::Bytes; +use axum::extract::{DefaultBodyLimit, Multipart, Path, Query, Request, State}; +use axum::http::HeaderMap; +use axum::middleware::Next; +use axum::response::IntoResponse; +use axum::routing::{get, post}; +use axum::{middleware, Json, Router}; +use j_db::database::Database; +use j_db::model::JdbModel; +use log::info; +use rand::prelude::SliceRandom; +use rand::thread_rng; +use std::sync::Arc; +use tokio::sync::RwLock; +pub mod models; + +async fn create_album( + State(context): State, + headers: HeaderMap, + Json(album): Json, +) -> Result, PicOxError> { + let user_id = get_user_id_from_headers(&headers)?; + let new_album = Album::new(&album.album_name, Vec::new(), user_id); + + info!( + "Creating new album '{}pub pub pub ' for user {}", + album.album_name, user_id + ); + let new_album = context.db.insert(new_album)?; + + Ok(Response(new_album)) +} + +async fn get_album( + album_id: Path, + State(context): State, +) -> Result, PicOxError> { + let album = context.db.get::(*album_id)?; + + Ok(Response(album)) +} + +async fn query_images( + image_query: Query, + State(context): State, +) -> Result>, PicOxError> { + let album_id = if let Some(album) = &image_query.album { + Some( + Album::find_album_by_query( + &context.db, + AlbumQuery { + album_name: Some(album.to_string()), + }, + ) + .ok_or(PicOxError::AlbumNotFound)? + .id() + .unwrap(), + ) + } else { + None + }; + + let mut images: Vec = context + .db + .filter(|_, img: &Image| { + if let Some(album_id) = album_id { + if img.album != album_id { + return false; + } + } + + if !image_query.tags.is_empty() { + let mut found = false; + for tag in &image_query.tags { + if img.tags.contains(tag) { + found = true; + break; + } + } + + if !found { + return false; + } + } + + true + })? + .collect(); + + match image_query.order { + ImageSort::Random => { + images.shuffle(&mut thread_rng()); + } + ImageSort::DateAscending => { + images.sort_by(|img_a, img_b| img_a.create_date.cmp(&img_b.create_date)) + } + ImageSort::DateDescending => { + images.sort_by(|img_a, img_b| img_b.create_date.cmp(&img_a.create_date)) + } + ImageSort::None => {} + } + + if images.len() > image_query.limit { + images.drain(image_query.limit..); + } + + Ok(Response(images)) +} + +fn get_user_id_from_headers(headers: &HeaderMap) -> Result { + let user = headers.get("user").ok_or(PicOxError::NoUserInHeader)?; + let user_str = user.to_str().unwrap(); + let user: u64 = user_str.parse().unwrap(); + + Ok(user) +} + +async fn add_image( + State(context): State, + headers: HeaderMap, + mut img_data: Multipart, +) -> Result, PicOxError> { + let mut data: Vec = Vec::new(); + let mut metadata: Option = None; + let mut file_name = None; + while let Some(field) = img_data.next_field().await.unwrap() { + let field_name = field.name(); + if let Some(field_name) = field_name { + if field_name == "metadata" { + let metadata_json = field.text().await.unwrap(); + metadata = Some(serde_json::from_str(&metadata_json).unwrap()) + } else if field_name == "img_data" { + file_name = Some(field.file_name().unwrap().to_string()); + let file_segment = field.bytes().await.unwrap_or(Bytes::new()); + data.extend_from_slice(file_segment.as_ref()); + } + } + } + + let user = headers.get("user").unwrap().to_str().unwrap(); + let user: u64 = user.parse().unwrap(); + + let mut album = + Album::find_album_by_query(&context.db, metadata.clone().unwrap().album).unwrap(); + + let mut store_manager = context.store_manager.write().await; + let img = store_manager + .store_img( + &context.db, + None, + ImageData::Bytes(data), + &file_name.unwrap(), + user, + album.id().unwrap(), + ) + .await?; + + info!( + "Creating new image id={} at {:?}", + img.id().unwrap(), + img.storage_location + ); + + album.images.push(img.id().unwrap()); + + context.db.insert::(album)?; + + Ok(Response(img)) +} + +async fn query_album( + album_query: Query, + State(context): State, +) -> Result>, PicOxError> { + let resp = Album::find_album_by_query(&context.db, album_query.0); + Ok(Response(resp)) +} + +async fn check_token_header( + State(context): State, + mut request: Request, + next: Next, +) -> Result { + let headers = request.headers(); + + if let Some(token) = headers.get("token") { + if let Some(api_key) = ApiKey::find_api_key_by_token(&context.db, token.to_str().unwrap())? + { + info!( + "Authenticated user {}: '{}'", + api_key.id().unwrap(), + api_key.description + ); + request + .headers_mut() + .insert("user", api_key.id().unwrap().into()); + return Ok(next.run(request).await); + } + } + + Err(PicOxError::TokenInvalid) +} + +pub async fn run_picox(db: Database, config: PicOxConfig) { + let store_manager = StorageManager::new(config.storage_config.clone()); + + let context = Context { + db, + config: config.clone(), + store_manager: RwLock::new(store_manager), + }; + + let context = Arc::new(context); + + let app = Router::new() + .route("/api/image/", post(add_image)) + .layer(DefaultBodyLimit::max(1024 * 1024 * 1024)) + .route("/api/album/create", post(create_album)) + .layer(middleware::from_fn_with_state( + context.clone(), + check_token_header, + )) + .route("/api/album/:id", get(get_album)) + .route("/api/album/", get(query_album)) + .route("/api/image/", get(query_images)) + .with_state(context); + + let listener = tokio::net::TcpListener::bind(&config.host).await.unwrap(); + info!("Serving at {}", config.host); + axum::serve(listener, app).await.unwrap(); +} diff --git a/src/api/models.rs b/src/api/models.rs new file mode 100644 index 0000000..2f1838c --- /dev/null +++ b/src/api/models.rs @@ -0,0 +1,118 @@ +use crate::state::Context; +use crate::storage_manager::StoreError; +use axum::http::StatusCode; +use axum::response::IntoResponse; +use axum::Json; +use axum_macros::FromRequest; +use serde::{Deserialize, Serialize}; +use std::sync::Arc; + +pub type PicContext = Arc; + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct CreateAlbum { + pub album_name: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct AlbumQuery { + pub album_name: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct AddImage { + pub album: AlbumQuery, + pub tags: Vec, +} + +#[derive(Debug, Clone, Serialize, Deserialize, Default)] +pub enum ImageSort { + Random, + DateAscending, + DateDescending, + #[default] + None, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ImageQuery { + pub album: Option, + #[serde(default)] + pub tags: Vec, + #[serde(default)] + pub order: ImageSort, + #[serde(default)] + pub limit: usize, +} + +#[derive(FromRequest)] +#[from_request(via(axum::Json), rejection(PicOxError))] +pub struct Response(pub T); + +impl IntoResponse for Response +where + Json: IntoResponse, +{ + fn into_response(self) -> axum::response::Response { + Json(self.0).into_response() + } +} + +#[allow(dead_code)] +pub enum PicOxError { + StoreError(StoreError), + DbError(j_db::error::JDbError), + AlbumNotFound, + ImageNotFound, + TokenInvalid, + NoUserInHeader, +} + +impl From for PicOxError { + fn from(value: StoreError) -> Self { + Self::StoreError(value) + } +} + +impl From for PicOxError { + fn from(value: j_db::error::JDbError) -> Self { + Self::DbError(value) + } +} + +#[derive(Serialize)] +pub struct ErrorResponse { + pub message: String, +} + +impl IntoResponse for PicOxError { + fn into_response(self) -> axum::response::Response { + let (status, message) = match self { + PicOxError::StoreError(err) => match err { + StoreError::InvalidFile => (StatusCode::BAD_REQUEST, err.to_string()), + StoreError::OutOfStorage => (StatusCode::INSUFFICIENT_STORAGE, err.to_string()), + StoreError::ImageTooBig => (StatusCode::UNAUTHORIZED, err.to_string()), + StoreError::IOError(_) => ( + StatusCode::INTERNAL_SERVER_ERROR, + "IO Error Has Occurred!".to_string(), + ), + }, + PicOxError::DbError(err) => (StatusCode::INTERNAL_SERVER_ERROR, err.to_string()), + PicOxError::AlbumNotFound => ( + StatusCode::INTERNAL_SERVER_ERROR, + "No album found!".to_string(), + ), + PicOxError::ImageNotFound => ( + StatusCode::INTERNAL_SERVER_ERROR, + "Image not found".to_string(), + ), + PicOxError::TokenInvalid => (StatusCode::UNAUTHORIZED, "Token is invalid".to_string()), + PicOxError::NoUserInHeader => ( + StatusCode::INTERNAL_SERVER_ERROR, + "User not found in header".to_string(), + ), + }; + + (status, Response(ErrorResponse { message })).into_response() + } +} diff --git a/src/main.rs b/src/main.rs index c081bd8..54d99a1 100644 --- a/src/main.rs +++ b/src/main.rs @@ -5,32 +5,12 @@ mod state; mod storage_manager; use crate::config::PicOxConfig; -use crate::model::album::Album; use crate::model::api_key::{ApiKey, ApiPermissions}; -use crate::model::image::{Image, ImageData}; -use crate::state::Context; -use crate::storage_manager::{StorageManager, StoreError}; -use axum::body::Bytes; -use axum::extract::{DefaultBodyLimit, Multipart, Path, Query, Request, State}; -use axum::http::{HeaderMap, StatusCode}; -use axum::middleware::Next; -use axum::response::IntoResponse; -use axum::routing::{get, post}; -use axum::{middleware, Json, Router}; -use axum_macros::FromRequest; use base64::Engine; use j_db::database::Database; -use j_db::model::JdbModel; use log::info; -use rand::seq::SliceRandom; -use rand::thread_rng; -use serde::{Deserialize, Serialize}; use std::path::PathBuf; -use std::sync::Arc; use structopt::StructOpt; -use tokio::sync::RwLock; - -type PicContext = Arc; #[derive(StructOpt, Debug, Clone)] #[structopt(about = "PicOx Commands")] @@ -65,319 +45,6 @@ struct Args { pub command: SubCommands, } -#[derive(Debug, Clone, Serialize, Deserialize)] -struct CreateAlbum { - pub album_name: String, -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -struct AlbumQuery { - pub album_name: Option, -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -struct AddImage { - pub album: AlbumQuery, - pub tags: Vec, -} - -#[derive(Debug, Clone, Serialize, Deserialize, Default)] -pub enum ImageSort { - Random, - DateAscending, - DateDescending, - #[default] - None, -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -struct ImageQuery { - pub album: Option, - #[serde(default)] - pub tags: Vec, - #[serde(default)] - pub order: ImageSort, - #[serde(default)] - pub limit: usize, -} - -#[derive(FromRequest)] -#[from_request(via(axum::Json), rejection(PicOxError))] -struct Response(T); - -impl IntoResponse for Response -where - Json: IntoResponse, -{ - fn into_response(self) -> axum::response::Response { - Json(self.0).into_response() - } -} - -#[allow(dead_code)] -pub enum PicOxError { - StoreError(StoreError), - DbError(j_db::error::JDbError), - AlbumNotFound, - ImageNotFound, - TokenInvalid, - NoUserInHeader, -} - -impl From for PicOxError { - fn from(value: StoreError) -> Self { - Self::StoreError(value) - } -} - -impl From for PicOxError { - fn from(value: j_db::error::JDbError) -> Self { - Self::DbError(value) - } -} - -#[derive(Serialize)] -struct ErrorResponse { - pub message: String, -} - -impl IntoResponse for PicOxError { - fn into_response(self) -> axum::response::Response { - let (status, message) = match self { - PicOxError::StoreError(err) => match err { - StoreError::InvalidFile => (StatusCode::BAD_REQUEST, err.to_string()), - StoreError::OutOfStorage => (StatusCode::INSUFFICIENT_STORAGE, err.to_string()), - StoreError::ImageTooBig => (StatusCode::UNAUTHORIZED, err.to_string()), - StoreError::IOError(_) => ( - StatusCode::INTERNAL_SERVER_ERROR, - "IO Error Has Occurred!".to_string(), - ), - }, - PicOxError::DbError(err) => (StatusCode::INTERNAL_SERVER_ERROR, err.to_string()), - PicOxError::AlbumNotFound => ( - StatusCode::INTERNAL_SERVER_ERROR, - "No album found!".to_string(), - ), - PicOxError::ImageNotFound => ( - StatusCode::INTERNAL_SERVER_ERROR, - "Image not found".to_string(), - ), - PicOxError::TokenInvalid => (StatusCode::UNAUTHORIZED, "Token is invalid".to_string()), - PicOxError::NoUserInHeader => ( - StatusCode::INTERNAL_SERVER_ERROR, - "User not found in header".to_string(), - ), - }; - - (status, Response(ErrorResponse { message })).into_response() - } -} - -async fn create_album( - State(context): State, - headers: HeaderMap, - Json(album): Json, -) -> Result, PicOxError> { - let user_id = get_user_id_from_headers(&headers)?; - let new_album = Album::new(&album.album_name, Vec::new(), user_id); - - info!( - "Creating new album '{}' for user {}", - album.album_name, user_id - ); - let new_album = context.db.insert(new_album)?; - - Ok(Response(new_album)) -} - -async fn get_album( - album_id: Path, - State(context): State, -) -> Result, PicOxError> { - let album = context.db.get::(*album_id)?; - - Ok(Response(album)) -} - -async fn query_images( - image_query: Query, - State(context): State, -) -> Result>, PicOxError> { - let album_id = if let Some(album) = &image_query.album { - Some( - Album::find_album_by_query( - &context.db, - AlbumQuery { - album_name: Some(album.to_string()), - }, - ) - .ok_or(PicOxError::AlbumNotFound)? - .id() - .unwrap(), - ) - } else { - None - }; - - let mut images: Vec = context - .db - .filter(|_, img: &Image| { - if let Some(album_id) = album_id { - if img.album != album_id { - return false; - } - } - - if !image_query.tags.is_empty() { - let mut found = false; - for tag in &image_query.tags { - if img.tags.contains(tag) { - found = true; - break; - } - } - - if !found { - return false; - } - } - - true - })? - .collect(); - - match image_query.order { - ImageSort::Random => { - images.shuffle(&mut thread_rng()); - } - ImageSort::DateAscending => { - images.sort_by(|img_a, img_b| img_a.create_date.cmp(&img_b.create_date)) - } - ImageSort::DateDescending => { - images.sort_by(|img_a, img_b| img_b.create_date.cmp(&img_a.create_date)) - } - ImageSort::None => {} - } - - if images.len() > image_query.limit { - images.drain(image_query.limit..); - } - - Ok(Response(images)) -} - -fn get_user_id_from_headers(headers: &HeaderMap) -> Result { - let user = headers.get("user").ok_or(PicOxError::NoUserInHeader)?; - let user_str = user.to_str().unwrap(); - let user: u64 = user_str.parse().unwrap(); - - Ok(user) -} - -async fn add_image( - State(context): State, - headers: HeaderMap, - mut img_data: Multipart, -) -> Result, PicOxError> { - let mut data: Vec = Vec::new(); - let mut metadata: Option = None; - let mut file_name = None; - while let Some(field) = img_data.next_field().await.unwrap() { - let field_name = field.name(); - if let Some(field_name) = field_name { - if field_name == "metadata" { - let metadata_json = field.text().await.unwrap(); - metadata = Some(serde_json::from_str(&metadata_json).unwrap()) - } else if field_name == "img_data" { - file_name = Some(field.file_name().unwrap().to_string()); - let file_segment = field.bytes().await.unwrap_or(Bytes::new()); - data.extend_from_slice(file_segment.as_ref()); - } - } - } - - let user = headers.get("user").unwrap().to_str().unwrap(); - let user: u64 = user.parse().unwrap(); - - let mut album = - Album::find_album_by_query(&context.db, metadata.clone().unwrap().album).unwrap(); - - let mut store_manager = context.store_manager.write().await; - let img = store_manager - .store_img( - &context.db, - None, - ImageData::Bytes(data), - &file_name.unwrap(), - user, - album.id().unwrap(), - ) - .await?; - - album.images.push(img.id().unwrap()); - - context.db.insert::(album)?; - - Ok(Response(img)) -} - -async fn query_album( - album_query: Query, - State(context): State, -) -> Result>, PicOxError> { - let resp = Album::find_album_by_query(&context.db, album_query.0); - Ok(Response(resp)) -} - -async fn check_token_header( - State(context): State, - mut request: Request, - next: Next, -) -> Result { - let headers = request.headers(); - - if let Some(token) = headers.get("token") { - if let Some(api_key) = ApiKey::find_api_key_by_token(&context.db, token.to_str().unwrap())? - { - request - .headers_mut() - .insert("user", api_key.id().unwrap().into()); - return Ok(next.run(request).await); - } - } - - Err(PicOxError::TokenInvalid) -} - -async fn run_picox(db: Database, config: PicOxConfig) { - let store_manager = StorageManager::new(config.storage_config.clone()); - - let context = Context { - db, - config: config.clone(), - store_manager: RwLock::new(store_manager), - }; - - let context = Arc::new(context); - - let app = Router::new() - .route("/api/image/", post(add_image)) - .layer(DefaultBodyLimit::max(1024 * 1024 * 1024)) - .route("/api/album/create", post(create_album)) - .layer(middleware::from_fn_with_state( - context.clone(), - check_token_header, - )) - .route("/api/album/:id", get(get_album)) - .route("/api/album/", get(query_album)) - .route("/api/image/", get(query_images)) - .with_state(context); - - let listener = tokio::net::TcpListener::bind(&config.host).await.unwrap(); - info!("Serving at {}", config.host); - axum::serve(listener, app).await.unwrap(); -} - #[tokio::main] async fn main() { let args = Args::from_args(); @@ -391,7 +58,7 @@ async fn main() { match args.command { SubCommands::Start => { - run_picox(db, config).await; + api::run_picox(db, config).await; } SubCommands::CreateKey { description, diff --git a/src/model/album.rs b/src/model/album.rs index fb01051..2a3e869 100644 --- a/src/model/album.rs +++ b/src/model/album.rs @@ -1,4 +1,4 @@ -use crate::AlbumQuery; +use crate::api::models::AlbumQuery; use chrono::{DateTime, Utc}; use j_db::database::Database; use j_db::model::JdbModel;