Skip to main content

max / makenotwork

11.1 KB · 381 lines History Blame Raw
1 //! Cookie-aware in-process HTTP client for integration tests.
2 //!
3 //! Carries a cookie jar across requests and scrapes the CSRF token out of every
4 //! HTML response, injecting it on each mutating method, so a test writes the
5 //! request it means rather than the session bookkeeping around it.
6 //! `post_form_no_csrf` and `post_form_with_token` opt out on purpose: they exist
7 //! to test that rejection works.
8
9 use axum::Router;
10 use axum::body::Body;
11 use axum::extract::ConnectInfo;
12 use axum::http::{Method, Request, StatusCode, header};
13 use http_body_util::BodyExt;
14 use std::collections::HashMap;
15 use std::net::SocketAddr;
16 use tower::ServiceExt;
17
18 pub(crate) struct TestClient {
19 app: Router,
20 cookies: HashMap<String, String>,
21 csrf_token: Option<String>,
22 }
23
24 impl TestClient {
25 pub(crate) fn new(app: Router) -> Self {
26 TestClient {
27 app,
28 cookies: HashMap::new(),
29 csrf_token: None,
30 }
31 }
32
33 pub(crate) async fn get(&mut self, uri: &str) -> TestResponse {
34 self.send(Method::GET, uri, None, None).await
35 }
36
37 pub(crate) async fn post_form(&mut self, uri: &str, body: &str) -> TestResponse {
38 self.send(
39 Method::POST,
40 uri,
41 Some("application/x-www-form-urlencoded"),
42 Some(body.to_string()),
43 )
44 .await
45 }
46
47 /// POST without injecting the CSRF token. Used to test CSRF rejection.
48 pub(crate) async fn post_form_no_csrf(&mut self, uri: &str, body: &str) -> TestResponse {
49 self.send_raw(
50 Method::POST,
51 uri,
52 Some("application/x-www-form-urlencoded"),
53 Some(body.to_string()),
54 false,
55 )
56 .await
57 }
58
59 /// POST with a specific (wrong) CSRF token.
60 pub(crate) async fn post_form_with_token(
61 &mut self,
62 uri: &str,
63 body: &str,
64 token: &str,
65 ) -> TestResponse {
66 self.send_raw_with_token(
67 Method::POST,
68 uri,
69 Some("application/x-www-form-urlencoded"),
70 Some(body.to_string()),
71 token,
72 )
73 .await
74 }
75
76 pub(crate) async fn post_json(&mut self, uri: &str, body: &str) -> TestResponse {
77 self.send(
78 Method::POST,
79 uri,
80 Some("application/json"),
81 Some(body.to_string()),
82 )
83 .await
84 }
85
86 pub(crate) async fn post_multipart(
87 &mut self,
88 uri: &str,
89 file_data: &[u8],
90 content_type: &str,
91 filename: &str,
92 ) -> TestResponse {
93 let boundary = "----TestBoundary1234567890";
94 let mut body = Vec::new();
95 body.extend_from_slice(format!("--{boundary}\r\n").as_bytes());
96 body.extend_from_slice(
97 format!(
98 "Content-Disposition: form-data; name=\"file\"; filename=\"{filename}\"\r\n\
99 Content-Type: {content_type}\r\n\r\n"
100 )
101 .as_bytes(),
102 );
103 body.extend_from_slice(file_data);
104 body.extend_from_slice(format!("\r\n--{boundary}--\r\n").as_bytes());
105
106 let ct = format!("multipart/form-data; boundary={boundary}");
107 self.send_bytes(Method::POST, uri, &ct, body).await
108 }
109
110 pub(crate) fn csrf_token(&self) -> Option<&str> {
111 self.csrf_token.as_deref()
112 }
113
114 async fn send(
115 &mut self,
116 method: Method,
117 uri: &str,
118 content_type: Option<&str>,
119 body: Option<String>,
120 ) -> TestResponse {
121 self.send_raw(method, uri, content_type, body, true).await
122 }
123
124 async fn send_bytes(
125 &mut self,
126 method: Method,
127 uri: &str,
128 content_type: &str,
129 body: Vec<u8>,
130 ) -> TestResponse {
131 let mut builder = Request::builder()
132 .method(&method)
133 .uri(uri)
134 .header(header::CONTENT_TYPE, content_type);
135
136 if matches!(
137 method,
138 Method::POST | Method::PUT | Method::PATCH | Method::DELETE
139 ) && let Some(ref token) = self.csrf_token
140 {
141 builder = builder.header("X-CSRF-Token", token.as_str());
142 }
143
144 if !self.cookies.is_empty() {
145 let cookie_header: String = self
146 .cookies
147 .iter()
148 .map(|(k, v)| format!("{k}={v}"))
149 .collect::<Vec<_>>()
150 .join("; ");
151 builder = builder.header(header::COOKIE, cookie_header);
152 }
153
154 let mut request = builder
155 .body(Body::from(body))
156 .expect("Failed to build request");
157 request
158 .extensions_mut()
159 .insert(ConnectInfo(SocketAddr::from(([127, 0, 0, 1], 0))));
160
161 let response = self
162 .app
163 .clone()
164 .oneshot(request)
165 .await
166 .expect("Failed to send request");
167 let status = response.status();
168 let headers = response.headers().clone();
169
170 for value in headers.get_all(header::SET_COOKIE) {
171 if let Ok(cookie_str) = value.to_str()
172 && let Some(nv) = cookie_str.split(';').next()
173 && let Some((name, val)) = nv.split_once('=')
174 {
175 self.cookies
176 .insert(name.trim().to_string(), val.trim().to_string());
177 }
178 }
179
180 let body_bytes = response
181 .into_body()
182 .collect()
183 .await
184 .expect("Failed to read response body")
185 .to_bytes();
186 let text = String::from_utf8_lossy(&body_bytes).to_string();
187
188 if let Some(token) = extract_csrf_from_html(&text) {
189 self.csrf_token = Some(token);
190 }
191
192 TestResponse {
193 status,
194 text,
195 headers,
196 }
197 }
198
199 async fn send_raw(
200 &mut self,
201 method: Method,
202 uri: &str,
203 content_type: Option<&str>,
204 body: Option<String>,
205 inject_csrf: bool,
206 ) -> TestResponse {
207 let body_data = body.unwrap_or_default();
208 let mut builder = Request::builder().method(&method).uri(uri);
209
210 if let Some(ct) = content_type {
211 builder = builder.header(header::CONTENT_TYPE, ct);
212 }
213
214 if inject_csrf
215 && matches!(
216 method,
217 Method::POST | Method::PUT | Method::PATCH | Method::DELETE
218 )
219 && let Some(ref token) = self.csrf_token
220 {
221 builder = builder.header("X-CSRF-Token", token.as_str());
222 }
223
224 if !self.cookies.is_empty() {
225 let cookie_header: String = self
226 .cookies
227 .iter()
228 .map(|(k, v)| format!("{k}={v}"))
229 .collect::<Vec<_>>()
230 .join("; ");
231 builder = builder.header(header::COOKIE, cookie_header);
232 }
233
234 let mut request = builder
235 .body(Body::from(body_data))
236 .expect("Failed to build request");
237 // TrustedProxyKeyExtractor keys on the direct peer, so without
238 // ConnectInfo every rate-limit test shares one bucket
239 request
240 .extensions_mut()
241 .insert(ConnectInfo(SocketAddr::from(([127, 0, 0, 1], 0))));
242
243 let response = self
244 .app
245 .clone()
246 .oneshot(request)
247 .await
248 .expect("Failed to send request");
249
250 let status = response.status();
251 let headers = response.headers().clone();
252
253 for value in headers.get_all(header::SET_COOKIE) {
254 if let Ok(cookie_str) = value.to_str()
255 && let Some(nv) = cookie_str.split(';').next()
256 && let Some((name, val)) = nv.split_once('=')
257 {
258 self.cookies
259 .insert(name.trim().to_string(), val.trim().to_string());
260 }
261 }
262
263 let body_bytes = response
264 .into_body()
265 .collect()
266 .await
267 .expect("Failed to read response body")
268 .to_bytes();
269 let text = String::from_utf8_lossy(&body_bytes).to_string();
270
271 if let Some(token) = extract_csrf_from_html(&text) {
272 self.csrf_token = Some(token);
273 }
274
275 TestResponse {
276 status,
277 text,
278 headers,
279 }
280 }
281
282 async fn send_raw_with_token(
283 &mut self,
284 method: Method,
285 uri: &str,
286 content_type: Option<&str>,
287 body: Option<String>,
288 token: &str,
289 ) -> TestResponse {
290 let body_data = body.unwrap_or_default();
291 let mut builder = Request::builder().method(&method).uri(uri);
292
293 if let Some(ct) = content_type {
294 builder = builder.header(header::CONTENT_TYPE, ct);
295 }
296
297 builder = builder.header("X-CSRF-Token", token);
298
299 if !self.cookies.is_empty() {
300 let cookie_header: String = self
301 .cookies
302 .iter()
303 .map(|(k, v)| format!("{k}={v}"))
304 .collect::<Vec<_>>()
305 .join("; ");
306 builder = builder.header(header::COOKIE, cookie_header);
307 }
308
309 let mut request = builder
310 .body(Body::from(body_data))
311 .expect("Failed to build request");
312 request
313 .extensions_mut()
314 .insert(ConnectInfo(SocketAddr::from(([127, 0, 0, 1], 0))));
315
316 let response = self
317 .app
318 .clone()
319 .oneshot(request)
320 .await
321 .expect("Failed to send request");
322
323 let status = response.status();
324 let headers = response.headers().clone();
325
326 for value in headers.get_all(header::SET_COOKIE) {
327 if let Ok(cookie_str) = value.to_str()
328 && let Some(nv) = cookie_str.split(';').next()
329 && let Some((name, val)) = nv.split_once('=')
330 {
331 self.cookies
332 .insert(name.trim().to_string(), val.trim().to_string());
333 }
334 }
335
336 let body_bytes = response
337 .into_body()
338 .collect()
339 .await
340 .expect("Failed to read response body")
341 .to_bytes();
342 let text = String::from_utf8_lossy(&body_bytes).to_string();
343
344 if let Some(token) = extract_csrf_from_html(&text) {
345 self.csrf_token = Some(token);
346 }
347
348 TestResponse {
349 status,
350 text,
351 headers,
352 }
353 }
354 }
355
356 #[allow(dead_code)]
357 pub(crate) struct TestResponse {
358 pub status: StatusCode,
359 pub text: String,
360 pub headers: axum::http::HeaderMap,
361 }
362
363 #[allow(dead_code)]
364 impl TestResponse {
365 pub(crate) fn json<T: serde::de::DeserializeOwned>(&self) -> T {
366 serde_json::from_str(&self.text)
367 .unwrap_or_else(|e| panic!("Failed to parse JSON: {}\nBody: {}", e, self.text))
368 }
369
370 pub(crate) fn header(&self, name: &str) -> Option<&str> {
371 self.headers.get(name).and_then(|v| v.to_str().ok())
372 }
373 }
374
375 fn extract_csrf_from_html(html: &str) -> Option<String> {
376 let marker = "csrf-token\" content=\"";
377 let start = html.find(marker)? + marker.len();
378 let end = html[start..].find('"')? + start;
379 Some(html[start..end].to_string())
380 }
381