Ops/rustocat
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Implementing udp support partially...in theory. Still untested :)

Browse Source
This commit is contained in:
Fabian Stamm 2022-09-21 22:22:10 +02:00
parent f7af1ac020
commit 7530ce3aae
7 changed files with 370 additions and 153 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
.editorconfig Normal file
View File

@ -0,0 +1,2 @@
[*.rs]
indent_size = 4

142
Cargo.lock generated
View File

@ -128,9 +128,9 @@ dependencies = [
[[package]] [[package]]
name = "hashbrown" name = "hashbrown"
version = "0.11.2" version = "0.12.3"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ab5ef0d4909ef3724cc8cce6ccc8572c5c817592e9285f5464f8e86f8bd3726e" checksum = "8a9ee70c43aaf417c914396645a0fa852624801b24ebb7ae78fe8272889ac888"
[[package]] [[package]]
name = "hermit-abi" name = "hermit-abi"
@ -143,9 +143,9 @@ dependencies = [
[[package]] [[package]]
name = "indexmap" name = "indexmap"
version = "1.8.0" version = "1.9.1"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "282a6247722caba404c065016bbfa522806e51714c34f5dfc3e4a3a46fcb4223" checksum = "10a35a97730320ffe8e2d410b5d3b69279b98d2c14bdb8b70ea89ecf7888d41e"
dependencies = [ dependencies = [
"autocfg", "autocfg",
"hashbrown", "hashbrown",
@ -163,12 +163,6 @@ version = "0.2.121"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "efaa7b300f3b5fe8eb6bf21ce3895e1751d9665086af2d64b42f19701015ff4f" checksum = "efaa7b300f3b5fe8eb6bf21ce3895e1751d9665086af2d64b42f19701015ff4f"
[[package]]
name = "linked-hash-map"
version = "0.5.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7fb9b38af92608140b86b693604b9ffcc5824240a484d1ecd4795bacb2fe88f3"
[[package]] [[package]]
name = "lock_api" name = "lock_api"
version = "0.4.6" version = "0.4.6"
@ -195,34 +189,14 @@ checksum = "308cc39be01b73d0d18f82a0e7b2a3df85245f84af96fdddc5d202d27e47b86a"
[[package]] [[package]]
name = "mio" name = "mio"
version = "0.8.2" version = "0.8.4"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "52da4364ffb0e4fe33a9841a98a3f3014fb964045ce4f7a45a398243c8d6b0c9" checksum = "57ee1c23c7c63b0c9250c339ffdc69255f110b298b901b9f6c82547b7b87caaf"
dependencies = [ dependencies = [
"libc", "libc",
"log", "log",
"miow",
"ntapi",
"wasi", "wasi",
"winapi", "windows-sys 0.36.1",
]
[[package]]
name = "miow"
version = "0.3.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b9f1c5b025cda876f66ef43a113f91ebc9f4ccef34843000e0adf6ebbab84e21"
dependencies = [
"winapi",
]
[[package]]
name = "ntapi"
version = "0.3.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c28774a7fd2fbb4f0babd8237ce554b73af68021b5f695a3cebd6c59bac0980f"
dependencies = [
"winapi",
] ]
[[package]] [[package]]
@ -261,7 +235,7 @@ dependencies = [
"libc", "libc",
"redox_syscall", "redox_syscall",
"smallvec", "smallvec",
"windows-sys", "windows-sys 0.32.0",
] ]
[[package]] [[package]]
@ -284,11 +258,11 @@ checksum = "eb9f9e6e233e5c4a35559a617bf40a4ec447db2e84c20b55a6f83167b7e57872"
[[package]] [[package]]
name = "proc-macro2" name = "proc-macro2"
version = "1.0.36" version = "1.0.43"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c7342d5883fbccae1cc37a2353b09c87c9b0f3afd73f5fb9bba687a1f733b029" checksum = "0a2ca2c61bc9f3d74d2886294ab7b9853abd9c1ad903a3ac7815c58989bb7bab"
dependencies = [ dependencies = [
"unicode-xid", "unicode-ident",
] ]
[[package]] [[package]]
@ -366,18 +340,18 @@ checksum = "d29ab0c6d3fc0ee92fe66e2d99f700eab17a8d57d1c1d3b748380fb20baa78cd"
[[package]] [[package]]
name = "serde" name = "serde"
version = "1.0.136" version = "1.0.144"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ce31e24b01e1e524df96f1c2fdd054405f8d7376249a5110886fb4b658484789" checksum = "0f747710de3dcd43b88c9168773254e809d8ddbdf9653b84e2554ab219f17860"
dependencies = [ dependencies = [
"serde_derive", "serde_derive",
] ]
[[package]] [[package]]
name = "serde_derive" name = "serde_derive"
version = "1.0.136" version = "1.0.144"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "08597e7152fcd306f41838ed3e37be9eaeed2b61c42e2117266a554fab4662f9" checksum = "94ed3a816fb1d101812f83e789f888322c34e291f894f19590dc310963e87a00"
dependencies = [ dependencies = [
"proc-macro2", "proc-macro2",
"quote", "quote",
@ -397,14 +371,15 @@ dependencies = [
[[package]] [[package]]
name = "serde_yaml" name = "serde_yaml"
version = "0.8.23" version = "0.9.13"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a4a521f2940385c165a24ee286aa8599633d162077a54bdcae2a6fd5a7bfa7a0" checksum = "8613d593412a0deb7bbd8de9d908efff5a0cb9ccd8f62c641e7b2ed2f57291d1"
dependencies = [ dependencies = [
"indexmap", "indexmap",
"itoa",
"ryu", "ryu",
"serde", "serde",
"yaml-rust", "unsafe-libyaml",
] ]
[[package]] [[package]]
@ -446,21 +421,22 @@ dependencies = [
[[package]] [[package]]
name = "syn" name = "syn"
version = "1.0.89" version = "1.0.100"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ea297be220d52398dcc07ce15a209fce436d361735ac1db700cab3b6cdfb9f54" checksum = "52205623b1b0f064a4e71182c3b18ae902267282930c6d5462c91b859668426e"
dependencies = [ dependencies = [
"proc-macro2", "proc-macro2",
"quote", "quote",
"unicode-xid", "unicode-ident",
] ]
[[package]] [[package]]
name = "tokio" name = "tokio"
version = "1.17.0" version = "1.21.1"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2af73ac49756f3f7c01172e34a23e5d0216f6c32333757c2c61feb2bbff5a5ee" checksum = "0020c875007ad96677dcc890298f4b942882c5d4eb7cc8f439fc3bf813dc9c95"
dependencies = [ dependencies = [
"autocfg",
"bytes", "bytes",
"libc", "libc",
"memchr", "memchr",
@ -487,10 +463,16 @@ dependencies = [
] ]
[[package]] [[package]]
name = "unicode-xid" name = "unicode-ident"
version = "0.2.2" version = "1.0.4"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8ccb82d61f80a663efe1f787a51b16b5a51e3314d6ac365b08639f52387b33f3" checksum = "dcc811dc4066ac62f84f11307873c4850cb653bfa9b1719cee2bd2204a4bc5dd"
[[package]]
name = "unsafe-libyaml"
version = "0.2.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c1e5fa573d8ac5f1a856f8d7be41d390ee973daf97c806b2c1a465e4e1406e68"
[[package]] [[package]]
name = "wasi" name = "wasi"
@ -526,11 +508,24 @@ version = "0.32.0"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3df6e476185f92a12c072be4a189a0210dcdcf512a1891d6dff9edb874deadc6" checksum = "3df6e476185f92a12c072be4a189a0210dcdcf512a1891d6dff9edb874deadc6"
dependencies = [ dependencies = [
"windows_aarch64_msvc", "windows_aarch64_msvc 0.32.0",
"windows_i686_gnu", "windows_i686_gnu 0.32.0",
"windows_i686_msvc", "windows_i686_msvc 0.32.0",
"windows_x86_64_gnu", "windows_x86_64_gnu 0.32.0",
"windows_x86_64_msvc", "windows_x86_64_msvc 0.32.0",
]
[[package]]
name = "windows-sys"
version = "0.36.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ea04155a16a59f9eab786fe12a4a450e75cdb175f9e0d80da1e17db09f55b8d2"
dependencies = [
"windows_aarch64_msvc 0.36.1",
"windows_i686_gnu 0.36.1",
"windows_i686_msvc 0.36.1",
"windows_x86_64_gnu 0.36.1",
"windows_x86_64_msvc 0.36.1",
] ]
[[package]] [[package]]
@ -539,24 +534,48 @@ version = "0.32.0"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d8e92753b1c443191654ec532f14c199742964a061be25d77d7a96f09db20bf5" checksum = "d8e92753b1c443191654ec532f14c199742964a061be25d77d7a96f09db20bf5"
[[package]]
name = "windows_aarch64_msvc"
version = "0.36.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9bb8c3fd39ade2d67e9874ac4f3db21f0d710bee00fe7cab16949ec184eeaa47"
[[package]] [[package]]
name = "windows_i686_gnu" name = "windows_i686_gnu"
version = "0.32.0" version = "0.32.0"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6a711c68811799e017b6038e0922cb27a5e2f43a2ddb609fe0b6f3eeda9de615" checksum = "6a711c68811799e017b6038e0922cb27a5e2f43a2ddb609fe0b6f3eeda9de615"
[[package]]
name = "windows_i686_gnu"
version = "0.36.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "180e6ccf01daf4c426b846dfc66db1fc518f074baa793aa7d9b9aaeffad6a3b6"
[[package]] [[package]]
name = "windows_i686_msvc" name = "windows_i686_msvc"
version = "0.32.0" version = "0.32.0"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "146c11bb1a02615db74680b32a68e2d61f553cc24c4eb5b4ca10311740e44172" checksum = "146c11bb1a02615db74680b32a68e2d61f553cc24c4eb5b4ca10311740e44172"
[[package]]
name = "windows_i686_msvc"
version = "0.36.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e2e7917148b2812d1eeafaeb22a97e4813dfa60a3f8f78ebe204bcc88f12f024"
[[package]] [[package]]
name = "windows_x86_64_gnu" name = "windows_x86_64_gnu"
version = "0.32.0" version = "0.32.0"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c912b12f7454c6620635bbff3450962753834be2a594819bd5e945af18ec64bc" checksum = "c912b12f7454c6620635bbff3450962753834be2a594819bd5e945af18ec64bc"
[[package]]
name = "windows_x86_64_gnu"
version = "0.36.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4dcd171b8776c41b97521e5da127a2d86ad280114807d0b2ab1e462bc764d9e1"
[[package]] [[package]]
name = "windows_x86_64_msvc" name = "windows_x86_64_msvc"
version = "0.32.0" version = "0.32.0"
@ -564,10 +583,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "504a2476202769977a040c6364301a3f65d0cc9e3fb08600b2bda150a0488316" checksum = "504a2476202769977a040c6364301a3f65d0cc9e3fb08600b2bda150a0488316"
[[package]] [[package]]
name = "yaml-rust" name = "windows_x86_64_msvc"
version = "0.4.5" version = "0.36.1"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "56c1936c4cc7a1c9ab21a1ebb602eb942ba868cbd44a99cb7cdc5892335e1c85" checksum = "c811ca4a8c853ef420abd8592ba53ddbbac90410fab6903b3e79972a631f7680"
dependencies = [
"linked-hash-map",
]

6
Cargo.toml
View File

@ -8,16 +8,16 @@ license = "ISC"
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
[dependencies] [dependencies]
tokio = { version = "1.17.0", features = ["full"] } tokio = { version = "1.21.1", features = ["full"] }
futures = "0.3.21" futures = "0.3.21"
serde = { version = "1.0.136", features = ["derive"] } serde = { version = "1.0.136", features = ["derive"] }
serde_json = "1.0.79" serde_json = "1.0.79"
serde_yaml = "0.8.23" serde_yaml = "0.9.13"
simple-error = "0.2.3" simple-error = "0.2.3"
rand = "0.8.5" rand = "0.8.5"
[profile.release] [profile.release]
opt-level = 3 # Optimize for size. opt-level = 3 # Optimize for size.
lto = true # Enable Link Time Optimization lto = true # Enable Link Time Optimization
codegen-units = 1 # Reduce number of codegen units to increase optimizations. codegen-units = 1 # Reduce number of codegen units to increase optimizations.
panic = 'abort' # Abort on panic panic = 'abort' # Abort on panic

58
src/main.rs
View File

@ -1,6 +1,7 @@
mod listener; mod listener;
mod shutdown; mod shutdown;
mod tcp; mod tcp;
mod udp;
use serde::Deserialize; use serde::Deserialize;
use simple_error::bail; use simple_error::bail;
@ -14,14 +15,14 @@ use tokio::sync::{broadcast, RwLock};
#[derive(Debug, Deserialize)] #[derive(Debug, Deserialize)]
struct Target { struct Target {
udp: Option<bool>,
source: String, source: String,
targets: Vec<String>, targets: Vec<String>,
} }
#[derive(Debug, Deserialize)] #[derive(Debug, Deserialize)]
struct Config { struct Config {
tcp: Vec<Target>, mappings: Vec<Target>,
// udp: Vec<Target>,
} }
fn load_yaml(path: &Path) -> Result<Config, Box<dyn Error>> { fn load_yaml(path: &Path) -> Result<Config, Box<dyn Error>> {
@ -62,6 +63,7 @@ fn load_config() -> Result<Config, Box<dyn Error>> {
// #[derive(Debug)] // #[derive(Debug)]
struct ActiveListener { struct ActiveListener {
udp: bool,
notify_shutdown: broadcast::Sender<()>, notify_shutdown: broadcast::Sender<()>,
targets: Arc<RwLock<Vec<String>>>, targets: Arc<RwLock<Vec<String>>>,
} }
@ -75,27 +77,34 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
let config = load_config().expect("config not found"); let config = load_config().expect("config not found");
let mut required_listeners: HashSet<String> = HashSet::new(); let mut required_listeners: HashSet<String> = HashSet::new();
for target in config.tcp { for target in config.mappings {
required_listeners.insert(target.source.clone()); let mut source_str = "".to_owned();
if let Some(listener) = listeners.get(&target.source) { if target.udp == None || target.udp == Some(false) {
source_str.push_str("udp:");
} else {
source_str.push_str("tcp:");
}
source_str.push_str(&target.source);
required_listeners.insert(source_str.clone());
if let Some(listener) = listeners.get(&source_str) {
let mut invalid = false; let mut invalid = false;
{
let targets = listener.targets.read().await; let targets = listener.targets.read().await;
for t in &target.targets { for t in &target.targets {
if !targets.iter().any(|e| *e == *t) { if !targets.iter().any(|e| *e == *t) {
invalid = true;
break;
}
}
if !invalid {
for t in targets.iter() {
if !target.targets.iter().any(|e| *e == *t) {
invalid = true; invalid = true;
break; break;
} }
} }
if !invalid {
for t in targets.iter() {
if !target.targets.iter().any(|e| *e == *t) {
invalid = true;
break;
}
}
}
} }
if invalid { if invalid {
@ -109,6 +118,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
let listener = ActiveListener { let listener = ActiveListener {
notify_shutdown: notify_shutdown, notify_shutdown: notify_shutdown,
targets: Arc::new(RwLock::new(target.targets)), targets: Arc::new(RwLock::new(target.targets)),
udp: target.udp == None || target.udp == Some(false),
}; };
let l = listener::Listener { let l = listener::Listener {
@ -118,12 +128,18 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
}; };
tokio::spawn(async move { tokio::spawn(async move {
if let Err(err) = tcp::start_tcp_listener(l).await { if target.udp == None || target.udp == Some(false) {
println!("listener error: {}", err); if let Err(err) = tcp::start_tcp_listener(l).await {
println!("tcp listener error: {}", err);
}
} else {
if let Err(err) = udp::start_udp_listener(l).await {
println!("udp listener error: {}", err);
}
} }
}); });
listeners.insert(target.source, listener); listeners.insert(source_str, listener);
} }
} }

63
src/shutdown.rs
View File

@ -11,39 +11,48 @@ use tokio::sync::broadcast;
/// received or not. /// received or not.
#[derive(Debug)] #[derive(Debug)]
pub(crate) struct Shutdown { pub(crate) struct Shutdown {
/// `true` if the shutdown signal has been received /// `true` if the shutdown signal has been received
shutdown: bool, shutdown: bool,
/// The receive half of the channel used to listen for shutdown. /// The receive half of the channel used to listen for shutdown.
notify: broadcast::Receiver<()>, notify: broadcast::Receiver<()>,
} }
impl Shutdown { impl Shutdown {
/// Create a new `Shutdown` backed by the given `broadcast::Receiver`. /// Create a new `Shutdown` backed by the given `broadcast::Receiver`.
pub(crate) fn new(notify: broadcast::Receiver<()>) -> Shutdown { pub(crate) fn new(notify: broadcast::Receiver<()>) -> Shutdown {
Shutdown { Shutdown {
shutdown: false, shutdown: false,
notify, notify,
} }
} }
/// Returns `true` if the shutdown signal has been received. /// Returns `true` if the shutdown signal has been received.
// pub(crate) fn is_shutdown(&self) -> bool { // pub(crate) fn is_shutdown(&self) -> bool {
// self.shutdown // self.shutdown
// } // }
/// Receive the shutdown notice, waiting if necessary. /// Receive the shutdown notice, waiting if necessary.
pub(crate) async fn recv(&mut self) { pub(crate) async fn recv(&mut self) {
// If the shutdown signal has already been received, then return // If the shutdown signal has already been received, then return
// immediately. // immediately.
if self.shutdown { if self.shutdown {
return; return;
} }
// Cannot receive a "lag error" as only one value is ever sent. // Cannot receive a "lag error" as only one value is ever sent.
let _ = self.notify.recv().await; let _ = self.notify.recv().await;
// Remember that the signal has been received. // Remember that the signal has been received.
self.shutdown = true; self.shutdown = true;
} }
}
impl Clone for Shutdown {
fn clone(&self) -> Self {
Shutdown {
shutdown: self.shutdown,
notify: self.notify.resubscribe(),
}
}
} }

74
src/tcp.rs
View File

@ -5,54 +5,54 @@ use tokio::net::{TcpListener, TcpStream};
#[derive(Debug)] #[derive(Debug)]
struct TcpHandler { struct TcpHandler {
stream: TcpStream, stream: TcpStream,
target: String, target: String,
} }
impl TcpHandler { impl TcpHandler {
async fn run(&mut self) -> Result<(), Box<dyn Error>> { async fn run(&mut self) -> Result<(), Box<dyn Error>> {
let mut stream = TcpStream::connect(&self.target).await?; let mut stream = TcpStream::connect(&self.target).await?;
tokio::io::copy_bidirectional(&mut self.stream, &mut stream).await?; tokio::io::copy_bidirectional(&mut self.stream, &mut stream).await?;
return Ok(()); return Ok(());
} }
} }
pub(crate) async fn start_tcp_listener( pub(crate) async fn start_tcp_listener(
mut listener_config: Listener, mut listener_config: Listener,
) -> Result<(), Box<dyn Error>> { ) -> Result<(), Box<dyn Error>> {
println!("start listening on {}", &listener_config.source); println!("start listening on {}", &listener_config.source);
let listener = TcpListener::bind(&listener_config.source).await?; let listener = TcpListener::bind(&listener_config.source).await?;
loop { loop {
let (next_socket, _) = tokio::select! { let (next_socket, _) = tokio::select! {
res = listener.accept() => res?, res = listener.accept() => res?,
_ = listener_config.shutdown.recv() => { _ = listener_config.shutdown.recv() => {
println!("Exiting listener!"); println!("Exiting listener!");
return Ok(()); return Ok(());
} }
}; };
let targets = listener_config.targets.read().await; let targets = listener_config.targets.read().await;
let mut rng = rand::thread_rng(); let mut rng = rand::thread_rng();
let selected_target = targets.choose(&mut rng).unwrap(); let selected_target = targets.choose(&mut rng).unwrap();
println!( println!(
"new connection from {} forwarding to {}", "new connection from {} forwarding to {}",
next_socket.peer_addr()?, next_socket.peer_addr()?,
&selected_target &selected_target
); );
let mut handler = TcpHandler { let mut handler = TcpHandler {
stream: next_socket, stream: next_socket,
target: selected_target.clone(), target: selected_target.clone(),
}; };
tokio::spawn(async move { tokio::spawn(async move {
// Process the connection. If an error is encountered, log it. // Process the connection. If an error is encountered, log it.
if let Err(err) = handler.run().await { if let Err(err) = handler.run().await {
println!("connection error {}", err); println!("connection error {}", err);
} }
}); });
} }
} }

178
src/udp.rs
View File

@ -1,2 +1,176 @@
// For UDP to work, there is some magic required with matches the returning packets use crate::listener::Listener;
// back to the original sender. Idk. how to do that right now, but maybe some day. use crate::shutdown::Shutdown;
use rand::seq::SliceRandom;
use std::collections::HashMap;
use std::error::Error;
use std::sync::Arc;
use tokio::net::UdpSocket;
use tokio::sync::broadcast;
use tokio::sync::mpsc::{channel, Sender};
use tokio::sync::Mutex;
use tokio::time::Instant;
async fn create_dual_udpsocket(
bind_addr: &String,
) -> Result<(UdpSocket, UdpSocket), Box<dyn Error>> {
let listener_std = std::net::UdpSocket::bind(bind_addr)?;
let responder_std = listener_std.try_clone()?;
let listener = UdpSocket::from_std(listener_std)?;
let responder = UdpSocket::from_std(responder_std)?;
return Ok((listener, responder));
}
fn get_udp_background_send(socket: UdpSocket, mut exit: Shutdown) -> Sender<(Vec<u8>, String)> {
let (tx, mut rx) = channel::<(Vec<u8>, String)>(1);
tokio::spawn(async move {
loop {
let (buf, dest) = (tokio::select! {
res = rx.recv() => Some(res.unwrap()),
_ = exit.recv() => {
println!("Exiting listener!");
return;
}
})
.unwrap();
let to_send = buf.as_slice();
socket.send_to(to_send, &dest).await.expect(&format!(
"Failed to forward response from upstream server to client {}",
dest
));
}
});
return tx;
}
struct UdpHandler {
last_packet: Arc<Mutex<Instant>>,
kill: broadcast::Sender<()>,
target: String,
sender: Sender<(Vec<u8>, String)>,
}
impl UdpHandler {
async fn start(
target: String,
source: String,
sender: Sender<(Vec<u8>, String)>,
) -> Result<UdpHandler, Box<dyn Error>> {
// Kill Channel
let (tx, mut rx) = broadcast::channel::<()>(1);
let (listener, responder) = create_dual_udpsocket(&"0.0.0.0:0".to_owned()).await?;
let s = get_udp_background_send(responder, Shutdown::new(tx.subscribe()));
let last_packet = Arc::new(Mutex::new(Instant::now()));
let handler = UdpHandler {
kill: tx,
last_packet: last_packet.clone(),
// source: source.clone(),
target: target.clone(),
sender: s,
};
listener.connect(target).await?;
tokio::spawn(async move {
let mut buf = [0; 64 * 1024];
loop {
let (num_bytes, _) = tokio::select! {
res = listener.recv_from(&mut buf) => res.unwrap(),
_ = rx.recv() => {
// FIXME: Source of memory leaks?
return ;
}
};
let mut n = last_packet.lock().await;
*n = Instant::now();
sender
.send((buf[0..num_bytes].to_vec(), source.clone()))
.await
.expect(&format!("Failed to send answer to sender {}", source));
}
});
return Ok(handler);
}
async fn exit(&self) -> Result<(), Box<dyn Error>> {
self.kill.send(())?;
Ok(())
}
async fn on_packet(&mut self, pkg: Vec<u8>) -> Result<(), Box<dyn Error>> {
let mut n = self.last_packet.lock().await;
*n = Instant::now();
self.sender.send((pkg, self.target.clone())).await?;
return Ok(());
}
}
pub(crate) async fn start_udp_listener(
mut listener_config: Listener,
) -> Result<(), Box<dyn Error>> {
println!("start listening on {}", &listener_config.source);
let (listener, responder) =
create_dual_udpsocket(&listener_config.source)
.await
.expect(&format!(
"Failed to clone primary listening address socket {}",
&listener_config.source,
));
let sender = get_udp_background_send(responder, listener_config.shutdown.clone());
let mut connections: HashMap<String, UdpHandler> = HashMap::new();
let mut buf = [0; 64 * 1024];
loop {
let (num_bytes, src_addr) = tokio::select! {
res = listener.recv_from(&mut buf) => res?,
_ = listener_config.shutdown.recv() => {
println!("Exiting listener!");
break;
}
};
let addr = src_addr.to_string();
let handler_opt = connections.get_mut(&addr);
let vec = buf[0..num_bytes].to_vec();
let handler = match handler_opt {
Some(handler) => handler,
None => {
let targets = listener_config.targets.read().await;
let selected_target = {
let mut rng = rand::thread_rng();
targets.choose(&mut rng).unwrap()
};
let handler =
UdpHandler::start(selected_target.clone(), addr.clone(), sender.clone())
.await?;
connections.insert(addr.clone(), handler);
connections.get_mut(&addr).unwrap()
}
};
match handler.on_packet(vec).await {
Ok(_) => (),
Err(err) => {
println!("Failed to forward request from client to server {}", err);
return Ok(());
}
}
}
for handler in connections.values() {
handler.exit().await?;
}
Ok(())
}