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
This commit is contained in:
Kristoffer Dalby 2026-02-18 13:18:09 +00:00
parent cfb308b4a7
commit aae2f7de71
7 changed files with 67 additions and 132 deletions

View File

@ -1,6 +1,7 @@
package cli package cli
import ( import (
"context"
"fmt" "fmt"
"strconv" "strconv"
"time" "time"
@ -46,13 +47,9 @@ var listAPIKeys = &cobra.Command{
Use: "list", Use: "list",
Short: "List the Api keys for headscale", Short: "List the Api keys for headscale",
Aliases: []string{"ls", "show"}, 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") output, _ := cmd.Flags().GetString("output")
ctx, client, conn, cancel := newHeadscaleCLIWithConfig()
defer cancel()
defer conn.Close()
request := &v1.ListApiKeysRequest{} request := &v1.ListApiKeysRequest{}
response, err := client.ListApiKeys(ctx, request) response, err := client.ListApiKeys(ctx, request)
@ -95,7 +92,7 @@ var listAPIKeys = &cobra.Command{
output, output,
) )
} }
}, }),
} }
var createAPIKeyCmd = &cobra.Command{ 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. and cannot be retrieved again.
If you loose a key, create a new one and revoke (expire) the old one.`, If you loose a key, create a new one and revoke (expire) the old one.`,
Aliases: []string{"c", "new"}, 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") output, _ := cmd.Flags().GetString("output")
request := &v1.CreateApiKeyRequest{} 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) request.Expiration = timestamppb.New(expiration)
ctx, client, conn, cancel := newHeadscaleCLIWithConfig()
defer cancel()
defer conn.Close()
response, err := client.CreateApiKey(ctx, request) response, err := client.CreateApiKey(ctx, request)
if err != nil { if err != nil {
ErrorOutput( 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) SuccessOutput(response.GetApiKey(), response.GetApiKey(), output)
}, }),
} }
var expireAPIKeyCmd = &cobra.Command{ var expireAPIKeyCmd = &cobra.Command{
Use: "expire", Use: "expire",
Short: "Expire an ApiKey", Short: "Expire an ApiKey",
Aliases: []string{"revoke", "exp", "e"}, 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") output, _ := cmd.Flags().GetString("output")
id, _ := cmd.Flags().GetUint64("id") 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{} request := &v1.ExpireApiKeyRequest{}
if id != 0 { if id != 0 {
request.Id = id request.Id = id
@ -189,14 +178,14 @@ var expireAPIKeyCmd = &cobra.Command{
} }
SuccessOutput(response, "Key expired", output) SuccessOutput(response, "Key expired", output)
}, }),
} }
var deleteAPIKeyCmd = &cobra.Command{ var deleteAPIKeyCmd = &cobra.Command{
Use: "delete", Use: "delete",
Short: "Delete an ApiKey", Short: "Delete an ApiKey",
Aliases: []string{"remove", "del"}, 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") output, _ := cmd.Flags().GetString("output")
id, _ := cmd.Flags().GetUint64("id") 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{} request := &v1.DeleteApiKeyRequest{}
if id != 0 { if id != 0 {
request.Id = id request.Id = id
@ -238,5 +223,5 @@ var deleteAPIKeyCmd = &cobra.Command{
} }
SuccessOutput(response, "Key deleted", output) SuccessOutput(response, "Key deleted", output)
}, }),
} }

View File

