cmd/xdpderper,derp/xdp: implement mode that drops STUN packets (#12527)

This is useful during maintenance as a method for shedding home client
load.

Updates tailscale/corp#20689

Signed-off-by: Jordan Whited <jordan@tailscale.com>

9 files changed, 128 insertions, 10 deletions

diff --git a/cmd/xdpderper/xdpderper.go b/cmd/xdpderper/xdpderper.go
index 8d52e9353..1af9c9d5a 100644
--- a/cmd/xdpderper/xdpderper.go
+++ b/cmd/xdpderper/xdpderper.go
@@ -5,6 +5,7 @@ package main
 
 import (
 	"flag"
+	"io"
 	"log"
 	"net/http"
 	"os"
@@ -57,7 +58,26 @@ func main() {
 	log.Println("XDP STUN server started")
 
 	mux := http.NewServeMux()
-	tsweb.Debugger(mux)
+	debug := tsweb.Debugger(mux)
+	debug.KVFunc("Drop STUN", func() any {
+		return server.GetDropSTUN()
+	})
+	debug.Handle("drop-stun-on", "Drop STUN packets", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		err := server.SetDropSTUN(true)
+		if err != nil {
+			http.Error(w, err.Error(), 500)
+		} else {
+			io.WriteString(w, "STUN packets are now being dropped.")
+		}
+	}))
+	debug.Handle("drop-stun-off", "Handle STUN packets", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		err := server.SetDropSTUN(false)
+		if err != nil {
+			http.Error(w, err.Error(), 500)
+		} else {
+			io.WriteString(w, "STUN packets are now being handled.")
+		}
+	}))
 	errCh := make(chan error, 1)
 	go func() {
 		err := http.ListenAndServe(*flagHTTP, mux)
diff --git a/derp/xdp/bpf_bpfeb.go b/derp/xdp/bpf_bpfeb.go
index 1883d52fe..e20228731 100644
--- a/derp/xdp/bpf_bpfeb.go
+++ b/derp/xdp/bpf_bpfeb.go
@@ -12,7 +12,10 @@ import (
 	"github.com/cilium/ebpf"
 )
 
-type bpfConfig struct{ DstPort uint16 }
+type bpfConfig struct {
+	DstPort  uint16
+	DropStun uint16
+}
 
 type bpfCounterKeyAf uint32
 
@@ -46,7 +49,8 @@ const (
 	bpfCounterKeyProgEndCOUNTER_KEY_END_INVALID_IP_CSUM            bpfCounterKeyProgEnd = 3
 	bpfCounterKeyProgEndCOUNTER_KEY_END_NOT_STUN_PORT              bpfCounterKeyProgEnd = 4
 	bpfCounterKeyProgEndCOUNTER_KEY_END_INVALID_SW_ATTR_VAL        bpfCounterKeyProgEnd = 5
-	bpfCounterKeyProgEndCOUNTER_KEY_END_LEN                        bpfCounterKeyProgEnd = 6
+	bpfCounterKeyProgEndCOUNTER_KEY_END_DROP_STUN                  bpfCounterKeyProgEnd = 6
+	bpfCounterKeyProgEndCOUNTER_KEY_END_LEN                        bpfCounterKeyProgEnd = 7
 )
 
 type bpfCountersKey struct {
diff --git a/derp/xdp/bpf_bpfeb.o b/derp/xdp/bpf_bpfeb.o
index 06ff73fc1..64bde142d 100644
--- a/derp/xdp/bpf_bpfeb.o
+++ b/derp/xdp/bpf_bpfeb.o
diff --git a/derp/xdp/bpf_bpfel.go b/derp/xdp/bpf_bpfel.go
index 8cb7dec56..aab06b041 100644
--- a/derp/xdp/bpf_bpfel.go
+++ b/derp/xdp/bpf_bpfel.go
@@ -12,7 +12,10 @@ import (
 	"github.com/cilium/ebpf"
 )
 
-type bpfConfig struct{ DstPort uint16 }
+type bpfConfig struct {
+	DstPort  uint16
+	DropStun uint16
+}
 
 type bpfCounterKeyAf uint32
 
@@ -46,7 +49,8 @@ const (
 	bpfCounterKeyProgEndCOUNTER_KEY_END_INVALID_IP_CSUM            bpfCounterKeyProgEnd = 3
 	bpfCounterKeyProgEndCOUNTER_KEY_END_NOT_STUN_PORT              bpfCounterKeyProgEnd = 4
 	bpfCounterKeyProgEndCOUNTER_KEY_END_INVALID_SW_ATTR_VAL        bpfCounterKeyProgEnd = 5
-	bpfCounterKeyProgEndCOUNTER_KEY_END_LEN                        bpfCounterKeyProgEnd = 6
+	bpfCounterKeyProgEndCOUNTER_KEY_END_DROP_STUN                  bpfCounterKeyProgEnd = 6
+	bpfCounterKeyProgEndCOUNTER_KEY_END_LEN                        bpfCounterKeyProgEnd = 7
 )
 
 type bpfCountersKey struct {
diff --git a/derp/xdp/bpf_bpfel.o b/derp/xdp/bpf_bpfel.o
index b4653efd9..04b909ac7 100644
--- a/derp/xdp/bpf_bpfel.o
+++ b/derp/xdp/bpf_bpfel.o
diff --git a/derp/xdp/xdp.c b/derp/xdp/xdp.c
index a72e67264..3ac6488a4 100644
--- a/derp/xdp/xdp.c
+++ b/derp/xdp/xdp.c
@@ -14,6 +14,10 @@ struct config {
 	// the context of the data. cilium/ebpf uses native endian encoding for map
 	// encoding even if we use big endian types here, e.g. __be16.
 	__u16 dst_port;
+	// If drop_stun is set to a nonzero value all UDP packets destined to
+	// dst_port will be dropped. This is useful for shedding home client load
+	// during maintenance.
+	__u16 drop_stun;
 };
 struct config *unused_config __attribute__((unused)); // required by bpf2go -type
 
@@ -60,6 +64,7 @@ enum counter_key_prog_end {
 	COUNTER_KEY_END_INVALID_IP_CSUM,
 	COUNTER_KEY_END_NOT_STUN_PORT,
 	COUNTER_KEY_END_INVALID_SW_ATTR_VAL,
+	COUNTER_KEY_END_DROP_STUN,
 	COUNTER_KEY_END_LEN
 };
 enum counter_key_prog_end *unused_counter_key_prog_end __attribute__((unused)); // required by bpf2go -type
@@ -334,6 +339,11 @@ static __always_inline int handle_packet(struct xdp_md *ctx, struct packet_conte
 		return XDP_PASS;
 	}
 
+	if (c->drop_stun) {
+		pctx->prog_end = COUNTER_KEY_END_DROP_STUN;
+		return XDP_DROP;
+	}
+
 	if (validate_udp_csum) {
 		__u16 cs;
 		__u32 pseudo_sum;
diff --git a/derp/xdp/xdp_default.go b/derp/xdp/xdp_default.go
index 35fd659ca..99bc30d2c 100644
--- a/derp/xdp/xdp_default.go
+++ b/derp/xdp/xdp_default.go
@@ -26,3 +26,11 @@ func (s *STUNServer) Close() error {
 func (s *STUNServer) Describe(descCh chan<- *prometheus.Desc) {}
 
 func (s *STUNServer) Collect(metricCh chan<- prometheus.Metric) {}
+
+func (s *STUNServer) SetDropSTUN(v bool) error {
+	return errors.New("unimplemented on this GOOS")
+}
+
+func (s *STUNServer) GetDropSTUN() bool {
+	return true
+}
diff --git a/derp/xdp/xdp_linux.go b/derp/xdp/xdp_linux.go
index f2d47a372..aa6559842 100644
--- a/derp/xdp/xdp_linux.go
+++ b/derp/xdp/xdp_linux.go
@@ -22,9 +22,11 @@ import (
 // the STUN protocol. It exports statistics for the XDP program via its
 // implementation of the prometheus.Collector interface.
 type STUNServer struct {
-	mu      sync.Mutex
-	objs    *bpfObjects
-	metrics *stunServerMetrics
+	mu       sync.Mutex
+	objs     *bpfObjects
+	metrics  *stunServerMetrics
+	dstPort  int
+	dropSTUN bool
 }
 
 //lint:ignore U1000 used in xdp_linux_test.go, which has a build tag
@@ -68,12 +70,13 @@ func NewSTUNServer(config *STUNServerConfig, opts ...STUNServerOption) (*STUNSer
 	server := &STUNServer{
 		objs:    objs,
 		metrics: newSTUNServerMetrics(),
+		dstPort: config.DstPort,
 	}
 	var key uint32
-	xdpConfig := bpfConfig{
+	xdpConfig := &bpfConfig{
 		DstPort: uint16(config.DstPort),
 	}
-	err = objs.ConfigMap.Put(key, &xdpConfig)
+	err = objs.ConfigMap.Put(key, xdpConfig)
 	if err != nil {
 		return nil, fmt.Errorf("error loading config in eBPF map: %w", err)
 	}
@@ -181,6 +184,7 @@ var (
 		bpfCounterKeyProgEndCOUNTER_KEY_END_INVALID_IP_CSUM:            "invalid_ip_csum",
 		bpfCounterKeyProgEndCOUNTER_KEY_END_NOT_STUN_PORT:              "not_stun_port",
 		bpfCounterKeyProgEndCOUNTER_KEY_END_INVALID_SW_ATTR_VAL:        "invalid_sw_attr_val",
+		bpfCounterKeyProgEndCOUNTER_KEY_END_DROP_STUN:                  "drop_stun",
 	}
 
 	packetCounterKeys = map[bpfCounterKeyPacketsBytesAction]bool{
@@ -262,6 +266,31 @@ func (s *STUNServer) Collect(metricCh chan<- prometheus.Metric) {
 	s.metrics.registry.Collect(metricCh)
 }
 
+func (s *STUNServer) SetDropSTUN(v bool) error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	dropSTUN := 0
+	if v {
+		dropSTUN = 1
+	}
+	xdpConfig := &bpfConfig{
+		DstPort:  uint16(s.dstPort),
+		DropStun: uint16(dropSTUN),
+	}
+	var key uint32
+	err := s.objs.ConfigMap.Put(key, xdpConfig)
+	if err == nil {
+		s.dropSTUN = v
+	}
+	return err
+}
+
+func (s *STUNServer) GetDropSTUN() bool {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	return s.dropSTUN
+}
+
 func (s *STUNServer) updateMetrics() error {
 	s.mu.Lock()
 	defer s.mu.Unlock()
diff --git a/derp/xdp/xdp_linux_test.go b/derp/xdp/xdp_linux_test.go
index 1a59f9444..07f11eff6 100644
--- a/derp/xdp/xdp_linux_test.go
+++ b/derp/xdp/xdp_linux_test.go
@@ -440,12 +440,51 @@ func TestXDP(t *testing.T) {
 
 	cases := []struct {
 		name          string
+		dropSTUN      bool
 		packetIn      []byte
 		wantCode      xdpAction
 		wantPacketOut []byte
 		wantMetrics   map[bpfCountersKey]uint64
 	}{
 		{
+			name:          "ipv4 STUN Binding Request Drop STUN",
+			dropSTUN:      true,
+			packetIn:      ipv4STUNBindingReqTX,
+			wantCode:      xdpActionDrop,
+			wantPacketOut: ipv4STUNBindingReqTX,
+			wantMetrics: map[bpfCountersKey]uint64{
+				{
+					Af:      uint8(bpfCounterKeyAfCOUNTER_KEY_AF_IPV4),
+					Pba:     uint8(bpfCounterKeyPacketsBytesActionCOUNTER_KEY_PACKETS_DROP_TOTAL),
+					ProgEnd: uint8(bpfCounterKeyProgEndCOUNTER_KEY_END_DROP_STUN),
+				}: 1,
+				{
+					Af:      uint8(bpfCounterKeyAfCOUNTER_KEY_AF_IPV4),
+					Pba:     uint8(bpfCounterKeyPacketsBytesActionCOUNTER_KEY_BYTES_DROP_TOTAL),
+					ProgEnd: uint8(bpfCounterKeyProgEndCOUNTER_KEY_END_DROP_STUN),
+				}: uint64(len(ipv4STUNBindingReqTX)),
+			},
+		},
+		{
+			name:          "ipv6 STUN Binding Request Drop STUN",
+			dropSTUN:      true,
+			packetIn:      ipv6STUNBindingReqTX,
+			wantCode:      xdpActionDrop,
+			wantPacketOut: ipv6STUNBindingReqTX,
+			wantMetrics: map[bpfCountersKey]uint64{
+				{
+					Af:      uint8(bpfCounterKeyAfCOUNTER_KEY_AF_IPV6),
+					Pba:     uint8(bpfCounterKeyPacketsBytesActionCOUNTER_KEY_PACKETS_DROP_TOTAL),
+					ProgEnd: uint8(bpfCounterKeyProgEndCOUNTER_KEY_END_DROP_STUN),
+				}: 1,
+				{
+					Af:      uint8(bpfCounterKeyAfCOUNTER_KEY_AF_IPV6),
+					Pba:     uint8(bpfCounterKeyPacketsBytesActionCOUNTER_KEY_BYTES_DROP_TOTAL),
+					ProgEnd: uint8(bpfCounterKeyProgEndCOUNTER_KEY_END_DROP_STUN),
+				}: uint64(len(ipv6STUNBindingReqTX)),
+			},
+		},
+		{
 			name:          "ipv4 STUN Binding Request TX",
 			packetIn:      ipv4STUNBindingReqTX,
 			wantCode:      xdpActionTX,
@@ -963,6 +1002,10 @@ func TestXDP(t *testing.T) {
 				Data:    c.packetIn,
 				DataOut: make([]byte, 1514),
 			}
+			err = server.SetDropSTUN(c.dropSTUN)
+			if err != nil {
+				t.Fatalf("error setting drop STUN: %v", err)
+			}
 			got, err := server.objs.XdpProgFunc.Run(&opts)
 			if err != nil {
 				t.Fatalf("error running program: %v", err)