wstunnel/src/socks5.rs

102 lines
3 KiB
Rust
Raw Normal View History

use anyhow::Context;
use fast_socks5::server::{Config, DenyAuthentication, Socks5Server};
use fast_socks5::util::target_addr::TargetAddr;
use futures_util::{stream, Stream, StreamExt};
use std::net::SocketAddr;
use std::pin::Pin;
use std::task::Poll;
use tokio::net::TcpStream;
use log::warn;
use tracing::{info, warn};
use url::Host;
pub struct Socks5Listener {
stream: Pin<Box<dyn Stream<Item = anyhow::Result<(TcpStream, Host, u16)>>>>,
}
impl Stream for Socks5Listener {
type Item = anyhow::Result<(TcpStream, Host, u16)>;
fn poll_next(
self: Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> Poll<Option<Self::Item>> {
unsafe { self.map_unchecked_mut(|x| &mut x.stream) }.poll_next(cx)
}
}
pub async fn run_server(bind: SocketAddr) -> Result<Socks5Listener, anyhow::Error> {
info!("Starting TCP server listening cnx on {}", bind);
let server = Socks5Server::<DenyAuthentication>::bind(bind)
.await
.with_context(|| format!("Cannot create socks5 server {:?}", bind))?;
let mut cfg = Config::<DenyAuthentication>::default();
cfg.set_allow_no_auth(true);
cfg.set_dns_resolve(false);
cfg.set_execute_command(false);
let server = server.with_config(cfg);
let stream = stream::unfold(server, move |server| async {
let mut acceptor = server.incoming();
loop {
let cnx = match acceptor.next().await {
None => return None,
Some(Err(err)) => {
drop(acceptor);
return Some((Err(anyhow::Error::new(err)), server));
}
Some(Ok(cnx)) => cnx,
};
let cnx = match cnx.upgrade_to_socks5().await {
Ok(cnx) => cnx,
Err(err) => {
warn!("Rejecting socks5 cnx: {}", err);
continue;
}
};
let Some(target) = cnx.target_addr() else {
warn!("Rejecting socks5 cnx: no target addr");
continue;
};
let (host, port) = match target {
TargetAddr::Ip(SocketAddr::V4(ip)) => (Host::Ipv4(*ip.ip()), ip.port()),
TargetAddr::Ip(SocketAddr::V6(ip)) => (Host::Ipv6(*ip.ip()), ip.port()),
TargetAddr::Domain(host, port) => (Host::Domain(host.clone()), *port),
};
drop(acceptor);
return Some((Ok((cnx.into_inner(), host, port)), server));
}
});
let listener = Socks5Listener {
stream: Box::pin(stream),
};
Ok(listener)
}
#[cfg(test)]
mod test {
use super::*;
use futures_util::StreamExt;
use std::str::FromStr;
#[tokio::test]
async fn socks5_server() {
let mut x = run_server(SocketAddr::from_str("[::]:4343").unwrap())
.await
.unwrap();
loop {
let cnx = x.next().await.unwrap().unwrap();
eprintln!("{:?}", cnx);
}
}
}