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

src/main.rs

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