Write new column

synapse/storage/databases/main/end_to_end_keys.py

@@ -1675,3 +1675,35 @@ async def store_e2e_cross_signing_signatures(
             ],
             desc="add_e2e_signing_key",
         )
+
+    async def allow_master_cross_signing_key_replacement_without_uia(
+        self, user_id: str, duration_ms: int
+    ) -> Optional[int]:
+        """Mark this user's latest master key as being replaceable without UIA.
+
+        Said replacement will only be permitted for a short time after calling this
+        function. That time period is controlled by the duration argument.
+
+        Returns:
+            None, if there is no such key.
+            Otherwise, the timestamp before which replacement is allowed without UIA.
+        """
+        timestamp = self._clock.time_msec() + duration_ms
+
+        def impl(txn: LoggingTransaction) -> Optional[int]:
+            txn.execute(
+                """
+                UPDATE e2e_cross_signing_keys
+                SET updatable_without_uia_before_ms = ?
+                WHERE user_id = ? AND keytype = 'master'
+            """,
+                (timestamp, user_id),
+            )
+            if txn.rowcount == 0:
+                return None
+            return timestamp
+
+        return await self.db_pool.runInteraction(
+            "allow_master_cross_signing_key_replacement_without_uia",
+            impl,
+        )

tests/handlers/test_e2e_keys.py

@@ -1626,12 +1626,10 @@ def test_check_cross_signing_setup(self) -> None:
         self.assertIs(replaceable_without_uia, False)
 
         # Set an expiry timestamp in the future.
-        now = self.clock.time_msec()
         self.get_success(
-            self.store.db_pool.simple_update_one(
-                "e2e_cross_signing_keys",
-                {"user_id": alice, "keytype": "master"},
-                {"updatable_without_uia_before_ms": now + 1000},
+            self.store.allow_master_cross_signing_key_replacement_without_uia(
+                alice,
+                1000,
             )
         )
 

tests/storage/databases/main/test_end_to_end_keys.py

@@ -46,12 +46,9 @@ def test_get_master_cross_signing_key_updatable_before(self) -> None:
         self.assertIsNone(timestamp)
 
         # Write an updateable_before timestamp.
-        written_timestamp = 123456789
-        self.get_success(
-            self.store.db_pool.simple_update_one(
-                "e2e_cross_signing_keys",
-                {"user_id": alice, "keytype": "master"},
-                {"updatable_without_uia_before_ms": written_timestamp},
+        written_timestamp = self.get_success(
+            self.store.allow_master_cross_signing_key_replacement_without_uia(
+                alice, 1000
             )
         )