cli: consolidate gRPC command helpers

cmd/headscale/cli/auth.go

@@ -50,44 +50,45 @@ var authRegisterCmd = &cobra.Command{
 		return printOutput(
 			cmd,
 			response.GetNode(),
-			fmt.Sprintf("Node %s registered", response.GetNode().GetGivenName()))
+			fmt.Sprintf("Node %s registered", response.GetNode().GetGivenName()),
+		)
 	}),
 }
 
+// authDecisionRunE builds a RunE for an auth decision command (approve or
+// reject) that reads the auth-id flag, invokes the given gRPC call, and prints
+// the response. errVerb is used in the error message; okMsg is printed on
+// success.
+func authDecisionRunE[Resp any](
+	errVerb, okMsg string,
+	call func(ctx context.Context, client v1.HeadscaleServiceClient, authID string) (Resp, error),
+) func(*cobra.Command, []string) error {
+	return grpcRunE(func(ctx context.Context, client v1.HeadscaleServiceClient, cmd *cobra.Command, args []string) error {
+		authID, _ := cmd.Flags().GetString("auth-id")
+
+		response, err := call(ctx, client, authID)
+		if err != nil {
+			return fmt.Errorf("%s auth request: %w", errVerb, err)
+		}
+
+		return printOutput(cmd, response, okMsg)
+	})
+}
+
 var authApproveCmd = &cobra.Command{
 	Use:   "approve",
 	Short: "Approve a pending authentication request",
-	RunE: grpcRunE(func(ctx context.Context, client v1.HeadscaleServiceClient, cmd *cobra.Command, args []string) error {
+	RunE: authDecisionRunE("approving", "Auth request approved",
-		authID, _ := cmd.Flags().GetString("auth-id")
+		func(ctx context.Context, client v1.HeadscaleServiceClient, authID string) (*v1.AuthApproveResponse, error) {
-
+			return client.AuthApprove(ctx, &v1.AuthApproveRequest{AuthId: authID})
-		request := &v1.AuthApproveRequest{
+		}),
-			AuthId: authID,
-		}
-
-		response, err := client.AuthApprove(ctx, request)
-		if err != nil {
-			return fmt.Errorf("approving auth request: %w", err)
-		}
-
-		return printOutput(cmd, response, "Auth request approved")
-	}),
 }
 
 var authRejectCmd = &cobra.Command{
 	Use:   "reject",
 	Short: "Reject a pending authentication request",
-	RunE: grpcRunE(func(ctx context.Context, client v1.HeadscaleServiceClient, cmd *cobra.Command, args []string) error {
+	RunE: authDecisionRunE("rejecting", "Auth request rejected",
-		authID, _ := cmd.Flags().GetString("auth-id")
+		func(ctx context.Context, client v1.HeadscaleServiceClient, authID string) (*v1.AuthRejectResponse, error) {
-
+			return client.AuthReject(ctx, &v1.AuthRejectRequest{AuthId: authID})
-		request := &v1.AuthRejectRequest{
+		}),
-			AuthId: authID,
-		}
-
-		response, err := client.AuthReject(ctx, request)
-		if err != nil {
-			return fmt.Errorf("rejecting auth request: %w", err)
-		}
-
-		return printOutput(cmd, response, "Auth request rejected")
-	}),
 }

cmd/headscale/cli/nodes.go

@@ -84,7 +84,8 @@ var registerNodeCmd = &cobra.Command{
 		return printOutput(
 			cmd,
 			response.GetNode(),
-			fmt.Sprintf("Node %s registered", response.GetNode().GetGivenName()))
+			fmt.Sprintf("Node %s registered", response.GetNode().GetGivenName()),
+		)
 	}),
 }
 
