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::RngCore;
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 err,
200 )]
201 async fn add(
202 &mut self,
203 rng: &mut (dyn RngCore + Send),
204 clock: &dyn Clock,
205 client: &Client,
206 redirect_uri: Url,
207 scope: Scope,
208 code: Option<AuthorizationCode>,
209 state: Option<String>,
210 nonce: Option<String>,
211 response_mode: ResponseMode,
212 response_type_id_token: bool,
213 login_hint: Option<String>,
214 locale: Option<String>,
215 raw_parameters: BTreeMap<String, String>,
216 ) -> Result<AuthorizationGrant, Self::Error> {
217 let code_challenge = code
218 .as_ref()
219 .and_then(|c| c.pkce.as_ref())
220 .map(|p| &p.challenge);
221 let code_challenge_method = code
222 .as_ref()
223 .and_then(|c| c.pkce.as_ref())
224 .map(|p| p.challenge_method.to_string());
225 let code_str = code.as_ref().map(|c| &c.code);
226
227 let created_at = clock.now();
228 let id = Ulid::from_datetime_with_rng(created_at, rng);
229 tracing::Span::current().record("grant.id", tracing::field::display(id));
230
231 sqlx::query!(
232 r#"
233 INSERT INTO oauth2_authorization_grants (
234 oauth2_authorization_grant_id,
235 oauth2_client_id,
236 redirect_uri,
237 scope,
238 state,
239 nonce,
240 response_mode,
241 code_challenge,
242 code_challenge_method,
243 response_type_code,
244 response_type_id_token,
245 authorization_code,
246 login_hint,
247 locale,
248 raw_parameters,
249 created_at
250 )
251 VALUES
252 ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16)
253 "#,
254 Uuid::from(id),
255 Uuid::from(client.id),
256 redirect_uri.to_string(),
257 scope.to_string(),
258 state,
259 nonce,
260 response_mode.to_string(),
261 code_challenge,
262 code_challenge_method,
263 code.is_some(),
264 response_type_id_token,
265 code_str,
266 login_hint,
267 locale,
268 Json(&raw_parameters) as _,
269 created_at,
270 )
271 .traced()
272 .execute(&mut *self.conn)
273 .await?;
274
275 Ok(AuthorizationGrant {
276 id,
277 stage: AuthorizationGrantStage::Pending,
278 code,
279 redirect_uri,
280 client_id: client.id,
281 scope,
282 state,
283 nonce,
284 response_mode,
285 created_at,
286 response_type_id_token,
287 login_hint,
288 locale,
289 raw_parameters,
290 })
291 }
292
293 #[tracing::instrument(
294 name = "db.oauth2_authorization_grant.lookup",
295 skip_all,
296 fields(
297 db.query.text,
298 grant.id = %id,
299 ),
300 err,
301 )]
302 async fn lookup(&mut self, id: Ulid) -> Result<Option<AuthorizationGrant>, Self::Error> {
303 let res = sqlx::query_as!(
304 GrantLookup,
305 r#"
306 SELECT oauth2_authorization_grant_id
307 , created_at
308 , cancelled_at
309 , fulfilled_at
310 , exchanged_at
311 , scope
312 , state
313 , redirect_uri
314 , response_mode
315 , nonce
316 , oauth2_client_id
317 , authorization_code
318 , response_type_code
319 , response_type_id_token
320 , code_challenge
321 , code_challenge_method
322 , login_hint
323 , locale
324 , raw_parameters AS "raw_parameters: Json<BTreeMap<String, String>>"
325 , user_session_id
326 , oauth2_session_id
327 FROM
328 oauth2_authorization_grants
329
330 WHERE oauth2_authorization_grant_id = $1
331 "#,
332 Uuid::from(id),
333 )
334 .traced()
335 .fetch_optional(&mut *self.conn)
336 .await?;
337
338 let Some(res) = res else { return Ok(None) };
339
340 Ok(Some(res.try_into()?))
341 }
342
343 #[tracing::instrument(
344 name = "db.oauth2_authorization_grant.find_by_code",
345 skip_all,
346 fields(
347 db.query.text,
348 ),
349 err,
350 )]
351 async fn find_by_code(
352 &mut self,
353 code: &str,
354 ) -> Result<Option<AuthorizationGrant>, Self::Error> {
355 let res = sqlx::query_as!(
356 GrantLookup,
357 r#"
358 SELECT oauth2_authorization_grant_id
359 , created_at
360 , cancelled_at
361 , fulfilled_at
362 , exchanged_at
363 , scope
364 , state
365 , redirect_uri
366 , response_mode
367 , nonce
368 , oauth2_client_id
369 , authorization_code
370 , response_type_code
371 , response_type_id_token
372 , code_challenge
373 , code_challenge_method
374 , login_hint
375 , locale
376 , raw_parameters AS "raw_parameters: Json<BTreeMap<String, String>>"
377 , user_session_id
378 , oauth2_session_id
379 FROM
380 oauth2_authorization_grants
381
382 WHERE authorization_code = $1
383 "#,
384 code,
385 )
386 .traced()
387 .fetch_optional(&mut *self.conn)
388 .await?;
389
390 let Some(res) = res else { return Ok(None) };
391
392 Ok(Some(res.try_into()?))
393 }
394
395 #[tracing::instrument(
396 name = "db.oauth2_authorization_grant.fulfill",
397 skip_all,
398 fields(
399 db.query.text,
400 %grant.id,
401 client.id = %grant.client_id,
402 %browser_session.id,
403 ),
404 err,
405 )]
406 async fn fulfill(
407 &mut self,
408 clock: &dyn Clock,
409 browser_session: &BrowserSession,
410 grant: AuthorizationGrant,
411 ) -> Result<AuthorizationGrant, Self::Error> {
412 let fulfilled_at = clock.now();
413 let res = sqlx::query!(
414 r#"
415 UPDATE oauth2_authorization_grants
416 SET fulfilled_at = $2
417 , user_session_id = $3
418 WHERE oauth2_authorization_grant_id = $1
419 "#,
420 Uuid::from(grant.id),
421 fulfilled_at,
422 Uuid::from(browser_session.id),
423 )
424 .traced()
425 .execute(&mut *self.conn)
426 .await?;
427
428 DatabaseError::ensure_affected_rows(&res, 1)?;
429
430 let grant = grant
431 .fulfill(fulfilled_at, browser_session)
432 .map_err(DatabaseError::to_invalid_operation)?;
433
434 Ok(grant)
435 }
436
437 #[tracing::instrument(
438 name = "db.oauth2_authorization_grant.exchange",
439 skip_all,
440 fields(
441 db.query.text,
442 %grant.id,
443 client.id = %grant.client_id,
444 %session.id,
445 ),
446 err,
447 )]
448 async fn exchange(
449 &mut self,
450 clock: &dyn Clock,
451 grant: AuthorizationGrant,
452 session: &Session,
453 ) -> Result<AuthorizationGrant, Self::Error> {
454 let exchanged_at = clock.now();
455 let res = sqlx::query!(
456 r#"
457 UPDATE oauth2_authorization_grants
458 SET exchanged_at = $2
459 , oauth2_session_id = $3
460 WHERE oauth2_authorization_grant_id = $1
461 "#,
462 Uuid::from(grant.id),
463 exchanged_at,
464 Uuid::from(session.id),
465 )
466 .traced()
467 .execute(&mut *self.conn)
468 .await?;
469
470 DatabaseError::ensure_affected_rows(&res, 1)?;
471
472 let grant = grant
473 .exchange(exchanged_at, session)
474 .map_err(DatabaseError::to_invalid_operation)?;
475
476 Ok(grant)
477 }
478
479 #[tracing::instrument(
480 name = "db.oauth2_authorization_grant.cleanup",
481 skip_all,
482 fields(
483 db.query.text,
484 since = since.map(tracing::field::display),
485 until = %until,
486 limit = limit,
487 ),
488 err,
489 )]
490 async fn cleanup(
491 &mut self,
492 since: Option<Ulid>,
493 until: Ulid,
494 limit: usize,
495 ) -> Result<(usize, Option<Ulid>), Self::Error> {
496 let res = sqlx::query_scalar!(
501 r#"
502 WITH to_delete AS (
503 SELECT oauth2_authorization_grant_id
504 FROM oauth2_authorization_grants
505 WHERE ($1::uuid IS NULL OR oauth2_authorization_grant_id > $1)
506 AND oauth2_authorization_grant_id <= $2
507 ORDER BY oauth2_authorization_grant_id
508 LIMIT $3
509 )
510 DELETE FROM oauth2_authorization_grants
511 USING to_delete
512 WHERE oauth2_authorization_grants.oauth2_authorization_grant_id = to_delete.oauth2_authorization_grant_id
513 RETURNING oauth2_authorization_grants.oauth2_authorization_grant_id
514 "#,
515 since.map(Uuid::from),
516 Uuid::from(until),
517 i64::try_from(limit).unwrap_or(i64::MAX)
518 )
519 .traced()
520 .fetch_all(&mut *self.conn)
521 .await?;
522
523 let count = res.len();
524 let max_id = res.into_iter().max();
525
526 Ok((count, max_id.map(Ulid::from)))
527 }
528}