summary refs log tree commit diff
path: root/tests/test_utils/oidc.py
blob: 9b204b196b8725901171e39d6291b12af8a04874 (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
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
#
# This file is licensed under the Affero General Public License (AGPL) version 3.
#
# Copyright (C) 2023 New Vector, Ltd
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# See the GNU Affero General Public License for more details:
# <https://www.gnu.org/licenses/agpl-3.0.html>.
#
# Originally licensed under the Apache License, Version 2.0:
# <http://www.apache.org/licenses/LICENSE-2.0>.
#
# [This file includes modifications made by New Vector Limited]
#
#


import json
from typing import Any, ContextManager, Dict, List, Optional, Tuple
from unittest.mock import Mock, patch
from urllib.parse import parse_qs

import attr

from twisted.web.http_headers import Headers
from twisted.web.iweb import IResponse

from synapse.server import HomeServer
from synapse.util import Clock
from synapse.util.stringutils import random_string

from tests.test_utils import FakeResponse


@attr.s(slots=True, frozen=True, auto_attribs=True)
class FakeAuthorizationGrant:
    userinfo: dict
    client_id: str
    redirect_uri: str
    scope: str
    nonce: Optional[str]
    sid: Optional[str]


class FakeOidcServer:
    """A fake OpenID Connect Provider."""

    # All methods here are mocks, so we can track when they are called, and override
    # their values
    request: Mock
    get_jwks_handler: Mock
    get_metadata_handler: Mock
    get_userinfo_handler: Mock
    post_token_handler: Mock

    sid_counter: int = 0

    def __init__(self, clock: Clock, issuer: str):
        from authlib.jose import ECKey, KeySet

        self._clock = clock
        self.issuer = issuer

        self.request = Mock(side_effect=self._request)
        self.get_jwks_handler = Mock(side_effect=self._get_jwks_handler)
        self.get_metadata_handler = Mock(side_effect=self._get_metadata_handler)
        self.get_userinfo_handler = Mock(side_effect=self._get_userinfo_handler)
        self.post_token_handler = Mock(side_effect=self._post_token_handler)

        # A code -> grant mapping
        self._authorization_grants: Dict[str, FakeAuthorizationGrant] = {}
        # An access token -> grant mapping
        self._sessions: Dict[str, FakeAuthorizationGrant] = {}

        # We generate here an ECDSA key with the P-256 curve (ES256 algorithm) used for
        # signing JWTs. ECDSA keys are really quick to generate compared to RSA.
        self._key = ECKey.generate_key(crv="P-256", is_private=True)
        self._jwks = KeySet([ECKey.import_key(self._key.as_pem(is_private=False))])

        self._id_token_overrides: Dict[str, Any] = {}

    def reset_mocks(self) -> None:
        self.request.reset_mock()
        self.get_jwks_handler.reset_mock()
        self.get_metadata_handler.reset_mock()
        self.get_userinfo_handler.reset_mock()
        self.post_token_handler.reset_mock()

    def patch_homeserver(self, hs: HomeServer) -> ContextManager[Mock]:
        """Patch the ``HomeServer`` HTTP client to handle requests through the ``FakeOidcServer``.

        This patch should be used whenever the HS is expected to perform request to the
        OIDC provider, e.g.::

            fake_oidc_server = self.helper.fake_oidc_server()
            with fake_oidc_server.patch_homeserver(hs):
                self.make_request("GET", "/_matrix/client/r0/login/sso/redirect")
        """
        return patch.object(hs.get_proxied_http_client(), "request", self.request)

    @property
    def authorization_endpoint(self) -> str:
        return self.issuer + "authorize"

    @property
    def token_endpoint(self) -> str:
        return self.issuer + "token"

    @property
    def userinfo_endpoint(self) -> str:
        return self.issuer + "userinfo"

    @property
    def metadata_endpoint(self) -> str:
        return self.issuer + ".well-known/openid-configuration"

    @property
    def jwks_uri(self) -> str:
        return self.issuer + "jwks"

    def get_metadata(self) -> dict:
        return {
            "issuer": self.issuer,
            "authorization_endpoint": self.authorization_endpoint,
            "token_endpoint": self.token_endpoint,
            "jwks_uri": self.jwks_uri,
            "userinfo_endpoint": self.userinfo_endpoint,
            "response_types_supported": ["code"],
            "subject_types_supported": ["public"],
            "id_token_signing_alg_values_supported": ["ES256"],
        }

    def get_jwks(self) -> dict:
        return self._jwks.as_dict()

    def get_userinfo(self, access_token: str) -> Optional[dict]:
        """Given an access token, get the userinfo of the associated session."""
        session = self._sessions.get(access_token, None)
        if session is None:
            return None
        return session.userinfo

    def _sign(self, payload: dict) -> str:
        from authlib.jose import JsonWebSignature

        jws = JsonWebSignature()
        kid = self.get_jwks()["keys"][0]["kid"]
        protected = {"alg": "ES256", "kid": kid}
        json_payload = json.dumps(payload)
        return jws.serialize_compact(protected, json_payload, self._key).decode("utf-8")

    def generate_id_token(self, grant: FakeAuthorizationGrant) -> str:
        now = int(self._clock.time())
        id_token = {
            **grant.userinfo,
            "iss": self.issuer,
            "aud": grant.client_id,
            "iat": now,
            "nbf": now,
            "exp": now + 600,
        }

        if grant.nonce is not None:
            id_token["nonce"] = grant.nonce

        if grant.sid is not None:
            id_token["sid"] = grant.sid

        id_token.update(self._id_token_overrides)

        return self._sign(id_token)

    def generate_logout_token(self, grant: FakeAuthorizationGrant) -> str:
        now = int(self._clock.time())
        logout_token = {
            "iss": self.issuer,
            "aud": grant.client_id,
            "iat": now,
            "jti": random_string(10),
            "events": {
                "http://schemas.openid.net/event/backchannel-logout": {},
            },
        }

        if grant.sid is not None:
            logout_token["sid"] = grant.sid

        if "sub" in grant.userinfo:
            logout_token["sub"] = grant.userinfo["sub"]

        return self._sign(logout_token)

    def id_token_override(self, overrides: dict) -> ContextManager[dict]:
        """Temporarily patch the ID token generated by the token endpoint."""
        return patch.object(self, "_id_token_overrides", overrides)

    def start_authorization(
        self,
        client_id: str,
        scope: str,
        redirect_uri: str,
        userinfo: dict,
        nonce: Optional[str] = None,
        with_sid: bool = False,
    ) -> Tuple[str, FakeAuthorizationGrant]:
        """Start an authorization request, and get back the code to use on the authorization endpoint."""
        code = random_string(10)
        sid = None
        if with_sid:
            sid = str(self.sid_counter)
            self.sid_counter += 1

        grant = FakeAuthorizationGrant(
            userinfo=userinfo,
            scope=scope,
            redirect_uri=redirect_uri,
            nonce=nonce,
            client_id=client_id,
            sid=sid,
        )
        self._authorization_grants[code] = grant

        return code, grant

    def exchange_code(self, code: str) -> Optional[Dict[str, Any]]:
        grant = self._authorization_grants.pop(code, None)
        if grant is None:
            return None

        access_token = random_string(10)
        self._sessions[access_token] = grant

        token = {
            "token_type": "Bearer",
            "access_token": access_token,
            "expires_in": 3600,
            "scope": grant.scope,
        }

        if "openid" in grant.scope:
            token["id_token"] = self.generate_id_token(grant)

        return dict(token)

    def buggy_endpoint(
        self,
        *,
        jwks: bool = False,
        metadata: bool = False,
        token: bool = False,
        userinfo: bool = False,
    ) -> ContextManager[Dict[str, Mock]]:
        """A context which makes a set of endpoints return a 500 error.

        Args:
            jwks: If True, makes the JWKS endpoint return a 500 error.
            metadata: If True, makes the OIDC Discovery endpoint return a 500 error.
            token: If True, makes the token endpoint return a 500 error.
            userinfo: If True, makes the userinfo endpoint return a 500 error.
        """
        buggy = FakeResponse(code=500, body=b"Internal server error")

        patches = {}
        if jwks:
            patches["get_jwks_handler"] = Mock(return_value=buggy)
        if metadata:
            patches["get_metadata_handler"] = Mock(return_value=buggy)
        if token:
            patches["post_token_handler"] = Mock(return_value=buggy)
        if userinfo:
            patches["get_userinfo_handler"] = Mock(return_value=buggy)

        return patch.multiple(self, **patches)

    async def _request(
        self,
        method: str,
        uri: str,
        data: Optional[bytes] = None,
        headers: Optional[Headers] = None,
    ) -> IResponse:
        """The override of the SimpleHttpClient#request() method"""
        access_token: Optional[str] = None

        if headers is None:
            headers = Headers()

        # Try to find the access token in the headers if any
        auth_headers = headers.getRawHeaders(b"Authorization")
        if auth_headers:
            parts = auth_headers[0].split(b" ")
            if parts[0] == b"Bearer" and len(parts) == 2:
                access_token = parts[1].decode("ascii")

        if method == "POST":
            # If the method is POST, assume it has an url-encoded body
            if data is None or headers.getRawHeaders(b"Content-Type") != [
                b"application/x-www-form-urlencoded"
            ]:
                return FakeResponse.json(code=400, payload={"error": "invalid_request"})

            params = parse_qs(data.decode("utf-8"))

            if uri == self.token_endpoint:
                # Even though this endpoint should be protected, this does not check
                # for client authentication. We're not checking it for simplicity,
                # and because client authentication is tested in other standalone tests.
                return self.post_token_handler(params)

        elif method == "GET":
            if uri == self.jwks_uri:
                return self.get_jwks_handler()
            elif uri == self.metadata_endpoint:
                return self.get_metadata_handler()
            elif uri == self.userinfo_endpoint:
                return self.get_userinfo_handler(access_token=access_token)

        return FakeResponse(code=404, body=b"404 not found")

    # Request handlers
    def _get_jwks_handler(self) -> IResponse:
        """Handles requests to the JWKS URI."""
        return FakeResponse.json(payload=self.get_jwks())

    def _get_metadata_handler(self) -> IResponse:
        """Handles requests to the OIDC well-known document."""
        return FakeResponse.json(payload=self.get_metadata())

    def _get_userinfo_handler(self, access_token: Optional[str]) -> IResponse:
        """Handles requests to the userinfo endpoint."""
        if access_token is None:
            return FakeResponse(code=401)
        user_info = self.get_userinfo(access_token)
        if user_info is None:
            return FakeResponse(code=401)

        return FakeResponse.json(payload=user_info)

    def _post_token_handler(self, params: Dict[str, List[str]]) -> IResponse:
        """Handles requests to the token endpoint."""
        code = params.get("code", [])

        if len(code) != 1:
            return FakeResponse.json(code=400, payload={"error": "invalid_request"})

        grant = self.exchange_code(code=code[0])
        if grant is None:
            return FakeResponse.json(code=400, payload={"error": "invalid_grant"})

        return FakeResponse.json(payload=grant)