@ -1,6 +1,7 @@
package cli package cli
import ( import (
"context"
"fmt" "fmt"
v1 "github.com/juanfont/headscale/gen/go/headscale/v1" v1 "github.com/juanfont/headscale/gen/go/headscale/v1"
@ -59,7 +60,7 @@ var debugCmd = &cobra.Command{
var createNodeCmd = &cobra.Command{ var createNodeCmd = &cobra.Command{
Use: "create-node", Use: "create-node",
Short: "Create a node that can be registered with `nodes register <>` command", 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") output, _ := cmd.Flags().GetString("output")
user, err := cmd.Flags().GetString("user") user, err := cmd.Flags().GetString("user")
@ -67,10 +68,6 @@ var createNodeCmd = &cobra.Command{
ErrorOutput(err, fmt.Sprintf("Error getting user: %s", err), output) 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") name, err := cmd.Flags().GetString("name")
if err != nil { if err != nil {
ErrorOutput( ErrorOutput(
@ -124,5 +121,5 @@ var createNodeCmd = &cobra.Command{
} }
SuccessOutput(response.GetNode(), "Node created", output) SuccessOutput(response.GetNode(), "Node created", output)
}, }),
} }

View File

@ -1,6 +1,8 @@
package cli package cli
import ( import (
"context"
v1 "github.com/juanfont/headscale/gen/go/headscale/v1" v1 "github.com/juanfont/headscale/gen/go/headscale/v1"
"github.com/spf13/cobra" "github.com/spf13/cobra"
) )
@ -13,18 +15,14 @@ var healthCmd = &cobra.Command{
Use: "health", Use: "health",
Short: "Check the health of the Headscale server", 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.", 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") output, _ := cmd.Flags().GetString("output")
ctx, client, conn, cancel := newHeadscaleCLIWithConfig()
defer cancel()
defer conn.Close()
response, err := client.Health(ctx, &v1.HealthRequest{}) response, err := client.Health(ctx, &v1.HealthRequest{})
if err != nil { if err != nil {
ErrorOutput(err, "Error checking health", output) ErrorOutput(err, "Error checking health", output)
} }
SuccessOutput(response, "", output) SuccessOutput(response, "", output)
}, }),
} }

View File

