Skip to main content

max / makenotwork

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