From 3f0639c87ddbb86fd9af3d210eafa98b80301ad1 Mon Sep 17 00:00:00 2001
From: Grigoriy Mikhalkin <grigoriymikhalkin@gmail.com>
Date: Mon, 11 Jul 2022 20:33:24 +0200
Subject: [PATCH] graceful shutdown lint fixes

---
 app.go    | 32 ++++++++++++++++++--------------
 config.go | 13 +++++++++----
 poll.go   |  4 ++--
 3 files changed, 29 insertions(+), 20 deletions(-)

diff --git a/app.go b/app.go
index de6ef669..5f5e2611 100644
--- a/app.go
+++ b/app.go
@@ -94,8 +94,8 @@ type Headscale struct {
 
 	ipAllocationMutex sync.Mutex
 
-	shutdownChan chan struct{}
-	wg           sync.WaitGroup
+	shutdownChan       chan struct{}
+	pollNetMapStreamWG sync.WaitGroup
 }
 
 // Look up the TLS constant relative to user-supplied TLS client
@@ -148,13 +148,13 @@ func NewHeadscale(cfg *Config) (*Headscale, error) {
 	)
 
 	app := Headscale{
-		cfg:               cfg,
-		dbType:            cfg.DBtype,
-		dbString:          dbString,
-		privateKey:        privKey,
-		aclRules:          tailcfg.FilterAllowAll, // default allowall
-		registrationCache: registrationCache,
-		wg:                sync.WaitGroup{},
+		cfg:                cfg,
+		dbType:             cfg.DBtype,
+		dbString:           dbString,
+		privateKey:         privKey,
+		aclRules:           tailcfg.FilterAllowAll, // default allowall
+		registrationCache:  registrationCache,
+		pollNetMapStreamWG: sync.WaitGroup{},
 	}
 
 	err = app.initDB()
@@ -672,7 +672,7 @@ func (h *Headscale) Serve() error {
 		syscall.SIGTERM,
 		syscall.SIGQUIT,
 		syscall.SIGHUP)
-	sig_func := func(c chan os.Signal) {
+	sigFunc := func(c chan os.Signal) {
 		// Wait for a SIGINT or SIGKILL:
 		for {
 			sig := <-c
@@ -703,7 +703,7 @@ func (h *Headscale) Serve() error {
 					Msg("Received signal to stop, shutting down gracefully")
 
 				close(h.shutdownChan)
-				h.wg.Wait()
+				h.pollNetMapStreamWG.Wait()
 
 				// Gracefully shut down servers
 				ctx, cancel := context.WithTimeout(context.Background(), HTTPShutdownTimeout)
@@ -747,7 +747,11 @@ func (h *Headscale) Serve() error {
 			}
 		}
 	}
-	errorGroup.Go(func() error { sig_func(sigc); return nil })
+	errorGroup.Go(func() error {
+		sigFunc(sigc)
+
+		return nil
+	})
 
 	return errorGroup.Wait()
 }
@@ -771,13 +775,13 @@ func (h *Headscale) getTLSSettings() (*tls.Config, error) {
 		}
 
 		switch h.cfg.TLS.LetsEncrypt.ChallengeType {
-		case "TLS-ALPN-01":
+		case tlsALPN01ChallengeType:
 			// Configuration via autocert with TLS-ALPN-01 (https://tools.ietf.org/html/rfc8737)
 			// The RFC requires that the validation is done on port 443; in other words, headscale
 			// must be reachable on port 443.
 			return certManager.TLSConfig(), nil
 
-		case "HTTP-01":
+		case http01ChallengeType:
 			// Configuration via autocert with HTTP-01. This requires listening on
 			// port 80 for the certificate validation in addition to the headscale
 			// service, which can be configured to run on any other port.
diff --git a/config.go b/config.go
index 6789f6f0..69358401 100644
--- a/config.go
+++ b/config.go
@@ -18,6 +18,11 @@ import (
 	"tailscale.com/types/dnstype"
 )
 
+const (
+	tlsALPN01ChallengeType = "TLS-ALPN-01"
+	http01ChallengeType    = "HTTP-01"
+)
+
 // Config contains the initial Headscale configuration.
 type Config struct {
 	ServerURL                      string
@@ -136,7 +141,7 @@ func LoadConfig(path string, isFile bool) error {
 	viper.AutomaticEnv()
 
 	viper.SetDefault("tls_letsencrypt_cache_dir", "/var/www/.cache")
-	viper.SetDefault("tls_letsencrypt_challenge_type", "HTTP-01")
+	viper.SetDefault("tls_letsencrypt_challenge_type", http01ChallengeType)
 	viper.SetDefault("tls_client_auth_mode", "relaxed")
 
 	viper.SetDefault("log_level", "info")
@@ -179,15 +184,15 @@ func LoadConfig(path string, isFile bool) error {
 	}
 
 	if (viper.GetString("tls_letsencrypt_hostname") != "") &&
-		(viper.GetString("tls_letsencrypt_challenge_type") == "TLS-ALPN-01") &&
+		(viper.GetString("tls_letsencrypt_challenge_type") == tlsALPN01ChallengeType) &&
 		(!strings.HasSuffix(viper.GetString("listen_addr"), ":443")) {
 		// this is only a warning because there could be something sitting in front of headscale that redirects the traffic (e.g. an iptables rule)
 		log.Warn().
 			Msg("Warning: when using tls_letsencrypt_hostname with TLS-ALPN-01 as challenge type, headscale must be reachable on port 443, i.e. listen_addr should probably end in :443")
 	}
 
-	if (viper.GetString("tls_letsencrypt_challenge_type") != "HTTP-01") &&
-		(viper.GetString("tls_letsencrypt_challenge_type") != "TLS-ALPN-01") {
+	if (viper.GetString("tls_letsencrypt_challenge_type") != http01ChallengeType) &&
+		(viper.GetString("tls_letsencrypt_challenge_type") != tlsALPN01ChallengeType) {
 		errorText += "Fatal config error: the only supported values for tls_letsencrypt_challenge_type are HTTP-01 and TLS-ALPN-01\n"
 	}
 
diff --git a/poll.go b/poll.go
index 94941aa3..b9a757aa 100644
--- a/poll.go
+++ b/poll.go
@@ -290,8 +290,8 @@ func (h *Headscale) PollNetMapStream(
 	keepAliveChan chan []byte,
 	updateChan chan struct{},
 ) {
-	h.wg.Add(1)
-	defer h.wg.Done()
+	h.pollNetMapStreamWG.Add(1)
+	defer h.pollNetMapStreamWG.Done()
 
 	ctx := context.WithValue(req.Context(), machineNameContextKey, machine.Hostname)