xmtp_api_backend/middleware/
auth.rs

1use crate::endpoints::backend::GET_CONFIGURATION_PATH;
2use arc_swap::ArcSwap;
3use prost::bytes::Bytes;
4use std::sync::{
5    Arc, Weak,
6    atomic::{AtomicBool, Ordering},
7};
8use tokio::sync::OnceCell;
9use xmtp_common::{BoxDynError, MaybeSend, MaybeSync, time::Instant};
10#[cfg(not(test))]
11use xmtp_configuration::AUTH_LOCKOUT_COOLDOWN;
12use xmtp_configuration::MAX_CONSECUTIVE_AUTH_FAILURES;
13use xmtp_events::{ClientEvent, EventWriter, LockoutChange, LockoutChanged};
14use xmtp_proto::api::{
15    ApiClientError, AuthError, BytesStream, Client, IsConnectedCheck, grpc_status,
16};
17
18#[cfg(test)]
19/// Longer than the transport's first reconnect delay (100 ms), so a reopen
20/// during a lockout reliably meets it instead of racing it.
21const AUTH_LOCKOUT_COOLDOWN: std::time::Duration = std::time::Duration::from_millis(500);
22
23#[cfg(not(test))]
24use xmtp_common::time::now_secs;
25// Use a fixed clock to keep expiry tests stable.
26#[cfg(test)]
27fn now_secs() -> i64 {
28    1_000_000
29}
30
31#[derive(Clone, Debug, PartialEq, Eq)]
32pub struct Credential {
33    name: http::header::HeaderName,
34    value: http::header::HeaderValue,
35    expires_at_seconds: i64,
36}
37
38impl Credential {
39    pub fn new(
40        name: Option<http::header::HeaderName>,
41        value: http::header::HeaderValue,
42        expires_at_seconds: i64,
43    ) -> Self {
44        Self {
45            name: name.unwrap_or(http::header::AUTHORIZATION),
46            value,
47            expires_at_seconds,
48        }
49    }
50}
51
52#[derive(Clone, Copy, Default)]
53struct AuthState {
54    generation: u64,
55    stale: bool,
56    failures: u32,
57    locked_until: Option<Instant>,
58}
59
60impl AuthState {
61    // implements: AUTH-023
62    fn fail(&mut self) -> bool {
63        self.failures = (self.failures + 1).min(MAX_CONSECUTIVE_AUTH_FAILURES);
64        if self.failures == MAX_CONSECUTIVE_AUTH_FAILURES && self.locked_until.is_none() {
65            self.locked_until = Some(Instant::now() + AUTH_LOCKOUT_COOLDOWN);
66            return true;
67        }
68        false
69    }
70}
71
72#[derive(Default)]
73struct AuthInner {
74    current: OnceCell<ArcSwap<Credential>>,
75    source_supplied: AtomicBool,
76    state: tokio::sync::Mutex<AuthState>,
77    /// Held across the callback so only one runs at a time. It is separate from
78    /// `state` because `AuthHandle::set` takes `state`, and a callback that
79    /// pushes its credential through the handle would deadlock on a lock this
80    /// function held across the await. Tokio's mutex is not reentrant.
81    refresh: tokio::sync::Mutex<()>,
82    event_writers: parking_lot::Mutex<Vec<Weak<dyn EventWriter<()>>>>,
83}
84
85impl AuthInner {
86    fn emit_lockout(&self, change: LockoutChange) {
87        let mut writers = self.event_writers.lock();
88        let live: Vec<_> = writers.iter().filter_map(Weak::upgrade).collect();
89        writers.retain(|writer| writer.strong_count() > 0);
90        drop(writers);
91        for writer in live {
92            writer.emit(
93                Some(ClientEvent::ClientLockoutChanged(LockoutChanged { change })),
94                None,
95            );
96        }
97    }
98    /// Store only while the state lock is held. This operation cannot be cancelled.
99    fn store(&self, credential: Credential) {
100        if let Some(current) = self.current.get() {
101            current.store(Arc::new(credential));
102        } else {
103            self.current
104                .set(ArcSwap::from_pointee(credential))
105                .unwrap_or_else(|_| unreachable!("state lock protects initialization"));
106        }
107    }
108}
109
110#[derive(Clone)]
111pub struct AuthHandle {
112    inner: Arc<AuthInner>,
113}
114
115impl Default for AuthHandle {
116    fn default() -> Self {
117        let inner = Arc::new(AuthInner::default());
118        inner.source_supplied.store(true, Ordering::Release);
119        Self { inner }
120    }
121}
122
123impl std::fmt::Debug for AuthHandle {
124    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
125        f.debug_struct("AuthHandle").finish_non_exhaustive()
126    }
127}
128
129impl AuthHandle {
130    pub fn new() -> Self {
131        Self::default()
132    }
133
134    /// Reserve a handle for an SDK credential set after client creation.
135    /// This empty handle does not satisfy CONF-051 by itself.
136    pub fn sdk_placeholder() -> Self {
137        Self {
138            inner: Arc::new(AuthInner::default()),
139        }
140    }
141
142    // implements: AUTH-024
143    pub async fn set(&self, credential: Credential) {
144        let mut state = self.inner.state.lock().await;
145        let was_locked = state.locked_until.is_some();
146        self.inner.store(credential);
147        self.inner.source_supplied.store(true, Ordering::Release);
148        *state = AuthState {
149            generation: state.generation + 1,
150            ..AuthState::default()
151        };
152        if was_locked {
153            self.inner.emit_lockout(LockoutChange::Left);
154        }
155    }
156
157    pub fn register_event_writer(&self, writer: &Arc<dyn EventWriter<()>>) {
158        let mut writers = self.inner.event_writers.lock();
159        writers.retain(|writer| writer.strong_count() > 0);
160        writers.push(Arc::downgrade(writer));
161    }
162
163    pub fn unregister_event_writer(&self, writer: &Arc<dyn EventWriter<()>>) {
164        self.inner.event_writers.lock().retain(|held| {
165            held.upgrade()
166                .is_some_and(|held| !Arc::ptr_eq(&held, writer))
167        });
168    }
169
170    pub fn id(&self) -> usize {
171        Arc::as_ptr(&self.inner) as usize
172    }
173}
174
175#[xmtp_common::async_trait]
176pub trait AuthCallback: MaybeSend + MaybeSync {
177    async fn on_auth_required(&self) -> Result<Credential, BoxDynError>;
178}
179
180/// Add credentials and refetch after expiry or rejection.
181/// Shared handles share credentials, callback serialization, and lockout state.
182/// Without a callback, expired or rejected credentials remain in use until `set`.
183#[derive(Clone)]
184pub struct AuthMiddleware<C> {
185    inner: C,
186    handle: AuthHandle,
187    callback: Option<Arc<dyn AuthCallback>>,
188}
189
190impl<C> AuthMiddleware<C> {
191    fn bypass_empty_sdk_placeholder(&self) -> bool {
192        !self.handle.inner.source_supplied.load(Ordering::Acquire)
193            && self.callback.is_none()
194            && self.handle.inner.current.get().is_none()
195    }
196
197    #[track_caller]
198    pub fn new(
199        inner: C,
200        callback: Option<Arc<dyn AuthCallback>>,
201        handle: Option<AuthHandle>,
202    ) -> Self {
203        assert!(
204            callback.is_some() || handle.is_some(),
205            "Either a callback or a handle must be provided"
206        );
207        Self {
208            inner,
209            handle: handle.unwrap_or_default(),
210            callback,
211        }
212    }
213
214    /// Keep the credential and its generation together across each network call.
215    /// Commit callback state only after the callback completes, including a probe.
216    // implements: AUTH-021
217    async fn get_credential(&self) -> Result<(Arc<Credential>, u64), AuthError> {
218        let inner = &self.handle.inner;
219        // Take the refresh lock first so only one callback runs at a time. The
220        // state lock is taken and released around it, never held across the
221        // await, because `AuthHandle::set` needs it while a callback runs.
222        let _refresh = inner.refresh.lock().await;
223        let mut state = *inner.state.lock().await;
224        let now = Instant::now();
225        let lockout_left = state.locked_until.is_some_and(|until| until <= now);
226        if let Some(until) = state.locked_until {
227            if until > now {
228                return Err(AuthError::Exhausted);
229            }
230            state.locked_until = None;
231            state.failures = MAX_CONSECUTIVE_AUTH_FAILURES - 1;
232            state.stale = true;
233        }
234        let mut credential = inner.current.get().map(|current| current.load_full());
235        let needs_refresh = state.stale
236            || credential
237                .as_ref()
238                .is_none_or(|credential| credential.expires_at_seconds <= now_secs());
239        if needs_refresh && let Some(callback) = &self.callback {
240            let generation_before = state.generation;
241            let result = callback.on_auth_required().await;
242            let mut guard = inner.state.lock().await;
243            // `AuthHandle::set` may have run during the callback. Its credential
244            // is newer, so keep it and leave its state alone.
245            if guard.generation != generation_before {
246                return Ok((
247                    inner
248                        .current
249                        .get()
250                        .map(|current| current.load_full())
251                        .ok_or(AuthError::MissingCredential)?,
252                    guard.generation,
253                ));
254            }
255            match result {
256                Ok(new) => {
257                    inner.store(new);
258                    credential = inner.current.get().map(|current| current.load_full());
259                    state.generation += 1;
260                    state.stale = false;
261                }
262                Err(_) => {
263                    let entered = state.fail();
264                    *guard = state;
265                    if lockout_left {
266                        inner.emit_lockout(LockoutChange::Left);
267                    }
268                    if entered {
269                        inner.emit_lockout(LockoutChange::Entered);
270                    }
271                    // Distinguish this failed refresh from later refusals.
272                    // Both keep the same public error code and cool-down.
273                    return Err(if state.locked_until.is_some() {
274                        AuthError::ExhaustedAfterAttempt
275                    } else {
276                        AuthError::CallbackFailed { retryable: true }
277                    });
278                }
279            }
280            *guard = state;
281            if lockout_left {
282                inner.emit_lockout(LockoutChange::Left);
283            }
284        } else {
285            *inner.state.lock().await = state;
286            if lockout_left {
287                inner.emit_lockout(LockoutChange::Left);
288            }
289        }
290        let credential = credential.ok_or(AuthError::MissingCredential)?;
291        Ok((credential, state.generation))
292    }
293
294    /// Return whether a current rejection permits one immediate replay.
295    async fn finish<T>(
296        &self,
297        generation: u64,
298        result: Result<T, ApiClientError>,
299    ) -> (Result<T, ApiClientError>, bool) {
300        let rejected = result
301            .as_ref()
302            .err()
303            .and_then(|error| grpc_status(error))
304            .is_some_and(|status| status.code() == tonic::Code::Unauthenticated);
305        let mut state = self.handle.inner.state.lock().await;
306        let current = generation == state.generation;
307        if current {
308            if result.is_ok() {
309                state.failures = 0;
310            } else if rejected {
311                state.stale = true;
312                if self.callback.is_some() && state.fail() {
313                    self.handle.inner.emit_lockout(LockoutChange::Entered);
314                }
315            }
316        }
317        if rejected {
318            // The failed request that starts lockout spends one recovery cycle.
319            // Later refusals spend none. Both keep the same public error code.
320            if self.callback.is_some()
321                && state
322                    .locked_until
323                    .is_some_and(|until| until > Instant::now())
324            {
325                return (Err(AuthError::ExhaustedAfterAttempt.into()), false);
326            }
327            let retryable = self.callback.is_some();
328            (
329                Err(AuthError::CredentialRejected { retryable }.into()),
330                current && retryable,
331            )
332        } else {
333            (result, false)
334        }
335    }
336
337    /// Rebuild all request parts. Replace the credential header instead of appending it.
338    // implements: AUTH-020
339    fn request_builder(
340        parts: &http::request::Parts,
341        credential: &Credential,
342    ) -> http::request::Builder {
343        let mut request = http::Request::builder()
344            .method(parts.method.clone())
345            .uri(parts.uri.clone())
346            .version(parts.version);
347        let headers = request.headers_mut().expect("validated request parts");
348        *headers = parts.headers.clone();
349        headers.insert(credential.name.clone(), credential.value.clone());
350        *request.extensions_mut().expect("validated request parts") = parts.extensions.clone();
351        request
352    }
353}
354
355#[xmtp_common::async_trait]
356impl<C: Client> Client for AuthMiddleware<C> {
357    /// An app-supplied handle is a source even before it holds a credential.
358    /// Only an empty handle reserved by the SDK is excluded from CONF-051.
359    fn has_credential_source(&self) -> bool {
360        self.callback.is_some() || self.handle.inner.source_supplied.load(Ordering::Acquire)
361    }
362
363    fn host(&self) -> &str {
364        self.inner.host()
365    }
366
367    // implements: AUTH-022
368    async fn request(
369        &self,
370        request: http::request::Builder,
371        path: http::uri::PathAndQuery,
372        body: Bytes,
373    ) -> Result<http::Response<Bytes>, ApiClientError> {
374        // The configuration read carries no credential and never
375        // invokes the app's callback. A client asks what the deployment
376        // requires before it can know whether it needs one.
377        // implements: CONF-029
378        if path.path() == GET_CONFIGURATION_PATH {
379            return self.inner.request(request, path, body).await;
380        }
381        if self.bypass_empty_sdk_placeholder() {
382            return self.inner.request(request, path, body).await;
383        }
384        let (parts, ()) = request.body(())?.into_parts();
385        let (credential, generation) = self.get_credential().await?;
386        let result = self
387            .inner
388            .request(
389                Self::request_builder(&parts, &credential),
390                path.clone(),
391                body.clone(),
392            )
393            .await;
394        let (result, replay) = self.finish(generation, result).await;
395        if !replay {
396            return result;
397        }
398        let (credential, generation) = self.get_credential().await?;
399        let result = self
400            .inner
401            .request(Self::request_builder(&parts, &credential), path, body)
402            .await;
403        self.finish(generation, result).await.0
404    }
405
406    // implements: AUTH-022
407    async fn stream(
408        &self,
409        request: http::request::Builder,
410        path: http::uri::PathAndQuery,
411        body: Bytes,
412    ) -> Result<http::Response<BytesStream>, ApiClientError> {
413        if self.bypass_empty_sdk_placeholder() {
414            return self.inner.stream(request, path, body).await;
415        }
416        let (parts, ()) = request.body(())?.into_parts();
417        let (credential, generation) = self.get_credential().await?;
418        let result = self
419            .inner
420            .stream(
421                Self::request_builder(&parts, &credential),
422                path.clone(),
423                body.clone(),
424            )
425            .await;
426        let (result, replay) = self.finish(generation, result).await;
427        if !replay {
428            return result;
429        }
430        let (credential, generation) = self.get_credential().await?;
431        let result = self
432            .inner
433            .stream(Self::request_builder(&parts, &credential), path, body)
434            .await;
435        self.finish(generation, result).await.0
436    }
437
438    async fn bidi_stream(
439        &self,
440        request: http::request::Builder,
441        path: http::uri::PathAndQuery,
442        body: xmtp_common::BoxDynStream<'static, Bytes>,
443    ) -> Result<http::Response<BytesStream>, ApiClientError> {
444        if self.bypass_empty_sdk_placeholder() {
445            return self.inner.bidi_stream(request, path, body).await;
446        }
447        let (parts, ()) = request.body(())?.into_parts();
448        let (credential, generation) = self.get_credential().await?;
449        let result = self
450            .inner
451            .bidi_stream(Self::request_builder(&parts, &credential), path, body)
452            .await;
453        self.finish(generation, result).await.0
454    }
455}
456
457#[xmtp_common::async_trait]
458impl<C: IsConnectedCheck> IsConnectedCheck for AuthMiddleware<C> {
459    async fn is_connected(&self) -> bool {
460        self.inner.is_connected().await
461    }
462}
463
464#[cfg(test)]
465mod tests;
466
467#[cfg(test)]
468mod refetch_tests;
xmtp_api_backend/middleware/auth.rs

xmtp_api_backend/middleware/
auth.rs
