Save the OIDC session ID (sid) with the device on login (#11482)

As a step towards allowing back-channel logout for OIDC.

4 files changed, 188 insertions, 33 deletions

diff --git a/tests/handlers/test_auth.py b/tests/handlers/test_auth.py
index 72e176da75..03b8b8615c 100644
--- a/tests/handlers/test_auth.py
+++ b/tests/handlers/test_auth.py
@@ -71,7 +71,7 @@ class AuthTestCase(unittest.HomeserverTestCase):
 
     def test_short_term_login_token_gives_user_id(self):
         token = self.macaroon_generator.generate_short_term_login_token(
-            self.user1, "", 5000
+            self.user1, "", duration_in_ms=5000
         )
         res = self.get_success(self.auth_handler.validate_short_term_login_token(token))
         self.assertEqual(self.user1, res.user_id)
@@ -94,7 +94,7 @@ class AuthTestCase(unittest.HomeserverTestCase):
 
     def test_short_term_login_token_cannot_replace_user_id(self):
         token = self.macaroon_generator.generate_short_term_login_token(
-            self.user1, "", 5000
+            self.user1, "", duration_in_ms=5000
         )
         macaroon = pymacaroons.Macaroon.deserialize(token)
 
@@ -213,6 +213,6 @@ class AuthTestCase(unittest.HomeserverTestCase):
 
     def _get_macaroon(self):
         token = self.macaroon_generator.generate_short_term_login_token(
-            self.user1, "", 5000
+            self.user1, "", duration_in_ms=5000
         )
         return pymacaroons.Macaroon.deserialize(token)
diff --git a/tests/handlers/test_cas.py b/tests/handlers/test_cas.py
index b625995d12..8705ff8943 100644
--- a/tests/handlers/test_cas.py
+++ b/tests/handlers/test_cas.py
@@ -66,7 +66,13 @@ class CasHandlerTestCase(HomeserverTestCase):
 
         # check that the auth handler got called as expected
         auth_handler.complete_sso_login.assert_called_once_with(
-            "@test_user:test", "cas", request, "redirect_uri", None, new_user=True
+            "@test_user:test",
+            "cas",
+            request,
+            "redirect_uri",
+            None,
+            new_user=True,
+            auth_provider_session_id=None,
         )
 
     def test_map_cas_user_to_existing_user(self):
@@ -89,7 +95,13 @@ class CasHandlerTestCase(HomeserverTestCase):
 
         # check that the auth handler got called as expected
         auth_handler.complete_sso_login.assert_called_once_with(
-            "@test_user:test", "cas", request, "redirect_uri", None, new_user=False
+            "@test_user:test",
+            "cas",
+            request,
+            "redirect_uri",
+            None,
+            new_user=False,
+            auth_provider_session_id=None,
         )
 
         # Subsequent calls should map to the same mxid.
@@ -98,7 +110,13 @@ class CasHandlerTestCase(HomeserverTestCase):
             self.handler._handle_cas_response(request, cas_response, "redirect_uri", "")
         )
         auth_handler.complete_sso_login.assert_called_once_with(
-            "@test_user:test", "cas", request, "redirect_uri", None, new_user=False
+            "@test_user:test",
+            "cas",
+            request,
+            "redirect_uri",
+            None,
+            new_user=False,
+            auth_provider_session_id=None,
         )
 
     def test_map_cas_user_to_invalid_localpart(self):
