use std::future::Future;
use std::net::SocketAddr;
use std::sync::Arc;
use bytes::Bytes;
use futures::future::BoxFuture;
use tokio::io::{AsyncRead, AsyncWrite};
pub use crate::rfc1928::util::target_addr::TargetAddr;
use crate::rfc1928::ReplyError;
pub struct TargetHandle {
pub read: Box<dyn AsyncRead + Send + Sync + Unpin>,
pub write: Box<dyn AsyncWrite + Send + Sync + Unpin>,
}
impl std::fmt::Debug for TargetHandle {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("TargetHandle").finish_non_exhaustive()
}
}
pub type DialFn = Arc<
dyn Fn(TargetAddr) -> BoxFuture<'static, Result<TargetHandle, SocksDialError>> + Send + Sync,
>;
#[derive(Debug, thiserror::Error)]
pub enum SocksDialError {
#[error("connection refused: {0}")]
Refused(String),
#[error("network unreachable: {0}")]
NetworkUnreachable(String),
#[error("host unreachable: {0}")]
HostUnreachable(String),
#[error("connection not allowed by ruleset: {0}")]
NotAllowed(String),
#[error("general failure: {0}")]
General(String),
}
impl SocksDialError {
pub fn to_reply_error(&self) -> ReplyError {
match self {
Self::Refused(_) => ReplyError::ConnectionRefused,
Self::NetworkUnreachable(_) => ReplyError::NetworkUnreachable,
Self::HostUnreachable(_) => ReplyError::HostUnreachable,
Self::NotAllowed(_) => ReplyError::ConnectionNotAllowed,
Self::General(_) => ReplyError::GeneralFailure,
}
}
}
#[derive(Debug, thiserror::Error)]
pub enum EgressError {
#[error("egress refused: {0}")]
Refused(String),
#[error("egress unreachable: {0}")]
Unreachable(String),
#[error("egress general failure: {0}")]
General(String),
}
impl EgressError {
pub fn to_reply_error(&self) -> ReplyError {
match self {
Self::Refused(_) => ReplyError::ConnectionRefused,
Self::Unreachable(_) => ReplyError::NetworkUnreachable,
Self::General(_) => ReplyError::GeneralFailure,
}
}
}
pub trait Egress {
fn send(&mut self, addr: TargetAddr, payload: Bytes) -> BoxFuture<'_, Result<(), EgressError>>;
fn next_reply(&mut self) -> BoxFuture<'_, Result<Option<Bytes>, EgressError>>;
}
pub type EgressFactory =
Arc<dyn Fn() -> BoxFuture<'static, Result<Box<dyn Egress + Send>, EgressError>> + Send + Sync>;
pub fn egress_factory<F>(f: F) -> EgressFactory
where
F: Fn() -> Result<Box<dyn Egress + Send>, EgressError> + Send + Sync + 'static,
{
Arc::new(move || {
let result = f();
Box::pin(async move { result })
})
}
pub struct BindTargetHandle {
pub peer: TargetAddr,
pub read: Box<dyn AsyncRead + Send + Sync + Unpin>,
pub write: Box<dyn AsyncWrite + Send + Sync + Unpin>,
}
impl std::fmt::Debug for BindTargetHandle {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("BindTargetHandle")
.field("peer", &self.peer)
.finish_non_exhaustive()
}
}
pub struct BindAccept {
bind_addr: SocketAddr,
accept: BoxFuture<'static, Result<BindTargetHandle, BindAcceptError>>,
}
impl BindAccept {
pub fn new(
bind_addr: SocketAddr,
accept: impl Future<Output = Result<BindTargetHandle, BindAcceptError>> + Send + 'static,
) -> Self {
Self {
bind_addr,
accept: Box::pin(accept),
}
}
pub fn bind_addr(&self) -> SocketAddr {
self.bind_addr
}
pub async fn accept(self) -> Result<BindTargetHandle, BindAcceptError> {
self.accept.await
}
}
impl std::fmt::Debug for BindAccept {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("BindAccept")
.field("bind_addr", &self.bind_addr)
.finish_non_exhaustive()
}
}
#[derive(Debug, thiserror::Error)]
pub enum BindAcceptError {
#[error("accept refused: {0}")]
Refused(String),
#[error("accept not allowed: {0}")]
NotAllowed(String),
#[error("accept general failure: {0}")]
General(String),
}
impl BindAcceptError {
pub fn to_reply_error(&self) -> ReplyError {
match self {
Self::Refused(_) => ReplyError::ConnectionRefused,
Self::NotAllowed(_) => ReplyError::ConnectionNotAllowed,
Self::General(_) => ReplyError::GeneralFailure,
}
}
}
#[derive(Debug, thiserror::Error)]
pub enum BindAcceptFactoryError {
#[error("bind refused: {0}")]
Refused(String),
#[error("bind not allowed: {0}")]
NotAllowed(String),
#[error("bind resource shortage: {0}")]
Shortage(String),
#[error("bind general failure: {0}")]
General(String),
}
impl BindAcceptFactoryError {
pub fn to_reply_error(&self) -> ReplyError {
match self {
Self::Refused(_) => ReplyError::ConnectionRefused,
Self::NotAllowed(_) => ReplyError::ConnectionNotAllowed,
Self::Shortage(_) => ReplyError::GeneralFailure,
Self::General(_) => ReplyError::GeneralFailure,
}
}
}
pub type BindAcceptFactory = Arc<
dyn Fn(TargetAddr) -> BoxFuture<'static, Result<BindAccept, BindAcceptFactoryError>>
+ Send
+ Sync,
>;
#[cfg(test)]
mod tests {
use super::*;
use std::collections::VecDeque;
fn domain(port: u16) -> TargetAddr {
TargetAddr::Domain("example.com".to_owned(), port)
}
#[test]
fn dial_error_maps_all_variants_to_rfc_reply_codes() {
for (err, code) in [
(
SocksDialError::Refused("x".into()),
ReplyError::ConnectionRefused,
),
(
SocksDialError::NetworkUnreachable("x".into()),
ReplyError::NetworkUnreachable,
),
(
SocksDialError::HostUnreachable("x".into()),
ReplyError::HostUnreachable,
),
(
SocksDialError::NotAllowed("x".into()),
ReplyError::ConnectionNotAllowed,
),
(
SocksDialError::General("x".into()),
ReplyError::GeneralFailure,
),
] {
assert_eq!(
err.to_reply_error().as_u8(),
code.as_u8(),
"{err} must map to {code:?}"
);
}
}
#[test]
fn egress_error_maps_all_variants_to_rfc_reply_codes() {
for (err, code) in [
(
EgressError::Refused("x".into()),
ReplyError::ConnectionRefused,
),
(
EgressError::Unreachable("x".into()),
ReplyError::NetworkUnreachable,
),
(EgressError::General("x".into()), ReplyError::GeneralFailure),
] {
assert_eq!(
err.to_reply_error().as_u8(),
code.as_u8(),
"{err} must map to {code:?}"
);
}
}
#[test]
fn bind_accept_error_maps_all_variants_to_rfc_reply_codes() {
for (err, code) in [
(
BindAcceptError::Refused("x".into()),
ReplyError::ConnectionRefused,
),
(
BindAcceptError::NotAllowed("x".into()),
ReplyError::ConnectionNotAllowed,
),
(
BindAcceptError::General("x".into()),
ReplyError::GeneralFailure,
),
] {
assert_eq!(
err.to_reply_error().as_u8(),
code.as_u8(),
"{err} must map to {code:?}"
);
}
}
#[test]
fn bind_accept_factory_error_maps_all_variants_to_rfc_reply_codes() {
for (err, code) in [
(
BindAcceptFactoryError::Refused("x".into()),
ReplyError::ConnectionRefused,
),
(
BindAcceptFactoryError::NotAllowed("x".into()),
ReplyError::ConnectionNotAllowed,
),
(
BindAcceptFactoryError::Shortage("x".into()),
ReplyError::GeneralFailure,
),
(
BindAcceptFactoryError::General("x".into()),
ReplyError::GeneralFailure,
),
] {
assert_eq!(
err.to_reply_error().as_u8(),
code.as_u8(),
"{err} must map to {code:?}"
);
}
}
struct QueueEgress {
replies: VecDeque<Bytes>,
}
impl Egress for QueueEgress {
fn send(
&mut self,
_addr: TargetAddr,
payload: Bytes,
) -> BoxFuture<'_, Result<(), EgressError>> {
Box::pin(async move {
self.replies.push_back(payload);
Ok(())
})
}
fn next_reply(&mut self) -> BoxFuture<'_, Result<Option<Bytes>, EgressError>> {
Box::pin(async move { Ok(self.replies.pop_front()) })
}
}
#[tokio::test]
async fn egress_factory_lifts_a_sync_factory_into_the_async_shape() {
let factory: EgressFactory = egress_factory(|| {
Ok(Box::new(QueueEgress {
replies: VecDeque::new(),
}) as Box<dyn Egress + Send>)
});
let mut egress = factory().await.expect("Ok lifts verbatim");
egress
.send(domain(53), Bytes::from_static(b"query"))
.await
.expect("send succeeds");
let reply = egress.next_reply().await.expect("next_reply succeeds");
assert_eq!(
reply.expect("the queued reply is ready"),
Bytes::from_static(b"query")
);
assert!(egress.next_reply().await.expect("drained").is_none());
}
#[tokio::test]
async fn egress_factory_delivers_a_sync_err_verbatim() {
let factory: EgressFactory =
egress_factory(|| Err(EgressError::Unreachable("no egress".into())));
let err = match factory().await {
Err(err) => err,
Ok(_) => panic!("expected the sync Err to lift verbatim"),
};
assert_eq!(
err.to_reply_error().as_u8(),
ReplyError::NetworkUnreachable.as_u8()
);
}
#[tokio::test]
async fn dial_fn_shape_accepts_a_closure_and_yields_the_handle() {
let dial: DialFn = Arc::new(|target: TargetAddr| {
Box::pin(async move {
assert_eq!(target, domain(443));
let (client_read, target_read) = tokio::io::duplex(64);
let (target_write, client_write) = tokio::io::duplex(64);
std::mem::forget((client_read, client_write));
Ok(TargetHandle {
read: Box::new(target_read) as Box<dyn AsyncRead + Send + Sync + Unpin>,
write: Box::new(target_write) as Box<dyn AsyncWrite + Send + Sync + Unpin>,
})
})
});
let handle = dial(domain(443)).await.expect("the dial resolves");
let _ = (&handle.read, &handle.write);
}
#[tokio::test]
async fn bind_accept_factory_shape_yields_bind_addr_then_the_handle() {
let factory: BindAcceptFactory = Arc::new(|_target: TargetAddr| {
Box::pin(async move {
Ok(BindAccept::new(
SocketAddr::from(([127, 0, 0, 1], 2121)),
async {
let (client_read, target_read) = tokio::io::duplex(64);
let (target_write, client_write) = tokio::io::duplex(64);
std::mem::forget((client_read, client_write));
Ok(BindTargetHandle {
peer: TargetAddr::Ip(SocketAddr::from(([10, 0, 0, 7], 44444))),
read: Box::new(target_read) as Box<dyn AsyncRead + Send + Sync + Unpin>,
write: Box::new(target_write)
as Box<dyn AsyncWrite + Send + Sync + Unpin>,
})
},
))
})
});
let accept = factory(domain(2121)).await.expect("the factory resolves");
assert_eq!(accept.bind_addr(), SocketAddr::from(([127, 0, 0, 1], 2121)));
let handle = accept.accept().await.expect("the accept resolves");
assert_eq!(
handle.peer,
TargetAddr::Ip(SocketAddr::from(([10, 0, 0, 7], 44444)))
);
}
#[test]
fn seam_types_have_the_poc_compile_shapes() {
fn assert_send<T: Send>() {}
fn assert_send_sync<T: Send + Sync>() {}
fn assert_unpin<T: Unpin>() {}
assert_send_sync::<TargetHandle>();
assert_send_sync::<BindTargetHandle>();
assert_unpin::<TargetHandle>();
assert_unpin::<BindTargetHandle>();
assert_send_sync::<DialFn>();
assert_send_sync::<EgressFactory>();
assert_send_sync::<BindAcceptFactory>();
assert_send::<BindAccept>();
assert_unpin::<BindAccept>();
assert_send_sync::<SocksDialError>();
assert_send_sync::<EgressError>();
assert_send_sync::<BindAcceptError>();
assert_send_sync::<BindAcceptFactoryError>();
}
#[test]
fn debug_impls_render_the_handle_surfaces() {
let dialed = TargetHandle {
read: Box::new(tokio::io::empty()),
write: Box::new(tokio::io::sink()),
};
let rendered = format!("{dialed:?}");
assert!(rendered.contains("TargetHandle"), "{rendered}");
assert!(rendered.contains(".."), "{rendered}");
let accepted = BindTargetHandle {
peer: TargetAddr::Ip(SocketAddr::from(([10, 0, 0, 7], 44444))),
read: Box::new(tokio::io::empty()),
write: Box::new(tokio::io::sink()),
};
let rendered = format!("{accepted:?}");
assert!(rendered.contains("BindTargetHandle"), "{rendered}");
assert!(
rendered.contains("10.0.0.7:44444"),
"the peer is named: {rendered}"
);
let accept = BindAccept::new(
SocketAddr::from(([127, 0, 0, 1], 2121)),
std::future::pending::<Result<BindTargetHandle, BindAcceptError>>(),
);
let rendered = format!("{accept:?}");
assert!(rendered.contains("BindAccept"), "{rendered}");
assert!(
rendered.contains("127.0.0.1:2121"),
"the bind address is named: {rendered}"
);
}
}