test(sdk): Handle more cases with ExpectedAccessToken

Signed-off-by: Kévin Commaille <[email protected]>

Author: zecakeh

crates/matrix-sdk/src/authentication/oauth/mod.rs

@@ -1173,6 +1173,7 @@ impl OAuth {
         if let Some(save_session_callback) = self.client.auth_ctx().save_session_callback.get() {
             // Satisfies the save_session_callback invariant: set_session_tokens has
             // been called just above.
+            tracing::debug!("call save_session_callback");
             if let Err(err) = save_session_callback(self.client.clone()) {
                 error!("when saving session after refresh: {err}");
             }
@@ -1183,6 +1184,7 @@ impl OAuth {
             lock.save_in_memory_and_db(&tokens_clone).await?;
         }
 
+        tracing::debug!("broadcast session changed");
         _ = self.client.auth_ctx().session_change_sender.send(SessionChange::TokensRefreshed);
 
         Ok(())

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

@@ -1861,7 +1861,7 @@ pub struct MockEndpoint<'a, T> {
 
 impl<'a, T> MockEndpoint<'a, T> {
     fn new(server: &'a MockServer, mock: MockBuilder, endpoint: T) -> Self {
-        Self { server, mock, endpoint, expected_access_token: ExpectedAccessToken::None }
+        Self { server, mock, endpoint, expected_access_token: ExpectedAccessToken::Ignore }
     }
 
     /// Expect authentication with the default access token on this endpoint.
@@ -1876,12 +1876,30 @@ impl<'a, T> MockEndpoint<'a, T> {
         self
     }
 
-    /// Don't expect authentication with an access token on this endpoint.
+    /// Expect authentication with any access token on this endpoint, regardless
+    /// of its value.
     ///
-    /// This should be used to override the default behavior of the endpoint,
-    /// when the access token is unknown for example.
-    pub fn do_not_expect_access_token(mut self) -> Self {
-        self.expected_access_token = ExpectedAccessToken::None;
+    /// This is useful if we don't want to track the value of the access token.
+    pub fn expect_any_access_token(mut self) -> Self {
+        self.expected_access_token = ExpectedAccessToken::Any;
+        self
+    }
+
+    /// Expect no authentication on this endpoint.
+    ///
+    /// This means that the endpoint will not match if an `AUTHENTICATION`
+    /// header is present.
+    pub fn expect_missing_access_token(mut self) -> Self {
+        self.expected_access_token = ExpectedAccessToken::Missing;
+        self
+    }
+
+    /// Ignore the access token on this endpoint.
+    ///
+    /// This should be used to override the default behavior of an endpoint that
+    /// requires access tokens.
+    pub fn ignore_access_token(mut self) -> Self {
+        self.expected_access_token = ExpectedAccessToken::Ignore;
         self
     }
 
@@ -1927,10 +1945,7 @@ impl<'a, T> MockEndpoint<'a, T> {
     /// # anyhow::Ok(()) });
     /// ```
     pub fn respond_with<R: Respond + 'static>(self, func: R) -> MatrixMock<'a> {
-        let mock = self
-            .expected_access_token
-            .maybe_match_authorization_header(self.mock)
-            .respond_with(func);
+        let mock = self.mock.and(self.expected_access_token).respond_with(func);
         MatrixMock { mock, server: self.server }
     }
 
@@ -1984,6 +1999,17 @@ impl<'a, T> MockEndpoint<'a, T> {
         })))
     }
 
+    /// Returns a mocked endpoint that emulates an unknown token error, i.e
+    /// responds with a 401 HTTP status code and an `M_UNKNOWN_TOKEN` Matrix
+    /// error code.
+    pub fn error_unknown_token(self, soft_logout: bool) -> MatrixMock<'a> {
+        self.respond_with(ResponseTemplate::new(401).set_body_json(json!({
+            "errcode": "M_UNKNOWN_TOKEN",
+            "error": "Unrecognized access token",
+            "soft_logout": soft_logout,
+        })))
+    }
+
     /// Internal helper to return an `{ event_id }` JSON struct along with a 200
     /// ok response.
     fn ok_with_event_id(self, event_id: OwnedEventId) -> MatrixMock<'a> {
