diff --git a/client/internal/conn_mgr.go b/client/internal/conn_mgr.go index 6fb56a93c..70e5e281c 100644 --- a/client/internal/conn_mgr.go +++ b/client/internal/conn_mgr.go @@ -71,7 +71,7 @@ func (e *ConnMgr) AddPeerConn(peerKey string, conn *peer.Conn) (exists bool) { return true } - if e.lazyConnMgr == nil { + if !e.isStarted() { conn.Open() return } @@ -101,7 +101,7 @@ func (e *ConnMgr) OnSignalMsg(peerKey string) (*peer.Conn, bool) { return nil, false } - if e.lazyConnMgr == nil { + if !e.isStarted() { return conn, true } @@ -119,7 +119,7 @@ func (e *ConnMgr) RemovePeerConn(peerKey string) { } defer conn.Close() - if e.lazyConnMgr == nil { + if !e.isStarted() { return } @@ -128,9 +128,10 @@ func (e *ConnMgr) RemovePeerConn(peerKey string) { } func (e *ConnMgr) Close() { - if e.lazyConnMgr == nil { + if !e.isStarted() { return } + e.ctxCancel() e.lazyConnMgr.Close() e.wg.Wait() @@ -147,3 +148,7 @@ func (e *ConnMgr) receiveLazyConnEvents(ctx context.Context) { e.peerStore.PeerConnOpen(peerID) } } + +func (e *ConnMgr) isStarted() bool { + return e.lazyConnMgr != nil && e.ctxCancel != nil +} diff --git a/client/internal/engine_test.go b/client/internal/engine_test.go index d55ae5383..d8a5cfddb 100644 --- a/client/internal/engine_test.go +++ b/client/internal/engine_test.go @@ -332,7 +332,6 @@ func TestEngine_SSH(t *testing.T) { } func TestEngine_UpdateNetworkMap(t *testing.T) { - t.Setenv(envDisableLazyConn, "true") // test setup key, err := wgtypes.GeneratePrivateKey() if err != nil { @@ -766,6 +765,7 @@ func TestEngine_UpdateNetworkMapWithRoutes(t *testing.T) { engine.routeManager = mockRouteManager engine.dnsServer = &dns.MockServer{} + engine.connMgr = NewConnMgr(engine.peerStore, engine.wgInterface) defer func() { exitErr := engine.Stop() @@ -962,6 +962,7 @@ func TestEngine_UpdateNetworkMapWithDNSUpdate(t *testing.T) { } engine.dnsServer = mockDNSServer + engine.connMgr = NewConnMgr(engine.peerStore, engine.wgInterface) defer func() { exitErr := engine.Stop()