mockforge_http/
op_middleware.rs

1//! Middleware/utilities to apply latency and failure injection per operation.
2use axum::body::Body;
3use axum::extract::State;
4use axum::http::{Request, StatusCode};
5use axum::middleware::Next;
6use axum::response::Response;
7
8use crate::latency_profiles::LatencyProfiles;
9use mockforge_chaos::core_failure_injection::FailureInjector;
10use mockforge_chaos::core_traffic_shaping::TrafficShaper;
11
12/// Metadata for the current OpenAPI operation
13#[derive(Clone)]
14pub struct OperationMeta {
15    /// OpenAPI operation ID
16    pub id: String,
17    /// Tags associated with this operation
18    pub tags: Vec<String>,
19    /// API path pattern
20    pub path: String,
21}
22
23/// Shared state for operation middleware
24#[derive(Clone)]
25pub struct Shared {
26    /// Latency profiles for request simulation
27    pub profiles: LatencyProfiles,
28    /// Optional failure injector for chaos engineering
29    pub failure_injector: Option<FailureInjector>,
30    /// Optional traffic shaper for bandwidth/loss simulation
31    pub traffic_shaper: Option<TrafficShaper>,
32    /// Whether traffic shaping is enabled
33    pub traffic_shaping_enabled: bool,
34}
35
36/// Middleware to add shared state to request extensions
37pub async fn add_shared_extension(
38    State(shared): State<Shared>,
39    mut req: Request<Body>,
40    next: Next,
41) -> Response {
42    req.extensions_mut().insert(shared);
43    next.run(req).await
44}
45
46/// Middleware to apply fault injection before processing request
47pub async fn fault_then_next(req: Request<Body>, next: Next) -> Response {
48    let shared = match req.extensions().get::<Shared>() {
49        Some(s) => s.clone(),
50        None => {
51            tracing::error!("Shared extension not found in request - ensure add_shared_extension middleware is configured");
52            let mut res =
53                Response::new(Body::from("Internal server error: middleware misconfiguration"));
54            *res.status_mut() = StatusCode::INTERNAL_SERVER_ERROR;
55            return res;
56        }
57    };
58    let op = req.extensions().get::<OperationMeta>().cloned();
59
60    // First, check the new enhanced failure injection system
61    if let Some(failure_injector) = &shared.failure_injector {
62        let tags = op.as_ref().map(|o| o.tags.as_slice()).unwrap_or(&[]);
63        if let Some((status_code, error_message)) = failure_injector.process_request(tags) {
64            let mut res = Response::new(Body::from(error_message));
65            *res.status_mut() =
66                StatusCode::from_u16(status_code).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
67            return res;
68        }
69    }
70
71    // Fallback to legacy latency profiles system for backward compatibility
72    if let Some(op) = &op {
73        if let Some((code, msg)) = shared
74            .profiles
75            .maybe_fault(&op.id, &op.tags.iter().map(|s| s.to_string()).collect::<Vec<_>>())
76            .await
77        {
78            let mut res = Response::new(Body::from(msg));
79            *res.status_mut() =
80                StatusCode::from_u16(code).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
81            return res;
82        }
83    }
84
85    // Apply traffic shaping (bandwidth throttling and burst loss) to the request
86    if shared.traffic_shaping_enabled {
87        if let Some(traffic_shaper) = &shared.traffic_shaper {
88            // Calculate request size for bandwidth throttling
89            let request_size = calculate_request_size(&req);
90
91            let tags = op.as_ref().map(|o| o.tags.as_slice()).unwrap_or(&[]);
92
93            // Apply traffic shaping
94            match traffic_shaper.process_transfer(request_size, tags).await {
95                Ok(Some(_timeout)) => {
96                    // Request was "lost" due to burst loss - return timeout error
97                    let mut res =
98                        Response::new(Body::from("Request timeout due to traffic shaping"));
99                    *res.status_mut() = StatusCode::REQUEST_TIMEOUT;
100                    return res;
101                }
102                Ok(None) => {
103                    // Transfer allowed, continue
104                }
105                Err(e) => {
106                    // Traffic shaping error - return internal server error
107                    let mut res =
108                        Response::new(Body::from(format!("Traffic shaping error: {}", e)));
109                    *res.status_mut() = StatusCode::INTERNAL_SERVER_ERROR;
110                    return res;
111                }
112            }
113        }
114    }
115
116    let (parts, body) = req.into_parts();
117    let req = Request::from_parts(parts, body);
118
119    let response = next.run(req).await;
120
121    // Apply traffic shaping to the response
122    if shared.traffic_shaping_enabled {
123        if let Some(traffic_shaper) = &shared.traffic_shaper {
124            // Calculate response size for bandwidth throttling
125            let response_size = calculate_response_size(&response);
126
127            let tags = op.as_ref().map(|o| o.tags.as_slice()).unwrap_or(&[]);
128
129            // Apply traffic shaping to response
130            match traffic_shaper.process_transfer(response_size, tags).await {
131                Ok(Some(_timeout)) => {
132                    // Response was "lost" due to burst loss - return timeout error
133                    let mut res =
134                        Response::new(Body::from("Response timeout due to traffic shaping"));
135                    *res.status_mut() = StatusCode::GATEWAY_TIMEOUT;
136                    return res;
137                }
138                Ok(None) => {
139                    // Transfer allowed, continue
140                }
141                Err(e) => {
142                    // Traffic shaping error - return internal server error
143                    let mut res =
144                        Response::new(Body::from(format!("Traffic shaping error: {}", e)));
145                    *res.status_mut() = StatusCode::INTERNAL_SERVER_ERROR;
146                    return res;
147                }
148            }
149        }
150    }
151
152    response
153}
154
155/// Calculate the approximate size of an HTTP request for bandwidth throttling
156fn calculate_request_size<B>(req: &Request<B>) -> u64 {
157    let mut size = 0u64;
158
159    // Add header sizes
160    for (name, value) in req.headers() {
161        size += name.as_str().len() as u64;
162        size += value.as_bytes().len() as u64;
163    }
164
165    // Add URI size
166    size += req.uri().to_string().len() as u64;
167
168    // Use Content-Length header for body size when available
169    if let Some(content_length) = req.headers().get(http::header::CONTENT_LENGTH) {
170        if let Ok(len_str) = content_length.to_str() {
171            if let Ok(len) = len_str.parse::<u64>() {
172                size += len;
173                return size;
174            }
175        }
176    }
177
178    // Fallback: estimate body size from method (GET/HEAD/DELETE typically have no body)
179    let method = req.method();
180    if method == http::Method::POST || method == http::Method::PUT || method == http::Method::PATCH
181    {
182        size += 256; // Conservative estimate for requests without Content-Length
183    }
184
185    size
186}
187
188/// Calculate the approximate size of an HTTP response for bandwidth throttling
189fn calculate_response_size(res: &Response) -> u64 {
190    let mut size = 0u64;
191
192    // Add header sizes
193    for (name, value) in res.headers() {
194        size += name.as_str().len() as u64;
195        size += value.as_bytes().len() as u64;
196    }
197
198    // Add status line size
199    size += 15; // "HTTP/1.1 200 OK\r\n"
200
201    // Use Content-Length header for body size when available
202    if let Some(content_length) = res.headers().get(http::header::CONTENT_LENGTH) {
203        if let Ok(len_str) = content_length.to_str() {
204            if let Ok(len) = len_str.parse::<u64>() {
205                size += len;
206                return size;
207            }
208        }
209    }
210
211    // Fallback: estimate from status code (204/304 have no body)
212    match res.status().as_u16() {
213        204 | 304 => {}   // No body
214        _ => size += 256, // Conservative estimate for responses without Content-Length
215    }
216
217    size
218}
219
220#[cfg(test)]
221mod tests {
222    use super::*;
223    use axum::http::{Request, Response, StatusCode};
224
225    #[test]
226    fn test_operation_meta_creation() {
227        let meta = OperationMeta {
228            id: "getUserById".to_string(),
229            tags: vec!["users".to_string(), "public".to_string()],
230            path: "/users/{id}".to_string(),
231        };
232
233        assert_eq!(meta.id, "getUserById");
234        assert_eq!(meta.tags.len(), 2);
235        assert_eq!(meta.path, "/users/{id}");
236    }
237
238    #[test]
239    fn test_shared_creation() {
240        let shared = Shared {
241            profiles: LatencyProfiles::default(),
242            failure_injector: None,
243            traffic_shaper: None,
244            traffic_shaping_enabled: false,
245        };
246
247        assert!(!shared.traffic_shaping_enabled);
248        assert!(shared.failure_injector.is_none());
249        assert!(shared.traffic_shaper.is_none());
250    }
251
252    #[test]
253    fn test_shared_with_failure_injector() {
254        let failure_injector = FailureInjector::new(None, true);
255        let shared = Shared {
256            profiles: LatencyProfiles::default(),
257            failure_injector: Some(failure_injector),
258            traffic_shaper: None,
259            traffic_shaping_enabled: false,
260        };
261
262        assert!(shared.failure_injector.is_some());
263    }
264
265    #[test]
266    fn test_calculate_request_size_basic() {
267        let req = Request::builder()
268            .uri("/test")
269            .header("content-type", "application/json")
270            .body(())
271            .unwrap();
272
273        let size = calculate_request_size(&req);
274
275        // Should be > 0 (includes headers + URI + body estimate)
276        assert!(size > 0);
277        // Should include at least the URI and header sizes
278        assert!(size >= "/test".len() as u64 + "content-type".len() as u64);
279    }
280
281    #[test]
282    fn test_calculate_request_size_with_multiple_headers() {
283        let req = Request::builder()
284            .uri("/api/users")
285            .header("content-type", "application/json")
286            .header("authorization", "Bearer token123")
287            .header("user-agent", "test-client")
288            .body(())
289            .unwrap();
290
291        let size = calculate_request_size(&req);
292
293        // Should account for all headers (URI + header names + header values)
294        assert!(size > 50);
295    }
296
297    #[test]
298    fn test_calculate_response_size_basic() {
299        let res = Response::builder()
300            .status(StatusCode::OK)
301            .header("content-type", "application/json")
302            .body(Body::empty())
303            .unwrap();
304
305        let size = calculate_response_size(&res);
306
307        // Should be > 0 (includes status line + headers + body estimate)
308        assert!(size > 0);
309        // Should include at least the status line estimate (50) and header sizes
310        assert!(size >= 50);
311    }
312
313    #[test]
314    fn test_calculate_response_size_with_multiple_headers() {
315        let res = Response::builder()
316            .status(StatusCode::OK)
317            .header("content-type", "application/json")
318            .header("cache-control", "no-cache")
319            .header("x-request-id", "123-456-789")
320            .body(Body::empty())
321            .unwrap();
322
323        let size = calculate_response_size(&res);
324
325        // Should account for all headers
326        assert!(size > 100);
327    }
328
329    #[test]
330    fn test_shared_clone() {
331        let shared = Shared {
332            profiles: LatencyProfiles::default(),
333            failure_injector: None,
334            traffic_shaper: None,
335            traffic_shaping_enabled: true,
336        };
337
338        let cloned = shared.clone();
339
340        assert_eq!(shared.traffic_shaping_enabled, cloned.traffic_shaping_enabled);
341    }
342
343    #[test]
344    fn test_operation_meta_clone() {
345        let meta = OperationMeta {
346            id: "testOp".to_string(),
347            tags: vec!["tag1".to_string()],
348            path: "/test".to_string(),
349        };
350
351        let cloned = meta.clone();
352
353        assert_eq!(meta.id, cloned.id);
354        assert_eq!(meta.tags, cloned.tags);
355        assert_eq!(meta.path, cloned.path);
356    }
357}
mockforge_http/op_middleware.rs

mockforge_http/
op_middleware.rs
