Skip to main content

brainwires_agent_network/middleware/
auth.rs

1use async_trait::async_trait;
2use brainwires_mcp::{JsonRpcError, JsonRpcRequest};
3
4use super::{Middleware, MiddlewareResult};
5use crate::connection::RequestContext;
6
7/// Token-based authentication middleware.
8pub struct AuthMiddleware {
9    token: String,
10}
11
12impl AuthMiddleware {
13    /// Create a new auth middleware with the given token.
14    pub fn new(token: impl Into<String>) -> Self {
15        Self {
16            token: token.into(),
17        }
18    }
19}
20
21#[async_trait]
22impl Middleware for AuthMiddleware {
23    async fn process_request(
24        &self,
25        request: &JsonRpcRequest,
26        ctx: &mut RequestContext,
27    ) -> MiddlewareResult {
28        // Skip auth for initialize - clients haven't authenticated yet
29        if request.method == "initialize" {
30            return MiddlewareResult::Continue;
31        }
32
33        // Check for token in metadata (set during initialize)
34        if let Some(serde_json::Value::String(token)) = ctx.get_metadata("auth_token")
35            && token == &self.token
36        {
37            return MiddlewareResult::Continue;
38        }
39
40        // Check params for auth token
41        if let Some(params) = &request.params
42            && let Some(token) = params.get("_auth_token").and_then(|v| v.as_str())
43            && token == self.token
44        {
45            ctx.set_metadata(
46                "auth_token".to_string(),
47                serde_json::Value::String(token.to_string()),
48            );
49            return MiddlewareResult::Continue;
50        }
51
52        MiddlewareResult::Reject(JsonRpcError {
53            code: -32003,
54            message: "Unauthorized: invalid or missing auth token".to_string(),
55            data: None,
56        })
57    }
58}