From 002e0b036fb93f3f67ffbdcd08ed67771e9e5065 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Zolt=C3=A1n=20Papp?= <zoltan.pmail@gmail.com>
Date: Tue, 16 Jun 2026 12:29:43 +0200
Subject: [PATCH] [client] let engine context unblock WaitStreamConnected

WaitStreamConnected only watched the signal client's own context, which
derives from the parent engineCtx rather than the engine's run context.
A Start blocked here (signal stream not yet up) could therefore not be
released by Engine.Stop, since Stop only cancels the engine's run
context.

Pass a context into WaitStreamConnected and select on it too, and have
the engine pass e.ctx, so Stop cancelling e.ctx unblocks a parked Start.
Update the Client interface, the mock, and callers accordingly.
---
 client/internal/engine.go           | 3 ++-
 shared/signal/client/client.go      | 2 +-
 shared/signal/client/client_test.go | 6 +++---
 shared/signal/client/grpc.go        | 3 ++-
 shared/signal/client/mock.go        | 2 +-
 5 files changed, 9 insertions(+), 7 deletions(-)

diff --git a/client/internal/engine.go b/client/internal/engine.go
index 11dccc036..68b80f7ab 100644
--- a/client/internal/engine.go
+++ b/client/internal/engine.go
@@ -1776,7 +1776,8 @@ func (e *Engine) receiveSignalEvents() {
 		}
 	}()
 
-	e.signal.WaitStreamConnected()
+	// todo: consider to remove this blocker. I do not see benefit to block the Start operations
+	e.signal.WaitStreamConnected(e.ctx)
 }
 
 func (e *Engine) parseNATExternalIPMappings() []string {
diff --git a/shared/signal/client/client.go b/shared/signal/client/client.go
index 9dc6ccd37..fb77cb90f 100644
--- a/shared/signal/client/client.go
+++ b/shared/signal/client/client.go
@@ -33,7 +33,7 @@ type Client interface {
 	Receive(ctx context.Context, msgHandler func(msg *proto.Message) error) error
 	Ready() bool
 	IsHealthy() bool
-	WaitStreamConnected()
+	WaitStreamConnected(context.Context)
 	SendToStream(msg *proto.EncryptedMessage) error
 	Send(msg *proto.Message) error
 	SetOnReconnectedListener(func())
diff --git a/shared/signal/client/client_test.go b/shared/signal/client/client_test.go
index 1af34e37a..b4ea2bee1 100644
--- a/shared/signal/client/client_test.go
+++ b/shared/signal/client/client_test.go
@@ -65,7 +65,7 @@ var _ = Describe("GrpcClient", func() {
 						return
 					}
 				}()
-				clientA.WaitStreamConnected()
+				clientA.WaitStreamConnected(context.Background())
 
 				// connect PeerB to Signal
 				keyB, _ := wgtypes.GenerateKey()
@@ -91,7 +91,7 @@ var _ = Describe("GrpcClient", func() {
 					}
 				}()
 
-				clientB.WaitStreamConnected()
+				clientB.WaitStreamConnected(context.Background())
 
 				// PeerA initiates ping-pong
 				err := clientA.Send(&sigProto.Message{
@@ -129,7 +129,7 @@ var _ = Describe("GrpcClient", func() {
 						return
 					}
 				}()
-				client.WaitStreamConnected()
+				client.WaitStreamConnected(context.Background())
 				Expect(client).NotTo(BeNil())
 			})
 		})
diff --git a/shared/signal/client/grpc.go b/shared/signal/client/grpc.go
index b245b2296..db23ed75b 100644
--- a/shared/signal/client/grpc.go
+++ b/shared/signal/client/grpc.go
@@ -282,7 +282,7 @@ func (c *GrpcClient) IsHealthy() bool {
 }
 
 // WaitStreamConnected waits until the client is connected to the Signal stream
-func (c *GrpcClient) WaitStreamConnected() {
+func (c *GrpcClient) WaitStreamConnected(ctx context.Context) {
 
 	if c.status == StreamConnected {
 		return
@@ -290,6 +290,7 @@ func (c *GrpcClient) WaitStreamConnected() {
 
 	ch := c.getStreamStatusChan()
 	select {
+	case <-ctx.Done():
 	case <-c.ctx.Done():
 	case <-ch:
 	}
diff --git a/shared/signal/client/mock.go b/shared/signal/client/mock.go
index 95381a5b0..0c8a083c5 100644
--- a/shared/signal/client/mock.go
+++ b/shared/signal/client/mock.go
@@ -55,7 +55,7 @@ func (sm *MockClient) Ready() bool {
 	return sm.ReadyFunc()
 }
 
-func (sm *MockClient) WaitStreamConnected() {
+func (sm *MockClient) WaitStreamConnected(context.Context) {
 	if sm.WaitStreamConnectedFunc == nil {
 		return
 	}