1use async_trait::async_trait;
9use chrono::{DateTime, Duration, Utc};
10use mas_data_model::{AccessToken, AccessTokenState, Clock, Session, UlidExt as _};
11use mas_storage::oauth2::OAuth2AccessTokenRepository;
12use rand::Rng;
13use sqlx::PgConnection;
14use ulid::Ulid;
15use uuid::Uuid;
16
17use crate::{DatabaseError, tracing::ExecuteExt};
18
19pub struct PgOAuth2AccessTokenRepository<'c> {
22 conn: &'c mut PgConnection,
23}
24
25impl<'c> PgOAuth2AccessTokenRepository<'c> {
26 pub fn new(conn: &'c mut PgConnection) -> Self {
29 Self { conn }
30 }
31}
32
33struct OAuth2AccessTokenLookup {
34 oauth2_access_token_id: Uuid,
35 oauth2_session_id: Uuid,
36 access_token: String,
37 created_at: DateTime<Utc>,
38 expires_at: Option<DateTime<Utc>>,
39 revoked_at: Option<DateTime<Utc>>,
40 first_used_at: Option<DateTime<Utc>>,
41}
42
43impl From<OAuth2AccessTokenLookup> for AccessToken {
44 fn from(value: OAuth2AccessTokenLookup) -> Self {
45 let state = match value.revoked_at {
46 None => AccessTokenState::Valid,
47 Some(revoked_at) => AccessTokenState::Revoked { revoked_at },
48 };
49
50 Self {
51 id: value.oauth2_access_token_id.into(),
52 state,
53 session_id: value.oauth2_session_id.into(),
54 access_token: value.access_token,
55 created_at: value.created_at,
56 expires_at: value.expires_at,
57 first_used_at: value.first_used_at,
58 }
59 }
60}
61
62#[async_trait]
63impl OAuth2AccessTokenRepository for PgOAuth2AccessTokenRepository<'_> {
64 type Error = DatabaseError;
65
66 async fn lookup(&mut self, id: Ulid) -> Result<Option<AccessToken>, Self::Error> {
67 let res = sqlx::query_as!(
68 OAuth2AccessTokenLookup,
69 r#"
70 SELECT oauth2_access_token_id
71 , access_token
72 , created_at
73 , expires_at
74 , revoked_at
75 , oauth2_session_id
76 , first_used_at
77
78 FROM oauth2_access_tokens
79
80 WHERE oauth2_access_token_id = $1
81 "#,
82 Uuid::from(id),
83 )
84 .fetch_optional(&mut *self.conn)
85 .await?;
86
87 let Some(res) = res else { return Ok(None) };
88
89 Ok(Some(res.into()))
90 }
91
92 #[tracing::instrument(
93 name = "db.oauth2_access_token.find_by_token",
94 skip_all,
95 fields(
96 db.query.text,
97 ),
98 )]
99 async fn find_by_token(
100 &mut self,
101 access_token: &str,
102 ) -> Result<Option<AccessToken>, Self::Error> {
103 let res = sqlx::query_as!(
104 OAuth2AccessTokenLookup,
105 r#"
106 SELECT oauth2_access_token_id
107 , access_token
108 , created_at
109 , expires_at
110 , revoked_at
111 , oauth2_session_id
112 , first_used_at
113
114 FROM oauth2_access_tokens
115
116 WHERE access_token = $1
117 "#,
118 access_token,
119 )
120 .fetch_optional(&mut *self.conn)
121 .await?;
122
123 let Some(res) = res else { return Ok(None) };
124
125 Ok(Some(res.into()))
126 }
127
128 #[tracing::instrument(
129 name = "db.oauth2_access_token.add",
130 skip_all,
131 fields(
132 db.query.text,
133 %session.id,
134 client.id = %session.client_id,
135 access_token.id,
136 ),
137 )]
138 async fn add(
139 &mut self,
140 rng: &mut (dyn Rng + Send),
141 clock: &dyn Clock,
142 session: &Session,
143 access_token: String,
144 expires_after: Option<Duration>,
145 ) -> Result<AccessToken, Self::Error> {
146 let created_at = clock.now();
147 let expires_at = expires_after.map(|d| created_at + d);
148 let id = Ulid::from_datetime_with_rng(created_at, rng);
149
150 tracing::Span::current().record("access_token.id", tracing::field::display(id));
151
152 sqlx::query!(
153 r#"
154 INSERT INTO oauth2_access_tokens
155 (oauth2_access_token_id, oauth2_session_id, access_token, created_at, expires_at)
156 VALUES
157 ($1, $2, $3, $4, $5)
158 "#,
159 Uuid::from(id),
160 Uuid::from(session.id),
161 &access_token,
162 created_at,
163 expires_at,
164 )
165 .traced()
166 .execute(&mut *self.conn)
167 .await?;
168
169 Ok(AccessToken {
170 id,
171 state: AccessTokenState::default(),
172 access_token,
173 session_id: session.id,
174 created_at,
175 expires_at,
176 first_used_at: None,
177 })
178 }
179
180 #[tracing::instrument(
181 name = "db.oauth2_access_token.revoke",
182 skip_all,
183 fields(
184 db.query.text,
185 session.id = %access_token.session_id,
186 %access_token.id,
187 ),
188 )]
189 async fn revoke(
190 &mut self,
191 clock: &dyn Clock,
192 access_token: AccessToken,
193 ) -> Result<AccessToken, Self::Error> {
194 let revoked_at = clock.now();
195 let res = sqlx::query!(
196 r#"
197 UPDATE oauth2_access_tokens
198 SET revoked_at = $2
199 WHERE oauth2_access_token_id = $1
200 "#,
201 Uuid::from(access_token.id),
202 revoked_at,
203 )
204 .traced()
205 .execute(&mut *self.conn)
206 .await?;
207
208 DatabaseError::ensure_affected_rows(&res, 1)?;
209
210 access_token
211 .revoke(revoked_at)
212 .map_err(DatabaseError::to_invalid_operation)
213 }
214
215 #[tracing::instrument(
216 name = "db.oauth2_access_token.mark_used",
217 skip_all,
218 fields(
219 db.query.text,
220 session.id = %access_token.session_id,
221 %access_token.id,
222 ),
223 )]
224 async fn mark_used(
225 &mut self,
226 clock: &dyn Clock,
227 mut access_token: AccessToken,
228 ) -> Result<AccessToken, Self::Error> {
229 let now = clock.now();
230 let res = sqlx::query!(
231 r#"
232 UPDATE oauth2_access_tokens
233 SET first_used_at = $2
234 WHERE oauth2_access_token_id = $1
235 "#,
236 Uuid::from(access_token.id),
237 now,
238 )
239 .execute(&mut *self.conn)
240 .await?;
241
242 DatabaseError::ensure_affected_rows(&res, 1)?;
243
244 access_token.first_used_at = Some(now);
245
246 Ok(access_token)
247 }
248
249 #[tracing::instrument(
250 name = "db.oauth2_access_token.cleanup_revoked",
251 skip_all,
252 fields(
253 db.query.text,
254 since = since.map(tracing::field::display),
255 until = %until,
256 limit = limit,
257 ),
258 )]
259 async fn cleanup_revoked(
260 &mut self,
261 since: Option<DateTime<Utc>>,
262 until: DateTime<Utc>,
263 limit: usize,
264 ) -> Result<(usize, Option<DateTime<Utc>>), Self::Error> {
265 let res = sqlx::query!(
266 r#"
267 WITH
268 to_delete AS (
269 SELECT oauth2_access_token_id
270 FROM oauth2_access_tokens
271 WHERE revoked_at IS NOT NULL
272 AND ($1::timestamptz IS NULL OR revoked_at >= $1::timestamptz)
273 AND revoked_at < $2::timestamptz
274 ORDER BY revoked_at ASC
275 LIMIT $3
276 FOR UPDATE
277 ),
278
279 deleted AS (
280 DELETE FROM oauth2_access_tokens
281 USING to_delete
282 WHERE oauth2_access_tokens.oauth2_access_token_id = to_delete.oauth2_access_token_id
283 RETURNING oauth2_access_tokens.revoked_at
284 )
285
286 SELECT
287 COUNT(*) as "count!",
288 MAX(revoked_at) as last_revoked_at
289 FROM deleted
290 "#,
291 since,
292 until,
293 i64::try_from(limit).unwrap_or(i64::MAX),
294 )
295 .traced()
296 .fetch_one(&mut *self.conn)
297 .await?;
298
299 Ok((
300 res.count.try_into().unwrap_or(usize::MAX),
301 res.last_revoked_at,
302 ))
303 }
304
305 #[tracing::instrument(
306 name = "db.oauth2_access_token.cleanup_expired",
307 skip_all,
308 fields(
309 db.query.text,
310 since = since.map(tracing::field::display),
311 until = %until,
312 limit = limit,
313 ),
314 )]
315 async fn cleanup_expired(
316 &mut self,
317 since: Option<DateTime<Utc>>,
318 until: DateTime<Utc>,
319 limit: usize,
320 ) -> Result<(usize, Option<DateTime<Utc>>), Self::Error> {
321 let res = sqlx::query!(
322 r#"
323 WITH
324 to_delete AS (
325 SELECT oauth2_access_token_id
326 FROM oauth2_access_tokens
327 WHERE expires_at IS NOT NULL
328 AND ($1::timestamptz IS NULL OR expires_at >= $1::timestamptz)
329 AND expires_at < $2::timestamptz
330 ORDER BY expires_at ASC
331 LIMIT $3
332 FOR UPDATE
333 ),
334
335 deleted AS (
336 DELETE FROM oauth2_access_tokens
337 USING to_delete
338 WHERE oauth2_access_tokens.oauth2_access_token_id = to_delete.oauth2_access_token_id
339 RETURNING oauth2_access_tokens.expires_at
340 )
341
342 SELECT
343 COUNT(*) as "count!",
344 MAX(expires_at) as last_expires_at
345 FROM deleted
346 "#,
347 since,
348 until,
349 i64::try_from(limit).unwrap_or(i64::MAX),
350 )
351 .traced()
352 .fetch_one(&mut *self.conn)
353 .await?;
354
355 Ok((
356 res.count.try_into().unwrap_or(usize::MAX),
357 res.last_expires_at,
358 ))
359 }
360}