mas_storage_pg/compat/
refresh_token.rs

1// Copyright 2025, 2026 Element Creations Ltd.
2// Copyright 2024, 2025 New Vector Ltd.
3// Copyright 2023, 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, Utc};
10use mas_data_model::{
11    Clock, CompatAccessToken, CompatRefreshToken, CompatRefreshTokenState, CompatSession,
12    UlidExt as _,
13};
14use mas_storage::compat::CompatRefreshTokenRepository;
15use rand::Rng;
16use sqlx::PgConnection;
17use ulid::Ulid;
18use uuid::Uuid;
19
20use crate::{DatabaseError, tracing::ExecuteExt};
21
22/// An implementation of [`CompatRefreshTokenRepository`] for a PostgreSQL
23/// connection
24pub struct PgCompatRefreshTokenRepository<'c> {
25    conn: &'c mut PgConnection,
26}
27
28impl<'c> PgCompatRefreshTokenRepository<'c> {
29    /// Create a new [`PgCompatRefreshTokenRepository`] from an active
30    /// PostgreSQL connection
31    pub fn new(conn: &'c mut PgConnection) -> Self {
32        Self { conn }
33    }
34}
35
36struct CompatRefreshTokenLookup {
37    compat_refresh_token_id: Uuid,
38    refresh_token: String,
39    created_at: DateTime<Utc>,
40    consumed_at: Option<DateTime<Utc>>,
41    compat_access_token_id: Uuid,
42    compat_session_id: Uuid,
43}
44
45impl From<CompatRefreshTokenLookup> for CompatRefreshToken {
46    fn from(value: CompatRefreshTokenLookup) -> Self {
47        let state = match value.consumed_at {
48            Some(consumed_at) => CompatRefreshTokenState::Consumed { consumed_at },
49            None => CompatRefreshTokenState::Valid,
50        };
51
52        Self {
53            id: value.compat_refresh_token_id.into(),
54            state,
55            session_id: value.compat_session_id.into(),
56            token: value.refresh_token,
57            created_at: value.created_at,
58            access_token_id: value.compat_access_token_id.into(),
59        }
60    }
61}
62
63#[async_trait]
64impl CompatRefreshTokenRepository for PgCompatRefreshTokenRepository<'_> {
65    type Error = DatabaseError;
66
67    #[tracing::instrument(
68        name = "db.compat_refresh_token.lookup",
69        skip_all,
70        fields(
71            db.query.text,
72            compat_refresh_token.id = %id,
73        ),
74    )]
75    async fn lookup(&mut self, id: Ulid) -> Result<Option<CompatRefreshToken>, Self::Error> {
76        let res = sqlx::query_as!(
77            CompatRefreshTokenLookup,
78            r#"
79                SELECT compat_refresh_token_id
80                     , refresh_token
81                     , created_at
82                     , consumed_at
83                     , compat_session_id
84                     , compat_access_token_id
85
86                FROM compat_refresh_tokens
87
88                WHERE compat_refresh_token_id = $1
89            "#,
90            Uuid::from(id),
91        )
92        .traced()
93        .fetch_optional(&mut *self.conn)
94        .await?;
95
96        let Some(res) = res else { return Ok(None) };
97
98        Ok(Some(res.into()))
99    }
100
101    #[tracing::instrument(
102        name = "db.compat_refresh_token.find_by_token",
103        skip_all,
104        fields(
105            db.query.text,
106        ),
107    )]
108    async fn find_by_token(
109        &mut self,
110        refresh_token: &str,
111    ) -> Result<Option<CompatRefreshToken>, Self::Error> {
112        let res = sqlx::query_as!(
113            CompatRefreshTokenLookup,
114            r#"
115                SELECT compat_refresh_token_id
116                     , refresh_token
117                     , created_at
118                     , consumed_at
119                     , compat_session_id
120                     , compat_access_token_id
121
122                FROM compat_refresh_tokens
123
124                WHERE refresh_token = $1
125            "#,
126            refresh_token,
127        )
128        .traced()
129        .fetch_optional(&mut *self.conn)
130        .await?;
131
132        let Some(res) = res else { return Ok(None) };
133
134        Ok(Some(res.into()))
135    }
136
137    #[tracing::instrument(
138        name = "db.compat_refresh_token.add",
139        skip_all,
140        fields(
141            db.query.text,
142            compat_refresh_token.id,
143            %compat_session.id,
144            user.id = %compat_session.user_id,
145        ),
146    )]
147    async fn add(
148        &mut self,
149        rng: &mut (dyn Rng + Send),
150        clock: &dyn Clock,
151        compat_session: &CompatSession,
152        compat_access_token: &CompatAccessToken,
153        token: String,
154    ) -> Result<CompatRefreshToken, Self::Error> {
155        let created_at = clock.now();
156        let id = Ulid::from_datetime_with_rng(created_at, rng);
157        tracing::Span::current().record("compat_refresh_token.id", tracing::field::display(id));
158
159        sqlx::query!(
160            r#"
161                INSERT INTO compat_refresh_tokens
162                    (compat_refresh_token_id, compat_session_id,
163                     compat_access_token_id, refresh_token, created_at)
164                VALUES ($1, $2, $3, $4, $5)
165            "#,
166            Uuid::from(id),
167            Uuid::from(compat_session.id),
168            Uuid::from(compat_access_token.id),
169            token,
170            created_at,
171        )
172        .traced()
173        .execute(&mut *self.conn)
174        .await?;
175
176        Ok(CompatRefreshToken {
177            id,
178            state: CompatRefreshTokenState::default(),
179            session_id: compat_session.id,
180            access_token_id: compat_access_token.id,
181            token,
182            created_at,
183        })
184    }
185
186    #[tracing::instrument(
187        name = "db.compat_refresh_token.consume_and_replace",
188        skip_all,
189        fields(
190            db.query.text,
191            %compat_refresh_token.id,
192            %successor_compat_refresh_token.id,
193            compat_session.id = %compat_refresh_token.session_id,
194        ),
195    )]
196    async fn consume_and_replace(
197        &mut self,
198        clock: &dyn Clock,
199        compat_refresh_token: CompatRefreshToken,
200        successor_compat_refresh_token: &CompatRefreshToken,
201    ) -> Result<CompatRefreshToken, Self::Error> {
202        if compat_refresh_token.session_id != successor_compat_refresh_token.session_id {
203            return Err(DatabaseError::invalid_operation());
204        }
205
206        let consumed_at = clock.now();
207        let res = sqlx::query!(
208            r#"
209                UPDATE compat_refresh_tokens
210                SET consumed_at = $2
211                WHERE compat_session_id = $1
212                  AND consumed_at IS NULL
213                  AND compat_refresh_token_id <> $3
214            "#,
215            Uuid::from(compat_refresh_token.session_id),
216            consumed_at,
217            Uuid::from(successor_compat_refresh_token.id),
218        )
219        .traced()
220        .execute(&mut *self.conn)
221        .await?;
222
223        // This can affect multiple rows in case we've imported refresh tokens
224        // from Synapse. What we care about is that it at least affected one,
225        // which is what we're checking here
226        if res.rows_affected() == 0 {
227            return Err(DatabaseError::RowsAffected {
228                expected: 1,
229                actual: 0,
230            });
231        }
232
233        let compat_refresh_token = compat_refresh_token
234            .consume(consumed_at)
235            .map_err(DatabaseError::to_invalid_operation)?;
236
237        Ok(compat_refresh_token)
238    }
239}
mas_storage_pg/compat/refresh_token.rs

mas_storage_pg/compat/
refresh_token.rs