@@ -116,7 +134,13 @@ class CasHandlerTestCase(HomeserverTestCase):
 
         # check that the auth handler got called as expected
         auth_handler.complete_sso_login.assert_called_once_with(
-            "@f=c3=b6=c3=b6:test", "cas", request, "redirect_uri", None, new_user=True
+            "@f=c3=b6=c3=b6:test",
+            "cas",
+            request,
+            "redirect_uri",
+            None,
+            new_user=True,
+            auth_provider_session_id=None,
         )
 
     @override_config(
@@ -160,7 +184,13 @@ class CasHandlerTestCase(HomeserverTestCase):
 
         # check that the auth handler got called as expected
         auth_handler.complete_sso_login.assert_called_once_with(
-            "@test_user:test", "cas", request, "redirect_uri", None, new_user=True
+            "@test_user:test",
+            "cas",
+            request,
+            "redirect_uri",
+            None,
+            new_user=True,
+            auth_provider_session_id=None,
         )
 
 
diff --git a/tests/handlers/test_oidc.py b/tests/handlers/test_oidc.py
index a25c89bd5b..cfe3de5266 100644
--- a/tests/handlers/test_oidc.py
+++ b/tests/handlers/test_oidc.py
@@ -252,13 +252,6 @@ class OidcHandlerTestCase(HomeserverTestCase):
         with patch.object(self.provider, "load_metadata", patched_load_metadata):
             self.get_failure(self.provider.load_jwks(force=True), RuntimeError)
 
-        # Return empty key set if JWKS are not used
-        self.provider._scopes = []  # not asking the openid scope
-        self.http_client.get_json.reset_mock()
-        jwks = self.get_success(self.provider.load_jwks(force=True))
-        self.http_client.get_json.assert_not_called()
-        self.assertEqual(jwks, {"keys": []})
-
     @override_config({"oidc_config": DEFAULT_CONFIG})
     def test_validate_config(self):
         """Provider metadatas are extensively validated."""
@@ -455,7 +448,13 @@ class OidcHandlerTestCase(HomeserverTestCase):
         self.get_success(self.handler.handle_oidc_callback(request))
 
         auth_handler.complete_sso_login.assert_called_once_with(
-            expected_user_id, "oidc", request, client_redirect_url, None, new_user=True
+            expected_user_id,
+            "oidc",
+            request,
+            client_redirect_url,
+            None,
+            new_user=True,
+            auth_provider_session_id=None,
         )
         self.provider._exchange_code.assert_called_once_with(code)
         self.provider._parse_id_token.assert_called_once_with(token, nonce=nonce)
@@ -482,17 +481,58 @@ class OidcHandlerTestCase(HomeserverTestCase):
         self.provider._fetch_userinfo.reset_mock()
 
         # With userinfo fetching
-        self.provider._scopes = []  # do not ask the "openid" scope
+        self.provider._user_profile_method = "userinfo_endpoint"
+        token = {
+            "type": "bearer",
+            "access_token": "access_token",
+        }
+        self.provider._exchange_code = simple_async_mock(return_value=token)
         self.get_success(self.handler.handle_oidc_callback(request))
 
         auth_handler.complete_sso_login.assert_called_once_with(
-            expected_user_id, "oidc", request, client_redirect_url, None, new_user=False
+            expected_user_id,
+            "oidc",
+            request,
+            client_redirect_url,
+            None,
+            new_user=False,
+            auth_provider_session_id=None,
         )
         self.provider._exchange_code.assert_called_once_with(code)
         self.provider._parse_id_token.assert_not_called()
         self.provider._fetch_userinfo.assert_called_once_with(token)
         self.render_error.assert_not_called()
 
+        # With an ID token, userinfo fetching and sid in the ID token
+        self.provider._user_profile_method = "userinfo_endpoint"
+        token = {
+            "type": "bearer",
+            "access_token": "access_token",
+            "id_token": "id_token",
+        }
+        id_token = {
+            "sid": "abcdefgh",
+        }
+        self.provider._parse_id_token = simple_async_mock(return_value=id_token)
+        self.provider._exchange_code = simple_async_mock(return_value=token)
+        auth_handler.complete_sso_login.reset_mock()
+        self.provider._fetch_userinfo.reset_mock()
+        self.get_success(self.handler.handle_oidc_callback(request))
+
+        auth_handler.complete_sso_login.assert_called_once_with(
+            expected_user_id,
+            "oidc",
+            request,
+            client_redirect_url,
+            None,
+            new_user=False,
+            auth_provider_session_id=id_token["sid"],
+        )
+        self.provider._exchange_code.assert_called_once_with(code)
+        self.provider._parse_id_token.assert_called_once_with(token, nonce=nonce)
+        self.provider._fetch_userinfo.assert_called_once_with(token)
+        self.render_error.assert_not_called()
+
         # Handle userinfo fetching error
         self.provider._fetch_userinfo = simple_async_mock(raises=Exception())
         self.get_success(self.handler.handle_oidc_callback(request))
@@ -776,6 +816,7 @@ class OidcHandlerTestCase(HomeserverTestCase):
             client_redirect_url,
             {"phone": "1234567"},
             new_user=True,
+            auth_provider_session_id=None,
         )
 
     @override_config({"oidc_config": DEFAULT_CONFIG})
