1 files changed, 16 insertions, 2 deletions
diff --git a/src/main.rs b/src/main.rs
index 96b6873..17a20be 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -11,6 +11,10 @@ struct Args {
     #[arg(short, long)]
     bind: String,
 
+    /// The local endpoint to read cleartext datagrams from & write to
+    #[arg(short, long)]
+    local: String,
+
     /// The destination to send obfuscated UDP datagrams to
     #[arg(short, long)]
     sendto: String,
@@ -28,8 +32,9 @@ async fn main() {
     let sock = UdpSocket::bind(&args.bind).await
         .expect("failed to bind UDP endpoint");
     let mut buf = [0; 65536];
+    let local = args.local.parse().expect("failed to parse local address");
     loop {
-        if let Ok((len, _)) = sock.recv_from(&mut buf).await {
+        if let Ok((len, addr)) = sock.recv_from(&mut buf).await {
             let mut i = 0;
             let n = key.len();
             for j in 0..len {
@@ -39,7 +44,16 @@ async fn main() {
                     i = 0;
                 }
             }
-            let _ = sock.send_to(&buf[..len], &args.sendto).await;
+            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;
         }
     }
 }