slinger_mitm/
interceptor.rs

1//! Traffic interception and modification interfaces
2
3use crate::error::Result;
4use bytes::Bytes;
5use slinger::{Body, Request, Response};
6use std::fmt;
7use std::net::SocketAddr;
8use std::sync::Arc;
9use std::time::{SystemTime, UNIX_EPOCH};
10
11/// MITM Request wrapper that wraps slinger::Request with connection metadata.
12/// Used for both HTTP and non-HTTP (raw TCP) traffic interception.
13#[derive(Clone)]
14pub struct MitmRequest {
15  /// Source address and port (client)
16  pub source: Option<SocketAddr>,
17  /// Destination address (host:port)
18  pub destination: String,
19  /// Timestamp when the request was intercepted
20  pub timestamp: u64,
21  /// Whether this is an HTTP request (true) or raw TCP (false)
22  is_http: bool,
23  /// The underlying request (contains body for both HTTP and raw TCP)
24  pub request: Request,
25}
26
27impl MitmRequest {
28  /// Create a new MITM request wrapper for HTTP traffic
29  pub fn new(destination: impl Into<String>, request: Request) -> Self {
30    Self {
31      source: None,
32      destination: destination.into(),
33      timestamp: SystemTime::now()
34        .duration_since(UNIX_EPOCH)
35        .map(|d| d.as_millis() as u64)
36        .unwrap_or(0),
37      is_http: true,
38      request,
39    }
40  }
41
42  /// Create a new MITM request with source address for HTTP traffic
43  pub fn with_source(source: SocketAddr, destination: impl Into<String>, request: Request) -> Self {
44    Self {
45      source: Some(source),
46      destination: destination.into(),
47      timestamp: SystemTime::now()
48        .duration_since(UNIX_EPOCH)
49        .map(|d| d.as_millis() as u64)
50        .unwrap_or(0),
51      is_http: true,
52      request,
53    }
54  }
55
56  /// Create a MITM request for raw TCP data (non-HTTP)
57  pub fn raw_tcp(destination: impl Into<String>, body: impl Into<Bytes>) -> Self {
58    let request = Request {
59      body: Some(Body::from(body.into())),
60      ..Default::default()
61    };
62    Self {
63      source: None,
64      destination: destination.into(),
65      timestamp: SystemTime::now()
66        .duration_since(UNIX_EPOCH)
67        .map(|d| d.as_millis() as u64)
68        .unwrap_or(0),
69      is_http: false,
70      request,
71    }
72  }
73
74  /// Create a MITM request for raw TCP data with source address
75  pub fn raw_tcp_with_source(
76    source: SocketAddr,
77    destination: impl Into<String>,
78    body: impl Into<Bytes>,
79  ) -> Self {
80    let request = Request {
81      body: Some(Body::from(body.into())),
82      ..Default::default()
83    };
84    Self {
85      source: Some(source),
86      destination: destination.into(),
87      timestamp: SystemTime::now()
88        .duration_since(UNIX_EPOCH)
89        .map(|d| d.as_millis() as u64)
90        .unwrap_or(0),
91      is_http: false,
92      request,
93    }
94  }
95
96  /// Get the source address
97  pub fn source(&self) -> Option<SocketAddr> {
98    self.source
99  }
100
101  /// Get the destination address
102  pub fn destination(&self) -> &str {
103    &self.destination
104  }
105
106  /// Get the timestamp
107  pub fn timestamp(&self) -> u64 {
108    self.timestamp
109  }
110
111  /// Get the underlying request
112  pub fn request(&self) -> &Request {
113    &self.request
114  }
115
116  /// Get a mutable reference to the underlying request
117  pub fn request_mut(&mut self) -> &mut Request {
118    &mut self.request
119  }
120
121  /// Get the body as bytes (for raw TCP traffic)
122  pub fn body(&self) -> Option<&Body> {
123    self.request.body.as_ref()
124  }
125
126  /// Set the body (for raw TCP traffic)
127  pub fn set_body(&mut self, body: impl Into<Bytes>) {
128    self.request.body = Some(Body::from(body.into()));
129  }
130
131  /// Check if this is an HTTP request (true) or raw TCP (false)
132  pub fn is_http(&self) -> bool {
133    self.is_http
134  }
135}
136
137impl fmt::Debug for MitmRequest {
138  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
139    f.debug_struct("MitmRequest")
140      .field("source", &self.source)
141      .field("destination", &self.destination)
142      .field("timestamp", &self.timestamp)
143      .field("is_http", &self.is_http())
144      .field("request", &self.request)
145      .finish()
146  }
147}
148
149/// MITM Response wrapper that wraps slinger::Response with connection metadata.
150/// Used for both HTTP and non-HTTP (raw TCP) traffic interception.
151#[derive(Clone)]
152pub struct MitmResponse {
153  /// Source address (where the response came from, host:port)
154  pub source: String,
155  /// Destination address and port (client)
156  pub destination: Option<SocketAddr>,
157  /// Timestamp when the response was intercepted
158  pub timestamp: u64,
159  /// Whether this is an HTTP response (true) or raw TCP (false)
160  is_http: bool,
161  /// The underlying response (contains body for both HTTP and raw TCP)
162  pub response: Response,
163}
164
165impl MitmResponse {
166  /// Create a new MITM response wrapper for HTTP traffic
167  pub fn new(source: impl Into<String>, response: Response) -> Self {
168    Self {
169      source: source.into(),
170      destination: None,
171      timestamp: SystemTime::now()
172        .duration_since(UNIX_EPOCH)
173        .map(|d| d.as_millis() as u64)
174        .unwrap_or(0),
175      is_http: true,
176      response,
177    }
178  }
179
180  /// Create a new MITM response with destination address for HTTP traffic
181  pub fn with_destination(
182    source: impl Into<String>,
183    destination: SocketAddr,
184    response: Response,
185  ) -> Self {
186    Self {
187      source: source.into(),
188      destination: Some(destination),
189      timestamp: SystemTime::now()
190        .duration_since(UNIX_EPOCH)
191        .map(|d| d.as_millis() as u64)
192        .unwrap_or(0),
193      is_http: true,
194      response,
195    }
196  }
197
198  /// Create a MITM response for raw TCP data (non-HTTP)
199  pub fn raw_tcp(source: impl Into<String>, body: impl Into<Bytes>) -> Self {
200    let response = Response {
201      body: Some(Body::from(body.into())),
202      ..Default::default()
203    };
204    Self {
205      source: source.into(),
206      destination: None,
207      timestamp: SystemTime::now()
208        .duration_since(UNIX_EPOCH)
209        .map(|d| d.as_millis() as u64)
210        .unwrap_or(0),
211      is_http: false,
212      response,
213    }
214  }
215
216  /// Create a MITM response for raw TCP data with destination address
217  pub fn raw_tcp_with_destination(
218    source: impl Into<String>,
219    destination: SocketAddr,
220    body: impl Into<Bytes>,
221  ) -> Self {
222    let response = Response {
223      body: Some(Body::from(body.into())),
224      ..Default::default()
225    };
226    Self {
227      source: source.into(),
228      destination: Some(destination),
229      timestamp: SystemTime::now()
230        .duration_since(UNIX_EPOCH)
231        .map(|d| d.as_millis() as u64)
232        .unwrap_or(0),
233      is_http: false,
234      response,
235    }
236  }
237
238  /// Get the source address
239  pub fn source(&self) -> &str {
240    &self.source
241  }
242
243  /// Get the destination address
244  pub fn destination(&self) -> Option<SocketAddr> {
245    self.destination
246  }
247
248  /// Get the timestamp
249  pub fn timestamp(&self) -> u64 {
250    self.timestamp
251  }
252
253  /// Get the underlying response
254  pub fn response(&self) -> &Response {
255    &self.response
256  }
257
258  /// Get a mutable reference to the underlying response
259  pub fn response_mut(&mut self) -> &mut Response {
260    &mut self.response
261  }
262
263  /// Get the body as bytes (for raw TCP traffic)
264  pub fn body(&self) -> Option<&Body> {
265    self.response.body.as_ref()
266  }
267
268  /// Set the body (for raw TCP traffic)
269  pub fn set_body(&mut self, body: impl Into<Bytes>) {
270    self.response.body = Some(Body::from(body.into()));
271  }
272
273  /// Check if this is an HTTP response (true) or raw TCP (false)
274  pub fn is_http(&self) -> bool {
275    self.is_http
276  }
277}
278
279impl fmt::Debug for MitmResponse {
280  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
281    f.debug_struct("MitmResponse")
282      .field("source", &self.source)
283      .field("destination", &self.destination)
284      .field("timestamp", &self.timestamp)
285      .field("is_http", &self.is_http())
286      .field("response", &self.response)
287      .finish()
288  }
289}
290
291/// Trait for intercepting and modifying requests (both HTTP and raw TCP)
292#[async_trait::async_trait]
293pub trait RequestInterceptor: Send + Sync {
294  /// Intercept and optionally modify a request
295  ///
296  /// Return `None` to block the request, or return a modified request
297  async fn intercept_request(&self, request: MitmRequest) -> Result<Option<MitmRequest>>;
298}
299
300/// Trait for intercepting and modifying responses (both HTTP and raw TCP)
301#[async_trait::async_trait]
302pub trait ResponseInterceptor: Send + Sync {
303  /// Intercept and optionally modify a response
304  ///
305  /// Return `None` to block the response, or return a modified response
306  async fn intercept_response(&self, response: MitmResponse) -> Result<Option<MitmResponse>>;
307}
308
309/// Combined interceptor handler for both HTTP and TCP traffic
310pub struct InterceptorHandler {
311  request_interceptors: Vec<Arc<dyn RequestInterceptor>>,
312  response_interceptors: Vec<Arc<dyn ResponseInterceptor>>,
313}
314
315impl InterceptorHandler {
316  /// Create a new interceptor handler
317  pub fn new() -> Self {
318    Self {
319      request_interceptors: Vec::new(),
320      response_interceptors: Vec::new(),
321    }
322  }
323
324  /// Add a request interceptor
325  pub fn add_request_interceptor(&mut self, interceptor: Arc<dyn RequestInterceptor>) {
326    self.request_interceptors.push(interceptor);
327  }
328
329  /// Add a response interceptor
330  pub fn add_response_interceptor(&mut self, interceptor: Arc<dyn ResponseInterceptor>) {
331    self.response_interceptors.push(interceptor);
332  }
333
334  /// Process a request through all interceptors
335  pub async fn process_request(&self, mut request: MitmRequest) -> Result<Option<MitmRequest>> {
336    for interceptor in &self.request_interceptors {
337      match interceptor.intercept_request(request).await? {
338        Some(modified) => request = modified,
339        None => return Ok(None), // Request blocked
340      }
341    }
342    Ok(Some(request))
343  }
344
345  /// Process a response through all interceptors
346  pub async fn process_response(&self, mut response: MitmResponse) -> Result<Option<MitmResponse>> {
347    for interceptor in &self.response_interceptors {
348      match interceptor.intercept_response(response).await? {
349        Some(modified) => response = modified,
350        None => return Ok(None), // Response blocked
351      }
352    }
353    Ok(Some(response))
354  }
355
356  /// Check if any interceptors are registered
357  pub fn has_interceptors(&self) -> bool {
358    !self.request_interceptors.is_empty() || !self.response_interceptors.is_empty()
359  }
360}
361
362impl Default for InterceptorHandler {
363  fn default() -> Self {
364    Self::new()
365  }
366}
367
368/// Default pass-through interceptor
369pub struct Interceptor;
370
371impl Interceptor {
372  /// Create a logging interceptor that prints requests/responses
373  pub fn logging() -> LoggingInterceptor {
374    LoggingInterceptor
375  }
376}
377
378/// Logging interceptor implementation that handles both HTTP and TCP traffic
379pub struct LoggingInterceptor;
380
381#[async_trait::async_trait]
382impl RequestInterceptor for LoggingInterceptor {
383  async fn intercept_request(&self, request: MitmRequest) -> Result<Option<MitmRequest>> {
384    if request.is_http() {
385      tracing::info!(
386        "[MITM] HTTP Request: {} {}",
387        request.request().method(),
388        request.request().uri()
389      );
390      for (name, value) in request.request().headers() {
391        tracing::info!("  {}: {:?}", name, value);
392      }
393    } else {
394      tracing::info!(
395        "[MITM] TCP Request to {}: {} bytes",
396        request.destination(),
397        request.body().map(|b| b.len()).unwrap_or(0)
398      );
399    }
400    if let Some(source) = request.source() {
401      tracing::info!("  From: {}", source);
402    }
403    tracing::info!("  Timestamp: {}", request.timestamp());
404    Ok(Some(request))
405  }
406}
407
408#[async_trait::async_trait]
409impl ResponseInterceptor for LoggingInterceptor {
410  async fn intercept_response(&self, response: MitmResponse) -> Result<Option<MitmResponse>> {
411    if response.is_http() {
412      tracing::info!(
413        "[MITM] HTTP Response: {}",
414        response.response().status_code()
415      );
416      for (name, value) in response.response().headers() {
417        tracing::info!("  {}: {:?}", name, value);
418      }
419    } else {
420      tracing::info!(
421        "[MITM] TCP Response from {}: {} bytes",
422        response.source(),
423        response.body().map(|b| b.len()).unwrap_or(0)
424      );
425    }
426    if let Some(destination) = response.destination() {
427      tracing::info!("  To: {}", destination);
428    }
429    tracing::info!("  Timestamp: {}", response.timestamp());
430    Ok(Some(response))
431  }
432}