Notify when keys uploaded

element-hq · Oct 10, 2024 · db63bd4 · db63bd4
1 parent 3182f6b
commit db63bd4
Show file tree

Hide file tree

Showing 3 changed files with 142 additions and 64 deletions.
diff --git a/synapse/handlers/e2e_keys.py b/synapse/handlers/e2e_keys.py
@@ -837,6 +837,7 @@ async def upload_keys_for_user(
  await self._upload_one_time_keys_for_user(
  user_id, device_id, time_now, one_time_keys
  )
+ self._notifier.notify_one_time_keys_changed([user_id])
  else:
  log_kv(
  {"message": "Did not update one_time_keys", "reason": "no keys given"}
@@ -853,6 +854,7 @@ async def upload_keys_for_user(
  }
  )
  await self.store.set_e2e_fallback_keys(user_id, device_id, fallback_keys)
+ self._notifier.notify_one_time_keys_changed([user_id])
  elif fallback_keys:
  log_kv({"message": "Did not update fallback_keys", "reason": "not a dict"})
  else:

diff --git a/synapse/notifier.py b/synapse/notifier.py
@@ -936,6 +936,9 @@ def notify_one_time_keys_changed(
  for listener in listeners:
  listener.callback(current_token)
 
+ # FIXME: How can we poke the replication so that other workers also see the
+ # one-time key change
+
 
 @attr.s(auto_attribs=True)
 class ReplicationNotifier:

diff --git a/tests/rest/client/sliding_sync/test_extension_e2ee.py b/tests/rest/client/sliding_sync/test_extension_e2ee.py
@@ -33,8 +33,10 @@
 
 class E2eeBumpAction(enum.Enum):
  device_lists = enum.auto()
- one_time_keys = enum.auto()
- fallback_one_time_keys = enum.auto()
+ claim_one_time_key = enum.auto()
+ claim_fallback_key = enum.auto()
+ upload_one_time_key = enum.auto()
+ upload_fallback_key = enum.auto()
 
 
 # FIXME: This can be removed once we bump `SCHEMA_COMPAT_VERSION` and run the
@@ -158,16 +160,24 @@ def test_no_data_incremental_sync(self) -> None:
  @parameterized.expand(
  [
  (
- "bump_device_lists",
+ "bump_via_device_list_change",
  E2eeBumpAction.device_lists,
  ),
  (
- "bump_one_time_keys",
- E2eeBumpAction.one_time_keys,
+ "bump_via_claim_one_time_keys",
+ E2eeBumpAction.claim_one_time_key,
  ),
  (
- "bump_fallback_one_time_keys",
- E2eeBumpAction.fallback_one_time_keys,
+ "bump_via_claim_fallback_keys",
+ E2eeBumpAction.claim_fallback_key,
+ ),
+ (
+ "bump_via_upload_one_time_keys",
+ E2eeBumpAction.upload_one_time_key,
+ ),
+ (
+ "bump_via_upload_fallback_keys",
+ E2eeBumpAction.upload_fallback_key,
  ),
  ]
  )
@@ -192,6 +202,53 @@ def test_wait_for_new_data(
  self.helper.join(room_id, user1_id, tok=user1_tok)
  self.helper.join(room_id, user3_id, tok=user3_tok)
 
+ # Upload one time keys for the user/device
+ upload_keys_response = self.get_success(
+ self.e2e_keys_handler.upload_keys_for_user(
+ user1_id,
+ test_device_id,
+ {
+ "one_time_keys": {
+ "alg1:k1": "key1",
+ "alg2:k2": {"key": "key2", "signatures": {"k1": "sig1"}},
+ "alg2:k3": {"key": "key3"},
+ },
+ # Upload fallback keys for the user/device
+ "fallback_keys": {
+ "fallback_alg1:k1": "fallback_key1",
+ "fallback_alg2:k2": "fallback_key2",
+ },
+ },
+ )
+ )
+ self.assertDictEqual(
+ upload_keys_response,
+ {
+ "one_time_key_counts": {
+ "alg1": 1,
+ "alg2": 2,
+ # Note that "signed_curve25519" is always returned in key count responses
+ # regardless of whether we uploaded any keys for it. This is necessary until
+ # https://github.com/matrix-org/matrix-doc/issues/3298 is fixed.
+ #
+ # Also related:
+ # https://github.com/element-hq/element-android/issues/3725 and
+ # https://github.com/matrix-org/synapse/issues/10456
+ "signed_curve25519": 0,
+ }
+ },
+ )
+ # We should now have an unused fallback_alg1 and fallback_alg2 key
+ fallback_res = self.get_success(
+ self.store.get_e2e_unused_fallback_key_types(user1_id, test_device_id)
+ )
+ self.assertIncludes(
+ set(fallback_res),
+ {"fallback_alg1", "fallback_alg2"},
+ exact=True,
+ message=str(fallback_res),
+ )
+
  sync_body = {
  "lists": {},
  "extensions": {
@@ -228,36 +285,7 @@ def test_wait_for_new_data(
  self.assertEqual(
  device_update_channel.code, 200, device_update_channel.json_body
  )
- elif bump_action == E2eeBumpAction.one_time_keys:
- # Upload one time keys for the user/device
- keys: JsonDict = {
- "alg1:k1": "key1",
- "alg2:k2": {"key": "key2", "signatures": {"k1": "sig1"}},
- "alg2:k3": {"key": "key3"},
- }
- upload_keys_response = self.get_success(
- self.e2e_keys_handler.upload_keys_for_user(
- user1_id, test_device_id, {"one_time_keys": keys}
- )
- )
- self.assertDictEqual(
- upload_keys_response,
- {
- "one_time_key_counts": {
- "alg1": 1,
- "alg2": 2,
- # Note that "signed_curve25519" is always returned in key count responses
- # regardless of whether we uploaded any keys for it. This is necessary until
- # https://github.com/matrix-org/matrix-doc/issues/3298 is fixed.
- #
- # Also related:
- # https://github.com/element-hq/element-android/issues/3725 and
- # https://github.com/matrix-org/synapse/issues/10456
- "signed_curve25519": 0,
- }
- },
- )
-
+ elif bump_action == E2eeBumpAction.claim_one_time_key:
  # Claim one of those new keys
  self.get_success(
  self.e2e_keys_handler.claim_local_one_time_keys(
@@ -270,31 +298,7 @@ def test_wait_for_new_data(
  # * no change in OTK count since the provided since token
  # * the server has zero OTKs left for this device
  # Spec issue: https://github.com/matrix-org/matrix-doc/issues/3298
- elif bump_action == E2eeBumpAction.fallback_one_time_keys:
- # Upload a fallback key for the user/device
- self.get_success(
- self.e2e_keys_handler.upload_keys_for_user(
- user1_id,
- test_device_id,
- {
- "fallback_keys": {
- "fallback_alg1:k1": "fallback_key1",
- "fallback_alg2:k2": "fallback_key2",
- }
- },
- )
- )
- # We should now have an unused alg1 and alg2 key
- fallback_res = self.get_success(
- self.store.get_e2e_unused_fallback_key_types(user1_id, test_device_id)
- )
- self.assertIncludes(
- set(fallback_res),
- {"fallback_alg1", "fallback_alg2"},
- exact=True,
- message=str(fallback_res),
- )
-
+ elif bump_action == E2eeBumpAction.claim_fallback_key:
  # Claim one of those fallback keys
  self.get_success(
  self.e2e_keys_handler.claim_local_one_time_keys(
@@ -307,6 +311,32 @@ def test_wait_for_new_data(
  # * no change in OTK count since the provided since token
  # * the server has zero OTKs left for this device
  # Spec issue: https://github.com/matrix-org/matrix-doc/issues/3298
+ elif bump_action == E2eeBumpAction.upload_one_time_key:
+ # Upload a new one-time key
+ upload_keys_response = self.get_success(
+ self.e2e_keys_handler.upload_keys_for_user(
+ user1_id,
+ test_device_id,
+ {
+ "one_time_keys": {
+ "new_alg3:k4": {"key": "key4"},
+ },
+ },
+ )
+ )
+ elif bump_action == E2eeBumpAction.upload_fallback_key:
+ # Upload a new fallback one-time key
+ upload_keys_response = self.get_success(
+ self.e2e_keys_handler.upload_keys_for_user(
+ user1_id,
+ test_device_id,
+ {
+ "fallback_keys": {
+ "fallback_new_alg3:k3": "fallback_key3",
+ },
+ },
+ )
+ )
  else:
  assert_never(bump_action)
 
@@ -330,7 +360,7 @@ def test_wait_for_new_data(
  .get("left"),
  [],
  )
- elif bump_action == E2eeBumpAction.one_time_keys:
+ elif bump_action == E2eeBumpAction.claim_one_time_key:
  # We should see the one-time key count change
  self.assertEqual(
  channel.json_body["extensions"]["e2ee"].get(
@@ -350,14 +380,15 @@ def test_wait_for_new_data(
  "signed_curve25519": 0,
  },
  )
- elif bump_action == E2eeBumpAction.fallback_one_time_keys:
+ elif bump_action == E2eeBumpAction.claim_fallback_key:
  # Check for the unused fallback key types
  self.assertIncludes(
  set(
  channel.json_body["extensions"]["e2ee"].get(
  "device_unused_fallback_key_types", []
  )
  ),
+ # `fallback_alg1` was claimed so it doesn't appear here anymore
  {"fallback_alg2"},
  exact=True,
  message=str(
@@ -366,6 +397,48 @@ def test_wait_for_new_data(
  )
  ),
  )
+ elif bump_action == E2eeBumpAction.upload_one_time_key:
+ # We should see the one-time key count change
+ self.assertEqual(
+ channel.json_body["extensions"]["e2ee"].get(
+ "device_one_time_keys_count"
+ ),
+ {
+ "alg1": 1,
+ "alg2": 2,
+ # One-time key we just uploaded
+ "new_alg3": 1,
+ # Note that "signed_curve25519" is always returned in key count responses
+ # regardless of whether we uploaded any keys for it. This is necessary until
+ # https://github.com/matrix-org/matrix-doc/issues/3298 is fixed.
+ #
+ # Also related:
+ # https://github.com/element-hq/element-android/issues/3725 and
+ # https://github.com/matrix-org/synapse/issues/10456
+ "signed_curve25519": 0,
+ },
+ )
+ elif bump_action == E2eeBumpAction.upload_fallback_key:
+ # Check for the unused fallback key types
+ self.assertIncludes(
+ set(
+ channel.json_body["extensions"]["e2ee"].get(
+ "device_unused_fallback_key_types", []
+ )
+ ),
+ {
+ "fallback_alg1",
+ "fallback_alg2",
+ # Fallback key we just uploaded
+ "fallback_new_alg3",
+ },
+ exact=True,
+ message=str(
+ channel.json_body["extensions"]["e2ee"].get(
+ "device_unused_fallback_key_types",
+ )
+ ),
+ )
  else:
  assert_never(bump_action)