Skip to main content

adk_server/
auth_bridge.rs

1//! Auth middleware bridge for flowing authenticated identity into agent execution.
2//!
3//! This module defines the [`RequestContextExtractor`] trait that server operators
4//! implement to extract identity from HTTP requests, and the [`RequestContextError`]
5//! enum for extraction failures.
6//!
7//! The extracted [`RequestContext`] (re-exported from `adk-core`) carries user_id,
8//! scopes, and metadata into the `InvocationContext`, making scopes available
9//! to tools via `ToolContext::user_scopes()`.
10//!
11//! # Example
12//!
13//! ```rust,ignore
14//! use adk_server::auth_bridge::{RequestContextExtractor, RequestContextError};
15//! use adk_core::RequestContext;
16//! use async_trait::async_trait;
17//!
18//! struct MyExtractor;
19//!
20//! #[async_trait]
21//! impl RequestContextExtractor for MyExtractor {
22//!     async fn extract(
23//!         &self,
24//!         parts: &axum::http::request::Parts,
25//!     ) -> Result<RequestContext, RequestContextError> {
26//!         let auth = parts.headers
27//!             .get("authorization")
28//!             .and_then(|v| v.to_str().ok())
29//!             .ok_or(RequestContextError::MissingAuth)?;
30//!         // ... validate token, build RequestContext ...
31//!         # todo!()
32//!     }
33//! }
34//! ```
35
36pub use adk_core::RequestContext;
37use async_trait::async_trait;
38use axum::extract::FromRequestParts;
39use axum::http::request::Parts;
40use std::convert::Infallible;
41
42/// Extracts authenticated identity from HTTP request headers.
43///
44/// Implementations typically parse a Bearer token from the `Authorization`
45/// header, validate it, and map claims to a [`RequestContext`].
46#[async_trait]
47pub trait RequestContextExtractor: Send + Sync {
48    /// Extract identity from the request parts (headers, URI, etc.).
49    async fn extract(
50        &self,
51        parts: &axum::http::request::Parts,
52    ) -> Result<RequestContext, RequestContextError>;
53}
54
55/// Errors that can occur during request context extraction.
56#[derive(Debug, thiserror::Error)]
57pub enum RequestContextError {
58    /// The `Authorization` header is missing from the request.
59    #[error("missing authorization header")]
60    MissingAuth,
61    /// The token was present but failed validation.
62    #[error("invalid token: {0}")]
63    InvalidToken(String),
64    /// An internal error occurred during extraction.
65    #[error("extraction failed: {0}")]
66    ExtractionFailed(String),
67}
68
69/// The authenticated principal of a request, read from the request extensions.
70///
71/// The server's authentication layer stores the extracted identity as an
72/// `Option<RequestContext>` extension; middleware written outside this crate
73/// can also store a bare [`RequestContext`]. This extractor accepts either and
74/// never rejects: it yields `None` when no authentication ran, which is how a
75/// handler tells an unauthenticated deployment from an authenticated caller.
76///
77/// # Example
78///
79/// ```rust
80/// use adk_server::auth_bridge::AuthenticatedCaller;
81///
82/// async fn whoami(AuthenticatedCaller(caller): AuthenticatedCaller) -> String {
83///     caller.map_or_else(|| "anonymous".to_string(), |context| context.user_id)
84/// }
85///
86/// let _app: axum::Router = axum::Router::new().route("/whoami", axum::routing::get(whoami));
87/// ```
88#[derive(Debug, Clone, Default)]
89pub struct AuthenticatedCaller(pub Option<RequestContext>);
90
91impl AuthenticatedCaller {
92    /// Returns the authenticated user ID, or `None` when no authentication ran.
93    pub fn user_id(&self) -> Option<&str> {
94        self.0.as_ref().map(|context| context.user_id.as_str())
95    }
96}
97
98impl<S: Send + Sync> FromRequestParts<S> for AuthenticatedCaller {
99    type Rejection = Infallible;
100
101    async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result<Self, Self::Rejection> {
102        let context = parts
103            .extensions
104            .get::<Option<RequestContext>>()
105            .cloned()
106            .flatten()
107            .or_else(|| parts.extensions.get::<RequestContext>().cloned());
108        Ok(Self(context))
109    }
110}
111
112#[cfg(test)]
113mod tests {
114    use super::*;
115
116    fn context(user_id: &str) -> RequestContext {
117        RequestContext {
118            user_id: user_id.to_string(),
119            scopes: vec![],
120            metadata: Default::default(),
121        }
122    }
123
124    async fn extract(parts: &mut Parts) -> Option<String> {
125        let caller = AuthenticatedCaller::from_request_parts(parts, &()).await.unwrap();
126        caller.user_id().map(str::to_string)
127    }
128
129    #[tokio::test]
130    async fn reads_the_auth_layer_extension() {
131        let (mut parts, ()) = axum::http::Request::new(()).into_parts();
132        parts.extensions.insert::<Option<RequestContext>>(Some(context("alice")));
133        assert_eq!(extract(&mut parts).await, Some("alice".to_string()));
134    }
135
136    #[tokio::test]
137    async fn reads_a_bare_request_context_extension() {
138        let (mut parts, ()) = axum::http::Request::new(()).into_parts();
139        parts.extensions.insert(context("bob"));
140        assert_eq!(extract(&mut parts).await, Some("bob".to_string()));
141    }
142
143    #[tokio::test]
144    async fn is_none_without_authentication() {
145        let (mut parts, ()) = axum::http::Request::new(()).into_parts();
146        assert_eq!(extract(&mut parts).await, None);
147        parts.extensions.insert::<Option<RequestContext>>(None);
148        assert_eq!(extract(&mut parts).await, None);
149    }
150}