feat: Send additional headers

zaunchat · Jul 15, 2022 · ea1f322 · ea1f322
2 parents 73b37e0 + 3d8021e
commit ea1f322
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 7 deletions.
diff --git a/src/middlewares/ratelimit.rs b/src/middlewares/ratelimit.rs
@@ -27,7 +27,7 @@ type RateLimiter =
     Limiter<String, DefaultKeyedStateStore<String>, DefaultClock, StateInformationMiddleware>;
 
 pub async fn ratelimit<B>(
-    mut req: Request<B>,
+    req: Request<B>,
     next: Next<B>,
     limiter: Arc<RateLimiter>,
 ) -> Result<Response, Error> {
@@ -59,11 +59,17 @@ pub async fn ratelimit<B>(
 
     if info.retry_after > 0 {
         log::info!("IP: {} has executed the rate limit", key);
-        req.extensions_mut().insert(info);
         return Err(Error::RateLimited(info));
     }
 
-    Ok(next.run(req).await)
+    let mut res = next.run(req).await;
+    let headers = res.headers_mut();
+
+    headers.insert("X-RateLimit-Remaining", info.remaining.into());
+    headers.insert("X-RateLimit-Limit", info.limit.into());
+    headers.insert("Retry-After", info.retry_after.into());
+
+    Ok(res)
 }
 
 #[macro_export]

diff --git a/src/utils/error.rs b/src/utils/error.rs
@@ -1,7 +1,7 @@
 use crate::middlewares::ratelimit::RateLimitInfo;
 use axum::{
     extract::rejection::JsonRejection,
-    http::StatusCode,
+    http::{HeaderMap, StatusCode},
     response::{IntoResponse, Json, Response},
 };
 use quick_error::quick_error;
@@ -10,6 +10,7 @@ use std::fmt::Debug;
 
 quick_error! {
     #[derive(Debug, Serialize, OpgModel)]
+    #[serde(tag = "type", rename_all = "snake_case")]
     pub enum Error {
         RateLimited(info: RateLimitInfo) {
             from(RateLimitInfo)
@@ -85,16 +86,24 @@ impl IntoResponse for Error {
         let status = match self {
             Error::RateLimited { .. } => StatusCode::TOO_MANY_REQUESTS,
             Error::InvalidToken => StatusCode::UNAUTHORIZED,
+            Error::InvalidBody => StatusCode::UNPROCESSABLE_ENTITY,
             _ => StatusCode::BAD_REQUEST,
         };
 
-        let mut body = serde_json::json!({ "type": self });
+        let mut headers = HeaderMap::new();
+        let mut body = serde_json::json!(self);
         let msg = self.to_string();
 
         if msg.contains(' ') {
-            body["message"] = serde_json::json!(msg);
+            body["message"] = msg.into();
         }
 
-        (status, Json(body)).into_response()
+        if let Error::RateLimited(info) = self {
+            headers.insert("X-RateLimit-Remaining", info.remaining.into());
+            headers.insert("X-RateLimit-Limit", info.limit.into());
+            headers.insert("Retry-After", info.retry_after.into());
+        }
+
+        (status, headers, Json(body)).into_response()
     }
 }