Skip to content

Instantly share code, notes, and snippets.

@fragsalat
Created November 28, 2019 15:24
Show Gist options
  • Star 4 You must be signed in to star a gist
  • Fork 1 You must be signed in to fork a gist
  • Save fragsalat/d55cc7a88ed56a97b89c2e13362a58a8 to your computer and use it in GitHub Desktop.
Save fragsalat/d55cc7a88ed56a97b89c2e13362a58a8 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.
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