wiremux_auth/providers/
azure.rs

1//! Azure AD client-credentials TokenProvider.
2
3use std::sync::Arc;
4use std::sync::atomic::{AtomicBool, Ordering};
5use std::time::{Duration, Instant};
6
7use tokio::sync::RwLock;
8
9use crate::TokenProvider;
10use crate::error::AuthError;
11use crate::helpers::{
12    InFlight, duration_from_expires_in_secs, lead_or_follow, oauth_http_client, post_form_url,
13    vendor_rejected_summary,
14};
15
16const DEFAULT_LIFETIME_SECS: u64 = 3600;
17const DEFAULT_SCOPE: &str = "https://management.azure.com/.default";
18
19struct CachedToken {
20    access_token: String,
21    acquired_at: Instant,
22    lifetime: Duration,
23}
24
25impl CachedToken {
26    fn needs_refresh(&self) -> bool {
27        self.acquired_at.elapsed() >= self.lifetime.mul_f64(0.8)
28    }
29}
30
31/// Azure AD client-credentials grant.
32#[derive(Clone)]
33pub struct AzureTokenProvider {
34    inner: Arc<Inner>,
35}
36
37struct Inner {
38    tenant_id: String,
39    client_id: String,
40    client_secret: String,
41    scope: String,
42    token_url: String,
43    http: reqwest::Client,
44    force_refresh: AtomicBool,
45    state: RwLock<Option<CachedToken>>,
46    inflight: InFlight,
47}
48
49impl std::fmt::Debug for AzureTokenProvider {
50    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
51        f.debug_struct("AzureTokenProvider")
52            .field("tenant_id", &self.inner.tenant_id)
53            .field("client_id", &self.inner.client_id)
54            .field("client_secret", &"[REDACTED]")
55            .field("scope", &self.inner.scope)
56            .field("token_url", &self.inner.token_url)
57            .finish()
58    }
59}
60
61impl AzureTokenProvider {
62    /// Client-credentials provider for `login.microsoftonline.com/{tenant}/oauth2/v2.0/token`.
63    pub fn new(
64        tenant_id: impl Into<String>,
65        client_id: impl Into<String>,
66        client_secret: impl Into<String>,
67        scope: impl Into<String>,
68    ) -> Result<Self, AuthError> {
69        let tenant_id = tenant_id.into();
70        let client_id = client_id.into();
71        let client_secret = client_secret.into();
72        let mut scope = scope.into();
73        if scope.trim().is_empty() {
74            scope = DEFAULT_SCOPE.to_string();
75        }
76        if tenant_id.trim().is_empty() {
77            return Err(AuthError::MissingField("Azure tenant_id".into()));
78        }
79        if client_id.trim().is_empty() {
80            return Err(AuthError::MissingField("Azure client_id".into()));
81        }
82        if client_secret.trim().is_empty() {
83            return Err(AuthError::MissingField("Azure client_secret".into()));
84        }
85        let token_url = format!("https://login.microsoftonline.com/{tenant_id}/oauth2/v2.0/token");
86        Ok(Self {
87            inner: Arc::new(Inner {
88                tenant_id,
89                client_id,
90                client_secret,
91                scope,
92                token_url,
93                http: oauth_http_client()?,
94                force_refresh: AtomicBool::new(false),
95                state: RwLock::new(None),
96                inflight: InFlight::new(),
97            }),
98        })
99    }
100
101    /// Override the token URL (tests, national clouds).
102    #[must_use]
103    pub fn with_token_url(self, token_url: impl Into<String>) -> Self {
104        let mut inner = Arc::try_unwrap(self.inner).unwrap_or_else(|arc| (*arc).clone_inner());
105        inner.token_url = token_url.into();
106        Self {
107            inner: Arc::new(inner),
108        }
109    }
110
111    async fn refresh(&self, force: bool) -> Result<String, AuthError> {
112        {
113            let state = self.inner.state.read().await;
114            if !force
115                && let Some(tok) = state.as_ref()
116                && !tok.needs_refresh()
117            {
118                return Ok(tok.access_token.clone());
119            }
120        }
121        let form = [
122            ("client_id", self.inner.client_id.as_str()),
123            ("client_secret", self.inner.client_secret.as_str()),
124            ("scope", self.inner.scope.as_str()),
125            ("grant_type", "client_credentials"),
126        ];
127        let (status, body) = post_form_url(&self.inner.http, &self.inner.token_url, &form).await?;
128        if !(200..300).contains(&status) {
129            return Err(AuthError::VendorRejected {
130                status,
131                summary: vendor_rejected_summary(
132                    "Azure token request failed",
133                    &body,
134                    &self.inner.token_url,
135                ),
136            });
137        }
138        let parsed = parse_azure_token(&body)?;
139        if parsed.access_token.trim().is_empty() {
140            return Err(AuthError::EmptyWriteRefused);
141        }
142        let lifetime =
143            duration_from_expires_in_secs(parsed.expires_in.unwrap_or(DEFAULT_LIFETIME_SECS));
144        let token = parsed.access_token.clone();
145        *self.inner.state.write().await = Some(CachedToken {
146            access_token: parsed.access_token,
147            acquired_at: Instant::now(),
148            lifetime,
149        });
150        Ok(token)
151    }
152}
153
154impl Inner {
155    fn clone_inner(&self) -> Self {
156        Self {
157            tenant_id: self.tenant_id.clone(),
158            client_id: self.client_id.clone(),
159            client_secret: self.client_secret.clone(),
160            scope: self.scope.clone(),
161            token_url: self.token_url.clone(),
162            http: self.http.clone(),
163            force_refresh: AtomicBool::new(self.force_refresh.load(Ordering::SeqCst)),
164            state: RwLock::new(None),
165            inflight: InFlight::new(),
166        }
167    }
168}
169
170struct ParsedAzureToken {
171    access_token: String,
172    expires_in: Option<u64>,
173}
174
175fn parse_azure_token(body: &str) -> Result<ParsedAzureToken, AuthError> {
176    let doc: serde_json::Value = serde_json::from_str(body)
177        .map_err(|e| AuthError::TokenProvider(format!("Azure token: invalid JSON: {e}")))?;
178    let access_token = doc
179        .get("access_token")
180        .and_then(|v| v.as_str())
181        .unwrap_or("")
182        .to_string();
183    if access_token.is_empty() {
184        return Err(AuthError::TokenProvider(
185            "Azure token response has no access_token".into(),
186        ));
187    }
188    let expires_in = doc.get("expires_in").and_then(|v| {
189        v.as_u64()
190            .or_else(|| v.as_str().and_then(|s| s.parse().ok()))
191    });
192    Ok(ParsedAzureToken {
193        access_token,
194        expires_in,
195    })
196}
197
198impl TokenProvider for AzureTokenProvider {
199    fn mark_stale(&self) {
200        self.inner.force_refresh.store(true, Ordering::SeqCst);
201    }
202
203    async fn get_token(&self) -> Result<String, AuthError> {
204        // load() not swap: a concurrent get_token must not see a cleared force
205        // flag before the InFlight claim.
206        {
207            let state = self.inner.state.read().await;
208            if !self.inner.force_refresh.load(Ordering::SeqCst)
209                && let Some(tok) = state.as_ref()
210                && !tok.needs_refresh()
211                && !self.inner.inflight.is_busy()
212            {
213                return Ok(tok.access_token.clone());
214            }
215        }
216        lead_or_follow(&self.inner.inflight, || async {
217            let force = self.inner.force_refresh.swap(false, Ordering::SeqCst);
218            let result = self.refresh(force).await;
219            if result.is_err() && force {
220                self.inner.force_refresh.store(true, Ordering::SeqCst);
221            }
222            result
223        })
224        .await
225    }
226}
227
228#[cfg(test)]
229mod tests {
230    use super::*;
231    use std::io::{Read, Write};
232    use std::net::TcpListener;
233    use std::thread;
234    use std::time::{Duration, Instant};
235
236    fn read_http_headers(stream: &mut impl Read) -> Vec<u8> {
237        let mut buf = Vec::new();
238        let mut tmp = [0u8; 1024];
239        loop {
240            match stream.read(&mut tmp) {
241                Ok(0) => break,
242                Ok(n) => {
243                    buf.extend_from_slice(&tmp[..n]);
244                    if buf.windows(4).any(|w| w == b"\r\n\r\n") {
245                        break;
246                    }
247                }
248                Err(_) => break,
249            }
250        }
251        buf
252    }
253
254    fn write_http_json(stream: &mut impl Write, status: u16, body: &str) {
255        let resp = format!(
256            "HTTP/1.1 {status} OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}",
257            body.len()
258        );
259        let _ = stream.write_all(resp.as_bytes());
260    }
261
262    fn spawn_http_server(status: u16, body: &str) -> (String, thread::JoinHandle<String>) {
263        let listener = TcpListener::bind("127.0.0.1:0").expect("bind mock");
264        let addr = listener.local_addr().expect("local_addr");
265        let body = body.to_owned();
266        let handle = thread::spawn(move || {
267            let (mut stream, _) = listener.accept().expect("accept");
268            stream.set_read_timeout(Some(Duration::from_secs(5))).ok();
269            let buf = read_http_headers(&mut stream);
270            write_http_json(&mut stream, status, &body);
271            String::from_utf8_lossy(&buf).into_owned()
272        });
273        (format!("http://{addr}/oauth/token"), handle)
274    }
275
276    fn spawn_http_script(responses: &[(u16, &str)]) -> (String, thread::JoinHandle<usize>) {
277        let listener = TcpListener::bind("127.0.0.1:0").expect("bind mock");
278        let addr = listener.local_addr().expect("local_addr");
279        let responses: Vec<(u16, String)> = responses
280            .iter()
281            .map(|(status, body)| (*status, (*body).to_owned()))
282            .collect();
283        let handle = thread::spawn(move || {
284            let mut served = 0usize;
285            for (status, body) in responses {
286                let (mut stream, _) = listener.accept().expect("accept");
287                stream.set_read_timeout(Some(Duration::from_secs(5))).ok();
288                let _ = read_http_headers(&mut stream);
289                write_http_json(&mut stream, status, &body);
290                served += 1;
291            }
292            served
293        });
294        (format!("http://{addr}/oauth/token"), handle)
295    }
296
297    /// Cache-fill 200, then up to two delayed refresh 200s (leader + extra POST).
298    fn spawn_http_cache_then_refresh(
299        first: &str,
300        refreshed: &str,
301        refresh_delay: Duration,
302        extra_idle: Duration,
303    ) -> (String, thread::JoinHandle<usize>) {
304        let listener = TcpListener::bind("127.0.0.1:0").expect("bind mock");
305        listener.set_nonblocking(true).expect("nonblocking");
306        let addr = listener.local_addr().expect("local_addr");
307        let first = first.to_owned();
308        let refreshed = refreshed.to_owned();
309        let handle = thread::spawn(move || {
310            let mut served = 0usize;
311            let first_deadline = Instant::now() + Duration::from_secs(5);
312            let mut extra_deadline: Option<Instant> = None;
313            while served < 3 {
314                match listener.accept() {
315                    Ok((mut stream, _)) => {
316                        stream.set_nonblocking(false).ok();
317                        stream.set_read_timeout(Some(Duration::from_secs(5))).ok();
318                        let _ = read_http_headers(&mut stream);
319                        if served == 0 {
320                            write_http_json(&mut stream, 200, &first);
321                        } else {
322                            if served == 1 && !refresh_delay.is_zero() {
323                                thread::sleep(refresh_delay);
324                            }
325                            write_http_json(&mut stream, 200, &refreshed);
326                            extra_deadline = Some(Instant::now() + extra_idle);
327                        }
328                        served += 1;
329                    }
330                    Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => {
331                        let limit = extra_deadline.unwrap_or(first_deadline);
332                        if Instant::now() >= limit {
333                            break;
334                        }
335                        thread::sleep(Duration::from_millis(5));
336                    }
337                    Err(_) => break,
338                }
339            }
340            served
341        });
342        (format!("http://{addr}/oauth/token"), handle)
343    }
344
345    #[tokio::test]
346    async fn azure_client_credentials_returns_access_token() {
347        let (url, handle) = spawn_http_server(
348            200,
349            r#"{"access_token":"az-tok","expires_in":3600,"token_type":"Bearer"}"#,
350        );
351        let p = AzureTokenProvider::new("tenant", "client", "secret", "https://example/.default")
352            .expect("new")
353            .with_token_url(url);
354        assert_eq!(p.get_token().await.expect("token"), "az-tok");
355        let req = handle.join().expect("join");
356        assert!(req.contains("grant_type=client_credentials"), "{req}");
357        assert!(req.contains("client_id=client"), "{req}");
358        assert!(req.contains("client_secret=secret"), "{req}");
359    }
360
361    #[tokio::test]
362    async fn azure_cached_token_skips_second_http() {
363        let (url, handle) =
364            spawn_http_server(200, r#"{"access_token":"az-tok","expires_in":3600}"#);
365        let p = AzureTokenProvider::new("tenant", "client", "secret", "scope")
366            .expect("new")
367            .with_token_url(url);
368        assert_eq!(p.get_token().await.expect("first"), "az-tok");
369        assert_eq!(p.get_token().await.expect("cached"), "az-tok");
370        let _ = handle.join();
371    }
372
373    #[test]
374    fn azure_debug_redacts_secret() {
375        let p = AzureTokenProvider::new("tenant", "client", "super-secret", "scope").expect("new");
376        let debug = format!("{p:?}");
377        assert!(!debug.contains("super-secret"), "{debug}");
378        assert!(debug.contains("[REDACTED]"), "{debug}");
379    }
380
381    #[test]
382    fn azure_empty_secret_fails() {
383        let err = AzureTokenProvider::new("t", "c", "", "s").expect_err("empty");
384        assert!(err.to_string().contains("client_secret"), "{err}");
385    }
386
387    #[tokio::test]
388    async fn forced_refresh_http_error_restores_force_flag() {
389        let (url, handle) = spawn_http_script(&[
390            (200, r#"{"access_token":"az-first","expires_in":3600}"#),
391            (401, r#"{"error":"invalid_client"}"#),
392            (200, r#"{"access_token":"az-retry","expires_in":3600}"#),
393        ]);
394        let p = AzureTokenProvider::new("tenant", "client", "secret", "scope")
395            .expect("new")
396            .with_token_url(url);
397        assert_eq!(p.get_token().await.expect("cache"), "az-first");
398        p.mark_stale();
399        let err = p.get_token().await.expect_err("forced 401");
400        let msg = err.to_string();
401        assert!(msg.contains("401") || msg.contains("Azure"), "{msg}");
402        assert_eq!(
403            msg.matches("(HTTP ").count(),
404            1,
405            "VendorRejected must not re-wrap HTTP: {msg}"
406        );
407        assert_eq!(
408            p.get_token().await.expect("retry after failed force"),
409            "az-retry"
410        );
411        assert_eq!(handle.join().expect("join"), 3);
412    }
413
414    #[tokio::test]
415    async fn concurrent_mark_stale_does_not_serve_stale_cache() {
416        let first = r#"{"access_token":"az-first","expires_in":3600}"#;
417        let second = r#"{"access_token":"az-second","expires_in":3600}"#;
418        let (url, handle) = spawn_http_cache_then_refresh(
419            first,
420            second,
421            Duration::from_millis(80),
422            Duration::from_millis(200),
423        );
424        let p = AzureTokenProvider::new("tenant", "client", "secret", "scope")
425            .expect("new")
426            .with_token_url(url);
427        assert_eq!(p.get_token().await.expect("cache"), "az-first");
428        p.mark_stale();
429        let a = p.clone();
430        let b = p.clone();
431        let (left, right) = tokio::join!(a.get_token(), b.get_token());
432        let left = left.expect("left");
433        let right = right.expect("right");
434        assert_ne!(left, "az-first", "must not return just-staled cache");
435        assert_ne!(right, "az-first", "must not return just-staled cache");
436        assert_eq!(left, "az-second");
437        assert_eq!(right, "az-second");
438        let served = handle.join().expect("join");
439        assert!(
440            served >= 2,
441            "cache fill plus at least one refresh, got {served}"
442        );
443    }
444}
wiremux_auth/providers/azure.rs

wiremux_auth/providers/
azure.rs
