Merge branch 'main' into refactor/relay-foreign-cache

# Conflicts:
#	shared/relay/client/manager.go

shared/management/client/grpc.go

@@ -33,10 +33,15 @@ const ConnectTimeout = 10 * time.Second
 const healthCheckTimeout = 5 * time.Second
 
 const (
-	// EnvMaxRecvMsgSize overrides the default gRPC max receive message size (4 MB)
+	// EnvMaxRecvMsgSize overrides the default gRPC max receive message size
 	// for the management client connection. Value is in bytes.
 	EnvMaxRecvMsgSize = "NB_MANAGEMENT_GRPC_MAX_MSG_SIZE"
 
+	// defaultMaxRecvMsgSize is the max gRPC receive message size used for the
+	// management client connection when EnvMaxRecvMsgSize is unset or invalid.
+	// It overrides the gRPC library default of 4 MB.
+	defaultMaxRecvMsgSize = 1024 * 1024 * 16
+
 	errMsgMgmtPublicKey    = "failed getting Management Service public key: %s"
 	errMsgNoMgmtConnection = "no connection to management"
 )
@@ -55,6 +60,14 @@ type GrpcClient struct {
 	connStateCallback     ConnStateNotifier
 	connStateCallbackLock sync.RWMutex
 	serverURL             string
+
+	// syncStreamErr holds the last Sync stream error, or nil while the stream
+	// is established and healthy. GetServerKey succeeds even when the peer
+	// cannot sync (e.g. the server returns "settings not found"), so the
+	// health probe must consult this to avoid reporting a healthy management
+	// connection while the Sync stream keeps failing.
+	syncStreamMu  sync.RWMutex
+	syncStreamErr error
 }
 
 type ExposeRequest struct {
@@ -76,22 +89,22 @@ type ExposeResponse struct {
 }
 
 // MaxRecvMsgSize returns the configured max gRPC receive message size from
-// the environment, or 0 if unset (which uses the gRPC default of 4 MB).
+// the environment, or defaultMaxRecvMsgSize (16 MB) if unset or invalid.
 func MaxRecvMsgSize() int {
 	val := os.Getenv(EnvMaxRecvMsgSize)
 	if val == "" {
-		return 0
+		return defaultMaxRecvMsgSize
 	}
 
 	size, err := strconv.Atoi(val)
 	if err != nil {
 		log.Warnf("invalid %s value %q, using default: %v", EnvMaxRecvMsgSize, val, err)
-		return 0
+		return defaultMaxRecvMsgSize
 	}
 
 	if size <= 0 {
 		log.Warnf("invalid %s value %d, must be positive, using default", EnvMaxRecvMsgSize, size)
-		return 0
+		return defaultMaxRecvMsgSize
 	}
 
 	return size
@@ -364,6 +377,8 @@ func (c *GrpcClient) handleSyncStream(ctx context.Context, serverPubKey wgtypes.
 	stream, err := c.connectToSyncStream(ctx, serverPubKey, sysInfo)
 	if err != nil {
 		log.Debugf("failed to open Management Service stream: %s", err)
+		c.notifyDisconnected(err)
+		c.setSyncStreamDisconnected(err)
 		if s, ok := gstatus.FromError(err); ok && s.Code() == codes.PermissionDenied {
 			return backoff.Permanent(err) // unrecoverable error, propagate to the upper layer
 		}
@@ -372,11 +387,13 @@ func (c *GrpcClient) handleSyncStream(ctx context.Context, serverPubKey wgtypes.
 
 	log.Infof("connected to the Management Service stream")
 	c.notifyConnected()
+	c.setSyncStreamConnected()
 
 	// blocking until error
 	err = c.receiveUpdatesEvents(stream, serverPubKey, msgHandler)
 	if err != nil {
 		c.notifyDisconnected(err)
+		c.setSyncStreamDisconnected(err)
 		if ctx.Err() != nil {
 			log.Debugf("management connection context has been canceled, this usually indicates shutdown")
 			return nil
@@ -524,12 +541,19 @@ func (c *GrpcClient) IsHealthy() bool {
 	ctx, cancel := context.WithTimeout(c.ctx, healthCheckTimeout)
 	defer cancel()
 
-	_, err := c.realClient.GetServerKey(ctx, &proto.Empty{})
+	_, err := c.realClient.IsHealthy(ctx, &proto.Empty{})
 	if err != nil {
 		c.notifyDisconnected(err)
 		log.Warnf("health check returned: %s", err)
 		return false
 	}
+
+	if syncErr := c.syncStreamError(); syncErr != nil {
+		c.notifyDisconnected(syncErr)
+		log.Warnf("management transport is up but the Sync stream is unhealthy: %s", syncErr)
+		return false
+	}
+
 	c.notifyConnected()
 	return true
 }
@@ -630,26 +654,14 @@ func (c *GrpcClient) ExtendAuthSession(sysInfo *system.Info, jwtToken string) (*
 		return nil, err
 	}
 
-	var resp *proto.EncryptedMessage
-	operation := func() error {
-		mgmCtx, cancel := context.WithTimeout(context.Background(), ConnectTimeout)
-		defer cancel()
+	mgmCtx, cancel := context.WithTimeout(c.ctx, ConnectTimeout)
+	defer cancel()
 
-		var err error
-		resp, err = c.realClient.ExtendAuthSession(mgmCtx, &proto.EncryptedMessage{
-			WgPubKey: c.key.PublicKey().String(),
-			Body:     reqBody,
-		})
-		if err != nil {
-			if s, ok := gstatus.FromError(err); ok && s.Code() == codes.Canceled {
-				return err
-			}
-			return backoff.Permanent(err)
-		}
-		return nil
-	}
-
-	if err := backoff.Retry(operation, nbgrpc.Backoff(c.ctx)); err != nil {
+	resp, err := c.realClient.ExtendAuthSession(mgmCtx, &proto.EncryptedMessage{
+		WgPubKey: c.key.PublicKey().String(),
+		Body:     reqBody,
+	})
+	if err != nil {
 		log.Errorf("failed to extend auth session on Management Service: %v", err)
 		return nil, err
 	}
@@ -771,6 +783,24 @@ func (c *GrpcClient) SyncMeta(sysInfo *system.Info) error {
 	return err
 }
 
+func (c *GrpcClient) setSyncStreamConnected() {
+	c.syncStreamMu.Lock()
+	defer c.syncStreamMu.Unlock()
+	c.syncStreamErr = nil
+}
+
+func (c *GrpcClient) setSyncStreamDisconnected(err error) {
+	c.syncStreamMu.Lock()
+	defer c.syncStreamMu.Unlock()
+	c.syncStreamErr = err
+}
+
+func (c *GrpcClient) syncStreamError() error {
+	c.syncStreamMu.RLock()
+	defer c.syncStreamMu.RUnlock()
+	return c.syncStreamErr
+}
+
 func (c *GrpcClient) notifyDisconnected(err error) {
 	c.connStateCallbackLock.RLock()
 	defer c.connStateCallbackLock.RUnlock()
@@ -993,8 +1023,6 @@ func infoToMetaData(info *system.Info) *proto.PeerSystemMeta {
 			BlockLANAccess:      info.BlockLANAccess,
 			BlockInbound:        info.BlockInbound,
 			DisableIPv6:         info.DisableIPv6,
-
-			LazyConnectionEnabled: info.LazyConnectionEnabled,
 		},
 
 		Capabilities: peerCapabilities(*info),

shared/management/client/grpc_test.go

@@ -21,11 +21,11 @@ func TestMaxRecvMsgSize(t *testing.T) {
 		envValue string
 		expected int
 	}{
-		{name: "unset returns 0", envValue: "", expected: 0},
+		{name: "unset returns default", envValue: "", expected: defaultMaxRecvMsgSize},
 		{name: "valid value", envValue: "10485760", expected: 10485760},
-		{name: "non-numeric returns 0", envValue: "abc", expected: 0},
-		{name: "negative returns 0", envValue: "-1", expected: 0},
-		{name: "zero returns 0", envValue: "0", expected: 0},
+		{name: "non-numeric returns default", envValue: "abc", expected: defaultMaxRecvMsgSize},
+		{name: "negative returns default", envValue: "-1", expected: defaultMaxRecvMsgSize},
+		{name: "zero returns default", envValue: "0", expected: defaultMaxRecvMsgSize},
 	}
 
 	for _, tt := range tests {

shared/management/proto/management.proto

@@ -312,6 +312,8 @@ message NetbirdConfig {
   RelayConfig relay = 4;
 
   FlowConfig flow = 5;
+
+  MetricsConfig metrics = 6;
 }
 
 // HostConfig describes connection properties of some server (e.g. STUN, Signal, Management)
@@ -350,6 +352,10 @@ message FlowConfig {
   bool dnsCollection = 8;
 }
 
+message MetricsConfig {
+  bool enabled = 1;
+}
+
 // JWTConfig represents JWT authentication configuration for validating tokens.
 message JWTConfig {
   string issuer = 1;

shared/management/proto/proxy_service.proto

@@ -43,6 +43,18 @@ service ProxyService {
   // issue a session cookie without redirecting through the OIDC flow.
   // Mirrors ValidateSession's response shape.
   rpc ValidateTunnelPeer(ValidateTunnelPeerRequest) returns (ValidateTunnelPeerResponse);
+
+  // CheckLLMPolicyLimits is the pre-flight RPC the proxy calls before each
+  // LLM request. Management runs the per-policy headroom selection across
+  // every policy authorising the caller's user / groups for the resolved
+  // provider and returns the chosen attribution policy + group, or a deny
+  // when no applicable policy has headroom > 0.
+  rpc CheckLLMPolicyLimits(CheckLLMPolicyLimitsRequest) returns (CheckLLMPolicyLimitsResponse);
+
+  // RecordLLMUsage is the post-flight RPC the proxy calls after the upstream
+  // returns. Increments the per-(dimension, window) counters for the
+  // attribution policy chosen by CheckLLMPolicyLimits.
+  rpc RecordLLMUsage(RecordLLMUsageRequest) returns (RecordLLMUsageResponse);
 }
 
 // ProxyCapabilities describes what a proxy can handle.
@@ -107,6 +119,59 @@ message PathTargetOptions {
   // reachable without WireGuard (public APIs, LAN services, localhost
   // sidecars). Defaults to false — embedded client is the standard path.
   bool direct_upstream = 7;
+  // Proxy clamps to [0, proxy-wide max (1 MiB)] at apply time. Agent-network
+  // synthesized targets only; private services leave these zero.
+  int64 capture_max_request_bytes = 8;
+  // Proxy clamps to [0, proxy-wide max (1 MiB)] at apply time.
+  int64 capture_max_response_bytes = 9;
+  // Content types eligible for body capture (e.g. "application/json").
+  repeated string capture_content_types = 10;
+  // Per-target middleware configurations populated by the agent-network
+  // synthesizer. Validated and clamped by the proxy at apply time.
+  repeated MiddlewareConfig middlewares = 11;
+  // When true, the proxy stamps agent_network=true on access-log entries
+  // for this target so management routes them to the agent-network log
+  // surface.
+  bool agent_network = 12;
+  // When true, the proxy suppresses the per-request access-log emission for
+  // this target. Defaults false to preserve existing access-log behavior for
+  // every non-agent-network target. The agent-network synth target sets this
+  // true only when the account's EnableLogCollection toggle is off.
+  bool disable_access_log = 13;
+}
+
+// MiddlewareSlot identifies where in the request lifecycle a middleware
+// runs. Mirrors proxy/internal/middleware.Slot.
+enum MiddlewareSlot {
+  MIDDLEWARE_SLOT_UNSPECIFIED = 0;
+  MIDDLEWARE_SLOT_ON_REQUEST = 1;
+  MIDDLEWARE_SLOT_ON_RESPONSE = 2;
+  MIDDLEWARE_SLOT_TERMINAL = 3;
+}
+
+// MiddlewareConfig is the per-target configuration for a single middleware.
+// The proxy validates every incoming MiddlewareConfig at apply time:
+// unknown ids are rejected, timeout is clamped to [10ms, 5s], and the
+// declared slot must match the registered middleware's slot.
+message MiddlewareConfig {
+  // Middleware id; must match the proxy-local compiled-in registry.
+  string id = 1;
+  bool enabled = 2;
+  MiddlewareSlot slot = 3;
+  // Free-form JSON unmarshalled by the middleware factory into its own typed
+  // config struct. Empty / null / {} are valid (zero-value config).
+  bytes config_json = 4;
+  enum FailMode {
+    FAIL_OPEN = 0;
+    FAIL_CLOSED = 1;
+  }
+  FailMode fail_mode = 5;
+  // Clamped to [10ms, 5s] at apply time; zero → 500ms default.
+  google.protobuf.Duration timeout = 6;
+  // When true, the middleware may mutate request headers or body (subject to
+  // policy). Honoured only when the implementation also declares
+  // MutationsSupported.
+  bool can_mutate = 7;
 }
 
 message PathMapping {
@@ -190,6 +255,10 @@ message AccessLog {
   string protocol = 16;
   // Extra key-value metadata for the access log entry (e.g. crowdsec_verdict, scenario).
   map<string, string> metadata = 17;
+  // When true, the entry was emitted by an agent-network synth service.
+  // Management routes these to the agent-network access-log surface instead
+  // of the standard service log.
+  bool agent_network = 18;
 }
 
 message AuthenticateRequest {
@@ -376,3 +445,59 @@ message SyncMappingsResponse {
   bool initial_sync_complete = 2;
 }
 
+// CheckLLMPolicyLimitsRequest carries the resolved caller identity and the
+// upstream provider already chosen by llm_router. Management computes which
+// policies authorise the request, picks the one with the most remaining
+// headroom, and returns the attribution decision.
+message CheckLLMPolicyLimitsRequest {
+  // account_id is the netbird account the request belongs to.
+  string account_id = 1;
+  // user_id is the netbird user id of the caller. May be empty when the
+  // principal is a tunnel-peer that isn't bound to a user; group membership
+  // still gates the request in that case.
+  string user_id = 2;
+  // group_ids is the caller's full group membership at request time.
+  repeated string group_ids = 3;
+  // provider_id is the agent-network provider record id chosen by llm_router.
+  string provider_id = 4;
+  // model is the upstream model identifier extracted from the request body.
+  string model = 5;
+}
+
+// CheckLLMPolicyLimitsResponse is management's allow-or-deny decision for a
+// pre-flight check.
+message CheckLLMPolicyLimitsResponse {
+  // decision is "allow" or "deny".
+  string decision = 1;
+  // selected_policy_id names the policy that paid for this request.
+  string selected_policy_id = 2;
+  // attribution_group_id is the source group the request booked against.
+  string attribution_group_id = 3;
+  // window_seconds is the cap window length the selected policy uses.
+  int64 window_seconds = 4;
+  // deny_code is set on decision="deny" with a stable label.
+  string deny_code = 5;
+  // deny_reason is a short human-readable explanation paired with deny_code.
+  string deny_reason = 6;
+}
+
+// RecordLLMUsageRequest is the post-flight increment the proxy posts after
+// the upstream call. Counters are keyed on (account, dimension, window).
+message RecordLLMUsageRequest {
+  string account_id = 1;
+  string user_id = 2;
+  // group_id is the selected policy's attribution group, recorded against the
+  // policy window (window_seconds).
+  string group_id = 3;
+  int64 window_seconds = 4;
+  int64 tokens_input = 5;
+  int64 tokens_output = 6;
+  double cost_usd = 7;
+  // group_ids is the caller's full group membership, used to fan the same
+  // usage out to every applicable account-level budget rule's own window.
+  repeated string group_ids = 8;
+}
+
+message RecordLLMUsageResponse {
+}
+

shared/management/proto/proxy_service_grpc.pb.go

@@ -43,6 +43,16 @@ type ProxyServiceClient interface {
 	// issue a session cookie without redirecting through the OIDC flow.
 	// Mirrors ValidateSession's response shape.
 	ValidateTunnelPeer(ctx context.Context, in *ValidateTunnelPeerRequest, opts ...grpc.CallOption) (*ValidateTunnelPeerResponse, error)
+	// CheckLLMPolicyLimits is the pre-flight RPC the proxy calls before each
+	// LLM request. Management runs the per-policy headroom selection across
+	// every policy authorising the caller's user / groups for the resolved
+	// provider and returns the chosen attribution policy + group, or a deny
+	// when no applicable policy has headroom > 0.
+	CheckLLMPolicyLimits(ctx context.Context, in *CheckLLMPolicyLimitsRequest, opts ...grpc.CallOption) (*CheckLLMPolicyLimitsResponse, error)
+	// RecordLLMUsage is the post-flight RPC the proxy calls after the upstream
+	// returns. Increments the per-(dimension, window) counters for the
+	// attribution policy chosen by CheckLLMPolicyLimits.
+	RecordLLMUsage(ctx context.Context, in *RecordLLMUsageRequest, opts ...grpc.CallOption) (*RecordLLMUsageResponse, error)
 }
 
 type proxyServiceClient struct {
@@ -179,6 +189,24 @@ func (c *proxyServiceClient) ValidateTunnelPeer(ctx context.Context, in *Validat
 	return out, nil
 }
 
+func (c *proxyServiceClient) CheckLLMPolicyLimits(ctx context.Context, in *CheckLLMPolicyLimitsRequest, opts ...grpc.CallOption) (*CheckLLMPolicyLimitsResponse, error) {
+	out := new(CheckLLMPolicyLimitsResponse)
+	err := c.cc.Invoke(ctx, "/management.ProxyService/CheckLLMPolicyLimits", in, out, opts...)
+	if err != nil {
+		return nil, err
+	}
+	return out, nil
+}
+
+func (c *proxyServiceClient) RecordLLMUsage(ctx context.Context, in *RecordLLMUsageRequest, opts ...grpc.CallOption) (*RecordLLMUsageResponse, error) {
+	out := new(RecordLLMUsageResponse)
+	err := c.cc.Invoke(ctx, "/management.ProxyService/RecordLLMUsage", in, out, opts...)
+	if err != nil {
+		return nil, err
+	}
+	return out, nil
+}
+
 // ProxyServiceServer is the server API for ProxyService service.
 // All implementations must embed UnimplementedProxyServiceServer
 // for forward compatibility
@@ -208,6 +236,16 @@ type ProxyServiceServer interface {
 	// issue a session cookie without redirecting through the OIDC flow.
 	// Mirrors ValidateSession's response shape.
 	ValidateTunnelPeer(context.Context, *ValidateTunnelPeerRequest) (*ValidateTunnelPeerResponse, error)
+	// CheckLLMPolicyLimits is the pre-flight RPC the proxy calls before each
+	// LLM request. Management runs the per-policy headroom selection across
+	// every policy authorising the caller's user / groups for the resolved
+	// provider and returns the chosen attribution policy + group, or a deny
+	// when no applicable policy has headroom > 0.
+	CheckLLMPolicyLimits(context.Context, *CheckLLMPolicyLimitsRequest) (*CheckLLMPolicyLimitsResponse, error)
+	// RecordLLMUsage is the post-flight RPC the proxy calls after the upstream
+	// returns. Increments the per-(dimension, window) counters for the
+	// attribution policy chosen by CheckLLMPolicyLimits.
+	RecordLLMUsage(context.Context, *RecordLLMUsageRequest) (*RecordLLMUsageResponse, error)
 	mustEmbedUnimplementedProxyServiceServer()
 }
 
@@ -242,6 +280,12 @@ func (UnimplementedProxyServiceServer) ValidateSession(context.Context, *Validat
 func (UnimplementedProxyServiceServer) ValidateTunnelPeer(context.Context, *ValidateTunnelPeerRequest) (*ValidateTunnelPeerResponse, error) {
 	return nil, status.Errorf(codes.Unimplemented, "method ValidateTunnelPeer not implemented")
 }
+func (UnimplementedProxyServiceServer) CheckLLMPolicyLimits(context.Context, *CheckLLMPolicyLimitsRequest) (*CheckLLMPolicyLimitsResponse, error) {
+	return nil, status.Errorf(codes.Unimplemented, "method CheckLLMPolicyLimits not implemented")
+}
+func (UnimplementedProxyServiceServer) RecordLLMUsage(context.Context, *RecordLLMUsageRequest) (*RecordLLMUsageResponse, error) {
+	return nil, status.Errorf(codes.Unimplemented, "method RecordLLMUsage not implemented")
+}
 func (UnimplementedProxyServiceServer) mustEmbedUnimplementedProxyServiceServer() {}
 
 // UnsafeProxyServiceServer may be embedded to opt out of forward compatibility for this service.
@@ -428,6 +472,42 @@ func _ProxyService_ValidateTunnelPeer_Handler(srv interface{}, ctx context.Conte
 	return interceptor(ctx, in, info, handler)
 }
 
+func _ProxyService_CheckLLMPolicyLimits_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
+	in := new(CheckLLMPolicyLimitsRequest)
+	if err := dec(in); err != nil {
+		return nil, err
+	}
+	if interceptor == nil {
+		return srv.(ProxyServiceServer).CheckLLMPolicyLimits(ctx, in)
+	}
+	info := &grpc.UnaryServerInfo{
+		Server:     srv,
+		FullMethod: "/management.ProxyService/CheckLLMPolicyLimits",
+	}
+	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
+		return srv.(ProxyServiceServer).CheckLLMPolicyLimits(ctx, req.(*CheckLLMPolicyLimitsRequest))
+	}
+	return interceptor(ctx, in, info, handler)
+}
+
+func _ProxyService_RecordLLMUsage_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
+	in := new(RecordLLMUsageRequest)
+	if err := dec(in); err != nil {
+		return nil, err
+	}
+	if interceptor == nil {
+		return srv.(ProxyServiceServer).RecordLLMUsage(ctx, in)
+	}
+	info := &grpc.UnaryServerInfo{
+		Server:     srv,
+		FullMethod: "/management.ProxyService/RecordLLMUsage",
+	}
+	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
+		return srv.(ProxyServiceServer).RecordLLMUsage(ctx, req.(*RecordLLMUsageRequest))
+	}
+	return interceptor(ctx, in, info, handler)
+}
+
 // ProxyService_ServiceDesc is the grpc.ServiceDesc for ProxyService service.
 // It's only intended for direct use with grpc.RegisterService,
 // and not to be introspected or modified (even as a copy)
@@ -463,6 +543,14 @@ var ProxyService_ServiceDesc = grpc.ServiceDesc{
 			MethodName: "ValidateTunnelPeer",
 			Handler:    _ProxyService_ValidateTunnelPeer_Handler,
 		},
+		{
+			MethodName: "CheckLLMPolicyLimits",
+			Handler:    _ProxyService_CheckLLMPolicyLimits_Handler,
+		},
+		{
+			MethodName: "RecordLLMUsage",
+			Handler:    _ProxyService_RecordLLMUsage_Handler,
+		},
 	},
 	Streams: []grpc.StreamDesc{
 		{

shared/management/status/error.go

@@ -219,6 +219,26 @@ func NewNetworkResourceNotFoundError(resourceID string) error {
 	return Errorf(NotFound, "network resource: %s not found", resourceID)
 }
 
+// NewAgentNetworkProviderNotFoundError creates a new Error with NotFound type for a missing Agent Network provider.
+func NewAgentNetworkProviderNotFoundError(providerID string) error {
+	return Errorf(NotFound, "agent network provider: %s not found", providerID)
+}
+
+// NewAgentNetworkPolicyNotFoundError creates a new Error with NotFound type for a missing Agent Network policy.
+func NewAgentNetworkPolicyNotFoundError(policyID string) error {
+	return Errorf(NotFound, "agent network policy: %s not found", policyID)
+}
+
+// NewAgentNetworkGuardrailNotFoundError creates a new Error with NotFound type for a missing Agent Network guardrail.
+func NewAgentNetworkGuardrailNotFoundError(guardrailID string) error {
+	return Errorf(NotFound, "agent network guardrail: %s not found", guardrailID)
+}
+
+// NewAgentNetworkBudgetRuleNotFoundError creates a new Error with NotFound type for a missing Agent Network budget rule.
+func NewAgentNetworkBudgetRuleNotFoundError(ruleID string) error {
+	return Errorf(NotFound, "agent network budget rule: %s not found", ruleID)
+}
+
 // NewPermissionDeniedError creates a new Error with PermissionDenied type for a permission denied error.
 func NewPermissionDeniedError() error {
 	return Errorf(PermissionDenied, "permission denied")

shared/relay/auth/allow/allow_all.go

@@ -8,7 +8,3 @@ type Auth struct {
 func (a *Auth) Validate(any) error {
 	return nil
 }
-
-func (a *Auth) ValidateHelloMsgType(any) error {
-	return nil
-}

shared/relay/auth/hmac/token.go

@@ -1,10 +1,8 @@
 package hmac
 
 import (
-	"bytes"
 	"crypto/hmac"
 	"encoding/base64"
-	"encoding/gob"
 	"fmt"
 	"hash"
 	"strconv"
@@ -18,14 +16,6 @@ type Token struct {
 	Signature string
 }
 
-func unmarshalToken(payload []byte) (Token, error) {
-	var creds Token
-	buffer := bytes.NewBuffer(payload)
-	decoder := gob.NewDecoder(buffer)
-	err := decoder.Decode(&creds)
-	return creds, err
-}
-
 // TimedHMAC generates a token with TTL and uses a pre-shared secret known to the relay server
 type TimedHMAC struct {
 	secret     string

shared/relay/auth/hmac/validator.go

@@ -1,33 +0,0 @@
-package hmac
-
-import (
-	"crypto/sha256"
-	"fmt"
-	"time"
-
-	log "github.com/sirupsen/logrus"
-)
-
-type TimedHMACValidator struct {
-	*TimedHMAC
-}
-
-func NewTimedHMACValidator(secret string, duration time.Duration) *TimedHMACValidator {
-	ta := NewTimedHMAC(secret, duration)
-	return &TimedHMACValidator{
-		ta,
-	}
-}
-
-func (a *TimedHMACValidator) Validate(credentials any) error {
-	b, ok := credentials.([]byte)
-	if !ok {
-		return fmt.Errorf("invalid credentials type")
-	}
-	c, err := unmarshalToken(b)
-	if err != nil {
-		log.Debugf("failed to unmarshal token: %s", err)
-		return err
-	}
-	return a.TimedHMAC.Validate(sha256.New, c)
-}

shared/relay/auth/validator.go

@@ -1,28 +1,19 @@
 package auth
 
 import (
-	"time"
-
-	auth "github.com/netbirdio/netbird/shared/relay/auth/hmac"
 	authv2 "github.com/netbirdio/netbird/shared/relay/auth/hmac/v2"
 )
 
 type TimedHMACValidator struct {
 	authenticatorV2 *authv2.Validator
-	authenticator   *auth.TimedHMACValidator
 }
 
-func NewTimedHMACValidator(secret []byte, duration time.Duration) *TimedHMACValidator {
+func NewTimedHMACValidator(secret []byte) *TimedHMACValidator {
 	return &TimedHMACValidator{
 		authenticatorV2: authv2.NewValidator(secret),
-		authenticator:   auth.NewTimedHMACValidator(string(secret), duration),
 	}
 }
 
 func (a *TimedHMACValidator) Validate(credentials any) error {
 	return a.authenticatorV2.Validate(credentials)
 }
-
-func (a *TimedHMACValidator) ValidateHelloMsgType(credentials any) error {
-	return a.authenticator.Validate(credentials)
-}

shared/relay/client/dialer/ws/close_generic.go

@@ -0,0 +1,9 @@
+//go:build !js
+
+package ws
+
+// closeConn closes the underlying WebSocket immediately, skipping the close
+// handshake.
+func (c *Conn) closeConn() error {
+	return c.Conn.CloseNow()
+}

shared/relay/client/dialer/ws/close_js.go

@@ -0,0 +1,25 @@
+//go:build js
+
+package ws
+
+import (
+	"github.com/coder/websocket"
+	log "github.com/sirupsen/logrus"
+)
+
+// closeConn closes the browser WebSocket without blocking the caller.
+//
+// The browser close API only accepts codes 1000 and 3000-4999, so CloseNow's
+// 1001 (going away) throws an InvalidAccessError. Close with a valid code
+// waits for the browser close event before returning, which can park the
+// calling goroutine (the relay teardown path holds its mutexes while closing)
+// until the close handshake finishes. Run the close in the background and
+// report success; a teardown close error is not actionable.
+func (c *Conn) closeConn() error {
+	go func() {
+		if err := c.Conn.Close(websocket.StatusNormalClosure, ""); err != nil {
+			log.Debugf("failed to close relay websocket: %v", err)
+		}
+	}()
+	return nil
+}

shared/relay/client/dialer/ws/conn.go

@@ -77,5 +77,5 @@ func (c *Conn) SetDeadline(t time.Time) error {
 }
 
 func (c *Conn) Close() error {
-	return c.Conn.CloseNow()
+	return c.closeConn()
 }

shared/relay/client/early_msg_buffer.go

@@ -10,7 +10,7 @@ import (
 
 const (
 	earlyMsgTTL      = 5 * time.Second
-	earlyMsgCapacity = 1000
+	earlyMsgCapacity = 10000
 )
 
 // earlyMsgBuffer buffers transport messages that arrive before the corresponding

shared/relay/client/manager_cleanup_test.go

@@ -0,0 +1,50 @@
+package client
+
+import (
+	"context"
+	"testing"
+	"time"
+)
+
+// TestCleanUpUnusedRelays_DoesNotBlockOnRealHangingDial drives a real, hanging foreign
+// relay dial and asserts the foreign store cleanup does not stall behind it.
+func TestCleanUpUnusedRelays_DoesNotBlockOnRealHangingDial(t *testing.T) {
+	serverAddr, accepted := stallingRelayListener(t)
+
+	mCtx, mCancel := context.WithCancel(context.Background())
+	t.Cleanup(mCancel)
+
+	m := NewManager(mCtx, nil, "alice", 1280)
+
+	dialDone := make(chan struct{})
+	go func() {
+		defer close(dialDone)
+		_, _ = m.foreign.OpenConn(mCtx, "peerKey", RelayServer{Addr: serverAddr})
+	}()
+
+	select {
+	case <-accepted:
+	case <-time.After(5 * time.Second):
+		t.Fatal("relay dial did not reach the listener")
+	}
+
+	cleanupDone := make(chan struct{})
+	go func() {
+		defer close(cleanupDone)
+		m.foreign.cleanupUnused()
+	}()
+
+	select {
+	case <-cleanupDone:
+	case <-time.After(2 * time.Second):
+		t.Fatal("cleanupUnused blocked on an in-progress relay dial")
+	}
+
+	// Release the hanging dial so the goroutine can exit cleanly.
+	mCancel()
+	select {
+	case <-dialDone:
+	case <-time.After(5 * time.Second):
+		t.Fatal("foreign OpenConn did not return after context cancellation")
+	}
+}

shared/relay/client/manager_relaystates_test.go

@@ -0,0 +1,96 @@
+package client
+
+import (
+	"context"
+	"net"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+)
+
+// stallingRelayListener accepts TCP connections and holds them open without ever
+// responding, so a relay handshake dialed against it blocks until its context is
+// cancelled. accepted is signalled once per incoming connection so a caller can
+// wait until a dial has actually reached the listener. It returns the
+// "rel://host:port" URL to dial.
+func stallingRelayListener(t *testing.T) (string, <-chan struct{}) {
+	t.Helper()
+
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	require.NoError(t, err)
+
+	accepted := make(chan struct{}, 1)
+	var mu sync.Mutex
+	var conns []net.Conn
+	go func() {
+		for {
+			c, err := ln.Accept()
+			if err != nil {
+				return
+			}
+			mu.Lock()
+			conns = append(conns, c)
+			mu.Unlock()
+			select {
+			case accepted <- struct{}{}:
+			default:
+			}
+		}
+	}()
+	t.Cleanup(func() {
+		_ = ln.Close()
+		mu.Lock()
+		for _, c := range conns {
+			_ = c.Close()
+		}
+		mu.Unlock()
+	})
+
+	return "rel://" + ln.Addr().String(), accepted
+}
+
+// TestRelayStates_DoesNotBlockOnRealHangingDial is a regression test for
+// RelayStates() called by a "status -d command" hanging behind an in-progress
+// foreign relay dial.
+func TestRelayStates_DoesNotBlockOnRealHangingDial(t *testing.T) {
+	serverAddr, accepted := stallingRelayListener(t)
+
+	mCtx, mCancel := context.WithCancel(context.Background())
+	t.Cleanup(mCancel)
+
+	m := NewManager(mCtx, nil, "alice", 1280)
+
+	dialDone := make(chan struct{})
+	go func() {
+		defer close(dialDone)
+		_, _ = m.foreign.OpenConn(mCtx, "peerKey", RelayServer{Addr: serverAddr})
+	}()
+
+	select {
+	case <-accepted:
+	case <-time.After(5 * time.Second):
+		t.Fatal("relay dial did not reach the listener")
+	}
+
+	done := make(chan []RelayConnState, 1)
+	go func() {
+		done <- m.RelayStates()
+	}()
+
+	select {
+	case states := <-done:
+		require.Empty(t, states, "a relay still being dialed carries no state and must be omitted")
+	case <-time.After(2 * time.Second):
+		t.Fatal("RelayStates blocked on a foreign relay whose Connect() is in progress")
+	}
+
+	// Release the hanging dial so the goroutine can exit cleanly.
+	mCancel()
+	select {
+	case <-dialDone:
+	case <-time.After(5 * time.Second):
+		t.Fatal("foreign OpenConn did not return after context cancellation")
+	}
+}

shared/relay/messages/address/address.go

@@ -1,21 +0,0 @@
-// Deprecated: This package is deprecated and will be removed in a future release.
-package address
-
-import (
-	"bytes"
-	"encoding/gob"
-	"fmt"
-)
-
-type Address struct {
-	URL string
-}
-
-func (addr *Address) Marshal() ([]byte, error) {
-	var buf bytes.Buffer
-	enc := gob.NewEncoder(&buf)
-	if err := enc.Encode(addr); err != nil {
-		return nil, fmt.Errorf("encode Address: %w", err)
-	}
-	return buf.Bytes(), nil
-}

shared/relay/messages/auth/auth.go

@@ -1,43 +0,0 @@
-// Deprecated: This package is deprecated and will be removed in a future release.
-package auth
-
-import (
-	"bytes"
-	"encoding/gob"
-	"fmt"
-)
-
-type Algorithm int
-
-const (
-	AlgoUnknown Algorithm = iota
-	AlgoHMACSHA256
-	AlgoHMACSHA512
-)
-
-func (a Algorithm) String() string {
-	switch a {
-	case AlgoHMACSHA256:
-		return "HMAC-SHA256"
-	case AlgoHMACSHA512:
-		return "HMAC-SHA512"
-	default:
-		return "Unknown"
-	}
-}
-
-type Msg struct {
-	AuthAlgorithm  Algorithm
-	AdditionalData []byte
-}
-
-func UnmarshalMsg(data []byte) (*Msg, error) {
-	var msg *Msg
-
-	buf := bytes.NewBuffer(data)
-	dec := gob.NewDecoder(buf)
-	if err := dec.Decode(&msg); err != nil {
-		return nil, fmt.Errorf("decode Msg: %w", err)
-	}
-	return msg, nil
-}

shared/relay/messages/message.go

@@ -14,9 +14,10 @@ const (
 	CurrentProtocolVersion = 1
 
 	MsgTypeUnknown MsgType = 0
-	// Deprecated: Use MsgTypeAuth instead.
-	MsgTypeHello = 1
-	// Deprecated: Use MsgTypeAuthResponse instead.
+	// MsgTypeHello and MsgTypeHelloResponse are the removed legacy handshake
+	// message types. They are retained only to reserve wire values 1 and 2 so
+	// the values are never reused; the server rejects both.
+	MsgTypeHello         = 1
 	MsgTypeHelloResponse = 2
 	MsgTypeTransport     = 3
 	MsgTypeClose         = 4
@@ -42,10 +43,6 @@ const (
 	offsetAuthPeerID    = sizeOfProtoHeader + sizeOfMagicByte
 	headerTotalSizeAuth = sizeOfProtoHeader + headerSizeAuth
 
-	// hello message
-	headerSizeHello     = sizeOfMagicByte + peerIDSize
-	headerSizeHelloResp = 0
-
 	// transport
 	headerSizeTransport      = peerIDSize
 	offsetTransportID        = sizeOfProtoHeader
@@ -113,7 +110,6 @@ func DetermineClientMessageType(msg []byte) (MsgType, error) {
 	msgType := MsgType(msg[1])
 	switch msgType {
 	case
-		MsgTypeHello,
 		MsgTypeAuth,
 		MsgTypeTransport,
 		MsgTypeClose,
@@ -135,7 +131,6 @@ func DetermineServerMessageType(msg []byte) (MsgType, error) {
 	msgType := MsgType(msg[1])
 	switch msgType {
 	case
-		MsgTypeHelloResponse,
 		MsgTypeAuthResponse,
 		MsgTypeTransport,
 		MsgTypeClose,
@@ -148,67 +143,6 @@ func DetermineServerMessageType(msg []byte) (MsgType, error) {
 	}
 }
 
-// Deprecated: Use MarshalAuthMsg instead.
-// MarshalHelloMsg initial hello message
-// The Hello message is the first message sent by a client after establishing a connection with the Relay server. This
-// message is used to authenticate the client with the server. The authentication is done using an HMAC method.
-// The protocol does not limit to use HMAC, it can be any other method. If the authentication failed the server will
-// close the network connection without any response.
-func MarshalHelloMsg(peerID PeerID, additions []byte) ([]byte, error) {
-	msg := make([]byte, sizeOfProtoHeader+sizeOfMagicByte, sizeOfProtoHeader+headerSizeHello+len(additions))
-
-	msg[0] = byte(CurrentProtocolVersion)
-	msg[1] = byte(MsgTypeHello)
-
-	copy(msg[sizeOfProtoHeader:sizeOfProtoHeader+sizeOfMagicByte], magicHeader)
-
-	msg = append(msg, peerID[:]...)
-	msg = append(msg, additions...)
-
-	return msg, nil
-}
-
-// Deprecated: Use UnmarshalAuthMsg instead.
-// UnmarshalHelloMsg extracts peerID and the additional data from the hello message. The Additional data is used to
-// authenticate the client with the server.
-func UnmarshalHelloMsg(msg []byte) (*PeerID, []byte, error) {
-	if len(msg) < sizeOfProtoHeader+headerSizeHello {
-		return nil, nil, ErrInvalidMessageLength
-	}
-	if !bytes.Equal(msg[sizeOfProtoHeader:sizeOfProtoHeader+sizeOfMagicByte], magicHeader) {
-		return nil, nil, errors.New("invalid magic header")
-	}
-
-	peerID := PeerID(msg[sizeOfProtoHeader+sizeOfMagicByte : sizeOfProtoHeader+headerSizeHello])
-
-	return &peerID, msg[headerSizeHello:], nil
-}
-
-// Deprecated: Use MarshalAuthResponse instead.
-// MarshalHelloResponse creates a response message to the hello message.
-// In case of success connection the server response with a Hello Response message. This message contains the server's
-// instance URL. This URL will be used by choose the common Relay server in case if the peers are in different Relay
-// servers.
-func MarshalHelloResponse(additionalData []byte) ([]byte, error) {
-	msg := make([]byte, sizeOfProtoHeader, sizeOfProtoHeader+headerSizeHelloResp+len(additionalData))
-
-	msg[0] = byte(CurrentProtocolVersion)
-	msg[1] = byte(MsgTypeHelloResponse)
-
-	msg = append(msg, additionalData...)
-
-	return msg, nil
-}
-
-// Deprecated: Use UnmarshalAuthResponse instead.
-// UnmarshalHelloResponse extracts the additional data from the hello response message.
-func UnmarshalHelloResponse(msg []byte) ([]byte, error) {
-	if len(msg) < sizeOfProtoHeader+headerSizeHelloResp {
-		return nil, ErrInvalidMessageLength
-	}
-	return msg, nil
-}
-
 // MarshalAuthMsg initial authentication message
 // The Auth message is the first message sent by a client after establishing a connection with the Relay server. This
 // message is used to authenticate the client with the server. The authentication is done using an HMAC method.

shared/relay/messages/message_test.go

@@ -4,28 +4,11 @@ import (
 	"testing"
 )
 
-func TestMarshalHelloMsg(t *testing.T) {
-	peerID := HashID("abdFAaBcawquEiCMzAabYosuUaGLtSNhKxz+")
-	msg, err := MarshalHelloMsg(peerID, nil)
-	if err != nil {
-		t.Fatalf("error: %v", err)
-	}
-
-	msgType, err := DetermineClientMessageType(msg)
-	if err != nil {
-		t.Fatalf("error: %v", err)
-	}
-
-	if msgType != MsgTypeHello {
-		t.Errorf("expected %d, got %d", MsgTypeHello, msgType)
-	}
-
-	receivedPeerID, _, err := UnmarshalHelloMsg(msg)
-	if err != nil {
-		t.Fatalf("error: %v", err)
-	}
-	if receivedPeerID.String() != peerID.String() {
-		t.Errorf("expected %s, got %s", peerID, receivedPeerID)
+func TestDetermineClientMessageTypeRejectsHello(t *testing.T) {
+	// The reserved legacy Hello message (type 1) must be rejected by the server.
+	msg := []byte{byte(CurrentProtocolVersion), byte(MsgTypeHello)}
+	if _, err := DetermineClientMessageType(msg); err == nil {
+		t.Fatalf("expected hello message type to be rejected")
 	}
 }
 

shared/signal/client/grpc.go

@@ -85,6 +85,7 @@ type GrpcClient struct {
 	// receive backpressure as a dead stream: reconnecting cannot help, since the
 	// new stream feeds the same worker, and only triggers a reconnect storm.
 	receiveHandoffBlocked atomic.Bool
+	watchdogWg            sync.WaitGroup
 }
 
 // NewClient creates a new Signal client
@@ -200,10 +201,18 @@ func (c *GrpcClient) Receive(ctx context.Context, msgHandler func(msg *proto.Mes
 		// Guard the receive direction: the transport can stay healthy while the
 		// server stops delivering messages. The watchdog reconnects via cancelStream.
 		c.markReceived()
-		go c.watchReceiveStream(streamCtx, cancelStream)
+		c.watchdogWg.Add(1)
+		go func() {
+			defer c.watchdogWg.Done()
+			c.watchReceiveStream(streamCtx, cancelStream)
+		}()
 
 		// start receiving messages from the Signal stream (from other peers through signal)
 		err = c.receive(stream)
+
+		cancelStream()
+		c.watchdogWg.Wait()
+
 		if err != nil {
 			// Check the parent context, not streamCtx: a watchdog-triggered
 			// cancelStream must reconnect, only a parent cancel is shutdown.
@@ -400,7 +409,12 @@ func (c *GrpcClient) encryptMessage(msg *proto.Message) (*proto.EncryptedMessage
 
 // Send sends a message to the remote Peer through the Signal Exchange.
 func (c *GrpcClient) Send(msg *proto.Message) error {
+	return c.send(c.ctx, msg)
+}
 
+// send delivers a message deriving per-attempt timeouts from parentCtx, so a
+// caller can abort an in-flight send by cancelling that context.
+func (c *GrpcClient) send(parentCtx context.Context, msg *proto.Message) error {
 	if !c.Ready() {
 		return fmt.Errorf("no connection to signal")
 	}
@@ -416,7 +430,7 @@ func (c *GrpcClient) Send(msg *proto.Message) error {
 		if attempt > 1 {
 			attemptTimeout = time.Duration(attempt) * 5 * time.Second
 		}
-		ctx, cancel := context.WithTimeout(c.ctx, attemptTimeout)
+		ctx, cancel := context.WithTimeout(parentCtx, attemptTimeout)
 
 		_, err = c.realClient.Send(ctx, encryptedMessage)
 
@@ -486,7 +500,7 @@ func (c *GrpcClient) watchReceiveStream(ctx context.Context, cancelStream contex
 			}
 
 			if probeSentAt.IsZero() {
-				if err := c.sendReceiveProbe(); err != nil {
+				if err := c.sendReceiveProbe(ctx); err != nil {
 					log.Debugf("failed to send signal receive probe: %v", err)
 				}
 				probeSentAt = time.Now()
@@ -495,11 +509,13 @@ func (c *GrpcClient) watchReceiveStream(ctx context.Context, cancelStream contex
 	}
 }
 
-// sendReceiveProbe sends a self-addressed heartbeat. The Signal server routes it
-// back to this client, exercising the exact receive path the watchdog guards.
-func (c *GrpcClient) sendReceiveProbe() error {
+// sendReceiveProbe sends a self-addressed heartbeat bound to ctx, so cancelStream
+// aborts an in-flight probe instead of leaving the watchdog blocked on send timeouts.
+// The Signal server routes it back to this client, exercising the exact receive
+// path the watchdog guards.
+func (c *GrpcClient) sendReceiveProbe(ctx context.Context) error {
 	self := c.key.PublicKey().String()
-	return c.Send(&proto.Message{
+	return c.send(ctx, &proto.Message{
 		Key:       self,
 		RemoteKey: self,
 		Body:      &proto.Body{Type: proto.Body_HEARTBEAT},
@@ -541,6 +557,9 @@ func (c *GrpcClient) receive(stream proto.SignalExchange_ConnectStreamClient) er
 		if err := c.decryptionWorker.AddMsg(c.ctx, msg); err != nil {
 			log.Errorf("failed to add message to decryption worker: %v", err)
 		}
+		// Refresh liveness before clearing the flag so the window between here and
+		// the next Recv does not read a stale timestamp as a dead stream.
+		c.markReceived()
 		c.receiveHandoffBlocked.Store(false)
 	}
 }

shared/signal/client/watchdog_test.go

@@ -2,6 +2,7 @@ package client
 
 import (
 	"context"
+	"io"
 	"net"
 	"testing"
 	"time"
@@ -74,7 +75,7 @@ func TestReceiveProbeRoundTrips(t *testing.T) {
 		t.Fatal("signal stream did not connect within timeout")
 	}
 
-	require.NoError(t, client.sendReceiveProbe())
+	require.NoError(t, client.sendReceiveProbe(ctx))
 
 	select {
 	case <-received:
@@ -106,3 +107,72 @@ func TestReceiveAliveTreatsHandoffBlockAsLiveness(t *testing.T) {
 	c.markReceived()
 	require.True(t, c.receiveAlive(), "a freshly received frame must keep the stream alive")
 }
+
+// fakeRecvStream feeds the receive loop frames from a channel and reports EOF
+// once the channel is closed. Only Recv is exercised by the loop.
+type fakeRecvStream struct {
+	sigProto.SignalExchange_ConnectStreamClient
+	frames chan *sigProto.EncryptedMessage
+}
+
+func (s *fakeRecvStream) Recv() (*sigProto.EncryptedMessage, error) {
+	msg, ok := <-s.frames
+	if !ok {
+		return nil, io.EOF
+	}
+	return msg, nil
+}
+
+// TestReceiveLoopRefreshesLivenessAfterBlockedHandoff drives the real receive
+// loop into a handoff that blocks past the inactivity threshold, then checks the
+// window after the handoff drains but before the next Recv. The loop must have
+// refreshed the timestamp on unblocking, otherwise that window reads the stale
+// pre-handoff timestamp as a dead stream and the watchdog tears down a healthy
+// connection.
+func TestReceiveLoopRefreshesLivenessAfterBlockedHandoff(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	t.Cleanup(cancel)
+	c := &GrpcClient{ctx: ctx}
+
+	handling := make(chan struct{}, 8)
+	gate := make(chan struct{})
+	decrypt := func(*sigProto.EncryptedMessage) (*sigProto.Message, error) { return &sigProto.Message{}, nil }
+	handler := func(*sigProto.Message) error {
+		handling <- struct{}{}
+		<-gate
+		return nil
+	}
+	c.decryptionWorker = NewWorker(decrypt, handler)
+	workerCtx, workerCancel := context.WithCancel(context.Background())
+	go c.decryptionWorker.Work(workerCtx)
+	t.Cleanup(workerCancel)
+
+	frames := make(chan *sigProto.EncryptedMessage)
+	t.Cleanup(func() { close(frames) })
+	go func() { _ = c.receive(&fakeRecvStream{frames: frames}) }()
+
+	// First frame: the worker drains it and parks in the blocking handler.
+	frames <- &sigProto.EncryptedMessage{}
+	<-handling
+	// Second frame fills the worker's single-slot pool.
+	frames <- &sigProto.EncryptedMessage{}
+	// Third frame: the pool is full, so the loop parks on the handoff.
+	frames <- &sigProto.EncryptedMessage{}
+
+	require.Eventually(t, c.receiveHandoffBlocked.Load, time.Second, time.Millisecond,
+		"receive loop should park on the worker handoff")
+
+	// Simulate the handoff having blocked past the inactivity threshold.
+	c.lastReceived.Store(time.Now().Add(-2 * receiveInactivityThreshold).UnixNano())
+	require.True(t, c.receiveAlive(), "a loop parked on the handoff must stay alive")
+
+	// Drain the worker so the handoff returns and the loop resumes reading.
+	close(gate)
+
+	// Once the handoff clears, the loop is parked on the next Recv with no frame
+	// pending. The stream must still read as alive in that window.
+	require.Eventually(t, func() bool { return !c.receiveHandoffBlocked.Load() }, time.Second, time.Millisecond,
+		"handoff should drain once the worker is released")
+	require.True(t, c.receiveAlive(),
+		"the loop must refresh liveness when the handoff drains, before the next Recv")
+}

shared/testing_helpers/populate_fields.go

@@ -0,0 +1,101 @@
+package testing_helpers
+
+import (
+	"fmt"
+	"net/netip"
+	"reflect"
+)
+
+type PopulateFields struct {
+	CustomFieldSetters map[reflect.Type]func(this *PopulateFields, field reflect.Value) (int, error)
+	TagsToSkip         map[string]string
+}
+
+func NewPopulateFields() *PopulateFields {
+	return &PopulateFields{CustomFieldSetters: defaultCustomFieldSetters(), TagsToSkip: make(map[string]string)}
+}
+
+func (p *PopulateFields) WithCustomFieldSetter(t reflect.Type, f func(this *PopulateFields, field reflect.Value) (int, error)) *PopulateFields {
+	p.CustomFieldSetters[t] = f
+	return p
+}
+
+func (p *PopulateFields) WithSkippedTag(tag, value string) *PopulateFields {
+	p.TagsToSkip[tag] = value
+	return p
+}
+
+func (p *PopulateFields) PopulateAll(v reflect.Value) (int, error) {
+	typ := v.Type()
+	totalExportedFields := 0
+	for i := 0; i < typ.NumField(); i++ {
+		f := typ.Field(i)
+		if f.PkgPath != "" { // unexported
+			continue
+		}
+
+		if p.skippedTagPresent(f.Tag) {
+			continue
+		}
+
+		numOfExportedFields, err := p.setNonZero(v.Field(i))
+		totalExportedFields += numOfExportedFields
+		if err != nil {
+			return totalExportedFields, err
+		}
+	}
+	return totalExportedFields, nil
+}
+
+// setNonZero assigns a deterministic non-zero value to a field based on its kind,
+// recursing into nested structs and populating one element of slice fields.
+func (p *PopulateFields) setNonZero(field reflect.Value) (int, error) {
+	if f, ok := p.CustomFieldSetters[field.Type()]; ok {
+		return f(p, field)
+	}
+
+	switch field.Kind() {
+	case reflect.String:
+		field.SetString("non-zero")
+	case reflect.Bool:
+		field.SetBool(true)
+	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
+		field.SetInt(7)
+	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
+		field.SetUint(7)
+	case reflect.Float32, reflect.Float64:
+		field.SetFloat(7)
+	case reflect.Struct:
+		n, err := p.PopulateAll(field)
+		return n + 1, err
+	case reflect.Slice:
+		s := reflect.MakeSlice(field.Type(), 1, 1)
+		_, err := p.setNonZero(s.Index(0))
+		if err != nil {
+			return 0, err
+		}
+		field.Set(s)
+	default:
+		return 0, fmt.Errorf("unhandled field kind %s; extend setNonZero", field.Kind())
+	}
+
+	return 1, nil
+}
+
+func defaultCustomFieldSetters() map[reflect.Type]func(this *PopulateFields, field reflect.Value) (int, error) {
+	return map[reflect.Type]func(this *PopulateFields, field reflect.Value) (int, error){
+		reflect.TypeOf(netip.Prefix{}): func(_ *PopulateFields, field reflect.Value) (int, error) {
+			field.Set(reflect.ValueOf(netip.MustParsePrefix("10.0.0.0/24")))
+			return 1, nil
+		},
+	}
+}
+
+func (p *PopulateFields) skippedTagPresent(t reflect.StructTag) bool {
+	for tag, value := range p.TagsToSkip {
+		if v := t.Get(tag); v == value {
+			return true
+		}
+	}
+	return false
+}