use std::io::ErrorKind;
use std::net::SocketAddr;
use std::sync::Arc;
use bytes::Bytes;
use futures::future::BoxFuture;
use tokio::io::{AsyncRead, AsyncWrite};
use tokio::net::{TcpListener, TcpStream, UdpSocket};
pub use crate::rfc1928::util::target_addr::{TargetAddr, ToTargetAddr};
use crate::{
BindAccept, BindAcceptError, BindAcceptFactory, BindAcceptFactoryError, BindTargetHandle,
DialFn, Egress, EgressError, EgressFactory, SocksDialError, TargetHandle,
};
pub(crate) async fn resolve_target(addr: &TargetAddr) -> Result<SocketAddr, EgressError> {
match addr {
TargetAddr::Ip(sa) => Ok(*sa),
TargetAddr::Domain(domain, port) => {
let addrs = tokio::net::lookup_host((domain.as_str(), *port))
.await
.map_err(|e| EgressError::Unreachable(format!("resolve {domain}: {e}")))?
.next()
.ok_or_else(|| {
EgressError::Unreachable(format!("resolve {domain}: no addresses"))
})?;
Ok(addrs)
}
}
}
pub(crate) fn dial_error_of(what: &str, err: &std::io::Error) -> SocksDialError {
match err.kind() {
ErrorKind::ConnectionRefused => SocksDialError::Refused(format!("{what}: {err}")),
ErrorKind::NetworkUnreachable => {
SocksDialError::NetworkUnreachable(format!("{what}: {err}"))
}
ErrorKind::HostUnreachable => SocksDialError::HostUnreachable(format!("{what}: {err}")),
ErrorKind::TimedOut => SocksDialError::NetworkUnreachable(format!("{what}: {err}")),
ErrorKind::PermissionDenied => SocksDialError::General(format!("{what}: {err}")),
_ => SocksDialError::General(format!("{what}: {err}")),
}
}
pub fn tcp_dial() -> DialFn {
Arc::new(|target: TargetAddr| Box::pin(async move { local_tcp_dial(target).await }))
}
async fn local_tcp_dial(target: TargetAddr) -> Result<TargetHandle, SocksDialError> {
let addrs: Vec<SocketAddr> = match &target {
TargetAddr::Ip(sa) => vec![*sa],
TargetAddr::Domain(domain, port) => tokio::net::lookup_host((domain.as_str(), *port))
.await
.map_err(|e| SocksDialError::HostUnreachable(format!("resolve {domain}: {e}")))?
.collect(),
};
if addrs.is_empty() {
return Err(SocksDialError::HostUnreachable(format!(
"resolve {target}: no addresses"
)));
}
let mut last: Option<SocksDialError> = None;
for addr in addrs {
match TcpStream::connect(addr).await {
Ok(stream) => {
let (read, write) = stream.into_split();
return Ok(TargetHandle {
read: Box::new(read) as Box<dyn AsyncRead + Send + Sync + Unpin>,
write: Box::new(write) as Box<dyn AsyncWrite + Send + Sync + Unpin>,
});
}
Err(err) => last = Some(dial_error_of(&format!("connect {target}"), &err)),
}
}
Err(last.unwrap_or_else(|| {
SocksDialError::General(format!("connect {target}: no addresses tried"))
}))
}
pub fn udp_egress() -> EgressFactory {
Arc::new(|| {
Box::pin(async {
LocalEgress::bind()
.await
.map(|e| Box::new(e) as Box<dyn Egress + Send>)
})
})
}
pub struct LocalEgress {
sock: Arc<UdpSocket>,
}
impl LocalEgress {
pub async fn bind() -> Result<Self, EgressError> {
let sock = UdpSocket::bind("127.0.0.1:0")
.await
.map_err(|e| EgressError::General(format!("udp bind: {e}")))?;
Ok(Self {
sock: Arc::new(sock),
})
}
}
impl Egress for LocalEgress {
fn send(&mut self, addr: TargetAddr, payload: Bytes) -> BoxFuture<'_, Result<(), EgressError>> {
let sock = Arc::clone(&self.sock);
Box::pin(async move {
let dest = resolve_target(&addr).await?;
sock.send_to(&payload, dest)
.await
.map(|_| ())
.map_err(|e| EgressError::General(format!("udp send_to {dest}: {e}")))
})
}
fn next_reply(&mut self) -> BoxFuture<'_, Result<Option<Bytes>, EgressError>> {
let sock = Arc::clone(&self.sock);
Box::pin(async move {
sock.readable()
.await
.map_err(|e| EgressError::General(format!("udp readable: {e}")))?;
let mut scratch = vec![0u8; 65535];
match sock.try_recv_from(&mut scratch) {
Ok((n, from)) => {
let mut inner = TargetAddr::Ip(from)
.to_be_bytes()
.map_err(|e| EgressError::General(format!("addr encode: {e}")))?;
inner.extend_from_slice(&scratch[..n]);
Ok(Some(Bytes::from(inner)))
}
Err(e) if e.kind() == ErrorKind::WouldBlock => Ok(None),
Err(e) => Err(EgressError::General(format!("udp try_recv: {e}"))),
}
})
}
}
pub fn bind_accept(bind_host: impl Into<std::net::IpAddr>) -> BindAcceptFactory {
let bind_host = bind_host.into();
Arc::new(move |target: TargetAddr| {
Box::pin(async move { local_bind_accept(bind_host, &target).await })
})
}
async fn local_bind_accept(
bind_host: std::net::IpAddr,
target: &TargetAddr,
) -> Result<BindAccept, BindAcceptFactoryError> {
let requested_port = match target {
TargetAddr::Ip(addr) => addr.port(),
TargetAddr::Domain(_, port) => *port,
};
let listener = TcpListener::bind(SocketAddr::new(bind_host, requested_port))
.await
.map_err(|e| match e.kind() {
ErrorKind::PermissionDenied | ErrorKind::AddrInUse => {
BindAcceptFactoryError::Shortage(format!("tcp bind: {e}"))
}
ErrorKind::AddrNotAvailable => {
BindAcceptFactoryError::Refused(format!("tcp bind: {e}"))
}
_ => BindAcceptFactoryError::General(format!("tcp bind: {e}")),
})?;
let bind_addr = listener
.local_addr()
.map_err(|e| BindAcceptFactoryError::General(format!("local_addr: {e}")))?;
Ok(BindAccept::new(bind_addr, async move {
let (stream, peer) = listener
.accept()
.await
.map_err(|e| BindAcceptError::General(format!("accept: {e}")))?;
let peer = peer
.to_target_addr()
.map_err(|e| BindAcceptError::General(format!("peer addr: {e}")))?;
let (read, write) = stream.into_split();
Ok(BindTargetHandle {
peer,
read: Box::new(read) as Box<dyn AsyncRead + Send + Sync + Unpin>,
write: Box::new(write) as Box<dyn AsyncWrite + Send + Sync + Unpin>,
})
}))
}
#[cfg(test)]
mod tests;
mod producer_door;
pub use producer_door::{
ProducerDoor, ProducerDoorBuildError, ProducerDoorBuilder, ProducerDoorError, ProducerDoorGuard,
};