@ -1,6 +1,7 @@
package cli package cli
import ( import (
"context"
"fmt" "fmt"
"log" "log"
"net/netip" "net/netip"
@ -103,7 +104,7 @@ var nodeCmd = &cobra.Command{
var registerNodeCmd = &cobra.Command{ var registerNodeCmd = &cobra.Command{
Use: "register", Use: "register",
Short: "Registers a node to your network", 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") output, _ := cmd.Flags().GetString("output")
user, err := cmd.Flags().GetString("user") user, err := cmd.Flags().GetString("user")
@ -111,10 +112,6 @@ var registerNodeCmd = &cobra.Command{
ErrorOutput(err, fmt.Sprintf("Error getting user: %s", err), output) 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") registrationID, err := cmd.Flags().GetString("key")
if err != nil { if err != nil {
ErrorOutput( ErrorOutput(
@ -144,14 +141,14 @@ var registerNodeCmd = &cobra.Command{
SuccessOutput( SuccessOutput(
response.GetNode(), response.GetNode(),
fmt.Sprintf("Node %s registered", response.GetNode().GetGivenName()), output) fmt.Sprintf("Node %s registered", response.GetNode().GetGivenName()), output)
}, }),
} }
var listNodesCmd = &cobra.Command{ var listNodesCmd = &cobra.Command{
Use: "list", Use: "list",
Short: "List nodes", Short: "List nodes",
Aliases: []string{"ls", "show"}, 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") output, _ := cmd.Flags().GetString("output")
user, err := cmd.Flags().GetString("user") user, err := cmd.Flags().GetString("user")
@ -159,10 +156,6 @@ var listNodesCmd = &cobra.Command{
ErrorOutput(err, fmt.Sprintf("Error getting user: %s", err), output) ErrorOutput(err, fmt.Sprintf("Error getting user: %s", err), output)
} }
ctx, client, conn, cancel := newHeadscaleCLIWithConfig()
defer cancel()
defer conn.Close()
request := &v1.ListNodesRequest{ request := &v1.ListNodesRequest{
User: user, User: user,
} }
@ -193,14 +186,14 @@ var listNodesCmd = &cobra.Command{
output, output,
) )
} }
}, }),
} }
var listNodeRoutesCmd = &cobra.Command{ var listNodeRoutesCmd = &cobra.Command{
Use: "list-routes", Use: "list-routes",
Short: "List routes available on nodes", Short: "List routes available on nodes",
Aliases: []string{"lsr", "routes"}, 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") output, _ := cmd.Flags().GetString("output")
identifier, err := cmd.Flags().GetUint64("identifier") 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{} request := &v1.ListNodesRequest{}
response, err := client.ListNodes(ctx, request) response, err := client.ListNodes(ctx, request)
@ -256,7 +245,7 @@ var listNodeRoutesCmd = &cobra.Command{
output, output,
) )
} }
}, }),
} }
var expireNodeCmd = &cobra.Command{ var expireNodeCmd = &cobra.Command{
@ -264,7 +253,7 @@ var expireNodeCmd = &cobra.Command{
Short: "Expire (log out) a node in your network", 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.", Long: "Expiring a node will keep the node in the database and force it to reauthenticate.",
Aliases: []string{"logout", "exp", "e"}, 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") output, _ := cmd.Flags().GetString("output")
identifier, err := cmd.Flags().GetUint64("identifier") 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{ request := &v1.ExpireNodeRequest{
NodeId: identifier, NodeId: identifier,
Expiry: timestamppb.New(expiryTime), Expiry: timestamppb.New(expiryTime),
@ -329,13 +314,13 @@ var expireNodeCmd = &cobra.Command{
} else { } else {
SuccessOutput(response.GetNode(), "Node expiration updated", output) SuccessOutput(response.GetNode(), "Node expiration updated", output)
} }
}, }),
} }
var renameNodeCmd = &cobra.Command{ var renameNodeCmd = &cobra.Command{
Use: "rename NEW_NAME", Use: "rename NEW_NAME",
Short: "Renames a node in your network", 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") output, _ := cmd.Flags().GetString("output")
identifier, err := cmd.Flags().GetUint64("identifier") 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 := "" newName := ""
if len(args) > 0 { if len(args) > 0 {
newName = args[0] newName = args[0]
@ -374,14 +355,14 @@ var renameNodeCmd = &cobra.Command{
} }
SuccessOutput(response.GetNode(), "Node renamed", output) SuccessOutput(response.GetNode(), "Node renamed", output)
}, }),
} }
var deleteNodeCmd = &cobra.Command{ var deleteNodeCmd = &cobra.Command{
Use: "delete", Use: "delete",
Short: "Delete a node", Short: "Delete a node",
Aliases: []string{"del"}, 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") output, _ := cmd.Flags().GetString("output")
identifier, err := cmd.Flags().GetUint64("identifier") 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{ getRequest := &v1.GetNodeRequest{
NodeId: identifier, NodeId: identifier,
} }
@ -448,7 +425,7 @@ var deleteNodeCmd = &cobra.Command{
} else { } else {
SuccessOutput(map[string]string{"Result": "Node not deleted"}, "Node not deleted", output) SuccessOutput(map[string]string{"Result": "Node not deleted"}, "Node not deleted", output)
} }
}, }),
} }
var backfillNodeIPsCmd = &cobra.Command{ var backfillNodeIPsCmd = &cobra.Command{
@ -666,13 +643,9 @@ var tagCmd = &cobra.Command{
Use: "tag", Use: "tag",
Short: "Manage the tags of a node", Short: "Manage the tags of a node",
Aliases: []string{"tags", "t"}, 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") output, _ := cmd.Flags().GetString("output")
ctx, client, conn, cancel := newHeadscaleCLIWithConfig()
defer cancel()
defer conn.Close()
// retrieve flags from CLI // retrieve flags from CLI
identifier, err := cmd.Flags().GetUint64("identifier") identifier, err := cmd.Flags().GetUint64("identifier")
if err != nil { if err != nil {
@ -714,19 +687,15 @@ var tagCmd = &cobra.Command{
output, output,
) )
} }
}, }),
} }
var approveRoutesCmd = &cobra.Command{ var approveRoutesCmd = &cobra.Command{
Use: "approve-routes", Use: "approve-routes",
Short: "Manage the approved routes of a node", 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") output, _ := cmd.Flags().GetString("output")
ctx, client, conn, cancel := newHeadscaleCLIWithConfig()
defer cancel()
defer conn.Close()
// retrieve flags from CLI // retrieve flags from CLI
identifier, err := cmd.Flags().GetUint64("identifier") identifier, err := cmd.Flags().GetUint64("identifier")
if err != nil { if err != nil {
@ -768,5 +737,5 @@ var approveRoutesCmd = &cobra.Command{
output, output,
) )
} }
}, }),
} }

View File

