Refactored API code into its own module

main
Joey Hines 2024-02-10 11:16:56 -07:00
parent 5d2011861b
commit 136405d19b
No known key found for this signature in database
GPG Key ID: 995E531F7A569DDB
4 changed files with 360 additions and 335 deletions

View File

@ -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<PicContext>,
headers: HeaderMap,
Json(album): Json<CreateAlbum>,
) -> Result<Response<Album>, 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<u64>,
State(context): State<PicContext>,
) -> Result<Response<Album>, PicOxError> {
let album = context.db.get::<Album>(*album_id)?;
Ok(Response(album))
}
async fn query_images(
image_query: Query<ImageQuery>,
State(context): State<PicContext>,
) -> Result<Response<Vec<Image>>, 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<Image> = 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<u64, PicOxError> {
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<PicContext>,
headers: HeaderMap,
mut img_data: Multipart,
) -> Result<Response<Image>, PicOxError> {
let mut data: Vec<u8> = Vec::new();
let mut metadata: Option<AddImage> = 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>(album)?;
Ok(Response(img))
}
async fn query_album(
album_query: Query<AlbumQuery>,
State(context): State<PicContext>,
) -> Result<Response<Option<Album>>, PicOxError> {
let resp = Album::find_album_by_query(&context.db, album_query.0);
Ok(Response(resp))
}
async fn check_token_header(
State(context): State<PicContext>,
mut request: Request,
next: Next,
) -> Result<impl IntoResponse, PicOxError> {
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();
}

118
src/api/models.rs 100644
View File

@ -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<Context>;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CreateAlbum {
pub album_name: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AlbumQuery {
pub album_name: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AddImage {
pub album: AlbumQuery,
pub tags: Vec<String>,
}
#[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<String>,
#[serde(default)]
pub tags: Vec<String>,
#[serde(default)]
pub order: ImageSort,
#[serde(default)]
pub limit: usize,
}
#[derive(FromRequest)]
#[from_request(via(axum::Json), rejection(PicOxError))]
pub struct Response<T>(pub T);
impl<T> IntoResponse for Response<T>
where
Json<T>: 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<StoreError> for PicOxError {
fn from(value: StoreError) -> Self {
Self::StoreError(value)
}
}
impl From<j_db::error::JDbError> 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()
}
}

View File

@ -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<Context>;
#[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<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
struct AddImage {
pub album: AlbumQuery,
pub tags: Vec<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub enum ImageSort {
Random,
DateAscending,
DateDescending,
#[default]
None,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
struct ImageQuery {
pub album: Option<String>,
#[serde(default)]
pub tags: Vec<String>,
#[serde(default)]
pub order: ImageSort,
#[serde(default)]
pub limit: usize,
}
#[derive(FromRequest)]
#[from_request(via(axum::Json), rejection(PicOxError))]
struct Response<T>(T);
impl<T> IntoResponse for Response<T>
where
Json<T>: 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<StoreError> for PicOxError {
fn from(value: StoreError) -> Self {
Self::StoreError(value)
}
}
impl From<j_db::error::JDbError> 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<PicContext>,
headers: HeaderMap,
Json(album): Json<CreateAlbum>,
) -> Result<Response<Album>, 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<u64>,
State(context): State<PicContext>,
) -> Result<Response<Album>, PicOxError> {
let album = context.db.get::<Album>(*album_id)?;
Ok(Response(album))
}
async fn query_images(
image_query: Query<ImageQuery>,
State(context): State<PicContext>,
) -> Result<Response<Vec<Image>>, 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<Image> = 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<u64, PicOxError> {
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<PicContext>,
headers: HeaderMap,
mut img_data: Multipart,
) -> Result<Response<Image>, PicOxError> {
let mut data: Vec<u8> = Vec::new();
let mut metadata: Option<AddImage> = 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>(album)?;
Ok(Response(img))
}
async fn query_album(
album_query: Query<AlbumQuery>,
State(context): State<PicContext>,
) -> Result<Response<Option<Album>>, PicOxError> {
let resp = Album::find_album_by_query(&context.db, album_query.0);
Ok(Response(resp))
}
async fn check_token_header(
State(context): State<PicContext>,
mut request: Request,
next: Next,
) -> Result<impl IntoResponse, PicOxError> {
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,

View File

@ -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;