Merge pull request #1237 from ersonp/fix/stun-datarace

Fix stun client datarace
skycoin · May 30, 2022 · 0c1aabe · 0c1aabe
2 parents 89b22ba + a7f0478
commit 0c1aabe
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 12 deletions.
diff --git a/pkg/visor/api.go b/pkg/visor/api.go
@@ -122,7 +122,8 @@ func (v *Visor) Overview() (*Overview, error) {
 			newTransportSummary(v.tpM, tp, true, v.router.SetupIsTrusted(tp.Remote())))
 		return true
 	})
-	if v.stunClient != nil {
+
+	if v.isStunReady() {
 		switch v.stunClient.NATType {
 		case stun.NATNone, stun.NATFull, stun.NATRestricted, stun.NATPortRestricted:
 			publicIP = v.stunClient.PublicIP.IP()

diff --git a/pkg/visor/init.go b/pkg/visor/init.go
@@ -279,12 +279,12 @@ func initDiscovery(ctx context.Context, v *Visor, log *logging.Logger) error {
 }
 
 func initStunClient(ctx context.Context, v *Visor, log *logging.Logger) error {
-	defer v.wgStunClient.Done()
 
 	sc := network.GetStunDetails(v.conf.StunServers, log)
 	v.initLock.Lock()
 	v.stunClient = sc
 	v.initLock.Unlock()
+	v.stunReadyOnce.Do(func() { close(v.stunReady) })
 	return nil
 }
 
@@ -517,7 +517,7 @@ func getRouteSetupHooks(ctx context.Context, v *Visor, log *logging.Logger) []ro
 				ntype := network.Type(trans)
 
 				// Wait until stun client is ready
-				v.wgStunClient.Wait()
+				<-v.stunReady
 
 				// skip if SUDPH is under symmetric NAT / under UDP firewall.
 				if ntype == network.SUDPH && (v.stunClient.NATType == stun.NATSymmetric ||
@@ -1157,7 +1157,7 @@ func getPublicIP(v *Visor, service string) (pIP string, err error) {
 	}
 
 	// Wait until stun client is ready
-	v.wgStunClient.Wait()
+	<-v.stunReady
 
 	pIP = v.stunClient.PublicIP.IP()
 	return pIP, nil

diff --git a/pkg/visor/visor.go b/pkg/visor/visor.go
@@ -70,12 +70,14 @@ type Visor struct {
 	trackers      *dmsgtracker.Manager
 	trackersReady bool
 
-	stunClient   *network.StunDetails
-	wgStunClient *sync.WaitGroup
-	tpM          *transport.Manager
-	arClient     addrresolver.APIClient
-	router       router.Router
-	rfClient     rfclient.Client
+	stunClient    *network.StunDetails
+	stunReady     chan struct{}
+	stunReadyOnce sync.Once
+
+	tpM      *transport.Manager
+	arClient addrresolver.APIClient
+	router   router.Router
+	rfClient rfclient.Client
 
 	procM       appserver.ProcManager // proc manager
 	appL        *launcher.Launcher    // app launcher
@@ -118,10 +120,9 @@ func NewVisor(ctx context.Context, conf *visorconfig.V1, restartCtx *restart.Con
 		restartCtx:        restartCtx,
 		initLock:          new(sync.RWMutex),
 		isServicesHealthy: newInternalHealthInfo(),
-		wgStunClient:      new(sync.WaitGroup),
+		stunReady:         make(chan struct{}),
 		trackersReady:     false,
 	}
-	v.wgStunClient.Add(1)
 	v.isServicesHealthy.init()
 
 	if logLvl, err := logging.LevelFromString(conf.LogLevel); err != nil {
@@ -196,6 +197,15 @@ func (v *Visor) processRuntimeErrs() bool {
 	}
 }
 
+func (v *Visor) isStunReady() bool {
+	select {
+	case <-v.stunReady:
+		return true
+	default:
+		return false
+	}
+}
+
 // Close safely stops spawned Apps and Visor.
 func (v *Visor) Close() error {
 	if v == nil {