From c2761162a002e65e2305f7570b2c54c561ac151f Mon Sep 17 00:00:00 2001
From: Mike O'Driscoll <quick.job4920@hidemail.ca>
Date: Fri, 6 Dec 2024 14:27:52 -0500
Subject: [PATCH] cmd/stunc: enforce read timeout deadline (#14309)

Make argparsing use flag for adding a new
parameter that requires parsing.

Enforce a read timeout deadline waiting for response
from the stun server provided in the args. Otherwise
the program will never exit.

Fixes #14267

Signed-off-by: Mike O'Driscoll <mikeo@tailscale.com>
---
 cmd/stunc/stunc.go | 34 +++++++++++++++++++++++++++-------
 1 file changed, 27 insertions(+), 7 deletions(-)

diff --git a/cmd/stunc/stunc.go b/cmd/stunc/stunc.go
index 9743a3300..c4b2eedd3 100644
--- a/cmd/stunc/stunc.go
+++ b/cmd/stunc/stunc.go
@@ -5,24 +5,40 @@
 package main
 
 import (
+	"flag"
 	"log"
 	"net"
 	"os"
 	"strconv"
+	"time"
 
 	"tailscale.com/net/stun"
 )
 
 func main() {
 	log.SetFlags(0)
-
-	if len(os.Args) < 2 || len(os.Args) > 3 {
-		log.Fatalf("usage: %s <hostname> [port]", os.Args[0])
-	}
-	host := os.Args[1]
+	var host string
 	port := "3478"
-	if len(os.Args) == 3 {
-		port = os.Args[2]
+
+	var readTimeout time.Duration
+	flag.DurationVar(&readTimeout, "timeout", 3*time.Second, "response wait timeout")
+
+	flag.Parse()
+
+	values := flag.Args()
+	if len(values) < 1 || len(values) > 2 {
+		log.Printf("usage: %s <hostname> [port]", os.Args[0])
+		flag.PrintDefaults()
+		os.Exit(1)
+	} else {
+		for i, value := range values {
+			switch i {
+			case 0:
+				host = value
+			case 1:
+				port = value
+			}
+		}
 	}
 	_, err := strconv.ParseUint(port, 10, 16)
 	if err != nil {
@@ -46,6 +62,10 @@ func main() {
 		log.Fatal(err)
 	}
 
+	err = c.SetReadDeadline(time.Now().Add(readTimeout))
+	if err != nil {
+		log.Fatal(err)
+	}
 	var buf [1024]byte
 	n, raddr, err := c.ReadFromUDPAddrPort(buf[:])
 	if err != nil {