From dce6b8d72e8670b18d329b9c874de0c1ac308b39 Mon Sep 17 00:00:00 2001 From: Kristoffer Dalby Date: Mon, 8 Nov 2021 08:02:01 +0000 Subject: [PATCH] Add test case and fix nil pointer in preauthkeys command without expiration --- cmd/headscale/cli/preauthkeys.go | 48 ++++++++++++++-------------- grpcv1.go | 6 +++- integration_cli_test.go | 54 ++++++++++++++++++++++++++++++++ 3 files changed, 82 insertions(+), 26 deletions(-) diff --git a/cmd/headscale/cli/preauthkeys.go b/cmd/headscale/cli/preauthkeys.go index 413dbdc5..06f61974 100644 --- a/cmd/headscale/cli/preauthkeys.go +++ b/cmd/headscale/cli/preauthkeys.go @@ -2,13 +2,12 @@ package cli import ( "fmt" - "log" "strconv" "time" - "github.com/hako/durafmt" v1 "github.com/juanfont/headscale/gen/go/headscale/v1" "github.com/pterm/pterm" + "github.com/rs/zerolog/log" "github.com/spf13/cobra" "google.golang.org/protobuf/types/known/timestamppb" ) @@ -18,7 +17,7 @@ func init() { preauthkeysCmd.PersistentFlags().StringP("namespace", "n", "", "Namespace") err := preauthkeysCmd.MarkPersistentFlagRequired("namespace") if err != nil { - log.Fatalf(err.Error()) + log.Fatal().Err(err).Msg("") } preauthkeysCmd.AddCommand(listPreAuthKeys) preauthkeysCmd.AddCommand(createPreAuthKeyCmd) @@ -26,7 +25,7 @@ func init() { createPreAuthKeyCmd.PersistentFlags().Bool("reusable", false, "Make the preauthkey reusable") createPreAuthKeyCmd.PersistentFlags().Bool("ephemeral", false, "Preauthkey for ephemeral nodes") createPreAuthKeyCmd.Flags(). - StringP("expiration", "e", "", "Human-readable expiration of the key (30m, 24h, 365d...)") + DurationP("expiration", "e", 24*time.Hour, "Human-readable expiration of the key (30m, 24h, 365d...)") } var preauthkeysCmd = &cobra.Command{ @@ -92,7 +91,8 @@ var listPreAuthKeys = &cobra.Command{ } err = pterm.DefaultTable.WithHasHeader().WithData(d).Render() if err != nil { - log.Fatal(err) + ErrorOutput(err, fmt.Sprintf("Failed to render pterm table: %s", err), output) + return } }, } @@ -103,7 +103,7 @@ var createPreAuthKeyCmd = &cobra.Command{ Run: func(cmd *cobra.Command, args []string) { output, _ := cmd.Flags().GetString("output") - n, err := cmd.Flags().GetString("namespace") + namespace, err := cmd.Flags().GetString("namespace") if err != nil { ErrorOutput(err, fmt.Sprintf("Error getting namespace: %s", err), output) return @@ -112,28 +112,25 @@ var createPreAuthKeyCmd = &cobra.Command{ reusable, _ := cmd.Flags().GetBool("reusable") ephemeral, _ := cmd.Flags().GetBool("ephemeral") - e, _ := cmd.Flags().GetString("expiration") - var expiration *time.Time - if e != "" { - duration, err := durafmt.ParseStringShort(e) - if err != nil { - log.Fatalf("Error parsing expiration: %s", err) - } - exp := time.Now().UTC().Add(duration.Duration()) - expiration = &exp + request := &v1.CreatePreAuthKeyRequest{ + Namespace: namespace, + Resuable: reusable, + Ephemeral: ephemeral, + } + + if cmd.Flags().Changed("expiration") { + duration, _ := cmd.Flags().GetDuration("expiration") + expiration := time.Now().UTC().Add(duration) + + log.Trace().Dur("expiration", duration).Msg("expiration has been set") + + request.Expiration = timestamppb.New(expiration) } ctx, client, conn, cancel := getHeadscaleCLIClient() defer cancel() defer conn.Close() - request := &v1.CreatePreAuthKeyRequest{ - Namespace: n, - Resuable: reusable, - Ephemeral: ephemeral, - Expiration: timestamppb.New(*expiration), - } - response, err := client.CreatePreAuthKey(ctx, request) if err != nil { ErrorOutput(err, fmt.Sprintf("Cannot create Pre Auth Key: %s\n", err), output) @@ -155,9 +152,10 @@ var expirePreAuthKeyCmd = &cobra.Command{ }, Run: func(cmd *cobra.Command, args []string) { output, _ := cmd.Flags().GetString("output") - n, err := cmd.Flags().GetString("namespace") + namespace, err := cmd.Flags().GetString("namespace") if err != nil { - log.Fatalf("Error getting namespace: %s", err) + ErrorOutput(err, fmt.Sprintf("Error getting namespace: %s", err), output) + return } ctx, client, conn, cancel := getHeadscaleCLIClient() @@ -165,7 +163,7 @@ var expirePreAuthKeyCmd = &cobra.Command{ defer conn.Close() request := &v1.ExpirePreAuthKeyRequest{ - Namespace: n, + Namespace: namespace, Key: args[0], } diff --git a/grpcv1.go b/grpcv1.go index 0a28fbe6..dfb86425 100644 --- a/grpcv1.go +++ b/grpcv1.go @@ -99,7 +99,11 @@ func (api headscaleV1APIServer) CreatePreAuthKey( ctx context.Context, request *v1.CreatePreAuthKeyRequest, ) (*v1.CreatePreAuthKeyResponse, error) { - expiration := request.GetExpiration().AsTime() + var expiration time.Time + if request.GetExpiration() != nil { + expiration = request.GetExpiration().AsTime() + } + preAuthKey, err := api.h.CreatePreAuthKey( request.GetNamespace(), request.GetResuable(), diff --git a/integration_cli_test.go b/integration_cli_test.go index 4b408093..20b6ce34 100644 --- a/integration_cli_test.go +++ b/integration_cli_test.go @@ -298,6 +298,12 @@ func (s *IntegrationCLITestSuite) TestPreAuthKeyCommand() { assert.True(s.T(), listedPreAuthKeys[3].Expiration.AsTime().After(time.Now())) assert.True(s.T(), listedPreAuthKeys[4].Expiration.AsTime().After(time.Now())) + assert.True(s.T(), listedPreAuthKeys[0].Expiration.AsTime().Before(time.Now().Add(time.Hour*26))) + assert.True(s.T(), listedPreAuthKeys[1].Expiration.AsTime().Before(time.Now().Add(time.Hour*26))) + assert.True(s.T(), listedPreAuthKeys[2].Expiration.AsTime().Before(time.Now().Add(time.Hour*26))) + assert.True(s.T(), listedPreAuthKeys[3].Expiration.AsTime().Before(time.Now().Add(time.Hour*26))) + assert.True(s.T(), listedPreAuthKeys[4].Expiration.AsTime().Before(time.Now().Add(time.Hour*26))) + // Expire three keys for i := 0; i < 3; i++ { _, err := ExecuteCommand( @@ -342,6 +348,54 @@ func (s *IntegrationCLITestSuite) TestPreAuthKeyCommand() { assert.True(s.T(), listedAfterExpirePreAuthKeys[4].Expiration.AsTime().After(time.Now())) } +func (s *IntegrationCLITestSuite) TestPreAuthKeyCommandWithoutExpiry() { + namespace, err := s.createNamespace("pre-auth-key-without-exp-namespace") + assert.Nil(s.T(), err) + + preAuthResult, err := ExecuteCommand( + &s.headscale, + []string{ + "headscale", + "preauthkeys", + "--namespace", + namespace.Name, + "create", + "--reusable", + "--output", + "json", + }, + []string{}, + ) + assert.Nil(s.T(), err) + + var preAuthKey v1.PreAuthKey + err = json.Unmarshal([]byte(preAuthResult), &preAuthKey) + assert.Nil(s.T(), err) + + // Test list of keys + listResult, err := ExecuteCommand( + &s.headscale, + []string{ + "headscale", + "preauthkeys", + "--namespace", + namespace.Name, + "list", + "--output", + "json", + }, + []string{}, + ) + assert.Nil(s.T(), err) + + var listedPreAuthKeys []v1.PreAuthKey + err = json.Unmarshal([]byte(listResult), &listedPreAuthKeys) + assert.Nil(s.T(), err) + + assert.Len(s.T(), listedPreAuthKeys, 1) + assert.True(s.T(), time.Time{}.Equal(listedPreAuthKeys[0].Expiration.AsTime())) +} + func (s *IntegrationCLITestSuite) TestNodeCommand() { namespace, err := s.createNamespace("machine-namespace") assert.Nil(s.T(), err)