From 58cb384afdcbeb43582cb3172c3209f2cbef1067 Mon Sep 17 00:00:00 2001 From: Adam Wick Date: Sun, 25 Jul 2021 17:18:29 -0700 Subject: [PATCH] Add `HasLocalAddress` for querying a socket's address, and a test. --- src/network/address.rs | 4 + src/network/datagram.rs | 10 ++- src/network/listener.rs | 11 +-- src/network/standard.rs | 176 +++++++++++++++++++--------------------- src/network/stream.rs | 10 ++- src/server.rs | 5 +- 6 files changed, 112 insertions(+), 104 deletions(-) diff --git a/src/network/address.rs b/src/network/address.rs index 153cfb0..c7ea036 100644 --- a/src/network/address.rs +++ b/src/network/address.rs @@ -168,6 +168,10 @@ impl SOCKSv5Address { } } +pub trait HasLocalAddress { + fn local_addr(&self) -> (SOCKSv5Address, u16); +} + #[cfg(test)] impl Arbitrary for SOCKSv5Address { fn arbitrary(g: &mut Gen) -> Self { diff --git a/src/network/datagram.rs b/src/network/datagram.rs index fd74f13..6f029d3 100644 --- a/src/network/datagram.rs +++ b/src/network/datagram.rs @@ -1,8 +1,8 @@ -use crate::network::address::SOCKSv5Address; +use crate::network::address::{HasLocalAddress, SOCKSv5Address}; use async_trait::async_trait; #[async_trait] -pub trait Datagramlike: Send + Sync { +pub trait Datagramlike: Send + Sync + HasLocalAddress { type Error; async fn send_to( @@ -35,3 +35,9 @@ impl Datagramlike for GenericDatagramSocket { Ok(self.internal.recv_from(buf).await?) } } + +impl HasLocalAddress for GenericDatagramSocket { + fn local_addr(&self) -> (SOCKSv5Address, u16) { + self.internal.local_addr() + } +} diff --git a/src/network/listener.rs b/src/network/listener.rs index f125948..231e43a 100644 --- a/src/network/listener.rs +++ b/src/network/listener.rs @@ -1,13 +1,12 @@ -use crate::network::address::SOCKSv5Address; +use crate::network::address::{HasLocalAddress, SOCKSv5Address}; use crate::network::stream::GenericStream; use async_trait::async_trait; #[async_trait] -pub trait Listenerlike: Send + Sync { +pub trait Listenerlike: Send + Sync + HasLocalAddress { type Error; async fn accept(&self) -> Result<(GenericStream, SOCKSv5Address, u16), Self::Error>; - fn info(&self) -> (SOCKSv5Address, u16); } pub struct GenericListener { @@ -21,8 +20,10 @@ impl Listenerlike for GenericListener { async fn accept(&self) -> Result<(GenericStream, SOCKSv5Address, u16), Self::Error> { Ok(self.internal.accept().await?) } +} - fn info(&self) -> (SOCKSv5Address, u16) { - self.internal.info() +impl HasLocalAddress for GenericListener { + fn local_addr(&self) -> (SOCKSv5Address, u16) { + self.internal.local_addr() } } diff --git a/src/network/standard.rs b/src/network/standard.rs index 2fc9628..abe06a9 100644 --- a/src/network/standard.rs +++ b/src/network/standard.rs @@ -1,14 +1,17 @@ use crate::messages::ServerResponseStatus; -use crate::network::address::SOCKSv5Address; +use crate::network::address::{HasLocalAddress, SOCKSv5Address}; use crate::network::datagram::{Datagramlike, GenericDatagramSocket}; use crate::network::generic::Networklike; use crate::network::listener::{GenericListener, Listenerlike}; use crate::network::stream::{GenericStream, Streamlike}; use async_std::io; +#[cfg(test)] +use async_std::io::ReadExt; use async_std::net::{TcpListener, TcpStream, UdpSocket}; use async_trait::async_trait; +#[cfg(test)] +use futures::AsyncWriteExt; use log::error; -use std::net::Ipv4Addr; pub struct Builtin {} @@ -18,6 +21,27 @@ impl Builtin { } } +macro_rules! local_address_impl { + ($t: ty) => { + impl HasLocalAddress for $t { + fn local_addr(&self) -> (SOCKSv5Address, u16) { + match self.local_addr() { + Ok(a) => + (SOCKSv5Address::from(a.ip()), a.port()), + Err(e) => { + error!("Couldn't translate (Streamlike) local address to SOCKS local address: {}", e); + (SOCKSv5Address::from("localhost"), 0) + } + } + } + } + }; +} + +local_address_impl!(TcpStream); +local_address_impl!(TcpListener); +local_address_impl!(UdpSocket); + impl Streamlike for TcpStream {} #[async_trait] @@ -28,21 +52,7 @@ impl Listenerlike for TcpListener { let (base, addrport) = self.accept().await?; let addr = addrport.ip(); let port = addrport.port(); - Ok((GenericStream::from(base), SOCKSv5Address::from(addr), port)) - } - - fn info(&self) -> (SOCKSv5Address, u16) { - match self.local_addr() { - Ok(x) => { - let addr = SOCKSv5Address::from(x.ip()); - let port = x.port(); - (addr, port) - } - Err(e) => { - error!("Someone asked for a listener address, and we got an error ({}); returning 0.0.0.0:0", e); - (SOCKSv5Address::IP4(Ipv4Addr::from(0)), 0) - } - } + Ok((GenericStream::new(base), SOCKSv5Address::from(addr), port)) } } @@ -88,7 +98,7 @@ impl Networklike for Builtin { SOCKSv5Address::Name(n) => TcpStream::connect((n.as_str(), port)).await?, }; - Ok(GenericStream::from(base_stream)) + Ok(GenericStream::new(base_stream)) } async fn listen>( @@ -128,14 +138,59 @@ impl Networklike for Builtin { } } -// pub struct StandardNetworking {} -// -// impl StandardNetworking { -// pub fn new() -> StandardNetworking { -// StandardNetworking {} -// } -// } -// +#[test] +fn check_sanity() { + async_std::task::block_on(async { + // Technically, this is UDP, and UDP is lossy. We're going to assume we're not + // going to get any dropped data along here ... which is a very questionable + // assumption, morally speaking, but probably fine for most purposes. + let mut network = Builtin::new(); + let receiver = network.bind("localhost", 0).await.expect("Failed to bind receiver socket."); + let sender = network.bind("localhost", 0).await.expect("Failed to bind sender socket."); + let buffer = [0xde, 0xea, 0xbe, 0xef]; + let (receiver_addr, receiver_port) = receiver.local_addr(); + sender.send_to(&buffer, receiver_addr, receiver_port).await.expect("Failure sending datagram!"); + let mut recvbuffer = [0; 4]; + let (s, f, p) = receiver.recv_from(&mut recvbuffer).await.expect("Didn't receive UDP message?"); + let (sender_addr, sender_port) = sender.local_addr(); + assert_eq!(s, 4); + assert_eq!(f, sender_addr); + assert_eq!(p, sender_port); + assert_eq!(recvbuffer, buffer); + }); + + // This whole block should be pretty solid, though, unless the system we're + // on is in a pretty weird place. + let mut network = Builtin::new(); + + let listener = async_std::task::block_on(network.listen("localhost", 0)).expect("Couldn't set up listener on localhost"); + let (listener_address, listener_port) = listener.local_addr(); + + let listener_task_handle = async_std::task::spawn(async move { + let (mut stream, addr, port) = listener.accept().await.expect("Didn't get connection"); + let mut result_buffer = [0u8; 4]; + println!("Starting read!"); + stream.read_exact(&mut result_buffer).await.expect("Read failure in TCP test"); + (result_buffer, addr, port) + }); + + let sender_task_handle = async_std::task::spawn(async move { + let mut sender = network.connect(listener_address, listener_port).await.expect("Coudln't connect to listener?"); + let (sender_address, sender_port) = sender.local_addr(); + let send_buffer = [0xa, 0xff, 0xab, 0x1e]; + sender.write_all(&send_buffer).await.expect("Couldn't send the write buffer"); + (sender_address, sender_port) + }); + + async_std::task::block_on(async { + let (result, result_from, result_from_port) = listener_task_handle.await; + assert_eq!(result, [0xa, 0xff, 0xab, 0x1e]); + let (sender_address, sender_port) = sender_task_handle.await; + assert_eq!(result_from, sender_address); + assert_eq!(result_from_port, sender_port); + }); +} + impl From for ServerResponseStatus { fn from(e: io::Error) -> ServerResponseStatus { match e.kind() { @@ -144,71 +199,4 @@ impl From for ServerResponseStatus { _ => ServerResponseStatus::GeneralFailure, } } -} -// -// #[async_trait] -// impl Network for StandardNetworking { -// type Stream = TcpStream; -// type Listener = TcpListener; -// type UdpSocket = UdpSocket; -// type Error = io::Error; -// -// async fn connect( -// &mut self, -// addr: A, -// port: u16, -// ) -> Result { -// let target = addr.into(); -// -// match target { -// SOCKSv5Address::IP4(a) => TcpStream::connect((a, port)).await, -// SOCKSv5Address::IP6(a) => TcpStream::connect((a, port)).await, -// SOCKSv5Address::Name(n) => TcpStream::connect((n.as_str(), port)).await, -// } -// } -// -// async fn udp_socket( -// &mut self, -// addr: A, -// port: Option, -// ) -> Result { -// let me = addr.into(); -// let real_port = port.unwrap_or(0); -// -// match me { -// SOCKSv5Address::IP4(a) => UdpSocket::bind((a, real_port)).await, -// SOCKSv5Address::IP6(a) => UdpSocket::bind((a, real_port)).await, -// SOCKSv5Address::Name(n) => UdpSocket::bind((n.as_str(), real_port)).await, -// } -// } -// -// async fn listen( -// &mut self, -// addr: A, -// port: Option, -// ) -> Result { -// let me = addr.into(); -// let real_port = port.unwrap_or(0); -// -// match me { -// SOCKSv5Address::IP4(a) => TcpListener::bind((a, real_port)).await, -// SOCKSv5Address::IP6(a) => TcpListener::bind((a, real_port)).await, -// SOCKSv5Address::Name(n) => TcpListener::bind((n.as_str(), real_port)).await, -// } -// } -// } -// -// #[async_trait] -// impl SingleShotListener for TcpListener { -// async fn accept(self) -> Result { -// self.accept().await -// } -// -// fn info(&self) -> Result<(SOCKSv5Address, u16), io::Error> { -// match self.local_addr()? { -// SocketAddr::V4(a) => Ok((SOCKSv5Address::IP4(*a.ip()), a.port())), -// SocketAddr::V6(a) => Ok((SOCKSv5Address::IP6(*a.ip()), a.port())), -// } -// } -// } -// +} \ No newline at end of file diff --git a/src/network/stream.rs b/src/network/stream.rs index f4d8d84..692b7b0 100644 --- a/src/network/stream.rs +++ b/src/network/stream.rs @@ -1,10 +1,13 @@ use async_std::task::{Context, Poll}; use futures::io; +use crate::network::SOCKSv5Address; use futures::io::{AsyncRead, AsyncWrite}; use std::pin::Pin; use std::sync::{Arc, Mutex}; -pub trait Streamlike: AsyncRead + AsyncWrite + Send + Sync + Unpin {} +use super::address::HasLocalAddress; + +pub trait Streamlike: AsyncRead + AsyncWrite + HasLocalAddress + Send + Sync + Unpin {} #[derive(Clone)] pub struct GenericStream { @@ -19,6 +22,11 @@ impl GenericStream { } } +impl HasLocalAddress for GenericStream { + fn local_addr(&self) -> (SOCKSv5Address, u16) { + let item = self.internal.lock().unwrap(); + item.local_addr() + } } impl AsyncRead for GenericStream { diff --git a/src/server.rs b/src/server.rs index 1d75603..060ebff 100644 --- a/src/server.rs +++ b/src/server.rs @@ -3,6 +3,7 @@ use crate::messages::{ AuthenticationMethod, ClientConnectionCommand, ClientConnectionRequest, ClientGreeting, ClientUsernamePassword, ServerChoice, ServerResponse, ServerResponseStatus, }; +use crate::network::address::HasLocalAddress; use crate::network::generic::Networklike; use crate::network::listener::{GenericListener, Listenerlike}; use crate::network::stream::GenericStream; @@ -45,7 +46,7 @@ impl SOCKSv5Server { } pub async fn run(self) -> Result<(), N::Error> { - let (my_addr, my_port) = self.listener.info(); + let (my_addr, my_port) = self.listener.local_addr(); info!("Starting SOCKSv5 server on {}:{}", my_addr, my_port); let locked_network = Arc::new(Mutex::new(self.network)); @@ -241,7 +242,7 @@ where continue; } }; - let (bound_address, bound_port) = incoming_listener.info(); + let (bound_address, bound_port) = incoming_listener.local_addr(); trace!( "Set up {}:{} to address request for {}:{}", bound_address,