cmd/headscale/cli: deduplicate expiration parsing and api-key flag validation

Add expirationFromFlag helper that parses the --expiration flag into a
timestamppb.Timestamp, replacing identical duration-parsing blocks in
api_key.go and preauthkeys.go.

Add apiKeyIDOrPrefix helper to validate the mutually-exclusive --id and
--prefix flags, replacing the duplicated switch block in expireAPIKeyCmd
and deleteAPIKeyCmd.
This commit is contained in:
Kristoffer Dalby
2026-02-18 14:53:30 +00:00
parent 6c08b49d63
commit 22fccae125
3 changed files with 60 additions and 63 deletions

View File

@@ -4,14 +4,11 @@ import (
"context"
"fmt"
"strconv"
"time"
v1 "github.com/juanfont/headscale/gen/go/headscale/v1"
"github.com/juanfont/headscale/hscontrol/util"
"github.com/prometheus/common/model"
"github.com/pterm/pterm"
"github.com/spf13/cobra"
"google.golang.org/protobuf/types/known/timestamppb"
)
const (
@@ -87,20 +84,14 @@ and cannot be retrieved again.
If you loose a key, create a new one and revoke (expire) the old one.`,
Aliases: []string{"c", "new"},
RunE: grpcRunE(func(ctx context.Context, client v1.HeadscaleServiceClient, cmd *cobra.Command, args []string) error {
request := &v1.CreateApiKeyRequest{}
durationStr, _ := cmd.Flags().GetString("expiration")
duration, err := model.ParseDuration(durationStr)
expiration, err := expirationFromFlag(cmd)
if err != nil {
return fmt.Errorf("parsing duration: %w", err)
return err
}
expiration := time.Now().UTC().Add(time.Duration(duration))
request.Expiration = timestamppb.New(expiration)
response, err := client.CreateApiKey(ctx, request)
response, err := client.CreateApiKey(ctx, &v1.CreateApiKeyRequest{
Expiration: expiration,
})
if err != nil {
return fmt.Errorf("creating api key: %w", err)
}
@@ -109,29 +100,36 @@ If you loose a key, create a new one and revoke (expire) the old one.`,
}),
}
// apiKeyIDOrPrefix reads --id and --prefix from cmd and validates that
// exactly one is provided.
func apiKeyIDOrPrefix(cmd *cobra.Command) (uint64, string, error) {
id, _ := cmd.Flags().GetUint64("id")
prefix, _ := cmd.Flags().GetString("prefix")
switch {
case id == 0 && prefix == "":
return 0, "", fmt.Errorf("either --id or --prefix must be provided: %w", errMissingParameter)
case id != 0 && prefix != "":
return 0, "", fmt.Errorf("only one of --id or --prefix can be provided: %w", errMissingParameter)
}
return id, prefix, nil
}
var expireAPIKeyCmd = &cobra.Command{
Use: "expire",
Short: "Expire an ApiKey",
Aliases: []string{"revoke", "exp", "e"},
RunE: grpcRunE(func(ctx context.Context, client v1.HeadscaleServiceClient, cmd *cobra.Command, args []string) error {
id, _ := cmd.Flags().GetUint64("id")
prefix, _ := cmd.Flags().GetString("prefix")
switch {
case id == 0 && prefix == "":
return fmt.Errorf("either --id or --prefix must be provided: %w", errMissingParameter)
case id != 0 && prefix != "":
return fmt.Errorf("only one of --id or --prefix can be provided: %w", errMissingParameter)
id, prefix, err := apiKeyIDOrPrefix(cmd)
if err != nil {
return err
}
request := &v1.ExpireApiKeyRequest{}
if id != 0 {
request.Id = id
} else {
request.Prefix = prefix
}
response, err := client.ExpireApiKey(ctx, request)
response, err := client.ExpireApiKey(ctx, &v1.ExpireApiKeyRequest{
Id: id,
Prefix: prefix,
})
if err != nil {
return fmt.Errorf("expiring api key: %w", err)
}
@@ -145,24 +143,15 @@ var deleteAPIKeyCmd = &cobra.Command{
Short: "Delete an ApiKey",
Aliases: []string{"remove", "del"},
RunE: grpcRunE(func(ctx context.Context, client v1.HeadscaleServiceClient, cmd *cobra.Command, args []string) error {
id, _ := cmd.Flags().GetUint64("id")
prefix, _ := cmd.Flags().GetString("prefix")
switch {
case id == 0 && prefix == "":
return fmt.Errorf("either --id or --prefix must be provided: %w", errMissingParameter)
case id != 0 && prefix != "":
return fmt.Errorf("only one of --id or --prefix can be provided: %w", errMissingParameter)
id, prefix, err := apiKeyIDOrPrefix(cmd)
if err != nil {
return err
}
request := &v1.DeleteApiKeyRequest{}
if id != 0 {
request.Id = id
} else {
request.Prefix = prefix
}
response, err := client.DeleteApiKey(ctx, request)
response, err := client.DeleteApiKey(ctx, &v1.DeleteApiKeyRequest{
Id: id,
Prefix: prefix,
})
if err != nil {
return fmt.Errorf("deleting api key: %w", err)
}

View File

@@ -5,13 +5,10 @@ import (
"fmt"
"strconv"
"strings"
"time"
v1 "github.com/juanfont/headscale/gen/go/headscale/v1"
"github.com/prometheus/common/model"
"github.com/pterm/pterm"
"github.com/spf13/cobra"
"google.golang.org/protobuf/types/known/timestamppb"
)
const (
@@ -109,23 +106,18 @@ var createPreAuthKeyCmd = &cobra.Command{
ephemeral, _ := cmd.Flags().GetBool("ephemeral")
tags, _ := cmd.Flags().GetStringSlice("tags")
request := &v1.CreatePreAuthKeyRequest{
User: user,
Reusable: reusable,
Ephemeral: ephemeral,
AclTags: tags,
}
durationStr, _ := cmd.Flags().GetString("expiration")
duration, err := model.ParseDuration(durationStr)
expiration, err := expirationFromFlag(cmd)
if err != nil {
return fmt.Errorf("parsing duration: %w", err)
return err
}
expiration := time.Now().UTC().Add(time.Duration(duration))
request.Expiration = timestamppb.New(expiration)
request := &v1.CreatePreAuthKeyRequest{
User: user,
Reusable: reusable,
Ephemeral: ephemeral,
AclTags: tags,
Expiration: expiration,
}
response, err := client.CreatePreAuthKey(ctx, request)
if err != nil {

View File

@@ -7,17 +7,20 @@ import (
"errors"
"fmt"
"os"
"time"
v1 "github.com/juanfont/headscale/gen/go/headscale/v1"
"github.com/juanfont/headscale/hscontrol"
"github.com/juanfont/headscale/hscontrol/types"
"github.com/juanfont/headscale/hscontrol/util"
"github.com/juanfont/headscale/hscontrol/util/zlog/zf"
"github.com/prometheus/common/model"
"github.com/rs/zerolog/log"
"github.com/spf13/cobra"
"google.golang.org/grpc"
"google.golang.org/grpc/credentials"
"google.golang.org/grpc/credentials/insecure"
"google.golang.org/protobuf/types/known/timestamppb"
"gopkg.in/yaml.v3"
)
@@ -221,6 +224,19 @@ func printOutput(cmd *cobra.Command, result any, override string) error {
return nil
}
// expirationFromFlag parses the --expiration flag as a Prometheus-style
// duration (e.g. "90d", "1h") and returns an absolute timestamp.
func expirationFromFlag(cmd *cobra.Command) (*timestamppb.Timestamp, error) {
durationStr, _ := cmd.Flags().GetString("expiration")
duration, err := model.ParseDuration(durationStr)
if err != nil {
return nil, fmt.Errorf("parsing duration: %w", err)
}
return timestamppb.New(time.Now().UTC().Add(time.Duration(duration))), nil
}
// confirmAction returns true when the user confirms a prompt, or when
// --force is set. Callers decide what to do when it returns false.
func confirmAction(cmd *cobra.Command, prompt string) bool {