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