mockforge_http/middleware/
response_buffer.rs

1//! Response body buffering middleware
2//!
3//! This middleware buffers response bodies so they can be read multiple times,
4//! enabling downstream middleware to access the response body for analysis.
5
6use axum::{body::Body, extract::Request, http::Response, middleware::Next};
7use serde_json::Value;
8
9/// Buffered response body
10#[derive(Clone)]
11pub struct BufferedResponse {
12    /// Response status
13    pub status: u16,
14    /// Response headers
15    pub headers: axum::http::HeaderMap,
16    /// Response body as bytes
17    pub body: axum::body::Bytes,
18}
19
20impl BufferedResponse {
21    /// Get response body as JSON value
22    pub fn json(&self) -> Option<Value> {
23        serde_json::from_slice(&self.body).ok()
24    }
25
26    /// Get response body as string
27    pub fn text(&self) -> String {
28        String::from_utf8_lossy(&self.body).to_string()
29    }
30}
31
32/// Middleware to buffer response bodies
33///
34/// This middleware reads the entire response body into memory so it can be
35/// accessed multiple times by downstream middleware. The buffered response
36/// is stored in request extensions.
37pub async fn buffer_response_middleware(req: Request, next: Next) -> Response<Body> {
38    // Process request
39    let response = next.run(req).await;
40
41    // Extract response parts
42    let (parts, body) = response.into_parts();
43
44    // Read body into bytes
45    let body_bytes = match axum::body::to_bytes(body, usize::MAX).await {
46        Ok(bytes) => bytes,
47        Err(e) => {
48            tracing::warn!("Failed to buffer response body: {}", e);
49            // Return error response if body buffering fails
50            return Response::builder()
51                .status(axum::http::StatusCode::INTERNAL_SERVER_ERROR)
52                .body(Body::from("Failed to buffer response"))
53                .unwrap();
54        }
55    };
56
57    // Create buffered response
58    let buffered = BufferedResponse {
59        status: parts.status.as_u16(),
60        headers: parts.headers.clone(),
61        body: body_bytes.clone(),
62    };
63
64    // Store in request extensions for downstream middleware
65    // Note: We can't modify request extensions after the response is created,
66    // so we'll store it in a way that can be accessed via a different mechanism
67    // For now, we'll just recreate the response with the buffered body
68
69    // Recreate response with buffered body
70    let mut response_builder = Response::builder().status(parts.status).version(parts.version);
71
72    // Copy headers
73    for (name, value) in parts.headers.iter() {
74        response_builder = response_builder.header(name, value);
75    }
76
77    // Add buffered response to response extensions
78    let mut response = response_builder.body(Body::from(body_bytes)).unwrap();
79
80    // Store buffered response in response extensions
81    response.extensions_mut().insert(buffered);
82
83    response
84}
85
86/// Extract buffered response from response extensions
87pub fn get_buffered_response(response: &Response<Body>) -> Option<BufferedResponse> {
88    response.extensions().get::<BufferedResponse>().cloned()
89}
mockforge_http/middleware/response_buffer.rs

mockforge_http/middleware/
response_buffer.rs