@ -1,6 +1,7 @@
package cli package cli
import ( import (
"context"
"fmt" "fmt"
"strconv" "strconv"
"strings" "strings"
@ -46,13 +47,9 @@ var listPreAuthKeys = &cobra.Command{
Use: "list", Use: "list",
Short: "List all preauthkeys", Short: "List all preauthkeys",
Aliases: []string{"ls", "show"}, 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") output, _ := cmd.Flags().GetString("output")
ctx, client, conn, cancel := newHeadscaleCLIWithConfig()
defer cancel()
defer conn.Close()
response, err := client.ListPreAuthKeys(ctx, &v1.ListPreAuthKeysRequest{}) response, err := client.ListPreAuthKeys(ctx, &v1.ListPreAuthKeysRequest{})
if err != nil { if err != nil {
ErrorOutput( ErrorOutput(
@ -116,14 +113,14 @@ var listPreAuthKeys = &cobra.Command{
output, output,
) )
} }
}, }),
} }
var createPreAuthKeyCmd = &cobra.Command{ var createPreAuthKeyCmd = &cobra.Command{
Use: "create", Use: "create",
Short: "Creates a new preauthkey", Short: "Creates a new preauthkey",
Aliases: []string{"c", "new"}, 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") output, _ := cmd.Flags().GetString("output")
user, _ := cmd.Flags().GetUint64("user") user, _ := cmd.Flags().GetUint64("user")
@ -153,10 +150,6 @@ var createPreAuthKeyCmd = &cobra.Command{
request.Expiration = timestamppb.New(expiration) request.Expiration = timestamppb.New(expiration)
ctx, client, conn, cancel := newHeadscaleCLIWithConfig()
defer cancel()
defer conn.Close()
response, err := client.CreatePreAuthKey(ctx, request) response, err := client.CreatePreAuthKey(ctx, request)
if err != nil { if err != nil {
ErrorOutput( ErrorOutput(
@ -167,14 +160,14 @@ var createPreAuthKeyCmd = &cobra.Command{
} }
SuccessOutput(response.GetPreAuthKey(), response.GetPreAuthKey().GetKey(), output) SuccessOutput(response.GetPreAuthKey(), response.GetPreAuthKey().GetKey(), output)
}, }),
} }
var expirePreAuthKeyCmd = &cobra.Command{ var expirePreAuthKeyCmd = &cobra.Command{
Use: "expire", Use: "expire",
Short: "Expire a preauthkey", Short: "Expire a preauthkey",
Aliases: []string{"revoke", "exp", "e"}, 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") output, _ := cmd.Flags().GetString("output")
id, _ := cmd.Flags().GetUint64("id") id, _ := cmd.Flags().GetUint64("id")
@ -188,10 +181,6 @@ var expirePreAuthKeyCmd = &cobra.Command{
return return
} }
ctx, client, conn, cancel := newHeadscaleCLIWithConfig()
defer cancel()
defer conn.Close()
request := &v1.ExpirePreAuthKeyRequest{ request := &v1.ExpirePreAuthKeyRequest{
Id: id, Id: id,
} }
@ -206,14 +195,14 @@ var expirePreAuthKeyCmd = &cobra.Command{
} }
SuccessOutput(response, "Key expired", output) SuccessOutput(response, "Key expired", output)
}, }),
} }
var deletePreAuthKeyCmd = &cobra.Command{ var deletePreAuthKeyCmd = &cobra.Command{
Use: "delete", Use: "delete",
Short: "Delete a preauthkey", Short: "Delete a preauthkey",
Aliases: []string{"del", "rm", "d"}, 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") output, _ := cmd.Flags().GetString("output")
id, _ := cmd.Flags().GetUint64("id") id, _ := cmd.Flags().GetUint64("id")
@ -227,10 +216,6 @@ var deletePreAuthKeyCmd = &cobra.Command{
return return
} }
ctx, client, conn, cancel := newHeadscaleCLIWithConfig()
defer cancel()
defer conn.Close()
request := &v1.DeletePreAuthKeyRequest{ request := &v1.DeletePreAuthKeyRequest{
Id: id, Id: id,
} }
@ -245,5 +230,5 @@ var deletePreAuthKeyCmd = &cobra.Command{
} }
SuccessOutput(response, "Key deleted", output) SuccessOutput(response, "Key deleted", output)
}, }),
} }

View File

