From de3eda92006bad897f67cfda2fc6e68ef7f0ae98 Mon Sep 17 00:00:00 2001
From: Riz Ashraf <reazul.ashraf@inseinc.com>
Date: Thu, 8 Oct 2026 22:12:48 +0100
Subject: [PATCH] Fix --wait timeout duration logic and iterator borrows

---
 src/main.rs | 21 ++++++++++++++++++---
 1 file changed, 18 insertions(+), 3 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 987246e..da16654 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -9,13 +9,25 @@ fn main() {
     let mut port_str = String::new();
     let mut payload = String::new();
     let mut positional_args = 0;
+    let mut timeout_secs = 30; // default 30s timeout
 
     // 1. Iterate args safely, avoiding heavy Vec allocations
-    for arg in env::args().skip(1) {
+    let args: Vec<String> = env::args().skip(1).collect();
+    let mut args_iter = args.into_iter().peekable();
+    
+    while let Some(arg) = args_iter.next() {
         if arg == "--debug" {
             debug = true;
         } else if arg == "--wait" {
             wait = true;
+            // Parse optional timeout if present, otherwise default to 30s
+            // Peek at next arg to see if it's a number
+            if let Some(next_arg) = args_iter.peek() {
+                if let Ok(secs) = next_arg.parse::<u64>() {
+                    timeout_secs = secs;
+                    args_iter.next(); // Consume the timeout arg
+                }
+            }
         } else if positional_args == 0 {
             port_str = arg;
             positional_args += 1;
@@ -61,8 +73,11 @@ fn main() {
                     process::exit(1);
                 }
             } else if wait {
-                // Set a timeout of 10 minutes for waiting
-                let _ = socket.set_read_timeout(Some(Duration::from_secs(600)));
+                // Set a dynamic timeout for waiting
+                if timeout_secs == 0 {
+                    if debug { eprintln!("cannot set a 0 duration timeout"); process::exit(1); } else { process::exit(1); }
+                }
+                let _ = socket.set_read_timeout(Some(Duration::from_secs(timeout_secs)));
                 let mut buf = [0; 1024];
                 match socket.recv_from(&mut buf) {
                     Ok((size, _src)) => {