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    )]
65    async fn lookup(&mut self, id: Ulid) -> Result<Option<CompatAccessToken>, Self::Error> {
66        let res = sqlx::query_as!(
67            CompatAccessTokenLookup,
68            r#"
69                SELECT compat_access_token_id
70                     , access_token
71                     , created_at
72                     , expires_at
73                     , compat_session_id
74
75                FROM compat_access_tokens
76
77                WHERE compat_access_token_id = $1
78            "#,
79            Uuid::from(id),
80        )
81        .traced()
82        .fetch_optional(&mut *self.conn)
83        .await?;
84
85        let Some(res) = res else { return Ok(None) };
86
87        Ok(Some(res.into()))
88    }
89
90    #[tracing::instrument(
91        name = "db.compat_access_token.find_by_token",
92        skip_all,
93        fields(
94            db.query.text,
95        ),
96    )]
97    async fn find_by_token(
98        &mut self,
99        access_token: &str,
100    ) -> Result<Option<CompatAccessToken>, Self::Error> {
101        let res = sqlx::query_as!(
102            CompatAccessTokenLookup,
103            r#"
104                SELECT compat_access_token_id
105                     , access_token
106                     , created_at
107                     , expires_at
108                     , compat_session_id
109
110                FROM compat_access_tokens
111
112                WHERE access_token = $1
113            "#,
114            access_token,
115        )
116        .traced()
117        .fetch_optional(&mut *self.conn)
118        .await?;
119
120        let Some(res) = res else { return Ok(None) };
121
122        Ok(Some(res.into()))
123    }
124
125    #[tracing::instrument(
126        name = "db.compat_access_token.add",
127        skip_all,
128        fields(
129            db.query.text,
130            compat_access_token.id,
131            %compat_session.id,
132            user.id = %compat_session.user_id,
133        ),
134    )]
135    async fn add(
136        &mut self,
137        rng: &mut (dyn Rng + Send),
138        clock: &dyn Clock,
139        compat_session: &CompatSession,
140        token: String,
141        expires_after: Option<Duration>,
142    ) -> Result<CompatAccessToken, Self::Error> {
143        let created_at = clock.now();
144        let id = Ulid::from_datetime_with_rng(created_at, rng);
145        tracing::Span::current().record("compat_access_token.id", tracing::field::display(id));
146
147        let expires_at = expires_after.map(|expires_after| created_at + expires_after);
148
149        sqlx::query!(
150            r#"
151                INSERT INTO compat_access_tokens
152                    (compat_access_token_id, compat_session_id, access_token, created_at, expires_at)
153                VALUES ($1, $2, $3, $4, $5)
154            "#,
155            Uuid::from(id),
156            Uuid::from(compat_session.id),
157            token,
158            created_at,
159            expires_at,
160        )
161        .traced()
162        .execute(&mut *self.conn)
163        .await?;
164
165        Ok(CompatAccessToken {
166            id,
167            session_id: compat_session.id,
168            token,
169            created_at,
170            expires_at,
171        })
172    }
173
174    #[tracing::instrument(
175        name = "db.compat_access_token.expire",
176        skip_all,
177        fields(
178            db.query.text,
179            %compat_access_token.id,
180            compat_session.id = %compat_access_token.session_id,
181        ),
182    )]
183    async fn expire(
184        &mut self,
185        clock: &dyn Clock,
186        mut compat_access_token: CompatAccessToken,
187    ) -> Result<CompatAccessToken, Self::Error> {
188        let expires_at = clock.now();
189        let res = sqlx::query!(
190            r#"
191                UPDATE compat_access_tokens
192                SET expires_at = $2
193                WHERE compat_access_token_id = $1
194            "#,
195            Uuid::from(compat_access_token.id),
196            expires_at,
197        )
198        .traced()
199        .execute(&mut *self.conn)
200        .await?;
201
202        DatabaseError::ensure_affected_rows(&res, 1)?;
203
204        compat_access_token.expires_at = Some(expires_at);
205        Ok(compat_access_token)
206    }
207}
mas_storage_pg/compat/access_token.rs

mas_storage_pg/compat/
access_token.rs
