test_keys.py 4.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137
  1. # Copyright 2017 Vector Creations Ltd
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. import signedjson.key
  15. import signedjson.types
  16. import unpaddedbase64
  17. from synapse.storage.keys import FetchKeyResult
  18. import tests.unittest
  19. def decode_verify_key_base64(
  20. key_id: str, key_base64: str
  21. ) -> signedjson.types.VerifyKey:
  22. key_bytes = unpaddedbase64.decode_base64(key_base64)
  23. return signedjson.key.decode_verify_key_bytes(key_id, key_bytes)
  24. KEY_1 = decode_verify_key_base64(
  25. "ed25519:key1", "fP5l4JzpZPq/zdbBg5xx6lQGAAOM9/3w94cqiJ5jPrw"
  26. )
  27. KEY_2 = decode_verify_key_base64(
  28. "ed25519:key2", "Noi6WqcDj0QmPxCNQqgezwTlBKrfqehY1u2FyWP9uYw"
  29. )
  30. class KeyStoreTestCase(tests.unittest.HomeserverTestCase):
  31. def test_get_server_signature_keys(self) -> None:
  32. store = self.hs.get_datastores().main
  33. key_id_1 = "ed25519:key1"
  34. key_id_2 = "ed25519:KEY_ID_2"
  35. self.get_success(
  36. store.store_server_signature_keys(
  37. "from_server",
  38. 10,
  39. {
  40. ("server1", key_id_1): FetchKeyResult(KEY_1, 100),
  41. ("server1", key_id_2): FetchKeyResult(KEY_2, 200),
  42. },
  43. )
  44. )
  45. res = self.get_success(
  46. store.get_server_signature_keys(
  47. [
  48. ("server1", key_id_1),
  49. ("server1", key_id_2),
  50. ("server1", "ed25519:key3"),
  51. ]
  52. )
  53. )
  54. self.assertEqual(len(res.keys()), 3)
  55. res1 = res[("server1", key_id_1)]
  56. self.assertEqual(res1.verify_key, KEY_1)
  57. self.assertEqual(res1.verify_key.version, "key1")
  58. self.assertEqual(res1.valid_until_ts, 100)
  59. res2 = res[("server1", key_id_2)]
  60. self.assertEqual(res2.verify_key, KEY_2)
  61. # version comes from the ID it was stored with
  62. self.assertEqual(res2.verify_key.version, "KEY_ID_2")
  63. self.assertEqual(res2.valid_until_ts, 200)
  64. # non-existent result gives None
  65. self.assertIsNone(res[("server1", "ed25519:key3")])
  66. def test_cache(self) -> None:
  67. """Check that updates correctly invalidate the cache."""
  68. store = self.hs.get_datastores().main
  69. key_id_1 = "ed25519:key1"
  70. key_id_2 = "ed25519:key2"
  71. self.get_success(
  72. store.store_server_signature_keys(
  73. "from_server",
  74. 0,
  75. {
  76. ("srv1", key_id_1): FetchKeyResult(KEY_1, 100),
  77. ("srv1", key_id_2): FetchKeyResult(KEY_2, 200),
  78. },
  79. )
  80. )
  81. res = self.get_success(
  82. store.get_server_signature_keys([("srv1", key_id_1), ("srv1", key_id_2)])
  83. )
  84. self.assertEqual(len(res.keys()), 2)
  85. res1 = res[("srv1", key_id_1)]
  86. self.assertEqual(res1.verify_key, KEY_1)
  87. self.assertEqual(res1.valid_until_ts, 100)
  88. res2 = res[("srv1", key_id_2)]
  89. self.assertEqual(res2.verify_key, KEY_2)
  90. self.assertEqual(res2.valid_until_ts, 200)
  91. # we should be able to look up the same thing again without a db hit
  92. res = self.get_success(store.get_server_signature_keys([("srv1", key_id_1)]))
  93. self.assertEqual(len(res.keys()), 1)
  94. self.assertEqual(res[("srv1", key_id_1)].verify_key, KEY_1)
  95. new_key_2 = signedjson.key.get_verify_key(
  96. signedjson.key.generate_signing_key("key2")
  97. )
  98. d = store.store_server_signature_keys(
  99. "from_server", 10, {("srv1", key_id_2): FetchKeyResult(new_key_2, 300)}
  100. )
  101. self.get_success(d)
  102. res = self.get_success(
  103. store.get_server_signature_keys([("srv1", key_id_1), ("srv1", key_id_2)])
  104. )
  105. self.assertEqual(len(res.keys()), 2)
  106. res1 = res[("srv1", key_id_1)]
  107. self.assertEqual(res1.verify_key, KEY_1)
  108. self.assertEqual(res1.valid_until_ts, 100)
  109. res2 = res[("srv1", key_id_2)]
  110. self.assertEqual(res2.verify_key, new_key_2)
  111. self.assertEqual(res2.valid_until_ts, 300)