blindly incorporate PR review - needs testing & fixing

5 files changed, 99 insertions, 79 deletions

diff --git a/synapse/api/errors.py b/synapse/api/errors.py
index 8c97e91ba1..d37bcb4082 100644
--- a/synapse/api/errors.py
+++ b/synapse/api/errors.py
@@ -289,9 +289,14 @@ class LimitExceededError(SynapseError):
 class RoomKeysVersionError(SynapseError):
     """A client has tried to upload to a non-current version of the room_keys store
     """
-    def __init__(self, code=403, msg="Wrong room_keys version", current_version=None,
-                 errcode=Codes.WRONG_ROOM_KEYS_VERSION):
-        super(RoomKeysVersionError, self).__init__(code, msg, errcode)
+    def __init__(self, current_version):
+        """
+        Args:
+            current_version (str): the current version of the store they should have used
+        """
+        super(RoomKeysVersionError, self).__init__(
+            403, "Wrong room_keys version", Codes.WRONG_ROOM_KEYS_VERSION
+        )
         self.current_version = current_version
 
     def error_dict(self):
diff --git a/synapse/handlers/e2e_room_keys.py b/synapse/handlers/e2e_room_keys.py
index bd58be6558..dda31fdd24 100644
--- a/synapse/handlers/e2e_room_keys.py
+++ b/synapse/handlers/e2e_room_keys.py
@@ -24,8 +24,21 @@ logger = logging.getLogger(__name__)
 
 
 class E2eRoomKeysHandler(object):
+    """
+    Implements an optional realtime backup mechanism for encrypted E2E megolm room keys.
+    This gives a way for users to store and recover their megolm keys if they lose all
+    their clients. It should also extend easily to future room key mechanisms.
+    The actual payload of the encrypted keys is completely opaque to the handler.
+    """
+
     def __init__(self, hs):
         self.store = hs.get_datastore()
+
+        # Used to lock whenever a client is uploading key data.  This prevents collisions
+        # between clients trying to upload the details of a new session, given all
+        # clients belonging to a user will receive and try to upload a new session at
+        # roughly the same time.  Also used to lock out uploads when the key is being
+        # changed.
         self._upload_linearizer = Linearizer("upload_room_keys_lock")
 
     @defer.inlineCallbacks
@@ -40,33 +53,34 @@ class E2eRoomKeysHandler(object):
 
     @defer.inlineCallbacks
     def delete_room_keys(self, user_id, version, room_id, session_id):
-        yield self.store.delete_e2e_room_keys(user_id, version, room_id, session_id)
+        # lock for consistency with uploading
+        with (yield self._upload_linearizer.queue(user_id)):
+            yield self.store.delete_e2e_room_keys(user_id, version, room_id, session_id)
 
     @defer.inlineCallbacks
     def upload_room_keys(self, user_id, version, room_keys):
 
         # TODO: Validate the JSON to make sure it has the right keys.
 
-        # Check that the version we're trying to upload is the current version
-
-        try:
-            version_info = yield self.get_version_info(user_id, version)
-        except StoreError as e:
-            if e.code == 404:
-                raise SynapseError(404, "Version '%s' not found" % (version,))
-            else:
-                raise e
-
-        if version_info['version'] != version:
-            raise RoomKeysVersionError(current_version=version_info.version)
-
         # XXX: perhaps we should use a finer grained lock here?
         with (yield self._upload_linearizer.queue(user_id)):
