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}