From 402a29e50cfe7203e51b5365fea2031528f13ba7 Mon Sep 17 00:00:00 2001
From: Jiang Zhu <jiang.moe@outlook.com>
Date: Sun, 5 Jun 2022 18:25:09 +0800
Subject: [PATCH] impl heascale -c to specify config file

---
 cmd/headscale/cli/root.go  | 57 ++++++++++++++++++++++++++++++++++++++
 cmd/headscale/headscale.go | 43 ----------------------------
 2 files changed, 57 insertions(+), 43 deletions(-)

diff --git a/cmd/headscale/cli/root.go b/cmd/headscale/cli/root.go
index 99b15140..ab312c81 100644
--- a/cmd/headscale/cli/root.go
+++ b/cmd/headscale/cli/root.go
@@ -3,17 +3,74 @@ package cli
 import (
 	"fmt"
 	"os"
+	"runtime"
 
+	"github.com/rs/zerolog"
+	"github.com/rs/zerolog/log"
 	"github.com/spf13/cobra"
+	"github.com/spf13/viper"
+	"github.com/tcnksm/go-latest"
 )
 
+var cfgFile string = ""
+
 func init() {
+	cobra.OnInitialize(initConfig)
+	rootCmd.PersistentFlags().
+		StringVarP(&cfgFile, "config", "c", "", "config file (default is /etc/headscale/config.yaml)")
 	rootCmd.PersistentFlags().
 		StringP("output", "o", "", "Output format. Empty for human-readable, 'json', 'json-line' or 'yaml'")
 	rootCmd.PersistentFlags().
 		Bool("force", false, "Disable prompts and forces the execution")
 }
 
+func initConfig() {
+	if cfgFile != "" {
+		if err := LoadConfig(cfgFile, true); err != nil {
+			log.Fatal().Caller().Err(err)
+		}
+	} else {
+		if err := LoadConfig("", false); err != nil {
+			log.Fatal().Caller().Err(err)
+		}
+	}
+
+	machineOutput := HasMachineOutputFlag()
+
+	logLevel := viper.GetString("log_level")
+	level, err := zerolog.ParseLevel(logLevel)
+	if err != nil {
+		zerolog.SetGlobalLevel(zerolog.DebugLevel)
+	} else {
+		zerolog.SetGlobalLevel(level)
+	}
+
+	// If the user has requested a "machine" readable format,
+	// then disable login so the output remains valid.
+	if machineOutput {
+		zerolog.SetGlobalLevel(zerolog.Disabled)
+	}
+
+	if !viper.GetBool("disable_check_updates") && !machineOutput {
+		if (runtime.GOOS == "linux" || runtime.GOOS == "darwin") &&
+			Version != "dev" {
+			githubTag := &latest.GithubTag{
+				Owner:      "juanfont",
+				Repository: "headscale",
+			}
+			res, err := latest.Check(githubTag, Version)
+			if err == nil && res.Outdated {
+				//nolint
+				fmt.Printf(
+					"An updated version of Headscale has been found (%s vs. your current %s). Check it out https://github.com/juanfont/headscale/releases\n",
+					res.Current,
+					Version,
+				)
+			}
+		}
+	}
+}
+
 var rootCmd = &cobra.Command{
 	Use:   "headscale",
 	Short: "headscale - a Tailscale control server",
diff --git a/cmd/headscale/headscale.go b/cmd/headscale/headscale.go
index 28b5f2ed..40772a93 100644
--- a/cmd/headscale/headscale.go
+++ b/cmd/headscale/headscale.go
@@ -1,17 +1,13 @@
 package main
 
 import (
-	"fmt"
 	"os"
-	"runtime"
 	"time"
 
 	"github.com/efekarakus/termcolor"
 	"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"
 )
 
 func main() {
@@ -43,44 +39,5 @@ func main() {
 		NoColor:    !colors,
 	})
 
-	if err := cli.LoadConfig("", false); 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)
-	}
-
-	// If the user has requested a "machine" readable format,
-	// then disable login so the output remains valid.
-	if machineOutput {
-		zerolog.SetGlobalLevel(zerolog.Disabled)
-	}
-
-	if !viper.GetBool("disable_check_updates") && !machineOutput {
-		if (runtime.GOOS == "linux" || runtime.GOOS == "darwin") &&
-			cli.Version != "dev" {
-			githubTag := &latest.GithubTag{
-				Owner:      "juanfont",
-				Repository: "headscale",
-			}
-			res, err := latest.Check(githubTag, cli.Version)
-			if err == nil && res.Outdated {
-				//nolint
-				fmt.Printf(
-					"An updated version of Headscale has been found (%s vs. your current %s). Check it out https://github.com/juanfont/headscale/releases\n",
-					res.Current,
-					cli.Version,
-				)
-			}
-		}
-	}
-
 	cli.Execute()
 }