1use async_trait::async_trait;
9use chrono::{DateTime, Utc};
10use mas_data_model::{
11 BrowserSession, Clock, UlidExt as _, UpstreamOAuthAuthorizationSession,
12 UpstreamOAuthAuthorizationSessionState, UpstreamOAuthLink, UpstreamOAuthProvider,
13};
14use mas_storage::{
15 Page, Pagination,
16 pagination::Node,
17 upstream_oauth2::{UpstreamOAuthSessionFilter, UpstreamOAuthSessionRepository},
18};
19use rand::Rng;
20use sea_query::{
21 Expr, ExprTrait, PostgresQueryBuilder, Query, enum_def, extension::postgres::PgExpr,
22};
23use sea_query_sqlx::SqlxBinder;
24use sqlx::PgConnection;
25use ulid::Ulid;
26use uuid::Uuid;
27
28use crate::{
29 DatabaseError, DatabaseInconsistencyError,
30 filter::{Filter, StatementExt},
31 iden::UpstreamOAuthAuthorizationSessions,
32 pagination::QueryBuilderExt,
33 tracing::ExecuteExt,
34};
35
36impl Filter for UpstreamOAuthSessionFilter<'_> {
37 fn generate_condition(&self, _has_joins: bool) -> impl sea_query::IntoCondition {
38 sea_query::Condition::all()
39 .add_option(self.provider().map(|provider| {
40 Expr::col((
41 UpstreamOAuthAuthorizationSessions::Table,
42 UpstreamOAuthAuthorizationSessions::UpstreamOAuthProviderId,
43 ))
44 .eq(Uuid::from(provider.id))
45 }))
46 .add_option(self.sub_claim().map(|sub| {
47 Expr::col((
48 UpstreamOAuthAuthorizationSessions::Table,
49 UpstreamOAuthAuthorizationSessions::IdTokenClaims,
50 ))
51 .cast_json_field("sub")
52 .eq(sub)
53 }))
54 .add_option(self.sid_claim().map(|sid| {
55 Expr::col((
56 UpstreamOAuthAuthorizationSessions::Table,
57 UpstreamOAuthAuthorizationSessions::IdTokenClaims,
58 ))
59 .cast_json_field("sid")
60 .eq(sid)
61 }))
62 }
63}
64
65pub struct PgUpstreamOAuthSessionRepository<'c> {
68 conn: &'c mut PgConnection,
69}
70
71impl<'c> PgUpstreamOAuthSessionRepository<'c> {
72 pub fn new(conn: &'c mut PgConnection) -> Self {
75 Self { conn }
76 }
77}
78
79#[derive(sqlx::FromRow)]
80#[enum_def]
81struct SessionLookup {
82 upstream_oauth_authorization_session_id: Uuid,
83 upstream_oauth_provider_id: Uuid,
84 upstream_oauth_link_id: Option<Uuid>,
85 state: String,
86 code_challenge_verifier: Option<String>,
87 nonce: Option<String>,
88 id_token: Option<String>,
89 id_token_claims: Option<serde_json::Value>,
90 userinfo: Option<serde_json::Value>,
91 created_at: DateTime<Utc>,
92 completed_at: Option<DateTime<Utc>>,
93 consumed_at: Option<DateTime<Utc>>,
94 extra_callback_parameters: Option<serde_json::Value>,
95 unlinked_at: Option<DateTime<Utc>>,
96}
97
98impl Node<Ulid> for SessionLookup {
99 fn cursor(&self) -> Ulid {
100 self.upstream_oauth_authorization_session_id.into()
101 }
102}
103
104impl TryFrom<SessionLookup> for UpstreamOAuthAuthorizationSession {
105 type Error = DatabaseInconsistencyError;
106
107 fn try_from(value: SessionLookup) -> Result<Self, Self::Error> {
108 let id = value.upstream_oauth_authorization_session_id.into();
109 let state = match (
110 value.upstream_oauth_link_id,
111 value.id_token,
112 value.id_token_claims,
113 value.extra_callback_parameters,
114 value.userinfo,
115 value.completed_at,
116 value.consumed_at,
117 value.unlinked_at,
118 ) {
119 (None, None, None, None, None, None, None, None) => {
120 UpstreamOAuthAuthorizationSessionState::Pending
121 }
122 (
123 Some(link_id),
124 id_token,
125 id_token_claims,
126 extra_callback_parameters,
127 userinfo,
128 Some(completed_at),
129 None,
130 None,
131 ) => UpstreamOAuthAuthorizationSessionState::Completed {
132 completed_at,
133 link_id: link_id.into(),
134 id_token,
135 id_token_claims,
136 extra_callback_parameters,
137 userinfo,
138 },
139 (
140 Some(link_id),
141 id_token,
142 id_token_claims,
143 extra_callback_parameters,
144 userinfo,
145 Some(completed_at),
146 Some(consumed_at),
147 None,
148 ) => UpstreamOAuthAuthorizationSessionState::Consumed {
149 completed_at,
150 link_id: link_id.into(),
151 id_token,
152 id_token_claims,
153 extra_callback_parameters,
154 userinfo,
155 consumed_at,
156 },
157 (
158 _,
159 id_token,
160 id_token_claims,
161 _,
162 _,
163 Some(completed_at),
164 consumed_at,
165 Some(unlinked_at),
166 ) => UpstreamOAuthAuthorizationSessionState::Unlinked {
167 completed_at,
168 id_token,
169 id_token_claims,
170 consumed_at,
171 unlinked_at,
172 },
173 _ => {
174 return Err(DatabaseInconsistencyError::on(
175 "upstream_oauth_authorization_sessions",
176 )
177 .row(id));
178 }
179 };
180
181 Ok(Self {
182 id,
183 provider_id: value.upstream_oauth_provider_id.into(),
184 state_str: value.state,
185 nonce: value.nonce,
186 code_challenge_verifier: value.code_challenge_verifier,
187 created_at: value.created_at,
188 state,
189 })
190 }
191}
192
193#[async_trait]
194impl UpstreamOAuthSessionRepository for PgUpstreamOAuthSessionRepository<'_> {
195 type Error = DatabaseError;
196
197 #[tracing::instrument(
198 name = "db.upstream_oauth_authorization_session.lookup",
199 skip_all,
200 fields(
201 db.query.text,
202 upstream_oauth_provider.id = %id,
203 ),
204 )]
205 async fn lookup(
206 &mut self,
207 id: Ulid,
208 ) -> Result<Option<UpstreamOAuthAuthorizationSession>, Self::Error> {
209 let res = sqlx::query_as!(
210 SessionLookup,
211 r#"
212 SELECT
213 upstream_oauth_authorization_session_id,
214 upstream_oauth_provider_id,
215 upstream_oauth_link_id,
216 state,
217 code_challenge_verifier,
218 nonce,
219 id_token,
220 id_token_claims,
221 extra_callback_parameters,
222 userinfo,
223 created_at,
224 completed_at,
225 consumed_at,
226 unlinked_at
227 FROM upstream_oauth_authorization_sessions
228 WHERE upstream_oauth_authorization_session_id = $1
229 "#,
230 Uuid::from(id),
231 )
232 .traced()
233 .fetch_optional(&mut *self.conn)
234 .await?;
235
236 let Some(res) = res else { return Ok(None) };
237
238 Ok(Some(res.try_into()?))
239 }
240
241 #[tracing::instrument(
242 name = "db.upstream_oauth_authorization_session.add",
243 skip_all,
244 fields(
245 db.query.text,
246 %upstream_oauth_provider.id,
247 upstream_oauth_provider.issuer = upstream_oauth_provider.issuer,
248 %upstream_oauth_provider.client_id,
249 upstream_oauth_authorization_session.id,
250 ),
251 )]
252 async fn add(
253 &mut self,
254 rng: &mut (dyn Rng + Send),
255 clock: &dyn Clock,
256 upstream_oauth_provider: &UpstreamOAuthProvider,
257 state_str: String,
258 code_challenge_verifier: Option<String>,
259 nonce: Option<String>,
260 ) -> Result<UpstreamOAuthAuthorizationSession, Self::Error> {
261 let created_at = clock.now();
262 let id = Ulid::from_datetime_with_rng(created_at, rng);
263 tracing::Span::current().record(
264 "upstream_oauth_authorization_session.id",
265 tracing::field::display(id),
266 );
267
268 sqlx::query!(
269 r#"
270 INSERT INTO upstream_oauth_authorization_sessions (
271 upstream_oauth_authorization_session_id,
272 upstream_oauth_provider_id,
273 state,
274 code_challenge_verifier,
275 nonce,
276 created_at,
277 completed_at,
278 consumed_at,
279 id_token,
280 userinfo
281 ) VALUES ($1, $2, $3, $4, $5, $6, NULL, NULL, NULL, NULL)
282 "#,
283 Uuid::from(id),
284 Uuid::from(upstream_oauth_provider.id),
285 &state_str,
286 code_challenge_verifier.as_deref(),
287 nonce,
288 created_at,
289 )
290 .traced()
291 .execute(&mut *self.conn)
292 .await?;
293
294 Ok(UpstreamOAuthAuthorizationSession {
295 id,
296 state: UpstreamOAuthAuthorizationSessionState::default(),
297 provider_id: upstream_oauth_provider.id,
298 state_str,
299 code_challenge_verifier,
300 nonce,
301 created_at,
302 })
303 }
304
305 #[tracing::instrument(
306 name = "db.upstream_oauth_authorization_session.complete_with_link",
307 skip_all,
308 fields(
309 db.query.text,
310 %upstream_oauth_authorization_session.id,
311 %upstream_oauth_link.id,
312 ),
313 )]
314 async fn complete_with_link(
315 &mut self,
316 clock: &dyn Clock,
317 upstream_oauth_authorization_session: UpstreamOAuthAuthorizationSession,
318 upstream_oauth_link: &UpstreamOAuthLink,
319 id_token: Option<String>,
320 id_token_claims: Option<serde_json::Value>,
321 extra_callback_parameters: Option<serde_json::Value>,
322 userinfo: Option<serde_json::Value>,
323 ) -> Result<UpstreamOAuthAuthorizationSession, Self::Error> {
324 let completed_at = clock.now();
325
326 sqlx::query!(
327 r#"
328 UPDATE upstream_oauth_authorization_sessions
329 SET upstream_oauth_link_id = $1
330 , completed_at = $2
331 , id_token = $3
332 , id_token_claims = $4
333 , extra_callback_parameters = $5
334 , userinfo = $6
335 WHERE upstream_oauth_authorization_session_id = $7
336 "#,
337 Uuid::from(upstream_oauth_link.id),
338 completed_at,
339 id_token,
340 id_token_claims,
341 extra_callback_parameters,
342 userinfo,
343 Uuid::from(upstream_oauth_authorization_session.id),
344 )
345 .traced()
346 .execute(&mut *self.conn)
347 .await?;
348
349 let upstream_oauth_authorization_session = upstream_oauth_authorization_session
350 .complete(
351 completed_at,
352 upstream_oauth_link,
353 id_token,
354 id_token_claims,
355 extra_callback_parameters,
356 userinfo,
357 )
358 .map_err(DatabaseError::to_invalid_operation)?;
359
360 Ok(upstream_oauth_authorization_session)
361 }
362
363 #[tracing::instrument(
365 name = "db.upstream_oauth_authorization_session.consume",
366 skip_all,
367 fields(
368 db.query.text,
369 %upstream_oauth_authorization_session.id,
370 ),
371 )]
372 async fn consume(
373 &mut self,
374 clock: &dyn Clock,
375 upstream_oauth_authorization_session: UpstreamOAuthAuthorizationSession,
376 browser_session: &BrowserSession,
377 ) -> Result<UpstreamOAuthAuthorizationSession, Self::Error> {
378 let consumed_at = clock.now();
379 sqlx::query!(
380 r#"
381 UPDATE upstream_oauth_authorization_sessions
382 SET consumed_at = $1,
383 user_session_id = $2
384 WHERE upstream_oauth_authorization_session_id = $3
385 "#,
386 consumed_at,
387 Uuid::from(browser_session.id),
388 Uuid::from(upstream_oauth_authorization_session.id),
389 )
390 .traced()
391 .execute(&mut *self.conn)
392 .await?;
393
394 let upstream_oauth_authorization_session = upstream_oauth_authorization_session
395 .consume(consumed_at)
396 .map_err(DatabaseError::to_invalid_operation)?;
397
398 Ok(upstream_oauth_authorization_session)
399 }
400
401 #[tracing::instrument(
402 name = "db.upstream_oauth_authorization_session.list",
403 skip_all,
404 fields(
405 db.query.text,
406 ),
407 )]
408 async fn list(
409 &mut self,
410 filter: UpstreamOAuthSessionFilter<'_>,
411 pagination: Pagination,
412 ) -> Result<Page<UpstreamOAuthAuthorizationSession>, Self::Error> {
413 let (sql, arguments) = Query::select()
414 .expr_as(
415 Expr::col((
416 UpstreamOAuthAuthorizationSessions::Table,
417 UpstreamOAuthAuthorizationSessions::UpstreamOAuthAuthorizationSessionId,
418 )),
419 SessionLookupIden::UpstreamOauthAuthorizationSessionId,
420 )
421 .expr_as(
422 Expr::col((
423 UpstreamOAuthAuthorizationSessions::Table,
424 UpstreamOAuthAuthorizationSessions::UpstreamOAuthProviderId,
425 )),
426 SessionLookupIden::UpstreamOauthProviderId,
427 )
428 .expr_as(
429 Expr::col((
430 UpstreamOAuthAuthorizationSessions::Table,
431 UpstreamOAuthAuthorizationSessions::UpstreamOAuthLinkId,
432 )),
433 SessionLookupIden::UpstreamOauthLinkId,
434 )
435 .expr_as(
436 Expr::col((
437 UpstreamOAuthAuthorizationSessions::Table,
438 UpstreamOAuthAuthorizationSessions::State,
439 )),
440 SessionLookupIden::State,
441 )
442 .expr_as(
443 Expr::col((
444 UpstreamOAuthAuthorizationSessions::Table,
445 UpstreamOAuthAuthorizationSessions::CodeChallengeVerifier,
446 )),
447 SessionLookupIden::CodeChallengeVerifier,
448 )
449 .expr_as(
450 Expr::col((
451 UpstreamOAuthAuthorizationSessions::Table,
452 UpstreamOAuthAuthorizationSessions::Nonce,
453 )),
454 SessionLookupIden::Nonce,
455 )
456 .expr_as(
457 Expr::col((
458 UpstreamOAuthAuthorizationSessions::Table,
459 UpstreamOAuthAuthorizationSessions::IdToken,
460 )),
461 SessionLookupIden::IdToken,
462 )
463 .expr_as(
464 Expr::col((
465 UpstreamOAuthAuthorizationSessions::Table,
466 UpstreamOAuthAuthorizationSessions::IdTokenClaims,
467 )),
468 SessionLookupIden::IdTokenClaims,
469 )
470 .expr_as(
471 Expr::col((
472 UpstreamOAuthAuthorizationSessions::Table,
473 UpstreamOAuthAuthorizationSessions::ExtraCallbackParameters,
474 )),
475 SessionLookupIden::ExtraCallbackParameters,
476 )
477 .expr_as(
478 Expr::col((
479 UpstreamOAuthAuthorizationSessions::Table,
480 UpstreamOAuthAuthorizationSessions::Userinfo,
481 )),
482 SessionLookupIden::Userinfo,
483 )
484 .expr_as(
485 Expr::col((
486 UpstreamOAuthAuthorizationSessions::Table,
487 UpstreamOAuthAuthorizationSessions::CreatedAt,
488 )),
489 SessionLookupIden::CreatedAt,
490 )
491 .expr_as(
492 Expr::col((
493 UpstreamOAuthAuthorizationSessions::Table,
494 UpstreamOAuthAuthorizationSessions::CompletedAt,
495 )),
496 SessionLookupIden::CompletedAt,
497 )
498 .expr_as(
499 Expr::col((
500 UpstreamOAuthAuthorizationSessions::Table,
501 UpstreamOAuthAuthorizationSessions::ConsumedAt,
502 )),
503 SessionLookupIden::ConsumedAt,
504 )
505 .expr_as(
506 Expr::col((
507 UpstreamOAuthAuthorizationSessions::Table,
508 UpstreamOAuthAuthorizationSessions::UnlinkedAt,
509 )),
510 SessionLookupIden::UnlinkedAt,
511 )
512 .from(UpstreamOAuthAuthorizationSessions::Table)
513 .apply_filter(filter)
514 .generate_pagination(
515 (
516 UpstreamOAuthAuthorizationSessions::Table,
517 UpstreamOAuthAuthorizationSessions::UpstreamOAuthAuthorizationSessionId,
518 ),
519 pagination,
520 )
521 .build_sqlx(PostgresQueryBuilder);
522
523 let edges: Vec<SessionLookup> = sqlx::query_as_with(&sql, arguments)
524 .traced()
525 .fetch_all(&mut *self.conn)
526 .await?;
527
528 let page = pagination
529 .process(edges)
530 .try_map(UpstreamOAuthAuthorizationSession::try_from)?;
531
532 Ok(page)
533 }
534
535 #[tracing::instrument(
536 name = "db.upstream_oauth_authorization_session.count",
537 skip_all,
538 fields(
539 db.query.text,
540 ),
541 )]
542 async fn count(
543 &mut self,
544 filter: UpstreamOAuthSessionFilter<'_>,
545 ) -> Result<usize, Self::Error> {
546 let (sql, arguments) = Query::select()
547 .expr(
548 Expr::col((
549 UpstreamOAuthAuthorizationSessions::Table,
550 UpstreamOAuthAuthorizationSessions::UpstreamOAuthAuthorizationSessionId,
551 ))
552 .count(),
553 )
554 .from(UpstreamOAuthAuthorizationSessions::Table)
555 .apply_filter(filter)
556 .build_sqlx(PostgresQueryBuilder);
557
558 let count: i64 = sqlx::query_scalar_with(&sql, arguments)
559 .traced()
560 .fetch_one(&mut *self.conn)
561 .await?;
562
563 count
564 .try_into()
565 .map_err(DatabaseError::to_invalid_operation)
566 }
567
568 #[tracing::instrument(
569 name = "db.upstream_oauth_authorization_session.cleanup",
570 skip_all,
571 fields(
572 db.query.text,
573 since = since.map(tracing::field::display),
574 until = %until,
575 limit = limit,
576 ),
577 )]
578 async fn cleanup_orphaned(
579 &mut self,
580 since: Option<Ulid>,
581 until: Ulid,
582 limit: usize,
583 ) -> Result<(usize, Option<Ulid>), Self::Error> {
584 let res = sqlx::query_scalar!(
588 r#"
589 WITH to_delete AS (
590 SELECT upstream_oauth_authorization_session_id
591 FROM upstream_oauth_authorization_sessions
592 WHERE ($1::uuid IS NULL OR upstream_oauth_authorization_session_id > $1)
593 AND upstream_oauth_authorization_session_id <= $2
594 AND user_session_id IS NULL
595 ORDER BY upstream_oauth_authorization_session_id
596 LIMIT $3
597 )
598 DELETE FROM upstream_oauth_authorization_sessions
599 USING to_delete
600 WHERE upstream_oauth_authorization_sessions.upstream_oauth_authorization_session_id = to_delete.upstream_oauth_authorization_session_id
601 RETURNING upstream_oauth_authorization_sessions.upstream_oauth_authorization_session_id
602 "#,
603 since.map(Uuid::from),
604 Uuid::from(until),
605 i64::try_from(limit).unwrap_or(i64::MAX)
606 )
607 .traced()
608 .fetch_all(&mut *self.conn)
609 .await?;
610
611 let count = res.len();
612 let max_id = res.into_iter().max();
613
614 Ok((count, max_id.map(Ulid::from)))
615 }
616}