net/tstun: instrument Wrapper with statistics gathering (#5847)

If Wrapper.StatisticsEnable is enabled,
then per-connection counters are maintained.
If enabled, Wrapper.StatisticsExtract must be periodically called
otherwise there is unbounded memory growth.

Signed-off-by: Joe Tsai <joetsai@digital-static.net>

cmd/tailscaled/depaware.txt

@@ -240,6 +240,7 @@ tailscale.com/cmd/tailscaled dependencies: (generated by github.com/tailscale/de
         tailscale.com/net/tsdial                                     from tailscale.com/control/controlclient+
      💣 tailscale.com/net/tshttpproxy                                from tailscale.com/control/controlclient+
         tailscale.com/net/tstun                                      from tailscale.com/net/dns+
+        tailscale.com/net/tunstats                                   from tailscale.com/net/tstun
         tailscale.com/paths                                          from tailscale.com/ipn/ipnlocal+
         tailscale.com/portlist                                       from tailscale.com/ipn/ipnlocal
         tailscale.com/safesocket                                     from tailscale.com/client/tailscale+

net/tstun/wrap.go

@@ -22,8 +22,10 @@
 	"golang.zx2c4.com/wireguard/tun"
 	"gvisor.dev/gvisor/pkg/tcpip/stack"
 	"tailscale.com/disco"
+	"tailscale.com/net/flowtrack"
 	"tailscale.com/net/packet"
 	"tailscale.com/net/tsaddr"
+	"tailscale.com/net/tunstats"
 	"tailscale.com/syncs"
 	"tailscale.com/tstime/mono"
 	"tailscale.com/types/ipproto"
@@ -166,6 +168,12 @@ type Wrapper struct {
 
 	// disableTSMPRejected disables TSMP rejected responses. For tests.
 	disableTSMPRejected bool
+
+	// stats maintains per-connection counters.
+	stats struct {
+		enabled atomic.Bool
+		tunstats.Statistics
+	}
 }
 
 // tunReadResult is the result of a TUN read, or an injected result pretending to be a TUN read.
@@ -560,6 +568,9 @@ func (t *Wrapper) Read(buf []byte, offset int) (int, error) {
 		}
 	}
 
+	if t.stats.enabled.Load() {
+		t.stats.UpdateTx(buf[offset:][:n])
+	}
 	t.noteActivity()
 	return n, nil
 }
@@ -690,6 +701,9 @@ func (t *Wrapper) Write(buf []byte, offset int) (int, error) {
 }
 
 func (t *Wrapper) tdevWrite(buf []byte, offset int) (int, error) {
+	if t.stats.enabled.Load() {
+		t.stats.UpdateRx(buf[offset:])
+	}
 	if t.isTAP {
 		return t.tapWrite(buf, offset)
 	}
@@ -829,6 +843,18 @@ func (t *Wrapper) Unwrap() tun.Device {
 	return t.tdev
 }
 
+// StatisticsEnable enables per-connections packet counters.
+// StatisticsExtract must be called periodically to avoid unbounded memory use.
+func (t *Wrapper) StatisticsEnable(enable bool) {
+	t.stats.enabled.Store(enable)
+}
+
+// StatisticsExtract extracts and resets the counters for all active connections.
+// It must be called periodically otherwise the memory used is unbounded.
+func (t *Wrapper) StatisticsExtract() map[flowtrack.Tuple]tunstats.Counts {
+	return t.stats.Extract()
+}
+
 var (
 	metricPacketIn              = clientmetric.NewCounter("tstun_in_from_wg")
 	metricPacketInDrop          = clientmetric.NewCounter("tstun_in_from_wg_drop")

net/tstun/wrap_test.go

@@ -9,6 +9,7 @@
 	"encoding/binary"
 	"fmt"
 	"net/netip"
+	"reflect"
 	"strconv"
 	"strings"
 	"testing"
@@ -18,8 +19,10 @@
 	"go4.org/netipx"
 	"golang.zx2c4.com/wireguard/tun/tuntest"
 	"tailscale.com/disco"
+	"tailscale.com/net/flowtrack"
 	"tailscale.com/net/netaddr"
 	"tailscale.com/net/packet"
+	"tailscale.com/net/tunstats"
 	"tailscale.com/tstest"
 	"tailscale.com/tstime/mono"
 	"tailscale.com/types/ipproto"
@@ -281,6 +284,11 @@ func TestWriteAndInject(t *testing.T) {
 			t.Errorf("%s not received", packet)
 		}
 	}
+
+	// Statistics gathering is disabled by default.
+	if stats := tun.StatisticsExtract(); len(stats) > 0 {
+		t.Errorf("tun.StatisticsExtract = %v, want {}", stats)
+	}
 }
 
 func TestFilter(t *testing.T) {
@@ -329,12 +337,17 @@ func TestFilter(t *testing.T) {
 	}()
 
 	var buf [MaxPacketSize]byte
+	tun.StatisticsEnable(true)
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			var n int
 			var err error
 			var filtered bool
 
+			if stats := tun.StatisticsExtract(); len(stats) > 0 {
+				t.Errorf("tun.StatisticsExtract = %v, want {}", stats)
+			}
+
 			if tt.dir == in {
 				// Use the side effect of updating the last
 				// activity atomic to determine whether the
@@ -364,6 +377,24 @@ func TestFilter(t *testing.T) {
 					t.Errorf("got accept; want drop")
 				}
 			}
+
+			got := tun.StatisticsExtract()
+			want := map[flowtrack.Tuple]tunstats.Counts{}
+			if !tt.drop {
+				var p packet.Parsed
+				p.Decode(tt.data)
+				switch tt.dir {
+				case in:
+					tuple := flowtrack.Tuple{Proto: ipproto.UDP, Src: p.Dst, Dst: p.Src}
+					want[tuple] = tunstats.Counts{RxPackets: 1, RxBytes: uint64(len(tt.data))}
+				case out:
+					tuple := flowtrack.Tuple{Proto: ipproto.UDP, Src: p.Src, Dst: p.Dst}
+					want[tuple] = tunstats.Counts{TxPackets: 1, TxBytes: uint64(len(tt.data))}
+				}
+			}
+			if !reflect.DeepEqual(got, want) {
+				t.Errorf("tun.StatisticsExtract = %v, want %v", got, want)
+			}
 		})
 	}
 }