add health check for agent

- Updated command-line flag parsing.
- Moved GetAddress and GetNetwork to server.go

beszel/cmd/agent/agent.go

@@ -6,8 +6,8 @@ import (
 	"flag"
 	"fmt"
 	"log"
+	"log/slog"
 	"os"
-	"strings"
 
 	"golang.org/x/crypto/ssh"
 )
@@ -18,39 +18,50 @@ type cmdOptions struct {
 	listen string // listen is the address or port to listen on.
 }
 
-// parseFlags parses the command line flags and populates the config struct.
-func (opts *cmdOptions) parseFlags() {
+// 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() {
-		fmt.Printf("Usage: %s [options] [subcommand]\n", os.Args[0])
-		fmt.Println("\nOptions:")
+		fmt.Printf("Usage: %s [command] [flags]\n", os.Args[0])
+		fmt.Println("\nCommands:")
+		fmt.Println("  health    Check if the agent is running")
+		fmt.Println("  help      Display this help message")
+		fmt.Println("  update    Update to the latest version")
+		fmt.Println("  version   Display the version")
+		fmt.Println("\nFlags:")
 		flag.PrintDefaults()
-		fmt.Println("\nSubcommands:")
-		fmt.Println("  version      Display the version")
-		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
+	subcommand := ""
+	if len(os.Args) > 1 {
+		subcommand = os.Args[1]
 	}
-	switch os.Args[1] {
-	case "version", "-v":
+
+	switch subcommand {
+	case "-v", "version":
 		fmt.Println(beszel.AppName+"-agent", beszel.Version)
-		os.Exit(0)
+		return true
 	case "help":
 		flag.Usage()
-		os.Exit(0)
+		return true
 	case "update":
 		agent.Update()
-		os.Exit(0)
+		return true
+	case "health":
+		// for health, we need to parse flags first to get the listen address
+		args := append(os.Args[2:], subcommand)
+		flag.CommandLine.Parse(args)
+		addr := opts.getAddress()
+		network := agent.GetNetwork(addr)
+		exitCode, err := agent.Health(addr, network)
+		slog.Info("Health", "code", exitCode, "err", err)
+		os.Exit(exitCode)
 	}
+
+	flag.Parse()
 	return false
 }
 
@@ -79,46 +90,18 @@ func (opts *cmdOptions) loadPublicKeys() ([]ssh.PublicKey, error) {
 	return agent.ParseKeys(string(pubKey))
 }
 
-// getAddress gets the address to listen on from the command line flag, environment variable, or default value.
 func (opts *cmdOptions) getAddress() string {
-	// Try command line flag first
-	if opts.listen != "" {
-		return opts.listen
-	}
-	// Try environment variables
-	if addr, ok := agent.GetEnv("LISTEN"); 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 (opts *cmdOptions) getNetwork() string {
-	if network, _ := agent.GetEnv("NETWORK"); network != "" {
-		return network
-	}
-	if strings.HasPrefix(opts.listen, "/") {
-		return "unix"
-	}
-	return "tcp"
+	return agent.GetAddress(opts.listen)
 }
 
 func main() {
 	var opts cmdOptions
-	opts.parseFlags()
+	subcommandHandled := opts.parse()
 
-	if handleSubcommand() {
+	if subcommandHandled {
 		return
 	}
 
-	flag.Parse()
-
-	opts.listen = opts.getAddress()
-
 	var serverConfig agent.ServerOptions
 	var err error
 	serverConfig.Keys, err = opts.loadPublicKeys()
@@ -126,8 +109,9 @@ func main() {
 		log.Fatal("Failed to load public keys:", err)
 	}
 
-	serverConfig.Addr = opts.listen
-	serverConfig.Network = opts.getNetwork()
+	addr := opts.getAddress()
+	serverConfig.Addr = addr
+	serverConfig.Network = agent.GetNetwork(addr)
 
 	agent := agent.NewAgent()
 	if err := agent.StartServer(serverConfig); err != nil {

beszel/cmd/agent/agent_test.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"beszel/internal/agent"
 	"crypto/ed25519"
 	"flag"
 	"os"
@@ -29,7 +30,7 @@ func TestGetAddress(t *testing.T) {
 			opts: cmdOptions{
 				listen: "8080",
 			},
-			expected: "8080",
+			expected: ":8080",
 		},
 		{
 			name: "use unix socket from flag",
@@ -52,7 +53,7 @@ func TestGetAddress(t *testing.T) {
 			envVars: map[string]string{
 				"PORT": "7070",
 			},
-			expected: "7070",
+			expected: ":7070",
 		},
 		{
 			name: "use unix socket from env var",
@@ -233,7 +234,7 @@ func TestGetNetwork(t *testing.T) {
 			for k, v := range tt.envVars {
 				t.Setenv(k, v)
 			}
-			network := tt.opts.getNetwork()
+			network := agent.GetNetwork(tt.opts.listen)
 			assert.Equal(t, tt.expected, network)
 		})
 	}
@@ -293,7 +294,7 @@ func TestParseFlags(t *testing.T) {
 			os.Args = tt.args
 
 			var opts cmdOptions
-			opts.parseFlags()
+			opts.parse()
 			flag.Parse()
 
 			assert.Equal(t, tt.expected, opts)