Skip to main content

mockforge_proxy/
body_transform.rs

1//! Body transformation middleware for proxy requests and responses
2//!
3//! This module provides functionality to transform request and response bodies
4//! using JSONPath expressions and template expansion. Useful for browser proxy
5//! mode where you want to inspect and replace values in intercepted traffic.
6
7use crate::config::{BodyTransform, BodyTransformRule, TransformOperation};
8use mockforge_core::templating::TemplateEngine;
9use mockforge_core::Result;
10use serde_json::Value;
11use tracing::{debug, error, warn};
12
13/// Body transformation middleware that applies JSONPath-based transformations
14pub struct BodyTransformationMiddleware {
15    /// Request transformation rules
16    request_rules: Vec<BodyTransformRule>,
17    /// Response transformation rules
18    response_rules: Vec<BodyTransformRule>,
19    /// Template engine for expanding template tokens
20    template_engine: TemplateEngine,
21}
22
23impl BodyTransformationMiddleware {
24    /// Create a new body transformation middleware
25    pub fn new(
26        request_rules: Vec<BodyTransformRule>,
27        response_rules: Vec<BodyTransformRule>,
28    ) -> Self {
29        Self {
30            request_rules,
31            response_rules,
32            template_engine: TemplateEngine::new(),
33        }
34    }
35
36    /// Transform a request body based on configured rules
37    pub fn transform_request_body(&self, url: &str, body: &mut Option<Vec<u8>>) -> Result<()> {
38        if body.is_none() || self.request_rules.is_empty() {
39            return Ok(());
40        }
41
42        // Find matching rules for this URL
43        let matching_rules: Vec<&BodyTransformRule> =
44            self.request_rules.iter().filter(|rule| rule.matches_url(url)).collect();
45
46        if matching_rules.is_empty() {
47            return Ok(());
48        }
49
50        // Try to parse as JSON
51        let body_str = match String::from_utf8(body.as_ref().unwrap().clone()) {
52            Ok(s) => s,
53            Err(_) => {
54                // Not UTF-8, skip transformation
55                return Ok(());
56            }
57        };
58
59        // Try to parse as JSON
60        let mut json: Value = match serde_json::from_str(&body_str) {
61            Ok(v) => v,
62            Err(_) => {
63                // Not JSON, skip transformation
64                debug!("Request body is not JSON, skipping transformation");
65                return Ok(());
66            }
67        };
68
69        // Apply all matching rules
70        for rule in matching_rules {
71            if let Err(e) = self.apply_transform_rule(&mut json, rule) {
72                warn!("Failed to apply request transformation rule: {}", e);
73            }
74        }
75
76        // Serialize back to bytes
77        let new_body = serde_json::to_vec(&json)?;
78        *body = Some(new_body);
79
80        Ok(())
81    }
82
83    /// Transform a response body based on configured rules
84    pub fn transform_response_body(
85        &self,
86        url: &str,
87        status_code: u16,
88        body: &mut Option<Vec<u8>>,
89    ) -> Result<()> {
90        if body.is_none() || self.response_rules.is_empty() {
91            return Ok(());
92        }
93
94        // Find matching rules for this URL
95        let matching_rules: Vec<&BodyTransformRule> = self
96            .response_rules
97            .iter()
98            .filter(|rule| rule.matches_url(url) && rule.matches_status_code(status_code))
99            .collect();
100
101        if matching_rules.is_empty() {
102            return Ok(());
103        }
104
105        // Try to parse as JSON
106        let body_str = match String::from_utf8(body.as_ref().unwrap().clone()) {
107            Ok(s) => s,
108            Err(_) => {
109                // Not UTF-8, skip transformation
110                return Ok(());
111            }
112        };
113
114        // Try to parse as JSON
115        let mut json: Value = match serde_json::from_str(&body_str) {
116            Ok(v) => v,
117            Err(_) => {
118                // Not JSON, skip transformation
119                debug!("Response body is not JSON, skipping transformation");
120                return Ok(());
121            }
122        };
123
124        // Apply all matching rules
125        for rule in matching_rules {
126            if let Err(e) = self.apply_transform_rule(&mut json, rule) {
127                warn!("Failed to apply response transformation rule: {}", e);
128            }
129        }
130
131        // Serialize back to bytes
132        let new_body = serde_json::to_vec(&json)?;
133        *body = Some(new_body);
134
135        Ok(())
136    }
137
138    /// Apply a single transformation rule to JSON
139    fn apply_transform_rule(&self, json: &mut Value, rule: &BodyTransformRule) -> Result<()> {
140        for transform in &rule.body_transforms {
141            match self.apply_single_transform(json, transform) {
142                Ok(_) => {
143                    debug!("Applied transformation: {} -> {}", transform.path, transform.replace);
144                }
145                Err(e) => {
146                    error!("Failed to apply transformation {}: {}", transform.path, e);
147                    // Continue with other transforms even if one fails
148                }
149            }
150        }
151        Ok(())
152    }
153
154    /// Apply a single transform to JSON using JSONPath
155    /// Uses a simplified path-based approach for common JSONPath expressions
156    fn apply_single_transform(&self, json: &mut Value, transform: &BodyTransform) -> Result<()> {
157        // For now, use the simplified path-based approach
158        // Full JSONPath support can be added later if needed
159        self.apply_single_transform_simple(json, transform)
160    }
161}
162
163// Simplified implementation that works with direct path access
164impl BodyTransformationMiddleware {
165    /// Apply a single transform using a simplified path-based approach
166    /// This works for simple paths like "$.field" or "$.field.subfield"
167    fn apply_single_transform_simple(
168        &self,
169        json: &mut Value,
170        transform: &BodyTransform,
171    ) -> Result<()> {
172        // Expand template in replacement value
173        let replacement_value = self.template_engine.expand_str(&transform.replace);
174
175        // Parse replacement value as JSON if possible, otherwise use as string
176        let replacement_json: Value = match serde_json::from_str(&replacement_value) {
177            Ok(v) => v,
178            Err(_) => Value::String(replacement_value.clone()),
179        };
180
181        // Extract path components (simplified - supports $.field.subfield)
182        let path = transform.path.trim_start_matches("$.");
183        let parts: Vec<&str> = path.split('.').collect();
184
185        if parts.is_empty() {
186            return Err(mockforge_core::Error::internal("Empty JSONPath".to_string()));
187        }
188
189        // Navigate to the target location
190        let mut current = json;
191        for (i, part) in parts.iter().enumerate() {
192            let is_last = i == parts.len() - 1;
193
194            if is_last {
195                // Apply the transformation
196                // Check type first to avoid multiple mutable borrows
197                match transform.operation {
198                    TransformOperation::Replace => {
199                        match current {
200                            Value::Object(ref mut obj) => {
201                                obj.insert(part.to_string(), replacement_json.clone());
202                            }
203                            Value::Array(ref mut arr) => {
204                                if let Ok(idx) = part.parse::<usize>() {
205                                    if idx < arr.len() {
206                                        arr[idx] = replacement_json.clone();
207                                    }
208                                }
209                            }
210                            _ => {}
211                        }
212                        // Break early since we're done
213                        break;
214                    }
215                    TransformOperation::Add => {
216                        if let Value::Object(ref mut obj) = current {
217                            obj.insert(part.to_string(), replacement_json.clone());
218                        }
219                        // Break early since we're done
220                        break;
221                    }
222                    TransformOperation::Remove => {
223                        match current {
224                            Value::Object(ref mut obj) => {
225                                // part is &str, which String can Borrow from
226                                obj.remove(*part);
227                            }
228                            Value::Array(ref mut arr) => {
229                                if let Ok(idx) = part.parse::<usize>() {
230                                    if idx < arr.len() {
231                                        arr.remove(idx);
232                                    }
233                                }
234                            }
235                            _ => {}
236                        }
237                        // Break early since we're done
238                        break;
239                    }
240                }
241            } else {
242                // Navigate deeper - use match to avoid borrow conflicts
243                match current {
244                    Value::Object(ref mut obj) => {
245                        current = obj
246                            .entry(part.to_string())
247                            .or_insert_with(|| Value::Object(serde_json::Map::new()));
248                    }
249                    Value::Array(ref mut arr) => {
250                        if let Ok(idx) = part.parse::<usize>() {
251                            if idx < arr.len() {
252                                current = &mut arr[idx];
253                            } else {
254                                return Err(mockforge_core::Error::internal(format!(
255                                    "Array index {} out of bounds",
256                                    idx
257                                )));
258                            }
259                        } else {
260                            return Err(mockforge_core::Error::internal(format!(
261                                "Invalid array index: {}",
262                                part
263                            )));
264                        }
265                    }
266                    _ => {
267                        // Create intermediate objects as needed
268                        *current = Value::Object(serde_json::Map::new());
269                        if let Value::Object(ref mut obj) = current {
270                            current = obj
271                                .entry(part.to_string())
272                                .or_insert_with(|| Value::Object(serde_json::Map::new()));
273                        }
274                    }
275                }
276            }
277        }
278
279        Ok(())
280    }
281}