From 894957934dc388ed511427c53155581a1ba5f90f Mon Sep 17 00:00:00 2001
From: Quentin Gliech <quenting@element.io>
Date: Tue, 19 Sep 2023 19:50:42 +0200
Subject: [PATCH] Test the activity tracker on the introspection endpoint

---
 crates/data-model/src/oauth2/session.rs       |  1 +
 crates/handlers/src/activity_tracker/mod.rs   | 16 ++++++--
 .../handlers/src/activity_tracker/worker.rs   |  3 +-
 crates/handlers/src/oauth2/introspection.rs   | 38 +++++++++++++++++++
 4 files changed, 53 insertions(+), 5 deletions(-)

diff --git a/crates/data-model/src/oauth2/session.rs b/crates/data-model/src/oauth2/session.rs
index 5d14a927..0d17ea6d 100644
--- a/crates/data-model/src/oauth2/session.rs
+++ b/crates/data-model/src/oauth2/session.rs
@@ -13,6 +13,7 @@
 // limitations under the License.
 
 use std::net::IpAddr;
+
 use chrono::{DateTime, Utc};
 use oauth2_types::scope::Scope;
 use serde::Serialize;
diff --git a/crates/handlers/src/activity_tracker/mod.rs b/crates/handlers/src/activity_tracker/mod.rs
index f867720c..702a1eb3 100644
--- a/crates/handlers/src/activity_tracker/mod.rs
+++ b/crates/handlers/src/activity_tracker/mod.rs
@@ -52,7 +52,7 @@ enum Message {
         date_time: DateTime<Utc>,
         ip: Option<IpAddr>,
     },
-    Flush,
+    Flush(tokio::sync::oneshot::Sender<()>),
     Shutdown(tokio::sync::oneshot::Sender<()>),
 }
 
@@ -150,10 +150,18 @@ impl ActivityTracker {
 
     /// Manually flush the activity tracker.
     pub async fn flush(&self) {
-        let res = self.channel.send(Message::Flush).await;
+        let (tx, rx) = tokio::sync::oneshot::channel();
+        let res = self.channel.send(Message::Flush(tx)).await;
 
-        if let Err(e) = res {
-            tracing::error!("Failed to flush activity tracker: {}", e);
+        match res {
+            Ok(_) => {
+                if let Err(e) = rx.await {
+                    tracing::error!("Failed to flush activity tracker: {}", e);
+                }
+            }
+            Err(e) => {
+                tracing::error!("Failed to flush activity tracker: {}", e);
+            }
         }
     }
 
diff --git a/crates/handlers/src/activity_tracker/worker.rs b/crates/handlers/src/activity_tracker/worker.rs
index 7fde6b44..978e7891 100644
--- a/crates/handlers/src/activity_tracker/worker.rs
+++ b/crates/handlers/src/activity_tracker/worker.rs
@@ -137,10 +137,11 @@ impl Worker {
 
                     record.end_time = date_time.max(record.end_time);
                 }
-                Message::Flush => {
+                Message::Flush(tx) => {
                     self.message_counter.add(1, &[TYPE.string("flush")]);
 
                     self.flush().await;
+                    let _ = tx.send(());
                 }
                 Message::Shutdown(tx) => {
                     self.message_counter.add(1, &[TYPE.string("shutdown")]);
diff --git a/crates/handlers/src/oauth2/introspection.rs b/crates/handlers/src/oauth2/introspection.rs
index 015d47e0..27af1875 100644
--- a/crates/handlers/src/oauth2/introspection.rs
+++ b/crates/handlers/src/oauth2/introspection.rs
@@ -463,6 +463,7 @@ mod tests {
     use mas_data_model::{AccessToken, RefreshToken};
     use mas_iana::oauth::OAuthTokenTypeHint;
     use mas_router::{OAuth2Introspection, OAuth2RegistrationEndpoint, SimpleRoute};
+    use mas_storage::Clock;
     use oauth2_types::{
         registration::ClientRegistrationResponse,
         requests::IntrospectionResponse,
@@ -618,7 +619,20 @@ mod tests {
         let response: IntrospectionResponse = response.json();
         assert!(!response.active); // It shouldn't be active
 
+        // We should have recorded the session last activity
+        state.activity_tracker.flush().await;
+        let mut repo = state.repository().await.unwrap();
+        let session = repo
+            .oauth2_session()
+            .lookup(session.id)
+            .await
+            .unwrap()
+            .unwrap();
+        assert_eq!(session.last_active_at, Some(state.clock.now()));
+        repo.cancel().await.unwrap();
+
         // Advance the clock to invalidate the access token
+        let old_now = state.clock.now();
         state.clock.advance(Duration::hours(1));
 
         let request = Request::post(OAuth2Introspection::PATH)
@@ -629,6 +643,18 @@ mod tests {
         let response: IntrospectionResponse = response.json();
         assert!(!response.active); // It shouldn't be active anymore
 
+        // That should not have updated the session last activity
+        state.activity_tracker.flush().await;
+        let mut repo = state.repository().await.unwrap();
+        let session = repo
+            .oauth2_session()
+            .lookup(session.id)
+            .await
+            .unwrap()
+            .unwrap();
+        assert_eq!(session.last_active_at, Some(old_now));
+        repo.cancel().await.unwrap();
+
         // But the refresh token should still be valid
         let request = Request::post(OAuth2Introspection::PATH)
             .basic_auth(&introspecting_client_id, &introspecting_client_secret)
@@ -637,6 +663,18 @@ mod tests {
         response.assert_status(StatusCode::OK);
         let response: IntrospectionResponse = response.json();
         assert!(response.active);
+
+        // But this time, we should have updated the session last activity
+        state.activity_tracker.flush().await;
+        let mut repo = state.repository().await.unwrap();
+        let session = repo
+            .oauth2_session()
+            .lookup(session.id)
+            .await
+            .unwrap()
+            .unwrap();
+        assert_eq!(session.last_active_at, Some(state.clock.now()));
+        repo.cancel().await.unwrap();
     }
 
     #[sqlx::test(migrator = "mas_storage_pg::MIGRATOR")]