Skip to main content

max / makenotwork

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