1use async_trait::async_trait;
8use chrono::{DateTime, Utc};
9use mas_data_model::{
10 Clock, UlidExt as _,
11 personal::{PersonalAccessToken, session::PersonalSession},
12};
13use mas_storage::personal::PersonalAccessTokenRepository;
14use rand::Rng;
15use sha2::{Digest, Sha256};
16use sqlx::PgConnection;
17use ulid::Ulid;
18use uuid::Uuid;
19
20use crate::{DatabaseError, tracing::ExecuteExt as _};
21
22pub struct PgPersonalAccessTokenRepository<'c> {
25 conn: &'c mut PgConnection,
26}
27
28impl<'c> PgPersonalAccessTokenRepository<'c> {
29 pub fn new(conn: &'c mut PgConnection) -> Self {
32 Self { conn }
33 }
34}
35
36struct PersonalAccessTokenLookup {
37 personal_access_token_id: Uuid,
38 personal_session_id: Uuid,
39 created_at: DateTime<Utc>,
40 expires_at: Option<DateTime<Utc>>,
41 revoked_at: Option<DateTime<Utc>>,
42}
43
44impl From<PersonalAccessTokenLookup> for PersonalAccessToken {
45 fn from(value: PersonalAccessTokenLookup) -> Self {
46 Self {
47 id: Ulid::from(value.personal_access_token_id),
48 session_id: Ulid::from(value.personal_session_id),
49 created_at: value.created_at,
50 expires_at: value.expires_at,
51 revoked_at: value.revoked_at,
52 }
53 }
54}
55
56#[async_trait]
57impl PersonalAccessTokenRepository for PgPersonalAccessTokenRepository<'_> {
58 type Error = DatabaseError;
59
60 #[tracing::instrument(
61 name = "db.personal_access_token.lookup",
62 skip_all,
63 fields(
64 db.query.text,
65 personal_access_token.id = %id,
66 ),
67 err,
68 )]
69 async fn lookup(&mut self, id: Ulid) -> Result<Option<PersonalAccessToken>, Self::Error> {
70 let res = sqlx::query_as!(
71 PersonalAccessTokenLookup,
72 r#"
73 SELECT personal_access_token_id
74 , personal_session_id
75 , created_at
76 , expires_at
77 , revoked_at
78
79 FROM personal_access_tokens
80
81 WHERE personal_access_token_id = $1
82 "#,
83 Uuid::from(id),
84 )
85 .traced()
86 .fetch_optional(&mut *self.conn)
87 .await?;
88
89 let Some(res) = res else { return Ok(None) };
90
91 Ok(Some(res.into()))
92 }
93
94 #[tracing::instrument(
95 name = "db.personal_access_token.find_by_token",
96 skip_all,
97 fields(
98 db.query.text,
99 ),
100 err,
101 )]
102 async fn find_by_token(
103 &mut self,
104 access_token: &str,
105 ) -> Result<Option<PersonalAccessToken>, Self::Error> {
106 let token_sha256 = Sha256::digest(access_token.as_bytes()).to_vec();
107
108 let res = sqlx::query_as!(
109 PersonalAccessTokenLookup,
110 r#"
111 SELECT personal_access_token_id
112 , personal_session_id
113 , created_at
114 , expires_at
115 , revoked_at
116
117 FROM personal_access_tokens
118
119 WHERE access_token_sha256 = $1
120 "#,
121 &token_sha256,
122 )
123 .traced()
124 .fetch_optional(&mut *self.conn)
125 .await?;
126
127 let Some(res) = res else { return Ok(None) };
128
129 Ok(Some(res.into()))
130 }
131
132 #[tracing::instrument(
133 name = "db.personal_access_token.find_active_for_session",
134 skip_all,
135 fields(
136 db.query.text,
137 ),
138 err,
139 )]
140 async fn find_active_for_session(
141 &mut self,
142 session: &PersonalSession,
143 ) -> Result<Option<PersonalAccessToken>, Self::Error> {
144 let res: Option<PersonalAccessTokenLookup> = sqlx::query_as!(
145 PersonalAccessTokenLookup,
146 r#"
147 SELECT personal_access_token_id
148 , personal_session_id
149 , created_at
150 , expires_at
151 , revoked_at
152
153 FROM personal_access_tokens
154
155 WHERE personal_session_id = $1
156 AND revoked_at IS NULL
157 "#,
158 Uuid::from(session.id),
159 )
160 .traced()
161 .fetch_optional(&mut *self.conn)
162 .await?;
163
164 let Some(res) = res else { return Ok(None) };
165
166 Ok(Some(res.into()))
167 }
168
169 #[tracing::instrument(
170 name = "db.personal_access_token.add",
171 skip_all,
172 fields(
173 db.query.text,
174 personal_access_token.id,
175 %session.id,
176 ),
177 err,
178 )]
179 async fn add(
180 &mut self,
181 rng: &mut (dyn Rng + Send),
182 clock: &dyn Clock,
183 session: &PersonalSession,
184 access_token: &str,
185 expires_after: Option<chrono::Duration>,
186 ) -> Result<PersonalAccessToken, Self::Error> {
187 let created_at = clock.now();
188 let id = Ulid::from_datetime_with_rng(created_at, rng);
189 tracing::Span::current().record("personal_access_token.id", tracing::field::display(id));
190
191 let token_sha256 = Sha256::digest(access_token.as_bytes()).to_vec();
192
193 let expires_at = expires_after.map(|expires_after| created_at + expires_after);
194
195 sqlx::query!(
196 r#"
197 INSERT INTO personal_access_tokens
198 (personal_access_token_id, personal_session_id, access_token_sha256, created_at, expires_at)
199 VALUES ($1, $2, $3, $4, $5)
200 "#,
201 Uuid::from(id),
202 Uuid::from(session.id),
203 &token_sha256,
204 created_at,
205 expires_at,
206 )
207 .traced()
208 .execute(&mut *self.conn)
209 .await?;
210
211 Ok(PersonalAccessToken {
212 id,
213 session_id: session.id,
214 created_at,
215 expires_at,
216 revoked_at: None,
217 })
218 }
219
220 #[tracing::instrument(
221 name = "db.personal_access_token.revoke",
222 skip_all,
223 fields(
224 db.query.text,
225 %access_token.id,
226 personal_session.id = %access_token.session_id,
227 ),
228 err,
229 )]
230 async fn revoke(
231 &mut self,
232 clock: &dyn Clock,
233 mut access_token: PersonalAccessToken,
234 ) -> Result<PersonalAccessToken, Self::Error> {
235 let revoked_at = clock.now();
236 let res = sqlx::query!(
237 r#"
238 UPDATE personal_access_tokens
239 SET revoked_at = $2
240 WHERE personal_access_token_id = $1
241 "#,
242 Uuid::from(access_token.id),
243 revoked_at,
244 )
245 .traced()
246 .execute(&mut *self.conn)
247 .await?;
248
249 DatabaseError::ensure_affected_rows(&res, 1)?;
250
251 access_token.revoked_at = Some(revoked_at);
252 Ok(access_token)
253 }
254}