Skip to main content

mas_storage_pg/personal/
access_token.rs

1// Copyright 2025, 2026 Element Creations Ltd.
2// Copyright 2025 New Vector Ltd.
3//
4// SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-Element-Commercial
5// Please see LICENSE files in the repository root for full details.
6
7use 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
22/// An implementation of [`PersonalAccessTokenRepository`] for a PostgreSQL
23/// connection
24pub struct PgPersonalAccessTokenRepository<'c> {
25    conn: &'c mut PgConnection,
26}
27
28impl<'c> PgPersonalAccessTokenRepository<'c> {
29    /// Create a new [`PgPersonalAccessTokenRepository`] from an active
30    /// PostgreSQL connection
31    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}