@@ -2036,25 +2062,44 @@ impl<'a, T> MockEndpoint<'a, T> {
 
 /// The access token to expect on an endpoint.
 enum ExpectedAccessToken {
-    /// We don't expect an access token.
-    None,
+    /// Ignore any access token or lack thereof.
+    Ignore,
 
     /// We expect the default access token.
     Default,
 
     /// We expect the given access token.
     Custom(&'static str),
+
+    /// We expect any access token.
+    Any,
+
+    /// We expect that there is no access token.
+    Missing,
 }
 
 impl ExpectedAccessToken {
-    /// Match an `Authorization` header on the given mock if one is expected.
-    fn maybe_match_authorization_header(&self, mock: MockBuilder) -> MockBuilder {
-        let token = match self {
-            Self::None => return mock,
-            Self::Default => "1234",
-            Self::Custom(token) => token,
-        };
-        mock.and(header(http::header::AUTHORIZATION, format!("Bearer {token}")))
+    /// Get the access token from the given request.
+    fn access_token(request: &Request) -> Option<&str> {
+        request
+            .headers
+            .get(&http::header::AUTHORIZATION)?
+            .to_str()
+            .ok()?
+            .strip_prefix("Bearer ")
+            .filter(|token| !token.is_empty())
+    }
+}
+
+impl wiremock::Match for ExpectedAccessToken {
+    fn matches(&self, request: &Request) -> bool {
+        match self {
+            Self::Ignore => true,
+            Self::Default => Self::access_token(request) == Some("1234"),
+            Self::Custom(token) => Self::access_token(request) == Some(token),
+            Self::Any => Self::access_token(request).is_some(),
+            Self::Missing => request.headers.get(&http::header::AUTHORIZATION).is_none(),
+        }
     }
 }
 
@@ -3483,14 +3528,6 @@ impl<'a> MockEndpoint<'a, WhoAmIEndpoint> {
             "device_id": device_id,
         })))
     }
-
-    /// Returns an error response with an `M_UNKNOWN_TOKEN`.
-    pub fn err_unknown_token(self) -> MatrixMock<'a> {
-        self.respond_with(ResponseTemplate::new(401).set_body_json(json!({
-            "errcode": "M_UNKNOWN_TOKEN",
-            "error": "Invalid token"
-        })))
-    }
 }
 
 /// A prebuilt mock for `POST /keys/upload` request.

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

