From 991f26e5f2ead55b52778532665c62c76e3eab6c Mon Sep 17 00:00:00 2001
From: malmert <marcus.almert@gmail.com>
Date: Fri, 6 Mar 2026 20:31:32 +0100
Subject: [PATCH] fix traffic gen keypress interrupt

---
 tools/traffic-generator.sh | 54 +++++++++++++++++++++++++++-----------
 1 file changed, 38 insertions(+), 16 deletions(-)

diff --git a/tools/traffic-generator.sh b/tools/traffic-generator.sh
index 8bab0de..ee9d212 100755
--- a/tools/traffic-generator.sh
+++ b/tools/traffic-generator.sh
@@ -101,6 +101,27 @@ cleanup() {
   fi
 }
 
+on_interrupt() {
+  STOP=1
+}
+
+run_interruptible() {
+  "$@" &
+  local pid=$!
+
+  while kill -0 "$pid" >/dev/null 2>&1; do
+    if [[ "$STOP" -eq 1 ]]; then
+      kill "$pid" >/dev/null 2>&1 || true
+      wait "$pid" >/dev/null 2>&1 || true
+      return 130
+    fi
+    sleep 0.1
+  done
+
+  wait "$pid" >/dev/null 2>&1 || true
+  return 0
+}
+
 parse_args() {
   while [[ $# -gt 0 ]]; do
     case "$1" in
@@ -190,7 +211,7 @@ run_arp() {
   has_cmd arping || return 0
   for _ in $(seq 1 "$REQUESTS"); do
     [[ "$STOP" -eq 1 ]] && return
-    arping -I "$IFACE" -c 1 "$TARGET" >/dev/null 2>&1 || true
+    run_interruptible arping -I "$IFACE" -c 1 "$TARGET" || true
     maybe_wait
   done
 }
@@ -199,7 +220,7 @@ run_icmp() {
   has_cmd ping || return 0
   for _ in $(seq 1 "$REQUESTS"); do
     [[ "$STOP" -eq 1 ]] && return
-    ping -c 1 -W 1 "$TARGET" >/dev/null 2>&1 || true
+    run_interruptible ping -c 1 -W 1 "$TARGET" || true
     maybe_wait
   done
 }
@@ -208,7 +229,7 @@ run_http() {
   has_cmd curl || return 0
   for _ in $(seq 1 "$REQUESTS"); do
     [[ "$STOP" -eq 1 ]] && return
-    curl -s "http://${TARGET}:${HTTP_PORT}/" >/dev/null 2>&1 || true
+    run_interruptible curl -s "http://${TARGET}:${HTTP_PORT}/" || true
     maybe_wait
   done
 }
@@ -217,7 +238,7 @@ run_https() {
   has_cmd curl || return 0
   for _ in $(seq 1 "$REQUESTS"); do
     [[ "$STOP" -eq 1 ]] && return
-    curl -sk "https://${TARGET}:${HTTPS_PORT}/" >/dev/null 2>&1 || true
+    run_interruptible curl -sk "https://${TARGET}:${HTTPS_PORT}/" || true
     maybe_wait
   done
 }
@@ -226,7 +247,7 @@ run_dns() {
   if has_cmd dig; then
     for _ in $(seq 1 "$REQUESTS"); do
       [[ "$STOP" -eq 1 ]] && return
-      dig @"$TARGET" -p "$DNS_PORT" example.com A +tries=1 +time=1 +short >/dev/null 2>&1 || true
+      run_interruptible dig @"$TARGET" -p "$DNS_PORT" example.com A +tries=1 +time=1 +short || true
       maybe_wait
     done
   fi
@@ -236,7 +257,7 @@ run_tcp_echo() {
   has_cmd nc || return 0
   for i in $(seq 1 "$REQUESTS"); do
     [[ "$STOP" -eq 1 ]] && return
-    printf 'tcp-echo-%s\n' "$i" | nc -w1 "$TARGET" "$TCP_ECHO_PORT" >/dev/null 2>&1 || true
+    printf 'tcp-echo-%s\n' "$i" | run_interruptible nc -w1 "$TARGET" "$TCP_ECHO_PORT" || true
     maybe_wait
   done
 }
@@ -245,26 +266,26 @@ run_udp_echo() {
   has_cmd nc || return 0
   for i in $(seq 1 "$REQUESTS"); do
     [[ "$STOP" -eq 1 ]] && return
-    printf 'udp-echo-%s\n' "$i" | nc -u -w1 "$TARGET" "$UDP_ECHO_PORT" >/dev/null 2>&1 || true
+    printf 'udp-echo-%s\n' "$i" | run_interruptible nc -u -w1 "$TARGET" "$UDP_ECHO_PORT" || true
     maybe_wait
   done
 }
 
 run_tcp_syn() {
   if has_cmd nping; then
-    nping --interface "$IFACE" --tcp -p "$TCP_ECHO_PORT" --flags syn -c "$REQUESTS" "$TARGET" >/dev/null 2>&1 || true
+    run_interruptible nping --interface "$IFACE" --tcp -p "$TCP_ECHO_PORT" --flags syn -c "$REQUESTS" "$TARGET" || true
   fi
 }
 
 run_tcp_flags() {
   if has_cmd nping; then
-    nping --interface "$IFACE" --tcp -p "$TCP_ECHO_PORT" --flags syn,ack,fin,rst,psh,urg -c "$REQUESTS" "$TARGET" >/dev/null 2>&1 || true
+    run_interruptible nping --interface "$IFACE" --tcp -p "$TCP_ECHO_PORT" --flags syn,ack,fin,rst,psh,urg -c "$REQUESTS" "$TARGET" || true
   fi
 }
 
 run_udp_raw() {
   if has_cmd nping; then
-    nping --interface "$IFACE" --udp -p "$UDP_ECHO_PORT" --data-length 256 -c "$REQUESTS" "$TARGET" >/dev/null 2>&1 || true
+    run_interruptible nping --interface "$IFACE" --udp -p "$UDP_ECHO_PORT" --data-length 256 -c "$REQUESTS" "$TARGET" || true
   fi
 }
 
@@ -272,19 +293,19 @@ run_traceroute() {
   has_cmd traceroute || return 0
   for _ in $(seq 1 "$REQUESTS"); do
     [[ "$STOP" -eq 1 ]] && return
-    traceroute -n -m 5 "$TARGET" >/dev/null 2>&1 || true
+    run_interruptible traceroute -n -m 5 "$TARGET" || true
     maybe_wait
   done
 }
 
 run_iperf_tcp() {
   has_cmd iperf3 || return 0
-  iperf3 -c "$TARGET" -p "$IPERF_TCP_PORT" -t 10 >/dev/null 2>&1 || true
+  run_interruptible iperf3 -c "$TARGET" -p "$IPERF_TCP_PORT" -t 10 || true
 }
 
 run_iperf_udp() {
   has_cmd iperf3 || return 0
-  iperf3 -c "$TARGET" -p "$IPERF_UDP_PORT" -u -b 100M -t 10 >/dev/null 2>&1 || true
+  run_interruptible iperf3 -c "$TARGET" -p "$IPERF_UDP_PORT" -u -b 100M -t 10 || true
 }
 
 run_port_probe() {
@@ -292,7 +313,7 @@ run_port_probe() {
   has_cmd nc || return 0
   for _ in $(seq 1 "$REQUESTS"); do
     [[ "$STOP" -eq 1 ]] && return
-    nc -z -w1 "$TARGET" "$port" >/dev/null 2>&1 || true
+    run_interruptible nc -z -w1 "$TARGET" "$port" || true
     maybe_wait
   done
 }
@@ -304,7 +325,7 @@ run_ntp_probe() {
   has_cmd nc || return 0
   for _ in $(seq 1 "$REQUESTS"); do
     [[ "$STOP" -eq 1 ]] && return
-    printf 'ntp?\n' | nc -u -w1 "$TARGET" 123 >/dev/null 2>&1 || true
+    printf 'ntp?\n' | run_interruptible nc -u -w1 "$TARGET" 123 || true
     maybe_wait
   done
 }
@@ -312,7 +333,7 @@ run_mdns_query() {
   if has_cmd dig; then
     for _ in $(seq 1 "$REQUESTS"); do
       [[ "$STOP" -eq 1 ]] && return
-      dig @224.0.0.251 -p 5353 _services._dns-sd._udp.local PTR +tries=1 +time=1 +short >/dev/null 2>&1 || true
+      run_interruptible dig @224.0.0.251 -p 5353 _services._dns-sd._udp.local PTR +tries=1 +time=1 +short || true
       maybe_wait
     done
   fi
@@ -357,6 +378,7 @@ main() {
   fi
 
   start_keypress_listener
+  trap on_interrupt INT TERM USR1
   trap cleanup EXIT
 
   log "Target=${TARGET}, iface=${IFACE}, requests=${REQUESTS}, protocols=${PROTOCOLS_CSV}, stop_on_keypress=${STOP_ON_KEYPRESS}"