1use 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
16const 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 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 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 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 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 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}