From af891808f6451f3cdbd3ed3a88f7d4b76852f08e Mon Sep 17 00:00:00 2001
From: Kristoffer Dalby <kradalby@kradalby.no>
Date: Sun, 5 Jun 2022 17:47:12 +0200
Subject: [PATCH] Make get config load the config, use config in main method

---
 cmd/headscale/cli/utils.go | 15 ++++++++++++---
 cmd/headscale/headscale.go | 14 ++++----------
 config.go                  | 11 ++++++++---
 3 files changed, 24 insertions(+), 16 deletions(-)

diff --git a/cmd/headscale/cli/utils.go b/cmd/headscale/cli/utils.go
index 0b447abd..f5c679c2 100644
--- a/cmd/headscale/cli/utils.go
+++ b/cmd/headscale/cli/utils.go
@@ -24,6 +24,11 @@ const (
 )
 
 func getHeadscaleApp() (*headscale.Headscale, error) {
+	cfg, err := headscale.GetHeadscaleConfig()
+	if err != nil {
+		return nil, fmt.Errorf("failed to load configuration while creating headscale instance: %w", err)
+	}
+
 	// Minimum inactivity time out is keepalive timeout (60s) plus a few seconds
 	// to avoid races
 	minInactivityTimeout, _ := time.ParseDuration("65s")
@@ -39,8 +44,6 @@ func getHeadscaleApp() (*headscale.Headscale, error) {
 		return nil, err
 	}
 
-	cfg := headscale.GetHeadscaleConfig()
-
 	app, err := headscale.NewHeadscale(cfg)
 	if err != nil {
 		return nil, err
@@ -63,7 +66,13 @@ func getHeadscaleApp() (*headscale.Headscale, error) {
 }
 
 func getHeadscaleCLIClient() (context.Context, v1.HeadscaleServiceClient, *grpc.ClientConn, context.CancelFunc) {
-	cfg := headscale.GetHeadscaleConfig()
+	cfg, err := headscale.GetHeadscaleConfig()
+	if err != nil {
+		log.Fatal().
+			Err(err).
+			Caller().
+			Msgf("Failed to load configuration")
+	}
 
 	log.Debug().
 		Dur("timeout", cfg.CLI.Timeout).
diff --git a/cmd/headscale/headscale.go b/cmd/headscale/headscale.go
index dde1c2e7..f5e28661 100644
--- a/cmd/headscale/headscale.go
+++ b/cmd/headscale/headscale.go
@@ -11,7 +11,6 @@ import (
 	"github.com/juanfont/headscale/cmd/headscale/cli"
 	"github.com/rs/zerolog"
 	"github.com/rs/zerolog/log"
-	"github.com/spf13/viper"
 	"github.com/tcnksm/go-latest"
 )
 
@@ -44,19 +43,14 @@ func main() {
 		NoColor:    !colors,
 	})
 
-	if err := headscale.LoadConfig(""); err != nil {
+	cfg, err := headscale.GetHeadscaleConfig()
+	if err != nil {
 		log.Fatal().Caller().Err(err)
 	}
 
 	machineOutput := cli.HasMachineOutputFlag()
 
-	logLevel := viper.GetString("log_level")
-	level, err := zerolog.ParseLevel(logLevel)
-	if err != nil {
-		zerolog.SetGlobalLevel(zerolog.DebugLevel)
-	} else {
-		zerolog.SetGlobalLevel(level)
-	}
+	zerolog.SetGlobalLevel(cfg.LogLevel)
 
 	// If the user has requested a "machine" readable format,
 	// then disable login so the output remains valid.
@@ -64,7 +58,7 @@ func main() {
 		zerolog.SetGlobalLevel(zerolog.Disabled)
 	}
 
-	if !viper.GetBool("disable_check_updates") && !machineOutput {
+	if !cfg.DisableUpdateCheck && !machineOutput {
 		if (runtime.GOOS == "linux" || runtime.GOOS == "darwin") &&
 			cli.Version != "dev" {
 			githubTag := &latest.GithubTag{
diff --git a/config.go b/config.go
index 045ce02a..909a48c4 100644
--- a/config.go
+++ b/config.go
@@ -376,7 +376,12 @@ func GetDNSConfig() (*tailcfg.DNSConfig, string) {
 	return nil, ""
 }
 
-func GetHeadscaleConfig() Config {
+func GetHeadscaleConfig() (*Config, error) {
+	err := LoadConfig("")
+	if err != nil {
+		return nil, err
+	}
+
 	dnsConfig, baseDomain := GetDNSConfig()
 	derpConfig := GetDERPConfig()
 	logConfig := GetLogTailConfig()
@@ -435,7 +440,7 @@ func GetHeadscaleConfig() Config {
 			Msgf("'ip_prefixes' not configured, falling back to default: %v", prefixes)
 	}
 
-	return Config{
+	return &Config{
 		ServerURL:          viper.GetString("server_url"),
 		Addr:               viper.GetString("listen_addr"),
 		MetricsAddr:        viper.GetString("metrics_listen_addr"),
@@ -495,5 +500,5 @@ func GetHeadscaleConfig() Config {
 		},
 
 		ACL: GetACLConfig(),
-	}
+	}, nil
 }