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 <[email protected]>

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 {