agentguard_server/
auth_layer.rs

1//! Authentication middleware for the AuthZEN HTTP server.
2//!
3//! Modes:
4//! - [`AuthConfig::Disabled`] — no auth. Suitable for development or
5//!   loopback-only deployments. The server refuses to start with
6//!   `Disabled` auth on a non-loopback bind unless
7//!   `AGENTGUARD_ALLOW_LOOPBACK_BYPASS=1` is set.
8//! - [`AuthConfig::ApiKey`] — `Authorization: Bearer <raw>` against
9//!   the configured `ApiKeyStore`. Argon2id is the verification path;
10//!   the cost (~150 ms) is acceptable on the auth path.
11//!
12//! Health probes (`/healthz`, `/readyz`) are always unauthenticated
13//! so Kubernetes can poll them without a credential.
14
15use crate::authzen::AppState;
16use crate::listener::AuthConfig;
17use agentguard_auth::{ApiKeyIdentity, ApiKeyStore};
18use axum::{
19    extract::{Request, State},
20    http::StatusCode,
21    middleware::Next,
22    response::{IntoResponse, Response},
23};
24use std::sync::Arc;
25use tokio::sync::Semaphore;
26
27/// Argon2 verification uses substantial CPU and memory. Keep its blocking
28/// work off async workers and cap concurrent verifications independently of
29/// request concurrency (two 64 MiB hashes fit the default PDP memory budget).
30const MAX_CONCURRENT_API_KEY_VERIFICATIONS: usize = 2;
31
32/// Authenticated identity and tenant asserted by a verified, bound API key.
33#[derive(Debug, Clone)]
34pub struct AuthenticatedIdentity {
35    pub identity: ApiKeyIdentity,
36    /// Stable API-key id for audit correlation; never the raw bearer secret.
37    pub credential_id: String,
38    /// True only for credentials granted the explicit `authorize:any` scope.
39    pub can_act_as: bool,
40}
41
42#[derive(Debug, Clone, Copy, PartialEq, Eq)]
43pub enum AuthenticationFailure {
44    Unauthenticated,
45    Forbidden,
46    Unavailable,
47}
48
49/// What the auth layer needs to validate requests. Built once at
50/// startup and shared across Axum workers.
51#[derive(Clone)]
52pub enum AuthLayer {
53    Disabled,
54    ApiKey(Arc<ApiKeyStore>),
55    ReloadableApiKey {
56        store: Arc<ApiKeyStore>,
57        path: std::path::PathBuf,
58    },
59}
60
61impl AuthLayer {
62    /// Build from the configured mode.
63    pub fn from_config(cfg: &AuthConfig, allow_loopback_bypass: bool) -> Result<Self, String> {
64        let layer = match cfg {
65            AuthConfig::Disabled => AuthLayer::Disabled,
66            AuthConfig::ApiKey { path } => {
67                let store = ApiKeyStore::load_from_file(path)
68                    .map_err(|e| format!("load api-key store {:?}: {}", path, e))?;
69                AuthLayer::ReloadableApiKey {
70                    store: Arc::new(store),
71                    path: path.clone(),
72                }
73            }
74        };
75        if allow_loopback_bypass && matches!(layer, AuthLayer::Disabled) {
76            tracing::warn!(
77                "AGENTGUARD_ALLOW_LOOPBACK_BYPASS=1: auth disabled; \
78                 operator assumes the listener is loopback-bound"
79            );
80        }
81        Ok(layer)
82    }
83
84    /// Return the key-store source used by the standalone lifecycle so it
85    /// can watch and atomically reload operator rotations/revocations.
86    pub fn reloadable_key_store(&self) -> Option<(std::path::PathBuf, Arc<ApiKeyStore>)> {
87        match self {
88            Self::ReloadableApiKey { store, path } => Some((path.clone(), store.clone())),
89            Self::Disabled | Self::ApiKey(_) => None,
90        }
91    }
92
93    /// Authenticate non-HTTP transports using the same scope and identity
94    /// requirements as the AuthZEN HTTP endpoints.
95    pub fn authenticate_bearer(
96        &self,
97        authorization: Option<&str>,
98        required_scope: &str,
99        require_identity: bool,
100    ) -> Result<Option<AuthenticatedIdentity>, AuthenticationFailure> {
101        match self {
102            Self::Disabled => Ok(None),
103            Self::ApiKey(store) | Self::ReloadableApiKey { store, .. } => {
104                let token =
105                    bearer_token(authorization).ok_or(AuthenticationFailure::Unauthenticated)?;
106                let key = store
107                    .verify(token)
108                    .map_err(|_| AuthenticationFailure::Unauthenticated)?;
109                let can_act_as = required_scope == "authorize" && has_act_as_scope(&key.scopes);
110                if !key.has_scope(required_scope) && !can_act_as {
111                    return Err(AuthenticationFailure::Forbidden);
112                }
113                let identity = key.identity.ok_or(if require_identity {
114                    AuthenticationFailure::Forbidden
115                } else {
116                    AuthenticationFailure::Unauthenticated
117                });
118                match identity {
119                    Ok(identity) => Ok(Some(AuthenticatedIdentity {
120                        identity,
121                        credential_id: key.id,
122                        can_act_as,
123                    })),
124                    Err(AuthenticationFailure::Unauthenticated) if !require_identity => Ok(None),
125                    Err(err) => Err(err),
126                }
127            }
128        }
129    }
130
131    /// Authenticate without running the memory-hard hash on a Tokio worker.
132    /// Saturation fails fast so hostile traffic cannot build an unbounded
133    /// queue of expensive password-hash jobs.
134    pub async fn authenticate_bearer_async(
135        &self,
136        authorization: Option<&str>,
137        required_scope: &str,
138        require_identity: bool,
139    ) -> Result<Option<AuthenticatedIdentity>, AuthenticationFailure> {
140        self.authenticate_bearer_async_with_slots(
141            authorization,
142            required_scope,
143            require_identity,
144            api_key_verification_slots(),
145        )
146        .await
147    }
148
149    async fn authenticate_bearer_async_with_slots(
150        &self,
151        authorization: Option<&str>,
152        required_scope: &str,
153        require_identity: bool,
154        slots: Arc<Semaphore>,
155    ) -> Result<Option<AuthenticatedIdentity>, AuthenticationFailure> {
156        if matches!(self, Self::Disabled) {
157            return self.authenticate_bearer(authorization, required_scope, require_identity);
158        }
159        // Reject absent or structurally malformed credentials before taking
160        // a scarce Argon2 slot. This keeps anonymous probes cheap under load.
161        let token = bearer_token(authorization).ok_or(AuthenticationFailure::Unauthenticated)?;
162        let mut parts = token.split(':');
163        if parts.next().is_none_or(str::is_empty)
164            || parts.next().is_none_or(str::is_empty)
165            || parts.next().is_none_or(str::is_empty)
166            || parts.next().is_some()
167        {
168            return Err(AuthenticationFailure::Unauthenticated);
169        }
170        let permit = slots
171            .try_acquire_owned()
172            .map_err(|_| AuthenticationFailure::Unavailable)?;
173        let layer = self.clone();
174        let authorization = authorization.map(str::to_owned);
175        let required_scope = required_scope.to_owned();
176        tokio::task::spawn_blocking(move || {
177            let _permit = permit;
178            layer.authenticate_bearer(authorization.as_deref(), &required_scope, require_identity)
179        })
180        .await
181        .map_err(|_| AuthenticationFailure::Unavailable)?
182    }
183}
184
185/// Parse the HTTP authentication scheme without imposing case sensitivity on
186/// the scheme name. Credentials remain opaque and must not contain whitespace.
187fn bearer_token(authorization: Option<&str>) -> Option<&str> {
188    let value = authorization?;
189    let separator = value.find(' ')?;
190    let (scheme, rest) = value.split_at(separator);
191    if !scheme.eq_ignore_ascii_case("Bearer") {
192        return None;
193    }
194    let token = rest.trim_start_matches(' ');
195    if token.is_empty() || token.bytes().any(|byte| byte.is_ascii_whitespace()) {
196        return None;
197    }
198    Some(token)
199}
200
201fn has_act_as_scope(scopes: &[String]) -> bool {
202    scopes.iter().any(|scope| scope == "authorize:any")
203}
204
205fn api_key_verification_slots() -> Arc<Semaphore> {
206    static SLOTS: std::sync::OnceLock<Arc<Semaphore>> = std::sync::OnceLock::new();
207    SLOTS
208        .get_or_init(|| Arc::new(Semaphore::new(MAX_CONCURRENT_API_KEY_VERIFICATIONS)))
209        .clone()
210}
211
212/// Middleware function. Wire via `axum::middleware::from_fn_with_state`
213/// because the router's state is needed for the auth layer reference.
214pub async fn auth_layer_fn(State(state): State<AppState>, req: Request, next: Next) -> Response {
215    let path = req.uri().path();
216    if path == "/healthz" || path == "/readyz" {
217        return next.run(req).await;
218    }
219    let method = req.method();
220    let (required_scope, require_identity) = match (method.as_str(), path) {
221        ("POST", "/access/v1/evaluation" | "/access/v1/evaluations") => ("authorize", true),
222        ("GET", "/metrics") => ("metrics:read", false),
223        _ => ("", false),
224    };
225    match &state.auth {
226        AuthLayer::Disabled => next.run(req).await,
227        AuthLayer::ApiKey(_) | AuthLayer::ReloadableApiKey { .. } => {
228            let identity = match state
229                .auth
230                .authenticate_bearer_async(
231                    req.headers()
232                        .get(axum::http::header::AUTHORIZATION)
233                        .and_then(|value| value.to_str().ok()),
234                    required_scope,
235                    require_identity,
236                )
237                .await
238            {
239                Ok(identity) => identity,
240                Err(AuthenticationFailure::Unauthenticated) => return unauthorized(),
241                Err(AuthenticationFailure::Forbidden) => return forbidden(),
242                Err(AuthenticationFailure::Unavailable) => {
243                    return (
244                        StatusCode::SERVICE_UNAVAILABLE,
245                        "authentication capacity exhausted\n",
246                    )
247                        .into_response()
248                }
249            };
250            let mut req = req;
251            if let Some(identity) = identity {
252                req.extensions_mut().insert(identity);
253            }
254            next.run(req).await
255        }
256    }
257}
258
259fn unauthorized() -> Response {
260    (StatusCode::UNAUTHORIZED, "unauthorized\n").into_response()
261}
262
263fn forbidden() -> Response {
264    (StatusCode::FORBIDDEN, "forbidden\n").into_response()
265}
266
267#[cfg(test)]
268mod async_auth_tests {
269    use super::{ApiKeyStore, AuthLayer, AuthenticationFailure};
270    use std::sync::Arc;
271    use tokio::sync::Semaphore;
272
273    #[test]
274    fn bearer_scheme_is_case_insensitive_and_allows_multiple_spaces() {
275        for header in [
276            "Bearer ag_test:key-id:secret",
277            "bearer ag_test:key-id:secret",
278            "bEaReR  ag_test:key-id:secret",
279        ] {
280            assert_eq!(
281                super::bearer_token(Some(header)),
282                Some("ag_test:key-id:secret")
283            );
284        }
285    }
286
287    #[test]
288    fn malformed_bearer_headers_are_rejected() {
289        for header in [
290            "Basic abc",
291            "Bearer",
292            "Bearer ",
293            "Bearer\tag_test:key-id:secret",
294            "Bearer ag_test:key-id: secret",
295        ] {
296            assert_eq!(super::bearer_token(Some(header)), None, "{header:?}");
297        }
298        assert_eq!(super::bearer_token(None), None);
299    }
300
301    #[test]
302    fn wildcard_authorization_does_not_implicitly_grant_act_as() {
303        assert!(!super::has_act_as_scope(&["*".into()]));
304        assert!(!super::has_act_as_scope(&["authorize".into()]));
305        assert!(super::has_act_as_scope(&["authorize:any".into()]));
306    }
307
308    #[tokio::test]
309    async fn saturated_verification_capacity_fails_fast() {
310        let auth = AuthLayer::ApiKey(Arc::new(ApiKeyStore::new()));
311        let result = auth
312            .authenticate_bearer_async_with_slots(
313                Some("Bearer ag_test:key-id:secret"),
314                "authorize",
315                true,
316                Arc::new(Semaphore::new(0)),
317            )
318            .await;
319        assert!(matches!(result, Err(AuthenticationFailure::Unavailable)));
320    }
321
322    #[tokio::test]
323    async fn disabled_auth_stays_nonblocking_and_does_not_need_a_slot() {
324        let result = AuthLayer::Disabled
325            .authenticate_bearer_async_with_slots(
326                None,
327                "authorize",
328                true,
329                Arc::new(Semaphore::new(0)),
330            )
331            .await;
332        assert!(matches!(result, Ok(None)));
333    }
334
335    #[tokio::test]
336    async fn missing_credentials_do_not_consume_verification_capacity() {
337        let auth = AuthLayer::ApiKey(Arc::new(ApiKeyStore::new()));
338        let result = auth
339            .authenticate_bearer_async_with_slots(
340                None,
341                "authorize",
342                true,
343                Arc::new(Semaphore::new(0)),
344            )
345            .await;
346        assert!(matches!(
347            result,
348            Err(AuthenticationFailure::Unauthenticated)
349        ));
350    }
351
352    #[tokio::test]
353    async fn argon_verification_does_not_block_async_workers() {
354        use agentguard_auth::ApiKeyIdentity;
355        use std::sync::atomic::{AtomicUsize, Ordering};
356        use std::time::Duration;
357
358        let store = ApiKeyStore::new();
359        let identity = ApiKeyIdentity::new("User", "alice", None).unwrap();
360        let (_, raw) = store
361            .create_bound("ag_test", vec!["authorize".into()], None, identity)
362            .unwrap();
363        let auth = AuthLayer::ApiKey(Arc::new(store));
364        let ticks = Arc::new(AtomicUsize::new(0));
365        let tick_count = ticks.clone();
366        let ticker = tokio::spawn(async move {
367            let mut interval = tokio::time::interval(Duration::from_millis(10));
368            loop {
369                interval.tick().await;
370                tick_count.fetch_add(1, Ordering::Relaxed);
371            }
372        });
373
374        let result = auth
375            .authenticate_bearer_async_with_slots(
376                Some(&format!("Bearer {raw}")),
377                "authorize",
378                true,
379                Arc::new(Semaphore::new(1)),
380            )
381            .await;
382        ticker.abort();
383
384        assert!(matches!(result, Ok(Some(identity)) if !identity.can_act_as));
385        assert!(
386            ticks.load(Ordering::Relaxed) >= 2,
387            "the async runtime should keep making progress during Argon2 verification"
388        );
389    }
390
391    #[tokio::test]
392    async fn act_as_scope_is_explicit_and_keeps_a_bound_service_identity() {
393        use agentguard_auth::ApiKeyIdentity;
394
395        let store = ApiKeyStore::new();
396        let identity = ApiKeyIdentity::new("Agent", "console-service", None).unwrap();
397        let (key, raw) = store
398            .create_bound("ag_test", vec!["authorize:any".into()], None, identity)
399            .unwrap();
400        let auth = AuthLayer::ApiKey(Arc::new(store));
401        let result = auth
402            .authenticate_bearer_async_with_slots(
403                Some(&format!("Bearer {raw}")),
404                "authorize",
405                true,
406                Arc::new(Semaphore::new(1)),
407            )
408            .await
409            .unwrap()
410            .unwrap();
411        assert_eq!(result.identity.subject_id, "console-service");
412        assert_eq!(result.credential_id, key.id);
413        assert!(
414            result.can_act_as,
415            "only authorize:any opts into acting as a request subject"
416        );
417    }
418}
agentguard_server/auth_layer.rs

agentguard_server/
auth_layer.rs
