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 get_image( image_id: Path, State(context): State, ) -> Result, PicOxError> { let image = context.db.get::(*image_id)?; Ok(Response(image)) } async fn query_images( State(context): State, Json(image_query): Json, ) -> 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()), }, ) .first() .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 album = Album::find_album_by_query(&context.db, metadata.clone().unwrap().album) .first() .cloned(); let mut album = if let Some(album) = album { album } else { let name = metadata .clone() .unwrap() .album .album_name .ok_or(PicOxError::AlbumNotFound)?; let album = Album::new(&name, Vec::new(), user); context.db.insert::(album)? }; 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)) .route("/api/image/:id", get(get_image)) .with_state(context); let listener = tokio::net::TcpListener::bind(&config.host).await.unwrap(); info!("Serving at {}", config.host); axum::serve(listener, app).await.unwrap(); }