public/headscale
1
0
Fork 0
mirror of https://github.com/juanfont/headscale.git synced 2026-07-08 00:50:20 +09:00
Code Issues Packages Projects Releases Wiki Activity

cli: consolidate gRPC command helpers

Browse Source
This commit is contained in:
Kristoffer Dalby
2026-06-16 08:45:59 +00:00
parent b88a3cb7b2
commit 0b7c8aa270
6 changed files with 132 additions and 112 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cmd/headscale/cli/auth.go
+30 -29
View File
@@ -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})
}),
}
cmd/headscale/cli/nodes.go
+3 -2
View File
@@ -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
+37 -41
View File
@@ -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
}
cmd/headscale/cli/preauthkeys.go
+16 -8
View File
@@ -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{
cmd/headscale/cli/users.go
+29 -32
View File
@@ -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
+17
View File
@@ -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 {