feat(anda_engine_server): support middleware

zensh · zensh · commit f1afb93cb2ea · 2026-01-16T19:17:37.000+08:00
diff --git a/anda_engine_server/Cargo.toml b/anda_engine_server/Cargo.toml
@@ -3,7 +3,7 @@ name = "anda_engine_server"
 description = "A http server to serve multiple Anda engines."
 repository = "https://github.com/ldclabs/anda/tree/main/anda_engine_server"
 publish = true
-version = "0.9.3"
+version = "0.9.4"
 edition.workspace = true
 keywords.workspace = true
 categories.workspace = true
diff --git a/anda_engine_server/src/handler.rs b/anda_engine_server/src/handler.rs
@@ -26,20 +26,6 @@ pub struct AppState {
     pub(crate) engines: Arc<BTreeMap<Principal, Engine>>,
     pub(crate) default_engine: Principal,
     pub(crate) start_time_ms: u64,
-    pub(crate) api_key: Option<String>,
-}
-
-impl AppState {
-    pub fn check_api_key(&self, headers: &http::HeaderMap) -> Result<(), String> {
-        if let Some(expected_key) = &self.api_key {
-            match headers.get("x-api-key") {
-                Some(provided_key) if provided_key == expected_key => Ok(()),
-                _ => Err("missing or invalid x-api-key in headers".to_string()),
-            }
-        } else {
-            Ok(())
-        }
-    }
 }
 
 /// GET /.well-known/information
