Client: Check for StateMain in disco callback

2 files changed, 31 insertions, 11 deletions

diff --git a/client/client.go b/client/client.go
index 0fd4230..393a6ed 100644
--- a/client/client.go
+++ b/client/client.go
@@ -58,7 +58,21 @@ type Client struct {
 }
 
 func (c *Client) DiscoveryStartup(cb func()) {
-	c.discoMan.Startup(context.Background(), cb)
+	fcb := func() {
+		if cb == nil {
+			return
+		}
+
+		c.mu.Lock()
+		defer c.mu.Unlock()
+
+		if c.FSM.Current != StateMain {
+			return
+		}
+
+		cb()
+	}
+	c.discoMan.Startup(context.Background(), fcb)
 }
 
 // GettingConfig is defined here to satisfy the server.Callbacks interface
@@ -456,11 +470,11 @@ func (c *Client) GetConfig(ck *cookie.Cookie, identifier string, _type srvtypes.
 	var disco *discovery.Discovery
 	disco, release = c.discoMan.Discovery(true)
 	if _type != srvtypes.TypeCustom {
-	    // make sure the servers are fetched fresh
-	    _, _, dserverr := disco.Servers(ctx)
-	    if dserverr != nil {
-		    log.Logger.Warningf("failed to fetch server discovery when getting config: %v", dserverr)
-	    }
+		// make sure the servers are fetched fresh
+		_, _, dserverr := disco.Servers(ctx)
+		if dserverr != nil {
+			log.Logger.Warningf("failed to fetch server discovery when getting config: %v", dserverr)
+		}
 	}
 
 	var srv *server.Server
diff --git a/client/discovery.go b/client/discovery.go
index 2758678..4d56c6d 100644
--- a/client/discovery.go
+++ b/client/discovery.go
@@ -6,9 +6,9 @@ import (
 	"strings"
 	"sync"
 
+	"github.com/eduvpn/eduvpn-common/i18nerr"
 	"github.com/eduvpn/eduvpn-common/internal/discovery"
 	"github.com/eduvpn/eduvpn-common/internal/log"
-	"github.com/eduvpn/eduvpn-common/i18nerr"
 	"github.com/eduvpn/eduvpn-common/types/cookie"
 	discotypes "github.com/eduvpn/eduvpn-common/types/discovery"
 )
@@ -119,8 +119,8 @@ type DiscoManager struct {
 	disco *discovery.Discovery
 
 	cancel context.CancelFunc
-	mu sync.RWMutex
-	wait sync.WaitGroup
+	mu     sync.RWMutex
+	wait   sync.WaitGroup
 }
 
 func (m *DiscoManager) lock(write bool) {
@@ -161,7 +161,6 @@ func (m *DiscoManager) Startup(ctx context.Context, cb func()) {
 	m.cancel = cancel
 	m.wait.Add(1)
 	go func() {
-		defer m.wait.Done()
 		m.lock(false)
 		discoCopy, err := m.disco.Copy()
 		if err != nil {
@@ -175,8 +174,15 @@ func (m *DiscoManager) Startup(ctx context.Context, cb func()) {
 		m.lock(true)
 		m.disco.UpdateServers(discoCopy)
 		m.unlock(true)
+		m.wait.Done()
 
-		if cb != nil {
+		select {
+		case <-ctx.Done():
+			return
+		default:
+			if cb == nil {
+				return
+			}
 			cb()
 		}
 	}()