custosm avatars and banners

5 files changed, 161 insertions, 10 deletions

diff --git a/src/api/image.rs b/src/api/image.rs
new file mode 100644
index 0000000..84eccc7
--- /dev/null
+++ b/src/api/image.rs
@@ -0,0 +1,54 @@
+use axum::{extract::Query, response::Response, routing::get, Router, http::StatusCode};
+use serde::Deserialize;
+
+use crate::types::http::ResponseCode;
+
+
+
+#[derive(Deserialize)]
+struct AvatarRequest {
+    user_id: u64,
+}
+
+async fn avatar(params: Option<Query<AvatarRequest>>) -> Response {
+    
+    let Some(params) = params else {
+        return ResponseCode::BadRequest.text("Missing query paramaters");
+    };
+
+    let custom = format!("/image/custom/avatar/{}.png", params.user_id);
+    let default = format!("/image/default/{}.png", params.user_id % 25);
+
+    let file = ResponseCode::Success.file(&custom).await;
+    if file.status() != StatusCode::OK {
+        return ResponseCode::Success.file(&default).await
+    }
+    file
+}
+
+#[derive(Deserialize)]
+struct BannerRequest {
+    user_id: u64,
+}
+
+async fn banner(params: Option<Query<BannerRequest>>) -> Response {
+    
+    let Some(params) = params else {
+        return ResponseCode::BadRequest.text("Missing query paramaters");
+    };
+
+    let custom = format!("/image/custom/banner/{}.png", params.user_id);
+
+    let file = ResponseCode::Success.file(&custom).await;
+    if file.status() != StatusCode::OK {
+        return ResponseCode::NotFound.text("User does not have a custom banner")
+    }
+    file
+}
+
+
+pub fn router() -> Router {
+    Router::new()
+        .route("/avatar", get(avatar))
+        .route("/banner", get(banner))
+}
diff --git a/src/api/mod.rs b/src/api/mod.rs
index d36b127..adc19d7 100644
--- a/src/api/mod.rs
+++ b/src/api/mod.rs
@@ -8,6 +8,7 @@ pub mod auth;
 pub mod pages;
 pub mod posts;
 pub mod users;
