mockforge_core/proxy/
middleware.rs

1//! Proxy middleware for request/response processing
2
3use crate::Result;
4use std::collections::HashMap;
5
6/// Middleware trait for processing proxy requests
7pub trait ProxyMiddleware {
8    /// Process a request before it's sent to the target
9    fn process_request(
10        &self,
11        method: &str,
12        url: &str,
13        headers: &mut HashMap<String, String>,
14        body: &mut Option<Vec<u8>>,
15    ) -> Result<()>;
16
17    /// Process a response before it's returned to the client
18    fn process_response(
19        &self,
20        status_code: u16,
21        headers: &mut HashMap<String, String>,
22        body: &mut Option<Vec<u8>>,
23    ) -> Result<()>;
24}
25
26/// Logging middleware that logs proxy requests and responses
27pub struct LoggingMiddleware;
28
29impl Default for LoggingMiddleware {
30    fn default() -> Self {
31        Self::new()
32    }
33}
34
35impl LoggingMiddleware {
36    /// Create a new logging middleware
37    pub fn new() -> Self {
38        Self
39    }
40}
41
42impl ProxyMiddleware for LoggingMiddleware {
43    fn process_request(
44        &self,
45        method: &str,
46        url: &str,
47        _headers: &mut HashMap<String, String>,
48        _body: &mut Option<Vec<u8>>,
49    ) -> Result<()> {
50        tracing::info!("Proxy request: {} {}", method, url);
51        Ok(())
52    }
53
54    fn process_response(
55        &self,
56        status_code: u16,
57        _headers: &mut HashMap<String, String>,
58        _body: &mut Option<Vec<u8>>,
59    ) -> Result<()> {
60        tracing::info!("Proxy response: {}", status_code);
61        Ok(())
62    }
63}
64
65/// Header modification middleware
66pub struct HeaderMiddleware {
67    /// Headers to add to requests
68    request_headers: HashMap<String, String>,
69    /// Headers to add to responses
70    response_headers: HashMap<String, String>,
71}
72
73impl Default for HeaderMiddleware {
74    fn default() -> Self {
75        Self::new()
76    }
77}
78
79impl HeaderMiddleware {
80    /// Create a new header middleware
81    pub fn new() -> Self {
82        Self {
83            request_headers: HashMap::new(),
84            response_headers: HashMap::new(),
85        }
86    }
87
88    /// Add a header to outgoing requests
89    pub fn add_request_header(mut self, key: String, value: String) -> Self {
90        self.request_headers.insert(key, value);
91        self
92    }
93
94    /// Add a header to outgoing responses
95    pub fn add_response_header(mut self, key: String, value: String) -> Self {
96        self.response_headers.insert(key, value);
97        self
98    }
99}
100
101impl ProxyMiddleware for HeaderMiddleware {
102    fn process_request(
103        &self,
104        _method: &str,
105        _url: &str,
106        headers: &mut HashMap<String, String>,
107        _body: &mut Option<Vec<u8>>,
108    ) -> Result<()> {
109        for (key, value) in &self.request_headers {
110            headers.insert(key.clone(), value.clone());
111        }
112        Ok(())
113    }
114
115    fn process_response(
116        &self,
117        _status_code: u16,
118        headers: &mut HashMap<String, String>,
119        _body: &mut Option<Vec<u8>>,
120    ) -> Result<()> {
121        for (key, value) in &self.response_headers {
122            headers.insert(key.clone(), value.clone());
123        }
124        Ok(())
125    }
126}