From 6258eca170eb18a577351131e6b341cb52dc7020 Mon Sep 17 00:00:00 2001
From: jwijenbergh <jeroenwijenbergh@protonmail.com>
Date: Mon, 15 Jul 2024 09:46:01 +0200
Subject: Client: Check for StateMain in disco callback

---
 client/discovery.go | 16 +++++++++++-----
 1 file changed, 11 insertions(+), 5 deletions(-)

(limited to 'client/discovery.go')

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()
 		}
 	}()
-- 
cgit v1.2.3