oauth2_test_server/
testkit.rs

1use base64::{engine::general_purpose, Engine};
2use jsonwebtoken::{encode, Algorithm, Header};
3use reqwest::{Client as HttpClient, StatusCode};
4use serde_json::{json, Value};
5use sha2::{Digest, Sha256};
6use std::{
7    collections::HashMap,
8    sync::{Arc, RwLock},
9};
10use tokio::task::JoinError;
11use url::Url;
12use uuid::Uuid;
13
14#[derive(Debug, Clone)]
15pub struct PkcePair {
16    pub code_verifier: String,
17    pub code_challenge: String,
18}
19
20#[derive(Debug, Default)]
21pub struct AuthorizeParams {
22    pub response_type: &'static str,
23    pub redirect_uri: String,
24    pub scope: String,
25    pub state: String,
26    pub pkce: Option<PkcePair>,
27    pub nonce: Option<String>,
28}
29
30impl AuthorizeParams {
31    pub fn new() -> Self {
32        Self {
33            response_type: "code",
34            redirect_uri: "http://localhost/cb".to_string(),
35            scope: "openid".to_string(),
36            state: Uuid::new_v4().to_string(),
37            pkce: None,
38            nonce: None,
39        }
40    }
41
42    pub fn redirect_uri(mut self, uri: impl Into<String>) -> Self {
43        self.redirect_uri = uri.into();
44        self
45    }
46
47    pub fn scope(mut self, scope: impl Into<String>) -> Self {
48        self.scope = scope.into();
49        self
50    }
51
52    pub fn state(mut self, state: impl Into<String>) -> Self {
53        self.state = state.into();
54        self
55    }
56
57    pub fn pkce(mut self, pkce: PkcePair) -> Self {
58        self.pkce = Some(pkce);
59        self
60    }
61
62    pub fn nonce(mut self, nonce: impl Into<String>) -> Self {
63        self.nonce = Some(nonce.into());
64        self
65    }
66}
67
68#[derive(Debug, Clone)]
69pub struct OauthEndpoints {
70    pub oauth_server: String,
71    pub discovery: String,
72    pub authorize: String,
73    pub token: String,
74    pub regsiter: String,
75    pub introspect: String,
76    pub revoke: String,
77    pub userinfo: String,
78}
79
80/// Start a test server with full programmatic control.
81///
82/// ```
83/// use oauth2_test_server::OAuthTestServer;
84///
85/// #[tokio::test]
86/// async fn test() {
87/// let server = OAuthTestServer::start().await;
88/// println!("server: {}", server.base_url());
89/// println!("authorize endpoint: {}", server.endpoints.authorize_url);
90/// // register a client
91/// let client = server.register_client(
92///     serde_json::json!({ "scope": "openid", "redirect_uris":["http://localhost:8080/callback"]}),
93/// );
94/// // generate a jwt
95/// let jwt = server.generate_jwt(&client, server.jwt_options().user_id("bob").build());
96/// assert_eq!(jwt.split('.').count(), 3);
97/// assert_eq!(server.clients().read().iter().len(), 1);
98/// assert_eq!(server.tokens().read().iter().len(), 1);
99/// }
100/// ```
101pub struct OAuthTestServer {
102    state: AppState,
103    pub base_url: url::Url,
104    pub endpoints: OauthEndpoints,
105    http: HttpClient,
106    _handle: tokio::task::JoinHandle<()>,
107}
108
109impl OAuthTestServer {
110    pub async fn start() -> Self {
111        let config = IssuerConfig {
112            port: 0,
113            ..Default::default()
114        };
115        Self::start_with_config(config).await
116    }
117
118    pub fn clients(&self) -> Arc<RwLock<HashMap<String, Client>>> {
119        self.state.clients.clone()
120    }
121
122    pub fn codes(&self) -> Arc<RwLock<HashMap<String, AuthorizationCode>>> {
123        self.state.codes.clone()
124    }
125
126    pub fn tokens(&self) -> Arc<RwLock<HashMap<String, Token>>> {
127        self.state.tokens.clone()
128    }
129
130    pub fn refresh_tokens(&self) -> Arc<RwLock<HashMap<String, Token>>> {
131        self.state.refresh_tokens.clone()
132    }
133
134    pub async fn start_with_config(config: IssuerConfig) -> Self {
135        // config.port = 0;
136        let state = AppState::new(config.clone());
137        let (addr, handle) = state.clone().start().await;
138        let base_url: Url = format!("http://{addr}").parse().unwrap();
139
140        let endpoints: OauthEndpoints = OauthEndpoints {
141            oauth_server: base_url.clone().to_string(),
142            discovery: format!("{base_url}.well-known/openid-configuration"),
143            authorize: format!("{base_url}register"),
144            regsiter: format!("{base_url}authorize"),
145            token: format!("{base_url}token"),
146            introspect: format!("{base_url}introspect"),
147            revoke: format!("{base_url}revoke"),
148            userinfo: format!("{base_url}userinfo"),
149        };
150
151        Self {
152            state,
153            base_url,
154            endpoints,
155            http: HttpClient::new(),
156            _handle: handle,
157        }
158    }
159
160    pub async fn wait_for_shutdown(self) -> Result<(), JoinError> {
161        self._handle.await
162    }
163
164    pub fn register_client(&self, metadata: serde_json::Value) -> Client {
165        self.state
166            .register_client(metadata)
167            .expect("client registration failed")
168    }
169
170    pub fn register_client_with_secret(&self, metadata: Value, force_secret: bool) -> Client {
171        let mut meta = metadata;
172        if let Some(obj) = meta.as_object_mut() {
173            obj.insert(
174                "generate_client_secret_for_dcr".to_string(),
175                json!(force_secret),
176            );
177        }
178        self.register_client(meta)
179    }
180
181    pub fn generate_jwt(&self, client: &Client, options: JwtOptions) -> String {
182        self.state
183            .generate_jwt(client, options)
184            .expect("JWT generation failed")
185    }
186
187    pub fn jwt_options(&self) -> JwtOptionsBuilder {
188        JwtOptionsBuilder::default()
189    }
190
191    pub fn pkce_pair(&self) -> PkcePair {
192        use rand::Rng;
193        let verifier_bytes: [u8; 32] = rand::thread_rng().r#gen();
194        let code_verifier = general_purpose::URL_SAFE_NO_PAD.encode(verifier_bytes);
195        let challenge =
196            general_purpose::URL_SAFE_NO_PAD.encode(Sha256::digest(code_verifier.as_bytes()));
197        PkcePair {
198            code_verifier,
199            code_challenge: challenge,
200        }
201    }
202
203    pub fn authorize_url(&self, client: &Client, params: AuthorizeParams) -> Url {
204        let mut url = self.base_url.join("authorize").unwrap();
205        let mut query = url.query_pairs_mut();
206
207        query
208            .append_pair("response_type", params.response_type)
209            .append_pair("client_id", &client.client_id)
210            .append_pair("redirect_uri", &params.redirect_uri)
211            .append_pair("scope", &params.scope)
212            .append_pair("state", &params.state);
213
214        if let Some(pkce) = params.pkce {
215            query
216                .append_pair("code_challenge", &pkce.code_challenge)
217                .append_pair("code_challenge_method", "S256");
218        }
219
220        if let Some(nonce) = params.nonce {
221            query.append_pair("nonce", &nonce);
222        }
223
224        drop(query);
225        url
226    }
227
228    pub fn rotate_keys(&self) {
229        // In real impl, regenerate KEYS and update JWKS_JSON
230        unimplemented!("Key rotation not implemented in test server")
231    }
232
233    pub async fn approve_consent(&self, auth_url: &Url, user_id: &str) -> String {
234        let resp = self.http.get(auth_url.clone()).send().await.unwrap();
235        assert_eq!(resp.status(), StatusCode::SEE_OTHER);
236
237        let location = resp.headers().get("location").unwrap().to_str().unwrap();
238        let redirect = Url::parse(location).unwrap();
239        let code = redirect
240            .query_pairs()
241            .find(|(k, _)| k == "code")
242            .map(|(_, v)| v.to_string())
243            .expect("no code in redirect");
244
245        // Store user_id for later token claims
246        let code_obj = self
247            .state
248            .codes
249            .read()
250            .unwrap()
251            .get(&code)
252            .cloned()
253            .unwrap();
254        let mut code_obj = code_obj;
255        code_obj.user_id = user_id.to_string();
256        self.state
257            .codes
258            .write()
259            .unwrap()
260            .insert(code.clone(), code_obj);
261
262        code
263    }
264
265    pub async fn exchange_code(
266        &self,
267        client: &Client,
268        code: &str,
269        pkce: Option<&PkcePair>,
270    ) -> Value {
271        let mut form = vec![
272            ("grant_type", "authorization_code"),
273            ("code", code),
274            ("redirect_uri", "http://localhost/cb"),
275        ];
276
277        if let Some(pkce) = pkce {
278            form.push(("code_verifier", &pkce.code_verifier));
279        }
280
281        let resp = self
282            .http
283            .post(self.base_url.join("token").unwrap())
284            .basic_auth(&client.client_id, client.client_secret.as_ref())
285            .form(&form)
286            .send()
287            .await
288            .unwrap();
289
290        assert_eq!(resp.status(), StatusCode::OK);
291        resp.json().await.unwrap()
292    }
293
294    pub async fn refresh_token(&self, client: &Client, refresh_token: &str) -> Value {
295        let resp = self
296            .http
297            .post(self.base_url.join("token").unwrap())
298            .basic_auth(&client.client_id, client.client_secret.as_ref())
299            .form(&[
300                ("grant_type", "refresh_token"),
301                ("refresh_token", refresh_token),
302            ])
303            .send()
304            .await
305            .unwrap();
306
307        resp.json().await.unwrap()
308    }
309
310    pub async fn introspect_token(&self, client: &Client, token: &str) -> Value {
311        let resp = self
312            .http
313            .post(self.base_url.join("introspect").unwrap())
314            .basic_auth(&client.client_id, client.client_secret.as_ref())
315            .form(&[("token", token)])
316            .send()
317            .await
318            .unwrap();
319
320        resp.json().await.unwrap()
321    }
322
323    pub async fn revoke_token(&self, client: &Client, token: &str) {
324        let resp = self
325            .http
326            .post(self.base_url.join("revoke").unwrap())
327            .basic_auth(&client.client_id, client.client_secret.as_ref())
328            .form(&[("token", token)])
329            .send()
330            .await
331            .unwrap();
332
333        assert!(resp.status().is_success());
334    }
335
336    pub fn client_assertion_jwt(&self, client: &Client) -> String {
337        let claims = json!({
338            "iss": client.client_id,
339            "sub": client.client_id,
340            "aud": self.issuer(),
341            "exp": (chrono::Utc::now() + chrono::Duration::minutes(5)).timestamp(),
342            "iat": chrono::Utc::now().timestamp(),
343            "jti": Uuid::new_v4().to_string(),
344        });
345
346        let mut header = Header::new(Algorithm::RS256);
347        header.kid = Some(KID.to_string());
348
349        encode(&header, &claims, &KEYS.encoding).unwrap()
350    }
351
352    pub fn base_url(&self) -> &url::Url {
353        &self.base_url
354    }
355
356    pub fn issuer(&self) -> &str {
357        self.state.issuer()
358    }
359}
360
361#[derive(Debug, Default)]
362pub struct JwtOptions {
363    pub user_id: String,
364    pub scope: Option<String>,
365    pub expires_in: i64,
366}
367
368#[derive(Default)]
369pub struct JwtOptionsBuilder {
370    user_id: Option<String>,
371    scope: Option<String>,
372    expires_in: Option<i64>,
373}
374
375impl JwtOptionsBuilder {
376    pub fn user_id(mut self, user_id: impl Into<String>) -> Self {
377        self.user_id = Some(user_id.into());
378        self
379    }
380
381    pub fn scope(mut self, scope: impl Into<String>) -> Self {
382        self.scope = Some(scope.into());
383        self
384    }
385
386    pub fn expires_in(mut self, seconds: i64) -> Self {
387        self.expires_in = Some(seconds);
388        self
389    }
390
391    pub fn build(self) -> JwtOptions {
392        JwtOptions {
393            user_id: self.user_id.unwrap_or("test-user-123".to_string()),
394            scope: self.scope,
395            expires_in: self.expires_in.unwrap_or(3600),
396        }
397    }
398}
399
400use crate::server::{AppState, AuthorizationCode, Client, IssuerConfig, Token, KEYS, KID};
oauth2_test_server/testkit.rs

oauth2_test_server/
testkit.rs
