[client] Bound the Signal stream registration wait and track its state based on stream health (#8175)

* [client] Move the Signal registration wait into its own helper

Pure move, no behavior change. The wait for the registration header and
the check on its value leave connect() for awaitRegistration(), so the
next commit can bound that wait in one place.

* [client] Split the Signal test server helper from its listener setup

Pure move, no behavior change. serveTestSignalExchange serves any
SignalExchangeServer on a local listener; startTestSignalServer keeps
serving the real server through it. Lets a test wrap the real server.

* [client] Bound the wait for the Signal registration header

After an outage the Signal client could wait forever for the registration
header of a new stream. Keepalive cannot end that wait when whatever
answers the transport (a proxy in front of Signal, or Signal itself)
accepts the stream and never answers it, and the receive watchdog only
starts once registration succeeds. Signal then stayed down until the
daemon was restarted, while Management and Relay recovered.

Cancel the stream when the header does not arrive within 10 seconds, so
connect() fails and the existing retry loop opens a new stream. A timer
is used instead of a context deadline because the stream context must
outlive registration.

* [client] Report Signal unhealthy while no stream is registered

IsHealthy judged Signal by the gRPC connection state alone and marked it
connected whenever the connection was Ready or Idle. With the connection
up but no stream registered, the client cannot receive offers or answers,
yet netbird status showed Signal: Connected, hiding a stuck registration
for as long as it lasted.

Require a registered stream before reporting healthy. When there is none,
return false without touching the recorder: the disconnect that dropped
the stream already reported its reason.

* [client] Name the server that reconnected in the SR watcher log

Signal and Relay shared one reconnect callback, which logged "reconnected
to Signal or Relay server" for either. A Relay reconnect while the Signal
stream was down read as Signal having recovered, which misled diagnosis
of a stuck Signal registration.

Register one callback per server so the log names the one that actually
reconnected. Which listeners get notified, and when, is unchanged.

* [client] Trim comments in the Signal registration and health code

Comment-only change, no behavior change. Drop comments that restate the
code and keep only the reason a timer is used instead of a context
deadline.

* [client] Check Signal stream registration before the connection state

IsHealthy returned true for a Connecting or Shutdown channel before the
registration check ran, so a client with no registered stream could
still be reported healthy while the channel was reconnecting. Check the
stream first; the connection-state cases now only apply to a client
that has a registered stream.

* [client] Log a plain connect in the SR watcher

The Signal and Relay callbacks also fire on the first connection, not
only after a disconnect, so "reconnected" was wrong at startup. Log
"connected to <server> server" for both cases.

* [client] Say what the SR watcher does when it logs

The line was read as a per-server connect log, which Signal and Relay
already emit themselves. What it marks is the peer guards being told to
reset their reconnection backoff, which happens only once both are
ready. Log that, naming the server whose event triggered it.

client/internal/peer/guard/sr_watcher.go

@@ -54,8 +54,8 @@ func (w *SRWatcher) Start(disableICEMonitor bool) {
 		iceMonitor := NewICEMonitor(w.iFaceDiscover, w.iceConfig, GetICEMonitorPeriod())
 		go iceMonitor.Start(ctx, w.onICEChanged)
 	}
-	w.signalClient.SetOnReconnectedListener(w.onReconnected)
-	w.relayManager.SetOnReconnectedListener(w.onReconnected)
+	w.signalClient.SetOnReconnectedListener(w.onSignalReconnected)
+	w.relayManager.SetOnReconnectedListener(w.onRelayReconnected)
 
 }
 
@@ -96,7 +96,15 @@ func (w *SRWatcher) onICEChanged() {
 	w.notify()
 }
 
