Record user agents on OAuth 2.0 and compat sessions

crates/data-model/src/compat/session.rs

@@ -83,6 +83,7 @@ pub struct CompatSession {
     pub user_session_id: Option<Ulid>,
     pub created_at: DateTime<Utc>,
     pub is_synapse_admin: bool,
+    pub user_agent: Option<String>,
     pub last_active_at: Option<DateTime<Utc>>,
     pub last_active_ip: Option<IpAddr>,
 }

crates/data-model/src/oauth2/session.rs

@@ -75,6 +75,7 @@ pub struct Session {
     pub user_session_id: Option<Ulid>,
     pub client_id: Ulid,
     pub scope: Scope,
+    pub user_agent: Option<String>,
     pub last_active_at: Option<DateTime<Utc>>,
     pub last_active_ip: Option<IpAddr>,
 }

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<MatrixHomeserver>,
     State(site_config): State<SiteConfig>,
+    user_agent: Option<TypedHeader<headers::UserAgent>>,
     Json(input): Json<RequestBody>,
 ) -> Result<impl IntoResponse, RouteError> {
-    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

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<SiteConfig>,
     State(encrypter): State<Encrypter>,
     policy: Policy,
+    user_agent: Option<TypedHeader<headers::UserAgent>>,
     client_authorization: ClientAuthorization<AccessTokenRequest>,
 ) -> Result<impl IntoResponse, RouteError> {
+    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<String>,
 ) -> 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<String>,
 ) -> 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<String>,
 ) -> 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<String>,
 ) -> 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);
 

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;

crates/storage-pg/src/app_session.rs

@@ -73,6 +73,7 @@ mod priv_ {
         pub(super) created_at: DateTime<Utc>,
         pub(super) finished_at: Option<DateTime<Utc>>,
         pub(super) is_synapse_admin: Option<bool>,
+        pub(super) user_agent: Option<String>,
         pub(super) last_active_at: Option<DateTime<Utc>>,
         pub(super) last_active_ip: Option<IpAddr>,
     }
@@ -98,6 +99,7 @@ impl TryFrom<AppSessionLookup> for AppSession {
             created_at,
             finished_at,
             is_synapse_admin,
+            user_agent,
             last_active_at,
             last_active_ip,
         } = value;
@@ -143,6 +145,7 @@ impl TryFrom<AppSessionLookup> for AppSession {
                     user_session_id,
                     created_at,
                     is_synapse_admin,
+                    user_agent,
                     last_active_at,
                     last_active_ip,
                 };
@@ -182,6 +185,7 @@ impl TryFrom<AppSessionLookup> 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,