Skip to main content

mockforge_http/middleware/
drift_tracking.rs

1//! Drift tracking middleware
2//!
3//! This middleware integrates drift budget evaluation and consumer usage tracking
4//! with contract diff analysis.
5
6// Uses DriftBudgetEngine + ContractDiffAnalyzer which stay in core.
7#![allow(deprecated)]
8
9use axum::{body::Body, extract::Request, http::Response, middleware::Next};
10use mockforge_contracts::consumer_contracts::{ConsumerBreakingChangeDetector, UsageRecorder};
11use mockforge_core::{
12    ai_contract_diff::ContractDiffAnalyzer,
13    contract_drift::DriftBudgetEngine,
14    incidents::{IncidentManager, IncidentSeverity, IncidentType},
15    openapi::OpenApiSpec,
16};
17use mockforge_foundation::contract_drift_types::DriftResult;
18use serde_json::Value;
19use std::collections::HashMap;
20use std::sync::Arc;
21use tracing::{debug, warn};
22
23/// Maximum request body size to buffer for drift tracking.
24///
25/// Issue #79 — see `contract_diff_middleware`. Same root cause: the
26/// old hard-coded 1 MiB cap, paired with the over-cap branch
27/// substituting `Body::empty()`, broke downstream handlers on large
28/// chunked uploads. Configurable via `MOCKFORGE_DRIFT_MAX_BODY_MB`.
29fn max_drift_body_size() -> usize {
30    const DEFAULT_MB: usize = 10;
31    std::env::var("MOCKFORGE_DRIFT_MAX_BODY_MB")
32        .ok()
33        .and_then(|v| v.parse::<usize>().ok())
34        .unwrap_or(DEFAULT_MB)
35        .saturating_mul(1024 * 1024)
36}
37
38/// State for drift tracking middleware
39#[derive(Clone)]
40pub struct DriftTrackingState {
41    /// Contract diff analyzer
42    pub diff_analyzer: Option<Arc<ContractDiffAnalyzer>>,
43    /// OpenAPI spec (if available)
44    pub spec: Option<Arc<OpenApiSpec>>,
45    /// Drift budget engine
46    pub drift_engine: Arc<DriftBudgetEngine>,
47    /// Incident manager
48    pub incident_manager: Arc<IncidentManager>,
49    /// Usage recorder for consumer contracts
50    pub usage_recorder: Arc<UsageRecorder>,
51    /// Consumer breaking change detector
52    pub consumer_detector: Arc<ConsumerBreakingChangeDetector>,
53    /// Whether drift tracking is enabled
54    pub enabled: bool,
55}
56
57/// Middleware to track drift and consumer usage (with state from extensions)
58///
59/// This middleware requires response body buffering middleware to be applied first.
60/// The response body is extracted from buffered response extensions.
61pub async fn drift_tracking_middleware_with_extensions(
62    req: Request<Body>,
63    next: Next,
64) -> Response<Body> {
65    // Extract state from request extensions
66    let state = req.extensions().get::<DriftTrackingState>().cloned();
67
68    let state = if let Some(state) = state {
69        state
70    } else {
71        // No state available, skip drift tracking
72        return next.run(req).await;
73    };
74
75    if !state.enabled {
76        return next.run(req).await;
77    }
78
79    let method = req.method().to_string();
80    let path = req.uri().path().to_string();
81    let max_body = max_drift_body_size();
82
83    // Issue #79 — pre-check Content-Length so we skip capture cleanly
84    // for over-cap requests instead of consuming the body and then
85    // substituting `Body::empty()` (which broke downstream handlers).
86    let content_length = req
87        .headers()
88        .get(axum::http::header::CONTENT_LENGTH)
89        .and_then(|v| v.to_str().ok())
90        .and_then(|s| s.parse::<usize>().ok());
91    if let Some(len) = content_length {
92        if len > max_body {
93            debug!(
94                "drift_tracking: skipping capture for {} {} — content-length {} > cap {}",
95                method, path, len, max_body
96            );
97            return next.run(req).await;
98        }
99    }
100
101    // Extract consumer identifier and headers from request
102    let consumer_id = extract_consumer_id(&req);
103
104    // Extract headers for capture
105    let captured_headers = extract_headers_for_capture(&req);
106
107    // Buffer the request body so we can capture it and still forward it
108    let (parts, body) = req.into_parts();
109    let body_bytes = match axum::body::to_bytes(body, max_body).await {
110        Ok(b) => b,
111        Err(_) => {
112            // Chunked over-cap body — body partially consumed and we
113            // cannot rebuild. 413 PayloadTooLarge instead of the old
114            // silent `Body::empty()` substitution. Issue #79.
115            return Response::builder()
116                .status(axum::http::StatusCode::PAYLOAD_TOO_LARGE)
117                .header(
118                    axum::http::header::CONTENT_TYPE,
119                    "application/json",
120                )
121                .body(Body::from(format!(
122                    r#"{{"error":"PAYLOAD_TOO_LARGE","message":"chunked request body exceeded drift_tracking capture cap (~{} MiB); raise MOCKFORGE_DRIFT_MAX_BODY_MB or send Content-Length"}}"#,
123                    max_body / (1024 * 1024)
124                )))
125                .unwrap_or_else(|_| Response::new(Body::from("payload too large")));
126        }
127    };
128
129    // Try to parse body as JSON for structured capture
130    let captured_body = if !body_bytes.is_empty() {
131        serde_json::from_slice::<Value>(&body_bytes).ok()
132    } else {
133        None
134    };
135
136    // Reconstruct the request with the buffered body
137    let req = Request::from_parts(parts, Body::from(body_bytes));
138
139    // Process request and get response
140    let response = next.run(req).await;
141
142    // Extract response body for consumer usage tracking
143    let response_body = extract_response_body(&response);
144
145    // Record consumer usage if consumer is identified
146    if let Some(ref consumer_id) = consumer_id {
147        if let Some(body) = &response_body {
148            state.usage_recorder.record_usage(consumer_id, &path, &method, Some(body)).await;
149        }
150    }
151
152    // Perform contract diff analysis if analyzer and spec are available
153    if let (Some(ref analyzer), Some(ref spec)) = (&state.diff_analyzer, &state.spec) {
154        // Create captured request with body and headers
155        let mut captured = mockforge_core::ai_contract_diff::CapturedRequest::new(
156            &method,
157            &path,
158            "drift_tracking",
159        )
160        .with_headers(captured_headers)
161        .with_response(response.status().as_u16(), response_body.clone());
162
163        if let Some(body_value) = captured_body {
164            captured = captured.with_body(body_value);
165        }
166
167        // Analyze request against contract
168        match analyzer.analyze(&captured, spec).await {
169            Ok(diff_result) => {
170                // Evaluate drift budget
171                let drift_result = state.drift_engine.evaluate(&diff_result, &path, &method);
172
173                // Record contracts pillar usage for drift detection
174                mockforge_core::pillar_tracking::record_contracts_usage(
175                    None, // workspace_id could be extracted from request if available
176                    None,
177                    "drift_detection",
178                    serde_json::json!({
179                        "endpoint": path,
180                        "method": method,
181                        "breaking_changes": drift_result.breaking_changes,
182                        "non_breaking_changes": drift_result.non_breaking_changes,
183                        "incident": drift_result.should_create_incident
184                    }),
185                )
186                .await;
187
188                // If this endpoint has a drift budget configured, record that so
189                // the Contracts pillar dashboard can count distinct budgeted endpoints.
190                let endpoint_key = format!("{} {}", method, path);
191                let budget_config = state.drift_engine.config();
192                if budget_config.enabled
193                    && (budget_config.per_endpoint_budgets.contains_key(&endpoint_key)
194                        || budget_config.default_budget.is_some())
195                {
196                    mockforge_core::pillar_tracking::record_contracts_usage(
197                        None,
198                        None,
199                        "drift_budget_configured",
200                        serde_json::json!({
201                            "endpoint": endpoint_key,
202                        }),
203                    )
204                    .await;
205                }
206
207                // Create incident if budget is exceeded or breaking changes detected
208                if drift_result.should_create_incident {
209                    let incident_type = if drift_result.breaking_changes > 0 {
210                        IncidentType::BreakingChange
211                    } else {
212                        IncidentType::ThresholdExceeded
213                    };
214
215                    let severity = determine_severity(&drift_result);
216
217                    let details = serde_json::json!({
218                        "breaking_changes": drift_result.breaking_changes,
219                        "non_breaking_changes": drift_result.non_breaking_changes,
220                        "breaking_mismatches": drift_result.breaking_mismatches,
221                        "non_breaking_mismatches": drift_result.non_breaking_mismatches,
222                        "budget_exceeded": drift_result.budget_exceeded,
223                    });
224
225                    // Create incident with before/after samples
226                    // Extract before/after samples from diff result if available
227                    let before_sample = Some(serde_json::json!({
228                        "contract_format": diff_result.metadata.contract_format,
229                        "contract_version": diff_result.metadata.contract_version,
230                        "endpoint": path,
231                        "method": method,
232                    }));
233
234                    let after_sample = Some(serde_json::json!({
235                        "mismatches": diff_result.mismatches,
236                        "recommendations": diff_result.recommendations,
237                        "corrections": diff_result.corrections,
238                    }));
239
240                    let _incident = state
241                        .incident_manager
242                        .create_incident_with_samples(
243                            path.clone(),
244                            method.clone(),
245                            incident_type,
246                            severity,
247                            details,
248                            None, // budget_id
249                            None, // workspace_id
250                            None, // sync_cycle_id
251                            None, // contract_diff_id (could be generated from diff_result)
252                            before_sample,
253                            after_sample,
254                            Some(drift_result.fitness_test_results.clone()), // fitness_test_results
255                            drift_result.consumer_impact.clone(),            // affected_consumers
256                            Some(mockforge_foundation::protocol::Protocol::Http), // protocol
257                        )
258                        .await;
259
260                    warn!(
261                        "Drift incident created: {} {} - {} breaking changes, {} non-breaking changes",
262                        method, path, drift_result.breaking_changes, drift_result.non_breaking_changes
263                    );
264                }
265
266                // Check for consumer-specific violations
267                if let Some(ref consumer_id) = consumer_id {
268                    let violations = state
269                        .consumer_detector
270                        .detect_violations(consumer_id, &path, &method, &diff_result, None)
271                        .await;
272
273                    if !violations.is_empty() {
274                        warn!(
275                            "Consumer {} has {} violations on {} {}",
276                            consumer_id,
277                            violations.len(),
278                            method,
279                            path
280                        );
281                    }
282                }
283            }
284            Err(e) => {
285                debug!("Contract diff analysis failed: {}", e);
286            }
287        }
288    }
289
290    response
291}
292
293/// Extract consumer identifier from request
294fn extract_consumer_id(req: &Request<Body>) -> Option<String> {
295    // Try to extract from various sources:
296    // 1. X-Consumer-ID header
297    if let Some(consumer_id) = req.headers().get("x-consumer-id").and_then(|h| h.to_str().ok()) {
298        return Some(consumer_id.to_string());
299    }
300
301    // 2. X-Workspace-ID header (for workspace-based consumers)
302    if let Some(workspace_id) = req.headers().get("x-workspace-id").and_then(|h| h.to_str().ok()) {
303        return Some(format!("workspace:{}", workspace_id));
304    }
305
306    // 3. API key from header
307    if let Some(api_key) = req
308        .headers()
309        .get("x-api-key")
310        .or_else(|| req.headers().get("authorization"))
311        .and_then(|h| h.to_str().ok())
312    {
313        // Hash the API key for privacy
314        use sha2::{Digest, Sha256};
315        let mut hasher = Sha256::new();
316        hasher.update(api_key.as_bytes());
317        let hash = format!("{:x}", hasher.finalize());
318        return Some(format!("api_key:{}", hash));
319    }
320
321    None
322}
323
324/// Extract safe headers for drift capture
325fn extract_headers_for_capture(req: &Request<Body>) -> HashMap<String, String> {
326    let safe_headers = [
327        "accept",
328        "accept-encoding",
329        "content-type",
330        "content-length",
331        "user-agent",
332    ];
333    let mut captured = HashMap::new();
334    for name in safe_headers {
335        if let Some(value) = req.headers().get(name).and_then(|v| v.to_str().ok()) {
336            captured.insert(name.to_string(), value.to_string());
337        }
338    }
339    captured
340}
341
342/// Extract response body as JSON value
343fn extract_response_body(response: &Response<Body>) -> Option<Value> {
344    // Try to get buffered response from extensions
345    if let Some(buffered) = crate::middleware::get_buffered_response(response) {
346        return buffered.json();
347    }
348
349    // If not buffered, try to parse from response body
350    // Note: This requires the response body to be buffered by upstream middleware
351    None
352}
353
354/// Determine incident severity from drift result
355fn determine_severity(drift_result: &DriftResult) -> IncidentSeverity {
356    if drift_result.breaking_changes > 0 {
357        // Check if any breaking mismatch is critical
358        if drift_result
359            .breaking_mismatches
360            .iter()
361            .any(|m| m.severity == mockforge_core::ai_contract_diff::MismatchSeverity::Critical)
362        {
363            return IncidentSeverity::Critical;
364        }
365        // Check if any breaking mismatch is high
366        if drift_result
367            .breaking_mismatches
368            .iter()
369            .any(|m| m.severity == mockforge_core::ai_contract_diff::MismatchSeverity::High)
370        {
371            return IncidentSeverity::High;
372        }
373        return IncidentSeverity::Medium;
374    }
375
376    // Non-breaking changes are lower severity
377    if drift_result.non_breaking_changes > 5 {
378        IncidentSeverity::Medium
379    } else {
380        IncidentSeverity::Low
381    }
382}