Persist user interactive authentication sessions

synapse/handlers/auth.py

@@ -111,7 +111,7 @@ def __init__(self, hs):
         self._clock = self.hs.get_clock()
 
         # Expire old UI auth sessions after a period of time.
-        if hs.config.worker_name is None:
+        if hs.config.worker_app is None:
             self._clock.looping_call(
                 run_as_background_process,
                 5 * 60 * 1000,

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

@@ -18,7 +18,7 @@
 import attr
 
 import synapse.util.stringutils as stringutils
-from synapse.api.errors import StoreError, SynapseError
+from synapse.api.errors import StoreError
 from synapse.storage._base import SQLBaseStore
 from synapse.types import JsonDict
 
@@ -28,7 +28,7 @@ class UIAuthSessionData:
     session_id = attr.ib(type=str)
     # The dictionary from the client root level, not the 'auth' key.
     clientdict = attr.ib(type=JsonDict)
-    # The URI and method the session was intiatied with, these are checked at
+    # The URI and method the session was intiatied with. These are checked at
     # each stage of the authentication to ensure that the asked for operation
     # has not changed.
     uri = attr.ib(type=str)
@@ -66,10 +66,10 @@ async def create_ui_auth_session(
             description:
                 A string description of the operation that the current
                 authentication is authorising.
-
         Returns:
             The newly created session.
-
+        Raise:
+            StoreError if a unique session ID cannot be generated.
         """
         # The clientdict gets stored as JSON.
         clientdict_json = json.dumps(clientdict)
@@ -109,17 +109,14 @@ async def get_ui_auth_session(self, session_id: str) -> UIAuthSessionData:
         Returns:
             A dict containing the device information.
         Raises:
-            SynapseError: if the session is not found.
+            StoreError if the session is not found.
         """
-        try:
-            result = await self.db.simple_select_one(
-                table="ui_auth_sessions",
-                keyvalues={"session_id": session_id},
-                retcols=("clientdict", "uri", "method", "description"),
-                desc="get_ui_auth_session",
-            )
-        except StoreError:
-            raise SynapseError(400, "Unknown session ID: %s" % session_id)
+        result = await self.db.simple_select_one(
+            table="ui_auth_sessions",
+            keyvalues={"session_id": session_id},
+            retcols=("clientdict", "uri", "method", "description"),
+            desc="get_ui_auth_session",
+        )
 
         result["clientdict"] = json.loads(result["clientdict"])
 
@@ -135,36 +132,22 @@ async def mark_ui_auth_stage_complete(
             session_id: The ID of the corresponding session.
             stage_type: The completed stage type.
             result: The result of the stage verification.
+        Raises:
+            StoreError if the session cannot be found.
         """
-        await self.db.runInteraction(
-            "mark_ui_auth_stage_complete",
-            self._mark_ui_auth_stage_complete,
-            session_id,
-            stage_type,
-            result,
-        )
-
-    def _mark_ui_auth_stage_complete(
-        self, txn, session_id: str, stage_type: str, result: Union[str, bool, JsonDict],
-    ):
+        # Add (or update) the results of the current stage to the database.
+        #
+        # Note that we need to allow for the same stage to complete multiple
+        # times here so that registration is idempotent.
         try:
-            # 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",),
+            await self.db.simple_upsert(
+                table="ui_auth_sessions_credentials",
+                keyvalues={"session_id": session_id, "stage_type": stage_type},
+                values={"result": json.dumps(result)},
+                desc="mark_ui_auth_stage_complete",
             )
-        except StoreError:
-            raise SynapseError(400, "Unknown session ID: %s" % 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)},
-        )
+        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
@@ -177,8 +160,6 @@ async def get_completed_ui_auth_stages(
         Returns:
             The completed stages mapped to the result of the verification of
             that auth-type.
-        Raises:
-            StoreError: if the session is not found.
         """
         results = {}
         for row in await self.db.simple_select_list(
@@ -201,16 +182,18 @@ async def set_ui_auth_session_data(self, session_id: str, key: str, value: Any):
             session_id: The ID of this session as returned from check_auth
             key: The key to store the data under
             value: The data to store
+        Raises:
+            StoreError if the session cannot be found.
         """
         await self.db.runInteraction(
             "set_ui_auth_session_data",
-            self._set_ui_auth_session_data,
+            self._set_ui_auth_session_data_txn,
             session_id,
             key,
             value,
         )
 
-    def _set_ui_auth_session_data(self, txn, session_id: str, key: str, value: Any):
+    def _set_ui_auth_session_data_txn(self, txn, session_id: str, key: str, value: Any):
         # Get the current value.
         result = self.db.simple_select_one_txn(
             txn,
@@ -240,6 +223,8 @@ async def get_ui_auth_session_data(
             session_id: The ID of this session as returned from check_auth
             key: The key to store the data under
             default: Value to return if the key has not been set
+        Raises:
+            StoreError if the session cannot be found.
         """
         result = await self.db.simple_select_one(
             table="ui_auth_sessions",
@@ -265,11 +250,11 @@ def delete_old_ui_auth_sessions(self, expiration_time: int):
         """
         return self.db.runInteraction(
             "delete_old_ui_auth_sessions",
-            self._delete_old_ui_auth_sessions,
+            self._delete_old_ui_auth_sessions_txn,
             expiration_time,
         )
 
-    def _delete_old_ui_auth_sessions(self, txn, expiration_time: int):
+    def _delete_old_ui_auth_sessions_txn(self, txn, expiration_time: int):
         # Get the expired sessions.
         sql = "SELECT session_id FROM ui_auth_sessions WHERE creation_time <= ?"
         txn.execute(sql, [expiration_time])

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,12 +211,13 @@ 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)
-        self.assertEqual(request.code, 400)
+        self.assertEqual(request.code, 404)