cmd/derper,derp: add --rate-config file with SIGHUP reload (#19314)

Add a --rate-config flag pointing to a JSON file for per-client receive
rate limits (bytes/sec and burst bytes). The config is reloaded on SIGHUP,
updating all existing client connections live. The --per-client-rate-limit
and --per-client-rate-burst flags are removed in favor of the config file.

In derpserver, rate limiting uses an atomic.Pointer[xrate.Limiter] per
client: nil when unlimited or mesh (zero overhead), non-nil when
rate-limited.

Document that clientSet.activeClient Store operations require Server.mu.

Updates tailscale/corp#38509

Signed-off-by: Mike O'Driscoll <mikeo@tailscale.com>

derp/derpserver/derpserver.go

@@ -209,11 +209,11 @@ type Server struct {
 	// perClientRecvBytesPerSec is the rate limit for receiving data from
 	// a single client connection, in bytes per second. 0 means unlimited.
 	// Mesh peers are exempt from this limit.
-	perClientRecvBytesPerSec uint
+	perClientRecvBytesPerSec uint64
 	// perClientRecvBurst is the burst size in bytes for the per-client
-	// receive rate limiter. perClientRecvBurst is only relevant when
-	// perClientRecvBytesPerSec is nonzero.
-	perClientRecvBurst uint
+	// receive rate limiter. Always at least [derp.MaxPacketSize] when
+	// set via [Server.UpdatePerClientRateLimit].
+	perClientRecvBurst uint64
 
 	clock tstime.Clock
 }
@@ -239,10 +239,6 @@ type clientSet struct {
 	// activeClient holds the currently active connection for the set. It's nil
 	// if there are no connections or the connection is disabled.
 	//
-	// A pointer to a clientSet can be held by peers for long periods of time
-	// without holding Server.mu to avoid mutex contention on Server.mu, only
-	// re-acquiring the mutex and checking the clients map if activeClient is
-	// nil.
 	activeClient atomic.Pointer[sclient]
 
 	// dup is non-nil if there are multiple connections for the
@@ -518,14 +514,62 @@ func (s *Server) SetTCPWriteTimeout(d time.Duration) {
 	s.tcpWriteTimeout = d
 }
 
-// SetPerClientRateLimit sets the per-client receive rate limit in bytes per
-// second and the burst size in bytes. Mesh peers are exempt from this limit.
-// The burst is at least [derp.MaxPacketSize], or burst, if burst is greater
-// than [derp.MaxPacketSize]. This ensures at least a full packet can
-// be received in a burst, even if the rate limit is low.
-func (s *Server) SetPerClientRateLimit(bytesPerSec, burst uint) {
+// RateConfig is a JSON-serializable configuration for per-client rate limits.
+// Values are in bytes.
+type RateConfig struct {
+	// PerClientRateLimitBytesPerSec represents the per-client
+	// rate limit in bytes per second. A zero value disables rate-limiting.
+	PerClientRateLimitBytesPerSec uint64 `json:",omitzero"`
+	// PerClientRateBurstBytes represents the per-client token bucket depth,
+	// or burst, in bytes. Any value lower than [derp.MaxPacketSize]
+	// will be increased to [derp.MaxPacketSize] before application.
+	PerClientRateBurstBytes uint64 `json:",omitzero"`
+}
+
+// LoadRateConfig reads and JSON-unmarshals a [RateConfig] from the file at path.
+func LoadRateConfig(path string) (RateConfig, error) {
+	if path == "" {
+		return RateConfig{}, errors.New("rate config path is empty")
+	}
+	b, err := os.ReadFile(path)
+	if err != nil {
+		return RateConfig{}, fmt.Errorf("reading rate config: %w", err)
+	}
+	var rc RateConfig
+	if err := json.Unmarshal(b, &rc); err != nil {
+		return RateConfig{}, fmt.Errorf("parsing rate config: %w", err)
+	}
+	return rc, nil
+}
+
+// LoadAndApplyRateConfig reads a [RateConfig] from the file at path and
+// applies it to the server via [Server.UpdatePerClientRateLimit].
+func (s *Server) LoadAndApplyRateConfig(path string) error {
+	rc, err := LoadRateConfig(path)
+	if err != nil {
+		return err
+	}
+	s.UpdatePerClientRateLimit(rc.PerClientRateLimitBytesPerSec, rc.PerClientRateBurstBytes)
+	s.logf("rate config applied: rate=%d bytes/sec, burst=%d bytes", rc.PerClientRateLimitBytesPerSec, rc.PerClientRateBurstBytes)
+	return nil
+}
+
+// UpdatePerClientRateLimit sets the per-client receive rate limit in bytes per
+// second and the burst size in bytes, updating all existing client connections.
+// The burst is at least [derp.MaxPacketSize], ensuring at least a full packet
+// can be received in a burst even if the rate limit is low. If bytesPerSec is
+// 0, rate limiting is set to infinity. Mesh peers are always exempt from rate
+// limiting.
+func (s *Server) UpdatePerClientRateLimit(bytesPerSec, burst uint64) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
 	s.perClientRecvBytesPerSec = bytesPerSec
 	s.perClientRecvBurst = max(burst, derp.MaxPacketSize)
+	for _, cs := range s.clients {
+		cs.ForeachClient(func(c *sclient) {
+			c.setRateLimit(s.perClientRecvBytesPerSec, s.perClientRecvBurst)
+		})
+	}
 }
 
 // HasMeshKey reports whether the server is configured with a mesh key.
@@ -690,6 +734,8 @@ func (s *Server) registerClient(c *sclient) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
+	c.setRateLimit(s.perClientRecvBytesPerSec, s.perClientRecvBurst)
+
 	cs, ok := s.clients[c.key]
 	if !ok {
 		c.debugLogf("register single client")
@@ -975,9 +1021,6 @@ func (s *Server) accept(ctx context.Context, nc derp.Conn, brw *bufio.ReadWriter
 		peerGoneLim:    rate.NewLimiter(rate.Every(time.Second), 3),
 	}
 
-	if s.perClientRecvBytesPerSec > 0 && !c.canMesh {
-		c.recvLim = xrate.NewLimiter(xrate.Limit(s.perClientRecvBytesPerSec), int(s.perClientRecvBurst))
-	}
 	if c.canMesh {
 		c.meshUpdate = make(chan struct{}, 1) // must be buffered; >1 is fine but wasteful
 	}
@@ -1267,19 +1310,38 @@ func (c *sclient) handleFrameSendPacket(_ derp.FrameType, fl uint32) error {
 	return c.sendPkt(dst, p)
 }
 
-// rateLimit applies the per-client receive rate limit, if configured.
+// setRateLimit updates the per-client receive rate limiter.
+// When bytesPerSec is 0 or the client is a mesh peer, the limiter is
+// set to nil so that [sclient.rateLimit] is a no-op.
+func (c *sclient) setRateLimit(bytesPerSec uint64, burst uint64) {
+	if bytesPerSec == 0 || c.canMesh {
+		c.recvLim.Store(nil)
+		return
+	}
+	if lim := c.recvLim.Load(); lim != nil {
+		// Update in place. SetBurst before SetLimit to avoid a transient
+		// state where a new higher rate exceeds the old lower burst.
+		lim.SetBurst(int(burst))
+		lim.SetLimit(xrate.Limit(bytesPerSec))
+		return
+	}
+	lim := xrate.NewLimiter(xrate.Limit(bytesPerSec), int(burst))
+	c.recvLim.Store(lim)
+}
+
+// rateLimit applies the per-client receive rate limit.
 // By limiting here we prevent reading from the buffered reader
 // [sclient.br] if the limit has been exceeded. Any reads done here provide space
 // within the buffered reader to fill back in with data from
 // the TCP socket. Pacing reads acts as a form of natural
 // backpressure via TCP flow control.
-// meshed clients are exempt from rate limits.
+// When rate limiting is disabled or the client is a mesh peer, recvLim is nil
+// and this is a no-op.
 func (c *sclient) rateLimit(n int) error {
-	if c.recvLim == nil || c.canMesh {
-		return nil
+	if lim := c.recvLim.Load(); lim != nil {
+		return lim.WaitN(c.ctx, n)
 	}
-
-	return c.recvLim.WaitN(c.ctx, n)
+	return nil
 }
 
 func (c *sclient) debugLogf(format string, v ...any) {
@@ -1714,10 +1776,15 @@ type sclient struct {
 	// through us with a peer we have no record of.
 	peerGoneLim *rate.Limiter
 
-	// recvLim is the per-connection receive rate limiter. If non-nil,
-	// the server calls WaitN per received DERP frame in order to
-	// apply TCP backpressure and throttle the sender.
-	recvLim *xrate.Limiter
+	// recvLim is the per-connection receive rate limiter. When rate
+	// limiting is enabled for a non-mesh client, it points to an
+	// [xrate.Limiter]. When rate limiting is disabled or the client is a
+	// mesh peer, it is nil and [sclient.rateLimit] is a no-op.
+	// Updated atomically by [sclient.setRateLimitLocked] so that
+	// [sclient.rateLimit] can load it without holding Server.mu.
+	// TODO(mikeodr): update to use mono time, requires updates
+	// to tstime/rate.Limiter
+	recvLim atomic.Pointer[xrate.Limiter]
 }
 
 func (c *sclient) presentFlags() derp.PeerPresentFlags {