chore: update tailscale (#6091)

tailnet/conn.go

@@ -18,6 +18,7 @@ import (
 	"gvisor.dev/gvisor/pkg/tcpip/adapters/gonet"
 	"tailscale.com/hostinfo"
 	"tailscale.com/ipn/ipnstate"
+	"tailscale.com/net/connstats"
 	"tailscale.com/net/dns"
 	"tailscale.com/net/netns"
 	"tailscale.com/net/tsdial"
@@ -55,11 +56,6 @@ type Options struct {
 	// If so, only DERPs can establish connections.
 	BlockEndpoints bool
 	Logger         slog.Logger
-
-	// EnableTrafficStats enables per-connection traffic statistics.
-	// ExtractTrafficStats must be called to reset the counters and be
-	// periodically called while enabled to avoid unbounded memory use.
-	EnableTrafficStats bool
 }
 
 // NewConn constructs a new Wireguard server that will accept connections from the addresses provided.
@@ -73,6 +69,7 @@ func NewConn(options *Options) (*Conn, error) {
 	if options.DERPMap == nil {
 		return nil, xerrors.New("DERPMap must be provided")
 	}
+
 	nodePrivateKey := key.NewNode()
 	nodePublicKey := nodePrivateKey.Public()
 
@@ -151,7 +148,6 @@ func NewConn(options *Options) (*Conn, error) {
 	if !ok {
 		return nil, xerrors.New("get wireguard internals")
 	}
-	tunDevice.SetStatisticsEnabled(options.EnableTrafficStats)
 
 	// Update the keys for the magic connection!
 	err = magicConn.SetPrivateKey(nodePrivateKey)
@@ -169,7 +165,7 @@ func NewConn(options *Options) (*Conn, error) {
 		return netStack.DialContextTCP(ctx, dst)
 	}
 	netStack.ProcessLocalIPs = true
-	err = netStack.Start()
+	err = netStack.Start(nil)
 	if err != nil {
 		return nil, xerrors.Errorf("start netstack: %w", err)
 	}
@@ -295,6 +291,8 @@ type Conn struct {
 	lastPreferredDERP int
 	lastDERPLatency   map[string]float64
 	nodeCallback      func(node *Node)
+
+	trafficStats *connstats.Statistics
 }
 
 // SetForwardTCPCallback is called every time a TCP connection is initiated inbound.
@@ -427,7 +425,7 @@ func (c *Conn) UpdateNodes(nodes []*Node) error {
 
 // Status returns the current ipnstate of a connection.
 func (c *Conn) Status() *ipnstate.Status {
-	sb := &ipnstate.StatusBuilder{}
+	sb := &ipnstate.StatusBuilder{WantPeers: true}
 	c.wireguardEngine.UpdateStatus(sb)
 	return sb.Status()
 }
@@ -538,6 +536,11 @@ func (c *Conn) Close() error {
 	_ = c.wireguardMonitor.Close()
 	_ = c.tunDevice.Close()
 	c.wireguardEngine.Close()
+	if c.trafficStats != nil {
+		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+		_ = c.trafficStats.Shutdown(ctx)
+	}
 	return nil
 }
 
@@ -701,11 +704,24 @@ func (c *Conn) forwardTCPToLocal(conn net.Conn, port uint16) {
 	c.logger.Debug(c.dialContext, "forwarded connection closed", slog.F("local_addr", dialAddrStr))
 }
 
-// ExtractTrafficStats extracts and resets the counters for all active
-// connections. It must be called periodically otherwise the memory used is
-// unbounded. EnableTrafficStats must be true when calling NewConn.
-func (c *Conn) ExtractTrafficStats() map[netlogtype.Connection]netlogtype.Counts {
-	return c.tunDevice.ExtractStatistics()
+// SetConnStatsCallback sets a callback to be called after maxPeriod or
+// maxConns, whichever comes first. Multiple calls overwrites the callback.
+func (c *Conn) SetConnStatsCallback(maxPeriod time.Duration, maxConns int, dump func(start, end time.Time, virtual, physical map[netlogtype.Connection]netlogtype.Counts)) {
+	connStats := connstats.NewStatistics(maxPeriod, maxConns, dump)
+
+	c.mutex.Lock()
+	old := c.trafficStats
+	c.trafficStats = connStats
+	c.mutex.Unlock()
+
+	// Make sure to shutdown the old callback.
+	if old != nil {
+		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+		_ = old.Shutdown(ctx)
+	}
+
+	c.tunDevice.SetStatistics(connStats)
 }
 
 type listenKey struct {