good_mitm_rule/
handler.rs

1use crate::Rule;
2use async_trait::async_trait;
3use hyper::{header, Body, Request, Response};
4use log::info;
5use mitm_core::{
6    handler::{CustomContextData, HttpHandler},
7    mitm::{HttpContext, RequestOrResponse},
8};
9use std::sync::Arc;
10
11#[derive(Clone)]
12pub struct RuleHttpHandler {
13    rules: Arc<Vec<Rule>>,
14}
15
16#[derive(Default, Clone)]
17pub struct RuleHandlerCtx {
18    rules: Vec<Rule>,
19}
20
21impl CustomContextData for RuleHandlerCtx {}
22
23impl RuleHttpHandler {
24    pub fn new(rules: Arc<Vec<Rule>>) -> Self {
25        Self { rules }
26    }
27
28    fn match_rules(&self, req: &Request<Body>) -> Vec<Rule> {
29        let mut matched = vec![];
30        for rule in self.rules.iter() {
31            for filter in &rule.filters {
32                if filter.is_match_req(req) {
33                    matched.push(rule.clone());
34                }
35            }
36        }
37        matched
38    }
39}
40
41#[async_trait]
42impl HttpHandler<RuleHandlerCtx> for RuleHttpHandler {
43    async fn handle_request(
44        &self,
45        ctx: &mut HttpContext<RuleHandlerCtx>,
46        req: Request<Body>,
47    ) -> RequestOrResponse {
48        ctx.uri = Some(req.uri().clone());
49
50        // remove accept-encoding to avoid encoded body
51        let mut req = req;
52        req.headers_mut().remove(header::ACCEPT_ENCODING);
53
54        let rules = self.match_rules(&req);
55        if !rules.is_empty() {
56            ctx.should_modify_response = true;
57        }
58
59        for mut rule in rules {
60            ctx.custom_data.rules.push(rule.clone());
61            let rt = rule.do_req(req).await;
62            if let RequestOrResponse::Request(r) = rt {
63                req = r;
64            } else {
65                return rt;
66            }
67        }
68
69        RequestOrResponse::Request(req)
70    }
71
72    async fn handle_response(
73        &self,
74        ctx: &mut HttpContext<RuleHandlerCtx>,
75        res: Response<Body>,
76    ) -> Response<Body> {
77        if !ctx.should_modify_response || ctx.custom_data.rules.is_empty() {
78            return res;
79        }
80        let uri = ctx.uri.as_ref().unwrap();
81        let content_type = match res.headers().get(header::CONTENT_TYPE) {
82            Some(content_type) => content_type.to_str().unwrap_or_default(),
83            None => "unknown",
84        };
85        info!(
86            "[Response] {} {} {}",
87            res.status(),
88            uri.host().unwrap_or_default(),
89            content_type
90        );
91
92        let mut res = res;
93        for rule in &ctx.custom_data.rules {
94            res = rule.do_res(res).await;
95        }
96        res
97    }
98}
good_mitm_rule/handler.rs

good_mitm_rule/
handler.rs
