Merge dc59464731 into 45f0721530

cmd/stunc/stunc.go

@@ -5,53 +5,121 @@
 package main
 
 import (
+	"encoding/json"
+	"flag"
+	"fmt"
 	"log"
 	"net"
+	"net/netip"
 	"os"
+	"time"
 
 	"tailscale.com/net/stun"
 )
 
+var (
+	rate    = flag.Duration("rate", time.Second, "rate at which to send probes (0 means as fast as possible)")
+	timeout = flag.Duration("timeout", time.Second, "time to wait for a response")
+	reuse   = flag.Bool("reuse", true, "reuse the same UDP socket for each probe")
+	jsonout = flag.Bool("json", false, "output in JSON format (default is human-readable)")
+)
+
 func main() {
-	log.SetFlags(0)
-
-	if len(os.Args) != 2 {
-		log.Fatalf("usage: %s <hostname>", os.Args[0])
+	flag.Usage = func() {
+		fmt.Printf("usage: %s [flags] <hostname>", os.Args[0])
+		flag.PrintDefaults()
 	}
-	host := os.Args[1]
 
-	uaddr, err := net.ResolveUDPAddr("udp", host+":3478")
-	if err != nil {
-		log.Fatal(err)
+	flag.Parse()
+	if flag.NArg() != 1 {
+		flag.Usage()
+		os.Exit(2)
 	}
-	c, err := net.ListenUDP("udp", nil)
+
+	host := flag.Args()[0]
+
+	naddr, err := net.ResolveIPAddr("ip", host)
 	if err != nil {
 		log.Fatal(err)
 	}
 
-	txID := stun.NewTxID()
-	req := stun.Request(txID)
-
-	_, err = c.WriteToUDP(req, uaddr)
+	nip, err := netip.ParseAddr(naddr.String())
 	if err != nil {
 		log.Fatal(err)
 	}
 
-	var buf [1024]byte
-	n, raddr, err := c.ReadFromUDPAddrPort(buf[:])
-	if err != nil {
-		log.Fatal(err)
+	uaddr := netip.AddrPortFrom(nip, 3478)
+
+	var c *net.UDPConn
+
+	var print = func(result map[string]string) {
+		r := result["status"]
+		if result["status"] == "ok" {
+			r = fmt.Sprintf("%s; %s < %s in %s", result["status"], result["from"], result["stun"], result["dur"])
+		}
+		fmt.Printf("%s > %s; %s\n", result["local"], result["to"], r)
+	}
+	if *jsonout {
+		j := json.NewEncoder(os.Stdout)
+		print = func(result map[string]string) {
+			if err := j.Encode(result); err != nil {
+				log.Fatal(err)
+			}
+		}
 	}
 
-	tid, saddr, err := stun.ParseResponse(buf[:n])
-	if err != nil {
-		log.Fatal(err)
-	}
-	if tid != txID {
-		log.Fatalf("txid mismatch: got %v, want %v", tid, txID)
-	}
+	for {
+		if c == nil || !*reuse {
+			if c != nil {
+				c.Close()
+			}
 
-	log.Printf("sent addr: %v", uaddr)
-	log.Printf("from addr: %v", raddr)
-	log.Printf("stun addr: %v", saddr)
+			c, err = net.ListenUDP("udp", nil)
+			if err != nil {
+				log.Fatal(err)
+			}
+		}
+
+		result := map[string]string{}
+		result["to"] = uaddr.String()
+		result["local"] = c.LocalAddr().String()
+
+		txID := stun.NewTxID()
+		req := stun.Request(txID)
+
+		t0 := time.Now()
+		result["at"] = t0.Format(time.RFC3339Nano)
+		_, err = c.WriteToUDPAddrPort(req, uaddr)
+		if err != nil {
+			log.Fatal(err)
+		}
+
+		c.SetReadDeadline(t0.Add(*timeout))
+		var buf [1024]byte
+		n, raddr, err := c.ReadFromUDPAddrPort(buf[:])
+		if err != nil {
+			if nerr, ok := err.(net.Error); ok && nerr.Timeout() {
+				result["status"] = "timeout"
+				print(result)
+				continue
+			}
+			log.Fatalf("%#v", err)
+		}
+		result["from"] = raddr.String()
+		result["dur"] = time.Since(t0).String()
+
+		tid, saddr, err := stun.ParseResponse(buf[:n])
+		if err != nil {
+			log.Fatal(err)
+		}
+		result["stun"] = saddr.String()
+		if tid != txID {
+			result["status"] = "badtxid"
+		} else {
+			result["status"] = "ok"
+		}
+
+		print(result)
+		time.Sleep(time.Until(t0.Add(*rate)))
+	}
 }