refactor: unify error handling, graceful shutdown, and constants across framework

src/middleware/api_key.rs

@@ -1,25 +1,36 @@
+//! API Key authentication middleware.
+//!
+//! Validates requests by checking for a valid API key in the X-API-Key header.
+
 use crate::{
     Responder,
     middleware::{Middleware, MiddlewareFuture, MiddlewareResult},
 };
-use http::Request;
-use hyper::body::Incoming;
+use http::{Request, Response};
+use hyper::body::{Bytes, Incoming};
 use log::warn;
 
+/// Middleware that validates API key authentication via X-API-Key header.
 pub struct ApiKeyMiddleware {
     api_key: String,
 }
 
 impl ApiKeyMiddleware {
+    /// Creates a new ApiKeyMiddleware with the specified expected API key.
     pub fn new(api_key: &str) -> Self {
         Self {
             api_key: api_key.to_string(),
         }
     }
+
+    /// Checks if the given API key is invalid.
+    pub fn is_invalid_key(&self, key: &str) -> bool {
+        key != self.api_key
+    }
 }
 
 impl Middleware for ApiKeyMiddleware {
-    fn run<'a>(&'a self, req: Request<Incoming>) -> MiddlewareFuture<'a> {
+    fn run(&self, req: Request<Incoming>) -> MiddlewareFuture<'_> {
         let expected_key = self.api_key.clone();
 
         Box::pin(async move {
@@ -28,15 +39,47 @@ impl Middleware for ApiKeyMiddleware {
                     if header == expected_key.as_str() {
                         MiddlewareResult::Continue(req)
                     } else {
-                        warn!("X-API-Key wrong");
-                        MiddlewareResult::Respond(Responder::unauthorized().unwrap())
+                        warn!("X-API-Key validation failed for request");
+                        // Return a default unauthorized response if Responder fails
+                        let response = Responder::unauthorized()
+                            .unwrap_or_else(|_| {
+                                // Fallback to a basic unauthorized response
+                                Response::builder()
+                                    .status(http::StatusCode::UNAUTHORIZED)
+                                    .body(http_body_util::Full::new(
+                                        Bytes::from("Unauthorized")
+                                    ))
+                                    .expect("Failed to build fallback response")
+                            });
+                        MiddlewareResult::Respond(response)
                     }
                 }
                 None => {
-                    warn!("X-API-Key missing");
-                    MiddlewareResult::Respond(Responder::unauthorized().unwrap())
+                    warn!("X-API-Key header missing from request");
+                    let response = Responder::unauthorized()
+                        .unwrap_or_else(|_| {
+                            Response::builder()
+                                .status(http::StatusCode::UNAUTHORIZED)
+                                .body(http_body_util::Full::new(
+                                    Bytes::from("Unauthorized")
+                                ))
+                                .expect("Failed to build fallback response")
+                        });
+                    MiddlewareResult::Respond(response)
                 }
             }
         })
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use http::Request;
+
+    #[test]
+    fn test_api_key_middleware_new() {
+        let middleware = ApiKeyMiddleware::new("test-key");
+        assert_eq!(middleware.api_key, "test-key");
+    }
+}