feat(agent): NETWORK env var and support for multiple keys

- merges agent.Run with agent.NewAgent
- separates StartServer method
- bumps go version to 1.24
- add tests

beszel/cmd/agent/agent.go

@@ -8,12 +8,19 @@ import (
 	"log"
 	"os"
 	"strings"
+
+	"golang.org/x/crypto/ssh"
 )
 
-func main() {
-	// Define flags for key and port
-	keyFlag := flag.String("key", "", "Public key")
-	portFlag := flag.String("port", "45876", "Port number")
+type cmdConfig struct {
+	key  string // key is the public key(s) for SSH authentication.
+	addr string // addr is the address or port to listen on.
+}
+
+// parseFlags parses the command line flags and populates the config struct.
+func parseFlags(cfg *cmdConfig) {
+	flag.StringVar(&cfg.key, "key", "", "Public key(s) for SSH authentication")
+	flag.StringVar(&cfg.addr, "addr", "", "Address or port to listen on")
 
 	flag.Usage = func() {
 		fmt.Printf("Usage: %s [options] [subcommand]\n", os.Args[0])
@@ -24,65 +31,103 @@ func main() {
 		fmt.Println("  help         Display this help message")
 		fmt.Println("  update       Update the agent to the latest version")
 	}
+}
+
+// handleSubcommand handles subcommands such as version, help, and update.
+// It returns true if a subcommand was handled, false otherwise.
+func handleSubcommand() bool {
+	if len(os.Args) <= 1 {
+		return false
+	}
+	switch os.Args[1] {
+	case "version", "-v":
+		fmt.Println(beszel.AppName+"-agent", beszel.Version)
+		os.Exit(0)
+	case "help":
+		flag.Usage()
+		os.Exit(0)
+	case "update":
+		agent.Update()
+		os.Exit(0)
+	}
+	return false
+}
+
+// loadPublicKeys loads the public keys from the command line flag, environment variable, or key file.
+func loadPublicKeys(cfg cmdConfig) ([]ssh.PublicKey, error) {
+	// Try command line flag first
+	if cfg.key != "" {
+		return agent.ParseKeys(cfg.key)
+	}
+
+	// Try environment variable
+	if key, ok := agent.GetEnv("KEY"); ok && key != "" {
+		return agent.ParseKeys(key)
+	}
+
+	// Try key file
+	keyFile, ok := agent.GetEnv("KEY_FILE")
+	if !ok {
+		return nil, fmt.Errorf("no key provided: must set -key flag, KEY env var, or KEY_FILE env var. ")
+	}
+
+	pubKey, err := os.ReadFile(keyFile)
+	if err != nil {
+		return nil, fmt.Errorf("failed to read key file: %w", err)
+	}
+	return agent.ParseKeys(string(pubKey))
+}
+
+// getAddress gets the address to listen on from the command line flag, environment variable, or default value.
+func getAddress(addr string) string {
+	// Try command line flag first
+	if addr != "" {
+		return addr
+	}
+	// Try environment variables
+	if addr, ok := agent.GetEnv("ADDR"); ok && addr != "" {
+		return addr
+	}
+	// Legacy PORT environment variable support
+	if port, ok := agent.GetEnv("PORT"); ok && port != "" {
+		return port
+	}
+	return ":45876"
+}
+
+// getNetwork returns the network type to use for the server.
+func getNetwork(addr string) string {
+	if network, _ := agent.GetEnv("NETWORK"); network != "" {
+		return network
+	}
+	if strings.HasPrefix(addr, "/") {
+		return "unix"
+	}
+	return "tcp"
+}
+
+func main() {
+	var cfg cmdConfig
+	parseFlags(&cfg)
+
+	if handleSubcommand() {
+		return
+	}
 
-	// Parse the flags
 	flag.Parse()
 
-	// handle flags / subcommands
-	if len(os.Args) > 1 {
-		switch os.Args[1] {
-		case "version":
-			fmt.Println(beszel.AppName+"-agent", beszel.Version)
-			os.Exit(0)
-		case "help":
-			flag.Usage()
-			os.Exit(0)
-		case "update":
-			agent.Update()
-			os.Exit(0)
-		}
+	var serverConfig agent.ServerConfig
+	var err error
+	serverConfig.Keys, err = loadPublicKeys(cfg)
+	if err != nil {
+		log.Fatal("Failed to load public keys:", err)
 	}
 
-	var pubKey []byte
-	// Override the key if the -key flag is provided
-	if *keyFlag != "" {
-		pubKey = []byte(*keyFlag)
-	} else {
-		// Try to get the key from the KEY environment variable.
-		key, _ := agent.GetEnv("KEY")
-		pubKey = []byte(key)
-	}
+	serverConfig.Addr = getAddress(cfg.addr)
+	serverConfig.Network = getNetwork(cfg.addr)
 
-	// If KEY is not set, try to read the key from the file specified by KEY_FILE.
-	if len(pubKey) == 0 {
-		keyFile, exists := agent.GetEnv("KEY_FILE")
-		if !exists {
-			log.Fatal("Must set KEY or KEY_FILE environment variable or supply as input argument. Use 'beszel-agent help' for more information.")
-		}
-		var err error
-		pubKey, err = os.ReadFile(keyFile)
-		if err != nil {
-			log.Fatal(err)
-		}
+	agent := agent.NewAgent()
+	if err := agent.StartServer(serverConfig); err != nil {
+		log.Fatal("Failed to start server:", err)
 	}
-
-	// Init with default port
-	addr := ":" + *portFlag
-	
-	//Use port from ENV if it exists
-	// TODO: change env var to ADDR
-	if portEnvVar, exists := agent.GetEnv("PORT"); exists {
-		// allow passing an address in the form of "127.0.0.1:45876"
-		if !strings.Contains(portEnvVar, ":") {
-			portEnvVar = ":" + portEnvVar
-		}
-		addr = portEnvVar
-	}
-
-	// Override the default and ENV port if the -port flag is provided and is non default
-	if *portFlag != "45876" {
-		addr = ":" + *portFlag
-	}
-
-	agent.NewAgent().Run(pubKey, addr)
 }