-
-
Save gauravssnl/9e4dceb88630ecb68a98e9641f6f7588 to your computer and use it in GitHub Desktop.
A simple and lightweight multi threaded TCP proxy written in Rust language. Not using any 3rd-party libraries but just standard libraries.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
use std::net::{TcpListener, SocketAddr, TcpStream, Shutdown, SocketAddrV4, Ipv4Addr}; | |
use std::str::FromStr; | |
use std::thread::{spawn, JoinHandle}; | |
use std::io::{BufReader, BufWriter, Read, Write}; | |
fn pipe(incoming: &mut TcpStream, outgoing: &mut TcpStream) -> Result<(), String> { | |
let mut buffer = [0; 1024]; | |
loop { | |
match incoming.read(&mut buffer) { | |
Ok(bytes_read) => { | |
// Socket is disconnected => Shutdown the other socket as well | |
if bytes_read == 0 { | |
outgoing.shutdown(Shutdown::Both); | |
break; | |
} | |
if outgoing.write(&buffer[..bytes_read]).is_ok() { | |
outgoing.flush(); | |
} | |
}, | |
Err(error) => return Err(format!("Could not read data: {}", error)) | |
} | |
} | |
Ok(()) | |
} | |
fn proxy_connection(mut incoming: TcpStream, target: &SocketAddr) -> Result<(), String> { | |
debug!("Client connected from: {:#?}", incoming.peer_addr()); | |
let mut outgoing = TcpStream::connect(target) | |
.map_err(|error| format!("Could not establish connection to {}: {}", target, error))?; | |
let mut incoming_clone = incoming.try_clone().map_err(|e| e.to_string())?; | |
let mut outgoing_clone = outgoing.try_clone().map_err(|e| e.to_string())?; | |
// Pipe for- and backward asynchronously | |
let forward = spawn(move || pipe(&mut incoming, &mut outgoing)); | |
let backward = spawn(move || pipe(&mut outgoing_clone, &mut incoming_clone)); | |
debug!("Proxying data..."); | |
forward.join().map_err(|error| format!("Forward failed: {}", error))?; | |
backward.join().map_err(|error| format!("Backward failed: {}", error))?; | |
debug!("Socket closed"); | |
Ok(()) | |
} | |
fn proxy(local_port: u16, target_addr: SocketAddr) -> JoinHandle<()> { | |
let listener = TcpListener::bind(local_port).unwrap(); | |
// One thread per port listener | |
spawn(move || { | |
for socket in listener.incoming() { | |
let socket = socket.unwrap_or_else(|error| { | |
error!("Could not handle connection: {}", error); | |
return; | |
}); | |
// One thread per connection | |
spawn(move || { | |
if let Err(error) = proxy_connection(socket, &target_addr) { | |
error!(error); | |
} | |
}); | |
} | |
}) | |
} | |
struct Mapping { | |
local_port: u16, | |
target_address: String | |
} | |
fn start_proxies(mappings: Vec<Mapping>) { | |
let mut handles: Vec<JoinHandle<()>> = Vec::new(); | |
for mapping in mappings { | |
&handles.push(proxy( | |
mapping.local_port, | |
mapping.target_address.parse::<SocketAddr>().unwrap() | |
)); | |
} | |
for handle in handles { | |
handle.join(); | |
} | |
} | |
fn main() { | |
let mappings = vec![Mapping { local_port: 1999, target_address: "127.0.0.1:22"}]; | |
start_proxies(mappings); | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment