From 061efa182239fb962acca74e41e90d645a84a654 Mon Sep 17 00:00:00 2001
From: Juan Font Alonso <juanfontalonso@gmail.com>
Date: Sun, 17 Oct 2021 11:57:53 +0200
Subject: [PATCH] Do not include BaseDomain as full route in DNSConfig + code
 reorg

---
 api.go | 14 +++++++-------
 dns.go | 31 +++++++++++++++++++++++++------
 2 files changed, 32 insertions(+), 13 deletions(-)

diff --git a/api.go b/api.go
index 0aad5ee9..519dbeae 100644
--- a/api.go
+++ b/api.go
@@ -258,13 +258,13 @@ func (h *Headscale) getMapResponse(mKey wgkey.Key, req tailcfg.MapRequest, m *Ma
 		return nil, err
 	}
 
-	var dnsConfig *tailcfg.DNSConfig
-	if h.cfg.DNSConfig != nil && h.cfg.DNSConfig.Proxied { // if MagicDNS is enabled
-		// Only inject the Search Domain of the current namespace - shared nodes should use their full FQDN
-		dnsConfig = h.cfg.DNSConfig.Clone()
-		dnsConfig.Domains = append(dnsConfig.Domains, fmt.Sprintf("%s.%s", m.Namespace.Name, h.cfg.BaseDomain))
-	} else {
-		dnsConfig = h.cfg.DNSConfig
+	dnsConfig, err := h.getMapResponseDNSConfig(*m, peers)
+	if err != nil {
+		log.Error().
+			Str("func", "getMapResponse").
+			Err(err).
+			Msg("Failed generate the DNSConfig")
+		return nil, err
 	}
 
 	resp := tailcfg.MapResponse{
diff --git a/dns.go b/dns.go
index 353e10be..ae7d290c 100644
--- a/dns.go
+++ b/dns.go
@@ -4,7 +4,9 @@ import (
 	"fmt"
 	"strings"
 
+	"github.com/fatih/set"
 	"inet.af/netaddr"
+	"tailscale.com/tailcfg"
 	"tailscale.com/util/dnsname"
 )
 
@@ -29,15 +31,10 @@ import (
 // From the netmask we can find out the wildcard bits (the bits that are not set in the netmask).
 // This allows us to then calculate the subnets included in the subsequent class block and generate the entries.
 func generateMagicDNSRootDomains(ipPrefix netaddr.IPPrefix, baseDomain string) ([]dnsname.FQDN, error) {
-	base, err := dnsname.ToFQDN(baseDomain)
-	if err != nil {
-		return nil, err
-	}
-
 	// TODO(juanfont): we are not handing out IPv6 addresses yet
 	// and in fact this is Tailscale.com's range (note the fd7a:115c:a1e0: range in the fc00::/7 network)
 	ipv6base := dnsname.FQDN("0.e.1.a.c.5.1.1.a.7.d.f.ip6.arpa.")
-	fqdns := []dnsname.FQDN{base, ipv6base}
+	fqdns := []dnsname.FQDN{ipv6base}
 
 	// Conversion to the std lib net.IPnet, a bit easier to operate
 	netRange := ipPrefix.IPNet()
@@ -71,3 +68,25 @@ func generateMagicDNSRootDomains(ipPrefix netaddr.IPPrefix, baseDomain string) (
 	}
 	return fqdns, nil
 }
+
+func (h *Headscale) getMapResponseDNSConfig(m Machine, peers Machines) (*tailcfg.DNSConfig, error) {
+	var dnsConfig *tailcfg.DNSConfig
+	if h.cfg.DNSConfig != nil && h.cfg.DNSConfig.Proxied { // if MagicDNS is enabled
+		// Only inject the Search Domain of the current namespace - shared nodes should use their full FQDN
+		dnsConfig = h.cfg.DNSConfig.Clone()
+		dnsConfig.Domains = append(dnsConfig.Domains, fmt.Sprintf("%s.%s", m.Namespace.Name, h.cfg.BaseDomain))
+
+		namespaceSet := set.New(set.ThreadSafe)
+		namespaceSet.Add(m.Namespace)
+		for _, p := range peers {
+			namespaceSet.Add(p.Namespace)
+		}
+		for _, namespace := range namespaceSet.List() {
+			dnsRoute := dnsname.FQDN(fmt.Sprintf("%s.%s", namespace.(Namespace).Name, h.cfg.BaseDomain))
+			dnsConfig.Routes[dnsRoute.WithoutTrailingDot()] = nil
+		}
+	} else {
+		dnsConfig = h.cfg.DNSConfig
+	}
+	return dnsConfig, nil
+}