From 0be26599cada3a6752952e39b105b041170074b3 Mon Sep 17 00:00:00 2001
From: Brad Fitzpatrick <bradfitz@tailscale.com>
Date: Thu, 2 Sep 2021 10:16:31 -0700
Subject: [PATCH] cmd/derper: refactor STUN path for testing, add serverSTUN
 benchmark

Real goal is to eliminate some allocs in the STUN path, but that requires
work in the standard library.

See comments in #2783.

Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>
---
 cmd/derper/derper.go      | 56 ++++++++++++++++++++++-----------------
 cmd/derper/derper_test.go | 34 ++++++++++++++++++++++++
 2 files changed, 66 insertions(+), 24 deletions(-)

diff --git a/cmd/derper/derper.go b/cmd/derper/derper.go
index 3c3cf7dce..09635af4c 100644
--- a/cmd/derper/derper.go
+++ b/cmd/derper/derper.go
@@ -49,6 +49,26 @@ var (
 	verifyClients = flag.Bool("verify-clients", false, "verify clients to this DERP server through a local tailscaled instance.")
 )
 
+var (
+	stats           = new(metrics.Set)
+	stunDisposition = &metrics.LabelMap{Label: "disposition"}
+	stunAddrFamily  = &metrics.LabelMap{Label: "family"}
+
+	stunReadError  = stunDisposition.Get("read_error")
+	stunNotSTUN    = stunDisposition.Get("not_stun")
+	stunWriteError = stunDisposition.Get("write_error")
+	stunSuccess    = stunDisposition.Get("success")
+
+	stunIPv4 = stunAddrFamily.Get("ipv4")
+	stunIPv6 = stunAddrFamily.Get("ipv6")
+)
+
+func init() {
+	stats.Set("counter_requests", stunDisposition)
+	stats.Set("counter_addrfamily", stunAddrFamily)
+	expvar.Publish("stun", stats)
+}
+
 type config struct {
 	PrivateKey wgkey.Private
 }
@@ -256,39 +276,27 @@ func serveSTUN(host string) {
 		log.Fatalf("failed to open STUN listener: %v", err)
 	}
 	log.Printf("running STUN server on %v", pc.LocalAddr())
+	serverSTUNListener(context.Background(), pc.(*net.UDPConn))
+}
 
+func serverSTUNListener(ctx context.Context, pc *net.UDPConn) {
+	var buf [64 << 10]byte
 	var (
-		stats           = new(metrics.Set)
-		stunDisposition = &metrics.LabelMap{Label: "disposition"}
-		stunAddrFamily  = &metrics.LabelMap{Label: "family"}
-
-		stunReadError  = stunDisposition.Get("read_error")
-		stunNotSTUN    = stunDisposition.Get("not_stun")
-		stunWriteError = stunDisposition.Get("write_error")
-		stunSuccess    = stunDisposition.Get("success")
-
-		stunIPv4 = stunAddrFamily.Get("ipv4")
-		stunIPv6 = stunAddrFamily.Get("ipv6")
+		n   int
+		ua  *net.UDPAddr
+		err error
 	)
-	stats.Set("counter_requests", stunDisposition)
-	stats.Set("counter_addrfamily", stunAddrFamily)
-	expvar.Publish("stun", stats)
-
-	var buf [64 << 10]byte
 	for {
-		n, addr, err := pc.ReadFrom(buf[:])
+		n, ua, err = pc.ReadFromUDP(buf[:])
 		if err != nil {
+			if ctx.Err() != nil {
+				return
+			}
 			log.Printf("STUN ReadFrom: %v", err)
 			time.Sleep(time.Second)
 			stunReadError.Add(1)
 			continue
 		}
-		ua, ok := addr.(*net.UDPAddr)
-		if !ok {
-			log.Printf("STUN unexpected address %T %v", addr, addr)
-			stunReadError.Add(1)
-			continue
-		}
 		pkt := buf[:n]
 		if !stun.Is(pkt) {
 			stunNotSTUN.Add(1)
@@ -305,7 +313,7 @@ func serveSTUN(host string) {
 			stunIPv6.Add(1)
 		}
 		res := stun.Response(txid, ua.IP, uint16(ua.Port))
-		_, err = pc.WriteTo(res, addr)
+		_, err = pc.WriteTo(res, ua)
 		if err != nil {
 			stunWriteError.Add(1)
 		} else {
diff --git a/cmd/derper/derper_test.go b/cmd/derper/derper_test.go
index e26864a32..4adab4e8b 100644
--- a/cmd/derper/derper_test.go
+++ b/cmd/derper/derper_test.go
@@ -6,7 +6,10 @@ package main
 
 import (
 	"context"
+	"net"
 	"testing"
+
+	"tailscale.com/net/stun"
 )
 
 func TestProdAutocertHostPolicy(t *testing.T) {
@@ -31,5 +34,36 @@ func TestProdAutocertHostPolicy(t *testing.T) {
 			t.Errorf("f(%q) = %v; want %v", tt.in, got, tt.wantOK)
 		}
 	}
+}
+
+func BenchmarkServerSTUN(b *testing.B) {
+	b.ReportAllocs()
+	pc, err := net.ListenPacket("udp", "127.0.0.1:0")
+	if err != nil {
+		b.Fatal(err)
+	}
+	defer pc.Close()
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	go serverSTUNListener(ctx, pc.(*net.UDPConn))
+	addr := pc.LocalAddr().(*net.UDPAddr)
+
+	var resBuf [1500]byte
+	cc, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.ParseIP("127.0.0.1")})
+	if err != nil {
+		b.Fatal(err)
+	}
+
+	tx := stun.NewTxID()
+	req := stun.Request(tx)
+	for i := 0; i < b.N; i++ {
+		if _, err := cc.WriteToUDP(req, addr); err != nil {
+			b.Fatal(err)
+		}
+		_, _, err := cc.ReadFromUDP(resBuf[:])
+		if err != nil {
+			b.Fatal(err)
+		}
+	}
 
 }