@@ -1534,25 +1534,11 @@ async fn test_server_version_without_auth() {
 
     // If we provide an access token, we encounter a failure, likely because the
     // token has expired.
-    Mock::given(method("GET"))
-        .and(path_regex(r"^/_matrix/client/versions"))
-        .and(header("authorization", "Bearer 1234"))
-        .respond_with(ResponseTemplate::new(401))
-        .mount(server.server())
-        .await;
+    server.mock_versions().expect_default_access_token().error_unknown_token(true).mount().await;
 
     // If we do not provide an access token, all is fine as the endpoint does not
     // require one.
-    Mock::given(method("GET"))
-        .and(path_regex(r"^/_matrix/client/versions"))
-        .respond_with(ResponseTemplate::new(200).set_body_json(json!({
-            "versions": [
-                "r0.0.1",
-                "v1.1"
-            ]
-        })))
-        .mount(server.server())
-        .await;
+    server.mock_versions().expect_missing_access_token().ok_with_unstable_features().mount().await;
 
     let request_config = RequestConfig::new().disable_retry();
     client

crates/matrix-sdk/tests/integration/encryption/shared_history.rs

@@ -140,7 +140,7 @@ async fn test_shared_history_out_of_order() {
     let bundle_info = bundle_info.await;
     matrix_mock_server
         .mock_authed_media_download()
-        .do_not_expect_access_token()
+        .expect_any_access_token()
         .ok_bytes(bundle)
         .mock_once()
         .named("media_download")

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

@@ -7,12 +7,15 @@ use assert_matches::assert_matches;
 use assert_matches2::assert_let;
 use matrix_sdk::{
     HttpError, RefreshTokenError, SessionChange, SessionTokens,
-    authentication::matrix::MatrixSession,
+    authentication::{matrix::MatrixSession, oauth::OAuthError},
     config::RequestConfig,
     executor::spawn,
     store::RoomLoadSettings,
     test_utils::{
-        client::mock_session_meta,
+        client::{
+            mock_prev_session_tokens_with_refresh, mock_session_meta,
+            mock_session_tokens_with_refresh, oauth::mock_session,
+        },
         logged_in_client_with_server,
         mocks::{LoginResponseTemplate200, MatrixMockServer},
         no_retry_test_client_with_server, test_client_builder_with_server,
@@ -30,7 +33,7 @@ use serde_json::json;
 use tokio::sync::{broadcast::error::TryRecvError, mpsc};
 use wiremock::{
     Mock, ResponseTemplate,
-    matchers::{body_partial_json, header, method, path, path_regex},
+    matchers::{body_partial_json, header, method, path},
 };
 
 fn session() -> MatrixSession {
@@ -563,17 +566,12 @@ async fn test_refresh_token_handled_other_error() {
 
 #[async_test]
 async fn test_oauth_refresh_token_handled_success() {
-    use matrix_sdk::test_utils::{
-        client::{mock_prev_session_tokens_with_refresh, oauth::mock_session},
-        mocks::MatrixMockServer,
-    };
-
     let server = MatrixMockServer::new().await;
     // Return an error first so the token is refreshed.
     server
         .mock_who_am_i()
         .expect_access_token("prev-access-token")
-        .err_unknown_token()
+        .error_unknown_token(true)
         .expect(1)
         .named("whoami_unknown_token")
         .mount()
@@ -619,20 +617,12 @@ async fn test_oauth_refresh_token_handled_success() {
 
 #[async_test]
 async fn test_oauth_refresh_token_handled_failure() {
-    use matrix_sdk::{
-        authentication::oauth::OAuthError,
-        test_utils::{
-            client::{mock_prev_session_tokens_with_refresh, oauth::mock_session},
-            mocks::MatrixMockServer,
-        },
-    };
-
     let server = MatrixMockServer::new().await;
     // Return an error first so the token is refreshed.
     server
         .mock_who_am_i()
         .expect_access_token("prev-access-token")
-        .err_unknown_token()
+        .error_unknown_token(false)
         .expect(1)
         .named("whoami_unknown_token")
         .mount()
@@ -688,14 +678,6 @@ async fn test_oauth_refresh_token_handled_failure() {
 
 #[async_test]
 async fn test_oauth_handle_refresh_tokens() {
-    use matrix_sdk::test_utils::{
-        client::{
-            mock_prev_session_tokens_with_refresh, mock_session_tokens_with_refresh,
-            oauth::mock_session,
-        },
-        mocks::MatrixMockServer,
-    };
-
     let server = MatrixMockServer::new().await;
     let oauth_server = server.oauth();
 
@@ -767,45 +749,47 @@ async fn test_oauth_handle_refresh_tokens() {
 
 #[async_test]
 async fn test_oauth_handle_refresh_tokens_without_versions() {
-    use matrix_sdk::test_utils::{
-        client::{
-            mock_prev_session_tokens_with_refresh, mock_session_tokens_with_refresh,
-            oauth::mock_session,
-        },
-        mocks::MatrixMockServer,
-    };
-
     let server = MatrixMockServer::new().await;
     let oauth_server = server.oauth();
 
     // If we provide an access token, we encounter a failure, likely because the
     // token has expired.
-    Mock::given(method("GET"))
-        .and(path_regex(r"^/_matrix/client/versions"))
-        .and(header("authorization", "Bearer prev-access-token"))
-        .respond_with(ResponseTemplate::new(401))
-        .mount(server.server())
+    server
+        .mock_versions()
+        .expect_access_token("prev-access-token")
+        .error_unknown_token(true)
+        .expect(1..)
+        .named("versions with expired token")
+        .mount()
         .await;
 
     // If we do not provide an access token, all is fine as the endpoint does not
     // require one.
-    Mock::given(method("GET"))
-        .and(path_regex(r"^/_matrix/client/versions"))
-        .respond_with(ResponseTemplate::new(200).set_body_json(json!({
-            "versions": [
-                "r0.0.1",
-                "v1.1"
-            ]
-        })))
+    server
+        .mock_versions()
+        .expect_missing_access_token()
+        .ok_with_unstable_features()
         .expect(1..)
-        .mount(server.server())
+        .named("unauthenticated versions")
+        .mount()
+        .await;
+
+    // If we provide the new access token, all is fine.
+    server
+        .mock_versions()
+        .expect_default_access_token()
+        .ok_with_unstable_features()
+        .expect(1)
+        .named("versions with fresh token")
+        .mount()
         .await;
 
     oauth_server.mock_server_metadata().ok().expect(1..).named("server_metadata").mount().await;
 
     let client = server
         .client_builder()
         .unlogged()
+        .no_server_versions()
         .on_builder(|builder| builder.handle_refresh_tokens())
         .build()
         .await;