Skip to main content

max / makenotwork

16.6 KB · 525 lines History Blame Raw
1 //! Cookie-aware in-process HTTP client.
2 //!
3 //! Wraps an Axum `Router` and uses `tower::ServiceExt::oneshot` for each
4 //! request. Manages cookies across requests and auto-injects CSRF tokens.
5
6 use axum::Router;
7 use axum::body::Body;
8 use axum::http::{Method, Request, StatusCode, header};
9 use http_body_util::BodyExt;
10 use std::collections::HashMap;
11 use std::fmt::Write as _;
12 use std::sync::atomic::{AtomicU32, Ordering};
13 use tower::ServiceExt;
14
15 /// Monotonic counter for unique per-test IPs (10.1.x.y).
16 static IP_COUNTER: AtomicU32 = AtomicU32::new(1);
17
18 /// A test HTTP client that talks to the app in-process.
19 pub(crate) struct TestClient {
20 app: Router,
21 cookies: HashMap<String, String>,
22 csrf_token: Option<String>,
23 forwarded_ip: String,
24 bearer_token: Option<String>,
25 actor_token: Option<String>,
26 }
27
28 impl TestClient {
29 pub(crate) fn new(app: Router) -> Self {
30 let n = IP_COUNTER.fetch_add(1, Ordering::Relaxed);
31 let octet3 = (n / 256) % 256;
32 let octet4 = n % 256;
33 TestClient {
34 app,
35 cookies: HashMap::new(),
36 csrf_token: None,
37 forwarded_ip: format!("10.1.{octet3}.{octet4}"),
38 bearer_token: None,
39 actor_token: None,
40 }
41 }
42
43 /// Fork a fresh client that talks to the same app but with its own empty
44 /// cookie jar, CSRF token, and unique IP, simulates a second, independent
45 /// browser session (e.g. the same user logged in on another device).
46 #[allow(dead_code)]
47 pub(crate) fn fork_fresh(&self) -> TestClient {
48 TestClient::new(self.app.clone())
49 }
50
51 /// Set the IP address used in the X-Forwarded-For header.
52 #[allow(dead_code)]
53 pub(crate) fn set_forwarded_ip(&mut self, ip: &str) {
54 self.forwarded_ip = ip.to_string();
55 }
56
57 /// Set a bearer token for subsequent requests (used by SyncKit JWT auth).
58 #[allow(dead_code)]
59 pub(crate) fn set_bearer_token(&mut self, token: &str) {
60 self.bearer_token = Some(token.to_string());
61 }
62
63 #[allow(dead_code)]
64 pub(crate) fn clear_bearer_token(&mut self) {
65 self.bearer_token = None;
66 }
67
68 /// Set the `X-MNW-Actor` assertion sent on internal-API requests (mirrors
69 /// what the CLI forwards after ssh-key-lookup).
70 #[allow(dead_code)]
71 pub(crate) fn set_actor_token(&mut self, token: &str) {
72 self.actor_token = Some(token.to_string());
73 }
74
75 /// Drop all stored cookies, simulates a fresh client with no session, e.g.
76 /// a CLI `git` request that authenticates via a token rather than a browser
77 /// cookie.
78 #[allow(dead_code)]
79 pub(crate) fn clear_cookies(&mut self) {
80 self.cookies.clear();
81 }
82
83 /// Access the current CSRF token (if any).
84 #[allow(dead_code)]
85 pub(crate) fn csrf_token(&self) -> Option<&str> {
86 self.csrf_token.as_deref()
87 }
88
89 /// GET request.
90 pub(crate) async fn get(&mut self, uri: &str) -> TestResponse {
91 self.request(Method::GET, uri, None, None).await
92 }
93
94 /// POST with form-encoded body.
95 pub(crate) async fn post_form(&mut self, uri: &str, body: &str) -> TestResponse {
96 self.request(
97 Method::POST,
98 uri,
99 Some("application/x-www-form-urlencoded"),
100 Some(body.to_string()),
101 )
102 .await
103 }
104
105 /// POST with JSON body.
106 #[allow(dead_code)]
107 pub(crate) async fn post_json(&mut self, uri: &str, body: &str) -> TestResponse {
108 self.request(
109 Method::POST,
110 uri,
111 Some("application/json"),
112 Some(body.to_string()),
113 )
114 .await
115 }
116
117 /// PUT with form-encoded body.
118 pub(crate) async fn put_form(&mut self, uri: &str, body: &str) -> TestResponse {
119 self.request(
120 Method::PUT,
121 uri,
122 Some("application/x-www-form-urlencoded"),
123 Some(body.to_string()),
124 )
125 .await
126 }
127
128 /// PUT with JSON body.
129 #[allow(dead_code)]
130 pub(crate) async fn put_json(&mut self, uri: &str, body: &str) -> TestResponse {
131 self.request(
132 Method::PUT,
133 uri,
134 Some("application/json"),
135 Some(body.to_string()),
136 )
137 .await
138 }
139
140 /// DELETE request.
141 #[allow(dead_code)]
142 pub(crate) async fn delete(&mut self, uri: &str) -> TestResponse {
143 self.request(Method::DELETE, uri, None, None).await
144 }
145
146 /// PATCH with JSON body.
147 #[allow(dead_code)]
148 pub(crate) async fn patch_json(&mut self, uri: &str, body: &str) -> TestResponse {
149 self.request(
150 Method::PATCH,
151 uri,
152 Some("application/json"),
153 Some(body.to_string()),
154 )
155 .await
156 }
157
158 /// DELETE with form-encoded body.
159 #[allow(dead_code)]
160 pub(crate) async fn delete_form(&mut self, uri: &str, body: &str) -> TestResponse {
161 self.request(
162 Method::DELETE,
163 uri,
164 Some("application/x-www-form-urlencoded"),
165 Some(body.to_string()),
166 )
167 .await
168 }
169
170 /// POST with multipart/form-data body. Fields are (name, value) pairs.
171 /// Supports repeated field names for Vec<T> deserialization.
172 #[allow(dead_code)]
173 pub(crate) async fn post_multipart(
174 &mut self,
175 uri: &str,
176 fields: &[(&str, &str)],
177 ) -> TestResponse {
178 let boundary = "----TestBoundary7MA4YWxkTrZu0gW";
179 let mut body = String::new();
180 for (name, value) in fields {
181 writeln!(body, "--{boundary}\r").unwrap();
182 writeln!(
183 body,
184 "Content-Disposition: form-data; name=\"{name}\"\r\n\r\n{value}\r"
185 )
186 .unwrap();
187 }
188 writeln!(body, "--{boundary}--\r").unwrap();
189
190 let content_type = format!("multipart/form-data; boundary={boundary}");
191 self.request(Method::POST, uri, Some(&content_type), Some(body))
192 .await
193 }
194
195 /// HTMX GET request (includes `HX-Request: true` header).
196 #[allow(dead_code)]
197 pub(crate) async fn htmx_get(&mut self, uri: &str) -> TestResponse {
198 self.request_htmx(Method::GET, uri, None, None).await
199 }
200
201 /// HTMX POST with form-encoded body.
202 #[allow(dead_code)]
203 pub(crate) async fn htmx_post_form(&mut self, uri: &str, body: &str) -> TestResponse {
204 self.request_htmx(
205 Method::POST,
206 uri,
207 Some("application/x-www-form-urlencoded"),
208 Some(body.to_string()),
209 )
210 .await
211 }
212
213 /// HTMX PUT with form-encoded body.
214 #[allow(dead_code)]
215 pub(crate) async fn htmx_put_form(&mut self, uri: &str, body: &str) -> TestResponse {
216 self.request_htmx(
217 Method::PUT,
218 uri,
219 Some("application/x-www-form-urlencoded"),
220 Some(body.to_string()),
221 )
222 .await
223 }
224
225 /// HTMX DELETE request (includes `HX-Request: true` header).
226 #[allow(dead_code)]
227 pub(crate) async fn htmx_delete(&mut self, uri: &str) -> TestResponse {
228 self.request_htmx(Method::DELETE, uri, None, None).await
229 }
230
231 /// Fetch the CSRF token by loading the /login page and extracting it
232 /// from `<meta name="csrf-token" content="...">`.
233 pub(crate) async fn fetch_csrf_token(&mut self) {
234 let resp = self.get("/login").await;
235 if let Some(token) = extract_csrf_from_html(&resp.text) {
236 self.csrf_token = Some(token);
237 }
238 }
239
240 /// Build and send a regular request (no HTMX header).
241 async fn request(
242 &mut self,
243 method: Method,
244 uri: &str,
245 content_type: Option<&str>,
246 body: Option<String>,
247 ) -> TestResponse {
248 self.send(method, uri, content_type, body, false).await
249 }
250
251 /// Build and send a request with the `HX-Request: true` header.
252 #[allow(dead_code)]
253 async fn request_htmx(
254 &mut self,
255 method: Method,
256 uri: &str,
257 content_type: Option<&str>,
258 body: Option<String>,
259 ) -> TestResponse {
260 self.send(method, uri, content_type, body, true).await
261 }
262
263 /// Build and send a request through `oneshot`, optionally with the HTMX header.
264 async fn send(
265 &mut self,
266 method: Method,
267 uri: &str,
268 content_type: Option<&str>,
269 body: Option<String>,
270 htmx: bool,
271 ) -> TestResponse {
272 let body_data = body.unwrap_or_default();
273 let mut builder = Request::builder()
274 .method(&method)
275 .uri(uri)
276 // Required: SmartIpKeyExtractor needs an IP; oneshot has no ConnectInfo
277 .header("X-Forwarded-For", &self.forwarded_ip)
278 // Production reads `CF-Connecting-IP` (the only header origin clients
279 // can't spoof through Caddy). Send both so the harness matches the
280 // production proxy chain and `extract_client_ip` finds a value.
281 .header("CF-Connecting-IP", &self.forwarded_ip);
282
283 if htmx {
284 builder = builder.header("HX-Request", "true");
285 }
286
287 // Inject bearer token if set
288 if let Some(ref token) = self.bearer_token {
289 builder = builder.header(header::AUTHORIZATION, format!("Bearer {token}"));
290 }
291 if let Some(ref actor) = self.actor_token {
292 builder = builder.header("X-MNW-Actor", actor);
293 }
294
295 // Set content type
296 if let Some(ct) = content_type {
297 builder = builder.header(header::CONTENT_TYPE, ct);
298 }
299
300 // Inject CSRF token for mutating methods
301 if matches!(
302 method,
303 Method::POST | Method::PUT | Method::PATCH | Method::DELETE
304 ) && let Some(ref token) = self.csrf_token
305 {
306 builder = builder.header("X-CSRF-Token", token.as_str());
307 }
308
309 // Attach cookies
310 if !self.cookies.is_empty() {
311 let cookie_header: String = self
312 .cookies
313 .iter()
314 .map(|(k, v)| format!("{k}={v}"))
315 .collect::<Vec<_>>()
316 .join("; ");
317 builder = builder.header(header::COOKIE, cookie_header);
318 }
319
320 let request = builder
321 .body(Body::from(body_data))
322 .expect("Failed to build request");
323
324 let response = self
325 .app
326 .clone()
327 .oneshot(request)
328 .await
329 .expect("Failed to send request");
330
331 // Collect set-cookie headers before consuming response
332 let status = response.status();
333 let headers = response.headers().clone();
334
335 // Store cookies from response
336 for value in headers.get_all(header::SET_COOKIE) {
337 if let Ok(cookie_str) = value.to_str() {
338 // Parse "name=value; ...", take only the name=value part
339 if let Some(nv) = cookie_str.split(';').next()
340 && let Some((name, val)) = nv.split_once('=')
341 {
342 self.cookies
343 .insert(name.trim().to_string(), val.trim().to_string());
344 }
345 }
346 }
347
348 // Read body
349 let body_bytes = response
350 .into_body()
351 .collect()
352 .await
353 .expect("Failed to read response body")
354 .to_bytes();
355 let text = String::from_utf8_lossy(&body_bytes).to_string();
356
357 // Auto-extract CSRF token from HTML responses for convenience
358 if let Some(token) = extract_csrf_from_html(&text) {
359 self.csrf_token = Some(token);
360 }
361
362 TestResponse {
363 status,
364 text,
365 headers,
366 }
367 }
368
369 /// Raw request with custom headers. Used for webhook tests where we need
370 /// to set the stripe-signature header and bypass CSRF.
371 #[allow(dead_code)]
372 pub(crate) async fn request_with_headers(
373 &mut self,
374 method: &str,
375 uri: &str,
376 body: Option<&str>,
377 extra_headers: &[(&str, &str)],
378 ) -> TestResponse {
379 let body_data = body.unwrap_or_default().to_string();
380 let mut builder = Request::builder()
381 .method(method)
382 .uri(uri)
383 .header("X-Forwarded-For", &self.forwarded_ip)
384 // Production reads `CF-Connecting-IP` (the only header origin clients
385 // can't spoof through Caddy). Send both so the harness matches the
386 // production proxy chain and `extract_client_ip` finds a value.
387 .header("CF-Connecting-IP", &self.forwarded_ip);
388
389 for (name, value) in extra_headers {
390 builder = builder.header(*name, *value);
391 }
392
393 // Attach cookies
394 if !self.cookies.is_empty() {
395 let cookie_header: String = self
396 .cookies
397 .iter()
398 .map(|(k, v)| format!("{k}={v}"))
399 .collect::<Vec<_>>()
400 .join("; ");
401 builder = builder.header(header::COOKIE, cookie_header);
402 }
403
404 let request = builder
405 .body(Body::from(body_data))
406 .expect("Failed to build request");
407
408 let response = self
409 .app
410 .clone()
411 .oneshot(request)
412 .await
413 .expect("Failed to send request");
414
415 let status = response.status();
416 let resp_headers = response.headers().clone();
417
418 for value in resp_headers.get_all(header::SET_COOKIE) {
419 if let Ok(cookie_str) = value.to_str()
420 && let Some(nv) = cookie_str.split(';').next()
421 && let Some((name, val)) = nv.split_once('=')
422 {
423 self.cookies
424 .insert(name.trim().to_string(), val.trim().to_string());
425 }
426 }
427
428 let body_bytes = response
429 .into_body()
430 .collect()
431 .await
432 .expect("Failed to read response body")
433 .to_bytes();
434 let text = String::from_utf8_lossy(&body_bytes).to_string();
435
436 TestResponse {
437 status,
438 text,
439 headers: resp_headers,
440 }
441 }
442
443 /// Send a GET request and return status + headers WITHOUT reading the body.
444 /// Use for streaming endpoints (SSE) where the body never ends.
445 #[allow(dead_code)]
446 pub(crate) async fn get_streaming(&mut self, uri: &str) -> TestResponse {
447 let mut builder = Request::builder()
448 .method(Method::GET)
449 .uri(uri)
450 .header("X-Forwarded-For", &self.forwarded_ip)
451 // Production reads `CF-Connecting-IP` (the only header origin clients
452 // can't spoof through Caddy). Send both so the harness matches the
453 // production proxy chain and `extract_client_ip` finds a value.
454 .header("CF-Connecting-IP", &self.forwarded_ip);
455
456 if let Some(ref token) = self.bearer_token {
457 builder = builder.header(header::AUTHORIZATION, format!("Bearer {token}"));
458 }
459 if let Some(ref actor) = self.actor_token {
460 builder = builder.header("X-MNW-Actor", actor);
461 }
462
463 if !self.cookies.is_empty() {
464 let cookie_header: String = self
465 .cookies
466 .iter()
467 .map(|(k, v)| format!("{k}={v}"))
468 .collect::<Vec<_>>()
469 .join("; ");
470 builder = builder.header(header::COOKIE, cookie_header);
471 }
472
473 let request = builder
474 .body(Body::empty())
475 .expect("Failed to build request");
476
477 let response = self
478 .app
479 .clone()
480 .oneshot(request)
481 .await
482 .expect("Failed to send request");
483
484 let status = response.status();
485 let headers = response.headers().clone();
486
487 // Do NOT read the body, return immediately
488 TestResponse {
489 status,
490 text: String::new(),
491 headers,
492 }
493 }
494 }
495
496 /// Response wrapper with convenience methods.
497 #[allow(dead_code)]
498 pub(crate) struct TestResponse {
499 pub status: StatusCode,
500 pub text: String,
501 pub headers: axum::http::HeaderMap,
502 }
503
504 #[allow(dead_code)]
505 impl TestResponse {
506 /// Parse the body as JSON.
507 pub(crate) fn json<T: serde::de::DeserializeOwned>(&self) -> T {
508 serde_json::from_str(&self.text)
509 .unwrap_or_else(|e| panic!("Failed to parse JSON: {}\nBody: {}", e, &self.text))
510 }
511
512 /// Get a header value as a string.
513 pub(crate) fn header(&self, name: &str) -> Option<&str> {
514 self.headers.get(name).and_then(|v| v.to_str().ok())
515 }
516 }
517
518 /// Extract CSRF token from `<meta name="csrf-token" content="...">`.
519 fn extract_csrf_from_html(html: &str) -> Option<String> {
520 let marker = "csrf-token\" content=\"";
521 let start = html.find(marker)? + marker.len();
522 let end = html[start..].find('"')? + start;
523 Some(html[start..end].to_string())
524 }
525