hatter/simple-rust-udp-proxy
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

chore: code style

Browse Source
This commit is contained in:
Hatter Jiang
2020-09-13 10:01:25 +08:00
parent 96b3ca62ea
commit 552f44ab10
2 changed files with 28 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

38
src/main.rs
View File

@@ -12,6 +12,7 @@ use std::thread;
use std::time::Duration;
use clap::{ Arg, App };
const LOCAL_ADDR: &str = "127.0.0.1";
const TIMEOUT: u64 = 3 * 60 * 100; //3 minutes
fn main() {
@@ -31,34 +32,37 @@ fn main() {
.arg(Arg::with_name("allowed_list").short("A").long("allowed-list").takes_value(true).multiple(true).help("Allowed IP list, e.g. 127.0.0.1, 120.0.0.0/8"))
.get_matches();
let local_port: i32 = matches.value_of("local_port").unwrap().parse().unwrap();
let remote_port: i32 = matches.value_of("remote_port").unwrap().parse().unwrap();
let local_port: u16 = matches.value_of("local_port").unwrap().parse().unwrap();
let remote_port: u16 = matches.value_of("remote_port").unwrap().parse().unwrap();
let remote_host = matches.value_of("host").unwrap();
let bind_addr = match matches.value_of("bind") {
Some(addr) => addr.to_owned(),
None => "127.0.0.1".to_owned(),
None => LOCAL_ADDR.to_owned(),
};
forward(&bind_addr, local_port, &remote_host, remote_port);
}
fn forward(bind_addr: &str, local_port: i32, remote_host: &str, remote_port: i32) {
fn forward(bind_addr: &str, local_port: u16, remote_host: &str, remote_port: u16) {
let local_addr = format!("{}:{}", bind_addr, local_port);
debugging!("Listen address and port: {}", local_addr);
let local = UdpSocket::bind(&local_addr).unwrap_or_else(|_| panic!("Unable to bind to {}", &local_addr));
information!("Listening on {}", local.local_addr().unwrap());
let remote_addr = format!("{}:{}", remote_host, remote_port);
let responder = local.try_clone()
.unwrap_or_else(|_| panic!("Failed to clone primary listening address socket {}", local.local_addr().unwrap()));
let responder = local.try_clone().unwrap_or_else(
|_| panic!("Failed to clone primary listening address socket {}", local.local_addr().unwrap())
);
let (main_sender, main_receiver) = channel::<(_, Vec<u8>)>();
thread::spawn(move || {
debugging!("Started new thread to deal out responses to clients");
loop {
let (dest, buf) = main_receiver.recv().unwrap();
let to_send = buf.as_slice();
responder.send_to(to_send, dest)
.unwrap_or_else(|_| panic!("Failed to forward response from upstream server to client {}", dest));
responder.send_to(to_send, dest).unwrap_or_else(
|_| panic!("Failed to forward response from upstream server to client {}", dest)
);
}
});
@@ -94,10 +98,12 @@ fn forward(bind_addr: &str, local_port: i32, remote_host: &str, remote_port: i32
//connection to 0.0.0.0 in all cases.
let temp_outgoing_addr = format!("0.0.0.0:{}", 1024 + rand::random::<u16>());
debugging!("Establishing new forwarder for client {} on {}", src_addr, &temp_outgoing_addr);
let upstream_send = UdpSocket::bind(&temp_outgoing_addr)
.unwrap_or_else(|_| panic!("Failed to bind to transient address {}", &temp_outgoing_addr));
let upstream_recv = upstream_send.try_clone()
.unwrap_or_else(|_| panic!("Failed to clone client-specific connection to upstream!"));
let upstream_send = UdpSocket::bind(&temp_outgoing_addr).unwrap_or_else(
|_| panic!("Failed to bind to transient address {}", &temp_outgoing_addr)
);
let upstream_recv = upstream_send.try_clone().unwrap_or_else(
|_| panic!("Failed to clone client-specific connection to upstream!")
);
let mut timeouts : u64 = 0;
let timed_out = Arc::new(AtomicBool::new(false));
@@ -110,8 +116,7 @@ fn forward(bind_addr: &str, local_port: i32, remote_host: &str, remote_port: i32
match upstream_recv.recv_from(&mut from_upstream) {
Ok((bytes_rcvd, _)) => {
let to_send = from_upstream[..bytes_rcvd].to_vec();
local_send_queue.send((src_addr, to_send))
.expect("Failed to queue response from upstream server for forwarding!");
local_send_queue.send((src_addr, to_send)).expect("Failed to queue response from upstream server for forwarding!");
},
Err(_) => {
if local_timed_out.load(Ordering::Relaxed) {
@@ -126,8 +131,9 @@ fn forward(bind_addr: &str, local_port: i32, remote_host: &str, remote_port: i32
loop {
match receiver.recv_timeout(Duration::from_millis(TIMEOUT)) {
Ok(from_client) => {
upstream_send.send_to(from_client.as_slice(), &remote_addr_copy)
.unwrap_or_else(|_| panic!("Failed to forward packet from client {} to upstream server!", src_addr));
upstream_send.send_to(from_client.as_slice(), &remote_addr_copy).unwrap_or_else(
|_| panic!("Failed to forward packet from client {} to upstream server!", src_addr)
);
timeouts = 0; //reset timeout count
},
Err(_) => {