Update changelog.d/15345.feature

Co-authored-by: Patrick Cloke <clokep@users.noreply.github.com>

changelog.d/15345.feature

@@ -0,0 +1 @@
+Follow-up to adding experimental feature flags per-user (#15344) which moves experimental features MSC3026 (busy presence), MSC3881 (remotely toggle push notifications for another client), and MSC3967 (Do not require UIA when first uploading cross signing keys) from the experimental config to per-user flags.

synapse/app/generic_worker.py

@@ -51,6 +51,7 @@ from synapse.rest.key.v2 import KeyResource
 from synapse.rest.synapse.client import build_synapse_client_resource_tree
 from synapse.rest.well_known import well_known_resource
 from synapse.server import HomeServer
+from synapse.storage.databases.main import ExperimentalFeaturesStore
 from synapse.storage.databases.main.account_data import AccountDataWorkerStore
 from synapse.storage.databases.main.appservice import (
     ApplicationServiceTransactionWorkerStore,
@@ -146,6 +147,7 @@ class GenericWorkerSlavedStore(
     TransactionWorkerStore,
     LockStore,
     SessionStore,
+    ExperimentalFeaturesStore,
 ):
     # Properties that multiple storage classes define. Tell mypy what the
     # expected type is.

synapse/handlers/presence.py

@@ -63,6 +63,7 @@ from synapse.replication.http.streams import ReplicationGetStreamUpdates
 from synapse.replication.tcp.commands import ClearUserSyncsCommand
 from synapse.replication.tcp.streams import PresenceFederationStream, PresenceStream
 from synapse.storage.databases.main import DataStore
+from synapse.storage.databases.main.experimental_features import ExperimentalFeature
 from synapse.streams import EventSource
 from synapse.types import (
     JsonDict,
@@ -148,8 +149,6 @@ class BasePresenceHandler(abc.ABC):
 
         self._federation_queue = PresenceFederationQueue(hs, self)
 
-        self._busy_presence_enabled = hs.config.experimental.msc3026_enabled
-
         active_presence = self.store.take_presence_startup_info()
         self.user_to_current_state = {state.user_id: state for state in active_presence}
 
@@ -422,8 +421,6 @@ class WorkerPresenceHandler(BasePresenceHandler):
             self.send_stop_syncing, UPDATE_SYNCING_USERS_MS
         )
 
-        self._busy_presence_enabled = hs.config.experimental.msc3026_enabled
-
         hs.get_reactor().addSystemEventTrigger(
             "before",
             "shutdown",
@@ -609,8 +606,12 @@ class WorkerPresenceHandler(BasePresenceHandler):
             PresenceState.BUSY,
         )
 
+        busy_presence_enabled = await self.hs.get_datastores().main.get_feature_enabled(
+            target_user.to_string(), ExperimentalFeature.MSC3026
+        )
+
         if presence not in valid_presence or (
-            presence == PresenceState.BUSY and not self._busy_presence_enabled
+            presence == PresenceState.BUSY and not busy_presence_enabled
         ):
             raise SynapseError(400, "Invalid presence state")
 
@@ -1238,8 +1239,12 @@ class PresenceHandler(BasePresenceHandler):
             PresenceState.BUSY,
         )
 
+        busy_presence_enabled = await self.hs.get_datastores().main.get_feature_enabled(
+            target_user.to_string(), ExperimentalFeature.MSC3026
+        )
+
         if presence not in valid_presence or (
-            presence == PresenceState.BUSY and not self._busy_presence_enabled
+            presence == PresenceState.BUSY and not busy_presence_enabled
         ):
             raise SynapseError(400, "Invalid presence state")
 
@@ -1257,7 +1262,7 @@ class PresenceHandler(BasePresenceHandler):
             new_fields["status_msg"] = status_msg
 
         if presence == PresenceState.ONLINE or (
-            presence == PresenceState.BUSY and self._busy_presence_enabled
+            presence == PresenceState.BUSY and busy_presence_enabled
         ):
             new_fields["last_active_ts"] = self.clock.time_msec()
 

synapse/rest/admin/experimental_features.py

@@ -13,7 +13,6 @@
 #  limitations under the License.
 
 
-from enum import Enum
 from http import HTTPStatus
 from typing import TYPE_CHECKING, Dict, Tuple
 
@@ -21,22 +20,13 @@ from synapse.api.errors import SynapseError
 from synapse.http.servlet import RestServlet, parse_json_object_from_request
 from synapse.http.site import SynapseRequest
 from synapse.rest.admin import admin_patterns, assert_requester_is_admin
+from synapse.storage.databases.main.experimental_features import ExperimentalFeature
 from synapse.types import JsonDict, UserID
 
 if TYPE_CHECKING:
     from synapse.server import HomeServer
 
 
-class ExperimentalFeature(str, Enum):
-    """
-    Currently supported per-user features
-    """
-
-    MSC3026 = "msc3026"
-    MSC3881 = "msc3881"
-    MSC3967 = "msc3967"
-
-
 class ExperimentalFeaturesRestServlet(RestServlet):
     """
     Enable or disable experimental features for a user or determine which features are enabled

synapse/rest/client/keys.py

@@ -31,6 +31,7 @@ from synapse.http.site import SynapseRequest
 from synapse.logging.opentracing import log_kv, set_tag
 from synapse.replication.http.devices import ReplicationUploadKeysForUserRestServlet
 from synapse.rest.client._base import client_patterns, interactive_auth_handler
+from synapse.storage.databases.main.experimental_features import ExperimentalFeature
 from synapse.types import JsonDict, StreamToken
 from synapse.util.cancellation import cancellable
 
@@ -375,7 +376,11 @@ class SigningKeyUploadServlet(RestServlet):
         user_id = requester.user.to_string()
         body = parse_json_object_from_request(request)
 
-        if self.hs.config.experimental.msc3967_enabled:
+        msc3967_enabled = await self.hs.get_datastores().main.get_feature_enabled(
+            user_id, ExperimentalFeature.MSC3967
+        )
+
+        if msc3967_enabled:
             if await self.e2e_keys_handler.is_cross_signing_set_up_for_user(user_id):
                 # If we already have a master key then cross signing is set up and we require UIA to reset
                 await self.auth_handler.validate_user_via_ui_auth(

synapse/rest/client/pusher.py

@@ -27,6 +27,7 @@ from synapse.http.site import SynapseRequest
 from synapse.push import PusherConfigException
 from synapse.rest.client._base import client_patterns
 from synapse.rest.synapse.client.unsubscribe import UnsubscribeResource
+from synapse.storage.databases.main.experimental_features import ExperimentalFeature
 from synapse.types import JsonDict
 
 if TYPE_CHECKING:
@@ -42,7 +43,6 @@ class PushersRestServlet(RestServlet):
         super().__init__()
         self.hs = hs
         self.auth = hs.get_auth()
-        self._msc3881_enabled = self.hs.config.experimental.msc3881_enabled
 
     async def on_GET(self, request: SynapseRequest) -> Tuple[int, JsonDict]:
         requester = await self.auth.get_user_by_req(request)
@@ -54,8 +54,12 @@ class PushersRestServlet(RestServlet):
 
         pusher_dicts = [p.as_dict() for p in pushers]
 
+        msc3881_enabled = await self.hs.get_datastores().main.get_feature_enabled(
+            user.to_string(), ExperimentalFeature.MSC3881
+        )
+
         for pusher in pusher_dicts:
-            if self._msc3881_enabled:
+            if msc3881_enabled:
                 pusher["org.matrix.msc3881.enabled"] = pusher["enabled"]
                 pusher["org.matrix.msc3881.device_id"] = pusher["device_id"]
             del pusher["enabled"]
@@ -73,7 +77,6 @@ class PushersSetRestServlet(RestServlet):
         self.auth = hs.get_auth()
         self.notifier = hs.get_notifier()
         self.pusher_pool = self.hs.get_pusherpool()
-        self._msc3881_enabled = self.hs.config.experimental.msc3881_enabled
 
     async def on_POST(self, request: SynapseRequest) -> Tuple[int, JsonDict]:
         requester = await self.auth.get_user_by_req(request)
@@ -113,7 +116,11 @@ class PushersSetRestServlet(RestServlet):
             append = content["append"]
 
         enabled = True
-        if self._msc3881_enabled and "org.matrix.msc3881.enabled" in content:
+        msc3881_enabled = await self.hs.get_datastores().main.get_feature_enabled(
+            user.to_string(), ExperimentalFeature.MSC3881
+        )
+
+        if msc3881_enabled and "org.matrix.msc3881.enabled" in content:
             enabled = content["org.matrix.msc3881.enabled"]
 
         if not append:

synapse/rest/client/versions.py

@@ -96,7 +96,7 @@ class VersionsRestServlet(RestServlet):
                     "io.element.e2ee_forced.private": self.e2ee_forced_private,
                     "io.element.e2ee_forced.trusted_private": self.e2ee_forced_trusted_private,
                     # Supports the busy presence state described in MSC3026.
-                    "org.matrix.msc3026.busy_presence": self.config.experimental.msc3026_enabled,
+                    "org.matrix.msc3026.busy_presence": True,
                     # Supports receiving private read receipts as per MSC2285
                     "org.matrix.msc2285.stable": True,  # TODO: Remove when MSC2285 becomes a part of the spec
                     # Supports filtering of /publicRooms by room type as per MSC3827
@@ -115,7 +115,7 @@ class VersionsRestServlet(RestServlet):
                     # Adds support for login token requests as per MSC3882
                     "org.matrix.msc3882": self.config.experimental.msc3882_enabled,
                     # Adds support for remotely enabling/disabling pushers, as per MSC3881
-                    "org.matrix.msc3881": self.config.experimental.msc3881_enabled,
+                    "org.matrix.msc3881": True,
                     # Adds support for filtering /messages by event relation.
                     "org.matrix.msc3874": self.config.experimental.msc3874_enabled,
                     # Adds support for simple HTTP rendezvous as per MSC3886

synapse/storage/databases/main/experimental_features.py

@@ -12,6 +12,7 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
+from enum import Enum
 from typing import TYPE_CHECKING, Dict
 
 from synapse.storage.database import DatabasePool, LoggingDatabaseConnection
@@ -20,10 +21,19 @@ from synapse.types import StrCollection
 from synapse.util.caches.descriptors import cached
 
 if TYPE_CHECKING:
-    from synapse.rest.admin.experimental_features import ExperimentalFeature
     from synapse.server import HomeServer
 
 
+class ExperimentalFeature(str, Enum):
+    """
+    Currently supported per-user features
+    """
+
+    MSC3026 = "msc3026"
+    MSC3881 = "msc3881"
+    MSC3967 = "msc3967"
+
+
 class ExperimentalFeaturesStore(CacheInvalidationWorkerStore):
     def __init__(
         self,
@@ -73,3 +83,41 @@ class ExperimentalFeaturesStore(CacheInvalidationWorkerStore):
             )
 
             await self.invalidate_cache_and_stream("list_enabled_features", (user,))
+
+    async def get_feature_enabled(
+        self, user_id: str, feature: "ExperimentalFeature"
+    ) -> bool:
+        """
+        Checks to see if a given feature is enabled for the user
+
+        Args:
+            user_id: the user to be queried on
+            feature: the feature in question
+        Returns:
+                True if the feature is enabled, False if it is not or if the feature was
+                not found.
+        """
+
+        # check first if feature is enabled in the config
+        if feature == ExperimentalFeature.MSC3026:
+            globally_enabled = self.hs.config.experimental.msc3026_enabled
+        elif feature == ExperimentalFeature.MSC3881:
+            globally_enabled = self.hs.config.experimental.msc3881_enabled
+        else:
+            globally_enabled = self.hs.config.experimental.msc3967_enabled
+
+        if globally_enabled:
+            return globally_enabled
+
+        # if it's not enabled globally, check if it is enabled per-user
+        res = await self.db_pool.simple_select_one(
+            "per_user_experimental_features",
+            {"user_id": user_id, "feature": feature},
+            ["enabled"],
+            allow_none=True,
+        )
+
+        # None and false are treated the same
+        db_enabled = bool(res)
+
+        return db_enabled

tests/handlers/test_presence.py

@@ -36,7 +36,7 @@ from synapse.handlers.presence import (
     handle_update,
 )
 from synapse.rest import admin
-from synapse.rest.client import room
+from synapse.rest.client import login, room
 from synapse.server import HomeServer
 from synapse.types import JsonDict, UserID, get_domain_from_id
 from synapse.util import Clock
@@ -514,9 +514,13 @@ class PresenceTimeoutTestCase(unittest.TestCase):
 
 
 class PresenceHandlerTestCase(BaseMultiWorkerStreamTestCase):
+    servlets = [admin.register_servlets, login.register_servlets]
+
     def prepare(self, reactor: MemoryReactor, clock: Clock, hs: HomeServer) -> None:
         self.presence_handler = hs.get_presence_handler()
         self.clock = hs.get_clock()
+        self.user = self.register_user("test", "pass", True)
+        self.admin_user_tok = self.login("test", "pass")
 
     def test_external_process_timeout(self) -> None:
         """Test that if an external process doesn't update the records for a while
@@ -724,6 +728,62 @@ class PresenceHandlerTestCase(BaseMultiWorkerStreamTestCase):
         # our status message should be the same as it was before
         self.assertEqual(state.status_msg, status_msg)
 
+    @parameterized.expand([(False,), (True,)])
+    def test_set_presence_from_syncing_keeps_busy_via_admin(
+        self, test_with_workers: bool
+    ) -> None:
+        """Test that presence set by syncing doesn't affect busy status, with the busy status
+        enabled via the admin api.
+
+        Args:
+            test_with_workers: If True, check the presence state of the user by calling
+                /sync against a worker, rather than the main process.
+        """
+        status_msg = "I'm busy!"
+
+        # activate busy state via admin api
+        url = f"/_synapse/admin/v1/experimental_features/{self.user}"
+        channel = self.make_request(
+            "PUT",
+            url,
+            content={
+                "features": {"msc3026": True},
+            },
+            access_token=self.admin_user_tok,
+        )
+        self.assertEqual(channel.code, 200)
+
+        # By default, we call /sync against the main process.
+        worker_to_sync_against = self.hs
+        if test_with_workers:
+            # Create a worker and use it to handle /sync traffic instead.
+            # This is used to test that presence changes get replicated from workers
+            # to the main process correctly.
+            worker_to_sync_against = self.make_worker_hs(
+                "synapse.app.generic_worker", {"worker_name": "presence_writer"}
+            )
+
+        # Set presence to BUSY
+        self._set_presencestate_with_status_msg(
+            self.user, PresenceState.BUSY, status_msg
+        )
+
+        # Perform a sync with a presence state other than busy. This should NOT change
+        # our presence status; we only change from busy if we explicitly set it via
+        # /presence/*.
+        self.get_success(
+            worker_to_sync_against.get_presence_handler().user_syncing(
+                self.user, True, PresenceState.ONLINE
+            )
+        )
+
+        # Check against the main process that the user's presence did not change.
+        state = self.get_success(
+            self.presence_handler.get_state(UserID.from_string(self.user))
+        )
+        # we should still be busy
+        self.assertEqual(state.state, PresenceState.BUSY)
+
     @parameterized.expand([(False,), (True,)])
     @unittest.override_config(
         {
@@ -732,16 +792,16 @@ class PresenceHandlerTestCase(BaseMultiWorkerStreamTestCase):
             },
         }
     )
