From 63f02c4deae5d3dbd05a86b768a1b27b941ac4db Mon Sep 17 00:00:00 2001 From: Quentin Gliech Date: Wed, 21 Jan 2026 12:18:38 +0100 Subject: [PATCH] Track user session authenticated through upstream auth sessions This will help us avoid clearing upstream authorization sessions that might still be useful to keep around for OIDC Backchannel Logouts --- crates/handlers/src/upstream_oauth2/link.rs | 8 +++---- .../src/views/register/steps/finish.rs | 2 +- ...727a305cbe4029cd4cebd5ecc274e3e32f533.json | 16 +++++++++++++ ...9bbe6b23a8861c09a7246dc1659c28f12bf8d.json | 15 ------------ ...3025_upstream_oauth_track_user_session.sql | 11 +++++++++ crates/storage-pg/src/upstream_oauth2/mod.rs | 24 ++++++++++++------- .../storage-pg/src/upstream_oauth2/session.rs | 11 +++++---- crates/storage/src/upstream_oauth2/session.rs | 7 +++++- 8 files changed, 61 insertions(+), 33 deletions(-) create mode 100644 crates/storage-pg/.sqlx/query-5a6b91660e4c12b4a1fe2cad08e727a305cbe4029cd4cebd5ecc274e3e32f533.json delete mode 100644 crates/storage-pg/.sqlx/query-689ffbfc5137ec788e89062ad679bbe6b23a8861c09a7246dc1659c28f12bf8d.json create mode 100644 crates/storage-pg/migrations/20260121103025_upstream_oauth_track_user_session.sql diff --git a/crates/handlers/src/upstream_oauth2/link.rs b/crates/handlers/src/upstream_oauth2/link.rs index ab73520c7..357a7a72f 100644 --- a/crates/handlers/src/upstream_oauth2/link.rs +++ b/crates/handlers/src/upstream_oauth2/link.rs @@ -278,7 +278,7 @@ pub(crate) async fn get( // user. Mark the session as consumed and renew the authentication. let upstream_session = repo .upstream_oauth_session() - .consume(&clock, upstream_session) + .consume(&clock, upstream_session, &session) .await?; repo.browser_session() @@ -358,7 +358,7 @@ pub(crate) async fn get( let upstream_session = repo .upstream_oauth_session() - .consume(&clock, upstream_session) + .consume(&clock, upstream_session, &session) .await?; repo.browser_session() @@ -697,7 +697,7 @@ pub(crate) async fn get( let upstream_session = repo .upstream_oauth_session() - .consume(&clock, upstream_session) + .consume(&clock, upstream_session, &session) .await?; repo.browser_session() @@ -905,7 +905,7 @@ pub(crate) async fn post( let upstream_session = repo .upstream_oauth_session() - .consume(&clock, upstream_session) + .consume(&clock, upstream_session, &session) .await?; repo.browser_session() diff --git a/crates/handlers/src/views/register/steps/finish.rs b/crates/handlers/src/views/register/steps/finish.rs index b1a58c76d..af0b8ef9f 100644 --- a/crates/handlers/src/views/register/steps/finish.rs +++ b/crates/handlers/src/views/register/steps/finish.rs @@ -321,7 +321,7 @@ pub(crate) async fn get( if let Some((upstream_session, upstream_link)) = upstream_oauth { let upstream_session = repo .upstream_oauth_session() - .consume(&clock, upstream_session) + .consume(&clock, upstream_session, &user_session) .await?; repo.upstream_oauth_link() diff --git a/crates/storage-pg/.sqlx/query-5a6b91660e4c12b4a1fe2cad08e727a305cbe4029cd4cebd5ecc274e3e32f533.json b/crates/storage-pg/.sqlx/query-5a6b91660e4c12b4a1fe2cad08e727a305cbe4029cd4cebd5ecc274e3e32f533.json new file mode 100644 index 000000000..e1aa9740e --- /dev/null +++ b/crates/storage-pg/.sqlx/query-5a6b91660e4c12b4a1fe2cad08e727a305cbe4029cd4cebd5ecc274e3e32f533.json @@ -0,0 +1,16 @@ +{ + "db_name": "PostgreSQL", + "query": "\n UPDATE upstream_oauth_authorization_sessions\n SET consumed_at = $1,\n user_session_id = $2\n WHERE upstream_oauth_authorization_session_id = $3\n ", + "describe": { + "columns": [], + "parameters": { + "Left": [ + "Timestamptz", + "Uuid", + "Uuid" + ] + }, + "nullable": [] + }, + "hash": "5a6b91660e4c12b4a1fe2cad08e727a305cbe4029cd4cebd5ecc274e3e32f533" +} diff --git a/crates/storage-pg/.sqlx/query-689ffbfc5137ec788e89062ad679bbe6b23a8861c09a7246dc1659c28f12bf8d.json b/crates/storage-pg/.sqlx/query-689ffbfc5137ec788e89062ad679bbe6b23a8861c09a7246dc1659c28f12bf8d.json deleted file mode 100644 index b122c5a40..000000000 --- a/crates/storage-pg/.sqlx/query-689ffbfc5137ec788e89062ad679bbe6b23a8861c09a7246dc1659c28f12bf8d.json +++ /dev/null @@ -1,15 +0,0 @@ -{ - "db_name": "PostgreSQL", - "query": "\n UPDATE upstream_oauth_authorization_sessions\n SET consumed_at = $1\n WHERE upstream_oauth_authorization_session_id = $2\n ", - "describe": { - "columns": [], - "parameters": { - "Left": [ - "Timestamptz", - "Uuid" - ] - }, - "nullable": [] - }, - "hash": "689ffbfc5137ec788e89062ad679bbe6b23a8861c09a7246dc1659c28f12bf8d" -} diff --git a/crates/storage-pg/migrations/20260121103025_upstream_oauth_track_user_session.sql b/crates/storage-pg/migrations/20260121103025_upstream_oauth_track_user_session.sql new file mode 100644 index 000000000..f71d3280e --- /dev/null +++ b/crates/storage-pg/migrations/20260121103025_upstream_oauth_track_user_session.sql @@ -0,0 +1,11 @@ +-- Copyright 2026 Element Creations Ltd. +-- +-- SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-Element-Commercial +-- Please see LICENSE files in the repository root for full details. + +-- Start tracking the associated `user_session` directly on the authorization session +-- This will be backfilled in a separate migration rolling in the next version +ALTER TABLE upstream_oauth_authorization_sessions + ADD COLUMN user_session_id UUID + REFERENCES user_sessions (user_session_id) + ON DELETE SET NULL; diff --git a/crates/storage-pg/src/upstream_oauth2/mod.rs b/crates/storage-pg/src/upstream_oauth2/mod.rs index d98e840b6..12df9d5f0 100644 --- a/crates/storage-pg/src/upstream_oauth2/mod.rs +++ b/crates/storage-pg/src/upstream_oauth2/mod.rs @@ -167,11 +167,24 @@ mod tests { assert!(!session.is_consumed()); assert_eq!(session.link_id(), Some(link.id)); - let session = repo - .upstream_oauth_session() - .consume(&clock, session) + // We need to create a user and start a browser session to consume the session + let user = repo + .user() + .add(&mut rng, &clock, "john".to_owned()) .await .unwrap(); + let browser_session = repo + .browser_session() + .add(&mut rng, &clock, &user, None) + .await + .unwrap(); + + let session = repo + .upstream_oauth_session() + .consume(&clock, session, &browser_session) + .await + .unwrap(); + // Reload the session let session = repo .upstream_oauth_session() @@ -181,11 +194,6 @@ mod tests { .expect("session to be found in the database"); assert!(session.is_consumed()); - let user = repo - .user() - .add(&mut rng, &clock, "john".to_owned()) - .await - .unwrap(); repo.upstream_oauth_link() .associate_to_user(&link, &user) .await diff --git a/crates/storage-pg/src/upstream_oauth2/session.rs b/crates/storage-pg/src/upstream_oauth2/session.rs index ebd45cb45..4476aea58 100644 --- a/crates/storage-pg/src/upstream_oauth2/session.rs +++ b/crates/storage-pg/src/upstream_oauth2/session.rs @@ -8,8 +8,8 @@ use async_trait::async_trait; use chrono::{DateTime, Utc}; use mas_data_model::{ - Clock, UpstreamOAuthAuthorizationSession, UpstreamOAuthAuthorizationSessionState, - UpstreamOAuthLink, UpstreamOAuthProvider, + BrowserSession, Clock, UpstreamOAuthAuthorizationSession, + UpstreamOAuthAuthorizationSessionState, UpstreamOAuthLink, UpstreamOAuthProvider, }; use mas_storage::{ Page, Pagination, @@ -375,15 +375,18 @@ impl UpstreamOAuthSessionRepository for PgUpstreamOAuthSessionRepository<'_> { &mut self, clock: &dyn Clock, upstream_oauth_authorization_session: UpstreamOAuthAuthorizationSession, + browser_session: &BrowserSession, ) -> Result { let consumed_at = clock.now(); sqlx::query!( r#" UPDATE upstream_oauth_authorization_sessions - SET consumed_at = $1 - WHERE upstream_oauth_authorization_session_id = $2 + SET consumed_at = $1, + user_session_id = $2 + WHERE upstream_oauth_authorization_session_id = $3 "#, consumed_at, + Uuid::from(browser_session.id), Uuid::from(upstream_oauth_authorization_session.id), ) .traced() diff --git a/crates/storage/src/upstream_oauth2/session.rs b/crates/storage/src/upstream_oauth2/session.rs index 04c5a90f1..7b05d11e4 100644 --- a/crates/storage/src/upstream_oauth2/session.rs +++ b/crates/storage/src/upstream_oauth2/session.rs @@ -7,7 +7,8 @@ use async_trait::async_trait; use mas_data_model::{ - Clock, UpstreamOAuthAuthorizationSession, UpstreamOAuthLink, UpstreamOAuthProvider, + BrowserSession, Clock, UpstreamOAuthAuthorizationSession, UpstreamOAuthLink, + UpstreamOAuthProvider, }; use rand_core::RngCore; use ulid::Ulid; @@ -167,6 +168,8 @@ pub trait UpstreamOAuthSessionRepository: Send + Sync { /// /// * `clock`: the clock source /// * `upstream_oauth_authorization_session`: the session to consume + /// * `browser_session`: the browser session that was authenticated with + /// this authorization session /// /// # Errors /// @@ -175,6 +178,7 @@ pub trait UpstreamOAuthSessionRepository: Send + Sync { &mut self, clock: &dyn Clock, upstream_oauth_authorization_session: UpstreamOAuthAuthorizationSession, + browser_session: &BrowserSession, ) -> Result; /// List [`UpstreamOAuthAuthorizationSession`] with the given filter and @@ -262,6 +266,7 @@ repository_impl!(UpstreamOAuthSessionRepository: &mut self, clock: &dyn Clock, upstream_oauth_authorization_session: UpstreamOAuthAuthorizationSession, + browser_session: &BrowserSession, ) -> Result; async fn list(