1use std::collections::BTreeMap;
9
10use async_trait::async_trait;
11use chrono::{DateTime, Utc};
12use mas_data_model::{
13 AuthorizationCode, AuthorizationGrant, AuthorizationGrantStage, BrowserSession, Client, Clock,
14 Pkce, Session, UlidExt as _,
15};
16use mas_iana::oauth::PkceCodeChallengeMethod;
17use mas_storage::oauth2::OAuth2AuthorizationGrantRepository;
18use oauth2_types::{requests::ResponseMode, scope::Scope};
19use rand::Rng;
20use sqlx::{PgConnection, types::Json};
21use ulid::Ulid;
22use url::Url;
23use uuid::Uuid;
24
25use crate::{DatabaseError, DatabaseInconsistencyError, tracing::ExecuteExt};
26
27pub struct PgOAuth2AuthorizationGrantRepository<'c> {
30 conn: &'c mut PgConnection,
31}
32
33impl<'c> PgOAuth2AuthorizationGrantRepository<'c> {
34 pub fn new(conn: &'c mut PgConnection) -> Self {
37 Self { conn }
38 }
39}
40
41struct GrantLookup {
42 oauth2_authorization_grant_id: Uuid,
43 created_at: DateTime<Utc>,
44 cancelled_at: Option<DateTime<Utc>>,
45 fulfilled_at: Option<DateTime<Utc>>,
46 exchanged_at: Option<DateTime<Utc>>,
47 scope: String,
48 state: Option<String>,
49 nonce: Option<String>,
50 redirect_uri: String,
51 response_mode: String,
52 response_type_code: bool,
53 response_type_id_token: bool,
54 authorization_code: Option<String>,
55 code_challenge: Option<String>,
56 code_challenge_method: Option<String>,
57 login_hint: Option<String>,
58 locale: Option<String>,
59 raw_parameters: Option<Json<BTreeMap<String, String>>>,
60 oauth2_client_id: Uuid,
61 user_session_id: Option<Uuid>,
62 oauth2_session_id: Option<Uuid>,
63}
64
65impl TryFrom<GrantLookup> for AuthorizationGrant {
66 type Error = DatabaseInconsistencyError;
67
68 fn try_from(value: GrantLookup) -> Result<Self, Self::Error> {
69 let id = value.oauth2_authorization_grant_id.into();
70 let scope: Scope = value.scope.parse().map_err(|e| {
71 DatabaseInconsistencyError::on("oauth2_authorization_grants")
72 .column("scope")
73 .row(id)
74 .source(e)
75 })?;
76
77 let stage = match (
78 value.fulfilled_at,
79 value.exchanged_at,
80 value.cancelled_at,
81 value.user_session_id,
82 value.oauth2_session_id,
83 ) {
84 (None, None, None, None, None) => AuthorizationGrantStage::Pending,
85 (Some(fulfilled_at), None, None, Some(browser_session_id), None) => {
86 AuthorizationGrantStage::Fulfilled {
87 browser_session_id: browser_session_id.into(),
88 fulfilled_at,
89 }
90 }
91 (Some(fulfilled_at), None, None, None, Some(_)) => {
92 tracing::warn!(
93 grant.id = %id,
94 "Grant was fulfilled without a user_session_id, it will be treated as cancelled"
95 );
96 AuthorizationGrantStage::Cancelled {
97 cancelled_at: fulfilled_at,
98 }
99 }
100 (Some(fulfilled_at), Some(exchanged_at), None, _, Some(session_id)) => {
101 AuthorizationGrantStage::Exchanged {
102 session_id: session_id.into(),
103 fulfilled_at,
104 exchanged_at,
105 }
106 }
107 (None, None, Some(cancelled_at), None, None) => {
108 AuthorizationGrantStage::Cancelled { cancelled_at }
109 }
110 _ => {
111 return Err(
112 DatabaseInconsistencyError::on("oauth2_authorization_grants")
113 .column("stage")
114 .row(id),
115 );
116 }
117 };
118
119 let pkce = match (value.code_challenge, value.code_challenge_method) {
120 (Some(challenge), Some(challenge_method)) if challenge_method == "plain" => {
121 Some(Pkce {
122 challenge_method: PkceCodeChallengeMethod::Plain,
123 challenge,
124 })
125 }
126 (Some(challenge), Some(challenge_method)) if challenge_method == "S256" => Some(Pkce {
127 challenge_method: PkceCodeChallengeMethod::S256,
128 challenge,
129 }),
130 (None, None) => None,
131 _ => {
132 return Err(
133 DatabaseInconsistencyError::on("oauth2_authorization_grants")
134 .column("code_challenge_method")
135 .row(id),
136 );
137 }
138 };
139
140 let code: Option<AuthorizationCode> =
141 match (value.response_type_code, value.authorization_code, pkce) {
142 (false, None, None) => None,
143 (true, Some(code), pkce) => Some(AuthorizationCode { code, pkce }),
144 _ => {
145 return Err(
146 DatabaseInconsistencyError::on("oauth2_authorization_grants")
147 .column("authorization_code")
148 .row(id),
149 );
150 }
151 };
152
153 let redirect_uri = value.redirect_uri.parse().map_err(|e| {
154 DatabaseInconsistencyError::on("oauth2_authorization_grants")
155 .column("redirect_uri")
156 .row(id)
157 .source(e)
158 })?;
159
160 let response_mode = value.response_mode.parse().map_err(|e| {
161 DatabaseInconsistencyError::on("oauth2_authorization_grants")
162 .column("response_mode")
163 .row(id)
164 .source(e)
165 })?;
166
167 Ok(AuthorizationGrant {
168 id,
169 stage,
170 client_id: value.oauth2_client_id.into(),
171 code,
172 scope,
173 state: value.state,
174 nonce: value.nonce,
175 response_mode,
176 redirect_uri,
177 created_at: value.created_at,
178 response_type_id_token: value.response_type_id_token,
179 login_hint: value.login_hint,
180 locale: value.locale,
181 raw_parameters: value.raw_parameters.map(|Json(x)| x).unwrap_or_default(),
182 })
183 }
184}
185
186#[async_trait]
187impl OAuth2AuthorizationGrantRepository for PgOAuth2AuthorizationGrantRepository<'_> {
188 type Error = DatabaseError;
189
190 #[tracing::instrument(
191 name = "db.oauth2_authorization_grant.add",
192 skip_all,
193 fields(
194 db.query.text,
195 grant.id,
196 grant.scope = %scope,
197 %client.id,
198 ),
199 )]
200 async fn add(
201 &mut self,
202 rng: &mut (dyn Rng + Send),
203 clock: &dyn Clock,
204 client: &Client,
205 redirect_uri: Url,
206 scope: Scope,
207 code: Option<AuthorizationCode>,
208 state: Option<String>,
209 nonce: Option<String>,
210 response_mode: ResponseMode,
211 response_type_id_token: bool,
212 login_hint: Option<String>,
213 locale: Option<String>,
214 raw_parameters: BTreeMap<String, String>,
215 ) -> Result<AuthorizationGrant, Self::Error> {
216 let code_challenge = code
217 .as_ref()
218 .and_then(|c| c.pkce.as_ref())
219 .map(|p| &p.challenge);
220 let code_challenge_method = code
221 .as_ref()
222 .and_then(|c| c.pkce.as_ref())
223 .map(|p| p.challenge_method.to_string());
224 let code_str = code.as_ref().map(|c| &c.code);
225
226 let created_at = clock.now();
227 let id = Ulid::from_datetime_with_rng(created_at, rng);
228 tracing::Span::current().record("grant.id", tracing::field::display(id));
229
230 sqlx::query!(
231 r#"
232 INSERT INTO oauth2_authorization_grants (
233 oauth2_authorization_grant_id,
234 oauth2_client_id,
235 redirect_uri,
236 scope,
237 state,
238 nonce,
239 response_mode,
240 code_challenge,
241 code_challenge_method,
242 response_type_code,
243 response_type_id_token,
244 authorization_code,
245 login_hint,
246 locale,
247 raw_parameters,
248 created_at
249 )
250 VALUES
251 ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16)
252 "#,
253 Uuid::from(id),
254 Uuid::from(client.id),
255 redirect_uri.to_string(),
256 scope.to_string(),
257 state,
258 nonce,
259 response_mode.to_string(),
260 code_challenge,
261 code_challenge_method,
262 code.is_some(),
263 response_type_id_token,
264 code_str,
265 login_hint,
266 locale,
267 Json(&raw_parameters) as _,
268 created_at,
269 )
270 .traced()
271 .execute(&mut *self.conn)
272 .await?;
273
274 Ok(AuthorizationGrant {
275 id,
276 stage: AuthorizationGrantStage::Pending,
277 code,
278 redirect_uri,
279 client_id: client.id,
280 scope,
281 state,
282 nonce,
283 response_mode,
284 created_at,
285 response_type_id_token,
286 login_hint,
287 locale,
288 raw_parameters,
289 })
290 }
291
292 #[tracing::instrument(
293 name = "db.oauth2_authorization_grant.lookup",
294 skip_all,
295 fields(
296 db.query.text,
297 grant.id = %id,
298 ),
299 )]
300 async fn lookup(&mut self, id: Ulid) -> Result<Option<AuthorizationGrant>, Self::Error> {
301 let res = sqlx::query_as!(
302 GrantLookup,
303 r#"
304 SELECT oauth2_authorization_grant_id
305 , created_at
306 , cancelled_at
307 , fulfilled_at
308 , exchanged_at
309 , scope
310 , state
311 , redirect_uri
312 , response_mode
313 , nonce
314 , oauth2_client_id
315 , authorization_code
316 , response_type_code
317 , response_type_id_token
318 , code_challenge
319 , code_challenge_method
320 , login_hint
321 , locale
322 , raw_parameters AS "raw_parameters: Json<BTreeMap<String, String>>"
323 , user_session_id
324 , oauth2_session_id
325 FROM
326 oauth2_authorization_grants
327
328 WHERE oauth2_authorization_grant_id = $1
329 "#,
330 Uuid::from(id),
331 )
332 .traced()
333 .fetch_optional(&mut *self.conn)
334 .await?;
335
336 let Some(res) = res else { return Ok(None) };
337
338 Ok(Some(res.try_into()?))
339 }
340
341 #[tracing::instrument(
342 name = "db.oauth2_authorization_grant.find_by_code",
343 skip_all,
344 fields(
345 db.query.text,
346 ),
347 )]
348 async fn find_by_code(
349 &mut self,
350 code: &str,
351 ) -> Result<Option<AuthorizationGrant>, Self::Error> {
352 let res = sqlx::query_as!(
353 GrantLookup,
354 r#"
355 SELECT oauth2_authorization_grant_id
356 , created_at
357 , cancelled_at
358 , fulfilled_at
359 , exchanged_at
360 , scope
361 , state
362 , redirect_uri
363 , response_mode
364 , nonce
365 , oauth2_client_id
366 , authorization_code
367 , response_type_code
368 , response_type_id_token
369 , code_challenge
370 , code_challenge_method
371 , login_hint
372 , locale
373 , raw_parameters AS "raw_parameters: Json<BTreeMap<String, String>>"
374 , user_session_id
375 , oauth2_session_id
376 FROM
377 oauth2_authorization_grants
378
379 WHERE authorization_code = $1
380 "#,
381 code,
382 )
383 .traced()
384 .fetch_optional(&mut *self.conn)
385 .await?;
386
387 let Some(res) = res else { return Ok(None) };
388
389 Ok(Some(res.try_into()?))
390 }
391
392 #[tracing::instrument(
393 name = "db.oauth2_authorization_grant.fulfill",
394 skip_all,
395 fields(
396 db.query.text,
397 %grant.id,
398 client.id = %grant.client_id,
399 %browser_session.id,
400 ),
401 )]
402 async fn fulfill(
403 &mut self,
404 clock: &dyn Clock,
405 browser_session: &BrowserSession,
406 grant: AuthorizationGrant,
407 ) -> Result<AuthorizationGrant, Self::Error> {
408 let fulfilled_at = clock.now();
409 let res = sqlx::query!(
410 r#"
411 UPDATE oauth2_authorization_grants
412 SET fulfilled_at = $2
413 , user_session_id = $3
414 WHERE oauth2_authorization_grant_id = $1
415 "#,
416 Uuid::from(grant.id),
417 fulfilled_at,
418 Uuid::from(browser_session.id),
419 )
420 .traced()
421 .execute(&mut *self.conn)
422 .await?;
423
424 DatabaseError::ensure_affected_rows(&res, 1)?;
425
426 let grant = grant
427 .fulfill(fulfilled_at, browser_session)
428 .map_err(DatabaseError::to_invalid_operation)?;
429
430 Ok(grant)
431 }
432
433 #[tracing::instrument(
434 name = "db.oauth2_authorization_grant.exchange",
435 skip_all,
436 fields(
437 db.query.text,
438 %grant.id,
439 client.id = %grant.client_id,
440 %session.id,
441 ),
442 )]
443 async fn exchange(
444 &mut self,
445 clock: &dyn Clock,
446 grant: AuthorizationGrant,
447 session: &Session,
448 ) -> Result<AuthorizationGrant, Self::Error> {
449 let exchanged_at = clock.now();
450 let res = sqlx::query!(
451 r#"
452 UPDATE oauth2_authorization_grants
453 SET exchanged_at = $2
454 , oauth2_session_id = $3
455 WHERE oauth2_authorization_grant_id = $1
456 "#,
457 Uuid::from(grant.id),
458 exchanged_at,
459 Uuid::from(session.id),
460 )
461 .traced()
462 .execute(&mut *self.conn)
463 .await?;
464
465 DatabaseError::ensure_affected_rows(&res, 1)?;
466
467 let grant = grant
468 .exchange(exchanged_at, session)
469 .map_err(DatabaseError::to_invalid_operation)?;
470
471 Ok(grant)
472 }
473
474 #[tracing::instrument(
475 name = "db.oauth2_authorization_grant.cleanup",
476 skip_all,
477 fields(
478 db.query.text,
479 since = since.map(tracing::field::display),
480 until = %until,
481 limit = limit,
482 ),
483 )]
484 async fn cleanup(
485 &mut self,
486 since: Option<Ulid>,
487 until: Ulid,
488 limit: usize,
489 ) -> Result<(usize, Option<Ulid>), Self::Error> {
490 let res = sqlx::query_scalar!(
495 r#"
496 WITH to_delete AS (
497 SELECT oauth2_authorization_grant_id
498 FROM oauth2_authorization_grants
499 WHERE ($1::uuid IS NULL OR oauth2_authorization_grant_id > $1)
500 AND oauth2_authorization_grant_id <= $2
501 ORDER BY oauth2_authorization_grant_id
502 LIMIT $3
503 )
504 DELETE FROM oauth2_authorization_grants
505 USING to_delete
506 WHERE oauth2_authorization_grants.oauth2_authorization_grant_id = to_delete.oauth2_authorization_grant_id
507 RETURNING oauth2_authorization_grants.oauth2_authorization_grant_id
508 "#,
509 since.map(Uuid::from),
510 Uuid::from(until),
511 i64::try_from(limit).unwrap_or(i64::MAX)
512 )
513 .traced()
514 .fetch_all(&mut *self.conn)
515 .await?;
516
517 let count = res.len();
518 let max_id = res.into_iter().max();
519
520 Ok((count, max_id.map(Ulid::from)))
521 }
522}