TheArchive/headscale
1
0
Fork 0
mirror of https://github.com/juanfont/headscale.git synced 2024-11-23 18:15:26 +00:00
Code Issues Packages Projects Releases Wiki Activity

Attempt to dry up CLI client, add proepr config

Browse Source 
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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
app.go
View File

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

6
cmd/headscale/cli/debug.go
View File

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

18
cmd/headscale/cli/namespaces.go
View File

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

25
cmd/headscale/cli/nodes.go
View File

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

13
cmd/headscale/cli/preauthkeys.go
View File

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

10
cmd/headscale/cli/routes.go
View File

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

46
cmd/headscale/cli/utils.go
View File

@ -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,18 +366,10 @@ 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 {
if cfg.CLI.Insecure {
grpcOptions = append(grpcOptions, grpc.WithInsecure())
}
}
}
log.Trace().Caller().Str("address", address).Msg("Connecting via gRPC")
conn, err := grpc.DialContext(ctx, address, grpcOptions...)
@ -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) {