Skip to main content

matrix_sdk/test_utils/mocks/
oauth.rs

1// Copyright 2024 The Matrix.org Foundation C.I.C.
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
15//! Helpers to mock an OAuth 2.0 server for the purpose of integration tests.
16
17use std::time::Duration;
18
19use ruma::{
20    api::client::discovery::get_authorization_server_metadata::v1::AuthorizationServerMetadata,
21    serde::Raw,
22};
23use serde_json::json;
24use url::Url;
25use wiremock::{
26    Mock, MockBuilder, ResponseTemplate,
27    matchers::{method, path_regex},
28};
29
30use super::{MatrixMock, MatrixMockServer, MockEndpoint};
31
32/// A [`wiremock`] [`MockServer`] along with useful methods to help mocking
33/// OAuth 2.0 API endpoints easily.
34///
35/// It implements mock endpoints, limiting the shared code as much as possible,
36/// so the mocks are still flexible to use as scoped/unscoped mounts, named, and
37/// so on.
38///
39/// It works like this:
40///
41/// * start by saying which endpoint you'd like to mock, e.g.
42///   [`Self::mock_server_metadata()`]. This returns a specialized
43///   [`MockEndpoint`] data structure, with its own impl. For this example, it's
44///   `MockEndpoint<ServerMetadataEndpoint>`.
45/// * configure the response on the endpoint-specific mock data structure. For
46///   instance, if you want the sending to result in a transient failure, call
47///   [`MockEndpoint::error500`]; if you want it to succeed and return the
48///   metadata, call [`MockEndpoint::ok()`]. It's still possible to call
49///   [`MockEndpoint::respond_with()`], as we do with wiremock MockBuilder, for
50///   maximum flexibility when the helpers aren't sufficient.
51/// * once the endpoint's response is configured, for any mock builder, you get
52///   a [`MatrixMock`]; this is a plain [`wiremock::Mock`] with the server
53///   curried, so one doesn't have to pass it around when calling
54///   [`MatrixMock::mount()`] or [`MatrixMock::mount_as_scoped()`]. As such, it
55///   mostly defers its implementations to [`wiremock::Mock`] under the hood.
56///
57/// [`MockServer`]: wiremock::MockServer
58pub struct OAuthMockServer<'a> {
59    server: &'a MatrixMockServer,
60}
61
62impl<'a> OAuthMockServer<'a> {
63    pub(super) fn new(server: &'a MatrixMockServer) -> Self {
64        Self { server }
65    }
66
67    /// Mock the given endpoint.
68    fn mock_endpoint<T>(&self, mock: MockBuilder, endpoint: T) -> MockEndpoint<'a, T> {
69        self.server.mock_endpoint(mock, endpoint)
70    }
71
72    /// Get the mock OAuth 2.0 server metadata.
73    pub fn server_metadata(&self) -> AuthorizationServerMetadata {
74        MockServerMetadataBuilder::new(&self.server.uri())
75            .build()
76            .deserialize()
77            .expect("mock OAuth 2.0 server metadata should deserialize successfully")
78    }
79}
80
81// Specific mount endpoints.
82impl OAuthMockServer<'_> {
83    /// Creates a prebuilt mock for the Matrix endpoint used to query the
84    /// authorization server's metadata.
85    ///
86    /// Contrary to all the other endpoints of [`OAuthMockServer`], this is an
87    /// endpoint from the Matrix API, but it is only used in the context of the
88    /// OAuth 2.0 API, which is why it is mocked here rather than on
89    /// [`MatrixMockServer`].
90    ///
91    /// [`MatrixMockServer`]: super::MatrixMockServer
92    pub fn mock_server_metadata(&self) -> MockEndpoint<'_, ServerMetadataEndpoint> {
93        let mock = Mock::given(method("GET"))
94            .and(path_regex(r"^/_matrix/client/unstable/org.matrix.msc2965/auth_metadata"));
95        self.mock_endpoint(mock, ServerMetadataEndpoint::default())
96    }
97
98    /// Creates a prebuilt mock for the OAuth 2.0 endpoint used to register a
99    /// new client.
100    pub fn mock_registration(&self) -> MockEndpoint<'_, RegistrationEndpoint> {
101        let mock = Mock::given(method("POST")).and(path_regex(r"^/oauth2/registration"));
102        self.mock_endpoint(mock, RegistrationEndpoint)
103    }
104
105    /// Creates a prebuilt mock for the OAuth 2.0 endpoint used to authorize a
106    /// device.
107    pub fn mock_device_authorization(&self) -> MockEndpoint<'_, DeviceAuthorizationEndpoint> {
108        let mock = Mock::given(method("POST")).and(path_regex(r"^/oauth2/device"));
109        self.mock_endpoint(mock, DeviceAuthorizationEndpoint)
110    }
111
112    /// Creates a prebuilt mock for the OAuth 2.0 endpoint used to request an
113    /// access token.
114    pub fn mock_token(&self) -> MockEndpoint<'_, TokenEndpoint> {
115        let mock = Mock::given(method("POST")).and(path_regex(r"^/oauth2/token"));
116        self.mock_endpoint(mock, TokenEndpoint)
117    }
118
119    /// Creates a prebuilt mock for the OAuth 2.0 endpoint used to revoke a
120    /// token.
121    pub fn mock_revocation(&self) -> MockEndpoint<'_, RevocationEndpoint> {
122        let mock = Mock::given(method("POST")).and(path_regex(r"^/oauth2/revoke"));
123        self.mock_endpoint(mock, RevocationEndpoint)
124    }
125}
126
127/// A prebuilt mock for a `GET /auth_metadata` request.
128#[derive(Default)]
129pub struct ServerMetadataEndpoint {
130    /// Optional delay to respond to the query.
131    delay: Option<Duration>,
132}
133
134impl<'a> MockEndpoint<'a, ServerMetadataEndpoint> {
135    /// Respond with a given delay to the query.
136    pub fn with_delay(mut self, delay: Duration) -> Self {
137        self.endpoint.delay = Some(delay);
138        self
139    }
140
141    /// Returns a successful response with the given metadata, honouring the
142    /// delay set with [`Self::with_delay`].
143    fn ok_with_metadata(self, metadata: Raw<AuthorizationServerMetadata>) -> MatrixMock<'a> {
144        let mut template = ResponseTemplate::new(200).set_body_json(metadata);
145
146        if let Some(delay) = self.endpoint.delay {
147            template = template.set_delay(delay);
148        }
149
150        self.respond_with(template)
151    }
152
153    /// Returns a successful metadata response with all the supported endpoints.
154    pub fn ok(self) -> MatrixMock<'a> {
155        let metadata = MockServerMetadataBuilder::new(&self.server.uri()).build();
156        self.ok_with_metadata(metadata)
157    }
158
159    /// Returns a successful metadata response with all the supported endpoints
160    /// using HTTPS URLs.
161    ///
162    /// This should be used with
163    /// `MockClientBuilder::insecure_rewrite_https_to_http()` to bypass checks
164    /// from the oauth2 crate.
165    pub fn ok_https(self) -> MatrixMock<'a> {
166        let issuer = self.server.uri().replace("http://", "https://");
167
168        let metadata = MockServerMetadataBuilder::new(&issuer).build();
169        self.ok_with_metadata(metadata)
170    }
171
172    /// Returns a successful metadata response without the device authorization
173    /// endpoint.
174    pub fn ok_without_device_authorization(self) -> MatrixMock<'a> {
175        let metadata = MockServerMetadataBuilder::new(&self.server.uri())
176            .without_device_authorization()
177            .build();
178        self.ok_with_metadata(metadata)
179    }
180
181    /// Returns a successful metadata response without the registration
182    /// endpoint.
183    pub fn ok_without_registration(self) -> MatrixMock<'a> {
184        let metadata =
185            MockServerMetadataBuilder::new(&self.server.uri()).without_registration().build();
186        self.ok_with_metadata(metadata)
187    }
188}
189
190/// Helper struct to construct an `AuthorizationServerMetadata` for integration
191/// tests.
192#[derive(Debug, Clone)]
193pub struct MockServerMetadataBuilder {
194    issuer: Url,
195    with_device_authorization: bool,
196    with_registration: bool,
197}
198
199impl MockServerMetadataBuilder {
200    /// Construct a `MockServerMetadataBuilder` that will generate all the
201    /// supported fields.
202    pub fn new(issuer: &str) -> Self {
203        let issuer = Url::parse(issuer).expect("We should be able to parse the issuer");
204
205        Self { issuer, with_device_authorization: true, with_registration: true }
206    }
207
208    /// Don't generate the field for the device authorization endpoint.
209    fn without_device_authorization(mut self) -> Self {
210        self.with_device_authorization = false;
211        self
212    }
213
214    /// Don't generate the field for the registration endpoint.
215    fn without_registration(mut self) -> Self {
216        self.with_registration = false;
217        self
218    }
219
220    /// The authorization endpoint of this server.
221    fn authorization_endpoint(&self) -> Url {
222        self.issuer.join("oauth2/authorize").unwrap()
223    }
224
225    /// The token endpoint of this server.
226    fn token_endpoint(&self) -> Url {
227        self.issuer.join("oauth2/token").unwrap()
228    }
229
230    /// The JWKS URI of this server.
231    fn jwks_uri(&self) -> Url {
232        self.issuer.join("oauth2/keys.json").unwrap()
233    }
234
235    /// The registration endpoint of this server.
236    fn registration_endpoint(&self) -> Url {
237        self.issuer.join("oauth2/registration").unwrap()
238    }
239
240    /// The account management URI of this server.
241    fn account_management_uri(&self) -> Url {
242        self.issuer.join("account").unwrap()
243    }
244
245    /// The device authorization endpoint of this server.
246    fn device_authorization_endpoint(&self) -> Url {
247        self.issuer.join("oauth2/device").unwrap()
248    }
249
250    /// The revocation endpoint of this server.
251    fn revocation_endpoint(&self) -> Url {
252        self.issuer.join("oauth2/revoke").unwrap()
253    }
254
255    /// Build the server metadata.
256    pub fn build(&self) -> Raw<AuthorizationServerMetadata> {
257        let mut json_metadata = json!({
258            "issuer": self.issuer,
259            "authorization_endpoint": self.authorization_endpoint(),
260            "token_endpoint": self.token_endpoint(),
261            "response_types_supported": ["code"],
262            "response_modes_supported": ["query", "fragment"],
263            "grant_types_supported": ["authorization_code", "refresh_token", "urn:ietf:params:oauth:grant-type:device_code"],
264            "revocation_endpoint": self.revocation_endpoint(),
265            "code_challenge_methods_supported": ["S256"],
266            "account_management_uri": self.account_management_uri(),
267            "account_management_actions_supported": ["org.matrix.profile", "org.matrix.sessions_list", "org.matrix.session_view", "org.matrix.session_end", "org.matrix.deactivateaccount", "org.matrix.cross_signing_reset"],
268            "prompt_values_supported": ["create"],
269        });
270        let json_metadata_object = json_metadata.as_object_mut().unwrap();
271
272        if self.with_device_authorization {
273            json_metadata_object.insert(
274                "device_authorization_endpoint".to_owned(),
275                self.device_authorization_endpoint().as_str().into(),
276            );
277        }
278
279        if self.with_registration {
280            json_metadata_object.insert(
281                "registration_endpoint".to_owned(),
282                self.registration_endpoint().as_str().into(),
283            );
284        }
285
286        serde_json::from_value(json_metadata).unwrap()
287    }
288}
289
290/// A prebuilt mock for a `POST /oauth/registration` request.
291pub struct RegistrationEndpoint;
292
293impl<'a> MockEndpoint<'a, RegistrationEndpoint> {
294    /// Returns a successful registration response.
295    pub fn ok(self) -> MatrixMock<'a> {
296        self.respond_with(ResponseTemplate::new(200).set_body_json(json!({
297            "client_id": "test_client_id",
298            "client_id_issued_at": 1716375696,
299        })))
300    }
301}
302
303/// A prebuilt mock for a `POST /oauth/device` request.
304pub struct DeviceAuthorizationEndpoint;
305
306impl<'a> MockEndpoint<'a, DeviceAuthorizationEndpoint> {
307    /// Returns a successful device authorization response.
308    pub fn ok(self) -> MatrixMock<'a> {
309        let issuer_url = Url::parse(&self.server.uri())
310            .expect("We should be able to parse the wiremock server URI");
311        let verification_uri = issuer_url.join("link").unwrap();
312        let mut verification_uri_complete = issuer_url.join("link").unwrap();
313        verification_uri_complete.set_query(Some("code=N32YVC"));
314
315        self.respond_with(ResponseTemplate::new(200).set_body_json(json!({
316            "device_code": "N8NAYD9fOhMulpm37mSthx0xSw2p7vdR",
317            "expires_in": 1200,
318            "interval": 5,
319            "user_code": "N32YVC",
320            "verification_uri": verification_uri,
321            "verification_uri_complete": verification_uri_complete,
322        })))
323    }
324}
325
326/// A prebuilt mock for a `POST /oauth/token` request.
327pub struct TokenEndpoint;
328
329impl<'a> MockEndpoint<'a, TokenEndpoint> {
330    /// Returns a successful token response with the default tokens.
331    pub fn ok(self) -> MatrixMock<'a> {
332        self.ok_with_tokens("1234", "ZYXWV")
333    }
334
335    /// Returns a successful token response with custom tokens.
336    pub fn ok_with_tokens(self, access_token: &str, refresh_token: &str) -> MatrixMock<'a> {
337        self.respond_with(ResponseTemplate::new(200).set_body_json(json!({
338            "access_token": access_token,
339            "expires_in": 300,
340            "refresh_token":  refresh_token,
341            "token_type": "Bearer"
342        })))
343    }
344
345    /// Returns an error response when the request was invalid.
346    pub fn access_denied(self) -> MatrixMock<'a> {
347        self.respond_with(ResponseTemplate::new(400).set_body_json(json!({
348            "error": "access_denied",
349        })))
350    }
351
352    /// Returns an error response when the token in the request has expired.
353    pub fn expired_token(self) -> MatrixMock<'a> {
354        self.respond_with(ResponseTemplate::new(400).set_body_json(json!({
355            "error": "expired_token",
356        })))
357    }
358
359    /// Returns an error response when the token in the request is invalid.
360    pub fn invalid_grant(self) -> MatrixMock<'a> {
361        self.respond_with(ResponseTemplate::new(400).set_body_json(json!({
362            "error": "invalid_grant",
363        })))
364    }
365}
366
367/// A prebuilt mock for a `POST /oauth/revoke` request.
368pub struct RevocationEndpoint;
369
370impl<'a> MockEndpoint<'a, RevocationEndpoint> {
371    /// Returns a successful revocation response.
372    pub fn ok(self) -> MatrixMock<'a> {
373        self.respond_with(ResponseTemplate::new(200).set_body_json(json!({})))
374    }
375}