fix update mirror and make opt-in with --china-mirrors (#1035)

beszel/cmd/agent/agent.go

@@ -4,12 +4,12 @@ import (
 	"beszel"
 	"beszel/internal/agent"
 	"beszel/internal/agent/health"
-	"flag"
 	"fmt"
 	"log"
 	"os"
 	"strings"
 
+	"github.com/spf13/pflag"
 	"golang.org/x/crypto/ssh"
 )
 
@@ -17,43 +17,24 @@ import (
 type cmdOptions struct {
 	key    string // key is the public key(s) for SSH authentication.
 	listen string // listen is the address or port to listen on.
+	// TODO: add hubURL and token
+	// hubURL string // hubURL is the URL of the hub to use.
+	// token  string // token is the token to use for authentication.
 }
 
 // parse parses the command line flags and populates the config struct.
 // It returns true if a subcommand was handled and the program should exit.
 func (opts *cmdOptions) parse() bool {
-	flag.StringVar(&opts.key, "key", "", "Public key(s) for SSH authentication")
-	flag.StringVar(&opts.listen, "listen", "", "Address or port to listen on")
-
-	flag.Usage = func() {
-		builder := strings.Builder{}
-		builder.WriteString("Usage: ")
-		builder.WriteString(os.Args[0])
-		builder.WriteString(" [command] [flags]\n")
-		builder.WriteString("\nCommands:\n")
-		builder.WriteString("  health    Check if the agent is running\n")
-		builder.WriteString("  help      Display this help message\n")
-		builder.WriteString("  update    Update to the latest version\n")
-		builder.WriteString("\nFlags:\n")
-		fmt.Print(builder.String())
-		flag.PrintDefaults()
-	}
-
 	subcommand := ""
 	if len(os.Args) > 1 {
 		subcommand = os.Args[1]
 	}
 
+	// Subcommands that don't require any pflag parsing
 	switch subcommand {
 	case "-v", "version":
 		fmt.Println(beszel.AppName+"-agent", beszel.Version)
 		return true
-	case "help":
-		flag.Usage()
-		return true
-	case "update":
-		agent.Update()
-		return true
 	case "health":
 		err := health.Check()
 		if err != nil {
@@ -63,7 +44,57 @@ func (opts *cmdOptions) parse() bool {
 		return true
 	}
 
-	flag.Parse()
+	// pflag.CommandLine.ParseErrorsWhitelist.UnknownFlags = true
+	pflag.StringVarP(&opts.key, "key", "k", "", "Public key(s) for SSH authentication")
+	pflag.StringVarP(&opts.listen, "listen", "l", "", "Address or port to listen on")
+	// pflag.StringVarP(&opts.hubURL, "hub-url", "u", "", "URL of the hub to use")
+	// pflag.StringVarP(&opts.token, "token", "t", "", "Token to use for authentication")
+	chinaMirrors := pflag.BoolP("china-mirrors", "c", false, "Use mirror for update (gh.beszel.dev) instead of GitHub")
+	help := pflag.BoolP("help", "h", false, "Show this help message")
+
+	// Convert old single-dash long flags to double-dash for backward compatibility
+	flagsToConvert := []string{"key", "listen"}
+	for i, arg := range os.Args {
+		for _, flag := range flagsToConvert {
+			singleDash := "-" + flag
+			doubleDash := "--" + flag
+			if arg == singleDash {
+				os.Args[i] = doubleDash
+				break
+			} else if strings.HasPrefix(arg, singleDash+"=") {
+				os.Args[i] = doubleDash + arg[len(singleDash):]
+				break
+			}
+		}
+	}
+
+	pflag.Usage = func() {
+		builder := strings.Builder{}
+		builder.WriteString("Usage: ")
+		builder.WriteString(os.Args[0])
+		builder.WriteString(" [command] [flags]\n")
+		builder.WriteString("\nCommands:\n")
+		builder.WriteString("  health    Check if the agent is running\n")
+		// builder.WriteString("  help      Display this help message\n")
+		builder.WriteString("  update    Update to the latest version\n")
+		builder.WriteString("\nFlags:\n")
+		fmt.Print(builder.String())
+		pflag.PrintDefaults()
+	}
+
+	// Parse all arguments with pflag
+	pflag.Parse()
+
+	// Must run after pflag.Parse()
+	switch {
+	case *help || subcommand == "help":
+		pflag.Usage()
+		return true
+	case subcommand == "update":
+		agent.Update(*chinaMirrors)
+		return true
+	}
+
 	return false
 }
 

beszel/cmd/agent/agent_test.go

@@ -3,11 +3,11 @@ package main
 import (
 	"beszel/internal/agent"
 	"crypto/ed25519"
-	"flag"
 	"os"
 	"path/filepath"
 	"testing"
 
+	"github.com/spf13/pflag"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"golang.org/x/crypto/ssh"
@@ -245,7 +245,7 @@ func TestParseFlags(t *testing.T) {
 	oldArgs := os.Args
 	defer func() {
 		os.Args = oldArgs
-		flag.CommandLine = flag.NewFlagSet(os.Args[0], flag.ExitOnError)
+		pflag.CommandLine = pflag.NewFlagSet(os.Args[0], pflag.ExitOnError)
 	}()
 
 	tests := []struct {
@@ -269,6 +269,22 @@ func TestParseFlags(t *testing.T) {
 				listen: "",
 			},
 		},
+		{
+			name: "key flag double dash",
+			args: []string{"cmd", "--key", "testkey"},
+			expected: cmdOptions{
+				key:    "testkey",
+				listen: "",
+			},
+		},
+		{
+			name: "key flag short",
+			args: []string{"cmd", "-k", "testkey"},
+			expected: cmdOptions{
+				key:    "testkey",
+				listen: "",
+			},
+		},
 		{
 			name: "addr flag only",
 			args: []string{"cmd", "-listen", ":8080"},
@@ -277,6 +293,22 @@ func TestParseFlags(t *testing.T) {
 				listen: ":8080",
 			},
 		},
+		{
+			name: "addr flag double dash",
+			args: []string{"cmd", "--listen", ":8080"},
+			expected: cmdOptions{
+				key:    "",
+				listen: ":8080",
+			},
+		},
+		{
+			name: "addr flag short",
+			args: []string{"cmd", "-l", ":8080"},
+			expected: cmdOptions{
+				key:    "",
+				listen: ":8080",
+			},
+		},
 		{
 			name: "both flags",
 			args: []string{"cmd", "-key", "testkey", "-listen", ":8080"},
@@ -290,12 +322,12 @@ func TestParseFlags(t *testing.T) {
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			// Reset flags for each test
-			flag.CommandLine = flag.NewFlagSet(tt.args[0], flag.ExitOnError)
+			pflag.CommandLine = pflag.NewFlagSet(tt.args[0], pflag.ExitOnError)
 			os.Args = tt.args
 
 			var opts cmdOptions
 			opts.parse()
-			flag.Parse()
+			pflag.Parse()
 
 			assert.Equal(t, tt.expected, opts)
 		})

beszel/cmd/hub/hub.go

@@ -45,11 +45,13 @@ func getBaseApp() *pocketbase.PocketBase {
 	baseApp.RootCmd.Use = beszel.AppName
 	baseApp.RootCmd.Short = ""
 	// add update command
-	baseApp.RootCmd.AddCommand(&cobra.Command{
+	updateCmd := &cobra.Command{
 		Use:   "update",
 		Short: "Update " + beszel.AppName + " to the latest version",
 		Run:   hub.Update,
-	})
+	}
+	updateCmd.Flags().Bool("china-mirrors", false, "Use mirror (gh.beszel.dev) instead of GitHub")
+	baseApp.RootCmd.AddCommand(updateCmd)
 	// add health command
 	baseApp.RootCmd.AddCommand(newHealthCmd())
 

beszel/internal/agent/update.go

@@ -60,7 +60,7 @@ func detectRestarter() restarter {
 
 // Update checks GitHub for a newer release of beszel-agent, applies it,
 // fixes SELinux context if needed, and restarts the service.
-func Update() error {
+func Update(useMirror bool) error {
 	exePath, _ := os.Executable()
 
 	dataDir, err := getDataDir()
@@ -70,6 +70,7 @@ func Update() error {
 	updated, err := ghupdate.Update(ghupdate.Config{
 		ArchiveExecutable: "beszel-agent",
 		DataDir:           dataDir,
+		UseMirror:         useMirror,
 	})
 	if err != nil {
 		log.Fatal(err)
@@ -99,6 +100,8 @@ func Update() error {
 		if err := r.Restart(); err != nil {
 			ghupdate.ColorPrintf(ghupdate.ColorYellow, "Warning: failed to restart service: %v", err)
 			ghupdate.ColorPrint(ghupdate.ColorYellow, "Please restart the service manually.")
+		} else {
+			ghupdate.ColorPrint(ghupdate.ColorGreen, "Service restarted successfully")
 		}
 	} else {
 		ghupdate.ColorPrint(ghupdate.ColorYellow, "No supported init system detected; please restart manually if needed.")

beszel/internal/ghupdate/ghupdate.go

@@ -23,7 +23,7 @@ import (
 const (
 	colorReset  = "\033[0m"
 	ColorYellow = "\033[33m"
-	colorGreen  = "\033[32m"
+	ColorGreen  = "\033[32m"
 	colorCyan   = "\033[36m"
 	colorGray   = "\033[90m"
 )
@@ -64,6 +64,10 @@ type Config struct {
 
 	// The data directory to use when fetching and downloading the latest release.
 	DataDir string
+
+	// UseMirror specifies whether to use the beszel.dev mirror instead of GitHub API.
+	// When false (default), always uses api.github.com. When true, uses gh.beszel.dev.
+	UseMirror bool
 }
 
 type updater struct {
@@ -106,21 +110,19 @@ func (p *updater) update() (updated bool, err error) {
 	var latest *release
 	var useMirror bool
 
+	// Determine the API endpoint based on UseMirror flag
+	apiURL := fmt.Sprintf("https://api.github.com/repos/%s/%s/releases/latest", p.config.Owner, p.config.Repo)
+	if p.config.UseMirror {
+		useMirror = true
+		apiURL = fmt.Sprintf("https://gh.beszel.dev/repos/%s/%s/releases/latest?api=true", p.config.Owner, p.config.Repo)
+		ColorPrint(ColorYellow, "Using mirror for update.")
+	}
+
 	latest, err = fetchLatestRelease(
 		p.config.Context,
 		p.config.HttpClient,
-		fmt.Sprintf("https://api.github.com/repos/%s/%s/releases/latest", p.config.Owner, p.config.Repo),
+		apiURL,
 	)
-	// if the first fetch fails, try the beszel.dev API (fallback for China)
-	if err != nil {
-		ColorPrint(ColorYellow, "Failed to fetch release. Trying beszel.dev mirror...")
-		useMirror = true
-		latest, err = fetchLatestRelease(
-			p.config.Context,
-			p.config.HttpClient,
-			fmt.Sprintf("https://gh.beszel.dev/repos/%s/%s/releases/latest?api=true", p.config.Owner, p.config.Repo),
-		)
-	}
 	if err != nil {
 		return false, err
 	}
@@ -129,7 +131,7 @@ func (p *updater) update() (updated bool, err error) {
 	newVersion := semver.MustParse(strings.TrimPrefix(latest.Tag, "v"))
 
 	if newVersion.LTE(currentVersion) {
-		ColorPrintf(colorGreen, "You already have the latest version %s.", p.currentVersion)
+		ColorPrintf(ColorGreen, "You already have the latest version %s.", p.currentVersion)
 		return false, nil
 	}
 
@@ -209,14 +211,11 @@ func (p *updater) update() (updated bool, err error) {
 	}
 
 	ColorPrint(colorGray, "---")
-	ColorPrint(colorGreen, "Update completed successfully! You can start the executable as usual.")
+	ColorPrint(ColorGreen, "Update completed successfully!")
 
 	// print the release notes
 	if latest.Body != "" {
 		fmt.Print("\n")
-		ColorPrintf(colorCyan, "Here is a list with some of the %s changes:", latest.Tag)
-		// remove the update command note to avoid "stuttering"
-		// (@todo consider moving to a config option)
 		releaseNotes := strings.TrimSpace(strings.Replace(latest.Body, "> _To update the prebuilt executable you can run `./"+p.config.ArchiveExecutable+" update`._", "", 1))
 		ColorPrint(colorCyan, releaseNotes)
 		fmt.Print("\n")

beszel/internal/hub/update.go

@@ -11,7 +11,7 @@ import (
 )
 
 // Update updates beszel to the latest version
-func Update(_ *cobra.Command, _ []string) {
+func Update(cmd *cobra.Command, _ []string) {
 	dataDir := os.TempDir()
 
 	// set dataDir to ./beszel_data if it exists
@@ -19,9 +19,13 @@ func Update(_ *cobra.Command, _ []string) {
 		dataDir = "./beszel_data"
 	}
 
+	// Check if china-mirrors flag is set
+	useMirror, _ := cmd.Flags().GetBool("china-mirrors")
+
 	updated, err := ghupdate.Update(ghupdate.Config{
 		ArchiveExecutable: "beszel",
 		DataDir:           dataDir,
+		UseMirror:         useMirror,
 	})
 	if err != nil {
 		log.Fatal(err)
@@ -49,13 +53,13 @@ func restartService() {
 		// Check if beszel service exists and is active
 		cmd := exec.Command("systemctl", "is-active", "beszel.service")
 		if err := cmd.Run(); err == nil {
-			fmt.Println("Restarting beszel service...")
+			ghupdate.ColorPrint(ghupdate.ColorYellow, "Restarting beszel service...")
 			restartCmd := exec.Command("systemctl", "restart", "beszel.service")
 			if err := restartCmd.Run(); err != nil {
-				fmt.Printf("Warning: Failed to restart service: %v\n", err)
+				ghupdate.ColorPrintf(ghupdate.ColorYellow, "Warning: Failed to restart service: %v\n", err)
-				fmt.Println("Please restart the service manually: sudo systemctl restart beszel")
+				ghupdate.ColorPrint(ghupdate.ColorYellow, "Please restart the service manually: sudo systemctl restart beszel")
 			} else {
-				fmt.Println("Service restarted successfully")
+				ghupdate.ColorPrint(ghupdate.ColorGreen, "Service restarted successfully")
 			}
 			return
 		}
@@ -65,17 +69,17 @@ func restartService() {
 	if _, err := exec.LookPath("rc-service"); err == nil {
 		cmd := exec.Command("rc-service", "beszel", "status")
 		if err := cmd.Run(); err == nil {
-			fmt.Println("Restarting beszel service...")
+			ghupdate.ColorPrint(ghupdate.ColorYellow, "Restarting beszel service...")
 			restartCmd := exec.Command("rc-service", "beszel", "restart")
 			if err := restartCmd.Run(); err != nil {
-				fmt.Printf("Warning: Failed to restart service: %v\n", err)
+				ghupdate.ColorPrintf(ghupdate.ColorYellow, "Warning: Failed to restart service: %v\n", err)
-				fmt.Println("Please restart the service manually: sudo rc-service beszel restart")
+				ghupdate.ColorPrint(ghupdate.ColorYellow, "Please restart the service manually: sudo rc-service beszel restart")
 			} else {
-				fmt.Println("Service restarted successfully")
+				ghupdate.ColorPrint(ghupdate.ColorGreen, "Service restarted successfully")
 			}
 			return
 		}
 	}
 
-	fmt.Println("Note: Service restart not attempted. If running as a service, restart manually.")
+	ghupdate.ColorPrint(ghupdate.ColorYellow, "Service restart not attempted. If running as a service, restart manually.")
 }