From abaa9e474eafb26861722e69785e33a0f04315a7 Mon Sep 17 00:00:00 2001
From: naison <895703375@qq.com>
Date: Thu, 23 Oct 2025 11:39:53 +0800
Subject: [PATCH] refactor: upgrade logic

---
 cmd/kubevpn/cmds/upgrade.go | 33 +--------------------------------
 pkg/daemon/handler/ssh.go   |  2 +-
 pkg/upgrade/upgrade.go      | 24 +++++++++++++++++++++---
 3 files changed, 23 insertions(+), 36 deletions(-)

diff --git a/cmd/kubevpn/cmds/upgrade.go b/cmd/kubevpn/cmds/upgrade.go
index 563965d5..1ff75db3 100644
--- a/cmd/kubevpn/cmds/upgrade.go
+++ b/cmd/kubevpn/cmds/upgrade.go
@@ -1,17 +1,11 @@
 package cmds
 
 import (
-	"fmt"
-	"net/http"
-	"os"
-
 	"github.com/spf13/cobra"
-	"golang.org/x/oauth2"
 	cmdutil "k8s.io/kubectl/pkg/cmd/util"
 	"k8s.io/kubectl/pkg/util/i18n"
 	"k8s.io/kubectl/pkg/util/templates"
 
-	"github.com/wencaiwulue/kubevpn/v2/pkg/config"
 	plog "github.com/wencaiwulue/kubevpn/v2/pkg/log"
 	"github.com/wencaiwulue/kubevpn/v2/pkg/upgrade"
 )
@@ -26,33 +20,8 @@ func CmdUpgrade(cmdutil.Factory) *cobra.Command {
 		upgrade local daemon grpc server to latest version.
 		`)),
 		RunE: func(cmd *cobra.Command, args []string) error {
-			const (
-				envLatestUrl = "KUBEVPN_LATEST_VERSION_URL"
-			)
 			plog.InitLoggerForClient()
-			var client = http.DefaultClient
-			if config.GitHubOAuthToken != "" {
-				client = oauth2.NewClient(cmd.Context(), oauth2.StaticTokenSource(&oauth2.Token{AccessToken: config.GitHubOAuthToken, TokenType: "Bearer"}))
-			}
-			var url = os.Getenv(envLatestUrl)
-			if url == "" {
-				var latestVersion string
-				var needsUpgrade bool
-				var err error
-				url, latestVersion, needsUpgrade, err = upgrade.NeedsUpgrade(cmd.Context(), client, config.Version)
-				if err != nil {
-					return err
-				}
-				if !needsUpgrade {
-					_, _ = fmt.Fprintln(os.Stdout, fmt.Sprintf("Already up to date, don't needs to upgrade, version: %s", latestVersion))
-					return nil
-				}
-				_, _ = fmt.Fprintln(os.Stdout, fmt.Sprintf("Current version is: %s less than latest version: %s, needs to upgrade", config.Version, latestVersion))
-				_ = os.Setenv(envLatestUrl, url)
-				_ = quit(cmd.Context(), true)
-				_ = quit(cmd.Context(), false)
-			}
-			return upgrade.Main(cmd.Context(), client, url)
+			return upgrade.Main(cmd.Context(), quit)
 		},
 	}
 	return cmd
diff --git a/pkg/daemon/handler/ssh.go b/pkg/daemon/handler/ssh.go
index 4d4755df..dfedaff8 100644
--- a/pkg/daemon/handler/ssh.go
+++ b/pkg/daemon/handler/ssh.go
@@ -290,7 +290,7 @@ func (w *wsHandler) installKubevpnOnRemote(ctx context.Context, sshClient *ssh.C
 		w.Log("Get latest kubevpn version failed: %v", err)
 		return err
 	}
-	w.Log("The latest version is: %s", latestVersion)
+	w.Log("The latest version is %s", latestVersion)
 	var temp *os.File
 	temp, err = os.CreateTemp("", "")
 	if err != nil {
diff --git a/pkg/upgrade/upgrade.go b/pkg/upgrade/upgrade.go
index 283f0a7f..656333ef 100644
--- a/pkg/upgrade/upgrade.go
+++ b/pkg/upgrade/upgrade.go
@@ -10,7 +10,9 @@ import (
 	"strings"
 
 	goversion "github.com/hashicorp/go-version"
+	"golang.org/x/oauth2"
 
+	"github.com/wencaiwulue/kubevpn/v2/pkg/config"
 	"github.com/wencaiwulue/kubevpn/v2/pkg/daemon"
 	"github.com/wencaiwulue/kubevpn/v2/pkg/daemon/elevate"
 	plog "github.com/wencaiwulue/kubevpn/v2/pkg/log"
@@ -25,21 +27,37 @@ import (
 // 5) unzip to temp file
 // 6) check permission of putting new kubevpn back
 // 7) chmod +x, move old to /temp, move new to CURRENT_FOLDER
-func Main(ctx context.Context, client *http.Client, url string) error {
+func Main(ctx context.Context, quit func(ctx context.Context, isSudo bool) error) error {
 	err := elevatePermission()
 	if err != nil {
 		return err
 	}
 
+	var client = http.DefaultClient
+	if config.GitHubOAuthToken != "" {
+		client = oauth2.NewClient(ctx, oauth2.StaticTokenSource(&oauth2.Token{AccessToken: config.GitHubOAuthToken, TokenType: "Bearer"}))
+	}
+	url, latestVersion, needsUpgrade, err := NeedsUpgrade(ctx, client, config.Version)
+	if err != nil {
+		return err
+	}
+	if !needsUpgrade {
+		_, _ = fmt.Fprintln(os.Stdout, fmt.Sprintf("Already up to date, don't needs to upgrade, version: %s", latestVersion))
+		return nil
+	}
+	_, _ = fmt.Fprintln(os.Stdout, fmt.Sprintf("Current version: %s less than latest version: %s, needs to upgrade", config.Version, latestVersion))
+	_ = quit(ctx, true)
+	_ = quit(ctx, false)
+
 	err = downloadAndInstall(client, url)
 	if err != nil {
 		return err
 	}
 
 	plog.G(ctx).Infof("Upgrade daemon...")
-	err = daemon.StartupDaemon(context.Background())
+	_ = daemon.StartupDaemon(context.Background())
 	plog.G(ctx).Info("Done")
-	return err
+	return nil
 }
 
 func downloadAndInstall(client *http.Client, url string) error {