Merge pull request #4563 from netbirdio/auto-upgrade-mod

Modify client-side behavior

client/internal/connect.go

@@ -25,6 +25,7 @@ import (
 	"github.com/netbirdio/netbird/client/internal/peer"
 	"github.com/netbirdio/netbird/client/internal/profilemanager"
 	"github.com/netbirdio/netbird/client/internal/stdnet"
+	nbnet "github.com/netbirdio/netbird/client/net"
 	cProto "github.com/netbirdio/netbird/client/proto"
 	"github.com/netbirdio/netbird/client/ssh"
 	"github.com/netbirdio/netbird/client/system"
@@ -34,7 +35,6 @@ import (
 	relayClient "github.com/netbirdio/netbird/shared/relay/client"
 	signal "github.com/netbirdio/netbird/shared/signal/client"
 	"github.com/netbirdio/netbird/util"
-	nbnet "github.com/netbirdio/netbird/client/net"
 	"github.com/netbirdio/netbird/version"
 )
 
@@ -272,6 +272,9 @@ 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 {
+			c.engine.handleAutoUpdateVersion(loginResp.PeerConfig.AutoUpdate)
+		}
 		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

@@ -16,6 +16,7 @@ import (
 	log "github.com/sirupsen/logrus"
 
 	"github.com/netbirdio/netbird/client/internal/peer"
+	"github.com/netbirdio/netbird/client/internal/statemanager"
 	cProto "github.com/netbirdio/netbird/client/proto"
 	"github.com/netbirdio/netbird/version"
 )
@@ -32,6 +33,15 @@ type UpdateInterface interface {
 	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 @@ type UpdateManager struct {
 	wg             sync.WaitGroup
 	currentVersion string
 	updateFunc     func(ctx context.Context, targetVersion string) error
+	stateManager   *statemanager.Manager
 
 	cancel context.CancelFunc
 	update UpdateInterface
@@ -49,7 +60,7 @@ type UpdateManager struct {
 	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),
@@ -57,17 +68,41 @@ func NewUpdateManager(statusRecorder *peer.Status) *UpdateManager {
 		currentVersion: version.NetbirdVersion(),
 		updateFunc:     triggerUpdate,
 		update:         version.NewUpdate("nb/client"),
+		stateManager:   stateManager,
 	}
+
 	return manager
 }
 
+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
 	}
 
-	go u.update.StartFetcher()
+	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) {
 	u.update.SetDaemonVersion(u.currentVersion)
 	u.update.SetOnUpdateListener(func() {
 		select {
@@ -75,12 +110,33 @@ func (u *UpdateManager) Start(ctx context.Context) {
 		default:
 		}
 	})
+	go u.update.StartFetcher()
 
-	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) {
@@ -129,12 +185,26 @@ func (u *UpdateManager) Stop() {
 	u.wg.Wait()
 }
 
+func (u *UpdateManager) onContextCancel() {
+	if u.cancel == nil {
+		return
+	}
+
+	u.expectedVersionMutex.Lock()
+	defer u.expectedVersionMutex.Unlock()
+	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 +259,46 @@ func (u *UpdateManager) handleUpdate(ctx context.Context) {
 		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"
 )
@@ -61,9 +64,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)
 
@@ -174,8 +178,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"
 )

client/ui/client_ui.go

@@ -94,13 +94,14 @@ func main() {
 		showLoginURL: flags.showLoginURL,
 		showDebug:    flags.showDebug,
 		showProfiles: flags.showProfiles,
+		showUpdate:   flags.showUpdate,
 	})
 
 	// Watch for theme/settings changes to update the icon.
 	go watchSettingsChanges(a, client)
 
 	// Run in window mode if any UI flag was set.
