summary refs log tree commit diff
path: root/synapse/storage/databases/main/keys.py
blob: 889c578b9c97fd5c28654c5969ea55145bd1d589 (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
# Copyright 2014-2016 OpenMarket Ltd
# Copyright 2019 New Vector Ltd.
#
# 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.

import itertools
import json
import logging
from typing import Dict, Iterable, Mapping, Optional, Tuple

from canonicaljson import encode_canonical_json
from signedjson.key import decode_verify_key_bytes
from unpaddedbase64 import decode_base64

from synapse.storage.database import LoggingTransaction
from synapse.storage.databases.main.cache import CacheInvalidationWorkerStore
from synapse.storage.keys import FetchKeyResult, FetchKeyResultForRemote
from synapse.storage.types import Cursor
from synapse.types import JsonDict
from synapse.util.caches.descriptors import cached, cachedList
from synapse.util.iterutils import batch_iter

logger = logging.getLogger(__name__)


db_binary_type = memoryview


class KeyStore(CacheInvalidationWorkerStore):
    """Persistence for signature verification keys"""

    async def store_server_keys_response(
        self,
        server_name: str,
        from_server: str,
        ts_added_ms: int,
        verify_keys: Dict[str, FetchKeyResult],
        response_json: JsonDict,
    ) -> None:
        """Stores the keys for the given server that we got from `from_server`.

        Args:
            server_name: The owner of the keys
            from_server: Which server we got the keys from
            ts_added_ms: When we're adding the keys
            verify_keys: The decoded keys
            response_json: The full *signed* response JSON that contains the keys.
        """

        key_json_bytes = encode_canonical_json(response_json)

        def store_server_keys_response_txn(txn: LoggingTransaction) -> None:
            self.db_pool.simple_upsert_many_txn(
                txn,
                table="server_signature_keys",
                key_names=("server_name", "key_id"),
                key_values=[(server_name, key_id) for key_id in verify_keys],
                value_names=(
                    "from_server",
                    "ts_added_ms",
                    "ts_valid_until_ms",
                    "verify_key",
                ),
                value_values=[
                    (
                        from_server,
                        ts_added_ms,
                        fetch_result.valid_until_ts,
                        db_binary_type(fetch_result.verify_key.encode()),
                    )
                    for fetch_result in verify_keys.values()
                ],
            )

            self.db_pool.simple_upsert_many_txn(
                txn,
                table="server_keys_json",
                key_names=("server_name", "key_id", "from_server"),
                key_values=[
                    (server_name, key_id, from_server) for key_id in verify_keys
                ],
                value_names=(
                    "ts_added_ms",
                    "ts_valid_until_ms",
                    "key_json",
                ),
                value_values=[
                    (
                        ts_added_ms,
                        fetch_result.valid_until_ts,
                        db_binary_type(key_json_bytes),
                    )
                    for fetch_result in verify_keys.values()
                ],
            )

            # invalidate takes a tuple corresponding to the params of
            # _get_server_keys_json. _get_server_keys_json only takes one
            # param, which is itself the 2-tuple (server_name, key_id).
            for key_id in verify_keys:
                self._invalidate_cache_and_stream(
                    txn, self._get_server_keys_json, ((server_name, key_id),)
                )
                self._invalidate_cache_and_stream(
                    txn, self.get_server_key_json_for_remote, (server_name, key_id)
                )

        await self.db_pool.runInteraction(
            "store_server_keys_response", store_server_keys_response_txn
        )

    @cached()
    def _get_server_keys_json(
        self, server_name_and_key_id: Tuple[str, str]
    ) -> FetchKeyResult:
        raise NotImplementedError()

    @cachedList(
        cached_method_name="_get_server_keys_json", list_name="server_name_and_key_ids"
    )
    async def get_server_keys_json(
        self, server_name_and_key_ids: Iterable[Tuple[str, str]]
    ) -> Mapping[Tuple[str, str], FetchKeyResult]:
        """
        Args:
            server_name_and_key_ids:
                iterable of (server_name, key-id) tuples to fetch keys for

        Returns:
            A map from (server_name, key_id) -> FetchKeyResult, or None if the
            key is unknown
        """
        keys = {}

        def _get_keys(txn: Cursor, batch: Tuple[Tuple[str, str], ...]) -> None:
            """Processes a batch of keys to fetch, and adds the result to `keys`."""

            # batch_iter always returns tuples so it's safe to do len(batch)
            where_clause = " OR (server_name=? AND key_id=?)" * len(batch)

            # `server_keys_json` can have multiple entries per server (one per
            # remote server we fetched from, if using perspectives). Order by
            # `ts_added_ms` so the most recently fetched one always wins.
            sql = f"""
                SELECT server_name, key_id, key_json, ts_valid_until_ms
                FROM server_keys_json WHERE 1=0
                {where_clause}
                ORDER BY ts_added_ms
            """

            txn.execute(sql, tuple(itertools.chain.from_iterable(batch)))

            for server_name, key_id, key_json_bytes, ts_valid_until_ms in txn:
                if ts_valid_until_ms is None:
                    # Old keys may be stored with a ts_valid_until_ms of null,
                    # in which case we treat this as if it was set to `0`, i.e.
                    # it won't match key requests that define a minimum
                    # `ts_valid_until_ms`.
                    ts_valid_until_ms = 0

                # The entire signed JSON response is stored in server_keys_json,
                # fetch out the bits needed.
                key_json = json.loads(bytes(key_json_bytes))
                key_base64 = key_json["verify_keys"][key_id]["key"]

                keys[(server_name, key_id)] = FetchKeyResult(
                    verify_key=decode_verify_key_bytes(
                        key_id, decode_base64(key_base64)
                    ),
                    valid_until_ts=ts_valid_until_ms,
                )

        def _txn(txn: Cursor) -> Dict[Tuple[str, str], FetchKeyResult]:
            for batch in batch_iter(server_name_and_key_ids, 50):
                _get_keys(txn, batch)
            return keys

        return await self.db_pool.runInteraction("get_server_keys_json", _txn)

    @cached()
    def get_server_key_json_for_remote(
        self,
        server_name: str,
        key_id: str,
    ) -> Optional[FetchKeyResultForRemote]:
        raise NotImplementedError()

    @cachedList(
        cached_method_name="get_server_key_json_for_remote", list_name="key_ids"
    )
    async def get_server_keys_json_for_remote(
        self, server_name: str, key_ids: Iterable[str]
    ) -> Mapping[str, Optional[FetchKeyResultForRemote]]:
        """Fetch the cached keys for the given server/key IDs.

        If we have multiple entries for a given key ID, returns the most recent.
        """
        rows = await self.db_pool.simple_select_many_batch(
            table="server_keys_json",
            column="key_id",
            iterable=key_ids,
            keyvalues={"server_name": server_name},
            retcols=(
                "key_id",
                "from_server",
                "ts_added_ms",
                "ts_valid_until_ms",
                "key_json",
            ),
            desc="get_server_keys_json_for_remote",
        )

        if not rows:
            return {}

        # We sort the rows so that the most recently added entry is picked up.
        rows.sort(key=lambda r: r["ts_added_ms"])

        return {
            row["key_id"]: FetchKeyResultForRemote(
                # Cast to bytes since postgresql returns a memoryview.
                key_json=bytes(row["key_json"]),
                valid_until_ts=row["ts_valid_until_ms"],
                added_ts=row["ts_added_ms"],
            )
            for row in rows
        }

    async def get_all_server_keys_json_for_remote(
        self,
        server_name: str,
    ) -> Dict[str, FetchKeyResultForRemote]:
        """Fetch the cached keys for the given server.

        If we have multiple entries for a given key ID, returns the most recent.
        """
        rows = await self.db_pool.simple_select_list(
            table="server_keys_json",
            keyvalues={"server_name": server_name},
            retcols=(
                "key_id",
                "from_server",
                "ts_added_ms",
                "ts_valid_until_ms",
                "key_json",
            ),
            desc="get_server_keys_json_for_remote",
        )

        if not rows:
            return {}

        rows.sort(key=lambda r: r["ts_added_ms"])

        return {
            row["key_id"]: FetchKeyResultForRemote(
                # Cast to bytes since postgresql returns a memoryview.
                key_json=bytes(row["key_json"]),
                valid_until_ts=row["ts_valid_until_ms"],
                added_ts=row["ts_added_ms"],
            )
            for row in rows
        }