424 lines
12 KiB
Rust
424 lines
12 KiB
Rust
mod api;
|
|
mod config;
|
|
mod model;
|
|
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")]
|
|
enum SubCommands {
|
|
/// Start the PicOx server
|
|
Start,
|
|
/// Create a Picox API Key
|
|
CreateKey {
|
|
/// Key description
|
|
description: String,
|
|
/// API Key permissions (WRITE, DELETE)
|
|
permissions: ApiPermissions,
|
|
},
|
|
/// Dump the database state to json
|
|
Dump {
|
|
/// Output path
|
|
out_path: PathBuf,
|
|
},
|
|
/// Import the database state from json
|
|
Import {
|
|
/// Path to json containing DB state
|
|
db_file: PathBuf,
|
|
},
|
|
}
|
|
|
|
#[derive(Debug, Clone, StructOpt)]
|
|
struct Args {
|
|
/// Path to the config file
|
|
#[structopt(short, long, env = "PICOX_CONFIG")]
|
|
pub config: PathBuf,
|
|
#[structopt(subcommand)]
|
|
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();
|
|
|
|
// initialize tracing
|
|
tracing_subscriber::fmt::init();
|
|
|
|
let config = PicOxConfig::new(args.config);
|
|
|
|
let db = Database::new(&config.db_path).unwrap();
|
|
|
|
match args.command {
|
|
SubCommands::Start => {
|
|
run_picox(db, config).await;
|
|
}
|
|
SubCommands::CreateKey {
|
|
description,
|
|
permissions,
|
|
} => {
|
|
let (token, api_key) = ApiKey::create_new_key(&db, description, permissions).unwrap();
|
|
|
|
info!("New Key info: {:?}", api_key);
|
|
info!("Token: {}", base64::prelude::BASE64_STANDARD.encode(token));
|
|
db.db.flush().unwrap();
|
|
}
|
|
SubCommands::Dump { out_path } => {
|
|
info!("Dumping database state to {:?}", out_path);
|
|
|
|
tokio::fs::write(out_path, db.dump_db().unwrap().pretty(4))
|
|
.await
|
|
.unwrap();
|
|
}
|
|
SubCommands::Import { db_file } => {
|
|
info!("Importing database state from {:?}", db_file);
|
|
let db_state = tokio::fs::read(db_file).await.unwrap();
|
|
let db_state = String::from_utf8(db_state).unwrap();
|
|
let db_state = json::parse(&db_state).unwrap();
|
|
|
|
db.import_db(db_state).unwrap();
|
|
db.db.flush().unwrap();
|
|
}
|
|
}
|
|
}
|