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