Skip to main content

soma_auth/
sqlite_upstream_oauth.rs

1//! Persistence for the three `upstream_oauth_*` tables: long-lived provider
2//! credentials, short-lived per-flow CSRF/PKCE state, and dynamic client
3//! registrations. Split out of `sqlite.rs` to stay under the PATTERNS.md
4//! module size hard limit.
5use rusqlite::{OptionalExtension, params};
6
7use crate::error::AuthError;
8use crate::types::{
9    UpstreamOauthCredentialRow, UpstreamOauthDynamicClientRow, UpstreamOauthStateRow,
10};
11use crate::util::now_unix;
12
13use super::sqlite_rows::{row_to_upstream_oauth_credentials, row_to_upstream_oauth_state};
14use super::{SqliteStore, sqlite_error};
15
16/// Upper bound on how long a pending `upstream_oauth_state` row may live.
17const UPSTREAM_OAUTH_STATE_MAX_TTL_SECS: i64 = 600;
18
19impl SqliteStore {
20    pub async fn upsert_upstream_oauth_credentials(
21        &self,
22        row: UpstreamOauthCredentialRow,
23    ) -> Result<(), AuthError> {
24        self.with_conn(move |conn| {
25            conn.execute(
26                "INSERT INTO upstream_oauth_credentials (
27                    upstream_name, subject, issuer, client_id, granted_scopes_json,
28                    token_blob, token_blob_nonce, token_received_at,
29                    access_token_expires_at, refresh_token_present
30                 ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10)
31                 ON CONFLICT(upstream_name, subject) DO UPDATE SET
32                    issuer = excluded.issuer,
33                    client_id = excluded.client_id,
34                    granted_scopes_json = excluded.granted_scopes_json,
35                    token_blob = excluded.token_blob,
36                    token_blob_nonce = excluded.token_blob_nonce,
37                    token_received_at = excluded.token_received_at,
38                    access_token_expires_at = excluded.access_token_expires_at,
39                    refresh_token_present = excluded.refresh_token_present",
40                params![
41                    row.upstream_name,
42                    row.subject,
43                    row.issuer,
44                    row.client_id,
45                    row.granted_scopes_json,
46                    row.token_blob,
47                    row.token_blob_nonce,
48                    row.token_received_at,
49                    row.access_token_expires_at,
50                    i64::from(row.refresh_token_present),
51                ],
52            )
53            .map_err(sqlite_error)?;
54            Ok(())
55        })
56        .await
57    }
58
59    pub async fn find_upstream_oauth_credentials(
60        &self,
61        upstream_name: &str,
62        subject: &str,
63    ) -> Result<Option<UpstreamOauthCredentialRow>, AuthError> {
64        let upstream_name = upstream_name.to_string();
65        let subject = subject.to_string();
66        self.with_conn(move |conn| {
67            conn.query_row(
68                "SELECT upstream_name, subject, issuer, client_id, granted_scopes_json,
69                        token_blob, token_blob_nonce, token_received_at,
70                        access_token_expires_at, refresh_token_present
71                 FROM upstream_oauth_credentials
72                 WHERE upstream_name = ?1 AND subject = ?2",
73                params![upstream_name, subject],
74                row_to_upstream_oauth_credentials,
75            )
76            .optional()
77            .map_err(sqlite_error)
78        })
79        .await
80    }
81
82    pub async fn delete_upstream_oauth_credentials(
83        &self,
84        upstream_name: &str,
85        subject: &str,
86    ) -> Result<(), AuthError> {
87        let upstream_name = upstream_name.to_string();
88        let subject = subject.to_string();
89        self.with_conn(move |conn| {
90            conn.execute(
91                "DELETE FROM upstream_oauth_credentials
92                 WHERE upstream_name = ?1 AND subject = ?2",
93                params![upstream_name, subject],
94            )
95            .map_err(sqlite_error)?;
96            Ok(())
97        })
98        .await
99    }
100
101    pub async fn save_upstream_oauth_state(
102        &self,
103        row: UpstreamOauthStateRow,
104    ) -> Result<(), AuthError> {
105        if row.expires_at <= row.created_at
106            || row.expires_at - row.created_at > UPSTREAM_OAUTH_STATE_MAX_TTL_SECS
107        {
108            return Err(AuthError::InvalidGrant(
109                "state TTL exceeds 600s".to_string(),
110            ));
111        }
112        self.with_conn(move |conn| {
113            conn.execute(
114                "INSERT INTO upstream_oauth_state (
115                    upstream_name, subject, csrf_token, pkce_verifier,
116                    expected_issuer, require_issuer, requested_scopes_json,
117                    created_at, expires_at
118                 ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9)",
119                params![
120                    row.upstream_name,
121                    row.subject,
122                    row.csrf_token,
123                    row.pkce_verifier,
124                    row.expected_issuer,
125                    i64::from(row.require_issuer),
126                    row.requested_scopes_json,
127                    row.created_at,
128                    row.expires_at,
129                ],
130            )
131            .map_err(sqlite_error)?;
132            Ok(())
133        })
134        .await
135    }
136
137    pub async fn find_upstream_oauth_state_subject(
138        &self,
139        upstream_name: &str,
140        csrf_token: &str,
141        now: i64,
142    ) -> Result<Option<String>, AuthError> {
143        let upstream_name = upstream_name.to_string();
144        let csrf_token = csrf_token.to_string();
145        self.with_conn(move |conn| {
146            conn.query_row(
147                "SELECT subject
148                 FROM upstream_oauth_state
149                 WHERE upstream_name = ?1
150                   AND csrf_token = ?2
151                   AND expires_at > ?3",
152                params![upstream_name, csrf_token, now],
153                |row| row.get(0),
154            )
155            .optional()
156            .map_err(sqlite_error)
157        })
158        .await
159    }
160
161    /// Look up `(upstream_name, subject)` by `csrf_token` alone.
162    ///
163    /// Used by the OAuth callback handler to recover the upstream identity from
164    /// the state parameter without requiring the caller to know it upfront.
165    pub async fn find_upstream_oauth_state_owner(
166        &self,
167        csrf_token: &str,
168        now: i64,
169    ) -> Result<Option<(String, String)>, AuthError> {
170        let csrf_token = csrf_token.to_string();
171        self.with_conn(move |conn| {
172            conn.query_row(
173                "SELECT upstream_name, subject
174                 FROM upstream_oauth_state
175                 WHERE csrf_token = ?1
176                   AND expires_at > ?2",
177                params![csrf_token, now],
178                |row| Ok((row.get::<_, String>(0)?, row.get::<_, String>(1)?)),
179            )
180            .optional()
181            .map_err(sqlite_error)
182        })
183        .await
184    }
185
186    /// Delete a pending OAuth state token by CSRF token to foreclose replay attacks after exchange failure.
187    pub async fn delete_upstream_oauth_state_by_csrf(
188        &self,
189        csrf_token: &str,
190        now: i64,
191    ) -> Result<(), AuthError> {
192        let csrf_token = csrf_token.to_string();
193        self.with_conn(move |conn| {
194            conn.execute(
195                "DELETE FROM upstream_oauth_state
196                 WHERE csrf_token = ?1
197                   AND expires_at > ?2",
198                params![csrf_token, now],
199            )
200            .map_err(sqlite_error)?;
201            Ok(())
202        })
203        .await
204    }
205
206    /// Bind a dynamic OAuth `client_id` to a pending CSRF state row.
207    ///
208    /// Called by `begin_authorization` after generating the authorization URL
209    /// so that `complete_authorization_callback` can later look up which
210    /// `client_id` was used for this specific flow (lab-77y5.15).
211    pub async fn set_upstream_oauth_state_client_id(
212        &self,
213        upstream_name: &str,
214        csrf_token: &str,
215        client_id: &str,
216    ) -> Result<(), AuthError> {
217        let upstream_name = upstream_name.to_string();
218        let csrf_token = csrf_token.to_string();
219        let client_id = client_id.to_string();
220        self.with_conn(move |conn| {
221            conn.execute(
222                "UPDATE upstream_oauth_state
223                 SET dynamic_client_id = ?1
224                 WHERE upstream_name = ?2
225                   AND csrf_token = ?3",
226                params![client_id, upstream_name, csrf_token],
227            )
228            .map_err(sqlite_error)?;
229            Ok(())
230        })
231        .await
232    }
233
234    /// Retrieve the `dynamic_client_id` bound to a pending CSRF state row.
235    ///
236    /// Returns `None` when no row matches or the row has expired. Used by
237    /// `complete_authorization_callback` to recover the exact `client_id` that
238    /// was used when the authorization URL was generated (lab-77y5.15).
239    pub async fn get_upstream_oauth_state_client_id(
240        &self,
241        upstream_name: &str,
242        csrf_token: &str,
243        now: i64,
244    ) -> Result<Option<String>, AuthError> {
245        let upstream_name = upstream_name.to_string();
246        let csrf_token = csrf_token.to_string();
247        self.with_conn(move |conn| {
248            conn.query_row(
249                "SELECT dynamic_client_id
250                 FROM upstream_oauth_state
251                 WHERE upstream_name = ?1
252                   AND csrf_token = ?2
253                   AND expires_at > ?3",
254                params![upstream_name, csrf_token, now],
255                |row| row.get::<_, Option<String>>(0),
256            )
257            .optional()
258            .map(|opt| opt.flatten())
259            .map_err(sqlite_error)
260        })
261        .await
262    }
263
264    /// Atomic take-once via `DELETE ... RETURNING`.
265    pub async fn take_upstream_oauth_state(
266        &self,
267        upstream_name: &str,
268        subject: &str,
269        csrf_token: &str,
270        now: i64,
271    ) -> Result<Option<UpstreamOauthStateRow>, AuthError> {
272        let upstream_name = upstream_name.to_string();
273        let subject = subject.to_string();
274        let csrf_token = csrf_token.to_string();
275        self.with_conn(move |conn| {
276            conn.query_row(
277                "DELETE FROM upstream_oauth_state
278                 WHERE upstream_name = ?1
279                   AND subject = ?2
280                   AND csrf_token = ?3
281                   AND expires_at > ?4
282                 RETURNING upstream_name, subject, csrf_token, pkce_verifier,
283                           expected_issuer, require_issuer, requested_scopes_json,
284                           created_at, expires_at",
285                params![upstream_name, subject, csrf_token, now],
286                row_to_upstream_oauth_state,
287            )
288            .optional()
289            .map_err(sqlite_error)
290        })
291        .await
292    }
293
294    pub async fn save_dynamic_client_registration(
295        &self,
296        upstream_name: &str,
297        subject: &str,
298        client_id: &str,
299        issuer: &str,
300    ) -> Result<(), AuthError> {
301        let upstream_name = upstream_name.to_string();
302        let subject = subject.to_string();
303        let client_id = client_id.to_string();
304        let issuer = issuer.to_string();
305        let now = now_unix();
306        self.with_conn(move |conn| {
307            conn.execute(
308                "INSERT INTO upstream_oauth_dynamic_clients (upstream_name, subject, client_id, issuer, created_at)
309                 VALUES (?1, ?2, ?3, ?4, ?5)
310                 ON CONFLICT(upstream_name, subject) DO UPDATE SET
311                    client_id = excluded.client_id,
312                    issuer = excluded.issuer,
313                    created_at = excluded.created_at",
314                params![upstream_name, subject, client_id, issuer, now],
315            )
316            .map_err(sqlite_error)?;
317            Ok(())
318        })
319        .await
320    }
321
322    pub async fn find_dynamic_client_registration(
323        &self,
324        upstream_name: &str,
325        subject: &str,
326    ) -> Result<Option<UpstreamOauthDynamicClientRow>, AuthError> {
327        let upstream_name = upstream_name.to_string();
328        let subject = subject.to_string();
329        self.with_conn(move |conn| {
330            conn.query_row(
331                "SELECT client_id, issuer
332                 FROM upstream_oauth_dynamic_clients
333                 WHERE upstream_name = ?1 AND subject = ?2",
334                params![upstream_name, subject],
335                |row| {
336                    Ok(UpstreamOauthDynamicClientRow {
337                        client_id: row.get(0)?,
338                        issuer: row.get(1)?,
339                    })
340                },
341            )
342            .optional()
343            .map_err(sqlite_error)
344        })
345        .await
346    }
347
348    pub async fn delete_dynamic_client_registration(
349        &self,
350        upstream_name: &str,
351        subject: &str,
352    ) -> Result<(), AuthError> {
353        let upstream_name = upstream_name.to_string();
354        let subject = subject.to_string();
355        self.with_conn(move |conn| {
356            conn.execute(
357                "DELETE FROM upstream_oauth_dynamic_clients
358                 WHERE upstream_name = ?1 AND subject = ?2",
359                params![upstream_name, subject],
360            )
361            .map_err(sqlite_error)?;
362            Ok(())
363        })
364        .await
365    }
366}