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 97333b1fb..c0f439818 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 8dfe0c41a..0e5d59b00 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 427df6241..77e0b4769 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 b60e0a514..4edeb855c 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 e8252419c..413dbdc59 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 4a935cf9f..d200160b7 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 90479bcd0..30f1a7892 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) {