From 45aa50ed863e7f326e9744183e1d4f30e28acb64 Mon Sep 17 00:00:00 2001 From: Kristoffer Dalby Date: Wed, 18 Feb 2026 13:18:09 +0000 Subject: [PATCH] cmd/headscale/cli: add grpcRun wrapper for gRPC client lifecycle Add a grpcRun helper that wraps cobra RunFuncs, injecting a ready gRPC client and context. The connection lifecycle (cancel, close) is managed by the wrapper, eliminating the duplicated 3-line boilerplate (newHeadscaleCLIWithConfig + defer cancel + defer conn.Close) from 22 command handlers across 7 files. Three call sites are intentionally left unconverted: - backfillNodeIPsCmd: creates the client only after user confirmation - getPolicy/setPolicy: conditionally use gRPC vs direct DB access --- cmd/headscale/cli/api_key.go | 33 +++++----------- cmd/headscale/cli/debug.go | 9 ++--- cmd/headscale/cli/health.go | 10 ++--- cmd/headscale/cli/nodes.go | 65 +++++++++----------------------- cmd/headscale/cli/preauthkeys.go | 33 +++++----------- cmd/headscale/cli/users.go | 33 +++++----------- cmd/headscale/cli/utils.go | 16 ++++++++ 7 files changed, 67 insertions(+), 132 deletions(-) diff --git a/cmd/headscale/cli/api_key.go b/cmd/headscale/cli/api_key.go index b0720e0a..346ecf55 100644 --- a/cmd/headscale/cli/api_key.go +++ b/cmd/headscale/cli/api_key.go @@ -1,6 +1,7 @@ package cli import ( + "context" "fmt" "strconv" "time" @@ -46,13 +47,9 @@ var listAPIKeys = &cobra.Command{ Use: "list", Short: "List the Api keys for headscale", Aliases: []string{"ls", "show"}, - Run: func(cmd *cobra.Command, args []string) { + Run: grpcRun(func(ctx context.Context, client v1.HeadscaleServiceClient, cmd *cobra.Command, args []string) { output, _ := cmd.Flags().GetString("output") - ctx, client, conn, cancel := newHeadscaleCLIWithConfig() - defer cancel() - defer conn.Close() - request := &v1.ListApiKeysRequest{} response, err := client.ListApiKeys(ctx, request) @@ -95,7 +92,7 @@ var listAPIKeys = &cobra.Command{ output, ) } - }, + }), } var createAPIKeyCmd = &cobra.Command{ @@ -106,7 +103,7 @@ Creates a new Api key, the Api key is only visible on creation and cannot be retrieved again. If you loose a key, create a new one and revoke (expire) the old one.`, Aliases: []string{"c", "new"}, - Run: func(cmd *cobra.Command, args []string) { + Run: grpcRun(func(ctx context.Context, client v1.HeadscaleServiceClient, cmd *cobra.Command, args []string) { output, _ := cmd.Flags().GetString("output") request := &v1.CreateApiKeyRequest{} @@ -126,10 +123,6 @@ If you loose a key, create a new one and revoke (expire) the old one.`, request.Expiration = timestamppb.New(expiration) - ctx, client, conn, cancel := newHeadscaleCLIWithConfig() - defer cancel() - defer conn.Close() - response, err := client.CreateApiKey(ctx, request) if err != nil { ErrorOutput( @@ -140,14 +133,14 @@ If you loose a key, create a new one and revoke (expire) the old one.`, } SuccessOutput(response.GetApiKey(), response.GetApiKey(), output) - }, + }), } var expireAPIKeyCmd = &cobra.Command{ Use: "expire", Short: "Expire an ApiKey", Aliases: []string{"revoke", "exp", "e"}, - Run: func(cmd *cobra.Command, args []string) { + Run: grpcRun(func(ctx context.Context, client v1.HeadscaleServiceClient, cmd *cobra.Command, args []string) { output, _ := cmd.Flags().GetString("output") id, _ := cmd.Flags().GetUint64("id") @@ -168,10 +161,6 @@ var expireAPIKeyCmd = &cobra.Command{ ) } - ctx, client, conn, cancel := newHeadscaleCLIWithConfig() - defer cancel() - defer conn.Close() - request := &v1.ExpireApiKeyRequest{} if id != 0 { request.Id = id @@ -189,14 +178,14 @@ var expireAPIKeyCmd = &cobra.Command{ } SuccessOutput(response, "Key expired", output) - }, + }), } var deleteAPIKeyCmd = &cobra.Command{ Use: "delete", Short: "Delete an ApiKey", Aliases: []string{"remove", "del"}, - Run: func(cmd *cobra.Command, args []string) { + Run: grpcRun(func(ctx context.Context, client v1.HeadscaleServiceClient, cmd *cobra.Command, args []string) { output, _ := cmd.Flags().GetString("output") id, _ := cmd.Flags().GetUint64("id") @@ -217,10 +206,6 @@ var deleteAPIKeyCmd = &cobra.Command{ ) } - ctx, client, conn, cancel := newHeadscaleCLIWithConfig() - defer cancel() - defer conn.Close() - request := &v1.DeleteApiKeyRequest{} if id != 0 { request.Id = id @@ -238,5 +223,5 @@ var deleteAPIKeyCmd = &cobra.Command{ } SuccessOutput(response, "Key deleted", output) - }, + }), } diff --git a/cmd/headscale/cli/debug.go b/cmd/headscale/cli/debug.go index c0a7d3d0..0e03457f 100644 --- a/cmd/headscale/cli/debug.go +++ b/cmd/headscale/cli/debug.go @@ -1,6 +1,7 @@ package cli import ( + "context" "fmt" v1 "github.com/juanfont/headscale/gen/go/headscale/v1" @@ -59,7 +60,7 @@ var debugCmd = &cobra.Command{ var createNodeCmd = &cobra.Command{ Use: "create-node", Short: "Create a node that can be registered with `nodes register <>` command", - Run: func(cmd *cobra.Command, args []string) { + Run: grpcRun(func(ctx context.Context, client v1.HeadscaleServiceClient, cmd *cobra.Command, args []string) { output, _ := cmd.Flags().GetString("output") user, err := cmd.Flags().GetString("user") @@ -67,10 +68,6 @@ var createNodeCmd = &cobra.Command{ ErrorOutput(err, fmt.Sprintf("Error getting user: %s", err), output) } - ctx, client, conn, cancel := newHeadscaleCLIWithConfig() - defer cancel() - defer conn.Close() - name, err := cmd.Flags().GetString("name") if err != nil { ErrorOutput( @@ -124,5 +121,5 @@ var createNodeCmd = &cobra.Command{ } SuccessOutput(response.GetNode(), "Node created", output) - }, + }), } diff --git a/cmd/headscale/cli/health.go b/cmd/headscale/cli/health.go index a0b7b91b..46769c00 100644 --- a/cmd/headscale/cli/health.go +++ b/cmd/headscale/cli/health.go @@ -1,6 +1,8 @@ package cli import ( + "context" + v1 "github.com/juanfont/headscale/gen/go/headscale/v1" "github.com/spf13/cobra" ) @@ -13,18 +15,14 @@ var healthCmd = &cobra.Command{ Use: "health", Short: "Check the health of the Headscale server", Long: "Check the health of the Headscale server. This command will return an exit code of 0 if the server is healthy, or 1 if it is not.", - Run: func(cmd *cobra.Command, args []string) { + Run: grpcRun(func(ctx context.Context, client v1.HeadscaleServiceClient, cmd *cobra.Command, args []string) { output, _ := cmd.Flags().GetString("output") - ctx, client, conn, cancel := newHeadscaleCLIWithConfig() - defer cancel() - defer conn.Close() - response, err := client.Health(ctx, &v1.HealthRequest{}) if err != nil { ErrorOutput(err, "Error checking health", output) } SuccessOutput(response, "", output) - }, + }), } diff --git a/cmd/headscale/cli/nodes.go b/cmd/headscale/cli/nodes.go index 642dc1c6..624b1c29 100644 --- a/cmd/headscale/cli/nodes.go +++ b/cmd/headscale/cli/nodes.go @@ -1,6 +1,7 @@ package cli import ( + "context" "fmt" "log" "net/netip" @@ -103,7 +104,7 @@ var nodeCmd = &cobra.Command{ var registerNodeCmd = &cobra.Command{ Use: "register", Short: "Registers a node to your network", - Run: func(cmd *cobra.Command, args []string) { + Run: grpcRun(func(ctx context.Context, client v1.HeadscaleServiceClient, cmd *cobra.Command, args []string) { output, _ := cmd.Flags().GetString("output") user, err := cmd.Flags().GetString("user") @@ -111,10 +112,6 @@ var registerNodeCmd = &cobra.Command{ ErrorOutput(err, fmt.Sprintf("Error getting user: %s", err), output) } - ctx, client, conn, cancel := newHeadscaleCLIWithConfig() - defer cancel() - defer conn.Close() - registrationID, err := cmd.Flags().GetString("key") if err != nil { ErrorOutput( @@ -144,14 +141,14 @@ var registerNodeCmd = &cobra.Command{ SuccessOutput( response.GetNode(), fmt.Sprintf("Node %s registered", response.GetNode().GetGivenName()), output) - }, + }), } var listNodesCmd = &cobra.Command{ Use: "list", Short: "List nodes", Aliases: []string{"ls", "show"}, - Run: func(cmd *cobra.Command, args []string) { + Run: grpcRun(func(ctx context.Context, client v1.HeadscaleServiceClient, cmd *cobra.Command, args []string) { output, _ := cmd.Flags().GetString("output") user, err := cmd.Flags().GetString("user") @@ -159,10 +156,6 @@ var listNodesCmd = &cobra.Command{ ErrorOutput(err, fmt.Sprintf("Error getting user: %s", err), output) } - ctx, client, conn, cancel := newHeadscaleCLIWithConfig() - defer cancel() - defer conn.Close() - request := &v1.ListNodesRequest{ User: user, } @@ -193,14 +186,14 @@ var listNodesCmd = &cobra.Command{ output, ) } - }, + }), } var listNodeRoutesCmd = &cobra.Command{ Use: "list-routes", Short: "List routes available on nodes", Aliases: []string{"lsr", "routes"}, - Run: func(cmd *cobra.Command, args []string) { + Run: grpcRun(func(ctx context.Context, client v1.HeadscaleServiceClient, cmd *cobra.Command, args []string) { output, _ := cmd.Flags().GetString("output") identifier, err := cmd.Flags().GetUint64("identifier") @@ -212,10 +205,6 @@ var listNodeRoutesCmd = &cobra.Command{ ) } - ctx, client, conn, cancel := newHeadscaleCLIWithConfig() - defer cancel() - defer conn.Close() - request := &v1.ListNodesRequest{} response, err := client.ListNodes(ctx, request) @@ -256,7 +245,7 @@ var listNodeRoutesCmd = &cobra.Command{ output, ) } - }, + }), } var expireNodeCmd = &cobra.Command{ @@ -264,7 +253,7 @@ var expireNodeCmd = &cobra.Command{ Short: "Expire (log out) a node in your network", Long: "Expiring a node will keep the node in the database and force it to reauthenticate.", Aliases: []string{"logout", "exp", "e"}, - Run: func(cmd *cobra.Command, args []string) { + Run: grpcRun(func(ctx context.Context, client v1.HeadscaleServiceClient, cmd *cobra.Command, args []string) { output, _ := cmd.Flags().GetString("output") identifier, err := cmd.Flags().GetUint64("identifier") @@ -303,10 +292,6 @@ var expireNodeCmd = &cobra.Command{ } } - ctx, client, conn, cancel := newHeadscaleCLIWithConfig() - defer cancel() - defer conn.Close() - request := &v1.ExpireNodeRequest{ NodeId: identifier, Expiry: timestamppb.New(expiryTime), @@ -329,13 +314,13 @@ var expireNodeCmd = &cobra.Command{ } else { SuccessOutput(response.GetNode(), "Node expiration updated", output) } - }, + }), } var renameNodeCmd = &cobra.Command{ Use: "rename NEW_NAME", Short: "Renames a node in your network", - Run: func(cmd *cobra.Command, args []string) { + Run: grpcRun(func(ctx context.Context, client v1.HeadscaleServiceClient, cmd *cobra.Command, args []string) { output, _ := cmd.Flags().GetString("output") identifier, err := cmd.Flags().GetUint64("identifier") @@ -347,10 +332,6 @@ var renameNodeCmd = &cobra.Command{ ) } - ctx, client, conn, cancel := newHeadscaleCLIWithConfig() - defer cancel() - defer conn.Close() - newName := "" if len(args) > 0 { newName = args[0] @@ -374,14 +355,14 @@ var renameNodeCmd = &cobra.Command{ } SuccessOutput(response.GetNode(), "Node renamed", output) - }, + }), } var deleteNodeCmd = &cobra.Command{ Use: "delete", Short: "Delete a node", Aliases: []string{"del"}, - Run: func(cmd *cobra.Command, args []string) { + Run: grpcRun(func(ctx context.Context, client v1.HeadscaleServiceClient, cmd *cobra.Command, args []string) { output, _ := cmd.Flags().GetString("output") identifier, err := cmd.Flags().GetUint64("identifier") @@ -393,10 +374,6 @@ var deleteNodeCmd = &cobra.Command{ ) } - ctx, client, conn, cancel := newHeadscaleCLIWithConfig() - defer cancel() - defer conn.Close() - getRequest := &v1.GetNodeRequest{ NodeId: identifier, } @@ -448,7 +425,7 @@ var deleteNodeCmd = &cobra.Command{ } else { SuccessOutput(map[string]string{"Result": "Node not deleted"}, "Node not deleted", output) } - }, + }), } var backfillNodeIPsCmd = &cobra.Command{ @@ -666,13 +643,9 @@ var tagCmd = &cobra.Command{ Use: "tag", Short: "Manage the tags of a node", Aliases: []string{"tags", "t"}, - Run: func(cmd *cobra.Command, args []string) { + Run: grpcRun(func(ctx context.Context, client v1.HeadscaleServiceClient, cmd *cobra.Command, args []string) { output, _ := cmd.Flags().GetString("output") - ctx, client, conn, cancel := newHeadscaleCLIWithConfig() - defer cancel() - defer conn.Close() - // retrieve flags from CLI identifier, err := cmd.Flags().GetUint64("identifier") if err != nil { @@ -714,19 +687,15 @@ var tagCmd = &cobra.Command{ output, ) } - }, + }), } var approveRoutesCmd = &cobra.Command{ Use: "approve-routes", Short: "Manage the approved routes of a node", - Run: func(cmd *cobra.Command, args []string) { + Run: grpcRun(func(ctx context.Context, client v1.HeadscaleServiceClient, cmd *cobra.Command, args []string) { output, _ := cmd.Flags().GetString("output") - ctx, client, conn, cancel := newHeadscaleCLIWithConfig() - defer cancel() - defer conn.Close() - // retrieve flags from CLI identifier, err := cmd.Flags().GetUint64("identifier") if err != nil { @@ -768,5 +737,5 @@ var approveRoutesCmd = &cobra.Command{ output, ) } - }, + }), } diff --git a/cmd/headscale/cli/preauthkeys.go b/cmd/headscale/cli/preauthkeys.go index 84e6a7f1..1d04693f 100644 --- a/cmd/headscale/cli/preauthkeys.go +++ b/cmd/headscale/cli/preauthkeys.go @@ -1,6 +1,7 @@ package cli import ( + "context" "fmt" "strconv" "strings" @@ -46,13 +47,9 @@ var listPreAuthKeys = &cobra.Command{ Use: "list", Short: "List all preauthkeys", Aliases: []string{"ls", "show"}, - Run: func(cmd *cobra.Command, args []string) { + Run: grpcRun(func(ctx context.Context, client v1.HeadscaleServiceClient, cmd *cobra.Command, args []string) { output, _ := cmd.Flags().GetString("output") - ctx, client, conn, cancel := newHeadscaleCLIWithConfig() - defer cancel() - defer conn.Close() - response, err := client.ListPreAuthKeys(ctx, &v1.ListPreAuthKeysRequest{}) if err != nil { ErrorOutput( @@ -116,14 +113,14 @@ var listPreAuthKeys = &cobra.Command{ output, ) } - }, + }), } var createPreAuthKeyCmd = &cobra.Command{ Use: "create", Short: "Creates a new preauthkey", Aliases: []string{"c", "new"}, - Run: func(cmd *cobra.Command, args []string) { + Run: grpcRun(func(ctx context.Context, client v1.HeadscaleServiceClient, cmd *cobra.Command, args []string) { output, _ := cmd.Flags().GetString("output") user, _ := cmd.Flags().GetUint64("user") @@ -153,10 +150,6 @@ var createPreAuthKeyCmd = &cobra.Command{ request.Expiration = timestamppb.New(expiration) - ctx, client, conn, cancel := newHeadscaleCLIWithConfig() - defer cancel() - defer conn.Close() - response, err := client.CreatePreAuthKey(ctx, request) if err != nil { ErrorOutput( @@ -167,14 +160,14 @@ var createPreAuthKeyCmd = &cobra.Command{ } SuccessOutput(response.GetPreAuthKey(), response.GetPreAuthKey().GetKey(), output) - }, + }), } var expirePreAuthKeyCmd = &cobra.Command{ Use: "expire", Short: "Expire a preauthkey", Aliases: []string{"revoke", "exp", "e"}, - Run: func(cmd *cobra.Command, args []string) { + Run: grpcRun(func(ctx context.Context, client v1.HeadscaleServiceClient, cmd *cobra.Command, args []string) { output, _ := cmd.Flags().GetString("output") id, _ := cmd.Flags().GetUint64("id") @@ -188,10 +181,6 @@ var expirePreAuthKeyCmd = &cobra.Command{ return } - ctx, client, conn, cancel := newHeadscaleCLIWithConfig() - defer cancel() - defer conn.Close() - request := &v1.ExpirePreAuthKeyRequest{ Id: id, } @@ -206,14 +195,14 @@ var expirePreAuthKeyCmd = &cobra.Command{ } SuccessOutput(response, "Key expired", output) - }, + }), } var deletePreAuthKeyCmd = &cobra.Command{ Use: "delete", Short: "Delete a preauthkey", Aliases: []string{"del", "rm", "d"}, - Run: func(cmd *cobra.Command, args []string) { + Run: grpcRun(func(ctx context.Context, client v1.HeadscaleServiceClient, cmd *cobra.Command, args []string) { output, _ := cmd.Flags().GetString("output") id, _ := cmd.Flags().GetUint64("id") @@ -227,10 +216,6 @@ var deletePreAuthKeyCmd = &cobra.Command{ return } - ctx, client, conn, cancel := newHeadscaleCLIWithConfig() - defer cancel() - defer conn.Close() - request := &v1.DeletePreAuthKeyRequest{ Id: id, } @@ -245,5 +230,5 @@ var deletePreAuthKeyCmd = &cobra.Command{ } SuccessOutput(response, "Key deleted", output) - }, + }), } diff --git a/cmd/headscale/cli/users.go b/cmd/headscale/cli/users.go index 6cd4417e..b5f28532 100644 --- a/cmd/headscale/cli/users.go +++ b/cmd/headscale/cli/users.go @@ -1,6 +1,7 @@ package cli import ( + "context" "errors" "fmt" "net/url" @@ -80,15 +81,11 @@ var createUserCmd = &cobra.Command{ return nil }, - Run: func(cmd *cobra.Command, args []string) { + Run: grpcRun(func(ctx context.Context, client v1.HeadscaleServiceClient, cmd *cobra.Command, args []string) { output, _ := cmd.Flags().GetString("output") userName := args[0] - ctx, client, conn, cancel := newHeadscaleCLIWithConfig() - defer cancel() - defer conn.Close() - log.Trace().Interface(zf.Client, client).Msg("obtained gRPC client") request := &v1.CreateUserRequest{Name: userName} @@ -128,14 +125,14 @@ var createUserCmd = &cobra.Command{ } SuccessOutput(response.GetUser(), "User created", output) - }, + }), } var destroyUserCmd = &cobra.Command{ Use: "destroy --identifier ID or --name NAME", Short: "Destroys a user", Aliases: []string{"delete"}, - Run: func(cmd *cobra.Command, args []string) { + Run: grpcRun(func(ctx context.Context, client v1.HeadscaleServiceClient, cmd *cobra.Command, args []string) { output, _ := cmd.Flags().GetString("output") id, username := usernameAndIDFromFlag(cmd) @@ -144,10 +141,6 @@ var destroyUserCmd = &cobra.Command{ Id: id, } - ctx, client, conn, cancel := newHeadscaleCLIWithConfig() - defer cancel() - defer conn.Close() - users, err := client.ListUsers(ctx, request) if err != nil { ErrorOutput( @@ -194,20 +187,16 @@ var destroyUserCmd = &cobra.Command{ } else { SuccessOutput(map[string]string{"Result": "User not destroyed"}, "User not destroyed", output) } - }, + }), } var listUsersCmd = &cobra.Command{ Use: "list", Short: "List all the users", Aliases: []string{"ls", "show"}, - Run: func(cmd *cobra.Command, args []string) { + Run: grpcRun(func(ctx context.Context, client v1.HeadscaleServiceClient, cmd *cobra.Command, args []string) { output, _ := cmd.Flags().GetString("output") - ctx, client, conn, cancel := newHeadscaleCLIWithConfig() - defer cancel() - defer conn.Close() - request := &v1.ListUsersRequest{} id, _ := cmd.Flags().GetInt64("identifier") @@ -259,20 +248,16 @@ var listUsersCmd = &cobra.Command{ output, ) } - }, + }), } var renameUserCmd = &cobra.Command{ Use: "rename", Short: "Renames a user", Aliases: []string{"mv"}, - Run: func(cmd *cobra.Command, args []string) { + Run: grpcRun(func(ctx context.Context, client v1.HeadscaleServiceClient, cmd *cobra.Command, args []string) { output, _ := cmd.Flags().GetString("output") - ctx, client, conn, cancel := newHeadscaleCLIWithConfig() - defer cancel() - defer conn.Close() - id, username := usernameAndIDFromFlag(cmd) listReq := &v1.ListUsersRequest{ Name: username, @@ -314,5 +299,5 @@ var renameUserCmd = &cobra.Command{ } SuccessOutput(response.GetUser(), "User renamed", output) - }, + }), } diff --git a/cmd/headscale/cli/utils.go b/cmd/headscale/cli/utils.go index 2cfbc466..31087569 100644 --- a/cmd/headscale/cli/utils.go +++ b/cmd/headscale/cli/utils.go @@ -13,6 +13,7 @@ import ( "github.com/juanfont/headscale/hscontrol/util" "github.com/juanfont/headscale/hscontrol/util/zlog/zf" "github.com/rs/zerolog/log" + "github.com/spf13/cobra" "google.golang.org/grpc" "google.golang.org/grpc/credentials" "google.golang.org/grpc/credentials/insecure" @@ -41,6 +42,21 @@ func newHeadscaleServerWithConfig() (*hscontrol.Headscale, error) { return app, nil } +// grpcRun wraps a cobra RunFunc, injecting a ready gRPC client and context. +// Connection lifecycle is managed by the wrapper — callers never see +// the underlying conn or cancel func. +func grpcRun( + fn func(ctx context.Context, client v1.HeadscaleServiceClient, cmd *cobra.Command, args []string), +) func(*cobra.Command, []string) { + return func(cmd *cobra.Command, args []string) { + ctx, client, conn, cancel := newHeadscaleCLIWithConfig() + defer cancel() + defer conn.Close() + + fn(ctx, client, cmd, args) + } +} + func newHeadscaleCLIWithConfig() (context.Context, v1.HeadscaleServiceClient, *grpc.ClientConn, context.CancelFunc) { cfg, err := types.LoadCLIConfig() if err != nil {