portforwarder/
connection_plugin.rs

1use crate::address_matcher::IpAddrMatcher;
2use hex;
3use regex::Regex;
4use std::collections::HashMap;
5use std::net::{SocketAddr, ToSocketAddrs};
6
7pub trait ConnectionPlugin {
8    fn onlySingleTarget(&self) -> Option<SocketAddr>;
9    fn decideTarget(&self, buf: &[u8], addr: SocketAddr) -> Option<SocketAddr>;
10    fn testipaddr(&self, addr: &SocketAddr) -> bool;
11    fn transform(&mut self, buf: &[u8]) -> Option<Vec<u8>>;
12}
13
14pub struct RegexMultiplexer {
15    utarget: Option<SocketAddr>,
16    rules: Vec<(Box<dyn Fn(&[u8]) -> bool + Send + Sync>, SocketAddr)>,
17    ip_matcher: IpAddrMatcher,
18}
19
20fn build_kmp_table(pattern: &[u8]) -> Vec<usize> {
21    let mut table = vec![0; pattern.len()];
22    let mut i = 1;
23    let mut j = 0;
24
25    while i < pattern.len() {
26        if pattern[i] == pattern[j] {
27            j += 1;
28            table[i] = j;
29            i += 1;
30        } else {
31            if j != 0 {
32                j = table[j - 1];
33            } else {
34                table[i] = 0;
35                i += 1;
36            }
37        }
38    }
39
40    table
41}
42
43fn kmp_search(text: &[u8], pattern: &[u8]) -> Option<usize> {
44    if pattern.is_empty() {
45        return Some(0);
46    }
47
48    let table = build_kmp_table(pattern);
49    let mut i = 0;
50    let mut j = 0;
51
52    while i < text.len() {
53        if text[i] == pattern[j] {
54            i += 1;
55            j += 1;
56
57            if j == pattern.len() {
58                return Some(i - j);
59            }
60        } else {
61            if j != 0 {
62                j = table[j - 1];
63            } else {
64                i += 1;
65            }
66        }
67    }
68
69    None
70}
71
72impl From<(Vec<(String, String)>, Vec<String>)> for RegexMultiplexer {
73    fn from(regexPlusAllowed: (Vec<(String, String)>, Vec<String>)) -> Self {
74        let proto2regex: HashMap<&str, &str> = vec![
75            ("[ssh]", "^SSH-2\\.0-.+"),
76            (
77                "[http]",
78                "^(GET|POST|PUT|DELETE|OPTIONS|HEAD|CONNECT|TRACE).*HTTP.*",
79            ),
80        ]
81        .into_iter()
82        .collect();
83
84        let rules = regexPlusAllowed
85            .0
86            .iter()
87            .map(|pair| {
88                let gexp = match proto2regex.get(&pair.0.as_str()) {
89                    Some(re) => *re,
90                    None => &pair.0,
91                };
92
93                let addr = pair.1.to_socket_addrs().unwrap().next().unwrap();
94                if gexp == "[socks5]" {
95                    let func: Box<dyn Fn(&[u8]) -> bool + Send + Sync> =
96                        Box::new(|buf: &[u8]| {
97                            if buf.len() < 3
98                                || buf[0] != 0x05
99                                || buf.len() != usize::from(buf[1]) + 2
100                            {
101                                false
102                            } else {
103                                let mut is_valid = true;
104                                for octet_ in &buf[2..] {
105                                    let octet = *octet_;
106                                    if octet != 0
107                                        && octet != 1
108                                        && octet != 2
109                                        && octet != 3
110                                        && octet != 0x80
111                                        && octet != 0xFF
112                                    {
113                                        is_valid = false;
114                                        break;
115                                    }
116                                }
117                                is_valid
118                            }
119                        });
120                    return (func, addr);
121                } else if gexp == "[rdp]" {
122                    let func: Box<dyn Fn(&[u8]) -> bool + Send + Sync> =
123                        Box::new(|buf: &[u8]| {
124                            if buf.len() < 11 || buf[0] != 0x03 {
125                                return false;
126                            }
127
128                            let length = u16::from_be_bytes([buf[2], buf[3]]) as usize;
129                            if length != buf.len() {
130                                return false;
131                            }
132
133                            if (buf[4] as usize + 5) != buf.len() {
134                                return false;
135                            }
136
137                            // connection request
138                            if buf[5] & 0xE0 != 0xE0 {
139                                return false;
140                            }
141
142                            // DST-REF
143                            if buf[6] != 0 || buf[7] != 0 {
144                                return false;
145                            }
146
147                            true
148                        });
149                    return (func, addr);
150                } else if gexp.starts_with("[https:") && gexp.ends_with("]") {
151                    let domain_name = gexp[7..gexp.len() - 1].to_string();
152                    let func: Box<dyn Fn(&[u8]) -> bool + Send + Sync> =
153                        Box::new(move |buf: &[u8]| {
154                            if buf.len() < 3 + domain_name.len() {
155                                return false;
156                            }
157                            if buf[0] != 0x16 && buf[1] != 0x03 && buf[2] != 0x01 {
158                                return false;
159                            }
160                            return kmp_search(buf, domain_name.as_bytes()).is_some();
161                        });
162                    (func, addr)
163                } else {
164                    let exp = if gexp.starts_with("[http:") && gexp.ends_with("]") {
165                        let domain_name = &gexp[6..gexp.len() - 1];
166                        "^(GET|POST|PUT|DELETE|OPTIONS|HEAD|CONNECT|TRACE).*HTTP.*(.\r\n.*)*"
167                            .to_string()
168                            + domain_name
169                    } else {
170                        gexp.to_string()
171                    };
172
173                    let regex = Regex::new(&exp).unwrap();
174                    let func: Box<dyn Fn(&[u8]) -> bool + Send + Sync> =
175                        Box::new(move |buf: &[u8]| {
176                            let s1 = hex::encode(buf);
177                            if regex.is_match(&s1) {
178                                return true;
179                            }
180                            let s2 = String::from_utf8_lossy(buf);
181                            if regex.is_match(&s2) {
182                                return true;
183                            }
184
185                            return false;
186                        });
187                    return (func, addr);
188                }
189            })
190            .collect();
191        let ip_matcher = IpAddrMatcher::from(&regexPlusAllowed.1);
192        let utarget =
193            if regexPlusAllowed.0.len() == 1 && regexPlusAllowed.0.get(0).unwrap().0 == ".*" {
194                Some(
195                    regexPlusAllowed
196                        .0
197                        .get(0)
198                        .unwrap()
199                        .1
200                        .to_socket_addrs()
201                        .unwrap()
202                        .next()
203                        .unwrap(),
204                )
205            } else {
206                None
207            };
208        RegexMultiplexer {
209            utarget,
210            rules,
211            ip_matcher,
212        }
213    }
214}
215
216impl ConnectionPlugin for RegexMultiplexer {
217    fn onlySingleTarget(&self) -> Option<SocketAddr> {
218        return self.utarget;
219    }
220
221    fn decideTarget(&self, buf: &[u8], _addr: SocketAddr) -> Option<SocketAddr> {
222        for rule in &self.rules {
223            if rule.0(&buf) {
224                return Some(rule.1);
225            }
226        }
227        None
228    }
229
230    fn testipaddr(&self, addr: &SocketAddr) -> bool {
231        self.ip_matcher.testipaddr(&addr.ip())
232    }
233
234    fn transform(&mut self, _: &[u8]) -> Option<Vec<u8>> {
235        None
236    }
237}