diff --git a/synapse/storage/databases/main/end_to_end_keys.py b/synapse/storage/databases/main/end_to_end_keys.py index 2e3bd42a7577..8ebccf5634ec 100644 --- a/synapse/storage/databases/main/end_to_end_keys.py +++ b/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, + ) diff --git a/tests/handlers/test_e2e_keys.py b/tests/handlers/test_e2e_keys.py index ebc505be0036..07eb63f95efb 100644 --- a/tests/handlers/test_e2e_keys.py +++ b/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, ) ) diff --git a/tests/storage/databases/main/test_end_to_end_keys.py b/tests/storage/databases/main/test_end_to_end_keys.py index 8eebecbd112f..7a610ea191d2 100644 --- a/tests/storage/databases/main/test_end_to_end_keys.py +++ b/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 ) )