Skip to main content

kairos_rs/middleware/
transform.rs

1use actix_web::http::{header::{HeaderMap, HeaderName, HeaderValue}, StatusCode};
2use regex::Regex;
3use serde::{Deserialize, Serialize};
4use std::collections::HashMap;
5use std::str::FromStr;
6
7/// Transformation action type for request/response modification.
8/// 
9/// Defines the type of transformation to apply to headers, paths, or other request/response components.
10/// 
11/// # Actions
12/// 
13/// - **Add**: Add a new header or parameter (won't override existing)
14/// - **Set**: Set a header or parameter (overrides if exists)
15/// - **Remove**: Remove a header or parameter
16/// - **Replace**: Replace using regex pattern matching
17#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
18#[serde(rename_all = "snake_case")]
19pub enum TransformAction {
20    /// Add a value without overriding existing
21    Add,
22    /// Set a value (override if exists)
23    Set,
24    /// Remove a value
25    Remove,
26    /// Replace using regex pattern
27    Replace,
28}
29
30/// Header transformation rule.
31///
32/// Defines how to transform HTTP headers in requests or responses.
33/// 
34/// # Examples
35/// 
36/// ```rust
37/// use kairos_rs::middleware::transform::{HeaderTransformation, TransformAction};
38/// 
39/// // Add a custom header
40/// let add_header = HeaderTransformation {
41///     action: TransformAction::Add,
42///     name: "X-Custom-Header".to_string(),
43///     value: Some("custom-value".to_string()),
44///     pattern: None,
45///     replacement: None,
46/// };
47/// 
48/// // Remove authorization header
49/// let remove_auth = HeaderTransformation {
50///     action: TransformAction::Remove,
51///     name: "Authorization".to_string(),
52///     value: None,
53///     pattern: None,
54///     replacement: None,
55/// };
56/// 
57/// // Replace header value using regex
58/// let replace_header = HeaderTransformation {
59///     action: TransformAction::Replace,
60///     name: "User-Agent".to_string(),
61///     value: None,
62///     pattern: Some(r"(\d+\.\d+)".to_string()),
63///     replacement: Some("v$1-proxy".to_string()),
64/// };
65/// ```
66#[derive(Debug, Clone, Serialize, Deserialize)]
67pub struct HeaderTransformation {
68    /// Action to perform
69    pub action: TransformAction,
70    
71    /// Header name
72    pub name: String,
73    
74    /// Value to set/add (required for Add and Set)
75    pub value: Option<String>,
76    
77    /// Regex pattern for Replace action
78    pub pattern: Option<String>,
79    
80    /// Replacement template for Replace action
81    pub replacement: Option<String>,
82}
83
84/// Path transformation rule.
85///
86/// Defines how to rewrite request paths before forwarding to backend.
87/// Supports regex-based pattern matching with capture groups and template-based replacement.
88/// 
89/// # Examples
90/// 
91/// ```rust
92/// use kairos_rs::middleware::transform::PathTransformation;
93/// 
94/// // Rewrite /api/v1/users/:id to /users/:id
95/// let path_rewrite = PathTransformation {
96///     pattern: r"^/api/v1/(.+)$".to_string(),
97///     replacement: "/$1".to_string(),
98/// };
99/// 
100/// // Add prefix to all paths
101/// let add_prefix = PathTransformation {
102///     pattern: r"^(.+)$".to_string(),
103///     replacement: "/v2$1".to_string(),
104/// };
105/// ```
106#[derive(Debug, Clone, Serialize, Deserialize)]
107pub struct PathTransformation {
108    /// Regex pattern to match against the path
109    pub pattern: String,
110    
111    /// Replacement template with capture group support ($1, $2, etc.)
112    pub replacement: String,
113}
114
115/// Query parameter transformation rule.
116///
117/// Defines how to transform URL query parameters before forwarding requests.
118/// 
119/// # Examples
120/// 
121/// ```rust
122/// use kairos_rs::middleware::transform::{QueryTransformation, TransformAction};
123/// 
124/// // Add API key parameter
125/// let add_key = QueryTransformation {
126///     action: TransformAction::Add,
127///     name: "api_key".to_string(),
128///     value: Some("secret123".to_string()),
129/// };
130/// 
131/// // Remove debug parameter
132/// let remove_debug = QueryTransformation {
133///     action: TransformAction::Remove,
134///     name: "debug".to_string(),
135///     value: None,
136/// };
137/// ```
138#[derive(Debug, Clone, Serialize, Deserialize)]
139pub struct QueryTransformation {
140    /// Action to perform
141    pub action: TransformAction,
142    
143    /// Parameter name
144    pub name: String,
145    
146    /// Value to set/add (required for Add and Set)
147    pub value: Option<String>,
148}
149
150/// Request transformation configuration.
151///
152/// Complete configuration for transforming incoming requests before forwarding to backends.
153/// Supports header manipulation, path rewriting, and query parameter transformation.
154/// 
155/// # Examples
156/// 
157/// ```rust
158/// use kairos_rs::middleware::transform::{RequestTransformation, HeaderTransformation, PathTransformation, TransformAction};
159/// 
160/// let transform = RequestTransformation {
161///     headers: vec![
162///         HeaderTransformation {
163///             action: TransformAction::Add,
164///             name: "X-Forwarded-By".to_string(),
165///             value: Some("kairos-gateway".to_string()),
166///             pattern: None,
167///             replacement: None,
168///         }
169///     ],
170///     path: Some(PathTransformation {
171///         pattern: r"^/api/v1/(.+)$".to_string(),
172///         replacement: "/$1".to_string(),
173///     }),
174///     query_params: vec![],
175/// };
176/// ```
177#[derive(Debug, Clone, Serialize, Deserialize, Default)]
178pub struct RequestTransformation {
179    /// Header transformations to apply
180    #[serde(default)]
181    pub headers: Vec<HeaderTransformation>,
182    
183    /// Optional path transformation
184    pub path: Option<PathTransformation>,
185    
186    /// Query parameter transformations
187    #[serde(default)]
188    pub query_params: Vec<QueryTransformation>,
189}
190
191/// Response transformation configuration.
192///
193/// Configuration for transforming backend responses before sending to clients.
194/// Supports header manipulation, status code mapping, and body transformation.
195/// 
196/// # Examples
197/// 
198/// ```rust
199/// use kairos_rs::middleware::transform::{ResponseTransformation, HeaderTransformation, StatusCodeMapping, TransformAction};
200/// use actix_web::http::StatusCode;
201/// 
202/// let transform = ResponseTransformation {
203///     headers: vec![
204///         HeaderTransformation {
205///             action: TransformAction::Remove,
206///             name: "Server".to_string(),
207///             value: None,
208///             pattern: None,
209///             replacement: None,
210///         }
211///     ],
212///     status_code_mappings: vec![
213///         StatusCodeMapping {
214///             from: StatusCode::NOT_FOUND,
215///             to: StatusCode::OK,
216///             condition: Some("path == '/health'".to_string()),
217///         }
218///     ],
219/// };
220/// ```
221#[derive(Debug, Clone, Serialize, Deserialize, Default)]
222pub struct ResponseTransformation {
223    /// Header transformations to apply
224    #[serde(default)]
225    pub headers: Vec<HeaderTransformation>,
226    
227    /// Status code mappings
228    #[serde(default)]
229    pub status_code_mappings: Vec<StatusCodeMapping>,
230}
231
232/// Status code mapping for response transformation.
233///
234/// Maps specific status codes to different codes based on optional conditions.
235/// Useful for normalizing error responses or masking backend errors.
236/// 
237/// # Examples
238/// 
239/// ```rust
240/// use kairos_rs::middleware::transform::StatusCodeMapping;
241/// use actix_web::http::StatusCode;
242/// 
243/// // Map 404 to 200 for health checks
244/// let mapping = StatusCodeMapping {
245///     from: StatusCode::NOT_FOUND,
246///     to: StatusCode::OK,
247///     condition: Some("path == '/health'".to_string()),
248/// };
249/// ```
250#[derive(Debug, Clone, Serialize, Deserialize)]
251pub struct StatusCodeMapping {
252    /// Status code to match
253    #[serde(with = "status_code_serde")]
254    pub from: StatusCode,
255    
256    /// Status code to replace with
257    #[serde(with = "status_code_serde")]
258    pub to: StatusCode,
259    
260    /// Optional condition expression (future: support simple conditions)
261    pub condition: Option<String>,
262}
263
264/// Custom serialization for StatusCode
265mod status_code_serde {
266    use actix_web::http::StatusCode;
267    use serde::{Deserialize, Deserializer, Serializer};
268
269    pub fn serialize<S>(status: &StatusCode, serializer: S) -> Result<S::Ok, S::Error>
270    where
271        S: Serializer,
272    {
273        serializer.serialize_u16(status.as_u16())
274    }
275
276    pub fn deserialize<'de, D>(deserializer: D) -> Result<StatusCode, D::Error>
277    where
278        D: Deserializer<'de>,
279    {
280        let code = u16::deserialize(deserializer)?;
281        StatusCode::from_u16(code).map_err(serde::de::Error::custom)
282    }
283}
284
285/// Request transformer service.
286///
287/// Applies transformation rules to incoming requests before forwarding to backends.
288/// 
289/// # Examples
290/// 
291/// ```rust
292/// use kairos_rs::middleware::transform::{RequestTransformer, RequestTransformation, HeaderTransformation, TransformAction};
293/// use actix_web::test::TestRequest;
294/// 
295/// let config = RequestTransformation {
296///     headers: vec![
297///         HeaderTransformation {
298///             action: TransformAction::Add,
299///             name: "X-Custom".to_string(),
300///             value: Some("test".to_string()),
301///             pattern: None,
302///             replacement: None,
303///         }
304///     ],
305///     path: None,
306///     query_params: vec![],
307/// };
308/// 
309/// let transformer = RequestTransformer::new(config);
310/// ```
311pub struct RequestTransformer {
312    config: RequestTransformation,
313    path_regex: Option<Regex>,
314}
315
316impl RequestTransformer {
317    /// Creates a new request transformer with the given configuration.
318    /// 
319    /// # Arguments
320    /// 
321    /// * `config` - Request transformation configuration
322    /// 
323    /// # Examples
324    /// 
325    /// ```rust
326    /// use kairos_rs::middleware::transform::{RequestTransformer, RequestTransformation};
327    /// 
328    /// let config = RequestTransformation::default();
329    /// let transformer = RequestTransformer::new(config);
330    /// ```
331    pub fn new(config: RequestTransformation) -> Self {
332        let path_regex = config.path.as_ref().and_then(|p| {
333            Regex::new(&p.pattern).ok()
334        });
335
336        Self {
337            config,
338            path_regex,
339        }
340    }
341
342    /// Transforms request headers according to configuration.
343    /// 
344    /// # Arguments
345    /// 
346    /// * `headers` - Mutable reference to request headers
347    pub fn transform_headers(&self, headers: &mut HeaderMap) {
348        for transform in &self.config.headers {
349            match transform.action {
350                TransformAction::Add | TransformAction::Set => {
351                    if let Some(value) = &transform.value {
352                        if let (Ok(name), Ok(val)) = (
353                            HeaderName::from_str(&transform.name),
354                            HeaderValue::from_str(value),
355                        ) {
356                            if transform.action == TransformAction::Add {
357                                if !headers.contains_key(&name) {
358                                    headers.insert(name, val);
359                                }
360                            } else {
361                                headers.insert(name, val);
362                            }
363                        }
364                    }
365                }
366                TransformAction::Remove => {
367                    if let Ok(name) = HeaderName::from_str(&transform.name) {
368                        headers.remove(&name);
369                    }
370                }
371                TransformAction::Replace => {
372                    if let (Some(pattern), Some(replacement)) =
373                        (&transform.pattern, &transform.replacement)
374                    {
375                        if let Ok(regex) = Regex::new(pattern) {
376                            if let Ok(name) = HeaderName::from_str(&transform.name) {
377                                if let Some(value) = headers.get(&name) {
378                                    if let Ok(value_str) = value.to_str() {
379                                        let new_value = regex.replace_all(value_str, replacement.as_str());
380                                        if let Ok(new_val) = HeaderValue::from_str(&new_value) {
381                                            headers.insert(name, new_val);
382                                        }
383                                    }
384                                }
385                            }
386                        }
387                    }
388                }
389            }
390        }
391    }
392
393    /// Transforms request path according to configuration.
394    /// 
395    /// # Arguments
396    /// 
397    /// * `path` - Original request path
398    /// 
399    /// # Returns
400    /// 
401    /// Transformed path string
402    pub fn transform_path(&self, path: &str) -> String {
403        if let (Some(regex), Some(path_config)) = (&self.path_regex, &self.config.path) {
404            regex.replace(path, path_config.replacement.as_str()).to_string()
405        } else {
406            path.to_string()
407        }
408    }
409
410    /// Transforms query parameters according to configuration.
411    /// 
412    /// # Arguments
413    /// 
414    /// * `params` - Mutable reference to query parameters map
415    pub fn transform_query_params(&self, params: &mut HashMap<String, String>) {
416        for transform in &self.config.query_params {
417            match transform.action {
418                TransformAction::Add => {
419                    if let Some(value) = &transform.value {
420                        params.entry(transform.name.clone()).or_insert(value.clone());
421                    }
422                }
423                TransformAction::Set => {
424                    if let Some(value) = &transform.value {
425                        params.insert(transform.name.clone(), value.clone());
426                    }
427                }
428                TransformAction::Remove => {
429                    params.remove(&transform.name);
430                }
431                TransformAction::Replace => {
432                    // Query param replace not implemented yet
433                }
434            }
435        }
436    }
437}
438
439/// Response transformer service.
440///
441/// Applies transformation rules to backend responses before sending to clients.
442/// 
443/// # Examples
444/// 
445/// ```rust
446/// use kairos_rs::middleware::transform::{ResponseTransformer, ResponseTransformation};
447/// 
448/// let config = ResponseTransformation::default();
449/// let transformer = ResponseTransformer::new(config);
450/// ```
451pub struct ResponseTransformer {
452    config: ResponseTransformation,
453}
454
455impl ResponseTransformer {
456    /// Creates a new response transformer with the given configuration.
457    /// 
458    /// # Arguments
459    /// 
460    /// * `config` - Response transformation configuration
461    pub fn new(config: ResponseTransformation) -> Self {
462        Self { config }
463    }
464
465    /// Transforms response headers according to configuration.
466    /// 
467    /// # Arguments
468    /// 
469    /// * `headers` - Mutable reference to response headers
470    pub fn transform_headers(&self, headers: &mut HeaderMap) {
471        for transform in &self.config.headers {
472            match transform.action {
473                TransformAction::Add | TransformAction::Set => {
474                    if let Some(value) = &transform.value {
475                        if let (Ok(name), Ok(val)) = (
476                            HeaderName::from_str(&transform.name),
477                            HeaderValue::from_str(value),
478                        ) {
479                            if transform.action == TransformAction::Add {
480                                if !headers.contains_key(&name) {
481                                    headers.insert(name, val);
482                                }
483                            } else {
484                                headers.insert(name, val);
485                            }
486                        }
487                    }
488                }
489                TransformAction::Remove => {
490                    if let Ok(name) = HeaderName::from_str(&transform.name) {
491                        headers.remove(&name);
492                    }
493                }
494                TransformAction::Replace => {
495                    if let (Some(pattern), Some(replacement)) =
496                        (&transform.pattern, &transform.replacement)
497                    {
498                        if let Ok(regex) = Regex::new(pattern) {
499                            if let Ok(name) = HeaderName::from_str(&transform.name) {
500                                if let Some(value) = headers.get(&name) {
501                                    if let Ok(value_str) = value.to_str() {
502                                        let new_value = regex.replace_all(value_str, replacement.as_str());
503                                        if let Ok(new_val) = HeaderValue::from_str(&new_value) {
504                                            headers.insert(name, new_val);
505                                        }
506                                    }
507                                }
508                            }
509                        }
510                    }
511                }
512            }
513        }
514    }
515
516    /// Transforms status code according to configuration mappings.
517    /// 
518    /// # Arguments
519    /// 
520    /// * `status` - Original status code
521    /// * `path` - Request path (for condition evaluation)
522    /// 
523    /// # Returns
524    /// 
525    /// Transformed status code
526    pub fn transform_status_code(&self, status: StatusCode, _path: &str) -> StatusCode {
527        for mapping in &self.config.status_code_mappings {
528            if status == mapping.from {
529                // TODO: Evaluate condition if present
530                // For now, always apply if status matches
531                return mapping.to;
532            }
533        }
534        status
535    }
536}
537
538#[cfg(test)]
539mod tests {
540    use super::*;
541    use actix_web::http::header::HeaderName;
542
543    #[test]
544    fn test_header_transformation_add() {
545        let config = RequestTransformation {
546            headers: vec![HeaderTransformation {
547                action: TransformAction::Add,
548                name: "X-Custom".to_string(),
549                value: Some("test-value".to_string()),
550                pattern: None,
551                replacement: None,
552            }],
553            path: None,
554            query_params: vec![],
555        };
556
557        let transformer = RequestTransformer::new(config);
558        let mut headers = HeaderMap::new();
559        
560        transformer.transform_headers(&mut headers);
561        
562        assert_eq!(
563            headers.get("X-Custom").unwrap().to_str().unwrap(),
564            "test-value"
565        );
566    }
567
568    #[test]
569    fn test_header_transformation_remove() {
570        let config = RequestTransformation {
571            headers: vec![HeaderTransformation {
572                action: TransformAction::Remove,
573                name: "Authorization".to_string(),
574                value: None,
575                pattern: None,
576                replacement: None,
577            }],
578            path: None,
579            query_params: vec![],
580        };
581
582        let transformer = RequestTransformer::new(config);
583        let mut headers = HeaderMap::new();
584        headers.insert(
585            HeaderName::from_static("authorization"),
586            HeaderValue::from_static("Bearer token")
587        );
588        
589        transformer.transform_headers(&mut headers);
590        
591        assert!(headers.get("Authorization").is_none());
592    }
593
594    #[test]
595    fn test_path_transformation() {
596        let config = RequestTransformation {
597            headers: vec![],
598            path: Some(PathTransformation {
599                pattern: r"^/api/v1/(.+)$".to_string(),
600                replacement: "/$1".to_string(),
601            }),
602            query_params: vec![],
603        };
604
605        let transformer = RequestTransformer::new(config);
606        let transformed = transformer.transform_path("/api/v1/users/123");
607        
608        assert_eq!(transformed, "/users/123");
609    }
610
611    #[test]
612    fn test_query_param_transformation() {
613        let config = RequestTransformation {
614            headers: vec![],
615            path: None,
616            query_params: vec![
617                QueryTransformation {
618                    action: TransformAction::Add,
619                    name: "api_key".to_string(),
620                    value: Some("secret123".to_string()),
621                },
622                QueryTransformation {
623                    action: TransformAction::Remove,
624                    name: "debug".to_string(),
625                    value: None,
626                },
627            ],
628        };
629
630        let transformer = RequestTransformer::new(config);
631        let mut params = HashMap::new();
632        params.insert("debug".to_string(), "true".to_string());
633        
634        transformer.transform_query_params(&mut params);
635        
636        assert!(params.contains_key("api_key"));
637        assert!(!params.contains_key("debug"));
638    }
639
640    #[test]
641    fn test_status_code_mapping() {
642        let config = ResponseTransformation {
643            headers: vec![],
644            status_code_mappings: vec![StatusCodeMapping {
645                from: StatusCode::NOT_FOUND,
646                to: StatusCode::OK,
647                condition: None,
648            }],
649        };
650
651        let transformer = ResponseTransformer::new(config);
652        let mapped = transformer.transform_status_code(StatusCode::NOT_FOUND, "/test");
653        
654        assert_eq!(mapped, StatusCode::OK);
655    }
656}