soma_auth/upstream/
store.rs

1//! SQLite-backed rmcp `CredentialStore` and `StateStore` adapters.
2//!
3//! Both adapters are per-`(upstream_name, subject)` — construct a fresh pair for each
4//! upstream × subject combination.  The underlying `SqliteStore` is `Clone` so the
5//! shared connection pool is cheap to duplicate.
6//!
7//! # `StateStore::load` is a consuming take
8//!
9//! `StateStore::load` calls `take_upstream_oauth_state` (DELETE … RETURNING) rather than
10//! a plain SELECT.  This is intentional: consuming state on first read closes the replay
11//! window.  `StateStore::delete` is therefore a no-op — the row is already gone after
12//! a successful `load`.
13//!
14//! # Lifetime pattern
15//!
16//! `#[async_trait]` expands `async fn foo(&self)` to a two-lifetime form:
17//! `fn foo<'life0, 'async_trait>(&'life0 self) where 'life0: 'async_trait, Self: 'async_trait`.
18//! All method impls below match that exact pattern without importing the `async-trait` crate.
19
20use std::future::Future;
21use std::pin::Pin;
22
23use oauth2::{CsrfToken, PkceCodeVerifier, TokenResponse as _};
24use rmcp::transport::auth::{
25    AuthError, CredentialStore, StateStore, StoredAuthorizationState, StoredCredentials,
26};
27use rmcp_client as rmcp;
28
29use crate::sqlite::SqliteStore;
30use crate::types::{UpstreamOauthCredentialRow, UpstreamOauthStateRow};
31use crate::upstream::encryption::{self, EncryptionKey};
32
33fn credential_aad(upstream_name: &str, subject: &str, issuer: &str, client_id: &str) -> Vec<u8> {
34    format!("upstream={upstream_name}\0subject={subject}\0issuer={issuer}\0client_id={client_id}")
35        .into_bytes()
36}
37
38fn now_unix() -> i64 {
39    std::time::SystemTime::now()
40        .duration_since(std::time::UNIX_EPOCH)
41        .unwrap_or_default()
42        .as_secs() as i64
43}
44
45/// Per-`(upstream_name, subject)` credential store backed by SQLite.
46///
47/// Tokens are encrypted at rest with ChaCha20-Poly1305.  Decryption failure
48/// surfaces as `AuthError::AuthorizationRequired` so callers re-initiate the
49/// authorization flow rather than hard-failing.
50pub struct SqliteCredentialStore {
51    store: SqliteStore,
52    key: EncryptionKey,
53    upstream_name: String,
54    subject: String,
55    issuer: String,
56}
57
58impl SqliteCredentialStore {
59    pub fn new(
60        store: SqliteStore,
61        key: EncryptionKey,
62        upstream_name: impl Into<String>,
63        subject: impl Into<String>,
64        issuer: impl Into<String>,
65    ) -> Self {
66        Self {
67            store,
68            key,
69            upstream_name: upstream_name.into(),
70            subject: subject.into(),
71            issuer: issuer.into(),
72        }
73    }
74}
75
76impl CredentialStore for SqliteCredentialStore {
77    fn load<'life0, 'async_trait>(
78        &'life0 self,
79    ) -> Pin<
80        Box<
81            dyn Future<Output = Result<Option<StoredCredentials>, AuthError>> + Send + 'async_trait,
82        >,
83    >
84    where
85        'life0: 'async_trait,
86        Self: 'async_trait,
87    {
88        Box::pin(async move {
89            let row = self
90                .store
91                .find_upstream_oauth_credentials(&self.upstream_name, &self.subject)
92                .await
93                .map_err(|e| AuthError::InternalError(e.to_string()))?;
94
95            let Some(row) = row else {
96                return Ok(None);
97            };
98            if row.issuer.is_empty() || row.issuer != self.issuer {
99                self.store
100                    .delete_upstream_oauth_credentials(&self.upstream_name, &self.subject)
101                    .await
102                    .map_err(|e| AuthError::InternalError(e.to_string()))?;
103                return Ok(None);
104            }
105
106            let aad = credential_aad(
107                &row.upstream_name,
108                &row.subject,
109                &row.issuer,
110                &row.client_id,
111            );
112            let plaintext =
113                encryption::open_with_aad(&self.key, &row.token_blob, &row.token_blob_nonce, &aad)
114                    .map_err(|_| AuthError::AuthorizationRequired)?;
115
116            let creds: StoredCredentials = serde_json::from_slice(&plaintext)
117                .map_err(|e| AuthError::InternalError(format!("deserialize credentials: {e}")))?;
118            if creds.issuer.as_deref() != Some(row.issuer.as_str()) {
119                self.store
120                    .delete_upstream_oauth_credentials(&self.upstream_name, &self.subject)
121                    .await
122                    .map_err(|e| AuthError::InternalError(e.to_string()))?;
123                return Ok(None);
124            }
125
126            Ok(Some(creds))
127        })
128    }
129
130    fn save<'life0, 'async_trait>(
131        &'life0 self,
132        credentials: StoredCredentials,
133    ) -> Pin<Box<dyn Future<Output = Result<(), AuthError>> + Send + 'async_trait>>
134    where
135        'life0: 'async_trait,
136        Self: 'async_trait,
137    {
138        Box::pin(async move {
139            let credentials = credentials.with_issuer(Some(self.issuer.clone()));
140            let token_received_at = credentials
141                .token_received_at
142                .map(|t| t as i64)
143                .unwrap_or_else(now_unix);
144
145            let (access_token_expires_at, refresh_token_present) =
146                if let Some(ref token) = credentials.token_response {
147                    let expires_in = token.expires_in().map(|d| d.as_secs()).unwrap_or(3600) as i64;
148                    (
149                        token_received_at + expires_in,
150                        token.refresh_token().is_some(),
151                    )
152                } else {
153                    (0, false)
154                };
155
156            let granted_scopes_json = serde_json::to_string(&credentials.granted_scopes)
157                .map_err(|e| AuthError::InternalError(format!("serialize scopes: {e}")))?;
158
159            let plaintext = serde_json::to_vec(&credentials)
160                .map_err(|e| AuthError::InternalError(format!("serialize credentials: {e}")))?;
161
162            let aad = credential_aad(
163                &self.upstream_name,
164                &self.subject,
165                &self.issuer,
166                &credentials.client_id,
167            );
168            let (token_blob, token_blob_nonce) =
169                encryption::seal_with_aad(&self.key, &plaintext, &aad)
170                    .map_err(|e| AuthError::InternalError(format!("encrypt credentials: {e}")))?;
171
172            let row = UpstreamOauthCredentialRow {
173                upstream_name: self.upstream_name.clone(),
174                subject: self.subject.clone(),
175                issuer: self.issuer.clone(),
176                client_id: credentials.client_id.clone(),
177                granted_scopes_json,
178                token_blob,
179                token_blob_nonce,
180                token_received_at,
181                access_token_expires_at,
182                refresh_token_present,
183            };
184
185            self.store
186                .upsert_upstream_oauth_credentials(row)
187                .await
188                .map_err(|e| AuthError::InternalError(e.to_string()))
189        })
190    }
191
192    fn clear<'life0, 'async_trait>(
193        &'life0 self,
194    ) -> Pin<Box<dyn Future<Output = Result<(), AuthError>> + Send + 'async_trait>>
195    where
196        'life0: 'async_trait,
197        Self: 'async_trait,
198    {
199        Box::pin(async move {
200            self.store
201                .delete_upstream_oauth_credentials(&self.upstream_name, &self.subject)
202                .await
203                .map_err(|e| AuthError::InternalError(e.to_string()))
204        })
205    }
206}
207
208/// Per-`(upstream_name, subject)` state store backed by SQLite.
209///
210/// The `load` method uses `take_upstream_oauth_state` (atomic DELETE … RETURNING)
211/// rather than a SELECT, consuming the row on first read.  `delete` is a no-op.
212pub struct SqliteStateStore {
213    store: SqliteStore,
214    upstream_name: String,
215    subject: String,
216}
217
218impl SqliteStateStore {
219    pub fn new(
220        store: SqliteStore,
221        upstream_name: impl Into<String>,
222        subject: impl Into<String>,
223    ) -> Self {
224        Self {
225            store,
226            upstream_name: upstream_name.into(),
227            subject: subject.into(),
228        }
229    }
230}
231
232/// TTL for pending authorization state (5 minutes).
233const STATE_TTL_SECS: i64 = 300;
234
235impl StateStore for SqliteStateStore {
236    fn save<'life0, 'life1, 'async_trait>(
237        &'life0 self,
238        csrf_token: &'life1 str,
239        state: StoredAuthorizationState,
240    ) -> Pin<Box<dyn Future<Output = Result<(), AuthError>> + Send + 'async_trait>>
241    where
242        'life0: 'async_trait,
243        'life1: 'async_trait,
244        Self: 'async_trait,
245    {
246        Box::pin(async move {
247            let now = now_unix();
248            let requested_scopes_json =
249                serde_json::to_string(&state.requested_scopes).map_err(|e| {
250                    AuthError::InternalError(format!("serialize requested scopes: {e}"))
251                })?;
252            let row = UpstreamOauthStateRow {
253                upstream_name: self.upstream_name.clone(),
254                subject: self.subject.clone(),
255                csrf_token: csrf_token.to_string(),
256                pkce_verifier: state.pkce_verifier,
257                expected_issuer: state.expected_issuer,
258                require_issuer: state.require_issuer,
259                requested_scopes_json,
260                created_at: now,
261                expires_at: now + STATE_TTL_SECS,
262            };
263            self.store
264                .save_upstream_oauth_state(row)
265                .await
266                .map_err(|e| AuthError::InternalError(e.to_string()))
267        })
268    }
269
270    fn load<'life0, 'life1, 'async_trait>(
271        &'life0 self,
272        csrf_token: &'life1 str,
273    ) -> Pin<
274        Box<
275            dyn Future<Output = Result<Option<StoredAuthorizationState>, AuthError>>
276                + Send
277                + 'async_trait,
278        >,
279    >
280    where
281        'life0: 'async_trait,
282        'life1: 'async_trait,
283        Self: 'async_trait,
284    {
285        Box::pin(async move {
286            let now = now_unix();
287            let row = self
288                .store
289                .take_upstream_oauth_state(&self.upstream_name, &self.subject, csrf_token, now)
290                .await
291                .map_err(|e| AuthError::InternalError(e.to_string()))?;
292
293            row.map(|r| {
294                let requested_scopes =
295                    serde_json::from_str(&r.requested_scopes_json).map_err(|e| {
296                        AuthError::InternalError(format!("deserialize requested scopes: {e}"))
297                    })?;
298                let mut state = StoredAuthorizationState::new_with_expected_issuer(
299                    &PkceCodeVerifier::new(r.pkce_verifier),
300                    &CsrfToken::new(r.csrf_token),
301                    r.expected_issuer,
302                    r.require_issuer,
303                );
304                state.requested_scopes = requested_scopes;
305                Ok(state)
306            })
307            .transpose()
308        })
309    }
310
311    fn delete<'life0, 'life1, 'async_trait>(
312        &'life0 self,
313        _csrf_token: &'life1 str,
314    ) -> Pin<Box<dyn Future<Output = Result<(), AuthError>> + Send + 'async_trait>>
315    where
316        'life0: 'async_trait,
317        'life1: 'async_trait,
318        Self: 'async_trait,
319    {
320        // `load` already performs an atomic DELETE … RETURNING; a separate
321        // delete call would be a double-delete with no effect.
322        Box::pin(async move { Ok(()) })
323    }
324}
325
326// ── T10: OAuth callback/state error-path tests ────────────────────────────────
327//
328// Verify that the `SqliteStateStore` adapter correctly rejects unknown,
329// replayed (already-consumed), and expired state tokens.  These paths were
330// previously un-asserted at the adapter layer; the underlying `SqliteStore`
331// primitive is tested in `soma-auth/src/sqlite.rs` but those tests exercise the
332// raw SQL, not the rmcp `StateStore` adapter behaviour.
333
334#[cfg(test)]
335mod tests {
336    use base64::Engine as _;
337    use oauth2::{CsrfToken, PkceCodeVerifier};
338    use rmcp_client::transport::auth::{
339        CredentialStore, StateStore, StoredAuthorizationState, StoredCredentials,
340    };
341
342    use super::{SqliteCredentialStore, SqliteStateStore};
343    use crate::upstream::encryption::load_key;
344
345    /// Open a disposable in-memory SQLite store for testing.
346    async fn temp_store() -> crate::sqlite::SqliteStore {
347        let path = tempfile::tempdir()
348            .expect("tempdir")
349            .keep()
350            .join("upstream_oauth_test.db");
351        crate::sqlite::SqliteStore::open(path)
352            .await
353            .expect("open test store")
354    }
355
356    fn make_state_store(
357        store: crate::sqlite::SqliteStore,
358        upstream: &str,
359        subject: &str,
360    ) -> SqliteStateStore {
361        SqliteStateStore::new(store, upstream, subject)
362    }
363
364    fn sample_stored_state(csrf: &str) -> StoredAuthorizationState {
365        StoredAuthorizationState::new(
366            &PkceCodeVerifier::new("verifier-value".to_string()),
367            &CsrfToken::new(csrf.to_string()),
368        )
369    }
370
371    /// Loading a state token that was never saved returns `None`.
372    #[tokio::test]
373    async fn unknown_state_returns_none() {
374        let store = make_state_store(temp_store().await, "acme", "alice");
375        let result = store
376            .load("nonexistent-csrf")
377            .await
378            .expect("load should not error");
379        assert!(result.is_none(), "unknown state token must return None");
380    }
381
382    /// Loading a state token a second time returns `None` (replay prevention).
383    ///
384    /// `SqliteStateStore::load` uses an atomic DELETE … RETURNING so the first
385    /// successful load consumes the row.  A subsequent call for the same token
386    /// must return `None` rather than re-authorizing the flow.
387    #[tokio::test]
388    async fn replayed_state_is_rejected() {
389        let sqlite = temp_store().await;
390        let store = make_state_store(sqlite, "acme", "alice");
391        let csrf = "csrf-replay-test";
392
393        // Save state once.
394        store
395            .save(csrf, sample_stored_state(csrf))
396            .await
397            .expect("save should succeed");
398
399        // First load consumes the row.
400        let first = store.load(csrf).await.expect("first load should not error");
401        assert!(first.is_some(), "first load must return the stored state");
402
403        // Second load must return None — the row no longer exists.
404        let second = store
405            .load(csrf)
406            .await
407            .expect("second load should not error");
408        assert!(
409            second.is_none(),
410            "replayed (already-consumed) state token must return None"
411        );
412    }
413
414    /// Loading a state token whose TTL has expired returns `None`.
415    ///
416    /// The underlying `take_upstream_oauth_state` query filters by `expires_at`,
417    /// so an expired row is treated the same as a missing row.  We simulate this
418    /// by writing a row with `expires_at = now - 1` directly via the raw store.
419
420    #[tokio::test]
421    async fn state_round_trip_preserves_rfc9207_issuer_and_requested_scopes() {
422        let sqlite = temp_store().await;
423        let store = make_state_store(sqlite, "acme", "alice");
424        let csrf = "csrf-issuer-round-trip";
425        let mut state = StoredAuthorizationState::new_with_expected_issuer(
426            &PkceCodeVerifier::new("verifier-value".to_string()),
427            &CsrfToken::new(csrf.to_string()),
428            Some("https://issuer.example".to_string()),
429            true,
430        );
431        state.requested_scopes = vec!["mcp.read".to_string(), "mcp.write".to_string()];
432
433        store.save(csrf, state).await.expect("save state");
434        let restored = store
435            .load(csrf)
436            .await
437            .expect("load state")
438            .expect("state exists");
439
440        assert_eq!(
441            restored.expected_issuer.as_deref(),
442            Some("https://issuer.example")
443        );
444        assert!(restored.require_issuer);
445        assert_eq!(restored.requested_scopes, ["mcp.read", "mcp.write"]);
446    }
447
448    fn test_encryption_key() -> crate::upstream::encryption::EncryptionKey {
449        load_key(&base64::engine::general_purpose::STANDARD.encode([7u8; 32])).expect("test key")
450    }
451
452    #[tokio::test]
453    async fn credential_store_rejects_and_deletes_credentials_from_another_issuer() {
454        let sqlite = temp_store().await;
455        let key = test_encryption_key();
456        let issuer_a = SqliteCredentialStore::new(
457            sqlite.clone(),
458            key.clone(),
459            "acme",
460            "alice",
461            "https://issuer-a.example",
462        );
463        issuer_a
464            .save(
465                StoredCredentials::new(
466                    "client-a".to_string(),
467                    None,
468                    vec!["mcp".to_string()],
469                    Some(1_700_000_000),
470                )
471                .with_issuer(Some("https://issuer-a.example".to_string())),
472            )
473            .await
474            .expect("save credentials");
475
476        let issuer_b = SqliteCredentialStore::new(
477            sqlite.clone(),
478            key,
479            "acme",
480            "alice",
481            "https://issuer-b.example",
482        );
483        assert!(
484            issuer_b
485                .load()
486                .await
487                .expect("load mismatched credentials")
488                .is_none(),
489            "credentials minted by another issuer must never be reused"
490        );
491        assert!(
492            sqlite
493                .find_upstream_oauth_credentials("acme", "alice")
494                .await
495                .expect("query credentials")
496                .is_none(),
497            "the stale issuer-bound row must be removed"
498        );
499    }
500
501    #[tokio::test]
502    async fn credential_store_rejects_legacy_rows_without_an_issuer() {
503        let sqlite = temp_store().await;
504        sqlite
505            .upsert_upstream_oauth_credentials(crate::types::UpstreamOauthCredentialRow {
506                upstream_name: "acme".to_string(),
507                subject: "alice".to_string(),
508                issuer: String::new(),
509                client_id: "legacy-client".to_string(),
510                granted_scopes_json: "[]".to_string(),
511                token_blob: vec![1, 2, 3],
512                token_blob_nonce: vec![0; 12],
513                token_received_at: 0,
514                access_token_expires_at: 0,
515                refresh_token_present: false,
516            })
517            .await
518            .expect("insert legacy row");
519        let store = SqliteCredentialStore::new(
520            sqlite.clone(),
521            test_encryption_key(),
522            "acme",
523            "alice",
524            "https://issuer.example",
525        );
526
527        assert!(store.load().await.expect("load legacy row").is_none());
528        assert!(
529            sqlite
530                .find_upstream_oauth_credentials("acme", "alice")
531                .await
532                .expect("query legacy row")
533                .is_none()
534        );
535    }
536
537    #[tokio::test]
538    async fn expired_state_is_rejected() {
539        use crate::types::UpstreamOauthStateRow;
540
541        let sqlite = temp_store().await;
542
543        // Insert an already-expired row directly (bypassing the adapter's TTL).
544        let now = std::time::SystemTime::now()
545            .duration_since(std::time::UNIX_EPOCH)
546            .unwrap_or_default()
547            .as_secs() as i64;
548        let row = UpstreamOauthStateRow {
549            upstream_name: "acme".to_string(),
550            subject: "alice".to_string(),
551            csrf_token: "csrf-expired".to_string(),
552            pkce_verifier: "verifier".to_string(),
553            expected_issuer: Some("https://issuer.example".to_string()),
554            require_issuer: true,
555            requested_scopes_json: "[]".to_string(),
556            created_at: now - 400,
557            expires_at: now - 1, // already expired
558        };
559        sqlite
560            .save_upstream_oauth_state(row)
561            .await
562            .expect("save expired row");
563
564        // The adapter must refuse to return an expired state.
565        let store = make_state_store(sqlite, "acme", "alice");
566        let result = store
567            .load("csrf-expired")
568            .await
569            .expect("load should not error on expired state");
570        assert!(
571            result.is_none(),
572            "expired state token must return None, not the stored state"
573        );
574    }
575}
soma_auth/upstream/store.rs

soma_auth/upstream/
store.rs
