nu_test_support/net/
proxy.rs

1use std::{
2    collections::HashMap,
3    hash::Hash,
4    io, mem,
5    net::{SocketAddr, TcpStream},
6    sync::{Arc, Mutex, mpsc},
7    thread::{self, JoinHandle},
8};
9
10use gatekeeper::{
11    Address, ServerConfig,
12    connector::{Connector, TcpUdpConnector},
13};
14
15pub struct Socks5Proxy {
16    addr: SocketAddr,
17    th: Option<JoinHandle<Result<(), gatekeeper::error::Error>>>,
18    tx: mpsc::Sender<gatekeeper::ServerCommand<TcpStream>>,
19}
20
21impl Socks5Proxy {
22    pub fn builder() -> io::Result<Socks5ProxyBuilder> {
23        let addr = nu_utils::net::reserve_local_addr()?;
24        let config = ServerConfig {
25            server_ip: addr.ip(),
26            server_port: addr.port(),
27            ..Default::default()
28        };
29        Ok(Socks5ProxyBuilder {
30            config,
31            redirects: HashMap::new(),
32        })
33    }
34
35    pub fn spawn() -> io::Result<Self> {
36        Self::builder()?.spawn()
37    }
38
39    fn spawn_from_builder(builder: Socks5ProxyBuilder) -> io::Result<Self> {
40        let Socks5ProxyBuilder { config, redirects } = builder;
41
42        let (tx_done, rx_done) = mpsc::sync_channel(1);
43        let (mut server, tx) = gatekeeper::Server::with_binder(
44            config.clone(),
45            gatekeeper::acceptor::TcpBinder::new(
46                config.client_rw_timeout,
47                Arc::new(Mutex::new(rx_done)),
48                config.accept_timeout,
49            ),
50            tx_done,
51            RedirectingTcpConnector::new(redirects),
52        );
53
54        let th = thread::spawn(move || server.serve());
55
56        let addr = SocketAddr::new(config.server_ip, config.server_port);
57        Ok(Self {
58            addr,
59            th: Some(th),
60            tx,
61        })
62    }
63
64    pub fn addr(&self) -> SocketAddr {
65        self.addr
66    }
67
68    pub fn uri(&self) -> String {
69        format!("socks5://{}", self.addr)
70    }
71}
72
73impl Drop for Socks5Proxy {
74    fn drop(&mut self) {
75        let _ = self.tx.send(gatekeeper::ServerCommand::Terminate);
76        let _ = self.th.take().map(|th| th.join());
77    }
78}
79
80#[derive(Debug, Clone)]
81pub struct Socks5ProxyBuilder {
82    config: ServerConfig,
83    redirects: HashMap<HashableAddress, Address>,
84}
85
86#[derive(Debug, Clone, Eq, PartialEq)]
87struct HashableAddress(Address);
88
89impl Hash for HashableAddress {
90    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
91        match &self.0 {
92            addr @ Address::IpAddr(ip_addr, port) => {
93                mem::discriminant(addr).hash(state);
94                ip_addr.hash(state);
95                port.hash(state);
96            }
97            addr @ Address::Domain(domain, port) => {
98                mem::discriminant(addr).hash(state);
99                domain.hash(state);
100                port.hash(state);
101            }
102        }
103    }
104}
105
106impl Socks5ProxyBuilder {
107    pub fn add_redirect(mut self, from: Address, to: Address) -> Self {
108        self.redirects.insert(HashableAddress(from), to);
109        self
110    }
111
112    pub fn spawn(self) -> io::Result<Socks5Proxy> {
113        Socks5Proxy::spawn_from_builder(self)
114    }
115}
116
117#[derive(Debug, Clone)]
118struct RedirectingTcpConnector {
119    redirects: HashMap<HashableAddress, Address>,
120    connector: TcpUdpConnector,
121}
122
123impl RedirectingTcpConnector {
124    pub fn new(redirects: HashMap<HashableAddress, Address>) -> Self {
125        Self {
126            redirects,
127            connector: TcpUdpConnector::new(None),
128        }
129    }
130}
131
132impl Connector for RedirectingTcpConnector {
133    type B = <TcpUdpConnector as Connector>::B;
134    type P = <TcpUdpConnector as Connector>::P;
135
136    fn connect_byte_stream(
137        &self,
138        addr: Address,
139    ) -> Result<(Self::B, SocketAddr), gatekeeper::model::Error> {
140        let addr = HashableAddress(addr);
141        match self.redirects.get(&addr) {
142            Some(addr) => self.connector.connect_byte_stream(addr.clone()),
143            None => self.connector.connect_byte_stream(addr.0),
144        }
145    }
146
147    fn connect_pkt_stream(
148        &self,
149        _addr: Address,
150    ) -> Result<(Self::P, SocketAddr), gatekeeper::model::Error> {
151        unimplemented!("only supports tcp")
152    }
153}
nu_test_support/net/proxy.rs

nu_test_support/net/
proxy.rs
