apply PR review suggestions

2 files changed, 11 insertions, 10 deletions

diff --git a/synapse/storage/devices.py b/synapse/storage/devices.py
index da23a350a1..6a5572e001 100644
--- a/synapse/storage/devices.py
+++ b/synapse/storage/devices.py
@@ -317,7 +317,7 @@ class DeviceWorkerStore(SQLBaseStore):
                 user_ids,
                 stream_id,
             )
-        defer.returnValue(stream_id)
+        return stream_id
 
     def _add_user_signature_change_txn(self, txn, from_user_id, user_ids, stream_id):
         txn.call_after(
@@ -491,9 +491,9 @@ class DeviceWorkerStore(SQLBaseStore):
             rows = yield self._execute(
                 "get_users_whose_signatures_changed", None, sql, user_id, from_key
             )
-            defer.returnValue(set(user for row in rows for user in json.loads(row[0])))
+            return set(user for row in rows for user in json.loads(row[0]))
         else:
-            defer.returnValue(set())
+            return set()
 
     def get_all_device_list_changes_for_remotes(self, from_key, to_key):
         """Return a list of `(stream_id, user_id, destination)` which is the
diff --git a/synapse/storage/end_to_end_keys.py b/synapse/storage/end_to_end_keys.py
index de2d1bbb9f..b218b7b2e8 100644
--- a/synapse/storage/end_to_end_keys.py
+++ b/synapse/storage/end_to_end_keys.py
@@ -14,8 +14,6 @@
 # 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.
-import time
-
 from six import iteritems
 
 from canonicaljson import encode_canonical_json, json
@@ -284,7 +282,7 @@ class EndToEndKeyStore(EndToEndKeyWorkerStore, SQLBaseStore):
             "delete_e2e_keys_by_device", delete_e2e_keys_by_device_txn
         )
 
-    def _set_e2e_cross_signing_key_txn(self, txn, user_id, key_type, key):
+    def _set_e2e_cross_signing_key_txn(self, txn, user_id, key_type, key, added_ts):
         """Set a user's cross-signing key.
 
         Args:
@@ -294,6 +292,7 @@ class EndToEndKeyStore(EndToEndKeyWorkerStore, SQLBaseStore):
                 for a master key, 'self_signing' for a self-signing key, or
                 'user_signing' for a user-signing key
             key (dict): the key data
+            added_ts (int): the timestamp for when the key was added
         """
         # the cross-signing keys need to occupy the same namespace as devices,
         # since signatures are identified by device ID.  So add an entry to the
@@ -334,18 +333,19 @@ class EndToEndKeyStore(EndToEndKeyWorkerStore, SQLBaseStore):
                 "user_id": user_id,
                 "keytype": key_type,
                 "keydata": json.dumps(key),
-                "added_ts": time.time() * 1000,
+                "added_ts": added_ts,
             },
             desc="store_master_key",
         )
 
-    def set_e2e_cross_signing_key(self, user_id, key_type, key):
+    def set_e2e_cross_signing_key(self, user_id, key_type, key, added_ts):
         """Set a user's cross-signing key.
 
         Args:
             user_id (str): the user to set the user-signing key for
             key_type (str): the type of cross-signing key to set
             key (dict): the key data
+            added_ts (int): the timestamp for when the key was added
         """
         return self.runInteraction(
             "add_e2e_cross_signing_key",
@@ -353,6 +353,7 @@ class EndToEndKeyStore(EndToEndKeyWorkerStore, SQLBaseStore):
             user_id,
             key_type,
             key,
+            added_ts,
         )
 
     def _get_e2e_cross_signing_key_txn(self, txn, user_id, key_type, from_user_id=None):
@@ -368,7 +369,7 @@ class EndToEndKeyStore(EndToEndKeyWorkerStore, SQLBaseStore):
                 the key will be included in the result
 
         Returns:
-            dict of the key data
+            dict of the key data or None if not found
         """
         sql = (
             "SELECT keydata "
@@ -412,7 +413,7 @@ class EndToEndKeyStore(EndToEndKeyWorkerStore, SQLBaseStore):
                 the self-signing key will be included in the result
 
         Returns:
-            dict of the key data
+            dict of the key data or None if not found
         """
         return self.runInteraction(
             "get_e2e_cross_signing_key",