cli client timeout

DNS record filtering to only include peers that a peer can connect to, reducing unnecessary DNS data in the peer's network map.

- Adds a new `filterZoneRecordsForPeers` function to filter DNS records based on peer connectivity
- Modifies `GetPeerNetworkMap` to use filtered DNS records instead of all records in the custom zone
- Includes comprehensive test coverage for the new filtering functionality

.github/workflows/golang-test-linux.yml

@@ -217,7 +217,7 @@ jobs:
           - arch: "386"
             raceFlag: ""
           - arch: "amd64"
-            raceFlag: ""
+            raceFlag: "-race"
     runs-on: ubuntu-22.04
     steps:
       - name: Install Go

client/android/login.go

@@ -33,6 +33,7 @@ type ErrListener interface {
 // the backend want to show an url for the user
 type URLOpener interface {
 	Open(string)
+	OnLoginSuccess()
 }
 
 // Auth can register or login new client
@@ -181,6 +182,11 @@ func (a *Auth) login(urlOpener URLOpener) error {
 
 	err = a.withBackOff(a.ctx, func() error {
 		err := internal.Login(a.ctx, a.config, "", jwtToken)
+
+		if err == nil {
+			go urlOpener.OnLoginSuccess()
+		}
+
 		if s, ok := gstatus.FromError(err); ok && (s.Code() == codes.InvalidArgument || s.Code() == codes.PermissionDenied) {
 			return nil
 		}

client/cmd/root.go

@@ -231,7 +231,7 @@ func FlagNameToEnvVar(cmdFlag string, prefix string) string {
 
 // DialClientGRPCServer returns client connection to the daemon server.
 func DialClientGRPCServer(ctx context.Context, addr string) (*grpc.ClientConn, error) {
-	ctx, cancel := context.WithTimeout(ctx, time.Second*3)
+	ctx, cancel := context.WithTimeout(ctx, time.Second*10)
 	defer cancel()
 
 	return grpc.DialContext(

client/cmd/up.go

@@ -230,7 +230,7 @@ func runInDaemonMode(ctx context.Context, cmd *cobra.Command, pm *profilemanager
 
 	client := proto.NewDaemonServiceClient(conn)
 
-	status, err := client.Status(ctx, &proto.StatusRequest{})
+	status, err := client.Status(ctx, &proto.StatusRequest{WaitForConnectingShift: true})
 	if err != nil {
 		return fmt.Errorf("unable to get daemon status: %v", err)
 	}

client/internal/peer/worker_ice.go

@@ -218,7 +218,9 @@ func (w *WorkerICE) reCreateAgent(dialerCancel context.CancelFunc, candidates []
 		return nil, err
 	}
 
-	if err := agent.OnSelectedCandidatePairChange(w.onICESelectedCandidatePair); err != nil {
+	if err := agent.OnSelectedCandidatePairChange(func(c1, c2 ice.Candidate) {
+		w.onICESelectedCandidatePair(agent, c1, c2)
+	}); err != nil {
 		return nil, err
 	}
 
@@ -365,26 +367,17 @@ func (w *WorkerICE) onICECandidate(candidate ice.Candidate) {
 	}()
 }
 
-func (w *WorkerICE) onICESelectedCandidatePair(c1 ice.Candidate, c2 ice.Candidate) {
+func (w *WorkerICE) onICESelectedCandidatePair(agent *icemaker.ThreadSafeAgent, c1, c2 ice.Candidate) {
 	w.log.Debugf("selected candidate pair [local <-> remote] -> [%s <-> %s], peer %s", c1.String(), c2.String(),
 		w.config.Key)
 
-	w.muxAgent.Lock()
-
-	pair, err := w.agent.GetSelectedCandidatePair()
-	if err != nil {
-		w.log.Warnf("failed to get selected candidate pair: %s", err)
-		w.muxAgent.Unlock()
+	pairStat, ok := agent.GetSelectedCandidatePairStats()
+	if !ok {
+		w.log.Warnf("failed to get selected candidate pair stats")
 		return
 	}
-	if pair == nil {
-		w.log.Warnf("selected candidate pair is nil, cannot proceed")
-		w.muxAgent.Unlock()
-		return
-	}
-	w.muxAgent.Unlock()
 
-	duration := time.Duration(pair.CurrentRoundTripTime() * float64(time.Second))
+	duration := time.Duration(pairStat.CurrentRoundTripTime * float64(time.Second))
 	if err := w.statusRecorder.UpdateLatency(w.config.Key, duration); err != nil {
 		w.log.Debugf("failed to update latency for peer: %s", err)
 		return

client/proto/daemon.pb.go

@@ -1,7 +1,7 @@
 // Code generated by protoc-gen-go. DO NOT EDIT.
 // versions:
 // 	protoc-gen-go v1.36.6
-// 	protoc        v5.29.3
+// 	protoc        v3.21.9
 // source: daemon.proto
 
 package proto
@@ -791,11 +791,12 @@ func (*UpResponse) Descriptor() ([]byte, []int) {
 }
 
 type StatusRequest struct {
-	state             protoimpl.MessageState `protogen:"open.v1"`
-	GetFullPeerStatus bool                   `protobuf:"varint,1,opt,name=getFullPeerStatus,proto3" json:"getFullPeerStatus,omitempty"`
-	ShouldRunProbes   bool                   `protobuf:"varint,2,opt,name=shouldRunProbes,proto3" json:"shouldRunProbes,omitempty"`
-	unknownFields     protoimpl.UnknownFields
-	sizeCache         protoimpl.SizeCache
+	state                  protoimpl.MessageState `protogen:"open.v1"`
+	GetFullPeerStatus      bool                   `protobuf:"varint,1,opt,name=getFullPeerStatus,proto3" json:"getFullPeerStatus,omitempty"`
+	ShouldRunProbes        bool                   `protobuf:"varint,2,opt,name=shouldRunProbes,proto3" json:"shouldRunProbes,omitempty"`
+	WaitForConnectingShift bool                   `protobuf:"varint,3,opt,name=waitForConnectingShift,proto3" json:"waitForConnectingShift,omitempty"`
+	unknownFields          protoimpl.UnknownFields
+	sizeCache              protoimpl.SizeCache
 }
 
 func (x *StatusRequest) Reset() {
@@ -842,6 +843,13 @@ func (x *StatusRequest) GetShouldRunProbes() bool {
 	return false
 }
 
+func (x *StatusRequest) GetWaitForConnectingShift() bool {
+	if x != nil {
+		return x.WaitForConnectingShift
+	}
+	return false
+}
+
 type StatusResponse struct {
 	state protoimpl.MessageState `protogen:"open.v1"`
 	// status of the server.
@@ -4673,10 +4681,11 @@ const file_daemon_proto_rawDesc = "" +
 	"\f_profileNameB\v\n" +
 	"\t_username\"\f\n" +
 	"\n" +
-	"UpResponse\"g\n" +
+	"UpResponse\"\x9f\x01\n" +
 	"\rStatusRequest\x12,\n" +
 	"\x11getFullPeerStatus\x18\x01 \x01(\bR\x11getFullPeerStatus\x12(\n" +
-	"\x0fshouldRunProbes\x18\x02 \x01(\bR\x0fshouldRunProbes\"\x82\x01\n" +
+	"\x0fshouldRunProbes\x18\x02 \x01(\bR\x0fshouldRunProbes\x126\n" +
+	"\x16waitForConnectingShift\x18\x03 \x01(\bR\x16waitForConnectingShift\"\x82\x01\n" +
 	"\x0eStatusResponse\x12\x16\n" +
 	"\x06status\x18\x01 \x01(\tR\x06status\x122\n" +
 	"\n" +

client/proto/daemon.proto

@@ -186,6 +186,7 @@ message UpResponse {}
 message StatusRequest{
   bool getFullPeerStatus = 1;
   bool shouldRunProbes   = 2;
+  bool waitForConnectingShift = 3;
 }
 
 message StatusResponse{

client/server/server.go

@@ -119,14 +119,12 @@ func (s *Server) Start() error {
 	// if current state contains any error, return it
 	// in all other cases we can continue execution only if status is idle and up command was
 	// not in the progress or already successfully established connection.
-	status, err := state.Status()
+	_, err := state.Status()
 	if err != nil {
 		return err
 	}
 
-	if status != internal.StatusIdle {
-		return nil
-	}
+	state.Set(internal.StatusConnecting)
 
 	ctx, cancel := context.WithCancel(s.rootCtx)
 	s.actCancel = cancel
@@ -961,6 +959,33 @@ func (s *Server) sendLogoutRequestWithConfig(ctx context.Context, config *profil
 	return mgmClient.Logout()
 }
 
+func waitStateShift(ctx context.Context) {
+	timer := time.NewTimer(5 * time.Second)
+	defer timer.Stop()
+	for {
+		select {
+		case <-ctx.Done():
+			log.Warnf("context done while waiting for state shift: %v", ctx.Err())
+			timer.Stop()
+			return
+		case <-timer.C:
+			log.Warnf("state shift timed out")
+			timer.Stop()
+			return
+		default:
+			status, err := internal.CtxGetState(ctx).Status()
+			if err != nil {
+				log.Errorf("failed to get status: %v", err)
+				return
+			}
+			if status != internal.StatusConnecting {
+				log.Infof("state shifting status: %v", status)
+				return
+			}
+		}
+	}
+}
+
 // Status returns the daemon status
 func (s *Server) Status(
 	ctx context.Context,
@@ -973,6 +998,10 @@ func (s *Server) Status(
 	s.mutex.Lock()
 	defer s.mutex.Unlock()
 
+	if msg.WaitForConnectingShift {
+		waitStateShift(s.rootCtx)
+	}
+
 	status, err := internal.CtxGetState(s.rootCtx).Status()
 	if err != nil {
 		return nil, err

client/ui/client_ui.go

@@ -529,7 +529,7 @@ func (s *serviceClient) getSettingsForm() *widget.Form {
 				var req proto.SetConfigRequest
 				req.ProfileName = activeProf.Name
 				req.Username = currUser.Username
-				
+
 				if iMngURL != "" {
 					req.ManagementUrl = iMngURL
 				}
@@ -563,27 +563,28 @@ func (s *serviceClient) getSettingsForm() *widget.Form {
 					return
 				}
 
-				status, err := conn.Status(s.ctx, &proto.StatusRequest{})
-				if err != nil {
-					log.Errorf("get service status: %v", err)
-					dialog.ShowError(fmt.Errorf("Failed to get service status: %v", err), s.wSettings)
-					return
-				}
-				if status.Status == string(internal.StatusConnected) {
-					// run down & up
-					_, err = conn.Down(s.ctx, &proto.DownRequest{})
+				go func() {
+					status, err := conn.Status(s.ctx, &proto.StatusRequest{})
 					if err != nil {
-						log.Errorf("down service: %v", err)
-					}
-
-					_, err = conn.Up(s.ctx, &proto.UpRequest{})
-					if err != nil {
-						log.Errorf("up service: %v", err)
-						dialog.ShowError(fmt.Errorf("Failed to reconnect: %v", err), s.wSettings)
+						log.Errorf("get service status: %v", err)
+						dialog.ShowError(fmt.Errorf("Failed to get service status: %v", err), s.wSettings)
 						return
 					}
-				}
+					if status.Status == string(internal.StatusConnected) {
+						// run down & up
+						_, err = conn.Down(s.ctx, &proto.DownRequest{})
+						if err != nil {
+							log.Errorf("down service: %v", err)
+						}
 
+						_, err = conn.Up(s.ctx, &proto.UpRequest{})
+						if err != nil {
+							log.Errorf("up service: %v", err)
+							dialog.ShowError(fmt.Errorf("Failed to reconnect: %v", err), s.wSettings)
+							return
+						}
+					}
+				}()
 			}
 		},
 		OnCancel: func() {

management/server/policy.go

@@ -167,10 +167,22 @@ func arePolicyChangesAffectPeers(ctx context.Context, transaction store.Store, a
 // validatePolicy validates the policy and its rules.
 func validatePolicy(ctx context.Context, transaction store.Store, accountID string, policy *types.Policy) error {
 	if policy.ID != "" {
-		_, err := transaction.GetPolicyByID(ctx, store.LockingStrengthNone, accountID, policy.ID)
+		existingPolicy, err := transaction.GetPolicyByID(ctx, store.LockingStrengthNone, accountID, policy.ID)
 		if err != nil {
 			return err
 		}
+
+		// TODO: Refactor to support multiple rules per policy
+		existingRuleIDs := make(map[string]bool)
+		for _, rule := range existingPolicy.Rules {
+			existingRuleIDs[rule.ID] = true
+		}
+
+		for _, rule := range policy.Rules {
+			if rule.ID != "" && !existingRuleIDs[rule.ID] {
+				return status.Errorf(status.InvalidArgument, "invalid rule ID: %s", rule.ID)
+			}
+		}
 	} else {
 		policy.ID = xid.New().String()
 		policy.AccountID = accountID

management/server/types/account.go

@@ -302,7 +302,11 @@ func (a *Account) GetPeerNetworkMap(
 		var zones []nbdns.CustomZone
 
 		if peersCustomZone.Domain != "" {
-			zones = append(zones, peersCustomZone)
+			records := filterZoneRecordsForPeers(peer, peersCustomZone, peersToConnect)
+			zones = append(zones, nbdns.CustomZone{
+				Domain:  peersCustomZone.Domain,
+				Records: records,
+			})
 		}
 		dnsUpdate.CustomZones = zones
 		dnsUpdate.NameServerGroups = getPeerNSGroups(a, peerID)
@@ -1651,3 +1655,24 @@ func peerSupportsPortRanges(peerVer string) bool {
 	meetMinVer, err := posture.MeetsMinVersion(firewallRuleMinPortRangesVer, peerVer)
 	return err == nil && meetMinVer
 }
+
+// filterZoneRecordsForPeers filters DNS records to only include peers to connect.
+func filterZoneRecordsForPeers(peer *nbpeer.Peer, customZone nbdns.CustomZone, peersToConnect []*nbpeer.Peer) []nbdns.SimpleRecord {
+	filteredRecords := make([]nbdns.SimpleRecord, 0, len(customZone.Records))
+	peerIPs := make(map[string]struct{})
+
+	// Add peer's own IP to include its own DNS records
+	peerIPs[peer.IP.String()] = struct{}{}
+
+	for _, peerToConnect := range peersToConnect {
+		peerIPs[peerToConnect.IP.String()] = struct{}{}
+	}
+
+	for _, record := range customZone.Records {
+		if _, exists := peerIPs[record.RData]; exists {
+			filteredRecords = append(filteredRecords, record)
+		}
+	}
+
+	return filteredRecords
+}

management/server/types/account_test.go

@@ -2,14 +2,17 @@ package types
 
 import (
 	"context"
+	"fmt"
 	"net"
 	"net/netip"
 	"slices"
 	"testing"
 
+	"github.com/miekg/dns"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
+	nbdns "github.com/netbirdio/netbird/dns"
 	resourceTypes "github.com/netbirdio/netbird/management/server/networks/resources/types"
 	routerTypes "github.com/netbirdio/netbird/management/server/networks/routers/types"
 	networkTypes "github.com/netbirdio/netbird/management/server/networks/types"
@@ -835,3 +838,109 @@ func Test_NetworksNetMapGenShouldExcludeOtherRouters(t *testing.T) {
 	assert.Len(t, networkResourcesRoutes, 1, "expected network resource route don't match")
 	assert.Len(t, sourcePeers, 2, "expected source peers don't match")
 }
+
+func Test_FilterZoneRecordsForPeers(t *testing.T) {
+	tests := []struct {
+		name            string
+		peer            *nbpeer.Peer
+		customZone      nbdns.CustomZone
+		peersToConnect  []*nbpeer.Peer
+		expectedRecords []nbdns.SimpleRecord
+	}{
+		{
+			name: "empty peers to connect",
+			customZone: nbdns.CustomZone{
+				Domain: "netbird.cloud.",
+				Records: []nbdns.SimpleRecord{
+					{Name: "peer1.netbird.cloud", Type: int(dns.TypeA), Class: nbdns.DefaultClass, TTL: 300, RData: "10.0.0.1"},
+					{Name: "router.netbird.cloud", Type: int(dns.TypeA), Class: nbdns.DefaultClass, TTL: 300, RData: "10.0.0.100"},
+				},
+			},
+			peersToConnect: []*nbpeer.Peer{},
+			peer:           &nbpeer.Peer{ID: "router", IP: net.ParseIP("10.0.0.100")},
+			expectedRecords: []nbdns.SimpleRecord{
+				{Name: "router.netbird.cloud", Type: int(dns.TypeA), Class: nbdns.DefaultClass, TTL: 300, RData: "10.0.0.100"},
+			},
+		},
+		{
+			name: "multiple peers multiple records match",
+			customZone: nbdns.CustomZone{
+				Domain: "netbird.cloud.",
+				Records: func() []nbdns.SimpleRecord {
+					var records []nbdns.SimpleRecord
+					for i := 1; i <= 100; i++ {
+						records = append(records, nbdns.SimpleRecord{
+							Name:  fmt.Sprintf("peer%d.netbird.cloud", i),
+							Type:  int(dns.TypeA),
+							Class: nbdns.DefaultClass,
+							TTL:   300,
+							RData: fmt.Sprintf("10.0.%d.%d", i/256, i%256),
+						})
+					}
+					return records
+				}(),
+			},
+			peersToConnect: func() []*nbpeer.Peer {
+				var peers []*nbpeer.Peer
+				for _, i := range []int{1, 5, 10, 25, 50, 75, 100} {
+					peers = append(peers, &nbpeer.Peer{
+						ID: fmt.Sprintf("peer%d", i),
+						IP: net.ParseIP(fmt.Sprintf("10.0.%d.%d", i/256, i%256)),
+					})
+				}
+				return peers
+			}(),
+			peer: &nbpeer.Peer{ID: "router", IP: net.ParseIP("10.0.0.100")},
+			expectedRecords: func() []nbdns.SimpleRecord {
+				var records []nbdns.SimpleRecord
+				for _, i := range []int{1, 5, 10, 25, 50, 75, 100} {
+					records = append(records, nbdns.SimpleRecord{
+						Name:  fmt.Sprintf("peer%d.netbird.cloud", i),
+						Type:  int(dns.TypeA),
+						Class: nbdns.DefaultClass,
+						TTL:   300,
+						RData: fmt.Sprintf("10.0.%d.%d", i/256, i%256),
+					})
+				}
+				return records
+			}(),
+		},
+		{
+			name: "peers with multiple DNS labels",
+			customZone: nbdns.CustomZone{
+				Domain: "netbird.cloud.",
+				Records: []nbdns.SimpleRecord{
+					{Name: "peer1.netbird.cloud", Type: int(dns.TypeA), Class: nbdns.DefaultClass, TTL: 300, RData: "10.0.0.1"},
+					{Name: "peer1-alt.netbird.cloud", Type: int(dns.TypeA), Class: nbdns.DefaultClass, TTL: 300, RData: "10.0.0.1"},
+					{Name: "peer1-backup.netbird.cloud", Type: int(dns.TypeA), Class: nbdns.DefaultClass, TTL: 300, RData: "10.0.0.1"},
+					{Name: "peer2.netbird.cloud", Type: int(dns.TypeA), Class: nbdns.DefaultClass, TTL: 300, RData: "10.0.0.2"},
+					{Name: "peer2-service.netbird.cloud", Type: int(dns.TypeA), Class: nbdns.DefaultClass, TTL: 300, RData: "10.0.0.2"},
+					{Name: "peer3.netbird.cloud", Type: int(dns.TypeA), Class: nbdns.DefaultClass, TTL: 300, RData: "10.0.0.3"},
+					{Name: "peer3-alt.netbird.cloud", Type: int(dns.TypeA), Class: nbdns.DefaultClass, TTL: 300, RData: "10.0.0.3"},
+					{Name: "router.netbird.cloud", Type: int(dns.TypeA), Class: nbdns.DefaultClass, TTL: 300, RData: "10.0.0.100"},
+				},
+			},
+			peersToConnect: []*nbpeer.Peer{
+				{ID: "peer1", IP: net.ParseIP("10.0.0.1"), DNSLabel: "peer1", ExtraDNSLabels: []string{"peer1-alt", "peer1-backup"}},
+				{ID: "peer2", IP: net.ParseIP("10.0.0.2"), DNSLabel: "peer2", ExtraDNSLabels: []string{"peer2-service"}},
+			},
+			peer: &nbpeer.Peer{ID: "router", IP: net.ParseIP("10.0.0.100")},
+			expectedRecords: []nbdns.SimpleRecord{
+				{Name: "peer1.netbird.cloud", Type: int(dns.TypeA), Class: nbdns.DefaultClass, TTL: 300, RData: "10.0.0.1"},
+				{Name: "peer1-alt.netbird.cloud", Type: int(dns.TypeA), Class: nbdns.DefaultClass, TTL: 300, RData: "10.0.0.1"},
+				{Name: "peer1-backup.netbird.cloud", Type: int(dns.TypeA), Class: nbdns.DefaultClass, TTL: 300, RData: "10.0.0.1"},
+				{Name: "peer2.netbird.cloud", Type: int(dns.TypeA), Class: nbdns.DefaultClass, TTL: 300, RData: "10.0.0.2"},
+				{Name: "peer2-service.netbird.cloud", Type: int(dns.TypeA), Class: nbdns.DefaultClass, TTL: 300, RData: "10.0.0.2"},
+				{Name: "router.netbird.cloud", Type: int(dns.TypeA), Class: nbdns.DefaultClass, TTL: 300, RData: "10.0.0.100"},
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			result := filterZoneRecordsForPeers(tt.peer, tt.customZone, tt.peersToConnect)
+			assert.Equal(t, len(tt.expectedRecords), len(result))
+			assert.ElementsMatch(t, tt.expectedRecords, result)
+		})
+	}
+}

shared/relay/client/manager.go

@@ -78,9 +78,10 @@ func NewManager(ctx context.Context, serverURLs []string, peerID string, mtu uin
 		tokenStore: tokenStore,
 		mtu:        mtu,
 		serverPicker: &ServerPicker{
-			TokenStore: tokenStore,
-			PeerID:     peerID,
-			MTU:        mtu,
+			TokenStore:        tokenStore,
+			PeerID:            peerID,
+			MTU:               mtu,
+			ConnectionTimeout: defaultConnectionTimeout,
 		},
 		relayClients:            make(map[string]*RelayTrack),
 		onDisconnectedListeners: make(map[string]*list.List),

shared/relay/client/picker.go

@@ -13,11 +13,8 @@ import (
 )
 
 const (
-	maxConcurrentServers = 7
-)
-
-var (
-	connectionTimeout = 30 * time.Second
+	maxConcurrentServers     = 7
+	defaultConnectionTimeout = 30 * time.Second
 )
 
 type connResult struct {
@@ -27,14 +24,15 @@ type connResult struct {
 }
 
 type ServerPicker struct {
-	TokenStore *auth.TokenStore
-	ServerURLs atomic.Value
-	PeerID     string
-	MTU        uint16
+	TokenStore        *auth.TokenStore
+	ServerURLs        atomic.Value
+	PeerID            string
+	MTU               uint16
+	ConnectionTimeout time.Duration
 }
 
 func (sp *ServerPicker) PickServer(parentCtx context.Context) (*Client, error) {
-	ctx, cancel := context.WithTimeout(parentCtx, connectionTimeout)
+	ctx, cancel := context.WithTimeout(parentCtx, sp.ConnectionTimeout)
 	defer cancel()
 
 	totalServers := len(sp.ServerURLs.Load().([]string))

shared/relay/client/picker_test.go

@@ -8,15 +8,15 @@ import (
 )
 
 func TestServerPicker_UnavailableServers(t *testing.T) {
-	connectionTimeout = 5 * time.Second
-
+	timeout := 5 * time.Second
 	sp := ServerPicker{
-		TokenStore: nil,
-		PeerID:     "test",
+		TokenStore:        nil,
+		PeerID:            "test",
+		ConnectionTimeout: timeout,
 	}
 	sp.ServerURLs.Store([]string{"rel://dummy1", "rel://dummy2"})
 
-	ctx, cancel := context.WithTimeout(context.Background(), connectionTimeout+1)
+	ctx, cancel := context.WithTimeout(context.Background(), timeout+1)
 	defer cancel()
 
 	go func() {

shared/relay/healthcheck/env.go

@@ -0,0 +1,24 @@
+package healthcheck
+
+import (
+	"os"
+	"strconv"
+
+	log "github.com/sirupsen/logrus"
+)
+
+const (
+	defaultAttemptThresholdEnv = "NB_RELAY_HC_ATTEMPT_THRESHOLD"
+)
+
+func getAttemptThresholdFromEnv() int {
+	if attemptThreshold := os.Getenv(defaultAttemptThresholdEnv); attemptThreshold != "" {
+		threshold, err := strconv.ParseInt(attemptThreshold, 10, 64)
+		if err != nil {
+			log.Errorf("Failed to parse attempt threshold from environment variable \"%s\" should be an integer. Using default value", attemptThreshold)
+			return defaultAttemptThreshold
+		}
+		return int(threshold)
+	}
+	return defaultAttemptThreshold
+}

shared/relay/healthcheck/env_test.go

@@ -0,0 +1,36 @@
+package healthcheck
+
+import (
+	"os"
+	"testing"
+)
+
+//nolint:tenv
+func TestGetAttemptThresholdFromEnv(t *testing.T) {
+	tests := []struct {
+		name     string
+		envValue string
+		expected int
+	}{
+		{"Default attempt threshold when env is not set", "", defaultAttemptThreshold},
+		{"Custom attempt threshold when env is set to a valid integer", "3", 3},
+		{"Default attempt threshold when env is set to an invalid value", "invalid", defaultAttemptThreshold},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if tt.envValue == "" {
+				os.Unsetenv(defaultAttemptThresholdEnv)
+			} else {
+				os.Setenv(defaultAttemptThresholdEnv, tt.envValue)
+			}
+
+			result := getAttemptThresholdFromEnv()
+			if result != tt.expected {
+				t.Fatalf("Expected %d, got %d", tt.expected, result)
+			}
+
+			os.Unsetenv(defaultAttemptThresholdEnv)
+		})
+	}
+}

shared/relay/healthcheck/receiver.go

@@ -7,10 +7,15 @@ import (
 	log "github.com/sirupsen/logrus"
 )
 
-var (
-	heartbeatTimeout = healthCheckInterval + 10*time.Second
+const (
+	defaultHeartbeatTimeout = defaultHealthCheckInterval + 10*time.Second
 )
 
+type ReceiverOptions struct {
+	HeartbeatTimeout time.Duration
+	AttemptThreshold int
+}
+
 // Receiver is a healthcheck receiver
 // It will listen for heartbeat and check if the heartbeat is not received in a certain time
 // If the heartbeat is not received in a certain time, it will send a timeout signal and stop to work
@@ -27,6 +32,23 @@ type Receiver struct {
 
 // NewReceiver creates a new healthcheck receiver and start the timer in the background
 func NewReceiver(log *log.Entry) *Receiver {
+	opts := ReceiverOptions{
+		HeartbeatTimeout: defaultHeartbeatTimeout,
+		AttemptThreshold: getAttemptThresholdFromEnv(),
+	}
+	return NewReceiverWithOpts(log, opts)
+}
+
+func NewReceiverWithOpts(log *log.Entry, opts ReceiverOptions) *Receiver {
+	heartbeatTimeout := opts.HeartbeatTimeout
+	if heartbeatTimeout <= 0 {
+		heartbeatTimeout = defaultHeartbeatTimeout
+	}
+	attemptThreshold := opts.AttemptThreshold
+	if attemptThreshold <= 0 {
+		attemptThreshold = defaultAttemptThreshold
+	}
+
 	ctx, ctxCancel := context.WithCancel(context.Background())
 
 	r := &Receiver{
@@ -35,10 +57,10 @@ func NewReceiver(log *log.Entry) *Receiver {
 		ctx:              ctx,
 		ctxCancel:        ctxCancel,
 		heartbeat:        make(chan struct{}, 1),
-		attemptThreshold: getAttemptThresholdFromEnv(),
+		attemptThreshold: attemptThreshold,
 	}
 
-	go r.waitForHealthcheck()
+	go r.waitForHealthcheck(heartbeatTimeout)
 	return r
 }
 
@@ -55,7 +77,7 @@ func (r *Receiver) Stop() {
 	r.ctxCancel()
 }
 
-func (r *Receiver) waitForHealthcheck() {
+func (r *Receiver) waitForHealthcheck(heartbeatTimeout time.Duration) {
 	ticker := time.NewTicker(heartbeatTimeout)
 	defer ticker.Stop()
 	defer r.ctxCancel()

shared/relay/healthcheck/receiver_test.go

@@ -2,31 +2,18 @@ package healthcheck
 
 import (
 	"context"
-	"fmt"
-	"os"
-	"sync"
 	"testing"
 	"time"
 
 	log "github.com/sirupsen/logrus"
 )
 
-// Mutex to protect global variable access in tests
-var testMutex sync.Mutex
-
 func TestNewReceiver(t *testing.T) {
-	testMutex.Lock()
-	originalTimeout := heartbeatTimeout
-	heartbeatTimeout = 5 * time.Second
-	testMutex.Unlock()
 
-	defer func() {
-		testMutex.Lock()
-		heartbeatTimeout = originalTimeout
-		testMutex.Unlock()
-	}()
-
-	r := NewReceiver(log.WithContext(context.Background()))
+	opts := ReceiverOptions{
+		HeartbeatTimeout: 5 * time.Second,
+	}
+	r := NewReceiverWithOpts(log.WithContext(context.Background()), opts)
 	defer r.Stop()
 
 	select {
@@ -38,18 +25,10 @@ func TestNewReceiver(t *testing.T) {
 }
 
 func TestNewReceiverNotReceive(t *testing.T) {
-	testMutex.Lock()
-	originalTimeout := heartbeatTimeout
-	heartbeatTimeout = 1 * time.Second
-	testMutex.Unlock()
-
-	defer func() {
-		testMutex.Lock()
-		heartbeatTimeout = originalTimeout
-		testMutex.Unlock()
-	}()
-
-	r := NewReceiver(log.WithContext(context.Background()))
+	opts := ReceiverOptions{
+		HeartbeatTimeout: 1 * time.Second,
+	}
+	r := NewReceiverWithOpts(log.WithContext(context.Background()), opts)
 	defer r.Stop()
 
 	select {
@@ -61,18 +40,10 @@ func TestNewReceiverNotReceive(t *testing.T) {
 }
 
 func TestNewReceiverAck(t *testing.T) {
-	testMutex.Lock()
-	originalTimeout := heartbeatTimeout
-	heartbeatTimeout = 2 * time.Second
-	testMutex.Unlock()
-
-	defer func() {
-		testMutex.Lock()
-		heartbeatTimeout = originalTimeout
-		testMutex.Unlock()
-	}()
-
-	r := NewReceiver(log.WithContext(context.Background()))
+	opts := ReceiverOptions{
+		HeartbeatTimeout: 2 * time.Second,
+	}
+	r := NewReceiverWithOpts(log.WithContext(context.Background()), opts)
 	defer r.Stop()
 
 	r.Heartbeat()
@@ -97,30 +68,19 @@ func TestReceiverHealthCheckAttemptThreshold(t *testing.T) {
 
 	for _, tc := range testsCases {
 		t.Run(tc.name, func(t *testing.T) {
-			testMutex.Lock()
-			originalInterval := healthCheckInterval
-			originalTimeout := heartbeatTimeout
-			healthCheckInterval = 1 * time.Second
-			heartbeatTimeout = healthCheckInterval + 500*time.Millisecond
-			testMutex.Unlock()
+			healthCheckInterval := 1 * time.Second
 
-			defer func() {
-				testMutex.Lock()
-				healthCheckInterval = originalInterval
-				heartbeatTimeout = originalTimeout
-				testMutex.Unlock()
-			}()
-			//nolint:tenv
-			os.Setenv(defaultAttemptThresholdEnv, fmt.Sprintf("%d", tc.threshold))
-			defer os.Unsetenv(defaultAttemptThresholdEnv)
+			opts := ReceiverOptions{
+				HeartbeatTimeout: healthCheckInterval + 500*time.Millisecond,
+				AttemptThreshold: tc.threshold,
+			}
 
-			receiver := NewReceiver(log.WithField("test_name", tc.name))
+			receiver := NewReceiverWithOpts(log.WithField("test_name", tc.name), opts)
 
-			testTimeout := heartbeatTimeout*time.Duration(tc.threshold) + healthCheckInterval
+			testTimeout := opts.HeartbeatTimeout*time.Duration(tc.threshold) + healthCheckInterval
 
 			if tc.resetCounterOnce {
 				receiver.Heartbeat()
-				t.Logf("reset counter once")
 			}
 
 			select {
@@ -134,7 +94,6 @@ func TestReceiverHealthCheckAttemptThreshold(t *testing.T) {
 				}
 				t.Fatalf("should have timed out before %s", testTimeout)
 			}
-
 		})
 	}
 }

shared/relay/healthcheck/sender.go

@@ -2,52 +2,76 @@ package healthcheck
 
 import (
 	"context"
-	"os"
-	"strconv"
 	"time"
 
 	log "github.com/sirupsen/logrus"
 )
 
 const (
-	defaultAttemptThreshold    = 1
-	defaultAttemptThresholdEnv = "NB_RELAY_HC_ATTEMPT_THRESHOLD"
+	defaultAttemptThreshold = 1
+
+	defaultHealthCheckInterval = 25 * time.Second
+	defaultHealthCheckTimeout  = 20 * time.Second
 )
 
-var (
-	healthCheckInterval = 25 * time.Second
-	healthCheckTimeout  = 20 * time.Second
-)
+type SenderOptions struct {
+	HealthCheckInterval time.Duration
+	HealthCheckTimeout  time.Duration
+	AttemptThreshold    int
+}
 
 // Sender is a healthcheck sender
 // It will send healthcheck signal to the receiver
 // If the receiver does not receive the signal in a certain time, it will send a timeout signal and stop to work
 // It will also stop if the context is canceled
 type Sender struct {
-	log *log.Entry
 	// HealthCheck is a channel to send health check signal to the peer
 	HealthCheck chan struct{}
 	// Timeout is a channel to the health check signal is not received in a certain time
 	Timeout chan struct{}
 
+	log                 *log.Entry
+	healthCheckInterval time.Duration
+	timeout             time.Duration
+
 	ack              chan struct{}
 	alive            bool
 	attemptThreshold int
 }
 
-// NewSender creates a new healthcheck sender
-func NewSender(log *log.Entry) *Sender {
+func NewSenderWithOpts(log *log.Entry, opts SenderOptions) *Sender {
+	if opts.HealthCheckInterval <= 0 {
+		opts.HealthCheckInterval = defaultHealthCheckInterval
+	}
+	if opts.HealthCheckTimeout <= 0 {
+		opts.HealthCheckTimeout = defaultHealthCheckTimeout
+	}
+	if opts.AttemptThreshold <= 0 {
+		opts.AttemptThreshold = defaultAttemptThreshold
+	}
 	hc := &Sender{
-		log:              log,
-		HealthCheck:      make(chan struct{}, 1),
-		Timeout:          make(chan struct{}, 1),
-		ack:              make(chan struct{}, 1),
-		attemptThreshold: getAttemptThresholdFromEnv(),
+		HealthCheck:         make(chan struct{}, 1),
+		Timeout:             make(chan struct{}, 1),
+		log:                 log,
+		healthCheckInterval: opts.HealthCheckInterval,
+		timeout:             opts.HealthCheckInterval + opts.HealthCheckTimeout,
+		ack:                 make(chan struct{}, 1),
+		attemptThreshold:    opts.AttemptThreshold,
 	}
 
 	return hc
 }
 
+// NewSender creates a new healthcheck sender
+func NewSender(log *log.Entry) *Sender {
+	opts := SenderOptions{
+		HealthCheckInterval: defaultHealthCheckInterval,
+		HealthCheckTimeout:  defaultHealthCheckTimeout,
+		AttemptThreshold:    getAttemptThresholdFromEnv(),
+	}
+	return NewSenderWithOpts(log, opts)
+}
+
 // OnHCResponse sends an acknowledgment signal to the sender
 func (hc *Sender) OnHCResponse() {
 	select {
@@ -57,10 +81,10 @@ func (hc *Sender) OnHCResponse() {
 }
 
 func (hc *Sender) StartHealthCheck(ctx context.Context) {
-	ticker := time.NewTicker(healthCheckInterval)
+	ticker := time.NewTicker(hc.healthCheckInterval)
 	defer ticker.Stop()
 
-	timeoutTicker := time.NewTicker(hc.getTimeoutTime())
+	timeoutTicker := time.NewTicker(hc.timeout)
 	defer timeoutTicker.Stop()
 
 	defer close(hc.HealthCheck)
@@ -92,19 +116,3 @@ func (hc *Sender) StartHealthCheck(ctx context.Context) {
 		}
 	}
 }
-
-func (hc *Sender) getTimeoutTime() time.Duration {
-	return healthCheckInterval + healthCheckTimeout
-}
-
-func getAttemptThresholdFromEnv() int {
-	if attemptThreshold := os.Getenv(defaultAttemptThresholdEnv); attemptThreshold != "" {
-		threshold, err := strconv.ParseInt(attemptThreshold, 10, 64)
-		if err != nil {
-			log.Errorf("Failed to parse attempt threshold from environment variable \"%s\" should be an integer. Using default value", attemptThreshold)
-			return defaultAttemptThreshold
-		}
-		return int(threshold)
-	}
-	return defaultAttemptThreshold
-}

shared/relay/healthcheck/sender_test.go

@@ -2,26 +2,23 @@ package healthcheck
 
 import (
 	"context"
-	"fmt"
-	"os"
 	"testing"
 	"time"
 
 	log "github.com/sirupsen/logrus"
 )
 
-func TestMain(m *testing.M) {
-	// override the health check interval to speed up the test
-	healthCheckInterval = 2 * time.Second
-	healthCheckTimeout = 100 * time.Millisecond
-	code := m.Run()
-	os.Exit(code)
-}
+var (
+	testOpts = SenderOptions{
+		HealthCheckInterval: 2 * time.Second,
+		HealthCheckTimeout:  100 * time.Millisecond,
+	}
+)
 
 func TestNewHealthPeriod(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
-	hc := NewSender(log.WithContext(ctx))
+	hc := NewSenderWithOpts(log.WithContext(ctx), testOpts)
 	go hc.StartHealthCheck(ctx)
 
 	iterations := 0
@@ -32,7 +29,7 @@ func TestNewHealthPeriod(t *testing.T) {
 			hc.OnHCResponse()
 		case <-hc.Timeout:
 			t.Fatalf("health check is timed out")
-		case <-time.After(healthCheckInterval + 100*time.Millisecond):
+		case <-time.After(testOpts.HealthCheckInterval + 100*time.Millisecond):
 			t.Fatalf("health check not received")
 		}
 	}
@@ -41,19 +38,19 @@ func TestNewHealthPeriod(t *testing.T) {
 func TestNewHealthFailed(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
-	hc := NewSender(log.WithContext(ctx))
+	hc := NewSenderWithOpts(log.WithContext(ctx), testOpts)
 	go hc.StartHealthCheck(ctx)
 
 	select {
 	case <-hc.Timeout:
-	case <-time.After(healthCheckInterval + healthCheckTimeout + 100*time.Millisecond):
+	case <-time.After(testOpts.HealthCheckInterval + testOpts.HealthCheckTimeout + 100*time.Millisecond):
 		t.Fatalf("health check is not timed out")
 	}
 }
 
 func TestNewHealthcheckStop(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())
-	hc := NewSender(log.WithContext(ctx))
+	hc := NewSenderWithOpts(log.WithContext(ctx), testOpts)
 	go hc.StartHealthCheck(ctx)
 
 	time.Sleep(100 * time.Millisecond)
@@ -78,7 +75,7 @@ func TestNewHealthcheckStop(t *testing.T) {
 func TestTimeoutReset(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
-	hc := NewSender(log.WithContext(ctx))
+	hc := NewSenderWithOpts(log.WithContext(ctx), testOpts)
 	go hc.StartHealthCheck(ctx)
 
 	iterations := 0
@@ -89,7 +86,7 @@ func TestTimeoutReset(t *testing.T) {
 			hc.OnHCResponse()
 		case <-hc.Timeout:
 			t.Fatalf("health check is timed out")
-		case <-time.After(healthCheckInterval + 100*time.Millisecond):
+		case <-time.After(testOpts.HealthCheckInterval + 100*time.Millisecond):
 			t.Fatalf("health check not received")
 		}
 	}
@@ -118,19 +115,16 @@ func TestSenderHealthCheckAttemptThreshold(t *testing.T) {
 
 	for _, tc := range testsCases {
 		t.Run(tc.name, func(t *testing.T) {
-			originalInterval := healthCheckInterval
-			originalTimeout := healthCheckTimeout
-			healthCheckInterval = 1 * time.Second
-			healthCheckTimeout = 500 * time.Millisecond
-
-			//nolint:tenv
-			os.Setenv(defaultAttemptThresholdEnv, fmt.Sprintf("%d", tc.threshold))
-			defer os.Unsetenv(defaultAttemptThresholdEnv)
+			opts := SenderOptions{
+				HealthCheckInterval: 1 * time.Second,
+				HealthCheckTimeout:  500 * time.Millisecond,
+				AttemptThreshold:    tc.threshold,
+			}
 
 			ctx, cancel := context.WithCancel(context.Background())
 			defer cancel()
 
-			sender := NewSender(log.WithField("test_name", tc.name))
+			sender := NewSenderWithOpts(log.WithField("test_name", tc.name), opts)
 			senderExit := make(chan struct{})
 			go func() {
 				sender.StartHealthCheck(ctx)
@@ -155,7 +149,7 @@ func TestSenderHealthCheckAttemptThreshold(t *testing.T) {
 				}
 			}()
 
-			testTimeout := sender.getTimeoutTime()*time.Duration(tc.threshold) + healthCheckInterval
+			testTimeout := (opts.HealthCheckInterval+opts.HealthCheckTimeout)*time.Duration(tc.threshold) + opts.HealthCheckInterval
 
 			select {
 			case <-sender.Timeout:
@@ -175,39 +169,7 @@ func TestSenderHealthCheckAttemptThreshold(t *testing.T) {
 			case <-time.After(2 * time.Second):
 				t.Fatalf("sender did not exit in time")
 			}
-			healthCheckInterval = originalInterval
-			healthCheckTimeout = originalTimeout
 		})
 	}
 
 }
-
-//nolint:tenv
-func TestGetAttemptThresholdFromEnv(t *testing.T) {
-	tests := []struct {
-		name     string
-		envValue string
-		expected int
-	}{
-		{"Default attempt threshold when env is not set", "", defaultAttemptThreshold},
-		{"Custom attempt threshold when env is set to a valid integer", "3", 3},
-		{"Default attempt threshold when env is set to an invalid value", "invalid", defaultAttemptThreshold},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			if tt.envValue == "" {
-				os.Unsetenv(defaultAttemptThresholdEnv)
-			} else {
-				os.Setenv(defaultAttemptThresholdEnv, tt.envValue)
-			}
-
-			result := getAttemptThresholdFromEnv()
-			if result != tt.expected {
-				t.Fatalf("Expected %d, got %d", tt.expected, result)
-			}
-
-			os.Unsetenv(defaultAttemptThresholdEnv)
-		})
-	}
-}