From 2dfd42f80c2829abb5366417669870b149d22f4d Mon Sep 17 00:00:00 2001 From: Kristoffer Dalby Date: Sun, 7 Nov 2021 09:41:14 +0000 Subject: [PATCH] Attempt to dry up CLI client, add proepr config This commit is trying to DRY up the initiation of the gRPC client in each command: It renames the function to CLI instead of GRPC as it actually set up a CLI client, not a generic grpc client It also moves the configuration of address, timeout (which is now consistent) and api to use Viper, allowing users to set it via env vars and configuration file --- app.go | 9 ++++++ cmd/headscale/cli/debug.go | 6 +--- cmd/headscale/cli/namespaces.go | 18 +++--------- cmd/headscale/cli/nodes.go | 25 ++++------------- cmd/headscale/cli/preauthkeys.go | 13 ++------- cmd/headscale/cli/routes.go | 10 ++----- cmd/headscale/cli/utils.go | 48 +++++++++++++++++++------------- 7 files changed, 54 insertions(+), 75 deletions(-) diff --git a/app.go b/app.go index 97333b1f..c0f43981 100644 --- a/app.go +++ b/app.go @@ -86,6 +86,8 @@ type Config struct { OIDC OIDCConfig + CLI CLIConfig + MaxMachineRegistrationDuration time.Duration DefaultMachineRegistrationDuration time.Duration } @@ -104,6 +106,13 @@ type DERPConfig struct { UpdateFrequency time.Duration } +type CLIConfig struct { + Address string + APIKey string + Insecure bool + Timeout time.Duration +} + // Headscale represents the base app of the service. type Headscale struct { cfg Config diff --git a/cmd/headscale/cli/debug.go b/cmd/headscale/cli/debug.go index 8dfe0c41..0e5d59b0 100644 --- a/cmd/headscale/cli/debug.go +++ b/cmd/headscale/cli/debug.go @@ -1,9 +1,7 @@ package cli import ( - "context" "fmt" - "time" v1 "github.com/juanfont/headscale/gen/go/headscale/v1" "github.com/rs/zerolog/log" @@ -51,10 +49,8 @@ var createNodeCmd = &cobra.Command{ return } - ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + ctx, client, conn, cancel := getHeadscaleCLIClient() defer cancel() - - client, conn := getHeadscaleGRPCClient(ctx) defer conn.Close() name, err := cmd.Flags().GetString("name") diff --git a/cmd/headscale/cli/namespaces.go b/cmd/headscale/cli/namespaces.go index 427df624..77e0b476 100644 --- a/cmd/headscale/cli/namespaces.go +++ b/cmd/headscale/cli/namespaces.go @@ -1,9 +1,7 @@ package cli import ( - "context" "fmt" - "time" v1 "github.com/juanfont/headscale/gen/go/headscale/v1" "github.com/pterm/pterm" @@ -38,10 +36,8 @@ var createNamespaceCmd = &cobra.Command{ namespaceName := args[0] - ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + ctx, client, conn, cancel := getHeadscaleCLIClient() defer cancel() - - client, conn := getHeadscaleGRPCClient(ctx) defer conn.Close() log.Trace().Interface("client", client).Msg("Obtained gRPC client") @@ -73,10 +69,8 @@ var destroyNamespaceCmd = &cobra.Command{ namespaceName := args[0] - ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + ctx, client, conn, cancel := getHeadscaleCLIClient() defer cancel() - - client, conn := getHeadscaleGRPCClient(ctx) defer conn.Close() request := &v1.DeleteNamespaceRequest{Name: namespaceName} @@ -97,10 +91,8 @@ var listNamespacesCmd = &cobra.Command{ Run: func(cmd *cobra.Command, args []string) { output, _ := cmd.Flags().GetString("output") - ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + ctx, client, conn, cancel := getHeadscaleCLIClient() defer cancel() - - client, conn := getHeadscaleGRPCClient(ctx) defer conn.Close() request := &v1.ListNamespacesRequest{} @@ -147,10 +139,8 @@ var renameNamespaceCmd = &cobra.Command{ Run: func(cmd *cobra.Command, args []string) { output, _ := cmd.Flags().GetString("output") - ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + ctx, client, conn, cancel := getHeadscaleCLIClient() defer cancel() - - client, conn := getHeadscaleGRPCClient(ctx) defer conn.Close() request := &v1.RenameNamespaceRequest{ diff --git a/cmd/headscale/cli/nodes.go b/cmd/headscale/cli/nodes.go index b60e0a51..4edeb855 100644 --- a/cmd/headscale/cli/nodes.go +++ b/cmd/headscale/cli/nodes.go @@ -1,7 +1,6 @@ package cli import ( - "context" "fmt" "log" "strconv" @@ -80,10 +79,8 @@ var registerNodeCmd = &cobra.Command{ return } - ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + ctx, client, conn, cancel := getHeadscaleCLIClient() defer cancel() - - client, conn := getHeadscaleGRPCClient(ctx) defer conn.Close() machineKey, err := cmd.Flags().GetString("key") @@ -118,10 +115,8 @@ var listNodesCmd = &cobra.Command{ return } - ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + ctx, client, conn, cancel := getHeadscaleCLIClient() defer cancel() - - client, conn := getHeadscaleGRPCClient(ctx) defer conn.Close() request := &v1.ListMachinesRequest{ @@ -165,10 +160,8 @@ var deleteNodeCmd = &cobra.Command{ return } - ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + ctx, client, conn, cancel := getHeadscaleCLIClient() defer cancel() - - client, conn := getHeadscaleGRPCClient(ctx) defer conn.Close() getRequest := &v1.GetMachineRequest{ @@ -225,10 +218,8 @@ func sharingWorker( return "", nil, nil, err } - ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + ctx, client, conn, cancel := getHeadscaleCLIClient() defer cancel() - - client, conn := getHeadscaleGRPCClient(ctx) defer conn.Close() id, err := cmd.Flags().GetInt("identifier") @@ -270,10 +261,8 @@ var shareMachineCmd = &cobra.Command{ return } - ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + ctx, client, conn, cancel := getHeadscaleCLIClient() defer cancel() - - client, conn := getHeadscaleGRPCClient(ctx) defer conn.Close() request := &v1.ShareMachineRequest{ @@ -301,10 +290,8 @@ var unshareMachineCmd = &cobra.Command{ return } - ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + ctx, client, conn, cancel := getHeadscaleCLIClient() defer cancel() - - client, conn := getHeadscaleGRPCClient(ctx) defer conn.Close() request := &v1.UnshareMachineRequest{ diff --git a/cmd/headscale/cli/preauthkeys.go b/cmd/headscale/cli/preauthkeys.go index e8252419..413dbdc5 100644 --- a/cmd/headscale/cli/preauthkeys.go +++ b/cmd/headscale/cli/preauthkeys.go @@ -1,7 +1,6 @@ package cli import ( - "context" "fmt" "log" "strconv" @@ -47,10 +46,8 @@ var listPreAuthKeys = &cobra.Command{ return } - ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + ctx, client, conn, cancel := getHeadscaleCLIClient() defer cancel() - - client, conn := getHeadscaleGRPCClient(ctx) defer conn.Close() request := &v1.ListPreAuthKeysRequest{ @@ -126,10 +123,8 @@ var createPreAuthKeyCmd = &cobra.Command{ expiration = &exp } - ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + ctx, client, conn, cancel := getHeadscaleCLIClient() defer cancel() - - client, conn := getHeadscaleGRPCClient(ctx) defer conn.Close() request := &v1.CreatePreAuthKeyRequest{ @@ -165,10 +160,8 @@ var expirePreAuthKeyCmd = &cobra.Command{ log.Fatalf("Error getting namespace: %s", err) } - ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + ctx, client, conn, cancel := getHeadscaleCLIClient() defer cancel() - - client, conn := getHeadscaleGRPCClient(ctx) defer conn.Close() request := &v1.ExpirePreAuthKeyRequest{ diff --git a/cmd/headscale/cli/routes.go b/cmd/headscale/cli/routes.go index 4a935cf9..d200160b 100644 --- a/cmd/headscale/cli/routes.go +++ b/cmd/headscale/cli/routes.go @@ -1,11 +1,9 @@ package cli import ( - "context" "fmt" "log" "strconv" - "time" v1 "github.com/juanfont/headscale/gen/go/headscale/v1" "github.com/pterm/pterm" @@ -51,10 +49,8 @@ var listRoutesCmd = &cobra.Command{ return } - ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + ctx, client, conn, cancel := getHeadscaleCLIClient() defer cancel() - - client, conn := getHeadscaleGRPCClient(ctx) defer conn.Close() request := &v1.GetMachineRouteRequest{ @@ -108,10 +104,8 @@ omit the route you do not want to enable. return } - ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + ctx, client, conn, cancel := getHeadscaleCLIClient() defer cancel() - - client, conn := getHeadscaleGRPCClient(ctx) defer conn.Close() request := &v1.EnableMachineRoutesRequest{ diff --git a/cmd/headscale/cli/utils.go b/cmd/headscale/cli/utils.go index 90479bcd..30f1a789 100644 --- a/cmd/headscale/cli/utils.go +++ b/cmd/headscale/cli/utils.go @@ -9,7 +9,6 @@ import ( "os" "path/filepath" "regexp" - "strconv" "strings" "time" @@ -34,6 +33,9 @@ func LoadConfig(path string) error { // For testing viper.AddConfigPath(path) } + + viper.SetEnvPrefix("headscale") + viper.SetEnvKeyReplacer(strings.NewReplacer(".", "_")) viper.AutomaticEnv() viper.SetDefault("tls_letsencrypt_cache_dir", "/var/www/.cache") @@ -47,6 +49,9 @@ func LoadConfig(path string) error { viper.SetDefault("unix_socket", "/var/run/headscale.sock") + viper.SetDefault("cli.insecure", false) + viper.SetDefault("cli.timeout", "5s") + err := viper.ReadInConfig() if err != nil { return fmt.Errorf("Fatal error reading config file: %s \n", err) @@ -270,6 +275,13 @@ func getHeadscaleConfig() headscale.Config { ClientSecret: viper.GetString("oidc.client_secret"), }, + CLI: headscale.CLIConfig{ + Address: viper.GetString("cli.address"), + APIKey: viper.GetString("cli.api_key"), + Insecure: viper.GetBool("cli.insecure"), + Timeout: viper.GetDuration("cli.timeout"), + }, + MaxMachineRegistrationDuration: maxMachineRegistrationDuration, DefaultMachineRegistrationDuration: defaultMachineRegistrationDuration, } @@ -313,21 +325,27 @@ func getHeadscaleApp() (*headscale.Headscale, error) { return h, nil } -func getHeadscaleGRPCClient(ctx context.Context) (v1.HeadscaleServiceClient, *grpc.ClientConn) { +func getHeadscaleCLIClient() (context.Context, v1.HeadscaleServiceClient, *grpc.ClientConn, context.CancelFunc) { + cfg := getHeadscaleConfig() + + log.Debug(). + Dur("timeout", cfg.CLI.Timeout). + Msgf("Setting timeout") + + ctx, cancel := context.WithTimeout(context.Background(), cfg.CLI.Timeout) + grpcOptions := []grpc.DialOption{ grpc.WithBlock(), } - address := os.Getenv("HEADSCALE_ADDRESS") + address := cfg.CLI.Address // If the address is not set, we assume that we are on the server hosting headscale. if address == "" { - cfg := getHeadscaleConfig() - log.Debug(). Str("socket", cfg.UnixSocket). - Msgf("HEADSCALE_ADDRESS environment is not set, connecting to unix socket.") + Msgf("HEADSCALE_CLI_ADDRESS environment is not set, connecting to unix socket.") address = cfg.UnixSocket @@ -338,9 +356,9 @@ func getHeadscaleGRPCClient(ctx context.Context) (v1.HeadscaleServiceClient, *gr ) } else { // If we are not connecting to a local server, require an API key for authentication - apiKey := os.Getenv("HEADSCALE_API_KEY") + apiKey := cfg.CLI.APIKey if apiKey == "" { - log.Fatal().Msgf("HEADSCALE_API_KEY environment variable needs to be set.") + log.Fatal().Msgf("HEADSCALE_CLI_API_KEY environment variable needs to be set.") } grpcOptions = append(grpcOptions, grpc.WithPerRPCCredentials(tokenAuth{ @@ -348,16 +366,8 @@ func getHeadscaleGRPCClient(ctx context.Context) (v1.HeadscaleServiceClient, *gr }), ) - insecureStr := os.Getenv("HEADSCALE_INSECURE") - if insecureStr != "" { - insecure, err := strconv.ParseBool(insecureStr) - if err != nil { - log.Fatal().Err(err).Msgf("Failed to parse HEADSCALE_INSECURE: %v", err) - } - - if insecure { - grpcOptions = append(grpcOptions, grpc.WithInsecure()) - } + if cfg.CLI.Insecure { + grpcOptions = append(grpcOptions, grpc.WithInsecure()) } } @@ -369,7 +379,7 @@ func getHeadscaleGRPCClient(ctx context.Context) (v1.HeadscaleServiceClient, *gr client := v1.NewHeadscaleServiceClient(conn) - return client, conn + return ctx, client, conn, cancel } func SuccessOutput(result interface{}, override string, outputFormat string) {