use anyhow::Context;
use log::{debug, error};
use socks5_proto::{
Address, Command, Reply, Request as SocksRequest, Response,
handshake::{self, Method},
};
use tokio::{
io::{AsyncWriteExt, copy_bidirectional},
net::TcpStream,
};
use crate::dial::Dial;
use crate::proto::http_connect;
pub struct ProxyConnection {
ts: TcpStream,
dial: Box<dyn Dial>,
}
impl ProxyConnection {
pub fn new(ts: TcpStream, dial: Box<dyn Dial>) -> Self {
Self { ts, dial }
}
pub async fn handle(self) {
let mut first_bit = [0u8];
if let Err(e) = self.ts.peek(&mut first_bit).await {
error!("can't peek first_bit err: {e}");
return;
}
let ret = if first_bit[0] == socks5_proto::SOCKS_VERSION {
self.handle_socks().await
} else {
self.handle_http().await
};
if let Err(e) = ret {
error!("proxy handle err: {e:?}");
};
}
async fn handle_socks(mut self) -> anyhow::Result<()> {
debug!(
"socks proxy connection {:?} to {:?}",
self.ts.peer_addr().ok(),
self.ts.local_addr().ok()
);
let _req = handshake::Request::read_from(&mut self.ts)
.await
.context("socks handshake failed")?;
let resp = handshake::Response::new(Method::NONE);
resp.write_to(&mut self.ts)
.await
.context("socks write response failed")?;
let req = SocksRequest::read_from(&mut self.ts)
.await
.context("socks read request failed")?;
let addr = req.address;
debug!("start connect {addr}");
match req.command {
Command::Connect => {
let target = self.dial.dial(addr.clone()).await;
match target {
Ok(mut target) => {
self.socks_reply(Reply::Succeeded, Address::unspecified())
.await?;
if let Ok((a, b)) = copy_bidirectional(&mut self.ts, &mut target).await {
debug!(
"socks copy end for {} traffic: {}<=>{} total: {}",
addr,
a,
b,
a + b
);
}
Ok(())
}
Err(e) => {
self.socks_reply(Reply::HostUnreachable, Address::unspecified())
.await
.context("socks reply failed.")?;
Err(e).context(format!("socks dial {addr} failed ."))
}
}
}
cmd => {
debug!("socks unsupported command {:?}", cmd);
self.socks_reply(Reply::CommandNotSupported, Address::unspecified())
.await?;
Ok(())
}
}
}
async fn handle_http(mut self) -> anyhow::Result<()> {
debug!(
"http proxy connection {:?} to {:?}",
self.ts.peer_addr().ok(),
self.ts.local_addr().ok()
);
let buf = http_connect::read_http_request_end(&mut self.ts)
.await
.context("http proxy read http request end failed")?;
debug!(
"http proxy read buf: \n{}",
String::from_utf8_lossy(buf.as_slice())
);
match http_connect::HttpConnectRequest::parse(buf.as_slice()) {
Ok(req) => {
let addr = req.addr().clone();
let mut target = self
.dial
.dial(addr.clone())
.await
.context(format!("http proxy connect addr {} failed", addr))?;
if let Some(data) = req.nugget() {
target
.write_all(data.data().as_slice())
.await
.context("http proxy target write_all buf failed")?;
target
.flush()
.await
.context("http proxy flush target failed")?;
} else {
self.ts
.write("HTTP/1.1 200 OK\r\n\r\n".as_bytes())
.await
.context("http proxy write response failed")?;
}
if let Ok((a, b)) = copy_bidirectional(&mut self.ts, &mut target).await {
debug!(
"http copy end for {} traffic: {}<=>{} total: {}",
addr,
a,
b,
a + b
);
};
Ok(())
}
Err(e) => {
debug!("http proxy BAD_REQUEST");
self.ts
.write("HTTP/1.1 400 BAD_REQUEST\r\n\r\n".as_bytes())
.await
.context("http proxy write response failed")?;
Err(e).context("http dial failed .".to_string())
}
}
}
async fn socks_reply(&mut self, reply: Reply, addr: Address) -> anyhow::Result<()> {
let resp = Response::new(reply, addr);
resp.write_to(&mut self.ts)
.await
.context("scoks write reply response failed")
}
}
#[cfg(test)]
mod tests {
use std::sync::{Arc, Mutex};
use anyhow::anyhow;
use async_trait::async_trait;
use socks5_proto::{
Address, Command, Reply, Request as SocksRequest, Response,
handshake::{Method, Request as HandshakeRequest, Response as HandshakeResponse},
};
use tokio::{
io::{AsyncReadExt, AsyncWriteExt, DuplexStream},
net::{TcpListener, TcpStream},
};
use super::ProxyConnection;
use crate::dial::{AsyncStream, Dial};
struct MockDial {
result: Mutex<Option<anyhow::Result<DuplexStream>>>,
seen_addrs: Arc<Mutex<Vec<Address>>>,
}
impl MockDial {
fn succeed(stream: DuplexStream, seen_addrs: Arc<Mutex<Vec<Address>>>) -> Self {
Self {
result: Mutex::new(Some(Ok(stream))),
seen_addrs,
}
}
fn fail(err: anyhow::Error, seen_addrs: Arc<Mutex<Vec<Address>>>) -> Self {
Self {
result: Mutex::new(Some(Err(err))),
seen_addrs,
}
}
}
#[async_trait]
impl Dial for MockDial {
async fn dial(&self, addr: Address) -> anyhow::Result<Box<dyn AsyncStream>> {
self.seen_addrs.lock().unwrap().push(addr);
self.result
.lock()
.unwrap()
.take()
.expect("dial should only be called once")
.map(|s| Box::new(s) as Box<dyn AsyncStream>)
}
}
async fn tcp_pair() -> (TcpStream, TcpStream) {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let client = TcpStream::connect(addr).await.unwrap();
let (server, _) = listener.accept().await.unwrap();
(server, client)
}
#[tokio::test]
async fn handle_http_connect_replies_ok_and_dials_target() {
let (server, mut client) = tcp_pair().await;
let (target, mut target_peer) = tokio::io::duplex(256);
let seen_addrs = Arc::new(Mutex::new(Vec::new()));
let dial = MockDial::succeed(target, seen_addrs.clone());
let proxy = ProxyConnection::new(server, Box::new(dial));
let proxy_task = tokio::spawn(async move { proxy.handle_http().await });
let target_task = tokio::spawn(async move {
let mut buf = Vec::new();
target_peer.read_to_end(&mut buf).await.unwrap();
buf
});
client
.write_all(b"CONNECT example.com:443 HTTP/1.1\r\nHost: example.com:443\r\n\r\n")
.await
.unwrap();
let mut response = [0u8; 19];
client.read_exact(&mut response).await.unwrap();
assert_eq!(&response, b"HTTP/1.1 200 OK\r\n\r\n");
client.shutdown().await.unwrap();
drop(client);
proxy_task.await.unwrap().unwrap();
assert!(target_task.await.unwrap().is_empty());
assert_eq!(
seen_addrs.lock().unwrap().as_slice(),
&[Address::DomainAddress(b"example.com".to_vec(), 443)]
);
}
#[tokio::test]
async fn handle_http_with_nugget_forwards_request_body_to_target() {
let (server, mut client) = tcp_pair().await;
let (target, mut target_peer) = tokio::io::duplex(512);
let seen_addrs = Arc::new(Mutex::new(Vec::new()));
let dial = MockDial::succeed(target, seen_addrs.clone());
let raw = b"GET https://upstream.example/path HTTP/1.1\r\nHost: service.internal\r\n\r\n";
let proxy = ProxyConnection::new(server, Box::new(dial));
let proxy_task = tokio::spawn(async move { proxy.handle_http().await });
let target_task = tokio::spawn(async move {
let mut received = vec![0; raw.len()];
target_peer.read_exact(&mut received).await.unwrap();
target_peer
.write_all(b"HTTP/1.1 204 No Content\r\n\r\n")
.await
.unwrap();
target_peer.shutdown().await.unwrap();
received
});
client.write_all(raw).await.unwrap();
let mut response = vec![0; 27];
client.read_exact(&mut response).await.unwrap();
assert_eq!(response, b"HTTP/1.1 204 No Content\r\n\r\n");
client.shutdown().await.unwrap();
drop(client);
proxy_task.await.unwrap().unwrap();
assert_eq!(target_task.await.unwrap(), raw);
assert_eq!(
seen_addrs.lock().unwrap().as_slice(),
&[Address::DomainAddress(b"service.internal".to_vec(), 443)]
);
}
#[tokio::test]
async fn handle_http_bad_request_returns_400_without_dialing() {
let (server, mut client) = tcp_pair().await;
let seen_addrs = Arc::new(Mutex::new(Vec::new()));
let dial = MockDial::fail(anyhow!("dial should not be called"), seen_addrs.clone());
let proxy = ProxyConnection::new(server, Box::new(dial));
let proxy_task = tokio::spawn(async move { proxy.handle_http().await });
client.write_all(b"BAD\r\n\r\n").await.unwrap();
let mut response = [0u8; 28];
client.read_exact(&mut response).await.unwrap();
assert_eq!(&response, b"HTTP/1.1 400 BAD_REQUEST\r\n\r\n");
client.shutdown().await.unwrap();
drop(client);
assert!(proxy_task.await.unwrap().is_err());
assert!(seen_addrs.lock().unwrap().is_empty());
}
#[tokio::test]
async fn handle_socks_connect_negotiates_and_replies_succeeded() {
let (server, mut client) = tcp_pair().await;
let (target, mut target_peer) = tokio::io::duplex(256);
let seen_addrs = Arc::new(Mutex::new(Vec::new()));
let dial = MockDial::succeed(target, seen_addrs.clone());
let proxy = ProxyConnection::new(server, Box::new(dial));
let proxy_task = tokio::spawn(async move { proxy.handle_socks().await });
let target_task = tokio::spawn(async move {
let mut buf = Vec::new();
target_peer.read_to_end(&mut buf).await.unwrap();
buf
});
HandshakeRequest::new(vec![Method::NONE])
.write_to(&mut client)
.await
.unwrap();
let handshake = HandshakeResponse::read_from(&mut client).await.unwrap();
assert_eq!(handshake.method, Method::NONE);
SocksRequest::new(
Command::Connect,
Address::DomainAddress(b"example.com".to_vec(), 1080),
)
.write_to(&mut client)
.await
.unwrap();
let response = Response::read_from(&mut client).await.unwrap();
assert_eq!(response.reply, Reply::Succeeded);
client.shutdown().await.unwrap();
drop(client);
proxy_task.await.unwrap().unwrap();
assert!(target_task.await.unwrap().is_empty());
assert_eq!(
seen_addrs.lock().unwrap().as_slice(),
&[Address::DomainAddress(b"example.com".to_vec(), 1080)]
);
}
}