A better Rust ATProto crate
1use std::{str::FromStr, sync::Arc};
2
3use chrono::TimeDelta;
4
5use crate::{
6 atproto::{AtprotoClientMetadata, atproto_client_metadata},
7 authstore::ClientAuthStore,
8 dpop::DpopExt,
9 keyset::Keyset,
10 request::{OAuthMetadata, refresh},
11 resolver::OAuthResolver,
12 scopes::Scopes,
13 types::TokenSet,
14};
15
16use dashmap::DashMap;
17use jacquard_common::{
18 IntoStatic,
19 bos::{BosStr, DefaultStr},
20 deps::fluent_uri::Uri,
21 http_client::HttpClient,
22 session::SessionStoreError,
23 types::{did::Did, string::Datetime},
24};
25use jose_jwk::Key;
26use serde::{Deserialize, Serialize};
27use smol_str::{SmolStr, ToSmolStr, format_smolstr};
28use tokio::sync::Mutex;
29
30/// Provides DPoP key material and per-server nonces to the DPoP proof-building machinery.
31///
32/// This trait abstracts over two different holders of DPoP state: [`DpopReqData`] (used
33/// during the initial authorization request, where only an authserver nonce is tracked) and
34/// [`DpopClientData`] (used in active sessions, where both authserver and host nonces are
35/// maintained). Implementors must store nonces durably so that the next request to the same
36/// server includes the most recently observed nonce.
37pub trait DpopDataSource {
38 /// Return the private JWK used to sign DPoP proofs.
39 fn key(&self) -> &Key;
40 /// Return the most recently observed nonce from the authorization server, if any.
41 fn authserver_nonce(&self) -> Option<&str>;
42 /// Persist a new nonce received from the authorization server.
43 fn set_authserver_nonce(&mut self, nonce: SmolStr);
44 /// Return the most recently observed nonce from the resource server (PDS), if any.
45 fn host_nonce(&self) -> Option<&str>;
46 /// Persist a new nonce received from the resource server (PDS).
47 fn set_host_nonce(&mut self, nonce: SmolStr);
48}
49
50/// Persisted information about an OAuth session. Used to resume an active session.
51#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
52#[serde(bound(
53 serialize = "S: serde::Serialize + BosStr + Ord",
54 deserialize = "S: serde::Deserialize<'de> + BosStr + AsRef<str>"
55))]
56pub struct ClientSessionData<S: BosStr = DefaultStr> {
57 /// DID of the authenticated account; serves as the primary key for session storage
58 /// because only one active session per account is assumed.
59 pub account_did: Did<S>,
60
61 /// Opaque identifier that distinguishes this session from other sessions for the same account.
62 ///
63 /// Reuses the random `state` token generated during the PAR flow.
64 pub session_id: S,
65
66 /// Base URL of the resource server (PDS): scheme, host, and port only
67 pub host_url: Uri<String>,
68
69 /// Base URL of the authorization server (PDS or entryway): scheme, host, and port only
70 pub authserver_url: S,
71
72 /// Full URL of the authorization server's token endpoint.
73 pub authserver_token_endpoint: S,
74
75 /// Full URL of the authorization server's revocation endpoint, if advertised.
76 #[serde(skip_serializing_if = "std::option::Option::is_none")]
77 pub authserver_revocation_endpoint: Option<S>,
78
79 /// The set of OAuth scopes approved for this session, as returned in the initial token response.
80 pub scopes: Scopes<S>,
81
82 /// DPoP key and nonce state for ongoing requests in this session.
83 #[serde(flatten)]
84 pub dpop_data: DpopClientData,
85
86 /// Current token set (access token, refresh token, expiry, etc.).
87 #[serde(flatten)]
88 pub token_set: TokenSet<S>,
89
90 /// Fully expanded scopes with include scopes resolved.
91 ///
92 /// This is populated eagerly at session creation when `scope-check` is enabled.
93 /// It is `None` when scope checking is not enabled, when no eager resolution was
94 /// performed, or when reading older persisted sessions that predate this field.
95 #[serde(default, skip_serializing_if = "Option::is_none")]
96 pub resolved_scopes: Option<Vec<crate::scopes::Scope<smol_str::SmolStr>>>,
97}
98
99impl<S: BosStr + Ord + IntoStatic + AsRef<str>> IntoStatic for ClientSessionData<S>
100where
101 S::Output: BosStr + Ord + AsRef<str>,
102{
103 type Output = ClientSessionData<S::Output>;
104
105 fn into_static(self) -> Self::Output {
106 let resolved_scopes = self.resolved_scopes;
107
108 ClientSessionData {
109 authserver_url: self.authserver_url.into_static(),
110 authserver_token_endpoint: self.authserver_token_endpoint.into_static(),
111 authserver_revocation_endpoint: self
112 .authserver_revocation_endpoint
113 .map(IntoStatic::into_static),
114 scopes: self.scopes.into_static(),
115 dpop_data: self.dpop_data,
116 token_set: self.token_set.into_static(),
117 account_did: self.account_did.into_static(),
118 session_id: self.session_id.into_static(),
119 host_url: self.host_url.clone(),
120 resolved_scopes,
121 }
122 }
123}
124
125impl<S: BosStr + Ord + AsRef<str>> ClientSessionData<S> {
126 /// Update this session's token set and, if the new token set includes scopes, replace the scope list.
127 ///
128 /// Called after a successful token refresh so that any scope changes returned by the server
129 /// are reflected in the persisted session without requiring a full re-authentication.
130 ///
131 /// This method is only available on `DefaultStr`-backed sessions (the common case for
132 /// in-memory sessions). Zero-copy borrowed sessions are read-only by nature and would
133 /// not be refreshed in place.
134 pub fn update_with_tokens(&mut self, token_set: &TokenSet<S>)
135 where
136 S: FromStr + Clone + From<SmolStr> + AsRef<str>,
137 S::Err: std::fmt::Debug,
138 {
139 if let Some(scope_str) = token_set.scope.as_ref() {
140 // Parse scopes from the returned scope string, converting to the appropriate backing type
141 let scopes_smol = Scopes::new(scope_str.as_ref().to_smolstr())
142 .expect("server returned invalid scopes in token refresh");
143 self.scopes = scopes_smol.convert();
144 }
145 self.token_set = token_set.clone();
146 }
147}
148
149/// DPoP state for an active OAuth session, persisted alongside the token set.
150///
151/// Both nonces must be written back to the store after each request so that the next
152/// request to the same server includes the correct replay-protection nonce.
153#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
154pub struct DpopClientData {
155 /// The private JWK bound to this session; used to sign all DPoP proofs.
156 pub dpop_key: Key,
157 /// Most recently observed DPoP nonce from the authorization server.
158 pub dpop_authserver_nonce: SmolStr,
159 /// Most recently observed DPoP nonce from the resource server (PDS).
160 pub dpop_host_nonce: SmolStr,
161}
162
163impl DpopDataSource for DpopClientData {
164 fn key(&self) -> &Key {
165 &self.dpop_key
166 }
167
168 fn authserver_nonce(&self) -> Option<&str> {
169 Some(self.dpop_authserver_nonce.as_ref())
170 }
171
172 fn host_nonce(&self) -> Option<&str> {
173 Some(self.dpop_host_nonce.as_ref())
174 }
175
176 fn set_authserver_nonce(&mut self, nonce: SmolStr) {
177 self.dpop_authserver_nonce = nonce;
178 }
179
180 fn set_host_nonce(&mut self, nonce: SmolStr) {
181 self.dpop_host_nonce = nonce;
182 }
183}
184
185/// Transient state created during the PAR flow and consumed by the callback handler.
186///
187/// This struct is persisted to the auth store between [`crate::request::par`] and
188/// [`crate::client::OAuthClient::callback`] so that the callback can verify the
189/// `state`, reconstruct the token exchange, and create a full [`ClientSessionData`].
190#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
191#[serde(bound(
192 serialize = "S: serde::Serialize + BosStr + Ord",
193 deserialize = "S: serde::Deserialize<'de> + BosStr + AsRef<str>"
194))]
195pub struct AuthRequestData<S: BosStr = DefaultStr> {
196 /// Random identifier generated for this authorization request; used as the primary key
197 /// for storing and looking up this record during the callback.
198 pub state: S,
199
200 /// Base URL of the authorization server that was selected for this flow.
201 pub authserver_url: S,
202
203 /// If the flow was initiated with a DID or handle, the resolved DID is stored here
204 /// so it can be compared against the `sub` in the token response.
205 #[serde(skip_serializing_if = "std::option::Option::is_none")]
206 pub account_did: Option<Did<S>>,
207
208 /// OAuth scopes requested for this authorization.
209 pub scopes: Scopes<S>,
210
211 /// The PAR `request_uri` returned by the authorization server; included in the redirect URL.
212 pub request_uri: S,
213
214 /// Full URL of the authorization server's token endpoint.
215 pub authserver_token_endpoint: S,
216
217 /// Full URL of the authorization server's revocation endpoint, if advertised.
218 #[serde(skip_serializing_if = "std::option::Option::is_none")]
219 pub authserver_revocation_endpoint: Option<S>,
220
221 /// The PKCE code verifier whose SHA-256 hash was sent as the code challenge; required
222 /// at the token exchange step to prove the initiator of the auth request.
223 pub pkce_verifier: S,
224
225 /// DPoP key and any authserver nonce observed during the PAR request.
226 #[serde(flatten)]
227 pub dpop_data: DpopReqData,
228}
229
230impl<S: BosStr + Ord + IntoStatic + AsRef<str>> IntoStatic for AuthRequestData<S>
231where
232 S::Output: BosStr + Ord + AsRef<str>,
233{
234 type Output = AuthRequestData<S::Output>;
235
236 fn into_static(self) -> AuthRequestData<S::Output> {
237 AuthRequestData {
238 request_uri: self.request_uri.into_static(),
239 authserver_token_endpoint: self.authserver_token_endpoint.into_static(),
240 authserver_revocation_endpoint: self
241 .authserver_revocation_endpoint
242 .map(|s| s.into_static()),
243 pkce_verifier: self.pkce_verifier.into_static(),
244 dpop_data: self.dpop_data,
245 state: self.state.into_static(),
246 authserver_url: self.authserver_url.into_static(),
247 account_did: self.account_did.into_static(),
248 scopes: self.scopes.into_static(),
249 }
250 }
251}
252
253/// DPoP state for an in-progress authorization request (PAR through code exchange).
254///
255/// Unlike [`DpopClientData`], this struct only tracks the authserver nonce—no resource-server
256/// nonce is needed until a full session is established.
257#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
258pub struct DpopReqData {
259 /// The private JWK generated fresh for this authorization request and session.
260 pub dpop_key: Key,
261 /// DPoP nonce received from the authorization server during the PAR exchange, if any.
262 pub dpop_authserver_nonce: Option<SmolStr>,
263}
264
265impl DpopDataSource for DpopReqData {
266 fn key(&self) -> &Key {
267 &self.dpop_key
268 }
269
270 fn authserver_nonce(&self) -> Option<&str> {
271 self.dpop_authserver_nonce.as_ref().map(|n| n.as_ref())
272 }
273
274 fn host_nonce(&self) -> Option<&str> {
275 None
276 }
277
278 fn set_authserver_nonce(&mut self, nonce: SmolStr) {
279 self.dpop_authserver_nonce = Some(nonce);
280 }
281
282 fn set_host_nonce(&mut self, _nonce: SmolStr) {}
283}
284
285/// Static configuration for an OAuth client: the signing keyset and registered client metadata.
286///
287/// `ClientData` is constructed once at startup and shared (via `Arc`) across all sessions
288/// managed by the same [`crate::client::OAuthClient`].
289#[derive(Clone, Debug)]
290pub struct ClientData<S>
291where
292 S: BosStr + FromStr + Ord,
293 <S as FromStr>::Err: core::fmt::Debug,
294{
295 /// Optional private key set used for `private_key_jwt` client authentication.
296 /// When `None`, the `none` authentication method is used instead.
297 pub keyset: Option<Keyset>,
298 /// AT Protocol-specific client registration metadata (redirect URIs, scopes, etc.).
299 pub config: AtprotoClientMetadata<S>,
300}
301
302impl<S> IntoStatic for ClientData<S>
303where
304 S: BosStr + FromStr + Ord + IntoStatic,
305 S::Output: BosStr + Ord + FromStr,
306 <S as FromStr>::Err: core::fmt::Debug,
307 <S::Output as FromStr>::Err: core::fmt::Debug,
308{
309 type Output = ClientData<S::Output>;
310 fn into_static(self) -> ClientData<S::Output> {
311 ClientData {
312 keyset: self.keyset,
313 config: self.config.into_static(),
314 }
315 }
316}
317
318impl<S: BosStr + FromStr + Ord> ClientData<S>
319where
320 <S as FromStr>::Err: core::fmt::Debug,
321{
322 /// Create `ClientData` with an optional signing keyset and the given client metadata.
323 pub fn new(keyset: Option<Keyset>, config: AtprotoClientMetadata<S>) -> Self {
324 Self { keyset, config }
325 }
326
327 /// Create `ClientData` without a signing keyset, relying on the `none` auth method.
328 ///
329 /// Suitable for public clients (e.g., single-page applications or native apps) that
330 /// cannot securely store a private key.
331 pub fn new_public(config: AtprotoClientMetadata<S>) -> Self {
332 Self {
333 keyset: None,
334 config,
335 }
336 }
337}
338
339/// A bundle of client configuration and an active session, used for operations that need both.
340///
341/// `ClientSession` is a convenience type that pairs a [`ClientData`] with a
342/// [`ClientSessionData`] so that methods like `metadata` can access both without requiring
343/// callers to pass them separately.
344pub struct ClientSession<S: BosStr = DefaultStr>
345where
346 S: FromStr + Ord,
347 <S as FromStr>::Err: core::fmt::Debug,
348{
349 /// Optional signing keyset, forwarded from [`ClientData`].
350 pub keyset: Option<Keyset>,
351 /// Client registration metadata, forwarded from [`ClientData`].
352 pub config: AtprotoClientMetadata<S>,
353 /// The session state for the authenticated account.
354 pub session_data: ClientSessionData<S>,
355}
356
357impl<S: BosStr> ClientSession<S>
358where
359 S: FromStr + Ord + Clone,
360 <S as FromStr>::Err: core::fmt::Debug,
361{
362 /// Construct a `ClientSession` from a [`ClientData`] and an active session.
363 pub fn new(
364 ClientData { keyset, config }: ClientData<S>,
365 session_data: ClientSessionData<S>,
366 ) -> Self {
367 Self {
368 keyset,
369 config,
370 session_data,
371 }
372 }
373
374 /// Fetch and assemble an [`OAuthMetadata`] for the authorization server of this session.
375 pub async fn metadata<T: HttpClient + OAuthResolver + Send + Sync>(
376 &self,
377 client: &T,
378 ) -> Result<OAuthMetadata<S>, Error>
379 where
380 S: IntoStatic,
381 {
382 Ok(OAuthMetadata {
383 server_metadata: client
384 .get_authorization_server_metadata(self.session_data.authserver_url.as_ref())
385 .await
386 .map_err(|e| Error::ServerAgent(crate::request::RequestError::resolver(e)))?,
387 client_metadata: atproto_client_metadata(&self.config, &self.keyset).unwrap(),
388 keyset: self.keyset.clone(),
389 })
390 }
391}
392
393/// Errors that can occur during OAuth session management.
394#[derive(thiserror::Error, Debug, miette::Diagnostic)]
395#[non_exhaustive]
396pub enum Error {
397 /// A token-endpoint or metadata operation failed.
398 #[error(transparent)]
399 #[diagnostic(code(jacquard_oauth::session::request))]
400 ServerAgent(#[from] crate::request::RequestError),
401 /// The backing session store returned an error.
402 #[error(transparent)]
403 #[diagnostic(code(jacquard_oauth::session::storage))]
404 Store(#[from] SessionStoreError),
405 /// The requested session does not exist in the store.
406 #[error("session does not exist")]
407 #[diagnostic(code(jacquard_oauth::session::not_found))]
408 SessionNotFound,
409 /// Token refresh failed with a permanent error (e.g., `invalid_grant`); the session
410 /// has already been removed from the store and the user must re-authenticate.
411 #[error("session refresh failed permanently")]
412 #[diagnostic(
413 code(jacquard_oauth::session::refresh_failed),
414 help("the session has been cleared - user must re-authenticate")
415 )]
416 RefreshFailed(#[source] crate::request::RequestError),
417}
418
419impl Error {
420 /// Returns true if this error indicates a permanent auth failure
421 /// where the user needs to re-authenticate.
422 pub fn is_permanent(&self) -> bool {
423 match self {
424 Error::RefreshFailed(_) => true,
425 Error::SessionNotFound => true,
426 Error::ServerAgent(e) => e.is_permanent(),
427 Error::Store(_) => false,
428 }
429 }
430}
431
432/// Central coordinator for OAuth session storage and token refresh.
433///
434/// `SessionRegistry` wraps the [`ClientAuthStore`] and provides serialized token refresh:
435/// concurrent refresh attempts for the same `(DID, session_id)` pair are coalesced behind
436/// a per-key `Mutex` stored in `pending`, so only one refresh request is issued to the
437/// authorization server even when many concurrent requests detect an expired token.
438pub struct SessionRegistry<T, S, Str>
439where
440 T: OAuthResolver,
441 S: ClientAuthStore,
442 Str: BosStr + FromStr + Ord,
443 <Str as FromStr>::Err: core::fmt::Debug,
444{
445 /// Backing store for persisting session data across process restarts.
446 pub store: Arc<S>,
447 /// Shared resolver used to fetch authorization server metadata during refresh.
448 pub client: Arc<T>,
449 /// Static client configuration (keyset and registration metadata).
450 pub client_data: ClientData<Str>,
451 /// Per-`(DID, session_id)` mutex that serializes concurrent refresh attempts.
452 pending: DashMap<SmolStr, Arc<Mutex<()>>>,
453}
454
455impl<T, S, Str> SessionRegistry<T, S, Str>
456where
457 S: ClientAuthStore,
458 T: OAuthResolver,
459 Str: BosStr + FromStr + Ord,
460 <Str as FromStr>::Err: core::fmt::Debug,
461{
462 /// Create a new registry, taking ownership of the store.
463 pub fn new(store: S, client: Arc<T>, client_data: ClientData<Str>) -> Self {
464 let store = Arc::new(store);
465 Self {
466 store,
467 client,
468 client_data,
469 pending: DashMap::new(),
470 }
471 }
472
473 /// Create a new registry from an already-`Arc`-wrapped store.
474 ///
475 /// Use this variant when the store needs to be accessed from outside the registry,
476 /// for example to expose session listing or administration functionality.
477 pub fn new_shared(store: Arc<S>, client: Arc<T>, client_data: ClientData<Str>) -> Self {
478 Self {
479 store,
480 client,
481 client_data,
482 pending: DashMap::new(),
483 }
484 }
485}
486
487impl<T, S, Str> SessionRegistry<T, S, Str>
488where
489 S: ClientAuthStore + Send + Sync + 'static,
490 T: OAuthResolver + DpopExt + Send + Sync + 'static,
491 Str: BosStr + FromStr + Ord + Clone,
492 <Str as FromStr>::Err: core::fmt::Debug,
493{
494 async fn get_refreshed<D: BosStr + Send + Sync>(
495 &self,
496 did: &Did<D>,
497 session_id: &str,
498 ) -> Result<ClientSessionData, Error> {
499 let key = format_smolstr!("{}_{}", did, session_id);
500 let lock = self
501 .pending
502 .entry(key)
503 .or_insert_with(|| Arc::new(Mutex::new(())))
504 .clone();
505 let _guard = lock.lock().await;
506
507 let session = self
508 .store
509 .get_session(did, session_id)
510 .await?
511 .ok_or(Error::SessionNotFound)?;
512
513 // Check if token is still valid with a 60-second buffer before expiry.
514 // This triggers proactive refresh before the token actually expires,
515 // avoiding the race condition where a token expires mid-request.
516 const EXPIRY_BUFFER_SECS: i64 = 60;
517 if let Some(expires_at) = &session.token_set.expires_at {
518 let now_with_buffer = Datetime::now()
519 .as_ref()
520 .checked_add_signed(TimeDelta::seconds(EXPIRY_BUFFER_SECS))
521 .map(Datetime::new)
522 .unwrap_or_else(Datetime::now);
523 if expires_at > &now_with_buffer {
524 return Ok(session);
525 }
526 }
527 let metadata =
528 OAuthMetadata::new(self.client.as_ref(), &self.client_data, &session).await?;
529 match refresh(self.client.as_ref(), session, &metadata).await {
530 Ok(refreshed) => {
531 self.store.upsert_session(refreshed.clone()).await?;
532 Ok(refreshed)
533 }
534 Err(e) if e.is_permanent() => {
535 // Session is permanently dead - clean it up
536 let _ = self.store.delete_session(did, session_id).await;
537 Err(Error::RefreshFailed(e))
538 }
539 Err(e) => Err(Error::ServerAgent(e)),
540 }
541 }
542 /// Retrieve a session from the store, optionally refreshing it first.
543 ///
544 /// When `refresh` is `true`, proactively
545 /// renews the token if it is within 60 seconds of expiry. When `false`, returns the session
546 /// data as-is without contacting the authorization server.
547 pub async fn get<D: BosStr + Send + Sync>(
548 &self,
549 did: &Did<D>,
550 session_id: &str,
551 refresh: bool,
552 ) -> Result<ClientSessionData, Error> {
553 if refresh {
554 self.get_refreshed(did, session_id).await
555 } else {
556 // TODO: cached?
557 self.store
558 .get_session(did, session_id)
559 .await?
560 .ok_or(Error::SessionNotFound)
561 }
562 }
563 /// Persist an updated session to the backing store.
564 pub async fn set(&self, value: ClientSessionData) -> Result<(), Error> {
565 self.store.upsert_session(value).await?;
566 Ok(())
567 }
568 /// Delete a session from the backing store.
569 pub async fn del<D: BosStr + Send + Sync>(
570 &self,
571 did: &Did<D>,
572 session_id: &str,
573 ) -> Result<(), Error> {
574 self.store.delete_session(did, session_id).await?;
575 Ok(())
576 }
577}