-    def test_set_presence_from_syncing_keeps_busy(
+    def test_set_presence_from_syncing_keeps_busy_via_config(
         self, test_with_workers: bool
     ) -> None:
-        """Test that presence set by syncing doesn't affect busy status
+        """Test that presence set by syncing doesn't affect busy status, with the busy status
+        enabled via the config
 
         Args:
             test_with_workers: If True, check the presence state of the user by calling
                 /sync against a worker, rather than the main process.
         """
-        user_id = "@test:server"
         status_msg = "I'm busy!"
 
         # By default, we call /sync against the main process.
@@ -755,20 +815,22 @@ class PresenceHandlerTestCase(BaseMultiWorkerStreamTestCase):
             )
 
         # Set presence to BUSY
-        self._set_presencestate_with_status_msg(user_id, PresenceState.BUSY, status_msg)
+        self._set_presencestate_with_status_msg(
+            self.user, PresenceState.BUSY, status_msg
+        )
 
         # Perform a sync with a presence state other than busy. This should NOT change
         # our presence status; we only change from busy if we explicitly set it via
         # /presence/*.
         self.get_success(
             worker_to_sync_against.get_presence_handler().user_syncing(
-                user_id, True, PresenceState.ONLINE
+                self.user, True, PresenceState.ONLINE
             )
         )
 
         # Check against the main process that the user's presence did not change.
         state = self.get_success(
-            self.presence_handler.get_state(UserID.from_string(user_id))
+            self.presence_handler.get_state(UserID.from_string(self.user))
         )
         # we should still be busy
         self.assertEqual(state.state, PresenceState.BUSY)

tests/push/test_http.py

@@ -22,6 +22,7 @@ from synapse.logging.context import make_deferred_yieldable
 from synapse.push import PusherConfig, PusherConfigException
 from synapse.rest.client import login, push_rule, pusher, receipts, room
 from synapse.server import HomeServer
+from synapse.storage.databases.main.experimental_features import ExperimentalFeature
 from synapse.types import JsonDict
 from synapse.util import Clock
 
@@ -36,6 +37,7 @@ class HTTPPusherTests(HomeserverTestCase):
         receipts.register_servlets,
         push_rule.register_servlets,
         pusher.register_servlets,
+        synapse.rest.admin.register_servlets,
     ]
     user_id = True
     hijack_auth = False
@@ -820,8 +822,10 @@ class HTTPPusherTests(HomeserverTestCase):
         self.assertEqual(len(self.push_attempts), 1)
 
     @override_config({"experimental_features": {"msc3881_enabled": True}})
-    def test_disable(self) -> None:
-        """Tests that disabling a pusher means it's not pushed to anymore."""
+    def test_disable_via_config(self) -> None:
+        """Tests that disabling a pusher means it's not pushed to anymore, with the
+        ability to disable a pusher enabled via the config.
+        """
         user_id, access_token = self._make_user_with_pusher("user")
         other_user_id, other_access_token = self._make_user_with_pusher("otheruser")
 
@@ -848,7 +852,50 @@ class HTTPPusherTests(HomeserverTestCase):
         self.assertFalse(enabled)
         self.assertTrue(isinstance(enabled, bool))
 
-    @override_config({"experimental_features": {"msc3881_enabled": True}})
+    def test_disable_via_admin(self) -> None:
+        """Tests that disabling a pusher means it's not pushed to anymore,
+        with the ability to disable a pusher enabled via the admin api.
+        """
+        user_id, access_token = self._make_user_with_pusher("user")
+        other_user_id, other_access_token = self._make_user_with_pusher("otheruser")
+        self.register_user("admin", "pass", True)
+        admin_tok = self.login("admin", "pass")
+
+        room = self.helper.create_room_as(user_id, tok=access_token)
+        self.helper.join(room=room, user=other_user_id, tok=other_access_token)
+
+        # enable msc3881 per_user flag via the admin api
+        url = f"/_synapse/admin/v1/experimental_features/{user_id}"
+        channel = self.make_request(
+            "PUT",
+            url,
+            content={
+                "features": {"msc3881": True},
+            },
+            access_token=admin_tok,
+        )
+        self.assertEqual(channel.code, 200)
+
+        # Send a message and check that it generated a push.
+        self.helper.send(room, body="Hi!", tok=other_access_token)
+        self.assertEqual(len(self.push_attempts), 1)
+
+        # Disable the pusher.
+        self._set_pusher(user_id, access_token, enabled=False)
+
+        # Send another message and check that it did not generate a push.
+        self.helper.send(room, body="Hi!", tok=other_access_token)
+        self.assertEqual(len(self.push_attempts), 1)
+
+        # Get the pushers for the user and check that it is marked as disabled.
+        channel = self.make_request("GET", "/pushers", access_token=access_token)
+        self.assertEqual(channel.code, 200)
+        self.assertEqual(len(channel.json_body["pushers"]), 1)
+
+        enabled = channel.json_body["pushers"][0]["org.matrix.msc3881.enabled"]
+        self.assertFalse(enabled)
+        self.assertTrue(isinstance(enabled, bool))
+
     def test_enable(self) -> None:
         """Tests that enabling a disabled pusher means it gets pushed to."""
         # Create the user with the pusher already disabled.
@@ -858,6 +905,13 @@ class HTTPPusherTests(HomeserverTestCase):
         room = self.helper.create_room_as(user_id, tok=access_token)
         self.helper.join(room=room, user=other_user_id, tok=other_access_token)
 
+        # enable msc3881 per_user flag
+        self.get_success(
+            self.hs.get_datastores().main.set_features_for_user(
+                user_id, {ExperimentalFeature.MSC3881: True}
+            )
+        )
+
         # Send a message and check that it did not generate a push.
         self.helper.send(room, body="Hi!", tok=other_access_token)
         self.assertEqual(len(self.push_attempts), 0)
@@ -878,7 +932,7 @@ class HTTPPusherTests(HomeserverTestCase):
         self.assertTrue(enabled)
         self.assertTrue(isinstance(enabled, bool))
 
-    @override_config({"experimental_features": {"msc3881_enabled": True}})
+    # @override_config({"experimental_features": {"msc3881_enabled": True}})
     def test_null_enabled(self) -> None:
         """Tests that a pusher that has an 'enabled' column set to NULL (eg pushers
         created before the column was introduced) is considered enabled.
@@ -887,6 +941,13 @@ class HTTPPusherTests(HomeserverTestCase):
         # database.
         user_id, access_token = self._make_user_with_pusher("user", enabled=None)  # type: ignore[arg-type]
 
+        # enable msc3881 per_user flag
+        self.get_success(
+            self.hs.get_datastores().main.set_features_for_user(
+                user_id, {ExperimentalFeature.MSC3881: True}
+            )
+        )
+
         channel = self.make_request("GET", "/pushers", access_token=access_token)
         self.assertEqual(channel.code, 200)
         self.assertEqual(len(channel.json_body["pushers"]), 1)
@@ -922,14 +983,20 @@ class HTTPPusherTests(HomeserverTestCase):
         self.assertEqual(len(pushers), 1)
         self.assertEqual(pushers[0].device_id, device_id)
 
-    @override_config({"experimental_features": {"msc3881_enabled": True}})
     def test_device_id(self) -> None:
         """Tests that a pusher created with a given device ID shows that device ID in
         GET /pushers requests.
         """
-        self.register_user("user", "pass")
+        user = self.register_user("user", "pass")
         access_token = self.login("user", "pass")
 
+        # enable msc3881 per_user flag
+        self.get_success(
+            self.hs.get_datastores().main.set_features_for_user(
+                user, {ExperimentalFeature.MSC3881: True}
+            )
+        )
+
         # We create the pusher with an HTTP request rather than with
         # _make_user_with_pusher so that we can test the device ID is correctly set when
         # creating a pusher via an API call.

tests/rest/client/test_keys.py

@@ -36,6 +36,7 @@ class KeyQueryTestCase(unittest.HomeserverTestCase):
         keys.register_servlets,
         admin.register_servlets_for_client_rest_resource,
         login.register_servlets,
+        admin.register_servlets,
     ]
 
     def test_rejects_device_id_ice_key_outside_of_list(self) -> None:
@@ -205,12 +206,12 @@ class KeyQueryTestCase(unittest.HomeserverTestCase):
 
     @override_config(
         {
-            "experimental_features": {"msc3967_enabled": True},
             "ui_auth": {"session_timeout": "15s"},
+            "experimental_features": {"msc3967_enabled": True},
         }
     )
-    def test_device_signing_with_msc3967(self) -> None:
-        """Device signing key follows MSC3967 behaviour when enabled."""
+    def test_device_signing_with_msc3967_via_config(self) -> None:
+        """Device signing key follows MSC3967 behaviour when enabled in config."""
         password = "wonderland"
         device_id = "ABCDEFGHI"
         alice_id = self.register_user("alice", password)
@@ -259,3 +260,72 @@ class KeyQueryTestCase(unittest.HomeserverTestCase):
             alice_token,
         )
         self.assertEqual(channel.code, HTTPStatus.OK, channel.result)
+
+    @override_config(
+        {
+            "ui_auth": {"session_timeout": "15s"},
+        }
+    )
+    def test_device_signing_with_msc3967_via_admin(self) -> None:
+        """Device signing key follows MSC3967 behaviour when enabled for user via admin api."""
+        password = "wonderland"
+        device_id = "ABCDEFGHI"
+        alice_id = self.register_user("alice", password)
+        alice_token = self.login("alice", password, device_id=device_id)
+        self.register_user("admin", "pass", True)
+        admin_tok = self.login("admin", "pass")
+
+        url = f"/_synapse/admin/v1/experimental_features/{alice_id}"
+        channel = self.make_request(
+            "PUT",
+            url,
+            content={
+                "features": {"msc3967": True},
+            },
+            access_token=admin_tok,
+        )
+        self.assertEqual(channel.code, 200)
+
+        keys1 = self.make_device_keys(alice_id, device_id)
+
+        # Initial request should succeed as no existing keys are present.
+        channel = self.make_request(
+            "POST",
+            "/_matrix/client/v3/keys/device_signing/upload",
+            keys1,
+            alice_token,
+        )
+        self.assertEqual(channel.code, HTTPStatus.OK, channel.result)
+
+        keys2 = self.make_device_keys(alice_id, device_id)
+
+        # Subsequent request should require UIA as keys already exist even though session_timeout is set.
+        channel = self.make_request(
+            "POST",
+            "/_matrix/client/v3/keys/device_signing/upload",
+            keys2,
+            alice_token,
+        )
+        self.assertEqual(channel.code, HTTPStatus.UNAUTHORIZED, channel.result)
+
+        # Grab the session
+        session = channel.json_body["session"]
+        # Ensure that flows are what is expected.
+        self.assertIn({"stages": ["m.login.password"]}, channel.json_body["flows"])
+
+        # add UI auth
+        keys2["auth"] = {
+            "type": "m.login.password",
+            "identifier": {"type": "m.id.user", "user": alice_id},
+            "password": password,
+            "session": session,
+        }
+
+        # Request should complete
+        channel = self.make_request(
+            "POST",
+            "/_matrix/client/v3/keys/device_signing/upload",
+            keys2,
+            alice_token,
+        )
+        self.assertEqual(channel.code, HTTPStatus.OK, channel.result)