-
-            # go through the room_keys
-            for room_id in room_keys['rooms']:
-                for session_id in room_keys['rooms'][room_id]['sessions']:
-                    room_key = room_keys['rooms'][room_id]['sessions'][session_id]
+            # Check that the version we're trying to upload is the current version
+            try:
+                version_info = yield self.get_version_info(user_id, version)
+            except StoreError as e:
+                if e.code == 404:
+                    raise SynapseError(404, "Version '%s' not found" % (version,))
+                else:
+                    raise e
+
+            if version_info['version'] != version:
+                raise RoomKeysVersionError(current_version=version_info.version)
+
+            # go through the room_keys.
+            # XXX: this should/could be done concurrently, given we're in a lock.
+            for room_id, room in room_keys['rooms'].iteritems():
+                for session_id, session in room['sessions'].iteritems():
+                    room_key = session[session_id]
 
                     yield self._upload_room_key(
                         user_id, version, room_id, session_id, room_key
@@ -86,10 +100,29 @@ class E2eRoomKeysHandler(object):
             else:
                 raise e
 
-        # check whether we merge or not. spelling it out with if/elifs rather
-        # than lots of booleans for legibility.
-        upsert = True
+        if _should_replace_room_key(current_room_key, room_key):
+            yield self.store.set_e2e_room_key(
+                user_id, version, room_id, session_id, room_key
+            )
+
+    def _should_replace_room_key(current_room_key, room_key):
+        """
+        Determine whether to replace the current_room_key in our backup for this
+        session (if any) with a new room_key that has been uploaded.
+
+        Args:
+            current_room_key (dict): Optional, the current room_key dict if any
+            room_key (dict): The new room_key dict which may or may not be fit to
+                replace the current_room_key
+
+        Returns:
+            True if current_room_key should be replaced by room_key in the backup
+        """
+
         if current_room_key:
+            # spelt out with if/elifs rather than nested boolean expressions
+            # purely for legibility.
+
             if room_key['is_verified'] and not current_room_key['is_verified']:
                 pass
             elif (
@@ -97,16 +130,11 @@ class E2eRoomKeysHandler(object):
                 current_room_key['first_message_index']
             ):
                 pass
-            elif room_key['forwarded_count'] < room_key['forwarded_count']:
+            elif room_key['forwarded_count'] < current_room_key['forwarded_count']:
                 pass
             else:
-                upsert = False
-
-        # if so, we set the new room_key
-        if upsert:
-            yield self.store.set_e2e_room_key(
-                user_id, version, room_id, session_id, room_key
-            )
+                return False
+        return True
 
     @defer.inlineCallbacks
     def create_version(self, user_id, version_info):
diff --git a/synapse/rest/client/v2_alpha/room_keys.py b/synapse/rest/client/v2_alpha/room_keys.py
index 128b732fb1..70b7b4573f 100644
--- a/synapse/rest/client/v2_alpha/room_keys.py
+++ b/synapse/rest/client/v2_alpha/room_keys.py
@@ -68,6 +68,8 @@ class RoomKeysServlet(RestServlet):
          * lower forwarded_count always wins over higher forwarded_count
 
         We trust the clients not to lie and corrupt their own backups.
+        It also means that if your access_token is stolen, the attacker could
+        delete your backup.
 
         POST /room_keys/keys/!abc:matrix.org/c0ff33?version=1 HTTP/1.1
         Content-Type: application/json
diff --git a/synapse/storage/e2e_room_keys.py b/synapse/storage/e2e_room_keys.py
index 8efca11a8c..c11417c415 100644
--- a/synapse/storage/e2e_room_keys.py
+++ b/synapse/storage/e2e_room_keys.py
@@ -44,30 +44,21 @@ class EndToEndRoomKeyStore(SQLBaseStore):
 
     def set_e2e_room_key(self, user_id, version, room_id, session_id, room_key):
 
-        def _set_e2e_room_key_txn(txn):
-
-            self._simple_upsert_txn(
-                txn,
-                table="e2e_room_keys",
-                keyvalues={
-                    "user_id": user_id,
-                    "room_id": room_id,
-                    "session_id": session_id,
-                },
-                values={
-                    "version": version,
-                    "first_message_index": room_key['first_message_index'],
-                    "forwarded_count": room_key['forwarded_count'],
-                    "is_verified": room_key['is_verified'],
-                    "session_data": room_key['session_data'],
-                },
-                lock=False,
-            )
-
-            return True
-
-        return self.runInteraction(
-            "set_e2e_room_key", _set_e2e_room_key_txn
+        yield self._simple_upsert(
+            table="e2e_room_keys",
+            keyvalues={
+                "user_id": user_id,
+                "room_id": room_id,
+                "session_id": session_id,
+            },
+            values={
+                "version": version,
+                "first_message_index": room_key['first_message_index'],
+                "forwarded_count": room_key['forwarded_count'],
+                "is_verified": room_key['is_verified'],
+                "session_data": room_key['session_data'],
+            },
+            lock=False,
         )
 
     # XXX: this isn't currently used and isn't tested anywhere
@@ -107,7 +98,9 @@ class EndToEndRoomKeyStore(SQLBaseStore):
         )
 
     @defer.inlineCallbacks
-    def get_e2e_room_keys(self, user_id, version, room_id, session_id):
+    def get_e2e_room_keys(
+        self, user_id, version, room_id=room_id, session_id=session_id
+    ):
 
         keyvalues = {
             "user_id": user_id,
@@ -115,8 +108,8 @@ class EndToEndRoomKeyStore(SQLBaseStore):
         }
         if room_id:
             keyvalues['room_id'] = room_id
-        if session_id:
-            keyvalues['session_id'] = session_id
+            if session_id:
+                keyvalues['session_id'] = session_id
 
         rows = yield self._simple_select_list(
             table="e2e_room_keys",
@@ -133,18 +126,10 @@ class EndToEndRoomKeyStore(SQLBaseStore):
             desc="get_e2e_room_keys",
         )
 
-        # perlesque autovivification from https://stackoverflow.com/a/19829714/6764493
-        class AutoVivification(dict):
-            def __getitem__(self, item):
-                try:
-                    return dict.__getitem__(self, item)
-                except KeyError:
-                    value = self[item] = type(self)()
-                    return value
-
-        sessions = AutoVivification()
+        sessions = {}
         for row in rows:
-            sessions['rooms'][row['room_id']]['sessions'][row['session_id']] = {
+            room_entry = sessions['rooms'].setdefault(row['room_id'], {"sessions": {}})
+            room_entry['sessions'][row['session_id']] = {
                 "first_message_index": row["first_message_index"],
                 "forwarded_count": row["forwarded_count"],
                 "is_verified": row["is_verified"],
@@ -154,7 +139,9 @@ class EndToEndRoomKeyStore(SQLBaseStore):
         defer.returnValue(sessions)
 
     @defer.inlineCallbacks
-    def delete_e2e_room_keys(self, user_id, version, room_id, session_id):
+    def delete_e2e_room_keys(
+        self, user_id, version, room_id=room_id, session_id=session_id
+    ):
 
         keyvalues = {
             "user_id": user_id,
@@ -162,8 +149,8 @@ class EndToEndRoomKeyStore(SQLBaseStore):
         }
         if room_id:
             keyvalues['room_id'] = room_id
-        if session_id:
-            keyvalues['session_id'] = session_id
+            if session_id:
+                keyvalues['session_id'] = session_id
 
         yield self._simple_delete(
             table="e2e_room_keys",
diff --git a/synapse/storage/schema/delta/46/e2e_room_keys.sql b/synapse/storage/schema/delta/46/e2e_room_keys.sql
index 0d2a85fbe6..16499ac34c 100644
--- a/synapse/storage/schema/delta/46/e2e_room_keys.sql
+++ b/synapse/storage/schema/delta/46/e2e_room_keys.sql
@@ -25,16 +25,14 @@ CREATE TABLE e2e_room_keys (
     session_data TEXT NOT NULL
 );
 
-CREATE UNIQUE INDEX e2e_room_keys_user_idx ON e2e_room_keys(user_id);
-CREATE UNIQUE INDEX e2e_room_keys_room_idx ON e2e_room_keys(room_id);
-CREATE UNIQUE INDEX e2e_room_keys_session_idx ON e2e_room_keys(session_id);
+CREATE UNIQUE INDEX e2e_room_keys_idx ON e2e_room_keys(user_id, room_id, session_id);
 
 -- the metadata for each generation of encrypted e2e session backups
-CREATE TABLE e2e_room_key_versions (
+CREATE TABLE e2e_room_keys_versions (
     user_id TEXT NOT NULL,
     version TEXT NOT NULL,
     algorithm TEXT NOT NULL,
     auth_data TEXT NOT NULL
 );
 
-CREATE UNIQUE INDEX e2e_room_key_user_idx ON e2e_room_keys(user_id);
+CREATE UNIQUE INDEX e2e_room_keys_versions_user_idx ON e2e_room_keys_versions(user_id);