94 lines
2.8 KiB
Rust
94 lines
2.8 KiB
Rust
use std::{
|
|
collections::HashMap,
|
|
env,
|
|
net::{IpAddr, Ipv4Addr, SocketAddr, UdpSocket},
|
|
str::FromStr,
|
|
thread,
|
|
time::Duration,
|
|
};
|
|
|
|
pub struct URedir {
|
|
true_server: IpAddr,
|
|
port: u16,
|
|
|
|
clients: UdpSocket,
|
|
table: HashMap<SocketAddr, UdpSocket>,
|
|
buf: Vec<u8>,
|
|
}
|
|
|
|
fn main() {
|
|
let mut args = env::args().skip(1);
|
|
println!("uredir by TudbuT: uredir [true server] [port] [lport] [mtu]");
|
|
let true_server: String = args.next().unwrap();
|
|
let true_server = IpAddr::from_str(&true_server).unwrap();
|
|
let port: u16 = args.next().unwrap().parse().unwrap();
|
|
let lport: u16 = args.next().unwrap().parse().unwrap();
|
|
let mtu: usize = args.next().unwrap().parse().unwrap();
|
|
|
|
println!("Start!");
|
|
URedir::run_server(true_server, port, lport, mtu);
|
|
}
|
|
|
|
impl URedir {
|
|
pub fn run_server(true_server: IpAddr, port: u16, lport: u16, mtu: usize) -> ! {
|
|
let clients = UdpSocket::bind(("0.0.0.0", lport)).unwrap();
|
|
clients.set_nonblocking(true).unwrap();
|
|
|
|
let mut server = URedir {
|
|
true_server,
|
|
port,
|
|
clients,
|
|
table: HashMap::new(),
|
|
buf: vec![0u8; mtu],
|
|
};
|
|
|
|
loop {
|
|
let did_something = server.read_from_clients() || server.read_from_server();
|
|
|
|
if !did_something {
|
|
thread::sleep(Duration::from_millis(1));
|
|
}
|
|
}
|
|
}
|
|
|
|
fn read_from_clients(&mut self) -> bool {
|
|
if let Ok((packet_length, origin)) = self.clients.recv_from(&mut self.buf) {
|
|
let connection = match self.table.get(&origin) {
|
|
Some(it) => it,
|
|
None => {
|
|
println!("Got a new connection from {origin}.");
|
|
self.create_server_connection(origin);
|
|
println!("Connection setup complete");
|
|
self.table.get(&origin).unwrap()
|
|
}
|
|
};
|
|
let _ = connection.send(&self.buf[..packet_length]);
|
|
|
|
true
|
|
} else {
|
|
false
|
|
}
|
|
}
|
|
|
|
fn read_from_server(&mut self) -> bool {
|
|
for (client_addr, server_conn) in &self.table {
|
|
if let Ok((packet_length, origin)) = server_conn.recv_from(&mut self.buf) {
|
|
if origin.ip() != self.true_server {
|
|
continue;
|
|
}
|
|
let _ = self
|
|
.clients
|
|
.send_to(&self.buf[..packet_length], client_addr);
|
|
return true;
|
|
}
|
|
}
|
|
false
|
|
}
|
|
|
|
fn create_server_connection(&mut self, origin: SocketAddr) {
|
|
let it = UdpSocket::bind((Ipv4Addr::UNSPECIFIED, 0)).unwrap();
|
|
it.connect((self.true_server, self.port)).unwrap();
|
|
it.set_nonblocking(true).unwrap();
|
|
self.table.insert(origin, it);
|
|
}
|
|
}
|