fix(tailnet): Ensure connstats are closed on conn close

mafredri · mafredri · commit 5667170ee47f · 2023-02-21T13:55:34.000Z
diff --git a/tailnet/conn.go b/tailnet/conn.go
@@ -744,16 +744,25 @@ func (c *Conn) forwardTCPToLocal(conn net.Conn, port uint16) {
 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)
 
+	closeStats := func(s *connstats.Statistics) {
+		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+		_ = s.Shutdown(ctx)
+	}
+
 	c.mutex.Lock()
+	if c.isClosed() {
+		c.mutex.Unlock()
+		closeStats(connStats)
+		return
+	}
 	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)
+		closeStats(old)
 	}
 
 	c.tunDevice.SetStatistics(connStats)
