a2a_client/
auth.rs

1// Copyright AGNTCY Contributors (https://github.com/agntcy)
2// Copyright A2A Contributors (https://github.com/a2aproject)
3// SPDX-License-Identifier: Apache-2.0
4use a2a::A2AError;
5use async_trait::async_trait;
6use std::collections::HashMap;
7use std::sync::RwLock;
8
9use crate::middleware::CallInterceptor;
10use crate::transport::ServiceParams;
11
12/// Trait for providing credentials for authentication.
13#[async_trait]
14pub trait CredentialsStore: Send + Sync {
15    /// Get credentials for the given scheme name.
16    async fn get(&self, scheme: &str) -> Option<String>;
17}
18
19/// Simple in-memory credentials store.
20pub struct InMemoryCredentialsStore {
21    credentials: RwLock<HashMap<String, String>>,
22}
23
24impl InMemoryCredentialsStore {
25    pub fn new() -> Self {
26        InMemoryCredentialsStore {
27            credentials: RwLock::new(HashMap::new()),
28        }
29    }
30
31    pub fn set(&self, scheme: &str, credential: &str) {
32        self.credentials
33            .write()
34            .unwrap()
35            .insert(scheme.to_string(), credential.to_string());
36    }
37}
38
39impl Default for InMemoryCredentialsStore {
40    fn default() -> Self {
41        Self::new()
42    }
43}
44
45#[async_trait]
46impl CredentialsStore for InMemoryCredentialsStore {
47    async fn get(&self, scheme: &str) -> Option<String> {
48        self.credentials.read().unwrap().get(scheme).cloned()
49    }
50}
51
52/// Interceptor that injects authentication credentials into requests.
53pub struct AuthInterceptor {
54    header_name: String,
55    header_value: String,
56}
57
58impl AuthInterceptor {
59    /// Create an interceptor that adds a Bearer token.
60    pub fn bearer(token: impl Into<String>) -> Self {
61        AuthInterceptor {
62            header_name: "Authorization".to_string(),
63            header_value: format!("Bearer {}", token.into()),
64        }
65    }
66
67    /// Create an interceptor with a custom header.
68    pub fn custom(header_name: impl Into<String>, header_value: impl Into<String>) -> Self {
69        AuthInterceptor {
70            header_name: header_name.into(),
71            header_value: header_value.into(),
72        }
73    }
74}
75
76#[async_trait]
77impl CallInterceptor for AuthInterceptor {
78    async fn before(&self, _method: &str, params: &mut ServiceParams) -> Result<(), A2AError> {
79        params
80            .entry(self.header_name.clone())
81            .or_default()
82            .push(self.header_value.clone());
83        Ok(())
84    }
85}
86
87#[cfg(test)]
88mod tests {
89    use super::*;
90
91    #[test]
92    fn test_in_memory_credentials_store_new() {
93        let store = InMemoryCredentialsStore::new();
94        let creds = store.credentials.read().unwrap();
95        assert!(creds.is_empty());
96    }
97
98    #[test]
99    fn test_in_memory_credentials_store_default() {
100        let store = InMemoryCredentialsStore::default();
101        let creds = store.credentials.read().unwrap();
102        assert!(creds.is_empty());
103    }
104
105    #[test]
106    fn test_in_memory_credentials_store_set_get() {
107        let store = InMemoryCredentialsStore::new();
108        store.set("bearer", "token123");
109        let creds = store.credentials.read().unwrap();
110        assert_eq!(creds.get("bearer").unwrap(), "token123");
111    }
112
113    #[tokio::test]
114    async fn test_credentials_store_get() {
115        let store = InMemoryCredentialsStore::new();
116        store.set("api-key", "secret");
117        assert_eq!(store.get("api-key").await, Some("secret".to_string()));
118        assert_eq!(store.get("nonexistent").await, None);
119    }
120
121    #[tokio::test]
122    async fn test_auth_interceptor_bearer() {
123        let interceptor = AuthInterceptor::bearer("mytoken");
124        let mut params = ServiceParams::new();
125        interceptor.before("test", &mut params).await.unwrap();
126        assert_eq!(
127            params.get("Authorization").unwrap(),
128            &vec!["Bearer mytoken".to_string()]
129        );
130    }
131
132    #[tokio::test]
133    async fn test_auth_interceptor_custom() {
134        let interceptor = AuthInterceptor::custom("X-API-Key", "key123");
135        let mut params = ServiceParams::new();
136        interceptor.before("test", &mut params).await.unwrap();
137        assert_eq!(
138            params.get("X-API-Key").unwrap(),
139            &vec!["key123".to_string()]
140        );
141    }
142}
a2a_client/auth.rs

a2a_client/
auth.rs
