1 files changed, 82 insertions, 21 deletions

diff --git a/src/main.rs b/src/main.rs
index 17a20be..06f0c05 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,8 +1,11 @@
+use std::io;
+use std::net::{SocketAddr, ToSocketAddrs};
 
-
-use clap::Parser;
+use clap::{arg, Parser};
 use tokio::net::UdpSocket;
 
+use crate::Destination::{Address, SockAddr};
+
 /// Simple program to greet a person
 #[derive(Parser, Debug)]
 #[command(author, version, about, long_about = None)]
@@ -17,13 +20,20 @@ struct Args {
 
     /// The destination to send obfuscated UDP datagrams to
     #[arg(short, long)]
-    sendto: String,
+    sendto: Option<String>,
 
     /// The hex representation of XOR key to use
     #[arg(short, long)]
     key: String,
 }
 
+#[derive(Debug, Eq, PartialEq)]
+enum Destination<'a> {
+    None,
+    Address(&'a String),
+    SockAddr(SocketAddr),
+}
+
 #[tokio::main]
 async fn main() {
     let args = Args::parse();
@@ -33,27 +43,78 @@ async fn main() {
         .expect("failed to bind UDP endpoint");
     let mut buf = [0; 65536];
     let local = args.local.parse().expect("failed to parse local address");
+    let mut sendto = Destination::None;
+    let args_sendto: String;
+    if let Some(s) = args.sendto {
+        args_sendto = s;
+        sendto = Address(&args_sendto);
+    }
     loop {
-        if let Ok((len, addr)) = sock.recv_from(&mut buf).await {
-            let mut i = 0;
-            let n = key.len();
-            for j in 0..len {
-                buf[j] ^= key[i];
-                i += 1;
-                if i == n {
-                    i = 0;
+        match sock.recv_from(&mut buf).await {
+            Ok((len, addr)) => {
+                eprintln!("recv_from OK, len {}, addr {}", len, addr);
+                let mut i = 0;
+                let n = key.len();
+                for j in 0..len {
+                    buf[j] ^= key[i];
+                    i += 1;
+                    if i == n {
+                        i = 0;
+                    }
                 }
+                if sendto == Destination::None && addr != local {
+                    // assume the packet is coming from target
+                    // update missing sendto parameter
+                    sendto = SockAddr(addr);
+                }
+                let sendto = match sendto {
+                    Destination::None => Destination::None,
+                    Address(addr) => match resolve_addr(&addr) {
+                        Ok(addr) => SockAddr(addr),
+                        Err(why) => {
+                            eprintln!("Failed to resolve address {}: {}", addr, why);
+                            Destination::None
+                        }
+                    },
+                    SockAddr(addr) => SockAddr(addr),
+                };
+                let target = if let SockAddr(sendto_addr) = sendto {
+                    if addr == sendto_addr {
+                        // remote -> local
+                        Some(local)
+                    } else {
+                        Some(sendto_addr)
+                    }
+                } else {
+                    None
+                };
+                match target {
+                    None => {
+                        eprintln!("cannot send_to: destination not determined");
+                        continue;
+                    }
+                    Some(target) => {
+                        if let Err(why) = sock.send_to(&buf[..len], target).await {
+                            eprintln!("send_to failed, dest: {}, error: {}", target, why);
+                        } else {
+                            eprintln!("send_to OK, len: {}, dest: {}", len, target);
+                        }
+                    }
+                }
+            }
+            Err(why) => {
+                eprintln!("recv_from failed: {}", why);
             }
-            let target = if addr == local {
-                // packet comes from local endpoint
-                // send to target
-                &args.sendto
-            } else {
-                // packets coming from target
-                // send to local
-                &args.local
-            };
-            let _ = sock.send_to(&buf[..len], target).await;
         }
     }
 }
+
+fn resolve_addr(addr: &String) -> io::Result<SocketAddr> {
+    match addr.to_socket_addrs() {
+        Ok(mut addrs) => match addrs.next() {
+            None => Err(io::Error::new(io::ErrorKind::Other, "no available address")),
+            Some(addr) => Ok(addr),
+        },
+        Err(why) => Err(why),
+    }
+}
\ No newline at end of file