From 0b7c8aa270978c1996fe8d4a511b9e8359dcf139 Mon Sep 17 00:00:00 2001 From: Kristoffer Dalby Date: Tue, 16 Jun 2026 08:45:59 +0000 Subject: [PATCH] cli: consolidate gRPC command helpers --- cmd/headscale/cli/auth.go | 59 ++++++++++++------------ cmd/headscale/cli/nodes.go | 5 +- cmd/headscale/cli/policy.go | 78 +++++++++++++++----------------- cmd/headscale/cli/preauthkeys.go | 24 ++++++---- cmd/headscale/cli/users.go | 61 ++++++++++++------------- cmd/headscale/cli/utils.go | 17 +++++++ 6 files changed, 132 insertions(+), 112 deletions(-) diff --git a/cmd/headscale/cli/auth.go b/cmd/headscale/cli/auth.go index 8a5476dd..85a17a3c 100644 --- a/cmd/headscale/cli/auth.go +++ b/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 { - authID, _ := cmd.Flags().GetString("auth-id") - - 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") - }), + RunE: authDecisionRunE("approving", "Auth request approved", + func(ctx context.Context, client v1.HeadscaleServiceClient, authID string) (*v1.AuthApproveResponse, error) { + return client.AuthApprove(ctx, &v1.AuthApproveRequest{AuthId: authID}) + }), } 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 { - authID, _ := cmd.Flags().GetString("auth-id") - - 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") - }), + RunE: authDecisionRunE("rejecting", "Auth request rejected", + func(ctx context.Context, client v1.HeadscaleServiceClient, authID string) (*v1.AuthRejectResponse, error) { + return client.AuthReject(ctx, &v1.AuthRejectRequest{AuthId: authID}) + }), } diff --git a/cmd/headscale/cli/nodes.go b/cmd/headscale/cli/nodes.go index 5cf61a43..4bfe1fb5 100644 --- a/cmd/headscale/cli/nodes.go +++ b/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 { diff --git a/cmd/headscale/cli/policy.go b/cmd/headscale/cli/policy.go index d6a1f573..039fedd0 100644 --- a/cmd/headscale/cli/policy.go +++ b/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?") { - return errAborted - } - - d, err := bypassDatabase() + d, err := openBypassDB(cmd) if err != nil { return err } @@ -83,19 +90,19 @@ var getPolicy = &cobra.Command{ policyData = pol.Data } else { - ctx, client, conn, cancel, err := newHeadscaleCLIWithConfig() - if err != nil { - return fmt.Errorf("connecting to headscale: %w", err) - } - defer cancel() - defer conn.Close() + err := withGRPC(func(ctx context.Context, client v1.HeadscaleServiceClient) error { + response, err := client.GetPolicy(ctx, &v1.GetPolicyRequest{}) + if err != nil { + return fmt.Errorf("loading ACL policy: %w", err) + } - response, err := client.GetPolicy(ctx, &v1.GetPolicyRequest{}) - if err != nil { - return fmt.Errorf("loading ACL policy: %w", err) - } + policyData = response.GetPolicy() - 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?") { - return errAborted - } - - d, err := bypassDatabase() + d, err := openBypassDB(cmd) 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() - if err != nil { - return fmt.Errorf("connecting to headscale: %w", err) - } - defer cancel() - defer conn.Close() + err := withGRPC(func(ctx context.Context, client v1.HeadscaleServiceClient) error { + _, err := client.SetPolicy(ctx, request) + if err != nil { + return fmt.Errorf("setting ACL policy: %w", err) + } - _, 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?") { - return errAborted - } - - d, err := bypassDatabase() + d, err := openBypassDB(cmd) if err != nil { return err } @@ -224,14 +223,11 @@ var checkPolicy = &cobra.Command{ return nil } - ctx, client, conn, cancel, err := newHeadscaleCLIWithConfig() - if err != nil { - return fmt.Errorf("connecting to headscale: %w", err) - } - defer cancel() - defer conn.Close() + err = withGRPC(func(ctx context.Context, client v1.HeadscaleServiceClient) error { + _, err := client.CheckPolicy(ctx, &v1.CheckPolicyRequest{Policy: string(policyBytes)}) - _, err = client.CheckPolicy(ctx, &v1.CheckPolicyRequest{Policy: string(policyBytes)}) + return err + }) if err != nil { return err } diff --git a/cmd/headscale/cli/preauthkeys.go b/cmd/headscale/cli/preauthkeys.go index f88df70c..7eb8cf51 100644 --- a/cmd/headscale/cli/preauthkeys.go +++ b/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") - - if id == 0 { - return fmt.Errorf("missing --id parameter: %w", errMissingParameter) + id, err := preAuthKeyID(cmd) + if err != nil { + return err } 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") - - if id == 0 { - return fmt.Errorf("missing --id parameter: %w", errMissingParameter) + id, err := preAuthKeyID(cmd) + if err != nil { + return err } request := &v1.DeletePreAuthKeyRequest{ diff --git a/cmd/headscale/cli/users.go b/cmd/headscale/cli/users.go index bcc8c913..6cac0051 100644 --- a/cmd/headscale/cli/users.go +++ b/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{ diff --git a/cmd/headscale/cli/utils.go b/cmd/headscale/cli/utils.go index 364779ad..002147b3 100644 --- a/cmd/headscale/cli/utils.go +++ b/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 {