From fc5153af3e92ff029a964b10af8f24ec0bb9efd4 Mon Sep 17 00:00:00 2001
From: Juan Font Alonso <juanfontalonso@gmail.com>
Date: Sat, 9 Oct 2021 12:22:13 +0200
Subject: [PATCH] Generate MagicDNS search domains for any tailnet range

---
 app.go |  4 ++--
 dns.go | 27 +++++++++++++++++++++------
 2 files changed, 23 insertions(+), 8 deletions(-)

diff --git a/app.go b/app.go
index 31e306dc..2afc75b6 100644
--- a/app.go
+++ b/app.go
@@ -12,7 +12,7 @@ import (
 	"github.com/rs/zerolog/log"
 
 	"github.com/gin-gonic/gin"
-	"github.com/zsais/go-gin-prometheus"
+	ginprometheus "github.com/zsais/go-gin-prometheus"
 	"golang.org/x/crypto/acme"
 	"golang.org/x/crypto/acme/autocert"
 	"gorm.io/gorm"
@@ -111,7 +111,7 @@ func NewHeadscale(cfg Config) (*Headscale, error) {
 	}
 
 	if h.cfg.DNSConfig != nil && h.cfg.DNSConfig.Proxied { // if MagicDNS
-		magicDNSDomains, err := h.generateMagicDNSRootDomains()
+		magicDNSDomains, err := generateMagicDNSRootDomains(h.cfg.IPPrefix, h.cfg.BaseDomain)
 		if err != nil {
 			return nil, err
 		}
diff --git a/dns.go b/dns.go
index 74a85aee..c53849c9 100644
--- a/dns.go
+++ b/dns.go
@@ -2,12 +2,14 @@ package headscale
 
 import (
 	"fmt"
+	"strings"
 
+	"inet.af/netaddr"
 	"tailscale.com/util/dnsname"
 )
 
-func (h *Headscale) generateMagicDNSRootDomains() (*[]dnsname.FQDN, error) {
-	base, err := dnsname.ToFQDN(h.cfg.BaseDomain)
+func generateMagicDNSRootDomains(ipPrefix netaddr.IPPrefix, baseDomain string) (*[]dnsname.FQDN, error) {
+	base, err := dnsname.ToFQDN(baseDomain)
 	if err != nil {
 		return nil, err
 	}
@@ -17,14 +19,27 @@ func (h *Headscale) generateMagicDNSRootDomains() (*[]dnsname.FQDN, error) {
 	ipv6base := dnsname.FQDN("0.e.1.a.c.5.1.1.a.7.d.f.ip6.arpa.")
 	fqdns := []dnsname.FQDN{base, ipv6base}
 
-	for i := 64; i <= 127; i++ {
-		fqdn, err := dnsname.ToFQDN(fmt.Sprintf("%d.100.in-addr.arpa.", i))
+	netRange := ipPrefix.IPNet()
+	maskBits, _ := netRange.Mask.Size()
+
+	lastByte := maskBits / 8
+	unmaskedBits := 8 - maskBits%8
+	min := uint(netRange.IP[lastByte])
+	max := uint((min + 1<<uint(unmaskedBits)) - 1)
+
+	rdnsSlice := []string{}
+	for i := lastByte - 1; i >= 0; i-- {
+		rdnsSlice = append(rdnsSlice, fmt.Sprintf("%d", netRange.IP[i]))
+	}
+	rdnsSlice = append(rdnsSlice, "in-addr.arpa.")
+	rdnsBase := strings.Join(rdnsSlice, ".")
+
+	for i := min; i <= max; i++ {
+		fqdn, err := dnsname.ToFQDN(fmt.Sprintf("%d.%s", i, rdnsBase))
 		if err != nil {
-			// TODO: propagate error
 			continue
 		}
 		fqdns = append(fqdns, fqdn)
 	}
-
 	return &fqdns, nil
 }