From b9f0dafd90df9f26939d1ed245dd32ccb106dff7 Mon Sep 17 00:00:00 2001
From: David Robertson <davidr@element.io>
Date: Tue, 14 Nov 2023 20:16:15 +0000
Subject: [PATCH] Write new column

---
 .../storage/databases/main/end_to_end_keys.py | 32 +++++++++++++++++++
 tests/handlers/test_e2e_keys.py               |  8 ++---
 .../databases/main/test_end_to_end_keys.py    |  9 ++----
 3 files changed, 38 insertions(+), 11 deletions(-)

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
             )
         )