Skip to main content

icebox/core/proxy/
tcp.rs

1use tokio::io::AsyncReadExt;
2use tokio::io::AsyncWriteExt;
3use tokio::net::{TcpListener, TcpStream};
4use tracing::{info, warn};
5
6use super::{NetworkIsolator, ProxyListener};
7use crate::modules::hex_encode;
8use std::net::SocketAddr;
9
10pub struct TcpProxyIsolator;
11
12#[async_trait::async_trait]
13impl NetworkIsolator for TcpProxyIsolator {
14    async fn setup(&self) -> anyhow::Result<()> {
15        Ok(())
16    }
17
18    async fn teardown(&self) -> anyhow::Result<()> {
19        Ok(())
20    }
21
22    async fn spawn_proxy(
23        &self,
24        target_ip: &str,
25        target_port: u16,
26    ) -> anyhow::Result<(ProxyListener, tokio::task::JoinHandle<()>)> {
27        let listener = TcpListener::bind("127.0.0.1:0").await?;
28        let local_addr = listener.local_addr()?;
29        let target_str = format!("{}:{}", target_ip, target_port);
30        let target_addr: SocketAddr = tokio::net::lookup_host(&target_str)
31            .await?
32            .next()
33            .ok_or_else(|| anyhow::anyhow!("could not resolve hostname: {}", target_ip))?;
34
35        let handle = tokio::spawn(async move {
36            loop {
37                if let Ok((mut client_stream, _client_addr)) = listener.accept().await {
38                    info!("Proxy accepted connection for {}", target_addr);
39                    match tokio::time::timeout(
40                        std::time::Duration::from_secs(5),
41                        TcpStream::connect(target_addr),
42                    )
43                    .await
44                    {
45                        Ok(Ok(mut target_stream)) => {
46                            tokio::spawn(async move {
47                                let (mut cr, mut cw) = client_stream.split();
48                                let (mut tr, mut tw) = target_stream.split();
49
50                                let client_to_target = async {
51                                    let mut buf = vec![0u8; 8192];
52                                    loop {
53                                        match cr.read(&mut buf).await {
54                                            Ok(0) => break,
55                                            Ok(n) => {
56                                                info!(
57                                                    "Proxy C->S {} bytes: {}",
58                                                    n,
59                                                    hex_encode(&buf[..n.min(32)])
60                                                );
61                                                if tw.write_all(&buf[..n]).await.is_err() {
62                                                    break;
63                                                }
64                                            }
65                                            Err(_) => break,
66                                        }
67                                    }
68                                };
69
70                                let target_to_client = async {
71                                    let mut buf = vec![0u8; 8192];
72                                    loop {
73                                        match tr.read(&mut buf).await {
74                                            Ok(0) => break,
75                                            Ok(n) => {
76                                                info!(
77                                                    "Proxy S->C {} bytes: {}",
78                                                    n,
79                                                    hex_encode(&buf[..n.min(32)])
80                                                );
81                                                if cw.write_all(&buf[..n]).await.is_err() {
82                                                    break;
83                                                }
84                                            }
85                                            Err(_) => break,
86                                        }
87                                    }
88                                };
89
90                                tokio::select! {
91                                    _ = client_to_target => {}
92                                    _ = target_to_client => {}
93                                }
94                            });
95                        }
96                        Ok(Err(e)) => {
97                            warn!("Proxy failed to connect to target {}: {}", target_addr, e)
98                        }
99                        Err(_) => warn!("Proxy connect timeout to target {}", target_addr),
100                    }
101                }
102            }
103        });
104
105        Ok((
106            ProxyListener {
107                local_addr,
108                target_addr,
109            },
110            handle,
111        ))
112    }
113}