1use async_trait::async_trait;
2use serde::Deserialize;
3use serde::de::Deserializer;
4use std::time::{Duration, Instant};
5use tokio::sync::{Mutex, RwLock};
6use zeroize::Zeroizing;
7
8use crate::http_client::build_ssrf_pinned_client;
9use crate::types::AuthError;
10
11fn deserialize_zeroizing_string<'de, D>(deserializer: D) -> Result<Zeroizing<String>, D::Error>
12where
13 D: Deserializer<'de>,
14{
15 let s = String::deserialize(deserializer)?;
16 Ok(Zeroizing::new(s))
17}
18
19const DEFAULT_SKEW: Duration = Duration::from_secs(30);
20
21#[async_trait]
22pub trait TokenProvider: Send + Sync + std::fmt::Debug {
23 async fn get_token(&self) -> Result<String, AuthError>;
24}
25
26#[derive(Debug, Deserialize)]
27struct TokenResponse {
28 #[serde(deserialize_with = "deserialize_zeroizing_string")]
29 access_token: Zeroizing<String>,
30 #[allow(dead_code)]
31 token_type: String,
32 expires_in: u64,
33}
34
35struct CachedToken {
36 access_token: Zeroizing<String>,
37 #[allow(dead_code)]
38 expires_at: Instant,
39 refresh_at: Instant,
40}
41
42impl CachedToken {
43 fn new(access_token: Zeroizing<String>, expires_in: Duration, skew: Duration) -> Self {
44 let expires_at = Instant::now() + expires_in;
45 Self {
46 access_token,
47 refresh_at: expires_at.checked_sub(skew).unwrap_or(expires_at),
48 expires_at,
49 }
50 }
51
52 fn is_usable(&self) -> bool {
53 Instant::now() < self.refresh_at
54 }
55}
56
57pub struct ClientCredentialsProvider {
58 token_endpoint: String,
59 client_id: String,
60 client_secret: Zeroizing<String>,
61 scope: Option<String>,
62 audience: Option<Vec<String>>,
63 cache: RwLock<Option<CachedToken>>,
64 refresh_lock: Mutex<()>,
65 http: reqwest::Client,
66}
67
68impl std::fmt::Debug for ClientCredentialsProvider {
69 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
70 f.debug_struct("ClientCredentialsProvider")
71 .field("token_endpoint", &self.token_endpoint)
72 .field("client_id", &self.client_id)
73 .field("scope", &self.scope)
74 .field("audience", &self.audience)
75 .finish_non_exhaustive()
76 }
77}
78
79impl ClientCredentialsProvider {
80 pub async fn new(
87 token_endpoint: String,
88 client_id: String,
89 client_secret: String,
90 scope: Option<String>,
91 audience: Option<Vec<String>>,
92 ) -> Result<Self, AuthError> {
93 let http = build_ssrf_pinned_client(
94 &token_endpoint,
95 "OAuth2 token endpoint",
96 Duration::from_secs(10),
97 Duration::from_secs(30),
98 )
99 .await?;
100 Ok(Self {
101 token_endpoint,
102 client_id,
103 client_secret: Zeroizing::new(client_secret),
104 scope,
105 audience,
106 cache: RwLock::new(None),
107 refresh_lock: Mutex::new(()),
108 http,
109 })
110 }
111
112 #[doc(hidden)]
117 pub fn new_unchecked_for_test(
118 token_endpoint: String,
119 client_id: String,
120 client_secret: String,
121 scope: Option<String>,
122 audience: Option<Vec<String>>,
123 http: reqwest::Client,
124 ) -> Self {
125 Self {
126 token_endpoint,
127 client_id,
128 client_secret: Zeroizing::new(client_secret),
129 scope,
130 audience,
131 cache: RwLock::new(None),
132 refresh_lock: Mutex::new(()),
133 http,
134 }
135 }
136
137 async fn fetch_token(&self) -> Result<CachedToken, AuthError> {
138 let secret = self.client_secret.as_str();
139 let mut params: Vec<(&str, &str)> = vec![
140 ("grant_type", "client_credentials"),
141 ("client_id", &self.client_id),
142 ("client_secret", secret),
143 ];
144 if let Some(ref scope) = self.scope {
145 params.push(("scope", scope));
146 }
147 if let Some(ref audience) = self.audience {
148 for aud in audience {
149 params.push(("resource", aud));
150 }
151 }
152
153 let resp = self
154 .http
155 .post(&self.token_endpoint)
156 .form(¶ms)
157 .send()
158 .await
159 .map_err(|e| AuthError::ProviderUnavailable(format!("OAuth2 request failed: {e}")))?;
160
161 if !resp.status().is_success() {
162 let status = resp.status();
163 let body = resp.text().await.unwrap_or_default();
164 let sanitized = if body.len() > 128 {
165 format!("{}...(truncated)", &body[..128])
166 } else {
167 body
168 };
169 let message = format!("token endpoint returned {status}: {sanitized}"); return Err(AuthError::ProviderUnavailable(message));
171 }
172
173 let token_resp: TokenResponse = resp
174 .json()
175 .await
176 .map_err(|e| AuthError::ProviderUnavailable(format!("invalid OAuth2 response: {e}")))?;
177
178 Ok(CachedToken::new(
179 token_resp.access_token,
180 Duration::from_secs(token_resp.expires_in),
181 DEFAULT_SKEW,
182 ))
183 }
184}
185
186#[async_trait]
187impl TokenProvider for ClientCredentialsProvider {
188 async fn get_token(&self) -> Result<String, AuthError> {
189 {
190 let cache = self.cache.read().await;
191 if let Some(ref cached) = *cache
192 && cached.is_usable()
193 {
194 return Ok(cached.access_token.as_str().to_owned());
195 }
196 }
197
198 let _guard = self.refresh_lock.lock().await;
199
200 {
201 let cache = self.cache.read().await;
202 if let Some(ref cached) = *cache
203 && cached.is_usable()
204 {
205 return Ok(cached.access_token.as_str().to_owned());
206 }
207 }
208
209 let cached = self.fetch_token().await?;
210 let token = cached.access_token.as_str().to_owned();
211 {
212 let mut cache = self.cache.write().await;
213 *cache = Some(cached);
214 }
215 Ok(token)
216 }
217}
218
219#[cfg(test)]
220mod tests {
221 use std::sync::Arc;
222
223 use super::*;
224 use wiremock::matchers::{body_string_contains, method, path};
225 use wiremock::{Mock, MockServer, ResponseTemplate};
226
227 #[tokio::test]
228 async fn oauth2_rejects_private_ip_token_endpoint() {
229 let result = ClientCredentialsProvider::new(
232 "https://169.254.169.254/token".into(),
233 "client".into(),
234 "secret".into(),
235 None,
236 None,
237 )
238 .await;
239 assert!(
240 result.is_err(),
241 "private IP token endpoint should be rejected"
242 );
243 }
244
245 fn token_response(access_token: &str, expires_in: u64) -> serde_json::Value {
246 serde_json::json!({
247 "access_token": access_token,
248 "token_type": "Bearer",
249 "expires_in": expires_in,
250 })
251 }
252
253 #[tokio::test]
254 async fn test_get_token_fresh() {
255 let server = MockServer::start().await;
256 Mock::given(method("POST"))
257 .and(path("/protocol/openid-connect/token"))
258 .respond_with(ResponseTemplate::new(200).set_body_json(token_response("abc123", 300)))
259 .mount(&server)
260 .await;
261
262 let provider = ClientCredentialsProvider::new_unchecked_for_test(
263 format!("{}/protocol/openid-connect/token", server.uri()), "test-client".into(),
265 "test-secret".into(),
266 None,
267 None,
268 reqwest::Client::new(),
269 );
270 let token = provider.get_token().await.unwrap();
271 assert_eq!(token, "abc123");
272 }
273
274 #[tokio::test]
275 async fn test_get_token_uses_cache() {
276 let server = MockServer::start().await;
277 Mock::given(method("POST"))
278 .respond_with(ResponseTemplate::new(200).set_body_json(token_response("cached", 300)))
279 .expect(1)
280 .mount(&server)
281 .await;
282
283 let provider = ClientCredentialsProvider::new_unchecked_for_test(
284 format!("{}/protocol/openid-connect/token", server.uri()), "c".into(),
286 "s".into(),
287 None,
288 None,
289 reqwest::Client::new(),
290 );
291 let t1 = provider.get_token().await.unwrap();
292 let t2 = provider.get_token().await.unwrap();
293 assert_eq!(t1, "cached");
294 assert_eq!(t2, "cached");
295 }
296
297 #[tokio::test]
298 async fn test_get_token_refreshes_when_stale() {
299 let server = MockServer::start().await;
300 Mock::given(method("POST"))
301 .respond_with(ResponseTemplate::new(200).set_body_json(token_response("first", 1)))
302 .up_to_n_times(1)
303 .mount(&server)
304 .await;
305 Mock::given(method("POST"))
306 .respond_with(ResponseTemplate::new(200).set_body_json(token_response("second", 300)))
307 .mount(&server)
308 .await;
309
310 let provider = ClientCredentialsProvider::new_unchecked_for_test(
311 format!("{}/protocol/openid-connect/token", server.uri()), "c".into(),
313 "s".into(),
314 None,
315 None,
316 reqwest::Client::new(),
317 );
318 let t1 = provider.get_token().await.unwrap();
319 assert_eq!(t1, "first");
320 tokio::time::sleep(Duration::from_millis(1100)).await;
321 let t2 = provider.get_token().await.unwrap();
322 assert_eq!(t2, "second");
323 }
324
325 #[tokio::test]
326 async fn test_get_token_server_error() {
327 let server = MockServer::start().await;
328 Mock::given(method("POST"))
329 .respond_with(ResponseTemplate::new(500))
330 .mount(&server)
331 .await;
332
333 let provider = ClientCredentialsProvider::new_unchecked_for_test(
334 format!("{}/protocol/openid-connect/token", server.uri()), "c".into(),
336 "s".into(),
337 None,
338 None,
339 reqwest::Client::new(),
340 );
341 let err = provider.get_token().await.unwrap_err();
342 assert!(matches!(err, AuthError::ProviderUnavailable(_)));
343 }
344
345 #[tokio::test]
346 async fn test_get_token_invalid_response() {
347 let server = MockServer::start().await;
348 Mock::given(method("POST"))
349 .respond_with(
350 ResponseTemplate::new(200)
351 .set_body_json(serde_json::json!({"error": "invalid_grant"})),
352 )
353 .mount(&server)
354 .await;
355
356 let provider = ClientCredentialsProvider::new_unchecked_for_test(
357 format!("{}/protocol/openid-connect/token", server.uri()), "c".into(),
359 "s".into(),
360 None,
361 None,
362 reqwest::Client::new(),
363 );
364 let err = provider.get_token().await.unwrap_err();
365 assert!(matches!(err, AuthError::ProviderUnavailable(_)));
366 }
367
368 #[tokio::test]
369 async fn test_get_token_sends_audience_as_resource() {
370 let server = MockServer::start().await;
371 Mock::given(method("POST"))
372 .and(body_string_contains(
373 "resource=https%3A%2F%2Fapi.example.com",
374 ))
375 .respond_with(
376 ResponseTemplate::new(200).set_body_json(token_response("aud-token", 300)),
377 )
378 .mount(&server)
379 .await;
380
381 let provider = ClientCredentialsProvider::new_unchecked_for_test(
382 format!("{}/protocol/openid-connect/token", server.uri()), "c".into(),
384 "s".into(),
385 None,
386 Some(vec!["https://api.example.com".into()]),
387 reqwest::Client::new(),
388 );
389 let token = provider.get_token().await.unwrap();
390 assert_eq!(token, "aud-token");
391 }
392
393 #[tokio::test]
394 async fn test_single_flight_concurrent_callers() {
395 let server = MockServer::start().await;
396 Mock::given(method("POST"))
397 .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
398 "access_token": "single-flight-token",
399 "token_type": "Bearer",
400 "expires_in": 300,
401 })))
402 .expect(1)
403 .mount(&server)
404 .await;
405
406 let provider = Arc::new(ClientCredentialsProvider::new_unchecked_for_test(
407 format!("{}/protocol/openid-connect/token", server.uri()), "c".into(),
409 "s".into(),
410 None,
411 None,
412 reqwest::Client::new(),
413 ));
414
415 let mut handles = vec![];
416 for _ in 0..5 {
417 let p = Arc::clone(&provider);
418 handles.push(tokio::spawn(async move { p.get_token().await }));
419 }
420 for h in handles {
421 let token = h.await.unwrap().unwrap();
422 assert_eq!(token, "single-flight-token");
423 }
424 }
425}