diff --git a/synapse/storage/databases/main/e2e_room_keys.py b/synapse/storage/databases/main/e2e_room_keys.py
new file mode 100644
index 0000000000..2eeb9f97dc
--- /dev/null
+++ b/synapse/storage/databases/main/e2e_room_keys.py
@@ -0,0 +1,436 @@
+# -*- coding: utf-8 -*-
+# Copyright 2017 New Vector Ltd
+# Copyright 2019 Matrix.org Foundation C.I.C.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from synapse.api.errors import StoreError
+from synapse.logging.opentracing import log_kv, trace
+from synapse.storage._base import SQLBaseStore, db_to_json
+from synapse.util import json_encoder
+
+
+class EndToEndRoomKeyStore(SQLBaseStore):
+ async def update_e2e_room_key(
+ self, user_id, version, room_id, session_id, room_key
+ ):
+ """Replaces the encrypted E2E room key for a given session in a given backup
+
+ Args:
+ user_id(str): the user whose backup we're setting
+ version(str): the version ID of the backup we're updating
+ room_id(str): the ID of the room whose keys we're setting
+ session_id(str): the session whose room_key we're setting
+ room_key(dict): the room_key being set
+ Raises:
+ StoreError
+ """
+
+ await self.db_pool.simple_update_one(
+ table="e2e_room_keys",
+ keyvalues={
+ "user_id": user_id,
+ "version": version,
+ "room_id": room_id,
+ "session_id": session_id,
+ },
+ updatevalues={
+ "first_message_index": room_key["first_message_index"],
+ "forwarded_count": room_key["forwarded_count"],
+ "is_verified": room_key["is_verified"],
+ "session_data": json_encoder.encode(room_key["session_data"]),
+ },
+ desc="update_e2e_room_key",
+ )
+
+ async def add_e2e_room_keys(self, user_id, version, room_keys):
+ """Bulk add room keys to a given backup.
+
+ Args:
+ user_id (str): the user whose backup we're adding to
+ version (str): the version ID of the backup for the set of keys we're adding to
+ room_keys (iterable[(str, str, dict)]): the keys to add, in the form
+ (roomID, sessionID, keyData)
+ """
+
+ values = []
+ for (room_id, session_id, room_key) in room_keys:
+ values.append(
+ {
+ "user_id": user_id,
+ "version": version,
+ "room_id": room_id,
+ "session_id": session_id,
+ "first_message_index": room_key["first_message_index"],
+ "forwarded_count": room_key["forwarded_count"],
+ "is_verified": room_key["is_verified"],
+ "session_data": json_encoder.encode(room_key["session_data"]),
+ }
+ )
+ log_kv(
+ {
+ "message": "Set room key",
+ "room_id": room_id,
+ "session_id": session_id,
+ "room_key": room_key,
+ }
+ )
+
+ await self.db_pool.simple_insert_many(
+ table="e2e_room_keys", values=values, desc="add_e2e_room_keys"
+ )
+
+ @trace
+ async def get_e2e_room_keys(self, user_id, version, room_id=None, session_id=None):
+ """Bulk get the E2E room keys for a given backup, optionally filtered to a given
+ room, or a given session.
+
+ Args:
+ user_id (str): the user whose backup we're querying
+ version (str): the version ID of the backup for the set of keys we're querying
+ room_id (str): Optional. the ID of the room whose keys we're querying, if any.
+ If not specified, we return the keys for all the rooms in the backup.
+ session_id (str): Optional. the session whose room_key we're querying, if any.
+ If specified, we also require the room_id to be specified.
+ If not specified, we return all the keys in this version of
+ the backup (or for the specified room)
+
+ Returns:
+ A list of dicts giving the session_data and message metadata for
+ these room keys.
+ """
+
+ try:
+ version = int(version)
+ except ValueError:
+ return {"rooms": {}}
+
+ keyvalues = {"user_id": user_id, "version": version}
+ if room_id:
+ keyvalues["room_id"] = room_id
+ if session_id:
+ keyvalues["session_id"] = session_id
+
+ rows = await self.db_pool.simple_select_list(
+ table="e2e_room_keys",
+ keyvalues=keyvalues,
+ retcols=(
+ "user_id",
+ "room_id",
+ "session_id",
+ "first_message_index",
+ "forwarded_count",
+ "is_verified",
+ "session_data",
+ ),
+ desc="get_e2e_room_keys",
+ )
+
+ sessions = {"rooms": {}}
+ for row in rows:
+ 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 must be returned to the client as a boolean
+ "is_verified": bool(row["is_verified"]),
+ "session_data": db_to_json(row["session_data"]),
+ }
+
+ return sessions
+
+ def get_e2e_room_keys_multi(self, user_id, version, room_keys):
+ """Get multiple room keys at a time. The difference between this function and
+ get_e2e_room_keys is that this function can be used to retrieve
+ multiple specific keys at a time, whereas get_e2e_room_keys is used for
+ getting all the keys in a backup version, all the keys for a room, or a
+ specific key.
+
+ Args:
+ user_id (str): the user whose backup we're querying
+ version (str): the version ID of the backup we're querying about
+ room_keys (dict[str, dict[str, iterable[str]]]): a map from
+ room ID -> {"session": [session ids]} indicating the session IDs
+ that we want to query
+
+ Returns:
+ Deferred[dict[str, dict[str, dict]]]: a map of room IDs to session IDs to room key
+ """
+
+ return self.db_pool.runInteraction(
+ "get_e2e_room_keys_multi",
+ self._get_e2e_room_keys_multi_txn,
+ user_id,
+ version,
+ room_keys,
+ )
+
+ @staticmethod
+ def _get_e2e_room_keys_multi_txn(txn, user_id, version, room_keys):
+ if not room_keys:
+ return {}
+
+ where_clauses = []
+ params = [user_id, version]
+ for room_id, room in room_keys.items():
+ sessions = list(room["sessions"])
+ if not sessions:
+ continue
+ params.append(room_id)
+ params.extend(sessions)
+ where_clauses.append(
+ "(room_id = ? AND session_id IN (%s))"
+ % (",".join(["?" for _ in sessions]),)
+ )
+
+ # check if we're actually querying something
+ if not where_clauses:
+ return {}
+
+ sql = """
+ SELECT room_id, session_id, first_message_index, forwarded_count,
+ is_verified, session_data
+ FROM e2e_room_keys
+ WHERE user_id = ? AND version = ? AND (%s)
+ """ % (
+ " OR ".join(where_clauses)
+ )
+
+ txn.execute(sql, params)
+
+ ret = {}
+
+ for row in txn:
+ room_id = row[0]
+ session_id = row[1]
+ ret.setdefault(room_id, {})
+ ret[room_id][session_id] = {
+ "first_message_index": row[2],
+ "forwarded_count": row[3],
+ "is_verified": row[4],
+ "session_data": db_to_json(row[5]),
+ }
+
+ return ret
+
+ def count_e2e_room_keys(self, user_id, version):
+ """Get the number of keys in a backup version.
+
+ Args:
+ user_id (str): the user whose backup we're querying
+ version (str): the version ID of the backup we're querying about
+ """
+
+ return self.db_pool.simple_select_one_onecol(
+ table="e2e_room_keys",
+ keyvalues={"user_id": user_id, "version": version},
+ retcol="COUNT(*)",
+ desc="count_e2e_room_keys",
+ )
+
+ @trace
+ async def delete_e2e_room_keys(
+ self, user_id, version, room_id=None, session_id=None
+ ):
+ """Bulk delete the E2E room keys for a given backup, optionally filtered to a given
+ room or a given session.
+
+ Args:
+ user_id(str): the user whose backup we're deleting from
+ version(str): the version ID of the backup for the set of keys we're deleting
+ room_id(str): Optional. the ID of the room whose keys we're deleting, if any.
+ If not specified, we delete the keys for all the rooms in the backup.
+ session_id(str): Optional. the session whose room_key we're querying, if any.
+ If specified, we also require the room_id to be specified.
+ If not specified, we delete all the keys in this version of
+ the backup (or for the specified room)
+
+ Returns:
+ The deletion transaction
+ """
+
+ keyvalues = {"user_id": user_id, "version": int(version)}
+ if room_id:
+ keyvalues["room_id"] = room_id
+ if session_id:
+ keyvalues["session_id"] = session_id
+
+ await self.db_pool.simple_delete(
+ table="e2e_room_keys", keyvalues=keyvalues, desc="delete_e2e_room_keys"
+ )
+
+ @staticmethod
+ def _get_current_version(txn, user_id):
+ txn.execute(
+ "SELECT MAX(version) FROM e2e_room_keys_versions "
+ "WHERE user_id=? AND deleted=0",
+ (user_id,),
+ )
+ row = txn.fetchone()
+ if not row:
+ raise StoreError(404, "No current backup version")
+ return row[0]
+
+ def get_e2e_room_keys_version_info(self, user_id, version=None):
+ """Get info metadata about a version of our room_keys backup.
+
+ Args:
+ user_id(str): the user whose backup we're querying
+ version(str): Optional. the version ID of the backup we're querying about
+ If missing, we return the information about the current version.
+ Raises:
+ StoreError: with code 404 if there are no e2e_room_keys_versions present
+ Returns:
+ A deferred dict giving the info metadata for this backup version, with
+ fields including:
+ version(str)
+ algorithm(str)
+ auth_data(object): opaque dict supplied by the client
+ etag(int): tag of the keys in the backup
+ """
+
+ def _get_e2e_room_keys_version_info_txn(txn):
+ if version is None:
+ this_version = self._get_current_version(txn, user_id)
+ else:
+ try:
+ this_version = int(version)
+ except ValueError:
+ # Our versions are all ints so if we can't convert it to an integer,
+ # it isn't there.
+ raise StoreError(404, "No row found")
+
+ result = self.db_pool.simple_select_one_txn(
+ txn,
+ table="e2e_room_keys_versions",
+ keyvalues={"user_id": user_id, "version": this_version, "deleted": 0},
+ retcols=("version", "algorithm", "auth_data", "etag"),
+ )
+ result["auth_data"] = db_to_json(result["auth_data"])
+ result["version"] = str(result["version"])
+ if result["etag"] is None:
+ result["etag"] = 0
+ return result
+
+ return self.db_pool.runInteraction(
+ "get_e2e_room_keys_version_info", _get_e2e_room_keys_version_info_txn
+ )
+
+ @trace
+ def create_e2e_room_keys_version(self, user_id, info):
+ """Atomically creates a new version of this user's e2e_room_keys store
+ with the given version info.
+
+ Args:
+ user_id(str): the user whose backup we're creating a version
+ info(dict): the info about the backup version to be created
+
+ Returns:
+ A deferred string for the newly created version ID
+ """
+
+ def _create_e2e_room_keys_version_txn(txn):
+ txn.execute(
+ "SELECT MAX(version) FROM e2e_room_keys_versions WHERE user_id=?",
+ (user_id,),
+ )
+ current_version = txn.fetchone()[0]
+ if current_version is None:
+ current_version = "0"
+
+ new_version = str(int(current_version) + 1)
+
+ self.db_pool.simple_insert_txn(
+ txn,
+ table="e2e_room_keys_versions",
+ values={
+ "user_id": user_id,
+ "version": new_version,
+ "algorithm": info["algorithm"],
+ "auth_data": json_encoder.encode(info["auth_data"]),
+ },
+ )
+
+ return new_version
+
+ return self.db_pool.runInteraction(
+ "create_e2e_room_keys_version_txn", _create_e2e_room_keys_version_txn
+ )
+
+ @trace
+ def update_e2e_room_keys_version(
+ self, user_id, version, info=None, version_etag=None
+ ):
+ """Update a given backup version
+
+ Args:
+ user_id(str): the user whose backup version we're updating
+ version(str): the version ID of the backup version we're updating
+ info (dict): the new backup version info to store. If None, then
+ the backup version info is not updated
+ version_etag (Optional[int]): etag of the keys in the backup. If
+ None, then the etag is not updated
+ """
+ updatevalues = {}
+
+ if info is not None and "auth_data" in info:
+ updatevalues["auth_data"] = json_encoder.encode(info["auth_data"])
+ if version_etag is not None:
+ updatevalues["etag"] = version_etag
+
+ if updatevalues:
+ return self.db_pool.simple_update(
+ table="e2e_room_keys_versions",
+ keyvalues={"user_id": user_id, "version": version},
+ updatevalues=updatevalues,
+ desc="update_e2e_room_keys_version",
+ )
+
+ @trace
+ def delete_e2e_room_keys_version(self, user_id, version=None):
+ """Delete a given backup version of the user's room keys.
+ Doesn't delete their actual key data.
+
+ Args:
+ user_id(str): the user whose backup version we're deleting
+ version(str): Optional. the version ID of the backup version we're deleting
+ If missing, we delete the current backup version info.
+ Raises:
+ StoreError: with code 404 if there are no e2e_room_keys_versions present,
+ or if the version requested doesn't exist.
+ """
+
+ def _delete_e2e_room_keys_version_txn(txn):
+ if version is None:
+ this_version = self._get_current_version(txn, user_id)
+ if this_version is None:
+ raise StoreError(404, "No current backup version")
+ else:
+ this_version = version
+
+ self.db_pool.simple_delete_txn(
+ txn,
+ table="e2e_room_keys",
+ keyvalues={"user_id": user_id, "version": this_version},
+ )
+
+ return self.db_pool.simple_update_one_txn(
+ txn,
+ table="e2e_room_keys_versions",
+ keyvalues={"user_id": user_id, "version": this_version},
+ updatevalues={"deleted": 1},
+ )
+
+ return self.db_pool.runInteraction(
+ "delete_e2e_room_keys_version", _delete_e2e_room_keys_version_txn
+ )
|