@@ -790,7 +831,13 @@ class OidcHandlerTestCase(HomeserverTestCase):
         }
         self.get_success(_make_callback_with_userinfo(self.hs, userinfo))
         auth_handler.complete_sso_login.assert_called_once_with(
-            "@test_user:test", "oidc", ANY, ANY, None, new_user=True
+            "@test_user:test",
+            "oidc",
+            ANY,
+            ANY,
+            None,
+            new_user=True,
+            auth_provider_session_id=None,
         )
         auth_handler.complete_sso_login.reset_mock()
 
@@ -801,7 +848,13 @@ class OidcHandlerTestCase(HomeserverTestCase):
         }
         self.get_success(_make_callback_with_userinfo(self.hs, userinfo))
         auth_handler.complete_sso_login.assert_called_once_with(
-            "@test_user_2:test", "oidc", ANY, ANY, None, new_user=True
+            "@test_user_2:test",
+            "oidc",
+            ANY,
+            ANY,
+            None,
+            new_user=True,
+            auth_provider_session_id=None,
         )
         auth_handler.complete_sso_login.reset_mock()
 
@@ -838,14 +891,26 @@ class OidcHandlerTestCase(HomeserverTestCase):
         }
         self.get_success(_make_callback_with_userinfo(self.hs, userinfo))
         auth_handler.complete_sso_login.assert_called_once_with(
-            user.to_string(), "oidc", ANY, ANY, None, new_user=False
+            user.to_string(),
+            "oidc",
+            ANY,
+            ANY,
+            None,
+            new_user=False,
+            auth_provider_session_id=None,
         )
         auth_handler.complete_sso_login.reset_mock()
 
         # Subsequent calls should map to the same mxid.
         self.get_success(_make_callback_with_userinfo(self.hs, userinfo))
         auth_handler.complete_sso_login.assert_called_once_with(
-            user.to_string(), "oidc", ANY, ANY, None, new_user=False
+            user.to_string(),
+            "oidc",
+            ANY,
+            ANY,
+            None,
+            new_user=False,
+            auth_provider_session_id=None,
         )
         auth_handler.complete_sso_login.reset_mock()
 
@@ -860,7 +925,13 @@ class OidcHandlerTestCase(HomeserverTestCase):
         }
         self.get_success(_make_callback_with_userinfo(self.hs, userinfo))
         auth_handler.complete_sso_login.assert_called_once_with(
-            user.to_string(), "oidc", ANY, ANY, None, new_user=False
+            user.to_string(),
+            "oidc",
+            ANY,
+            ANY,
+            None,
+            new_user=False,
+            auth_provider_session_id=None,
         )
         auth_handler.complete_sso_login.reset_mock()
 
@@ -896,7 +967,13 @@ class OidcHandlerTestCase(HomeserverTestCase):
 
         self.get_success(_make_callback_with_userinfo(self.hs, userinfo))
         auth_handler.complete_sso_login.assert_called_once_with(
-            "@TEST_USER_2:test", "oidc", ANY, ANY, None, new_user=False
+            "@TEST_USER_2:test",
+            "oidc",
+            ANY,
+            ANY,
+            None,
+            new_user=False,
+            auth_provider_session_id=None,
         )
 
     @override_config({"oidc_config": DEFAULT_CONFIG})
@@ -934,7 +1011,13 @@ class OidcHandlerTestCase(HomeserverTestCase):
 
         # test_user is already taken, so test_user1 gets registered instead.
         auth_handler.complete_sso_login.assert_called_once_with(
-            "@test_user1:test", "oidc", ANY, ANY, None, new_user=True
+            "@test_user1:test",
+            "oidc",
+            ANY,
+            ANY,
+            None,
+            new_user=True,
+            auth_provider_session_id=None,
         )
         auth_handler.complete_sso_login.reset_mock()
 
