From 13314675a71beb5de55e0cabe1ab9495e1bbf1a0 Mon Sep 17 00:00:00 2001 From: moparisthebest Date: Wed, 17 Jul 2019 00:13:08 -0400 Subject: [PATCH] Refactor into lib, Implement `udp-test -is` to run full test in the same process --- README.md | 3 +- src/bin/udp-test.rs | 53 +++++++++++++- src/bin/wireguard-proxy.rs | 58 ++------------- src/bin/wireguard-proxyd.rs | 96 ++---------------------- src/lib.rs | 142 +++++++++++++++++++++++++++++++++++- 5 files changed, 207 insertions(+), 145 deletions(-) diff --git a/README.md b/README.md index 4f9b57a..ec5c6dc 100644 --- a/README.md +++ b/README.md @@ -9,7 +9,8 @@ Testing: `udp-test` is a utility to send a UDP packet and then receive a UDP packet and ensure they are the same, this verifies packets sent through proxy/proxyd are unmolested `test.sh` runs udp-test against itself and then through proxyd/proxy -`udp-test -s` runs udp-test against itself through proxyd/proxy +`udp-test -s` runs udp-test against itself through proxyd/proxy by spawning actual binaries +`udp-test -is` runs udp-test against itself through proxyd/proxy in same executable by using library, so does not test command line parsing etc Testing with GNU netcat: diff --git a/src/bin/udp-test.rs b/src/bin/udp-test.rs index 73403dd..58ca235 100644 --- a/src/bin/udp-test.rs +++ b/src/bin/udp-test.rs @@ -3,7 +3,7 @@ use std::time::Duration; use std::process::{exit, Command}; use std::{env, thread}; -use wireguard_proxy::Args; +use wireguard_proxy::{Args, ProxyClient, ProxyServer}; const PONG: [u8; 246] = [ 0x6A, 0x2, 0x6B, 0xC, 0x6C, 0x3F, 0x6D, 0xC, 0xA2, 0xEA, 0xDA, 0xB6, 0xDC, 0xD6, 0x6E, 0x0, @@ -70,10 +70,10 @@ impl Server { fn main() { let raw_args = env::args().collect(); let args = Args::new(&raw_args); - let first_arg = args.get_str(1, "127.0.0.1:51821"); + let mut first_arg = args.get_str(1, "127.0.0.1:51821"); if first_arg.contains("-h") { println!( - "usage: {} [-h] [-s run a self test through proxy/proxyd] [udp_host, 127.0.0.1:51821] [udp_target, 127.0.0.1:51821] [socket_timeout, 1]", + "usage: {} [-h] [-s run a self test through proxy/proxyd] [-is run a self test through proxy/proxyd without spawning other processes] [udp_host, 127.0.0.1:51821] [udp_target, 127.0.0.1:51821] [socket_timeout, 10]", args.get_str(0, "udp-test") ); return; @@ -123,12 +123,57 @@ fn main() { .code() .expect("could not get udp-test exit code"), ); + } else if first_arg.contains("-is") { + let host = "127.0.0.1:51822"; + let tcp_host = "127.0.0.1:5555"; + let sleep = Duration::from_secs(5); + + let proxy_server = ProxyServer::new( + tcp_host.to_owned(), + host.to_owned(), + "127.0.0.1".to_owned(), + 30000, + 30100, + 0, + ); + + println!( + "udp_target: {}, udp_bind_host_range: 127.0.0.1:30000-30100, socket_timeout: {:?}", + proxy_server.client_handler.udp_target, proxy_server.client_handler.socket_timeout, + ); + + println!("executing: wireguard-proxyd '{}' '{}'", tcp_host, host); + thread::spawn(move || proxy_server.start().expect("error running proxy_server")); + println!("waiting: {:?} for wireguard-proxyd to come up.....", sleep); + thread::sleep(sleep); + + let proxy_client = ProxyClient::new( + "127.0.0.1:51821".to_owned(), + "127.0.0.1:51820".to_owned(), + tcp_host.to_owned().to_owned(), + 15, + ); + + println!( + "udp_host: {}, udp_target: {}, tcp_target: {}, socket_timeout: {:?}", + proxy_client.udp_host, + proxy_client.udp_target, + proxy_client.tcp_target, + proxy_client.socket_timeout, + ); + + println!("executing: wireguard-proxy"); + thread::spawn(move || proxy_client.start().expect("error running proxy_client")); + println!("waiting: {:?} for wireguard-proxy to come up.....", sleep); + thread::sleep(sleep); + + first_arg = host; } let server = Server::new( first_arg.to_owned(), args.get_str(2, "127.0.0.1:51821").to_owned(), - args.get(3, 1), + args.get(3, 10), ); println!( diff --git a/src/bin/wireguard-proxy.rs b/src/bin/wireguard-proxy.rs index 81e4627..a065772 100644 --- a/src/bin/wireguard-proxy.rs +++ b/src/bin/wireguard-proxy.rs @@ -1,52 +1,5 @@ -use std::net::{TcpStream, UdpSocket}; -use std::time::Duration; - use std::env; -use std::thread; -use wireguard_proxy::{Args, TcpUdpPipe}; - -struct Server { - udp_host: String, - udp_target: String, - tcp_target: String, - socket_timeout: Option, -} - -impl Server { - fn new(udp_host: String, udp_target: String, tcp_target: String, secs: u64) -> Server { - Server { - udp_host, - udp_target, - tcp_target, - socket_timeout: match secs { - 0 => None, - x => Some(Duration::from_secs(x)), - }, - } - } - - fn start(&self) -> std::io::Result { - let tcp_stream = TcpStream::connect(&self.tcp_target)?; - - tcp_stream.set_read_timeout(self.socket_timeout)?; - - let udp_socket = UdpSocket::bind(&self.udp_host)?; - udp_socket.set_read_timeout(self.socket_timeout)?; - //udp_socket.connect(&self.udp_target)?; // this isn't strictly needed... just filters who we can receive from - - let mut udp_pipe = TcpUdpPipe::new(tcp_stream, udp_socket); - let mut udp_pipe_clone = udp_pipe.try_clone()?; - thread::spawn(move || loop { - udp_pipe_clone - .udp_to_tcp() - .expect("cannot write to tcp_clone"); - }); - - loop { - udp_pipe.tcp_to_udp()?; - } - } -} +use wireguard_proxy::{Args, ProxyClient}; fn main() { let raw_args = env::args().collect(); @@ -59,7 +12,7 @@ fn main() { return; } - let server = Server::new( + let proxy_client = ProxyClient::new( args.get_str(1, "127.0.0.1:51821").to_owned(), args.get_str(2, "127.0.0.1:51820").to_owned(), args.get_str(3, "127.0.0.1:5555").to_owned(), @@ -68,8 +21,11 @@ fn main() { println!( "udp_host: {}, udp_target: {}, tcp_target: {}, socket_timeout: {:?}", - server.udp_host, server.udp_target, server.tcp_target, server.socket_timeout, + proxy_client.udp_host, + proxy_client.udp_target, + proxy_client.tcp_target, + proxy_client.socket_timeout, ); - server.start().expect("error running server"); + proxy_client.start().expect("error running proxy_client"); } diff --git a/src/bin/wireguard-proxyd.rs b/src/bin/wireguard-proxyd.rs index 4d3dd05..300a9f7 100644 --- a/src/bin/wireguard-proxyd.rs +++ b/src/bin/wireguard-proxyd.rs @@ -1,70 +1,5 @@ -use std::net::{TcpListener, TcpStream, UdpSocket}; -use std::time::Duration; - use std::env; -use std::sync::Arc; -use std::thread; -use wireguard_proxy::{Args, TcpUdpPipe}; - -struct Server { - udp_target: String, - udp_host: String, - udp_low_port: u16, - udp_high_port: u16, - socket_timeout: Option, -} - -impl Server { - fn new( - udp_target: String, - udp_host: String, - udp_low_port: u16, - udp_high_port: u16, - secs: u64, - ) -> Server { - Server { - udp_target, - udp_host, - udp_low_port, - udp_high_port, - socket_timeout: match secs { - 0 => None, - x => Some(Duration::from_secs(x)), - }, - } - } - - fn handle_client(&self, tcp_stream: TcpStream) -> std::io::Result { - tcp_stream.set_read_timeout(self.socket_timeout)?; - - let mut port = self.udp_low_port; - let udp_socket = loop { - match UdpSocket::bind((&self.udp_host[..], port)) { - Ok(sock) => break sock, - Err(_) => { - port += 1; - if port > self.udp_high_port { - panic!("cannot find free port, increase range?"); - } - } - } - }; - udp_socket.set_read_timeout(self.socket_timeout)?; - udp_socket.connect(&self.udp_target)?; - - let mut udp_pipe = TcpUdpPipe::new(tcp_stream, udp_socket); - let mut udp_pipe_clone = udp_pipe.try_clone()?; - thread::spawn(move || loop { - udp_pipe_clone - .udp_to_tcp() - .expect("cannot write to tcp_clone"); - }); - - loop { - udp_pipe.tcp_to_udp()?; - } - } -} +use wireguard_proxy::{Args, ProxyServer}; fn main() { let raw_args = env::args().collect(); @@ -76,7 +11,6 @@ fn main() { ); return; } - let host = args.get_str(1, "127.0.0.1:5555"); let udp_bind_host_range_str = args.get_str(3, "127.0.0.1:30000-40000"); let mut udp_bind_host_range = udp_bind_host_range_str.split(":"); @@ -100,35 +34,21 @@ fn main() { .parse::() .expect("udp_bind_host_range low port invalid"); - let server = Arc::new(Server::new( + let proxy_server = ProxyServer::new( + args.get_str(1, "127.0.0.1:5555").to_owned(), args.get_str(2, "127.0.0.1:51820").to_owned(), udp_host.to_string(), udp_low_port, udp_high_port, args.get(4, 0), - )); + ); println!( "udp_target: {}, udp_bind_host_range: {}, socket_timeout: {:?}", - server.udp_target, udp_bind_host_range_str, server.socket_timeout, + proxy_server.client_handler.udp_target, + udp_bind_host_range_str, + proxy_server.client_handler.socket_timeout, ); - let listener = TcpListener::bind(&host).unwrap(); - println!("Listening for connections on {}", &host); - - for stream in listener.incoming() { - match stream { - Ok(stream) => { - let server = server.clone(); - thread::spawn(move || { - server - .handle_client(stream) - .expect("error handling connection") - }); - } - Err(e) => { - println!("Unable to connect: {}", e); - } - } - } + proxy_server.start().expect("error running proxy_server"); } diff --git a/src/lib.rs b/src/lib.rs index 09f0374..b20371a 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,6 +1,9 @@ use std::io::{Read, Write}; -use std::net::{TcpStream, UdpSocket}; +use std::net::{TcpListener, TcpStream, UdpSocket}; use std::str::FromStr; +use std::sync::Arc; +use std::thread; +use std::time::Duration; pub struct Args<'a> { args: &'a Vec, @@ -74,3 +77,140 @@ impl TcpUdpPipe { //assert_eq!(sent, len); } } + +pub struct ProxyClient { + pub udp_host: String, + pub udp_target: String, + pub tcp_target: String, + pub socket_timeout: Option, +} + +impl ProxyClient { + pub fn new(udp_host: String, udp_target: String, tcp_target: String, secs: u64) -> ProxyClient { + ProxyClient { + udp_host, + udp_target, + tcp_target, + socket_timeout: match secs { + 0 => None, + x => Some(Duration::from_secs(x)), + }, + } + } + + pub fn start(&self) -> std::io::Result { + let tcp_stream = TcpStream::connect(&self.tcp_target)?; + + tcp_stream.set_read_timeout(self.socket_timeout)?; + + let udp_socket = UdpSocket::bind(&self.udp_host)?; + udp_socket.set_read_timeout(self.socket_timeout)?; + //udp_socket.connect(&self.udp_target)?; // this isn't strictly needed... just filters who we can receive from + + let mut udp_pipe = TcpUdpPipe::new(tcp_stream, udp_socket); + let mut udp_pipe_clone = udp_pipe.try_clone()?; + thread::spawn(move || loop { + udp_pipe_clone + .udp_to_tcp() + .expect("cannot write to tcp_clone"); + }); + + loop { + udp_pipe.tcp_to_udp()?; + } + } +} + +pub struct ProxyServer { + pub tcp_host: String, + pub client_handler: Arc, +} + +impl ProxyServer { + pub fn new( + tcp_host: String, + udp_target: String, + udp_host: String, + udp_low_port: u16, + udp_high_port: u16, + secs: u64, + ) -> ProxyServer { + let client_handler = Arc::new(ProxyServerClientHandler { + udp_target, + udp_host, + udp_low_port, + udp_high_port, + socket_timeout: match secs { + 0 => None, + x => Some(Duration::from_secs(x)), + }, + }); + ProxyServer { + tcp_host, + client_handler, + } + } + + pub fn start(&self) -> std::io::Result<()> { + let listener = TcpListener::bind(&self.tcp_host)?; + println!("Listening for connections on {}", &self.tcp_host); + + for stream in listener.incoming() { + match stream { + Ok(stream) => { + let client_handler = self.client_handler.clone(); + thread::spawn(move || { + client_handler + .handle_client(stream) + .expect("error handling connection") + }); + } + Err(e) => { + println!("Unable to connect: {}", e); + } + } + } + Ok(()) + } +} + +pub struct ProxyServerClientHandler { + pub udp_target: String, + pub udp_host: String, + pub udp_low_port: u16, + pub udp_high_port: u16, + pub socket_timeout: Option, +} + +impl ProxyServerClientHandler { + pub fn handle_client(&self, tcp_stream: TcpStream) -> std::io::Result { + tcp_stream.set_read_timeout(self.socket_timeout)?; + + let mut port = self.udp_low_port; + let udp_socket = loop { + match UdpSocket::bind((&self.udp_host[..], port)) { + Ok(sock) => break sock, + Err(_) => { + port += 1; + if port > self.udp_high_port { + panic!("cannot find free port, increase range?"); + } + } + } + }; + udp_socket.set_read_timeout(self.socket_timeout)?; + udp_socket.connect(&self.udp_target)?; + + let mut udp_pipe = TcpUdpPipe::new(tcp_stream, udp_socket); + let mut udp_pipe_clone = udp_pipe.try_clone()?; + thread::spawn(move || loop { + udp_pipe_clone + .udp_to_tcp() + .expect("cannot write to tcp_clone"); + }); + + loop { + udp_pipe.tcp_to_udp()?; + } + } +}