test(sdk): Add tests for handle_refresh_tokens and Oidc

Signed-off-by: Kévin Commaille <zecakeh@tedomum.fr>

crates/matrix-sdk/src/authentication/oidc/cross_process.rs

@@ -260,10 +260,10 @@ mod tests {
 
     use super::compute_session_hash;
     use crate::{
-        authentication::oidc::{cross_process::SessionHash, tests::prev_session_tokens},
+        authentication::oidc::cross_process::SessionHash,
         test_utils::{
             client::{
-                oauth::{mock_session, mock_session_tokens},
+                oauth::{mock_prev_session_tokens, mock_session, mock_session_tokens},
                 MockClientBuilder,
             },
             mocks::MatrixMockServer,
@@ -392,7 +392,8 @@ mod tests {
         oidc.enable_cross_process_refresh_lock("lock".to_owned()).await?;
 
         // Restore the session.
-        oidc.restore_session(mock_session(prev_session_tokens(), server.server().uri())).await?;
+        oidc.restore_session(mock_session(mock_prev_session_tokens(), server.server().uri()))
+            .await?;
 
         // Immediately try to refresh the access token twice in parallel.
         for result in join_all([oidc.refresh_access_token(), oidc.refresh_access_token()]).await {
@@ -423,7 +424,7 @@ mod tests {
         oauth_server.mock_server_metadata().ok().expect(1..).named("server_metadata").mount().await;
         oauth_server.mock_token().ok().expect(1).named("token").mount().await;
 
-        let prev_tokens = prev_session_tokens();
+        let prev_tokens = mock_prev_session_tokens();
         let next_tokens = mock_session_tokens();
 
         // Create the first client.

crates/matrix-sdk/src/authentication/oidc/tests.rs

@@ -24,7 +24,7 @@ use wiremock::{
 
 use super::{
     registrations::OidcRegistrations, AuthorizationCode, AuthorizationError, AuthorizationResponse,
-    Oidc, OidcAuthorizationData, OidcError, OidcSessionTokens, RedirectUriQueryParseError,
+    Oidc, OidcAuthorizationData, OidcError, RedirectUriQueryParseError,
 };
 use crate::{
     authentication::oidc::{
@@ -33,7 +33,9 @@ use crate::{
     },
     test_utils::{
         client::{
-            oauth::{mock_client_metadata, mock_session, mock_session_tokens},
+            oauth::{
+                mock_client_metadata, mock_prev_session_tokens, mock_session, mock_session_tokens,
+            },
             MockClientBuilder,
         },
         mocks::{oauth::MockServerMetadataBuilder, MatrixMockServer},
@@ -43,15 +45,6 @@ use crate::{
 
 const REDIRECT_URI_STRING: &str = "http://127.0.0.1:6778/oidc/callback";
 
-/// Different session tokens than the ones returned by the mock server's
-/// token endpoint.
-pub(crate) fn prev_session_tokens() -> OidcSessionTokens {
-    OidcSessionTokens {
-        access_token: "prev-access-token".to_owned(),
-        refresh_token: Some("prev-refresh-token".to_owned()),
-    }
-}
-
 async fn mock_environment(
 ) -> anyhow::Result<(Oidc, MatrixMockServer, VerifiedClientMetadata, OidcRegistrations)> {
     let server = MatrixMockServer::new().await;
@@ -439,7 +432,7 @@ async fn test_insecure_clients() -> anyhow::Result<()> {
     oauth_server.mock_server_metadata().ok().expect(2..).named("server_metadata").mount().await;
     oauth_server.mock_token().ok().expect(2).named("token").mount().await;
 
-    let prev_tokens = prev_session_tokens();
+    let prev_tokens = mock_prev_session_tokens();
     let next_tokens = mock_session_tokens();
 
     for client in [

crates/matrix-sdk/src/test_utils/client.rs

@@ -79,6 +79,12 @@ impl MockClientBuilder {
         self
     }
 
+    /// Handle refreshing access tokens automatically.
+    pub fn handle_refresh_tokens(mut self) -> Self {
+        self.builder = self.builder.handle_refresh_tokens();
+        self
+    }
+
     /// Finish building the client into the final [`Client`] instance.
     pub async fn build(self) -> Client {
         let client = self.builder.build().await.expect("building client failed");
@@ -197,6 +203,15 @@ pub mod oauth {
         }
     }
 
+    /// Different session tokens than the ones returned by
+    /// [`mock_session_tokens()`].
+    pub fn mock_prev_session_tokens() -> OidcSessionTokens {
+        OidcSessionTokens {
+            access_token: "prev-access-token".to_owned(),
+            refresh_token: Some("prev-refresh-token".to_owned()),
+        }
+    }
+
     /// An [`OidcSession`] to restore, for unit or integration tests.
     pub fn mock_session(tokens: OidcSessionTokens, issuer: String) -> OidcSession {
         OidcSession {

crates/matrix-sdk/src/test_utils/mocks/mod.rs

@@ -951,11 +951,20 @@ impl MatrixMockServer {
     }
 
     /// Creates a prebuilt mock for the endpoint used to get information about
-    /// the owner of the current access token.
+    /// the owner of the default access token.
     pub fn mock_who_am_i(&self) -> MockEndpoint<'_, WhoAmIEndpoint> {
+        self.mock_who_am_i_with_access_token("1234")
+    }
+
+    /// Creates a prebuilt mock for the endpoint used to get information about
+    /// the owner of the given access token.
+    pub fn mock_who_am_i_with_access_token(
+        &self,
+        access_token: &str,
+    ) -> MockEndpoint<'_, WhoAmIEndpoint> {
         let mock = Mock::given(method("GET"))
             .and(path_regex(r"^/_matrix/client/v3/account/whoami"))
-            .and(header("authorization", "Bearer 1234"));
+            .and(header("authorization", format!("Bearer {access_token}")));
         MockEndpoint { mock, server: &self.server, endpoint: WhoAmIEndpoint }
     }
 
@@ -2421,6 +2430,16 @@ impl<'a> MockEndpoint<'a, WhoAmIEndpoint> {
 
         MatrixMock { server: self.server, mock }
     }
+
+    /// Returns an error response with an `M_UNKNOWN_TOKEN`.
+    pub fn err_unknown_token(self) -> MatrixMock<'a> {
+        let mock = self.mock.respond_with(ResponseTemplate::new(401).set_body_json(json!({
+            "errcode": "M_UNKNOWN_TOKEN",
+            "error": "Invalid token"
+        })));
+
+        MatrixMock { server: self.server, mock }
+    }
 }
 
 /// A prebuilt mock for `POST /keys/upload` request.

crates/matrix-sdk/src/test_utils/mocks/oauth.rs

@@ -318,6 +318,15 @@ impl<'a> MockEndpoint<'a, TokenEndpoint> {
 
         MatrixMock { server: self.server, mock }
     }
+
+    /// Returns an error response when the token in the request is invalid.
+    pub fn invalid_grant(self) -> MatrixMock<'a> {
+        let mock = self.mock.respond_with(ResponseTemplate::new(400).set_body_json(json!({
+            "error": "invalid_grant",
+        })));
+
+        MatrixMock { server: self.server, mock }
+    }
 }
 
 /// A prebuilt mock for a `POST /oauth/revoke` request.

crates/matrix-sdk/tests/integration/refresh_token.rs

@@ -298,6 +298,8 @@ async fn test_refresh_token_handled_success() {
         tokens_changed_stream.next().await.unwrap();
     });
 
+    let mut session_changes = client.subscribe_to_session_changes();
+
     Mock::given(method("POST"))
         .and(path("/_matrix/client/v3/refresh"))
         .respond_with(ResponseTemplate::new(200).set_body_json(&*test_json::REFRESH_TOKEN))
@@ -329,6 +331,9 @@ async fn test_refresh_token_handled_success() {
     client.whoami().await.unwrap();
     tokens_join_handle.await.unwrap();
     changed_join_handle.await.unwrap();
+
+    assert_eq!(session_changes.try_recv(), Ok(SessionChange::TokensRefreshed));
+    assert_eq!(session_changes.try_recv(), Err(TryRecvError::Empty));
 }
 
 #[async_test]
@@ -346,6 +351,8 @@ async fn test_refresh_token_handled_failure() {
     let session = session();
     auth.restore_session(session).await.unwrap();
 
+    let mut session_changes = client.subscribe_to_session_changes();
+
     Mock::given(method("POST"))
         .and(path("/_matrix/client/v3/refresh"))
         .respond_with(
@@ -378,7 +385,13 @@ async fn test_refresh_token_handled_failure() {
 
     let res = client.whoami().await;
     assert_let!(Err(HttpError::RefreshToken(RefreshTokenError::MatrixAuth(http_err))) = res);
-    assert_matches!(http_err.client_api_error_kind(), Some(ErrorKind::UnknownToken { .. }))
+    assert_matches!(
+        http_err.client_api_error_kind(),
+        Some(ErrorKind::UnknownToken { soft_logout: true })
+    );
+
+    assert_eq!(session_changes.try_recv(), Ok(SessionChange::UnknownToken { soft_logout: true }));
+    assert_eq!(session_changes.try_recv(), Err(TryRecvError::Empty));
 }
 
 #[async_test]
@@ -560,3 +573,137 @@ async fn test_refresh_token_handled_other_error() {
 
     client.whoami().await.unwrap_err();
 }
+
+#[cfg(feature = "experimental-oidc")]
+#[async_test]
+async fn test_oauth_refresh_token_handled_success() {
+    use matrix_sdk::test_utils::{
+        client::oauth::{mock_prev_session_tokens, mock_session, mock_session_tokens},
+        mocks::MatrixMockServer,
+    };
+
+    let server = MatrixMockServer::new().await;
+    // Return an error first so the token is refreshed.
+    server
+        .mock_who_am_i_with_access_token("prev-access-token")
+        .err_unknown_token()
+        .expect(1)
+        .named("whoami_unknown_token")
+        .mount()
+        .await;
+    server.mock_who_am_i().ok().expect(1).named("whoami_ok").mount().await;
+
+    let oauth_server = server.oauth();
+    oauth_server.mock_server_metadata().ok().expect(1..).named("server_metadata").mount().await;
+    oauth_server.mock_token().ok().expect(1).named("token").mount().await;
+
+    let issuer = server.server().uri();
+    let client = server.client_builder().unlogged().handle_refresh_tokens().build().await;
+    let oidc = client.oidc();
+
+    oidc.restore_session(mock_session(mock_prev_session_tokens(), issuer)).await.unwrap();
+
+    let mut tokens_stream = oidc.session_tokens_stream().unwrap();
+    let tokens_join_handle = spawn(async move {
+        let tokens = tokens_stream.next().await.unwrap();
+        assert_eq!(tokens, mock_session_tokens());
+    });
+
+    let mut session_changes = client.subscribe_to_session_changes();
+
+    client.whoami().await.unwrap();
+
+    tokens_join_handle.await.unwrap();
+
+    assert_eq!(session_changes.try_recv(), Ok(SessionChange::TokensRefreshed));
+    assert_eq!(session_changes.try_recv(), Err(TryRecvError::Empty));
+}
+
+#[cfg(feature = "experimental-oidc")]
+#[async_test]
+async fn test_oauth_refresh_token_handled_failure() {
+    use matrix_sdk::{
+        authentication::oidc::OidcError,
+        test_utils::{
+            client::oauth::{mock_prev_session_tokens, mock_session},
+            mocks::MatrixMockServer,
+        },
+    };
+
+    let server = MatrixMockServer::new().await;
+    // Return an error first so the token is refreshed.
+    server
+        .mock_who_am_i_with_access_token("prev-access-token")
+        .err_unknown_token()
+        .expect(1)
+        .named("whoami_unknown_token")
+        .mount()
+        .await;
+
+    let oauth_server = server.oauth();
+    oauth_server.mock_server_metadata().ok().expect(1..).named("server_metadata").mount().await;
+    // Return an error to fail the token refresh.
+    oauth_server.mock_token().invalid_grant().expect(1).named("token").mount().await;
+
+    let issuer = server.server().uri();
+    let client = server.client_builder().unlogged().handle_refresh_tokens().build().await;
+    let oidc = client.oidc();
+
+    oidc.restore_session(mock_session(mock_prev_session_tokens(), issuer)).await.unwrap();
+
+    let mut session_changes = client.subscribe_to_session_changes();
+
+    let res = client.whoami().await;
+    assert_let!(Err(HttpError::RefreshToken(RefreshTokenError::Oidc(oidc_err))) = res);
+    assert_matches!(*oidc_err, OidcError::RefreshToken(_));
+
+    assert_eq!(session_changes.try_recv(), Ok(SessionChange::UnknownToken { soft_logout: false }));
+    assert_eq!(session_changes.try_recv(), Err(TryRecvError::Empty));
+}
+
+#[cfg(feature = "experimental-oidc")]
+#[async_test]
+async fn test_oauth_handle_refresh_tokens() {
+    use matrix_sdk::test_utils::{
+        client::oauth::{mock_prev_session_tokens, mock_session, mock_session_tokens},
+        mocks::MatrixMockServer,
+    };
+
+    let server = MatrixMockServer::new().await;
+    let oauth_server = server.oauth();
+
+    oauth_server.mock_server_metadata().ok().expect(1..).named("server_metadata").mount().await;
+
+    let issuer = server.server().uri();
+    let client = server.client_builder().unlogged().handle_refresh_tokens().build().await;
+
+    let oidc = client.oidc();
+    oidc.restore_session(mock_session(mock_prev_session_tokens(), issuer)).await.unwrap();
+
+    let num_save_session_callback_calls = Arc::new(Mutex::new(0));
+    client
+        .set_session_callbacks(Box::new(|_| panic!("reload session never called")), {
+            let num_save_session_callback_calls = num_save_session_callback_calls.clone();
+            Box::new(move |_client| {
+                *num_save_session_callback_calls.lock().unwrap() += 1;
+                Ok(())
+            })
+        })
+        .unwrap();
+
+    let mut session_changes = client.subscribe_to_session_changes();
+    assert_eq!(session_changes.try_recv(), Err(TryRecvError::Empty));
+
+    assert_eq!(oidc.session_tokens(), Some(mock_prev_session_tokens()));
+
+    // Refresh token successfully.
+    oauth_server.mock_token().ok().mock_once().named("refresh_token").mount().await;
+
+    oidc.refresh_access_token().await.unwrap();
+    assert_eq!(oidc.session_tokens(), Some(mock_session_tokens()));
+
+    assert_eq!(*num_save_session_callback_calls.lock().unwrap(), 1);
+
+    assert_eq!(session_changes.try_recv(), Ok(SessionChange::TokensRefreshed));
+    assert_eq!(session_changes.try_recv(), Err(TryRecvError::Empty));
+}