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 )]
68 async fn lookup(&mut self, id: Ulid) -> Result<Option<PersonalAccessToken>, Self::Error> {
69 let res = sqlx::query_as!(
70 PersonalAccessTokenLookup,
71 r#"
72 SELECT personal_access_token_id
73 , personal_session_id
74 , created_at
75 , expires_at
76 , revoked_at
77
78 FROM personal_access_tokens
79
80 WHERE personal_access_token_id = $1
81 "#,
82 Uuid::from(id),
83 )
84 .traced()
85 .fetch_optional(&mut *self.conn)
86 .await?;
87
88 let Some(res) = res else { return Ok(None) };
89
90 Ok(Some(res.into()))
91 }
92
93 #[tracing::instrument(
94 name = "db.personal_access_token.find_by_token",
95 skip_all,
96 fields(
97 db.query.text,
98 ),
99 )]
100 async fn find_by_token(
101 &mut self,
102 access_token: &str,
103 ) -> Result<Option<PersonalAccessToken>, Self::Error> {
104 let token_sha256 = Sha256::digest(access_token.as_bytes()).to_vec();
105
106 let res = sqlx::query_as!(
107 PersonalAccessTokenLookup,
108 r#"
109 SELECT personal_access_token_id
110 , personal_session_id
111 , created_at
112 , expires_at
113 , revoked_at
114
115 FROM personal_access_tokens
116
117 WHERE access_token_sha256 = $1
118 "#,
119 &token_sha256,
120 )
121 .traced()
122 .fetch_optional(&mut *self.conn)
123 .await?;
124
125 let Some(res) = res else { return Ok(None) };
126
127 Ok(Some(res.into()))
128 }
129
130 #[tracing::instrument(
131 name = "db.personal_access_token.find_active_for_session",
132 skip_all,
133 fields(
134 db.query.text,
135 ),
136 )]
137 async fn find_active_for_session(
138 &mut self,
139 session: &PersonalSession,
140 ) -> Result<Option<PersonalAccessToken>, Self::Error> {
141 let res: Option<PersonalAccessTokenLookup> = sqlx::query_as!(
142 PersonalAccessTokenLookup,
143 r#"
144 SELECT personal_access_token_id
145 , personal_session_id
146 , created_at
147 , expires_at
148 , revoked_at
149
150 FROM personal_access_tokens
151
152 WHERE personal_session_id = $1
153 AND revoked_at IS NULL
154 "#,
155 Uuid::from(session.id),
156 )
157 .traced()
158 .fetch_optional(&mut *self.conn)
159 .await?;
160
161 let Some(res) = res else { return Ok(None) };
162
163 Ok(Some(res.into()))
164 }
165
166 #[tracing::instrument(
167 name = "db.personal_access_token.add",
168 skip_all,
169 fields(
170 db.query.text,
171 personal_access_token.id,
172 %session.id,
173 ),
174 )]
175 async fn add(
176 &mut self,
177 rng: &mut (dyn Rng + Send),
178 clock: &dyn Clock,
179 session: &PersonalSession,
180 access_token: &str,
181 expires_after: Option<chrono::Duration>,
182 ) -> Result<PersonalAccessToken, Self::Error> {
183 let created_at = clock.now();
184 let id = Ulid::from_datetime_with_rng(created_at, rng);
185 tracing::Span::current().record("personal_access_token.id", tracing::field::display(id));
186
187 let token_sha256 = Sha256::digest(access_token.as_bytes()).to_vec();
188
189 let expires_at = expires_after.map(|expires_after| created_at + expires_after);
190
191 sqlx::query!(
192 r#"
193 INSERT INTO personal_access_tokens
194 (personal_access_token_id, personal_session_id, access_token_sha256, created_at, expires_at)
195 VALUES ($1, $2, $3, $4, $5)
196 "#,
197 Uuid::from(id),
198 Uuid::from(session.id),
199 &token_sha256,
200 created_at,
201 expires_at,
202 )
203 .traced()
204 .execute(&mut *self.conn)
205 .await?;
206
207 Ok(PersonalAccessToken {
208 id,
209 session_id: session.id,
210 created_at,
211 expires_at,
212 revoked_at: None,
213 })
214 }
215
216 #[tracing::instrument(
217 name = "db.personal_access_token.revoke",
218 skip_all,
219 fields(
220 db.query.text,
221 %access_token.id,
222 personal_session.id = %access_token.session_id,
223 ),
224 )]
225 async fn revoke(
226 &mut self,
227 clock: &dyn Clock,
228 mut access_token: PersonalAccessToken,
229 ) -> Result<PersonalAccessToken, Self::Error> {
230 let revoked_at = clock.now();
231 let res = sqlx::query!(
232 r#"
233 UPDATE personal_access_tokens
234 SET revoked_at = $2
235 WHERE personal_access_token_id = $1
236 "#,
237 Uuid::from(access_token.id),
238 revoked_at,
239 )
240 .traced()
241 .execute(&mut *self.conn)
242 .await?;
243
244 DatabaseError::ensure_affected_rows(&res, 1)?;
245
246 access_token.revoked_at = Some(revoked_at);
247 Ok(access_token)
248 }
249}