@@ -135,7 +136,7 @@ var listNodeRoutesCmd = &cobra.Command{
 		}
 
 		nodes = lo.Filter(nodes, func(n *v1.Node, _ int) bool {
-			return (n.GetSubnetRoutes() != nil && len(n.GetSubnetRoutes()) > 0) || (n.GetApprovedRoutes() != nil && len(n.GetApprovedRoutes()) > 0) || (n.GetAvailableRoutes() != nil && len(n.GetAvailableRoutes()) > 0)
+			return len(n.GetSubnetRoutes()) > 0 || len(n.GetApprovedRoutes()) > 0 || len(n.GetAvailableRoutes()) > 0
 		})
 
 		return printListOutput(cmd, nodes, func() error {

cmd/headscale/cli/policy.go

@@ -1,6 +1,7 @@
 package cli
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"os"
@@ -36,6 +37,16 @@ func bypassDatabase() (*db.HSDatabase, error) {
 	return d, nil
 }
 
+// openBypassDB confirms the destructive bypass action and opens the database
+// directly. The caller is responsible for closing the returned handle.
+func openBypassDB(cmd *cobra.Command) (*db.HSDatabase, error) {
+	if !confirmAction(cmd, "DO NOT run this command if an instance of headscale is running, are you sure headscale is not running?") {
+		return nil, errAborted
+	}
+
+	return bypassDatabase()
+}
+
 func init() {
 	rootCmd.AddCommand(policyCmd)
 
@@ -66,11 +77,7 @@ var getPolicy = &cobra.Command{
 		var policyData string
 
 		if bypass, _ := cmd.Flags().GetBool(bypassFlag); bypass {
-			if !confirmAction(cmd, "DO NOT run this command if an instance of headscale is running, are you sure headscale is not running?") {
+			d, err := openBypassDB(cmd)
-				return errAborted
-			}
-
-			d, err := bypassDatabase()
 			if err != nil {
 				return err
 			}
@@ -83,19 +90,19 @@ var getPolicy = &cobra.Command{
 
 			policyData = pol.Data
 		} else {
-			ctx, client, conn, cancel, err := newHeadscaleCLIWithConfig()
+			err := withGRPC(func(ctx context.Context, client v1.HeadscaleServiceClient) error {
-			if err != nil {
+				response, err := client.GetPolicy(ctx, &v1.GetPolicyRequest{})
-				return fmt.Errorf("connecting to headscale: %w", err)
+				if err != nil {
-			}
+					return fmt.Errorf("loading ACL policy: %w", err)
-			defer cancel()
+				}
-			defer conn.Close()
 
-			response, err := client.GetPolicy(ctx, &v1.GetPolicyRequest{})
+				policyData = response.GetPolicy()
-			if err != nil {
-				return fmt.Errorf("loading ACL policy: %w", err)
-			}
 
-			policyData = response.GetPolicy()
+				return nil
+			})
+			if err != nil {
+				return err
+			}
 		}
 
 		// This does not pass output format as we don't support yaml, json or
@@ -122,11 +129,7 @@ var setPolicy = &cobra.Command{
 		}
 
 		if bypass, _ := cmd.Flags().GetBool(bypassFlag); bypass {
-			if !confirmAction(cmd, "DO NOT run this command if an instance of headscale is running, are you sure headscale is not running?") {
+			d, err := openBypassDB(cmd)
-				return errAborted
-			}
-
-			d, err := bypassDatabase()
 			if err != nil {
 				return err
 			}
@@ -149,16 +152,16 @@ var setPolicy = &cobra.Command{
 		} else {
 			request := &v1.SetPolicyRequest{Policy: string(policyBytes)}
 
-			ctx, client, conn, cancel, err := newHeadscaleCLIWithConfig()
+			err := withGRPC(func(ctx context.Context, client v1.HeadscaleServiceClient) error {
-			if err != nil {
+				_, err := client.SetPolicy(ctx, request)
-				return fmt.Errorf("connecting to headscale: %w", err)
+				if err != nil {
-			}
+					return fmt.Errorf("setting ACL policy: %w", err)
-			defer cancel()
+				}
-			defer conn.Close()
 
-			_, err = client.SetPolicy(ctx, request)
+				return nil
+			})
 			if err != nil {
-				return fmt.Errorf("setting ACL policy: %w", err)
+				return err
 			}
 		}
 
@@ -185,11 +188,7 @@ var checkPolicy = &cobra.Command{
 		}
 
 		if bypass, _ := cmd.Flags().GetBool(bypassFlag); bypass {
-			if !confirmAction(cmd, "DO NOT run this command if an instance of headscale is running, are you sure headscale is not running?") {
+			d, err := openBypassDB(cmd)
-				return errAborted
-			}
-
-			d, err := bypassDatabase()
 			if err != nil {
 				return err
 			}
@@ -224,14 +223,11 @@ var checkPolicy = &cobra.Command{
 			return nil
 		}
 
-		ctx, client, conn, cancel, err := newHeadscaleCLIWithConfig()
+		err = withGRPC(func(ctx context.Context, client v1.HeadscaleServiceClient) error {
-		if err != nil {
+			_, err := client.CheckPolicy(ctx, &v1.CheckPolicyRequest{Policy: string(policyBytes)})
-			return fmt.Errorf("connecting to headscale: %w", err)
-		}
-		defer cancel()
-		defer conn.Close()
 
-		_, err = client.CheckPolicy(ctx, &v1.CheckPolicyRequest{Policy: string(policyBytes)})
+			return err
+		})
 		if err != nil {
 			return err
 		}

cmd/headscale/cli/preauthkeys.go

@@ -128,15 +128,24 @@ var createPreAuthKeyCmd = &cobra.Command{
 	}),
 }
 
+// preAuthKeyID reads the required --id flag for preauthkey commands.
+func preAuthKeyID(cmd *cobra.Command) (uint64, error) {
+	id, _ := cmd.Flags().GetUint64("id")
+	if id == 0 {
+		return 0, fmt.Errorf("missing --id parameter: %w", errMissingParameter)
+	}
+
+	return id, nil
+}
+
 var expirePreAuthKeyCmd = &cobra.Command{
 	Use:     cmdExpire,
 	Short:   "Expire a preauthkey",
 	Aliases: []string{"revoke", aliasExp, "e"},
 	RunE: grpcRunE(func(ctx context.Context, client v1.HeadscaleServiceClient, cmd *cobra.Command, args []string) error {
-		id, _ := cmd.Flags().GetUint64("id")
+		id, err := preAuthKeyID(cmd)
-
+		if err != nil {
-		if id == 0 {
+			return err
-			return fmt.Errorf("missing --id parameter: %w", errMissingParameter)
 		}
 
 		request := &v1.ExpirePreAuthKeyRequest{
@@ -157,10 +166,9 @@ var deletePreAuthKeyCmd = &cobra.Command{
 	Short:   "Delete a preauthkey",
 	Aliases: []string{aliasDel, "rm", "d"},
 	RunE: grpcRunE(func(ctx context.Context, client v1.HeadscaleServiceClient, cmd *cobra.Command, args []string) error {
-		id, _ := cmd.Flags().GetUint64("id")
+		id, err := preAuthKeyID(cmd)
-
+		if err != nil {
-		if id == 0 {
+			return err
-			return fmt.Errorf("missing --id parameter: %w", errMissingParameter)
 		}
 
 		request := &v1.DeletePreAuthKeyRequest{

cmd/headscale/cli/users.go

@@ -44,6 +44,33 @@ func usernameAndIDFromFlag(cmd *cobra.Command) (uint64, string, error) {
 	return uint64(identifier), username, nil //nolint:gosec // identifier is clamped to >= 0 above
 }
 
+// resolveSingleUser resolves exactly one user from the --name/--id flags,
+// returning the raw flag id and the matched user.
+func resolveSingleUser(
+	ctx context.Context,
+	client v1.HeadscaleServiceClient,
+	cmd *cobra.Command,
+) (uint64, *v1.User, error) {
+	id, username, err := usernameAndIDFromFlag(cmd)
+	if err != nil {
+		return 0, nil, err
+	}
+
+	users, err := client.ListUsers(ctx, &v1.ListUsersRequest{
+		Name: username,
+		Id:   id,
+	})
+	if err != nil {
+		return 0, nil, fmt.Errorf("listing users: %w", err)
+	}
+
+	if len(users.GetUsers()) != 1 {
+		return 0, nil, errMultipleUsersMatch
+	}
+
+	return id, users.GetUsers()[0], nil
+}
+
 func init() {
 	rootCmd.AddCommand(userCmd)
 	userCmd.AddCommand(createUserCmd)
@@ -117,27 +144,11 @@ var destroyUserCmd = &cobra.Command{
 	Short:   "Destroys a user",
 	Aliases: []string{cmdDelete},
 	RunE: grpcRunE(func(ctx context.Context, client v1.HeadscaleServiceClient, cmd *cobra.Command, args []string) error {
-		id, username, err := usernameAndIDFromFlag(cmd)
+		_, user, err := resolveSingleUser(ctx, client, cmd)
 		if err != nil {
 			return err
 		}
 
-		request := &v1.ListUsersRequest{
-			Name: username,
-			Id:   id,
-		}
-
-		users, err := client.ListUsers(ctx, request)
-		if err != nil {
-			return fmt.Errorf("listing users: %w", err)
-		}
-
-		if len(users.GetUsers()) != 1 {
-			return errMultipleUsersMatch
-		}
-
-		user := users.GetUsers()[0]
-
 		if !confirmAction(cmd, fmt.Sprintf(
 			"Do you want to remove the user %q (%d) and any associated preauthkeys?",
 			user.GetName(), user.GetId(),
@@ -209,25 +220,11 @@ var renameUserCmd = &cobra.Command{
 	Short:   "Renames a user",
 	Aliases: []string{"mv"},
 	RunE: grpcRunE(func(ctx context.Context, client v1.HeadscaleServiceClient, cmd *cobra.Command, args []string) error {
-		id, username, err := usernameAndIDFromFlag(cmd)
+		id, _, err := resolveSingleUser(ctx, client, cmd)
 		if err != nil {
 			return err
 		}
 
-		listReq := &v1.ListUsersRequest{
-			Name: username,
-			Id:   id,
-		}
-
-		users, err := client.ListUsers(ctx, listReq)
-		if err != nil {
-			return fmt.Errorf("listing users: %w", err)
-		}
-
-		if len(users.GetUsers()) != 1 {
-			return errMultipleUsersMatch
-		}
-
 		newName, _ := cmd.Flags().GetString("new-name")
 
 		renameReq := &v1.RenameUserRequest{

cmd/headscale/cli/utils.go

@@ -85,6 +85,23 @@ func grpcRunE(
 	}
 }
 
+// withGRPC opens a gRPC client, runs fn with it, and tears the
+// connection down afterwards. It is the building block for commands
+// that branch on a flag before deciding to talk to the server, where
+// grpcRunE's whole-RunE wrapping does not fit.
+func withGRPC(
+	fn func(ctx context.Context, client v1.HeadscaleServiceClient) error,
+) error {
+	ctx, client, conn, cancel, err := newHeadscaleCLIWithConfig()
+	if err != nil {
+		return fmt.Errorf("connecting to headscale: %w", err)
+	}
+	defer cancel()
+	defer conn.Close()
+
+	return fn(ctx, client)
+}
+
 func newHeadscaleCLIWithConfig() (context.Context, v1.HeadscaleServiceClient, *grpc.ClientConn, context.CancelFunc, error) {
 	cfg, err := types.LoadCLIConfig()
 	if err != nil {