mas_storage_pg/compat/
access_token.rs

1// Copyright 2025, 2026 Element Creations Ltd.
2// Copyright 2024, 2025 New Vector Ltd.
3// Copyright 2022-2024 The Matrix.org Foundation C.I.C.
4//
5// SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-Element-Commercial
6// Please see LICENSE files in the repository root for full details.
7
8use async_trait::async_trait;
9use chrono::{DateTime, Duration, Utc};
10use mas_data_model::{Clock, CompatAccessToken, CompatSession, UlidExt as _};
11use mas_storage::compat::CompatAccessTokenRepository;
12use rand::Rng;
13use sqlx::PgConnection;
14use ulid::Ulid;
15use uuid::Uuid;
16
17use crate::{DatabaseError, tracing::ExecuteExt};
18
19/// An implementation of [`CompatAccessTokenRepository`] for a PostgreSQL
20/// connection
21pub struct PgCompatAccessTokenRepository<'c> {
22    conn: &'c mut PgConnection,
23}
24
25impl<'c> PgCompatAccessTokenRepository<'c> {
26    /// Create a new [`PgCompatAccessTokenRepository`] from an active PostgreSQL
27    /// connection
28    pub fn new(conn: &'c mut PgConnection) -> Self {
29        Self { conn }
30    }
31}
32
33struct CompatAccessTokenLookup {
34    compat_access_token_id: Uuid,
35    access_token: String,
36    created_at: DateTime<Utc>,
37    expires_at: Option<DateTime<Utc>>,
38    compat_session_id: Uuid,
39}
40
41impl From<CompatAccessTokenLookup> for CompatAccessToken {
42    fn from(value: CompatAccessTokenLookup) -> Self {
43        Self {
44            id: value.compat_access_token_id.into(),
45            session_id: value.compat_session_id.into(),
46            token: value.access_token,
47            created_at: value.created_at,
48            expires_at: value.expires_at,
49        }
50    }
51}
52
53#[async_trait]
54impl CompatAccessTokenRepository for PgCompatAccessTokenRepository<'_> {
55    type Error = DatabaseError;
56
57    #[tracing::instrument(
58        name = "db.compat_access_token.lookup",
59        skip_all,
60        fields(
61            db.query.text,
62            compat_session.id = %id,
63        ),
64        err,
65    )]
66    async fn lookup(&mut self, id: Ulid) -> Result<Option<CompatAccessToken>, Self::Error> {
67        let res = sqlx::query_as!(
68            CompatAccessTokenLookup,
69            r#"
70                SELECT compat_access_token_id
71                     , access_token
72                     , created_at
73                     , expires_at
74                     , compat_session_id
75
76                FROM compat_access_tokens
77
78                WHERE compat_access_token_id = $1
79            "#,
80            Uuid::from(id),
81        )
82        .traced()
83        .fetch_optional(&mut *self.conn)
84        .await?;
85
86        let Some(res) = res else { return Ok(None) };
87
88        Ok(Some(res.into()))
89    }
90
91    #[tracing::instrument(
92        name = "db.compat_access_token.find_by_token",
93        skip_all,
94        fields(
95            db.query.text,
96        ),
97        err,
98    )]
99    async fn find_by_token(
100        &mut self,
101        access_token: &str,
102    ) -> Result<Option<CompatAccessToken>, Self::Error> {
103        let res = sqlx::query_as!(
104            CompatAccessTokenLookup,
105            r#"
106                SELECT compat_access_token_id
107                     , access_token
108                     , created_at
109                     , expires_at
110                     , compat_session_id
111
112                FROM compat_access_tokens
113
114                WHERE access_token = $1
115            "#,
116            access_token,
117        )
118        .traced()
119        .fetch_optional(&mut *self.conn)
120        .await?;
121
122        let Some(res) = res else { return Ok(None) };
123
124        Ok(Some(res.into()))
125    }
126
127    #[tracing::instrument(
128        name = "db.compat_access_token.add",
129        skip_all,
130        fields(
131            db.query.text,
132            compat_access_token.id,
133            %compat_session.id,
134            user.id = %compat_session.user_id,
135        ),
136        err,
137    )]
138    async fn add(
139        &mut self,
140        rng: &mut (dyn Rng + Send),
141        clock: &dyn Clock,
142        compat_session: &CompatSession,
143        token: String,
144        expires_after: Option<Duration>,
145    ) -> Result<CompatAccessToken, Self::Error> {
146        let created_at = clock.now();
147        let id = Ulid::from_datetime_with_rng(created_at, rng);
148        tracing::Span::current().record("compat_access_token.id", tracing::field::display(id));
149
150        let expires_at = expires_after.map(|expires_after| created_at + expires_after);
151
152        sqlx::query!(
153            r#"
154                INSERT INTO compat_access_tokens
155                    (compat_access_token_id, compat_session_id, access_token, created_at, expires_at)
156                VALUES ($1, $2, $3, $4, $5)
157            "#,
158            Uuid::from(id),
159            Uuid::from(compat_session.id),
160            token,
161            created_at,
162            expires_at,
163        )
164        .traced()
165        .execute(&mut *self.conn)
166        .await?;
167
168        Ok(CompatAccessToken {
169            id,
170            session_id: compat_session.id,
171            token,
172            created_at,
173            expires_at,
174        })
175    }
176
177    #[tracing::instrument(
178        name = "db.compat_access_token.expire",
179        skip_all,
180        fields(
181            db.query.text,
182            %compat_access_token.id,
183            compat_session.id = %compat_access_token.session_id,
184        ),
185        err,
186    )]
187    async fn expire(
188        &mut self,
189        clock: &dyn Clock,
190        mut compat_access_token: CompatAccessToken,
191    ) -> Result<CompatAccessToken, Self::Error> {
192        let expires_at = clock.now();
193        let res = sqlx::query!(
194            r#"
195                UPDATE compat_access_tokens
196                SET expires_at = $2
197                WHERE compat_access_token_id = $1
198            "#,
199            Uuid::from(compat_access_token.id),
200            expires_at,
201        )
202        .traced()
203        .execute(&mut *self.conn)
204        .await?;
205
206        DatabaseError::ensure_affected_rows(&res, 1)?;
207
208        compat_access_token.expires_at = Some(expires_at);
209        Ok(compat_access_token)
210    }
211}
mas_storage_pg/compat/access_token.rs

mas_storage_pg/compat/
access_token.rs
