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    )]
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}