From f171d76dc555718086a47fe957dda1751e7c8fd4 Mon Sep 17 00:00:00 2001 From: Quentin Gliech Date: Thu, 22 Feb 2024 10:01:32 +0100 Subject: [PATCH] Record user agents on OAuth 2.0 and compat sessions (#2386) * Record user agents on OAuth 2.0 and compat sessions * Add tests for recording user agent in sessions --- crates/data-model/src/compat/session.rs | 1 + crates/data-model/src/oauth2/session.rs | 1 + crates/handlers/src/compat/login.rs | 13 ++++- crates/handlers/src/oauth2/token.rs | 51 +++++++++++++++++-- ...3353004f458c85f7b4f361802f86651900fbc.json | 15 ++++++ ...e3847e376e443ccd841f76b17a81f53fafc3a.json | 15 ++++++ ...72b36ab46a4377b46618205151ea041886d5.json} | 12 +++-- ...6b4d308302bb1fa3accb12932d1e5ce457e9.json} | 12 +++-- .../20240221164945_sessions_user_agent.sql | 17 +++++++ crates/storage-pg/src/app_session.rs | 12 +++++ crates/storage-pg/src/compat/mod.rs | 18 +++++++ crates/storage-pg/src/compat/session.rs | 44 ++++++++++++++++ crates/storage-pg/src/iden.rs | 2 + crates/storage-pg/src/oauth2/mod.rs | 18 +++++++ crates/storage-pg/src/oauth2/session.rs | 45 ++++++++++++++++ crates/storage/src/compat/session.rs | 22 ++++++++ crates/storage/src/oauth2/session.rs | 18 +++++++ 17 files changed, 303 insertions(+), 13 deletions(-) create mode 100644 crates/storage-pg/.sqlx/query-1919d402fd6f148d14417f633be3353004f458c85f7b4f361802f86651900fbc.json create mode 100644 crates/storage-pg/.sqlx/query-29148548d592046f7d711676911e3847e376e443ccd841f76b17a81f53fafc3a.json rename crates/storage-pg/.sqlx/{query-31aace373b20b5dbf65fa51d8663da7571d85b6a7d2d544d69e7d04260cdffc9.json => query-5a2e9b5002c1927c0035c22e393172b36ab46a4377b46618205151ea041886d5.json} (76%) rename crates/storage-pg/.sqlx/{query-04e25c9267bf2eb143a6445345229081e7b386743a93b3833ef8ad9d09972f3b.json => query-bb6f55a4cc10bec8ec0fc138485f6b4d308302bb1fa3accb12932d1e5ce457e9.json} (76%) create mode 100644 crates/storage-pg/migrations/20240221164945_sessions_user_agent.sql diff --git a/crates/data-model/src/compat/session.rs b/crates/data-model/src/compat/session.rs index 76516072..9d4582f7 100644 --- a/crates/data-model/src/compat/session.rs +++ b/crates/data-model/src/compat/session.rs @@ -83,6 +83,7 @@ pub struct CompatSession { pub user_session_id: Option, pub created_at: DateTime, pub is_synapse_admin: bool, + pub user_agent: Option, pub last_active_at: Option>, pub last_active_ip: Option, } diff --git a/crates/data-model/src/oauth2/session.rs b/crates/data-model/src/oauth2/session.rs index 0d17ea6d..fb27ab0a 100644 --- a/crates/data-model/src/oauth2/session.rs +++ b/crates/data-model/src/oauth2/session.rs @@ -75,6 +75,7 @@ pub struct Session { pub user_session_id: Option, pub client_id: Ulid, pub scope: Scope, + pub user_agent: Option, pub last_active_at: Option>, pub last_active_ip: Option, } diff --git a/crates/handlers/src/compat/login.rs b/crates/handlers/src/compat/login.rs index a3363b6c..5f27123c 100644 --- a/crates/handlers/src/compat/login.rs +++ b/crates/handlers/src/compat/login.rs @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use axum::{extract::State, response::IntoResponse, Json}; +use axum::{extract::State, response::IntoResponse, Json, TypedHeader}; use chrono::Duration; use hyper::StatusCode; use mas_axum_utils::sentry::SentryEventID; @@ -217,9 +217,11 @@ pub(crate) async fn post( activity_tracker: BoundActivityTracker, State(homeserver): State, State(site_config): State, + user_agent: Option>, Json(input): Json, ) -> Result { - let (session, user) = match (password_manager.is_enabled(), input.credentials) { + let user_agent = user_agent.map(|ua| ua.to_string()); + let (mut session, user) = match (password_manager.is_enabled(), input.credentials) { ( true, Credentials::Password { @@ -245,6 +247,13 @@ pub(crate) async fn post( } }; + if let Some(user_agent) = user_agent { + session = repo + .compat_session() + .record_user_agent(session, user_agent) + .await?; + } + let user_id = format!("@{username}:{homeserver}", username = user.username); // If the client asked for a refreshable token, make it expire diff --git a/crates/handlers/src/oauth2/token.rs b/crates/handlers/src/oauth2/token.rs index 6af48b68..f2e2df42 100644 --- a/crates/handlers/src/oauth2/token.rs +++ b/crates/handlers/src/oauth2/token.rs @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use axum::{extract::State, response::IntoResponse, Json}; +use axum::{extract::State, response::IntoResponse, Json, TypedHeader}; use chrono::{DateTime, Duration, Utc}; use headers::{CacheControl, HeaderMap, HeaderMapExt, Pragma}; use hyper::StatusCode; @@ -230,8 +230,10 @@ pub(crate) async fn post( State(site_config): State, State(encrypter): State, policy: Policy, + user_agent: Option>, client_authorization: ClientAuthorization, ) -> Result { + let user_agent = user_agent.map(|ua| ua.to_string()); let client = client_authorization .credentials .fetch(&mut repo) @@ -262,6 +264,7 @@ pub(crate) async fn post( &url_builder, &site_config, repo, + user_agent, ) .await? } @@ -274,6 +277,7 @@ pub(crate) async fn post( &client, &site_config, repo, + user_agent, ) .await? } @@ -287,6 +291,7 @@ pub(crate) async fn post( &site_config, repo, policy, + user_agent, ) .await? } @@ -301,6 +306,7 @@ pub(crate) async fn post( &url_builder, &site_config, repo, + user_agent, ) .await? } @@ -329,6 +335,7 @@ async fn authorization_code_grant( url_builder: &UrlBuilder, site_config: &SiteConfig, mut repo: BoxRepository, + user_agent: Option, ) -> Result<(AccessTokenResponse, BoxRepository), RouteError> { // Check that the client is allowed to use this grant type if !client.grant_types.contains(&GrantType::AuthorizationCode) { @@ -386,12 +393,19 @@ async fn authorization_code_grant( } }; - let session = repo + let mut session = repo .oauth2_session() .lookup(session_id) .await? .ok_or(RouteError::NoSuchOAuthSession)?; + if let Some(user_agent) = user_agent { + session = repo + .oauth2_session() + .record_user_agent(session, user_agent) + .await?; + } + // This should never happen, since we looked up in the database using the code let code = authz_grant.code.as_ref().ok_or(RouteError::InvalidGrant)?; @@ -490,6 +504,7 @@ async fn refresh_token_grant( client: &Client, site_config: &SiteConfig, mut repo: BoxRepository, + user_agent: Option, ) -> Result<(AccessTokenResponse, BoxRepository), RouteError> { // Check that the client is allowed to use this grant type if !client.grant_types.contains(&GrantType::RefreshToken) { @@ -502,12 +517,21 @@ async fn refresh_token_grant( .await? .ok_or(RouteError::RefreshTokenNotFound)?; - let session = repo + let mut session = repo .oauth2_session() .lookup(refresh_token.session_id) .await? .ok_or(RouteError::NoSuchOAuthSession)?; + // Let's for now record the user agent on each refresh, that should be + // responsive enough and not too much of a burden on the database. + if let Some(user_agent) = user_agent { + session = repo + .oauth2_session() + .record_user_agent(session, user_agent) + .await?; + } + if !refresh_token.is_valid() { return Err(RouteError::RefreshTokenInvalid(refresh_token.id)); } @@ -563,6 +587,7 @@ async fn client_credentials_grant( site_config: &SiteConfig, mut repo: BoxRepository, mut policy: Policy, + user_agent: Option, ) -> Result<(AccessTokenResponse, BoxRepository), RouteError> { // Check that the client is allowed to use this grant type if !client.grant_types.contains(&GrantType::ClientCredentials) { @@ -584,11 +609,18 @@ async fn client_credentials_grant( } // Start the session - let session = repo + let mut session = repo .oauth2_session() .add_from_client_credentials(rng, clock, client, scope) .await?; + if let Some(user_agent) = user_agent { + session = repo + .oauth2_session() + .record_user_agent(session, user_agent) + .await?; + } + let ttl = site_config.access_token_ttl; let access_token_str = TokenType::AccessToken.generate(rng); @@ -624,6 +656,7 @@ async fn device_code_grant( url_builder: &UrlBuilder, site_config: &SiteConfig, mut repo: BoxRepository, + user_agent: Option, ) -> Result<(AccessTokenResponse, BoxRepository), RouteError> { // Check that the client is allowed to use this grant type if !client.grant_types.contains(&GrantType::DeviceCode) { @@ -670,11 +703,19 @@ async fn device_code_grant( .ok_or(RouteError::NoSuchBrowserSession)?; // Start the session - let session = repo + let mut session = repo .oauth2_session() .add_from_browser_session(rng, clock, client, &browser_session, grant.scope) .await?; + // XXX: should we get the user agent from the device code grant instead? + if let Some(user_agent) = user_agent { + session = repo + .oauth2_session() + .record_user_agent(session, user_agent) + .await?; + } + let ttl = site_config.access_token_ttl; let access_token_str = TokenType::AccessToken.generate(rng); diff --git a/crates/storage-pg/.sqlx/query-1919d402fd6f148d14417f633be3353004f458c85f7b4f361802f86651900fbc.json b/crates/storage-pg/.sqlx/query-1919d402fd6f148d14417f633be3353004f458c85f7b4f361802f86651900fbc.json new file mode 100644 index 00000000..326ab111 --- /dev/null +++ b/crates/storage-pg/.sqlx/query-1919d402fd6f148d14417f633be3353004f458c85f7b4f361802f86651900fbc.json @@ -0,0 +1,15 @@ +{ + "db_name": "PostgreSQL", + "query": "\n UPDATE oauth2_sessions\n SET user_agent = $2\n WHERE oauth2_session_id = $1\n ", + "describe": { + "columns": [], + "parameters": { + "Left": [ + "Uuid", + "Text" + ] + }, + "nullable": [] + }, + "hash": "1919d402fd6f148d14417f633be3353004f458c85f7b4f361802f86651900fbc" +} diff --git a/crates/storage-pg/.sqlx/query-29148548d592046f7d711676911e3847e376e443ccd841f76b17a81f53fafc3a.json b/crates/storage-pg/.sqlx/query-29148548d592046f7d711676911e3847e376e443ccd841f76b17a81f53fafc3a.json new file mode 100644 index 00000000..6c84cb92 --- /dev/null +++ b/crates/storage-pg/.sqlx/query-29148548d592046f7d711676911e3847e376e443ccd841f76b17a81f53fafc3a.json @@ -0,0 +1,15 @@ +{ + "db_name": "PostgreSQL", + "query": "\n UPDATE compat_sessions\n SET user_agent = $2\n WHERE compat_session_id = $1\n ", + "describe": { + "columns": [], + "parameters": { + "Left": [ + "Uuid", + "Text" + ] + }, + "nullable": [] + }, + "hash": "29148548d592046f7d711676911e3847e376e443ccd841f76b17a81f53fafc3a" +} diff --git a/crates/storage-pg/.sqlx/query-31aace373b20b5dbf65fa51d8663da7571d85b6a7d2d544d69e7d04260cdffc9.json b/crates/storage-pg/.sqlx/query-5a2e9b5002c1927c0035c22e393172b36ab46a4377b46618205151ea041886d5.json similarity index 76% rename from crates/storage-pg/.sqlx/query-31aace373b20b5dbf65fa51d8663da7571d85b6a7d2d544d69e7d04260cdffc9.json rename to crates/storage-pg/.sqlx/query-5a2e9b5002c1927c0035c22e393172b36ab46a4377b46618205151ea041886d5.json index ea477bcb..5fae1ffa 100644 --- a/crates/storage-pg/.sqlx/query-31aace373b20b5dbf65fa51d8663da7571d85b6a7d2d544d69e7d04260cdffc9.json +++ b/crates/storage-pg/.sqlx/query-5a2e9b5002c1927c0035c22e393172b36ab46a4377b46618205151ea041886d5.json @@ -1,6 +1,6 @@ { "db_name": "PostgreSQL", - "query": "\n SELECT oauth2_session_id\n , user_id\n , user_session_id\n , oauth2_client_id\n , scope_list\n , created_at\n , finished_at\n , last_active_at\n , last_active_ip as \"last_active_ip: IpAddr\"\n FROM oauth2_sessions\n\n WHERE oauth2_session_id = $1\n ", + "query": "\n SELECT oauth2_session_id\n , user_id\n , user_session_id\n , oauth2_client_id\n , scope_list\n , created_at\n , finished_at\n , user_agent\n , last_active_at\n , last_active_ip as \"last_active_ip: IpAddr\"\n FROM oauth2_sessions\n\n WHERE oauth2_session_id = $1\n ", "describe": { "columns": [ { @@ -40,11 +40,16 @@ }, { "ordinal": 7, + "name": "user_agent", + "type_info": "Text" + }, + { + "ordinal": 8, "name": "last_active_at", "type_info": "Timestamptz" }, { - "ordinal": 8, + "ordinal": 9, "name": "last_active_ip: IpAddr", "type_info": "Inet" } @@ -63,8 +68,9 @@ false, true, true, + true, true ] }, - "hash": "31aace373b20b5dbf65fa51d8663da7571d85b6a7d2d544d69e7d04260cdffc9" + "hash": "5a2e9b5002c1927c0035c22e393172b36ab46a4377b46618205151ea041886d5" } diff --git a/crates/storage-pg/.sqlx/query-04e25c9267bf2eb143a6445345229081e7b386743a93b3833ef8ad9d09972f3b.json b/crates/storage-pg/.sqlx/query-bb6f55a4cc10bec8ec0fc138485f6b4d308302bb1fa3accb12932d1e5ce457e9.json similarity index 76% rename from crates/storage-pg/.sqlx/query-04e25c9267bf2eb143a6445345229081e7b386743a93b3833ef8ad9d09972f3b.json rename to crates/storage-pg/.sqlx/query-bb6f55a4cc10bec8ec0fc138485f6b4d308302bb1fa3accb12932d1e5ce457e9.json index 599fa963..2c81606b 100644 --- a/crates/storage-pg/.sqlx/query-04e25c9267bf2eb143a6445345229081e7b386743a93b3833ef8ad9d09972f3b.json +++ b/crates/storage-pg/.sqlx/query-bb6f55a4cc10bec8ec0fc138485f6b4d308302bb1fa3accb12932d1e5ce457e9.json @@ -1,6 +1,6 @@ { "db_name": "PostgreSQL", - "query": "\n SELECT compat_session_id\n , device_id\n , user_id\n , user_session_id\n , created_at\n , finished_at\n , is_synapse_admin\n , last_active_at\n , last_active_ip as \"last_active_ip: IpAddr\"\n FROM compat_sessions\n WHERE compat_session_id = $1\n ", + "query": "\n SELECT compat_session_id\n , device_id\n , user_id\n , user_session_id\n , created_at\n , finished_at\n , is_synapse_admin\n , user_agent\n , last_active_at\n , last_active_ip as \"last_active_ip: IpAddr\"\n FROM compat_sessions\n WHERE compat_session_id = $1\n ", "describe": { "columns": [ { @@ -40,11 +40,16 @@ }, { "ordinal": 7, + "name": "user_agent", + "type_info": "Text" + }, + { + "ordinal": 8, "name": "last_active_at", "type_info": "Timestamptz" }, { - "ordinal": 8, + "ordinal": 9, "name": "last_active_ip: IpAddr", "type_info": "Inet" } @@ -63,8 +68,9 @@ true, false, true, + true, true ] }, - "hash": "04e25c9267bf2eb143a6445345229081e7b386743a93b3833ef8ad9d09972f3b" + "hash": "bb6f55a4cc10bec8ec0fc138485f6b4d308302bb1fa3accb12932d1e5ce457e9" } diff --git a/crates/storage-pg/migrations/20240221164945_sessions_user_agent.sql b/crates/storage-pg/migrations/20240221164945_sessions_user_agent.sql new file mode 100644 index 00000000..5b1ce269 --- /dev/null +++ b/crates/storage-pg/migrations/20240221164945_sessions_user_agent.sql @@ -0,0 +1,17 @@ +-- Copyright 2024 The Matrix.org Foundation C.I.C. +-- +-- Licensed under the Apache License, Version 2.0 (the "License"); +-- you may not use this file except in compliance with the License. +-- You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, software +-- distributed under the License is distributed on an "AS IS" BASIS, +-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +-- See the License for the specific language governing permissions and +-- limitations under the License. + +-- Adds user agent columns to oauth and compat sessions tables +ALTER TABLE oauth2_sessions ADD COLUMN user_agent TEXT; +ALTER TABLE compat_sessions ADD COLUMN user_agent TEXT; diff --git a/crates/storage-pg/src/app_session.rs b/crates/storage-pg/src/app_session.rs index 94c83671..cd2786fd 100644 --- a/crates/storage-pg/src/app_session.rs +++ b/crates/storage-pg/src/app_session.rs @@ -73,6 +73,7 @@ mod priv_ { pub(super) created_at: DateTime, pub(super) finished_at: Option>, pub(super) is_synapse_admin: Option, + pub(super) user_agent: Option, pub(super) last_active_at: Option>, pub(super) last_active_ip: Option, } @@ -98,6 +99,7 @@ impl TryFrom for AppSession { created_at, finished_at, is_synapse_admin, + user_agent, last_active_at, last_active_ip, } = value; @@ -143,6 +145,7 @@ impl TryFrom for AppSession { user_session_id, created_at, is_synapse_admin, + user_agent, last_active_at, last_active_ip, }; @@ -182,6 +185,7 @@ impl TryFrom for AppSession { user_id: user_id.map(Ulid::from), user_session_id, scope, + user_agent, last_active_at, last_active_ip, }; @@ -250,6 +254,10 @@ impl<'c> AppSessionRepository for PgAppSessionRepository<'c> { AppSessionLookupIden::FinishedAt, ) .expr_as(Expr::cust("NULL"), AppSessionLookupIden::IsSynapseAdmin) + .expr_as( + Expr::col((OAuth2Sessions::Table, OAuth2Sessions::UserAgent)), + AppSessionLookupIden::UserAgent, + ) .expr_as( Expr::col((OAuth2Sessions::Table, OAuth2Sessions::LastActiveAt)), AppSessionLookupIden::LastActiveAt, @@ -317,6 +325,10 @@ impl<'c> AppSessionRepository for PgAppSessionRepository<'c> { Expr::col((CompatSessions::Table, CompatSessions::IsSynapseAdmin)), AppSessionLookupIden::IsSynapseAdmin, ) + .expr_as( + Expr::col((CompatSessions::Table, CompatSessions::UserAgent)), + AppSessionLookupIden::UserAgent, + ) .expr_as( Expr::col((CompatSessions::Table, CompatSessions::LastActiveAt)), AppSessionLookupIden::LastActiveAt, diff --git a/crates/storage-pg/src/compat/mod.rs b/crates/storage-pg/src/compat/mod.rs index 528624e6..f2c8894b 100644 --- a/crates/storage-pg/src/compat/mod.rs +++ b/crates/storage-pg/src/compat/mod.rs @@ -129,6 +129,24 @@ mod tests { assert!(session_lookup.is_valid()); assert!(!session_lookup.is_finished()); + // Record a user-agent for the session + assert!(session_lookup.user_agent.is_none()); + let session = repo + .compat_session() + .record_user_agent(session_lookup, "Mozilla/5.0".to_owned()) + .await + .unwrap(); + assert_eq!(session.user_agent.as_deref(), Some("Mozilla/5.0")); + + // Reload the session and check again + let session_lookup = repo + .compat_session() + .lookup(session.id) + .await + .unwrap() + .expect("compat session not found"); + assert_eq!(session_lookup.user_agent.as_deref(), Some("Mozilla/5.0")); + // Look up the session by device let list = repo .compat_session() diff --git a/crates/storage-pg/src/compat/session.rs b/crates/storage-pg/src/compat/session.rs index 71df6eed..b980bcbb 100644 --- a/crates/storage-pg/src/compat/session.rs +++ b/crates/storage-pg/src/compat/session.rs @@ -60,6 +60,7 @@ struct CompatSessionLookup { created_at: DateTime, finished_at: Option>, is_synapse_admin: bool, + user_agent: Option, last_active_at: Option>, last_active_ip: Option, } @@ -89,6 +90,7 @@ impl TryFrom for CompatSession { device, created_at: value.created_at, is_synapse_admin: value.is_synapse_admin, + user_agent: value.user_agent, last_active_at: value.last_active_at, last_active_ip: value.last_active_ip, }; @@ -107,6 +109,7 @@ struct CompatSessionAndSsoLoginLookup { created_at: DateTime, finished_at: Option>, is_synapse_admin: bool, + user_agent: Option, last_active_at: Option>, last_active_ip: Option, compat_sso_login_id: Option, @@ -142,6 +145,7 @@ impl TryFrom for (CompatSession, Option CompatSessionRepository for PgCompatSessionRepository<'c> { , created_at , finished_at , is_synapse_admin + , user_agent , last_active_at , last_active_ip as "last_active_ip: IpAddr" FROM compat_sessions @@ -290,6 +295,7 @@ impl<'c> CompatSessionRepository for PgCompatSessionRepository<'c> { user_session_id: browser_session.map(|s| s.id), created_at, is_synapse_admin, + user_agent: None, last_active_at: None, last_active_ip: None, }) @@ -377,6 +383,10 @@ impl<'c> CompatSessionRepository for PgCompatSessionRepository<'c> { Expr::col((CompatSessions::Table, CompatSessions::IsSynapseAdmin)), CompatSessionAndSsoLoginLookupIden::IsSynapseAdmin, ) + .expr_as( + Expr::col((CompatSessions::Table, CompatSessions::UserAgent)), + CompatSessionAndSsoLoginLookupIden::UserAgent, + ) .expr_as( Expr::col((CompatSessions::Table, CompatSessions::LastActiveAt)), CompatSessionAndSsoLoginLookupIden::LastActiveAt, @@ -552,4 +562,38 @@ impl<'c> CompatSessionRepository for PgCompatSessionRepository<'c> { Ok(()) } + + #[tracing::instrument( + name = "db.compat_session.record_user_agent", + skip_all, + fields( + db.statement, + %compat_session.id, + ), + err, + )] + async fn record_user_agent( + &mut self, + mut compat_session: CompatSession, + user_agent: String, + ) -> Result { + let res = sqlx::query!( + r#" + UPDATE compat_sessions + SET user_agent = $2 + WHERE compat_session_id = $1 + "#, + Uuid::from(compat_session.id), + user_agent, + ) + .traced() + .execute(&mut *self.conn) + .await?; + + compat_session.user_agent = Some(user_agent); + + DatabaseError::ensure_affected_rows(&res, 1)?; + + Ok(compat_session) + } } diff --git a/crates/storage-pg/src/iden.rs b/crates/storage-pg/src/iden.rs index c45f867b..194d335c 100644 --- a/crates/storage-pg/src/iden.rs +++ b/crates/storage-pg/src/iden.rs @@ -57,6 +57,7 @@ pub enum CompatSessions { CreatedAt, FinishedAt, IsSynapseAdmin, + UserAgent, LastActiveAt, LastActiveIp, } @@ -86,6 +87,7 @@ pub enum OAuth2Sessions { ScopeList, CreatedAt, FinishedAt, + UserAgent, LastActiveAt, LastActiveIp, } diff --git a/crates/storage-pg/src/oauth2/mod.rs b/crates/storage-pg/src/oauth2/mod.rs index 7c37085b..0f824ddc 100644 --- a/crates/storage-pg/src/oauth2/mod.rs +++ b/crates/storage-pg/src/oauth2/mod.rs @@ -367,6 +367,24 @@ mod tests { .unwrap(); assert!(!refresh_token.is_valid()); + // Record the user-agent on the session + assert!(session.user_agent.is_none()); + let session = repo + .oauth2_session() + .record_user_agent(session, "Mozilla/5.0".to_owned()) + .await + .unwrap(); + assert_eq!(session.user_agent.as_deref(), Some("Mozilla/5.0")); + + // Reload the session and check the user-agent + let session = repo + .oauth2_session() + .lookup(session.id) + .await + .unwrap() + .expect("session not found"); + assert_eq!(session.user_agent.as_deref(), Some("Mozilla/5.0")); + // Mark the session as finished assert!(session.is_valid()); let session = repo.oauth2_session().finish(&clock, session).await.unwrap(); diff --git a/crates/storage-pg/src/oauth2/session.rs b/crates/storage-pg/src/oauth2/session.rs index 8adfb4b8..e80523e2 100644 --- a/crates/storage-pg/src/oauth2/session.rs +++ b/crates/storage-pg/src/oauth2/session.rs @@ -59,6 +59,7 @@ struct OAuthSessionLookup { scope_list: Vec, created_at: DateTime, finished_at: Option>, + user_agent: Option, last_active_at: Option>, last_active_ip: Option, } @@ -93,6 +94,7 @@ impl TryFrom for Session { user_id: value.user_id.map(Ulid::from), user_session_id: value.user_session_id.map(Ulid::from), scope, + user_agent: value.user_agent, last_active_at: value.last_active_at, last_active_ip: value.last_active_ip, }) @@ -123,6 +125,7 @@ impl<'c> OAuth2SessionRepository for PgOAuth2SessionRepository<'c> { , scope_list , created_at , finished_at + , user_agent , last_active_at , last_active_ip as "last_active_ip: IpAddr" FROM oauth2_sessions @@ -197,6 +200,7 @@ impl<'c> OAuth2SessionRepository for PgOAuth2SessionRepository<'c> { user_session_id: user_session.map(|s| s.id), client_id: client.id, scope, + user_agent: None, last_active_at: None, last_active_ip: None, }) @@ -281,6 +285,10 @@ impl<'c> OAuth2SessionRepository for PgOAuth2SessionRepository<'c> { Expr::col((OAuth2Sessions::Table, OAuth2Sessions::FinishedAt)), OAuthSessionLookupIden::FinishedAt, ) + .expr_as( + Expr::col((OAuth2Sessions::Table, OAuth2Sessions::UserAgent)), + OAuthSessionLookupIden::UserAgent, + ) .expr_as( Expr::col((OAuth2Sessions::Table, OAuth2Sessions::LastActiveAt)), OAuthSessionLookupIden::LastActiveAt, @@ -427,4 +435,41 @@ impl<'c> OAuth2SessionRepository for PgOAuth2SessionRepository<'c> { Ok(()) } + + #[tracing::instrument( + name = "db.oauth2_session.record_user_agent", + skip_all, + fields( + db.statement, + %session.id, + %session.scope, + client.id = %session.client_id, + session.user_agent = %user_agent, + ), + err, + )] + async fn record_user_agent( + &mut self, + mut session: Session, + user_agent: String, + ) -> Result { + let res = sqlx::query!( + r#" + UPDATE oauth2_sessions + SET user_agent = $2 + WHERE oauth2_session_id = $1 + "#, + Uuid::from(session.id), + user_agent, + ) + .traced() + .execute(&mut *self.conn) + .await?; + + session.user_agent = Some(user_agent); + + DatabaseError::ensure_affected_rows(&res, 1)?; + + Ok(session) + } } diff --git a/crates/storage/src/compat/session.rs b/crates/storage/src/compat/session.rs index f829187f..d676d2ae 100644 --- a/crates/storage/src/compat/session.rs +++ b/crates/storage/src/compat/session.rs @@ -252,6 +252,22 @@ pub trait CompatSessionRepository: Send + Sync { &mut self, activity: Vec<(Ulid, DateTime, Option)>, ) -> Result<(), Self::Error>; + + /// Record the user agent of a compat session + /// + /// # Parameters + /// + /// * `compat_session`: The compat session to record the user agent for + /// * `user_agent`: The user agent to record + /// + /// # Errors + /// + /// Returns [`Self::Error`] if the underlying repository fails + async fn record_user_agent( + &mut self, + compat_session: CompatSession, + user_agent: String, + ) -> Result; } repository_impl!(CompatSessionRepository: @@ -285,4 +301,10 @@ repository_impl!(CompatSessionRepository: &mut self, activity: Vec<(Ulid, DateTime, Option)>, ) -> Result<(), Self::Error>; + + async fn record_user_agent( + &mut self, + compat_session: CompatSession, + user_agent: String, + ) -> Result; ); diff --git a/crates/storage/src/oauth2/session.rs b/crates/storage/src/oauth2/session.rs index 4ccc0d08..28a1b3a6 100644 --- a/crates/storage/src/oauth2/session.rs +++ b/crates/storage/src/oauth2/session.rs @@ -286,6 +286,18 @@ pub trait OAuth2SessionRepository: Send + Sync { &mut self, activity: Vec<(Ulid, DateTime, Option)>, ) -> Result<(), Self::Error>; + + /// Record the user agent of a [`Session`] + /// + /// # Parameters + /// + /// * `session`: The [`Session`] to record the user agent for + /// * `user_agent`: The user agent to record + async fn record_user_agent( + &mut self, + session: Session, + user_agent: String, + ) -> Result; } repository_impl!(OAuth2SessionRepository: @@ -333,4 +345,10 @@ repository_impl!(OAuth2SessionRepository: &mut self, activity: Vec<(Ulid, DateTime, Option)>, ) -> Result<(), Self::Error>; + + async fn record_user_agent( + &mut self, + session: Session, + user_agent: String, + ) -> Result; );