@@ -112,10 +98,6 @@ pub async fn anda_engine(
     Path(id): Path<String>,
     ct: ContentWithSHA3<RPCRequest>,
 ) -> impl IntoResponse {
-    if let Err(err) = app.check_api_key(&headers) {
-        return (StatusCode::UNAUTHORIZED, err).into_response();
-    }
-
     let id = if &id == "default" {
         app.default_engine
     } else if let Ok(id) = Principal::from_text(&id) {
diff --git a/anda_engine_server/src/lib.rs b/anda_engine_server/src/lib.rs
@@ -8,9 +8,11 @@ use tokio::signal;
 use tokio_util::sync::CancellationToken;
 
 mod handler;
+mod middleware;
 mod types;
 
 use handler::*;
+pub use middleware::{ApiKeyMiddleware, HttpMiddleware};
 
 const APP_NAME: &str = env!("CARGO_PKG_NAME");
 const APP_VERSION: &str = env!("CARGO_PKG_VERSION");
@@ -22,7 +24,7 @@ pub struct ServerBuilder {
     origin: String,
     engines: BTreeMap<Principal, Engine>,
     default_engine: Option<Principal>,
-    api_key: Option<String>,
+    middlewares: Vec<Arc<dyn HttpMiddleware>>,
 }
 
 impl Default for ServerBuilder {
@@ -43,7 +45,7 @@ impl ServerBuilder {
             origin: "https://localhost:8443".to_string(),
             engines: BTreeMap::new(),
             default_engine: None,
-            api_key: None,
+            middlewares: Vec::new(),
         }
     }
 
@@ -67,11 +69,6 @@ impl ServerBuilder {
         self
     }
 
-    pub fn with_api_key(mut self, api_key: String) -> Self {
-        self.api_key = Some(api_key);
-        self
-    }
-
     pub fn with_engines(
         mut self,
         mut engines: BTreeMap<Principal, Engine>,
@@ -86,6 +83,67 @@ impl ServerBuilder {
         self
     }
 
+    /// Register a router middleware.
+    ///
+    /// This is the low-level API. The middleware will be applied to the internal
+    /// axum `Router` (typically via `router.layer(...)`). Middlewares are applied
+    /// in the order they are added.
+    ///
+    /// More details: https://docs.rs/axum/latest/axum/middleware/index.html#ordering
+    ///
+    /// If you want a middleware that looks like `axum::middleware::from_fn`
+    /// (i.e. can operate on `(req, next)`), prefer [`with_request_middleware`].
+    ///
+    /// Example:
+    /// ```ignore
+    /// let server = ServerBuilder::new()
+    ///   .with_middleware(|router| {
+    ///     router.layer(axum::middleware::from_fn(|req, next| async move {
+    ///       // custom auth / param checks here
+    ///       next.run(req).await
+    ///     }))
+    ///   });
+    /// ```
+    pub fn with_middleware<M>(mut self, middleware: M) -> Self
+    where
+        M: HttpMiddleware,
+    {
+        self.middlewares.push(Arc::new(middleware));
+        self
+    }
+
+    /// Register a request middleware like `axum::middleware::from_fn`.
+    ///
+    /// The middleware function runs for every incoming request, and can decide
+    /// to short-circuit with a response or call `next.run(req)`.
+    ///
+    /// Example:
+    /// ```ignore
+    /// use axum::http::StatusCode;
+    /// use axum::response::IntoResponse;
+    ///
+    /// let server = ServerBuilder::new()
+    ///   .with_request_middleware(|req, next| async move {
+    ///     // custom auth / param checks here
+    ///     if req.headers().get("x-allow").is_none() {
+    ///       return (StatusCode::UNAUTHORIZED, "missing x-allow").into_response();
+    ///     }
+    ///
+    ///     next.run(req).await
+    ///   });
+    /// ```
+    pub fn with_request_middleware<F, Fut>(self, f: F) -> Self
+    where
+        F: Fn(axum::extract::Request, axum::middleware::Next) -> Fut
+            + Clone
+            + Send
+            + Sync
+            + 'static,
+        Fut: Future<Output = axum::response::Response> + Send + 'static,
+    {
+        self.with_middleware(middleware::RequestFnMiddleware::new(f))
+    }
+
     pub async fn serve(
         self,
         signal: impl Future<Output = ()> + Send + 'static,
@@ -105,18 +163,26 @@ impl ServerBuilder {
             engines: Arc::new(self.engines),
             default_engine,
             start_time_ms: unix_ms(),
-            api_key: self.api_key,
         };
-        let app = Router::new()
+
+        // Build a router that is still "missing" an `AppState`.
+        // We'll provide the state at the end (after applying middlewares) so we
+        // end up with a `Router<()>` that can be passed to `axum::serve`.
+        let mut app: Router<AppState> = Router::new()
             .route("/", routing::get(get_information))
             .route("/.well-known/information", routing::get(get_information))
             .route("/.well-known/agents", routing::get(get_information))
             .route(
                 "/.well-known/agents/{id}",
                 routing::get(get_engine_information),
             )
-            .route("/{*id}", routing::post(anda_engine))
-            .with_state(state);
+            .route("/{*id}", routing::post(anda_engine));
+
+        for middleware in &self.middlewares {
+            app = middleware.apply(app);
+        }
+
+        let app = app.with_state(state);
 
         let addr: SocketAddr = self.addr.parse()?;
         let listener = create_reuse_port_listener(addr).await?;
diff --git a/anda_engine_server/src/middleware.rs b/anda_engine_server/src/middleware.rs
@@ -0,0 +1,103 @@
+use std::future::Future;
+
+use axum::{
+    Router,
+    extract::Request,
+    http::StatusCode,
+    middleware::Next,
+    response::{IntoResponse, Response},
+};
+
+use crate::handler::AppState;
+
+/// Object-safe middleware trait for applying HTTP middleware to the server `Router`.
+///
+/// This is intentionally type-erased so callers can register arbitrary axum/tower
+/// middleware without turning `ServerBuilder` into a giant generic type.
+pub trait HttpMiddleware: Send + Sync + 'static {
+    fn apply(&self, router: Router<AppState>) -> Router<AppState>;
+}
+
+impl<F> HttpMiddleware for F
+where
+    F: Fn(Router<AppState>) -> Router<AppState> + Send + Sync + 'static,
+{
+    fn apply(&self, router: Router<AppState>) -> Router<AppState> {
+        (self)(router)
+    }
+}
+
+/// Middleware built from a function like axum::middleware::from_fn.
+#[derive(Clone)]
+pub struct RequestFnMiddleware<F> {
+    f: F,
+}
+
+impl<F> RequestFnMiddleware<F> {
+    pub fn new(f: F) -> Self {
+        Self { f }
+    }
+}
+
+impl<F, Fut> HttpMiddleware for RequestFnMiddleware<F>
+where
+    F: Fn(Request, Next) -> Fut + Clone + Send + Sync + 'static,
+    Fut: Future<Output = Response> + Send + 'static,
+{
+    fn apply(&self, router: Router<AppState>) -> Router<AppState> {
+        router.layer(axum::middleware::from_fn(self.f.clone()))
+    }
+}
+
+/// A simple API key middleware that validates `x-api-key` on every request.
+///
+/// Use `exempt_path` to allow unauthenticated endpoints (e.g. health/info routes).
+#[derive(Clone, Default)]
+pub struct ApiKeyMiddleware {
+    expected_key: String,
+    exempt_paths: Vec<String>,
+}
+
+impl ApiKeyMiddleware {
+    pub fn new(expected_key: impl Into<String>) -> Self {
+        Self {
+            expected_key: expected_key.into(),
+            exempt_paths: Vec::new(),
+        }
+    }
+
+    pub fn exempt_path(mut self, path: impl Into<String>) -> Self {
+        self.exempt_paths.push(path.into());
+        self
+    }
+}
+
+impl HttpMiddleware for ApiKeyMiddleware {
+    fn apply(&self, router: Router<AppState>) -> Router<AppState> {
+        let expected_key = self.expected_key.clone();
+        let exempt_paths = self.exempt_paths.clone();
+
+        router.layer(axum::middleware::from_fn(
+            move |req: Request, next: Next| {
+                let expected_key = expected_key.clone();
+                let exempt_paths = exempt_paths.clone();
+
+                async move {
+                    let path = req.uri().path();
+                    if exempt_paths.iter().any(|p| p == path) {
+                        return next.run(req).await;
+                    }
+
+                    match req.headers().get("x-api-key") {
+                        Some(provided_key) if provided_key == &expected_key => next.run(req).await,
+                        _ => (
+                            StatusCode::UNAUTHORIZED,
+                            "missing or invalid x-api-key in headers",
+                        )
+                            .into_response(),
+                    }
+                }
+            },
+        ))
+    }
+}