@ -1,6 +1,7 @@
package cli package cli
import ( import (
"context"
"errors" "errors"
"fmt" "fmt"
"net/url" "net/url"
@ -80,15 +81,11 @@ var createUserCmd = &cobra.Command{
return nil 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") output, _ := cmd.Flags().GetString("output")
userName := args[0] userName := args[0]
ctx, client, conn, cancel := newHeadscaleCLIWithConfig()
defer cancel()
defer conn.Close()
log.Trace().Interface(zf.Client, client).Msg("obtained gRPC client") log.Trace().Interface(zf.Client, client).Msg("obtained gRPC client")
request := &v1.CreateUserRequest{Name: userName} request := &v1.CreateUserRequest{Name: userName}
@ -128,14 +125,14 @@ var createUserCmd = &cobra.Command{
} }
SuccessOutput(response.GetUser(), "User created", output) SuccessOutput(response.GetUser(), "User created", output)
}, }),
} }
var destroyUserCmd = &cobra.Command{ var destroyUserCmd = &cobra.Command{
Use: "destroy --identifier ID or --name NAME", Use: "destroy --identifier ID or --name NAME",
Short: "Destroys a user", Short: "Destroys a user",
Aliases: []string{"delete"}, 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") output, _ := cmd.Flags().GetString("output")
id, username := usernameAndIDFromFlag(cmd) id, username := usernameAndIDFromFlag(cmd)
@ -144,10 +141,6 @@ var destroyUserCmd = &cobra.Command{
Id: id, Id: id,
} }
ctx, client, conn, cancel := newHeadscaleCLIWithConfig()
defer cancel()
defer conn.Close()
users, err := client.ListUsers(ctx, request) users, err := client.ListUsers(ctx, request)
if err != nil { if err != nil {
ErrorOutput( ErrorOutput(
@ -194,20 +187,16 @@ var destroyUserCmd = &cobra.Command{
} else { } else {
SuccessOutput(map[string]string{"Result": "User not destroyed"}, "User not destroyed", output) SuccessOutput(map[string]string{"Result": "User not destroyed"}, "User not destroyed", output)
} }
}, }),
} }
var listUsersCmd = &cobra.Command{ var listUsersCmd = &cobra.Command{
Use: "list", Use: "list",
Short: "List all the users", Short: "List all the users",
Aliases: []string{"ls", "show"}, 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") output, _ := cmd.Flags().GetString("output")
ctx, client, conn, cancel := newHeadscaleCLIWithConfig()
defer cancel()
defer conn.Close()
request := &v1.ListUsersRequest{} request := &v1.ListUsersRequest{}
id, _ := cmd.Flags().GetInt64("identifier") id, _ := cmd.Flags().GetInt64("identifier")
@ -259,20 +248,16 @@ var listUsersCmd = &cobra.Command{
output, output,
) )
} }
}, }),
} }
var renameUserCmd = &cobra.Command{ var renameUserCmd = &cobra.Command{
Use: "rename", Use: "rename",
Short: "Renames a user", Short: "Renames a user",
Aliases: []string{"mv"}, 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") output, _ := cmd.Flags().GetString("output")
ctx, client, conn, cancel := newHeadscaleCLIWithConfig()
defer cancel()
defer conn.Close()
id, username := usernameAndIDFromFlag(cmd) id, username := usernameAndIDFromFlag(cmd)
listReq := &v1.ListUsersRequest{ listReq := &v1.ListUsersRequest{
Name: username, Name: username,
@ -314,5 +299,5 @@ var renameUserCmd = &cobra.Command{
} }
SuccessOutput(response.GetUser(), "User renamed", output) SuccessOutput(response.GetUser(), "User renamed", output)
}, }),
} }

View File

@ -13,6 +13,7 @@ import (
"github.com/juanfont/headscale/hscontrol/util" "github.com/juanfont/headscale/hscontrol/util"
"github.com/juanfont/headscale/hscontrol/util/zlog/zf" "github.com/juanfont/headscale/hscontrol/util/zlog/zf"
"github.com/rs/zerolog/log" "github.com/rs/zerolog/log"
"github.com/spf13/cobra"
"google.golang.org/grpc" "google.golang.org/grpc"
"google.golang.org/grpc/credentials" "google.golang.org/grpc/credentials"
"google.golang.org/grpc/credentials/insecure" "google.golang.org/grpc/credentials/insecure"
@ -41,6 +42,21 @@ func newHeadscaleServerWithConfig() (*hscontrol.Headscale, error) {
return app, nil 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) { func newHeadscaleCLIWithConfig() (context.Context, v1.HeadscaleServiceClient, *grpc.ClientConn, context.CancelFunc) {
cfg, err := types.LoadCLIConfig() cfg, err := types.LoadCLIConfig()
if err != nil { if err != nil {