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
This commit is contained in:
Kristoffer Dalby 2021-11-07 09:41:14 +00:00
parent ce3f79a3bf
commit 2dfd42f80c
7 changed files with 54 additions and 75 deletions

9
app.go
View File

@ -86,6 +86,8 @@ type Config struct {
OIDC OIDCConfig OIDC OIDCConfig
CLI CLIConfig
MaxMachineRegistrationDuration time.Duration MaxMachineRegistrationDuration time.Duration
DefaultMachineRegistrationDuration time.Duration DefaultMachineRegistrationDuration time.Duration
} }
@ -104,6 +106,13 @@ type DERPConfig struct {
UpdateFrequency time.Duration UpdateFrequency time.Duration
} }
type CLIConfig struct {
Address string
APIKey string
Insecure bool
Timeout time.Duration
}
// Headscale represents the base app of the service. // Headscale represents the base app of the service.
type Headscale struct { type Headscale struct {
cfg Config cfg Config

View File

@ -1,9 +1,7 @@
package cli package cli
import ( import (
"context"
"fmt" "fmt"
"time"
v1 "github.com/juanfont/headscale/gen/go/headscale/v1" v1 "github.com/juanfont/headscale/gen/go/headscale/v1"
"github.com/rs/zerolog/log" "github.com/rs/zerolog/log"
@ -51,10 +49,8 @@ var createNodeCmd = &cobra.Command{
return return
} }
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) ctx, client, conn, cancel := getHeadscaleCLIClient()
defer cancel() defer cancel()
client, conn := getHeadscaleGRPCClient(ctx)
defer conn.Close() defer conn.Close()
name, err := cmd.Flags().GetString("name") name, err := cmd.Flags().GetString("name")

View File

@ -1,9 +1,7 @@
package cli package cli
import ( import (
"context"
"fmt" "fmt"
"time"
v1 "github.com/juanfont/headscale/gen/go/headscale/v1" v1 "github.com/juanfont/headscale/gen/go/headscale/v1"
"github.com/pterm/pterm" "github.com/pterm/pterm"
@ -38,10 +36,8 @@ var createNamespaceCmd = &cobra.Command{
namespaceName := args[0] namespaceName := args[0]
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) ctx, client, conn, cancel := getHeadscaleCLIClient()
defer cancel() defer cancel()
client, conn := getHeadscaleGRPCClient(ctx)
defer conn.Close() defer conn.Close()
log.Trace().Interface("client", client).Msg("Obtained gRPC client") log.Trace().Interface("client", client).Msg("Obtained gRPC client")
@ -73,10 +69,8 @@ var destroyNamespaceCmd = &cobra.Command{
namespaceName := args[0] namespaceName := args[0]
ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) ctx, client, conn, cancel := getHeadscaleCLIClient()
defer cancel() defer cancel()
client, conn := getHeadscaleGRPCClient(ctx)
defer conn.Close() defer conn.Close()
request := &v1.DeleteNamespaceRequest{Name: namespaceName} request := &v1.DeleteNamespaceRequest{Name: namespaceName}
@ -97,10 +91,8 @@ var listNamespacesCmd = &cobra.Command{
Run: func(cmd *cobra.Command, args []string) { Run: func(cmd *cobra.Command, args []string) {
output, _ := cmd.Flags().GetString("output") output, _ := cmd.Flags().GetString("output")
ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) ctx, client, conn, cancel := getHeadscaleCLIClient()
defer cancel() defer cancel()
client, conn := getHeadscaleGRPCClient(ctx)
defer conn.Close() defer conn.Close()
request := &v1.ListNamespacesRequest{} request := &v1.ListNamespacesRequest{}
@ -147,10 +139,8 @@ var renameNamespaceCmd = &cobra.Command{
Run: func(cmd *cobra.Command, args []string) { Run: func(cmd *cobra.Command, args []string) {
output, _ := cmd.Flags().GetString("output") output, _ := cmd.Flags().GetString("output")
ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) ctx, client, conn, cancel := getHeadscaleCLIClient()
defer cancel() defer cancel()
client, conn := getHeadscaleGRPCClient(ctx)
defer conn.Close() defer conn.Close()
request := &v1.RenameNamespaceRequest{ request := &v1.RenameNamespaceRequest{

View File

@ -1,7 +1,6 @@
package cli package cli
import ( import (
"context"
"fmt" "fmt"
"log" "log"
"strconv" "strconv"
@ -80,10 +79,8 @@ var registerNodeCmd = &cobra.Command{
return return
} }
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) ctx, client, conn, cancel := getHeadscaleCLIClient()
defer cancel() defer cancel()
client, conn := getHeadscaleGRPCClient(ctx)
defer conn.Close() defer conn.Close()
machineKey, err := cmd.Flags().GetString("key") machineKey, err := cmd.Flags().GetString("key")
@ -118,10 +115,8 @@ var listNodesCmd = &cobra.Command{
return return
} }
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) ctx, client, conn, cancel := getHeadscaleCLIClient()
defer cancel() defer cancel()
client, conn := getHeadscaleGRPCClient(ctx)
defer conn.Close() defer conn.Close()
request := &v1.ListMachinesRequest{ request := &v1.ListMachinesRequest{
@ -165,10 +160,8 @@ var deleteNodeCmd = &cobra.Command{
return return
} }
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) ctx, client, conn, cancel := getHeadscaleCLIClient()
defer cancel() defer cancel()
client, conn := getHeadscaleGRPCClient(ctx)
defer conn.Close() defer conn.Close()
getRequest := &v1.GetMachineRequest{ getRequest := &v1.GetMachineRequest{
@ -225,10 +218,8 @@ func sharingWorker(
return "", nil, nil, err return "", nil, nil, err
} }
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) ctx, client, conn, cancel := getHeadscaleCLIClient()
defer cancel() defer cancel()
client, conn := getHeadscaleGRPCClient(ctx)
defer conn.Close() defer conn.Close()
id, err := cmd.Flags().GetInt("identifier") id, err := cmd.Flags().GetInt("identifier")
@ -270,10 +261,8 @@ var shareMachineCmd = &cobra.Command{
return return
} }
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) ctx, client, conn, cancel := getHeadscaleCLIClient()
defer cancel() defer cancel()
client, conn := getHeadscaleGRPCClient(ctx)
defer conn.Close() defer conn.Close()
request := &v1.ShareMachineRequest{ request := &v1.ShareMachineRequest{
@ -301,10 +290,8 @@ var unshareMachineCmd = &cobra.Command{
return return
} }
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) ctx, client, conn, cancel := getHeadscaleCLIClient()
defer cancel() defer cancel()
client, conn := getHeadscaleGRPCClient(ctx)
defer conn.Close() defer conn.Close()
request := &v1.UnshareMachineRequest{ request := &v1.UnshareMachineRequest{

View File

@ -1,7 +1,6 @@
package cli package cli
import ( import (
"context"
"fmt" "fmt"
"log" "log"
"strconv" "strconv"
@ -47,10 +46,8 @@ var listPreAuthKeys = &cobra.Command{
return return
} }
ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) ctx, client, conn, cancel := getHeadscaleCLIClient()
defer cancel() defer cancel()
client, conn := getHeadscaleGRPCClient(ctx)
defer conn.Close() defer conn.Close()
request := &v1.ListPreAuthKeysRequest{ request := &v1.ListPreAuthKeysRequest{
@ -126,10 +123,8 @@ var createPreAuthKeyCmd = &cobra.Command{
expiration = &exp expiration = &exp
} }
ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) ctx, client, conn, cancel := getHeadscaleCLIClient()
defer cancel() defer cancel()
client, conn := getHeadscaleGRPCClient(ctx)
defer conn.Close() defer conn.Close()
request := &v1.CreatePreAuthKeyRequest{ request := &v1.CreatePreAuthKeyRequest{
@ -165,10 +160,8 @@ var expirePreAuthKeyCmd = &cobra.Command{
log.Fatalf("Error getting namespace: %s", err) log.Fatalf("Error getting namespace: %s", err)
} }
ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) ctx, client, conn, cancel := getHeadscaleCLIClient()
defer cancel() defer cancel()
client, conn := getHeadscaleGRPCClient(ctx)
defer conn.Close() defer conn.Close()
request := &v1.ExpirePreAuthKeyRequest{ request := &v1.ExpirePreAuthKeyRequest{

View File

@ -1,11 +1,9 @@
package cli package cli
import ( import (
"context"
"fmt" "fmt"
"log" "log"
"strconv" "strconv"
"time"
v1 "github.com/juanfont/headscale/gen/go/headscale/v1" v1 "github.com/juanfont/headscale/gen/go/headscale/v1"
"github.com/pterm/pterm" "github.com/pterm/pterm"
@ -51,10 +49,8 @@ var listRoutesCmd = &cobra.Command{
return return
} }
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) ctx, client, conn, cancel := getHeadscaleCLIClient()
defer cancel() defer cancel()
client, conn := getHeadscaleGRPCClient(ctx)
defer conn.Close() defer conn.Close()
request := &v1.GetMachineRouteRequest{ request := &v1.GetMachineRouteRequest{
@ -108,10 +104,8 @@ omit the route you do not want to enable.
return return
} }
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) ctx, client, conn, cancel := getHeadscaleCLIClient()
defer cancel() defer cancel()
client, conn := getHeadscaleGRPCClient(ctx)
defer conn.Close() defer conn.Close()
request := &v1.EnableMachineRoutesRequest{ request := &v1.EnableMachineRoutesRequest{

View File

@ -9,7 +9,6 @@ import (
"os" "os"
"path/filepath" "path/filepath"
"regexp" "regexp"
"strconv"
"strings" "strings"
"time" "time"
@ -34,6 +33,9 @@ func LoadConfig(path string) error {
// For testing // For testing
viper.AddConfigPath(path) viper.AddConfigPath(path)
} }
viper.SetEnvPrefix("headscale")
viper.SetEnvKeyReplacer(strings.NewReplacer(".", "_"))
viper.AutomaticEnv() viper.AutomaticEnv()
viper.SetDefault("tls_letsencrypt_cache_dir", "/var/www/.cache") 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("unix_socket", "/var/run/headscale.sock")
viper.SetDefault("cli.insecure", false)
viper.SetDefault("cli.timeout", "5s")
err := viper.ReadInConfig() err := viper.ReadInConfig()
if err != nil { if err != nil {
return fmt.Errorf("Fatal error reading config file: %s \n", err) 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"), 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, MaxMachineRegistrationDuration: maxMachineRegistrationDuration,
DefaultMachineRegistrationDuration: defaultMachineRegistrationDuration, DefaultMachineRegistrationDuration: defaultMachineRegistrationDuration,
} }
@ -313,21 +325,27 @@ func getHeadscaleApp() (*headscale.Headscale, error) {
return h, nil 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{ grpcOptions := []grpc.DialOption{
grpc.WithBlock(), 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 the address is not set, we assume that we are on the server hosting headscale.
if address == "" { if address == "" {
cfg := getHeadscaleConfig()
log.Debug(). log.Debug().
Str("socket", cfg.UnixSocket). 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 address = cfg.UnixSocket
@ -338,9 +356,9 @@ func getHeadscaleGRPCClient(ctx context.Context) (v1.HeadscaleServiceClient, *gr
) )
} else { } else {
// If we are not connecting to a local server, require an API key for authentication // 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 == "" { 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, grpcOptions = append(grpcOptions,
grpc.WithPerRPCCredentials(tokenAuth{ grpc.WithPerRPCCredentials(tokenAuth{
@ -348,16 +366,8 @@ func getHeadscaleGRPCClient(ctx context.Context) (v1.HeadscaleServiceClient, *gr
}), }),
) )
insecureStr := os.Getenv("HEADSCALE_INSECURE") if cfg.CLI.Insecure {
if insecureStr != "" { grpcOptions = append(grpcOptions, grpc.WithInsecure())
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())
}
} }
} }
@ -369,7 +379,7 @@ func getHeadscaleGRPCClient(ctx context.Context) (v1.HeadscaleServiceClient, *gr
client := v1.NewHeadscaleServiceClient(conn) client := v1.NewHeadscaleServiceClient(conn)
return client, conn return ctx, client, conn, cancel
} }
func SuccessOutput(result interface{}, override string, outputFormat string) { func SuccessOutput(result interface{}, override string, outputFormat string) {