From 987bbee1dbe9de611bc409e0f39647e460fda25e Mon Sep 17 00:00:00 2001
From: Kristoffer Dalby <kradalby@kradalby.no>
Date: Tue, 24 Aug 2021 07:09:47 +0100
Subject: [PATCH] Add DNSConfig field to configuration

---
 app.go                     |  2 ++
 cmd/headscale/cli/utils.go | 38 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 40 insertions(+)

diff --git a/app.go b/app.go
index e5f44103..c903d83f 100644
--- a/app.go
+++ b/app.go
@@ -43,6 +43,8 @@ type Config struct {
 
 	TLSCertPath string
 	TLSKeyPath  string
+
+	DNSConfig *tailcfg.DNSConfig
 }
 
 // Headscale represents the base app of the service
diff --git a/cmd/headscale/cli/utils.go b/cmd/headscale/cli/utils.go
index 7e7e8f96..7ba78648 100644
--- a/cmd/headscale/cli/utils.go
+++ b/cmd/headscale/cli/utils.go
@@ -41,6 +41,8 @@ func LoadConfig(path string) error {
 
 	viper.SetDefault("log_level", "info")
 
+	viper.SetDefault("dns_config", nil)
+
 	err := viper.ReadInConfig()
 	if err != nil {
 		return fmt.Errorf("Fatal error reading config file: %s \n", err)
@@ -70,6 +72,40 @@ func LoadConfig(path string) error {
 	} else {
 		return nil
 	}
+
+}
+
+func getDNSConfig() *tailcfg.DNSConfig {
+	if viper.IsSet("dns_config") {
+		dnsConfig := &tailcfg.DNSConfig{}
+
+		if viper.IsSet("dns_config.nameservers") {
+			nameserversStr := viper.GetStringSlice("dns_config.nameservers")
+
+			nameservers := make([]netaddr.IP, len(nameserversStr))
+
+			for index, nameserverStr := range nameserversStr {
+				nameserver, err := netaddr.ParseIP(nameserverStr)
+				if err != nil {
+					log.Error().
+						Str("func", "getDNSConfig").
+						Err(err).
+						Msgf("Could not parse nameserver IP: %s", nameserverStr)
+				}
+
+				nameservers[index] = nameserver
+			}
+
+			dnsConfig.Nameservers = nameservers
+		}
+		if viper.IsSet("dns_config.domains") {
+			dnsConfig.Domains = viper.GetStringSlice("dns_config.domains")
+		}
+
+		return dnsConfig
+	}
+
+	return nil
 }
 
 func absPath(path string) string {
@@ -126,6 +162,8 @@ func getHeadscaleApp() (*headscale.Headscale, error) {
 
 		TLSCertPath: absPath(viper.GetString("tls_cert_path")),
 		TLSKeyPath:  absPath(viper.GetString("tls_key_path")),
+
+		DNSConfig: getDNSConfig(),
 	}
 
 	h, err := headscale.NewHeadscale(cfg)