-	if flags.showSettings || flags.showNetworks || flags.showDebug || flags.showLoginURL || flags.showProfiles {
+	if flags.showSettings || flags.showNetworks || flags.showDebug || flags.showLoginURL || flags.showProfiles || flags.showUpdate {
 		a.Run()
 		return
 	}
@@ -128,6 +129,7 @@ type cliFlags struct {
 	showDebug      bool
 	showLoginURL   bool
 	errorMsg       string
+	showUpdate     bool
 	saveLogsInFile bool
 }
 
@@ -147,6 +149,7 @@ func parseFlags() *cliFlags {
 	flag.StringVar(&flags.errorMsg, "error-msg", "", "displays an error message window")
 	flag.BoolVar(&flags.saveLogsInFile, "use-log-file", false, fmt.Sprintf("save logs in a file: %s/netbird-ui-PID.log", os.TempDir()))
 	flag.BoolVar(&flags.showLoginURL, "login-url", false, "show login URL in a popup window")
+	flag.BoolVar(&flags.showUpdate, "update", false, "show update progress window")
 	flag.Parse()
 	return &flags
 }
@@ -297,6 +300,8 @@ type serviceClient struct {
 	mExitNodeDeselectAll *systray.MenuItem
 	logFile              string
 	wLoginURL            fyne.Window
+	wUpdateProgress      fyne.Window
+	updateContextCancel  context.CancelFunc
 }
 
 type menuHandler struct {
@@ -313,6 +318,7 @@ type newServiceClientArgs struct {
 	showDebug    bool
 	showLoginURL bool
 	showProfiles bool
+	showUpdate   bool
 }
 
 // newServiceClient instance constructor
@@ -348,6 +354,8 @@ func newServiceClient(args *newServiceClientArgs) *serviceClient {
 		s.showDebugUI()
 	case args.showProfiles:
 		s.showProfilesUI()
+	case args.showUpdate:
+		s.showUpdateProgress(ctx)
 	}
 
 	return s
@@ -393,6 +401,30 @@ func (s *serviceClient) updateIcon() {
 	s.updateIndicationLock.Unlock()
 }
 
+func (s *serviceClient) showUpdateProgress(ctx context.Context) {
+	s.wUpdateProgress = s.app.NewWindow("Automatically updating client")
+	loadingLabel := widget.NewLabel("Updating")
+	s.wUpdateProgress.SetContent(container.NewGridWithRows(2, widget.NewLabel("Your client version is older than auto-update version set in Management, updating client now."), loadingLabel))
+	s.wUpdateProgress.Show()
+	go func() {
+		dotCount := 0
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			case <-time.After(time.Second):
+				dotCount++
+				dotCount %= 4
+				loadingLabel.SetText(fmt.Sprintf("Updating%s", strings.Repeat(".", dotCount)))
+			}
+		}
+	}()
+	s.wUpdateProgress.CenterOnScreen()
+	s.wUpdateProgress.SetFixedSize(true)
+	s.wUpdateProgress.SetCloseIntercept(func() {})
+	s.wUpdateProgress.RequestFocus()
+}
+
 func (s *serviceClient) showSettingsUI() {
 	// Check if update settings are disabled by daemon
 	features, err := s.getFeatures()
@@ -951,6 +983,29 @@ func (s *serviceClient) onTrayReady() {
 			s.updateExitNodes()
 		}
 	})
+	s.eventManager.AddHandler(func(event *proto.SystemEvent) {
+		if windowAction, ok := event.Metadata["progress_window"]; ok {
+			log.Debugf("window action: %v", windowAction)
+			if windowAction == "show" {
+				log.Debugf("Inside show")
+				if s.updateContextCancel != nil {
+					s.updateContextCancel()
+					s.updateContextCancel = nil
+				}
+
+				subCtx, cancel := context.WithCancel(s.ctx)
+				go s.eventHandler.runSelfCommand(subCtx, "update", "true")
+				s.updateContextCancel = cancel
+			}
+			if windowAction == "hide" {
+				log.Debugf("Inside hide")
+				if s.updateContextCancel != nil {
+					s.updateContextCancel()
+					s.updateContextCancel = nil
+				}
+			}
+		}
+	})
 
 	go s.eventManager.Start(s.ctx)
 	go s.eventHandler.listen(s.ctx)

management/server/grpcserver.go

@@ -712,7 +712,9 @@ func toPeerConfig(peer *nbpeer.Peer, network *types.Network, dnsName string, set
 		Fqdn:                            fqdn,
 		RoutingPeerDnsResolutionEnabled: settings.RoutingPeerDNSResolutionEnabled,
 		LazyConnectionEnabled:           settings.LazyConnectionEnabled,
-		AutoUpdateVersion:               settings.AutoUpdateVersion,
+		AutoUpdate: &proto.AutoUpdateSettings{
+			Version: settings.AutoUpdateVersion,
+		},
 	}
 }
 

management/server/types/settings.go

@@ -54,7 +54,7 @@ type Settings struct {
 	LazyConnectionEnabled bool `gorm:"default:false"`
 
 	// AutoUpdateVersion client auto-update version
-	AutoUpdateVersion string
+	AutoUpdateVersion string `gorm:"default:'latest'"`
 }
 
 // Copy copies the Settings struct

shared/management/proto/management.proto

@@ -268,7 +268,13 @@ message PeerConfig {
   int32 mtu = 7;
 
   // Auto-update config
-  string autoUpdateVersion = 8;
+  AutoUpdateSettings autoUpdate = 8;
+}
+
+message AutoUpdateSettings {
+  string version = 1;
+  // When false, only update if the connection started < 1 minute ago
+  bool alwaysUpdate = 2;
 }
 
 // NetworkMap represents a network state of the peer with the corresponding configuration parameters to establish peer-to-peer connections