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)
 		})