Modify client-side behavior

client/internal/connect.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"github.com/netbirdio/netbird/client/internal/updatemanager"
 	"net"
 	"net/netip"
 	"runtime"
@@ -272,6 +273,18 @@ func (c *ConnectClient) run(mobileDependency MobileDependency, runningChan chan
 
 		c.engineMutex.Lock()
 		c.engine = NewEngine(engineCtx, cancel, signalClient, mgmClient, relayManager, engineConfig, mobileDependency, c.statusRecorder, checks)
+		if loginResp.PeerConfig != nil && loginResp.PeerConfig.AutoUpdate != nil {
+			if c.engine.updateManager == nil && loginResp.PeerConfig.AutoUpdate.Version != "disabled" {
+				c.engine.updateManager = updatemanager.NewUpdateManager(c.statusRecorder, c.engine.stateManager)
+				c.engine.updateManager.StartWithTimeout(engineCtx, time.Minute)
+			} else if c.engine.updateManager != nil && loginResp.PeerConfig.AutoUpdate.Version == "disabled" {
+				c.engine.updateManager.Stop()
+				c.engine.updateManager = nil
+			}
+			if c.engine.updateManager != nil {
+				c.engine.updateManager.SetVersion(loginResp.PeerConfig.AutoUpdate.Version)
+			}
+		}
 		c.engine.SetSyncResponsePersistence(c.persistSyncResponse)
 		c.engineMutex.Unlock()
 

client/internal/engine.go

@@ -721,16 +721,19 @@ func (e *Engine) PopulateNetbirdConfig(netbirdConfig *mgmProto.NetbirdConfig, mg
 	return nil
 }
 
-func (e *Engine) handleAutoUpdateVersion(autoUpdateVersion string) {
-	if e.updateManager == nil && autoUpdateVersion != disableAutoUpdate {
-		e.updateManager = updatemanager.NewUpdateManager(e.statusRecorder)
+func (e *Engine) handleAutoUpdateVersion(autoUpdateSettings *mgmProto.AutoUpdateSettings) {
+	if autoUpdateSettings == nil {
+		return
+	}
+	if e.updateManager == nil && autoUpdateSettings.Version != disableAutoUpdate && autoUpdateSettings.AlwaysUpdate {
+		e.updateManager = updatemanager.NewUpdateManager(e.statusRecorder, e.stateManager)
 		e.updateManager.Start(e.ctx)
-	} else if e.updateManager != nil && autoUpdateVersion == disableAutoUpdate {
+	} else if e.updateManager != nil && autoUpdateSettings.Version == disableAutoUpdate {
 		e.updateManager.Stop()
 		e.updateManager = nil
 	}
-	if e.updateManager != nil {
-		e.updateManager.SetVersion(autoUpdateVersion)
+	if e.updateManager != nil && autoUpdateSettings.AlwaysUpdate {
+		e.updateManager.SetVersion(autoUpdateSettings.Version)
 	}
 }
 
@@ -739,7 +742,7 @@ func (e *Engine) handleSync(update *mgmProto.SyncResponse) error {
 	defer e.syncMsgMux.Unlock()
 
 	if update.NetworkMap != nil && update.NetworkMap.PeerConfig != nil {
-		e.handleAutoUpdateVersion(update.NetworkMap.PeerConfig.AutoUpdateVersion)
+		e.handleAutoUpdateVersion(update.NetworkMap.PeerConfig.AutoUpdate)
 	}
 	if update.GetNetbirdConfig() != nil {
 		wCfg := update.GetNetbirdConfig()

client/internal/updatemanager/manager.go

@@ -3,6 +3,7 @@
 import (
 	"context"
 	"fmt"
+	"github.com/netbirdio/netbird/client/internal/statemanager"
 	"io"
 	"net/http"
 	"os"
@@ -32,6 +33,15 @@
 	StartFetcher()
 }
 
+type UpdateState struct {
+	PreUpdateVersion string
+	TargetVersion    string
+}
+
+func (u UpdateState) Name() string {
+	return "autoUpdate"
+}
+
 type UpdateManager struct {
 	lastTrigger    time.Time
 	statusRecorder *peer.Status
@@ -40,6 +50,7 @@
 	wg             sync.WaitGroup
 	currentVersion string
 	updateFunc     func(ctx context.Context, targetVersion string) error
+	stateManager   *statemanager.Manager
 
 	cancel context.CancelFunc
 	update UpdateInterface
@@ -49,15 +60,17 @@
 	expectedVersionMutex  sync.Mutex
 }
 
-func NewUpdateManager(statusRecorder *peer.Status) *UpdateManager {
+func NewUpdateManager(statusRecorder *peer.Status, stateManager *statemanager.Manager) *UpdateManager {
 	manager := &UpdateManager{
 		statusRecorder: statusRecorder,
 		mgmUpdateChan:  make(chan struct{}, 1),
 		updateChannel:  make(chan struct{}, 1),
 		currentVersion: version.NetbirdVersion(),
 		updateFunc:     triggerUpdate,
 		update:         version.NewUpdate("nb/client"),
+		stateManager:   stateManager,
 	}
+
 	return manager
 }
 
@@ -66,12 +79,35 @@
 	return u
 }
 
+func (u *UpdateManager) StartWithTimeout(ctx context.Context, timeout time.Duration) {
+	if u.cancel != nil {
+		log.Errorf("UpdateManager already started")
+		return
+	}
+
+	u.startInit(ctx)
+	ctx, cancel := context.WithTimeout(ctx, timeout)
+	u.cancel = cancel
+
+	u.wg.Add(1)
+	go u.updateLoop(ctx)
+}
+
 func (u *UpdateManager) Start(ctx context.Context) {
 	if u.cancel != nil {
 		log.Errorf("UpdateManager already started")
 		return
 	}
 
+	u.startInit(ctx)
+	ctx, cancel := context.WithCancel(ctx)
+	u.cancel = cancel
+
+	u.wg.Add(1)
+	go u.updateLoop(ctx)
+}
+
+func (u *UpdateManager) startInit(ctx context.Context) {
 	go u.update.StartFetcher()
 	u.update.SetDaemonVersion(u.currentVersion)
 	u.update.SetOnUpdateListener(func() {
@@ -81,11 +117,31 @@
 		}
 	})
 
-	ctx, cancel := context.WithCancel(ctx)
-	u.cancel = cancel
-
-	u.wg.Add(1)
-	go u.updateLoop(ctx)
+	u.stateManager.RegisterState(&UpdateState{})
+	if err := u.stateManager.LoadState(&UpdateState{}); err != nil {
+		log.Warnf("failed to load state: %v", err)
+		return
+	}
+	if u.stateManager.GetState(&UpdateState{}) == nil {
+		return
+	}
+	updateState := u.stateManager.GetState(&UpdateState{}).(*UpdateState)
+	log.Warnf("autoUpdate state loaded, %v", *updateState)
+	if updateState.TargetVersion == u.currentVersion {
+		log.Warnf("published notification event")
+		u.statusRecorder.PublishEvent(
+			cProto.SystemEvent_INFO,
+			cProto.SystemEvent_SYSTEM,
+			"Auto-update completed",
+			fmt.Sprintf("Your NetBird Client was auto-updated to version %s", u.currentVersion),
+			nil,
+		)
+	}
+	if err := u.stateManager.DeleteState(updateState); err != nil {
+		log.Warnf("failed to delete state: %v", err)
+	} else if err = u.stateManager.PersistState(ctx); err != nil {
+		log.Warnf("failed to persist state: %v", err)
+	}
 }
 
 func (u *UpdateManager) SetVersion(expectedVersion string) {
@@ -132,12 +188,24 @@
 	u.wg.Wait()
 }
 
+func (u *UpdateManager) onContextCancel() {
+	if u.cancel == nil {
+		return
+	}
+
+	if u.update != nil {
+		u.update.StopWatch()
+		u.update = nil
+	}
+}
+
 func (u *UpdateManager) updateLoop(ctx context.Context) {
 	defer u.wg.Done()
 
 	for {
 		select {
 		case <-ctx.Done():
+			u.onContextCancel()
 			return
 		case <-u.mgmUpdateChan:
 		case <-u.updateChannel:
@@ -189,9 +257,46 @@
 		nil,
 	)
 
-	err := u.updateFunc(ctx, updateVersion.String())
+	u.statusRecorder.PublishEvent(
+		cProto.SystemEvent_INFO,
+		cProto.SystemEvent_SYSTEM,
+		"",
+		"",
+		map[string]string{"progress_window": "show"},
+	)
+
+	updateState := UpdateState{
+		PreUpdateVersion: u.currentVersion,
+		TargetVersion:    updateVersion.String(),
+	}
+	err := u.stateManager.UpdateState(updateState)
+	if err != nil {
+		log.Warnf("failed to update state: %v", err)
+	} else {
+		err = u.stateManager.PersistState(ctx)
+		if err != nil {
+			log.Warnf("failed to persist state: %v", err)
+		}
+	}
+
+	err = u.updateFunc(ctx, updateVersion.String())
+
 	if err != nil {
 		log.Errorf("Error triggering auto-update: %v", err)
+		u.statusRecorder.PublishEvent(
+			cProto.SystemEvent_ERROR,
+			cProto.SystemEvent_SYSTEM,
+			"Auto-update failed",
+			fmt.Sprintf("Auto-update failed: %v", err),
+			nil,
+		)
+		u.statusRecorder.PublishEvent(
+			cProto.SystemEvent_INFO,
+			cProto.SystemEvent_SYSTEM,
+			"",
+			"",
+			map[string]string{"progress_window": "hide"},
+		)
 	}
 }
 

client/internal/updatemanager/manager_test.go

@@ -2,8 +2,11 @@ package updatemanager
 
 import (
 	"context"
+	"fmt"
 	v "github.com/hashicorp/go-version"
 	"github.com/netbirdio/netbird/client/internal/peer"
+	"github.com/netbirdio/netbird/client/internal/statemanager"
+	"path"
 	"testing"
 	"time"
 )
@@ -56,9 +59,10 @@ func Test_LatestVersion(t *testing.T) {
 		},
 	}
 
-	for _, c := range testMatrix {
+	for idx, c := range testMatrix {
 		mockUpdate := &versionUpdateMock{latestVersion: c.initialLatestVersion}
-		m := NewUpdateManager(peer.NewRecorder("")).WithCustomVersionUpdate(mockUpdate)
+		tmpFile := path.Join(t.TempDir(), fmt.Sprintf("update-test-%d.json", idx))
+		m := NewUpdateManager(peer.NewRecorder(""), statemanager.New(tmpFile)).WithCustomVersionUpdate(mockUpdate)
 
 		targetVersionChan := make(chan string, 1)
 
@@ -169,8 +173,9 @@ func Test_HandleUpdate(t *testing.T) {
 			shouldUpdate:    false,
 		},
 	}
-	for _, c := range testMatrix {
-		m := NewUpdateManager(peer.NewRecorder("")).WithCustomVersionUpdate(&versionUpdateMock{latestVersion: c.latestVersion})
+	for idx, c := range testMatrix {
+		tmpFile := path.Join(t.TempDir(), fmt.Sprintf("update-test-%d.json", idx))
+		m := NewUpdateManager(peer.NewRecorder(""), statemanager.New(tmpFile)).WithCustomVersionUpdate(&versionUpdateMock{latestVersion: c.latestVersion})
 		targetVersionChan := make(chan string, 1)
 
 		m.updateFunc = func(ctx context.Context, targetVersion string) error {

client/internal/updatemanager/update_windows.go

@@ -5,9 +5,8 @@ package updatemanager
 import (
 	"context"
 	"fmt"
-	"os/exec"
-
 	"golang.org/x/sys/windows/registry"
+	"os/exec"
 
 	log "github.com/sirupsen/logrus"
 )