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

@@ -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

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")

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{

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{

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{

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{

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) {