@@ -1018,7 +1101,13 @@ class OidcHandlerTestCase(HomeserverTestCase):
 
         # check that the auth handler got called as expected
         auth_handler.complete_sso_login.assert_called_once_with(
-            "@tester:test", "oidc", ANY, ANY, None, new_user=True
+            "@tester:test",
+            "oidc",
+            ANY,
+            ANY,
+            None,
+            new_user=True,
+            auth_provider_session_id=None,
         )
 
     @override_config(
@@ -1043,7 +1132,13 @@ class OidcHandlerTestCase(HomeserverTestCase):
 
         # check that the auth handler got called as expected
         auth_handler.complete_sso_login.assert_called_once_with(
-            "@tester:test", "oidc", ANY, ANY, None, new_user=True
+            "@tester:test",
+            "oidc",
+            ANY,
+            ANY,
+            None,
+            new_user=True,
+            auth_provider_session_id=None,
         )
 
     @override_config(
@@ -1156,7 +1251,7 @@ async def _make_callback_with_userinfo(
 
     handler = hs.get_oidc_handler()
     provider = handler._providers["oidc"]
-    provider._exchange_code = simple_async_mock(return_value={})
+    provider._exchange_code = simple_async_mock(return_value={"id_token": ""})
     provider._parse_id_token = simple_async_mock(return_value=userinfo)
     provider._fetch_userinfo = simple_async_mock(return_value=userinfo)
 
diff --git a/tests/handlers/test_saml.py b/tests/handlers/test_saml.py
index 8cfc184fef..50551aa6e3 100644
--- a/tests/handlers/test_saml.py
+++ b/tests/handlers/test_saml.py
@@ -130,7 +130,13 @@ class SamlHandlerTestCase(HomeserverTestCase):
 
         # check that the auth handler got called as expected
         auth_handler.complete_sso_login.assert_called_once_with(
-            "@test_user:test", "saml", request, "redirect_uri", None, new_user=True
+            "@test_user:test",
+            "saml",
+            request,
+            "redirect_uri",
+            None,
+            new_user=True,
+            auth_provider_session_id=None,
         )
 
     @override_config({"saml2_config": {"grandfathered_mxid_source_attribute": "mxid"}})
@@ -156,7 +162,13 @@ class SamlHandlerTestCase(HomeserverTestCase):
 
         # check that the auth handler got called as expected
         auth_handler.complete_sso_login.assert_called_once_with(
-            "@test_user:test", "saml", request, "", None, new_user=False
+            "@test_user:test",
+            "saml",
+            request,
+            "",
+            None,
+            new_user=False,
+            auth_provider_session_id=None,
         )
 
         # Subsequent calls should map to the same mxid.
@@ -165,7 +177,13 @@ class SamlHandlerTestCase(HomeserverTestCase):
             self.handler._handle_authn_response(request, saml_response, "")
         )
         auth_handler.complete_sso_login.assert_called_once_with(
-            "@test_user:test", "saml", request, "", None, new_user=False
+            "@test_user:test",
+            "saml",
+            request,
+            "",
+            None,
+            new_user=False,
+            auth_provider_session_id=None,
         )
 
     def test_map_saml_response_to_invalid_localpart(self):
@@ -213,7 +231,13 @@ class SamlHandlerTestCase(HomeserverTestCase):
 
         # test_user is already taken, so test_user1 gets registered instead.
         auth_handler.complete_sso_login.assert_called_once_with(
-            "@test_user1:test", "saml", request, "", None, new_user=True
+            "@test_user1:test",
+            "saml",
+            request,
+            "",
+            None,
+            new_user=True,
+            auth_provider_session_id=None,
         )
         auth_handler.complete_sso_login.reset_mock()
 
@@ -309,7 +333,13 @@ class SamlHandlerTestCase(HomeserverTestCase):
 
         # check that the auth handler got called as expected
         auth_handler.complete_sso_login.assert_called_once_with(
-            "@test_user:test", "saml", request, "redirect_uri", None, new_user=True
+            "@test_user:test",
+            "saml",
+            request,
+            "redirect_uri",
+            None,
+            new_user=True,
+            auth_provider_session_id=None,
         )