+pub mod image;
 
 pub fn router() -> Router {
     let governor_conf = Box::new(
diff --git a/src/api/users.rs b/src/api/users.rs
index afcdddd..83a0d4e 100644
--- a/src/api/users.rs
+++ b/src/api/users.rs
@@ -1,9 +1,9 @@
 use crate::types::{
-    extract::{AuthorizedUser, Check, CheckResult, Json},
+    extract::{AuthorizedUser, Check, CheckResult, Json, Png},
     http::ResponseCode,
     user::User,
 };
-use axum::{response::Response, routing::post, Router};
+use axum::{response::Response, routing::{post, put}, Router};
 use serde::Deserialize;
 
 #[derive(Deserialize)]
@@ -63,9 +63,33 @@ async fn load_self(AuthorizedUser(user): AuthorizedUser) -> Response {
     ResponseCode::Success.json(&json)
 }
 
+async fn avatar(AuthorizedUser(user): AuthorizedUser, Png(img): Png) -> Response {
+    
+    let path = format!("./public/image/custom/avatar/{}.png", user.user_id);
+
+    if img.save(path).is_err() {
+        return ResponseCode::InternalServerError.text("Failed to update avatar");
+    }
+
+    ResponseCode::Success.text("Successfully updated avatar")
+}
+
+async fn banner(AuthorizedUser(user): AuthorizedUser, Png(img): Png) -> Response {
+    
+    let path = format!("./public/image/custom/banner/{}.png", user.user_id);
+
+    if img.save(path).is_err() {
+        return ResponseCode::InternalServerError.text("Failed to update banner");
+    }
+
+    ResponseCode::Success.text("Successfully updated banner")
+}
+
 pub fn router() -> Router {
     Router::new()
         .route("/load", post(load_batch))
         .route("/self", post(load_self))
         .route("/page", post(load_page))
+        .route("/avatar", put(avatar))
+        .route("/banner", put(banner))
 }
diff --git a/src/main.rs b/src/main.rs
index a72ec5f..74f0a0b 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -3,10 +3,10 @@ use axum::{
     http::{Request, StatusCode},
     middleware::{self, Next},
     response::Response,
-    RequestExt, Router,
+    RequestExt, Router, extract::DefaultBodyLimit,
 };
 use axum_client_ip::ClientIp;
-use std::{net::SocketAddr, process::exit};
+use std::{net::SocketAddr, process::exit, fs};
 use tower_cookies::CookieManagerLayer;
 use tracing::{error, info, metadata::LevelFilter};
 use tracing_subscriber::{
@@ -14,7 +14,7 @@ use tracing_subscriber::{
 };
 use types::http::ResponseCode;
 
-use crate::api::pages;
+use crate::api::{pages, image};
 
 mod admin;
 mod api;
@@ -69,13 +69,19 @@ async fn main() {
         exit(1)
     };
 
+    fs::create_dir_all("./public/image/custom").expect("Coudn't make custom data directory");
+    fs::create_dir_all("./public/image/custom/avatar").expect("Coudn't make custom avatar directory");
+    fs::create_dir_all("./public/image/custom/banner").expect("Coudn't make custom banner directory");
+
     let app = Router::new()
         .fallback(not_found)
         .layer(middleware::from_fn(log))
         .layer(middleware::from_fn(serve))
         .nest("/", pages::router())
         .nest("/api", api::router())
-        .layer(CookieManagerLayer::new());
+        .nest("/image", image::router())
+        .layer(CookieManagerLayer::new())
+        .layer(DefaultBodyLimit::max(512_000));
 
     let Ok(addr) = "[::]:8080".parse::<std::net::SocketAddr>() else {
         error!("Failed to parse port binding");
diff --git a/src/types/extract.rs b/src/types/extract.rs
index 50c413b..54f250a 100644
--- a/src/types/extract.rs
+++ b/src/types/extract.rs
@@ -1,4 +1,4 @@
-use std::io::Read;
+use std::io::{Read, Cursor};
 
 use axum::{
     async_trait,
@@ -10,6 +10,7 @@ use axum::{
 };
 use axum_client_ip::ClientIp;
 use bytes::Bytes;
+use image::{io::Reader, ImageFormat, DynamicImage};
 use serde::de::DeserializeOwned;
 use tower_cookies::Cookies;
 
@@ -99,6 +100,36 @@ where
     }
 }
 
+pub struct Png(pub DynamicImage);
+
+#[async_trait]
+impl<S, B> FromRequest<S, B> for Png
+where
+    B: HttpBody + Sync + Send + 'static,
+    B::Data: Send,
+    B::Error: Into<BoxError>,
+    S: Send + Sync,
+{
+    type Rejection = Response;
+
+    async fn from_request(req: Request<B>, state: &S) -> Result<Self> {
+
+        let bytes = match read_body(req, state).await {
+            Ok(body) => body,
+            Err(err) => return Err(err),
+        };
+
+        let mut reader = Reader::new(Cursor::new(bytes));
+        reader.set_format(ImageFormat::Png);
+        
+        let Ok(img) = reader.decode()  else {
+            return Err(ResponseCode::BadRequest.text("Failed to decode png image"))
+        };
+
+        Ok(Self(img))
+    }
+}
+
 pub struct Json<T>(pub T);
 
 #[async_trait]
@@ -150,7 +181,43 @@ pub trait Check {
     }
 }
 
-pub async fn parse_body<S, B>(mut req: Request<B>, state: &S) -> Result<String>
+async fn read_body<S, B>(mut req: Request<B>, state: &S) -> Result<Vec<u8>>
+where
+    B: HttpBody + Sync + Send + 'static,
+    B::Data: Send,
+    B::Error: Into<BoxError>,
+    S: Send + Sync,
+{
+
+    let Ok(ClientIp(ip)) = req.extract_parts::<ClientIp>().await else {
+        tracing::error!("Failed to read client ip");
+        return Err(ResponseCode::InternalServerError.text("Failed to read client ip"));
+    };
+
+    let method = req.method().clone();
+    let uri = req.uri().clone();
+    let path = req
+        .extensions()
+        .get::<RouterURI>()
+        .map_or("", |path| path.0);
+
+    let Ok(bytes) = Bytes::from_request(req, state).await else {
+        return Err(ResponseCode::BadRequest.text("Request can be at most 512kb"));
+    };
+
+    console::log(
+        ip,
+        method,
+        uri,
+        Some(path.to_string()),
+        None,
+    )
+    .await;
+
+    Ok(bytes.bytes().flatten().collect())
+}
+
+async fn parse_body<S, B>(mut req: Request<B>, state: &S) -> Result<String>
 where
     B: HttpBody + Sync + Send + 'static,
     B::Data: Send,
@@ -170,8 +237,7 @@ where
         .map_or("", |path| path.0);
 
     let Ok(bytes) = Bytes::from_request(req, state).await else {
-        tracing::error!("Failed to read request body");
-        return Err(ResponseCode::InternalServerError.text("Failed to read request body"));
+        return Err(ResponseCode::BadRequest.text("Request can be at most 512kb"));
     };
 
     let Ok(body) = String::from_utf8(bytes.bytes().flatten().collect()) else {