1 files changed, 33 insertions, 4 deletions

diff --git a/src/public/mod.rs b/src/public/mod.rs
index cf8156d..eb846fb 100644
--- a/src/public/mod.rs
+++ b/src/public/mod.rs
@@ -1,11 +1,17 @@
 use axum::{
     body::Body,
-    http::{Request, StatusCode},
+    error_handling::HandleErrorLayer,
+    headers::HeaderName,
+    http::{HeaderValue, Request, StatusCode},
     response::{IntoResponse, Response},
     routing::get,
-    Router,
+    BoxError, Router,
+};
+use tower::{ServiceBuilder, ServiceExt};
+use tower_governor::{
+    errors::display_error, governor::GovernorConfigBuilder, key_extractor::SmartIpKeyExtractor,
+    GovernorLayer,
 };
-use tower::ServiceExt;
 use tower_http::services::ServeFile;
 
 use crate::types::http::ResponseCode;
@@ -16,6 +22,15 @@ pub mod file;
 pub mod pages;
 
 pub fn router() -> Router {
+    let governor_conf = Box::new(
+        GovernorConfigBuilder::default()
+            .burst_size(20)
+            .per_second(1)
+            .key_extractor(SmartIpKeyExtractor)
+            .finish()
+            .expect("Failed to create rate limiter"),
+    );
+
     Router::new()
         .nest("/", pages::router())
         .route("/js/*path", get(file::js))
@@ -24,6 +39,15 @@ pub fn router() -> Router {
         .route("/image/*path", get(file::image))
         .route("/image/avatar", get(file::avatar))
         .route("/image/banner", get(file::banner))
+        .layer(
+            ServiceBuilder::new()
+                .layer(HandleErrorLayer::new(|e: BoxError| async move {
+                    display_error(e)
+                }))
+                .layer(GovernorLayer {
+                    config: Box::leak(governor_conf),
+                }),
+        )
 }
 
 pub async fn serve(path: &str) -> Response {
@@ -34,7 +58,7 @@ pub async fn serve(path: &str) -> Response {
     let path = format!("public{path}");
     let file = ServeFile::new(path);
 
-    let Ok(res) = file.oneshot(Request::new(Body::empty())).await else {
+    let Ok(mut res) = file.oneshot(Request::new(Body::empty())).await else {
         tracing::error!("Error while fetching file");
         return ResponseCode::InternalServerError.text("Error while fetching file");
     };
@@ -43,5 +67,10 @@ pub async fn serve(path: &str) -> Response {
         return ResponseCode::NotFound.text("File not found");
     }
 
+    res.headers_mut().insert(
+        HeaderName::from_static("cache-control"),
+        HeaderValue::from_static("public max-age=300"),
+    );
+
     res.into_response()
 }