add health check for agent

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

beszel/internal/agent/server.go

@@ -23,20 +23,14 @@ func (a *Agent) StartServer(opts ServerOptions) error {
 
 	slog.Info("Starting SSH server", "addr", opts.Addr, "network", opts.Network)
 
-	switch opts.Network {
-	case "unix":
+	if opts.Network == "unix" {
 		// remove existing socket file if it exists
 		if err := os.Remove(opts.Addr); err != nil && !os.IsNotExist(err) {
 			return err
 		}
-	default:
-		// prefix with : if only port was provided
-		if !strings.Contains(opts.Addr, ":") {
-			opts.Addr = ":" + opts.Addr
-		}
 	}
 
-	// Listen on the address
+	// start listening on the address
 	ln, err := net.Listen(opts.Network, opts.Addr)
 	if err != nil {
 		return err
@@ -44,7 +38,7 @@ func (a *Agent) StartServer(opts ServerOptions) error {
 	defer ln.Close()
 
 	// Start SSH server on the listener
-	err = sshServer.Serve(ln, nil, sshServer.NoPty(),
+	return sshServer.Serve(ln, nil, sshServer.NoPty(),
 		sshServer.PublicKeyAuth(func(ctx sshServer.Context, key sshServer.PublicKey) bool {
 			for _, pubKey := range opts.Keys {
 				if sshServer.KeysEqual(key, pubKey) {
@@ -54,10 +48,6 @@ func (a *Agent) StartServer(opts ServerOptions) error {
 			return false
 		}),
 	)
-	if err != nil {
-		return err
-	}
-	return nil
 }
 
 func (a *Agent) handleSession(s sshServer.Session) {
@@ -89,3 +79,33 @@ func ParseKeys(input string) ([]ssh.PublicKey, error) {
 	}
 	return parsedKeys, nil
 }
+
+// GetAddress gets the address to listen on or connect to from environment variables or default value.
+func GetAddress(addr string) string {
+	if addr == "" {
+		addr, _ = GetEnv("LISTEN")
+	}
+	if addr == "" {
+		// Legacy PORT environment variable support
+		addr, _ = GetEnv("PORT")
+	}
+	if addr == "" {
+		return ":45876"
+	}
+	// prefix with : if only port was provided
+	if GetNetwork(addr) != "unix" && !strings.Contains(addr, ":") {
+		addr = ":" + addr
+	}
+	return addr
+}
+
+// GetNetwork returns the network type to use based on the address
+func GetNetwork(addr string) string {
+	if network, ok := GetEnv("NETWORK"); ok && network != "" {
+		return network
+	}
+	if strings.HasPrefix(addr, "/") {
+		return "unix"
+	}
+	return "tcp"
+}