Use foreign keys to simplify logic.

synapse/storage/data_stores/main/schema/delta/58/03persist_ui_auth.sql

@@ -30,5 +30,7 @@ CREATE TABLE IF NOT EXISTS ui_auth_sessions_credentials(
     session_id TEXT NOT NULL,  -- The corresponding UI Auth session.
     stage_type TEXT NOT NULL,  -- The stage type.
     result TEXT NOT NULL,  -- The result of the stage verification, stored as JSON.
-    UNIQUE (session_id, stage_type)
+    UNIQUE (session_id, stage_type),
+    FOREIGN KEY (session_id)
+        REFERENCES ui_auth_sessions (session_id)
 );

synapse/storage/data_stores/main/ui_auth.py

@@ -135,32 +135,19 @@ async def mark_ui_auth_stage_complete(
         Raises:
             StoreError if the session cannot be found.
         """
-        await self.db.runInteraction(
-            "mark_ui_auth_stage_complete",
-            self._mark_ui_auth_stage_complete_txn,
-            session_id,
-            stage_type,
-            result,
-        )
-
-    def _mark_ui_auth_stage_complete_txn(
-        self, txn, session_id: str, stage_type: str, result: Union[str, bool, JsonDict],
-    ):
-        # Get the session to ensure it exists.
-        self.db.simple_select_one_txn(
-            txn,
-            table="ui_auth_sessions",
-            keyvalues={"session_id": session_id},
-            retcols=("session_id",),
-        )
-
         # Add (or update) the results of the current stage to the database.
-        self.db.simple_upsert_txn(
-            txn,
-            table="ui_auth_sessions_credentials",
-            keyvalues={"session_id": session_id, "stage_type": stage_type},
-            values={"result": json.dumps(result)},
-        )
+        try:
+            await self.db.simple_insert(
+                table="ui_auth_sessions_credentials",
+                values={
+                    "session_id": session_id,
+                    "stage_type": stage_type,
+                    "result": json.dumps(result),
+                },
+                desc="mark_ui_auth_stage_complete",
+            )
+        except self.db.engine.module.IntegrityError:
+            raise StoreError(404, "Unknown session ID: %s" % (session_id,))
 
     async def get_completed_ui_auth_stages(
         self, session_id: str

synapse/storage/engines/sqlite.py

@@ -85,6 +85,7 @@ def on_new_connection(self, db_conn):
             prepare_database(db_conn, self, config=None)
 
         db_conn.create_function("rank", 1, _rank)
+        db_conn.execute("PRAGMA foreign_keys = ON;")
 
     def is_deadlock(self, error):
         return False

tests/rest/client/v2_alpha/test_auth.py

@@ -211,11 +211,12 @@ def test_complete_operation_unknown_session(self):
         self.render(request)
         self.assertEqual(request.code, 200)
 
-        invalid_session = session + "invalid"
+        # Attempt to complete an unknown session, which should return an error.
+        unknown_session = session + "unknown"
         request, channel = self.make_request(
             "POST",
             "auth/m.login.recaptcha/fallback/web?session="
-            + invalid_session
+            + unknown_session
             + "&g-recaptcha-response=a",
         )
         self.render(request)