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        err,
75    )]
76    async fn lookup(&mut self, id: Ulid) -> Result<Option<CompatRefreshToken>, Self::Error> {
77        let res = sqlx::query_as!(
78            CompatRefreshTokenLookup,
79            r#"
80                SELECT compat_refresh_token_id
81                     , refresh_token
82                     , created_at
83                     , consumed_at
84                     , compat_session_id
85                     , compat_access_token_id
86
87                FROM compat_refresh_tokens
88
89                WHERE compat_refresh_token_id = $1
90            "#,
91            Uuid::from(id),
92        )
93        .traced()
94        .fetch_optional(&mut *self.conn)
95        .await?;
96
97        let Some(res) = res else { return Ok(None) };
98
99        Ok(Some(res.into()))
100    }
101
102    #[tracing::instrument(
103        name = "db.compat_refresh_token.find_by_token",
104        skip_all,
105        fields(
106            db.query.text,
107        ),
108        err,
109    )]
110    async fn find_by_token(
111        &mut self,
112        refresh_token: &str,
113    ) -> Result<Option<CompatRefreshToken>, Self::Error> {
114        let res = sqlx::query_as!(
115            CompatRefreshTokenLookup,
116            r#"
117                SELECT compat_refresh_token_id
118                     , refresh_token
119                     , created_at
120                     , consumed_at
121                     , compat_session_id
122                     , compat_access_token_id
123
124                FROM compat_refresh_tokens
125
126                WHERE refresh_token = $1
127            "#,
128            refresh_token,
129        )
130        .traced()
131        .fetch_optional(&mut *self.conn)
132        .await?;
133
134        let Some(res) = res else { return Ok(None) };
135
136        Ok(Some(res.into()))
137    }
138
139    #[tracing::instrument(
140        name = "db.compat_refresh_token.add",
141        skip_all,
142        fields(
143            db.query.text,
144            compat_refresh_token.id,
145            %compat_session.id,
146            user.id = %compat_session.user_id,
147        ),
148        err,
149    )]
150    async fn add(
151        &mut self,
152        rng: &mut (dyn Rng + Send),
153        clock: &dyn Clock,
154        compat_session: &CompatSession,
155        compat_access_token: &CompatAccessToken,
156        token: String,
157    ) -> Result<CompatRefreshToken, Self::Error> {
158        let created_at = clock.now();
159        let id = Ulid::from_datetime_with_rng(created_at, rng);
160        tracing::Span::current().record("compat_refresh_token.id", tracing::field::display(id));
161
162        sqlx::query!(
163            r#"
164                INSERT INTO compat_refresh_tokens
165                    (compat_refresh_token_id, compat_session_id,
166                     compat_access_token_id, refresh_token, created_at)
167                VALUES ($1, $2, $3, $4, $5)
168            "#,
169            Uuid::from(id),
170            Uuid::from(compat_session.id),
171            Uuid::from(compat_access_token.id),
172            token,
173            created_at,
174        )
175        .traced()
176        .execute(&mut *self.conn)
177        .await?;
178
179        Ok(CompatRefreshToken {
180            id,
181            state: CompatRefreshTokenState::default(),
182            session_id: compat_session.id,
183            access_token_id: compat_access_token.id,
184            token,
185            created_at,
186        })
187    }
188
189    #[tracing::instrument(
190        name = "db.compat_refresh_token.consume_and_replace",
191        skip_all,
192        fields(
193            db.query.text,
194            %compat_refresh_token.id,
195            %successor_compat_refresh_token.id,
196            compat_session.id = %compat_refresh_token.session_id,
197        ),
198        err,
199    )]
200    async fn consume_and_replace(
201        &mut self,
202        clock: &dyn Clock,
203        compat_refresh_token: CompatRefreshToken,
204        successor_compat_refresh_token: &CompatRefreshToken,
205    ) -> Result<CompatRefreshToken, Self::Error> {
206        if compat_refresh_token.session_id != successor_compat_refresh_token.session_id {
207            return Err(DatabaseError::invalid_operation());
208        }
209
210        let consumed_at = clock.now();
211        let res = sqlx::query!(
212            r#"
213                UPDATE compat_refresh_tokens
214                SET consumed_at = $2
215                WHERE compat_session_id = $1
216                  AND consumed_at IS NULL
217                  AND compat_refresh_token_id <> $3
218            "#,
219            Uuid::from(compat_refresh_token.session_id),
220            consumed_at,
221            Uuid::from(successor_compat_refresh_token.id),
222        )
223        .traced()
224        .execute(&mut *self.conn)
225        .await?;
226
227        // This can affect multiple rows in case we've imported refresh tokens
228        // from Synapse. What we care about is that it at least affected one,
229        // which is what we're checking here
230        if res.rows_affected() == 0 {
231            return Err(DatabaseError::RowsAffected {
232                expected: 1,
233                actual: 0,
234            });
235        }
236
237        let compat_refresh_token = compat_refresh_token
238            .consume(consumed_at)
239            .map_err(DatabaseError::to_invalid_operation)?;
240
241        Ok(compat_refresh_token)
242    }
243}
mas_storage_pg/compat/refresh_token.rs

mas_storage_pg/compat/
refresh_token.rs