-func (w *SRWatcher) onReconnected() {
+func (w *SRWatcher) onSignalReconnected() {
+	w.onReconnected("Signal")
+}
+
+func (w *SRWatcher) onRelayReconnected() {
+	w.onReconnected("Relay")
+}
+
+func (w *SRWatcher) onReconnected(server string) {
 	if !w.signalClient.Ready() {
 		return
 	}
@@ -104,7 +112,7 @@ func (w *SRWatcher) onReconnected() {
 		return
 	}
 
-	log.Infof("reconnected to Signal or Relay server")
+	log.Infof("Signal and Relay ready (triggered by %s), resetting peer reconnection backoff", server)
 	w.notify()
 }
 

shared/signal/client/grpc.go

@@ -38,6 +38,8 @@ const (
 	receiveProbeTimeout = 10 * time.Second
 	// receiveWatchdogInterval is how often the watchdog evaluates the stream.
 	receiveWatchdogInterval = 10 * time.Second
+	// defaultRegistrationTimeout bounds the wait for a new stream's registration header.
+	defaultRegistrationTimeout = 10 * time.Second
 )
 
 // errReceiveStreamStalled is reported when the receive stream is transport-alive
@@ -92,6 +94,8 @@ type GrpcClient struct {
 	// new stream feeds the same worker, and only triggers a reconnect storm.
 	receiveHandoffBlocked atomic.Bool
 	watchdogWg            sync.WaitGroup
+
+	registrationTimeout time.Duration
 }
 
 // Option configures optional GrpcClient behavior.
@@ -112,6 +116,7 @@ func NewClient(ctx context.Context, addr string, key wgtypes.Key, tlsEnabled boo
 		mux:                   sync.Mutex{},
 		status:                StreamDisconnected,
 		connStateCallbackLock: sync.RWMutex{},
+		registrationTimeout:   defaultRegistrationTimeout,
 	}
 	for _, opt := range opts {
 		opt(c)
@@ -230,7 +235,7 @@ func (c *GrpcClient) Receive(ctx context.Context, msgHandler func(msg *proto.Mes
 		// todo once the key rotation logic has been implemented, consider changing to some other identifier (received from management)
 		streamCtx, cancelStream := context.WithCancel(ctx)
 		defer cancelStream()
-		stream, err := c.connect(streamCtx, c.key.PublicKey().String())
+		stream, err := c.connect(streamCtx, cancelStream, c.key.PublicKey().String())
 		if err != nil {
 			log.Warnf("disconnected from the Signal Exchange due to an error: %v", err)
 			return err
@@ -308,7 +313,7 @@ func (c *GrpcClient) notifyStreamConnected() {
 	}
 }
 
-func (c *GrpcClient) connect(ctx context.Context, key string) (proto.SignalExchange_ConnectStreamClient, error) {
+func (c *GrpcClient) connect(ctx context.Context, cancel context.CancelFunc, key string) (proto.SignalExchange_ConnectStreamClient, error) {
 	c.stream = nil
 
 	// add key fingerprint to the request header to be identified on the server side
@@ -319,19 +324,30 @@ func (c *GrpcClient) connect(ctx context.Context, key string) (proto.SignalExcha
 	if err != nil {
 		return nil, err
 	}
-	// blocks
-	header, err := c.stream.Header()
-	if err != nil {
+	if err := awaitRegistration(stream, c.registrationTimeout, cancel); err != nil {
 		return nil, err
 	}
-	registered := header.Get(proto.HeaderRegistered)
-	if len(registered) == 0 {
-		return nil, fmt.Errorf("didn't receive a registration header from the Signal server whille connecting to the streams")
-	}
 
 	return stream, nil
 }
 
+func awaitRegistration(stream proto.SignalExchange_ConnectStreamClient, timeout time.Duration, cancel context.CancelFunc) error {
+	// A timer rather than a context deadline: the stream context must outlive registration.
+	timer := time.AfterFunc(timeout, cancel)
+	header, err := stream.Header()
+	if !timer.Stop() {
+		return fmt.Errorf("no registration header from the Signal server within %s", timeout)
+	}
+	if err != nil {
+		return err
+	}
+	registered := header.Get(proto.HeaderRegistered)
+	if len(registered) == 0 {
+		return fmt.Errorf("didn't receive a registration header from the Signal server whille connecting to the streams")
+	}
+	return nil
+}
+
 // Ready indicates whether the client is okay and Ready to be used
 // for now it just checks whether gRPC connection to the service is in state Ready
 func (c *GrpcClient) Ready() bool {
@@ -343,6 +359,13 @@ func (c *GrpcClient) Ready() bool {
 // recorder accordingly. It does not actively probe: the watchdog
 // (watchReceiveStream) owns probing the receive path and reconnecting.
 func (c *GrpcClient) IsHealthy() bool {
+	c.mux.Lock()
+	registered := c.status == StreamConnected
+	c.mux.Unlock()
+	if !registered {
+		return false
+	}
+
 	switch c.signalConn.GetState() {
 	case connectivity.TransientFailure:
 		return false

shared/signal/client/health_test.go

@@ -0,0 +1,130 @@
+package client
+
+import (
+	"context"
+	"net"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/connectivity"
+	"google.golang.org/grpc/credentials/insecure"
+
+	sigProto "github.com/netbirdio/netbird/shared/signal/proto"
+)
+
+type recordingNotifier struct {
+	mu        sync.Mutex
+	connected bool
+}
+
+func (n *recordingNotifier) MarkSignalDisconnected(error) {
+	n.mu.Lock()
+	defer n.mu.Unlock()
+	n.connected = false
+}
+
+func (n *recordingNotifier) MarkSignalConnected() {
+	n.mu.Lock()
+	defer n.mu.Unlock()
+	n.connected = true
+}
+
+func (n *recordingNotifier) isConnected() bool {
+	n.mu.Lock()
+	defer n.mu.Unlock()
+	return n.connected
+}
+
+// TestIsHealthyRequiresRegisteredStream covers the status reporting half of a
+// stuck registration: the gRPC connection to Signal is up but no stream is
+// registered, so no peer can reach this client. The health check used to look
+// at the connection only, reporting Signal as connected throughout.
+func TestIsHealthyRequiresRegisteredStream(t *testing.T) {
+	addr := startTestSignalServer(t)
+
+	key, err := wgtypes.GenerateKey()
+	require.NoError(t, err)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	t.Cleanup(cancel)
+
+	client, err := NewClient(ctx, addr, key, false)
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = client.Close() })
+
+	notifier := &recordingNotifier{}
+	client.SetConnStateListener(notifier)
+
+	assert.False(t, client.IsHealthy(), "no registered stream must not be reported healthy")
+	assert.False(t, notifier.isConnected(), "no registered stream must not mark Signal connected")
+
+	// Stop Receive before Close runs (cleanups are LIFO), so the two never overlap.
+	receiveDone := make(chan struct{})
+	go func() {
+		defer close(receiveDone)
+		_ = client.Receive(ctx, func(*sigProto.Message) error { return nil })
+	}()
+	t.Cleanup(func() {
+		cancel()
+		<-receiveDone
+	})
+
+	streamReady := make(chan struct{})
+	go func() {
+		client.WaitStreamConnected(ctx)
+		close(streamReady)
+	}()
+	select {
+	case <-streamReady:
+	case <-time.After(5 * time.Second):
+		t.Fatal("signal stream did not connect within timeout")
+	}
+
+	assert.True(t, client.IsHealthy(), "a registered stream on a usable connection must be healthy")
+	assert.True(t, notifier.isConnected(), "a registered stream must mark Signal connected")
+}
+
+// TestIsHealthyConnectingWithoutStream covers the window while the connection is
+// still being established: a TCP listener that accepts but never speaks HTTP/2
+// keeps the channel in Connecting, and with no registered stream the client
+// cannot receive anything yet, so it must not be reported healthy.
+func TestIsHealthyConnectingWithoutStream(t *testing.T) {
+	lis, err := net.Listen("tcp", "127.0.0.1:0")
+	require.NoError(t, err)
+	var mu sync.Mutex
+	var accepted []net.Conn
+	t.Cleanup(func() {
+		_ = lis.Close()
+		mu.Lock()
+		defer mu.Unlock()
+		for _, c := range accepted {
+			_ = c.Close()
+		}
+	})
+	go func() {
+		for {
+			c, err := lis.Accept()
+			if err != nil {
+				return
+			}
+			mu.Lock()
+			accepted = append(accepted, c)
+			mu.Unlock()
+		}
+	}()
+
+	conn, err := grpc.NewClient(lis.Addr().String(), grpc.WithTransportCredentials(insecure.NewCredentials()))
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = conn.Close() })
+	conn.Connect()
+	require.Eventually(t, func() bool { return conn.GetState() == connectivity.Connecting }, 5*time.Second, 10*time.Millisecond,
+		"channel should be stuck connecting to a listener that never answers")
+
+	client := &GrpcClient{signalConn: conn, status: StreamDisconnected}
+	assert.False(t, client.IsHealthy(), "a connecting channel with no registered stream must not be reported healthy")
+}

shared/signal/client/registration_test.go

@@ -0,0 +1,86 @@
+package client
+
+import (
+	"context"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+	"go.opentelemetry.io/otel"
+	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
+
+	sigProto "github.com/netbirdio/netbird/shared/signal/proto"
+	"github.com/netbirdio/netbird/signal/server"
+)
+
+// stallingSignalServer holds the first ConnectStream open without ever sending
+// the registration header, while the transport keeps answering keepalive pings.
+// Later streams go to the real server. This is what a responder in front of or
+// inside Signal looks like when it accepts a stream and never answers it.
+type stallingSignalServer struct {
+	*server.Server
+	stalled atomic.Bool
+	held    chan struct{}
+}
+
+func (s *stallingSignalServer) ConnectStream(stream sigProto.SignalExchange_ConnectStreamServer) error {
+	if s.stalled.CompareAndSwap(false, true) {
+		close(s.held)
+		<-stream.Context().Done()
+		return stream.Context().Err()
+	}
+	return s.Server.ConnectStream(stream)
+}
+
+// TestReceiveRetriesWhenRegistrationHeaderNeverArrives reproduces a client stuck
+// forever waiting for the registration header of a stream the server side never
+// answers: the receive watchdog only starts after registration, so nothing ended
+// the wait and Signal never came back until the daemon was restarted. The client
+// must give up on the silent stream and register on a new one.
+func TestReceiveRetriesWhenRegistrationHeaderNeverArrives(t *testing.T) {
+	realServer, err := server.NewServer(context.Background(), otel.Meter(""))
+	require.NoError(t, err)
+	stalling := &stallingSignalServer{Server: realServer, held: make(chan struct{})}
+	addr := serveTestSignalExchange(t, stalling)
+
+	key, err := wgtypes.GenerateKey()
+	require.NoError(t, err)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	t.Cleanup(cancel)
+
+	client, err := NewClient(ctx, addr, key, false)
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = client.Close() })
+	client.registrationTimeout = 200 * time.Millisecond
+
+	// Stop Receive before Close runs (cleanups are LIFO), so the two never overlap.
+	receiveDone := make(chan struct{})
+	go func() {
+		defer close(receiveDone)
+		_ = client.Receive(ctx, func(*sigProto.Message) error { return nil })
+	}()
+	t.Cleanup(func() {
+		cancel()
+		<-receiveDone
+	})
+
+	select {
+	case <-stalling.held:
+	case <-time.After(5 * time.Second):
+		t.Fatal("first stream did not reach the stalling server")
+	}
+
+	streamReady := make(chan struct{})
+	go func() {
+		client.WaitStreamConnected(ctx)
+		close(streamReady)
+	}()
+	select {
+	case <-streamReady:
+	case <-time.After(5 * time.Second):
+		t.Fatal("client never registered after the first stream went unanswered")
+	}
+	require.True(t, client.StreamConnected(), "client should report a registered stream")
+}

shared/signal/client/watchdog_test.go

@@ -19,13 +19,20 @@ import (
 func startTestSignalServer(t *testing.T) string {
 	t.Helper()
 
+	srv, err := server.NewServer(context.Background(), otel.Meter(""))
+	require.NoError(t, err)
+
+	return serveTestSignalExchange(t, srv)
+}
+
+func serveTestSignalExchange(t *testing.T, impl sigProto.SignalExchangeServer) string {
+	t.Helper()
+
 	lis, err := net.Listen("tcp", "127.0.0.1:0")
 	require.NoError(t, err)
 
 	s := grpc.NewServer()
-	srv, err := server.NewServer(context.Background(), otel.Meter(""))
-	require.NoError(t, err)
-	sigProto.RegisterSignalExchangeServer(s, srv)
+	sigProto.RegisterSignalExchangeServer(s, impl)
 
 	go func() {
 		_ = s.Serve(lis)