Implement log throttling, loop prevention, and empty upstream handling in DNS discovery server

2026-08-18 00:26:29 +00:00 · 2026-02-16 21:02:42 +01:00
parent 7d140b3e2a
commit 025e15d65c
3 changed files with 161 additions and 6 deletions
@@ -305,7 +305,7 @@ When enabled, the DNS server:
 You can enable and configure the DNS server via the Web UI or environment variables:
 - `ENABLE_DNS_DISCOVERY=true`: Turns on the DNS server.
 - `DNS_BIND_ADDR=:53`: The port to listen on (requires root privileges for port 53).
- `DNS_UPSTREAM=1.1.1.1`: Your preferred upstream DNS provider.
+- `DNS_UPSTREAM=1.1.1.1`: Your preferred upstream DNS provider. **Note:** Ensure this is not set to the same address as the DNS server itself (loopback or local IP) to avoid forwarding loops. The server includes built-in loop prevention, but misconfiguration will cause forwarding to fail. If left empty, forwarding will be disabled and non-intercepted queries will return a failure.

 #### Manual Discovery via DNS
 Even without migrating a device, you can use the DNS server to discover what a device is querying by manually setting your router's DNS or the device's DNS to point to the AfterTouch service.
@@ -27,6 +27,13 @@ type DNSDiscovery struct {
 	// Servers for Shutdown
 	udpServer *dns.Server
 	tcpServer *dns.Server
+
+	// Address for loop prevention
+	bindAddr string
+
+	// Log throttling
+	lastLog   map[string]time.Time
+	lastLogMu sync.Mutex
 }

 // DiscoveredHost represents a host discovered via DNS queries.
@@ -46,6 +53,7 @@ func NewDNSDiscovery(upstreamDNS, serviceIP string) *DNSDiscovery {
 		upstreamDNS: upstreamDNS,
 		serviceIP:   serviceIP,
 		discovered:  make(map[string]*DiscoveredHost),
+		lastLog:     make(map[string]time.Time),
 	}
 }

@@ -72,14 +80,38 @@ func (d *DNSDiscovery) ServeDNS(w dns.ResponseWriter, r *dns.Msg) {
 	if isIntercepted {
 		// Return your service IP
 		d.respondWithIP(w, r, d.serviceIP)
-		log.Printf("[DNS] Intercepting %s (type %d) -> %s", hostname, q.Qtype, d.serviceIP)
+		d.throttledLog(fmt.Sprintf("[DNS] Intercepting %s (type %d) -> %s", hostname, q.Qtype, d.serviceIP))
 	} else {
 		// Forward to real DNS
-		log.Printf("[DNS] Forwarding %s (type %d) to %s", hostname, q.Qtype, d.upstreamDNS)
+		if d.upstreamDNS == "" {
+			d.throttledLog(fmt.Sprintf("[DNS ERROR] No upstream DNS configured, cannot forward %s", hostname))
+
+			m := new(dns.Msg)
+			m.SetReply(r)
+			m.Rcode = dns.RcodeServerFailure
+			_ = w.WriteMsg(m)
+
+			return
+		}
+
+		d.throttledLog(fmt.Sprintf("[DNS] Forwarding %s (type %d) to %s", hostname, q.Qtype, d.upstreamDNS))
 		d.forward(w, r)
 	}
 }

+func (d *DNSDiscovery) throttledLog(msg string) {
+	d.lastLogMu.Lock()
+	defer d.lastLogMu.Unlock()
+
+	now := time.Now()
+	if last, ok := d.lastLog[msg]; ok && now.Sub(last) < 10*time.Second {
+		return
+	}
+
+	d.lastLog[msg] = now
+	log.Print(msg)
+}
+
 // recordQuery logs a DNS query and updates the internal state.
 func (d *DNSDiscovery) recordQuery(hostname string, isIntercepted bool, remoteAddr string) {
 	d.mu.Lock()
@@ -183,8 +215,10 @@ func (d *DNSDiscovery) forward(w dns.ResponseWriter, r *dns.Msg) {
 		return
 	}

+	q := r.Question[0]
+
 	// Don't forward PTR queries for our own service IP to avoid loops or slow timeouts
-	if r.Question[0].Qtype == dns.TypePTR {
+	if q.Qtype == dns.TypePTR {
 		m := new(dns.Msg)
 		m.SetReply(r)

@@ -196,16 +230,30 @@ func (d *DNSDiscovery) forward(w dns.ResponseWriter, r *dns.Msg) {
 		return
 	}

-	c := new(dns.Client)
 	// Add port 53 if not present
 	upstream := d.upstreamDNS
 	if !strings.Contains(upstream, ":") {
 		upstream += ":53"
 	}

+	// Loop prevention: don't forward to ourselves
+	if upstream == d.bindAddr || (strings.HasPrefix(upstream, "127.0.0.1:") && strings.HasSuffix(d.bindAddr, upstream[9:])) {
+		d.throttledLog(fmt.Sprintf("[DNS ERROR] Refusing to forward %s to ourselves (%s)", q.Name, upstream))
+
+		m := new(dns.Msg)
+		m.SetReply(r)
+		m.Rcode = dns.RcodeServerFailure
+		_ = w.WriteMsg(m)
+
+		return
+	}
+
+	c := new(dns.Client)
+	c.Timeout = 2 * time.Second
+
 	in, _, err := c.Exchange(r, upstream)
 	if err != nil {
-		log.Printf("[DNS ERROR] Forward failed for %s (type %d): %v", r.Question[0].Name, r.Question[0].Qtype, err)
+		d.throttledLog(fmt.Sprintf("[DNS ERROR] Forward failed for %s (type %d): %v", q.Name, q.Qtype, err))
 		// Return a failure response instead of just dropping
 		m := new(dns.Msg)
 		m.SetReply(r)
@@ -267,6 +315,7 @@ func (d *DNSDiscovery) Start(addr string) error {
 	mux.HandleFunc(".", d.ServeDNS)

 	d.mu.Lock()
+	d.bindAddr = addr
 	d.udpServer = &dns.Server{
 		Addr:    addr,
 		Net:     "udp",
@@ -1,7 +1,9 @@
 package discovery

 import (
+	"log"
 	"net"
+	"strings"
 	"testing"
 	"time"

@@ -192,3 +194,107 @@ func (m *mockResponseWriter) Close() error                { return nil }
 func (m *mockResponseWriter) TsigStatus() error           { return nil }
 func (m *mockResponseWriter) TsigTimersOnly(bool)         {}
 func (m *mockResponseWriter) Hijack()                     {}
+
+func TestDNSDiscovery_LogThrottling(t *testing.T) {
+	d := NewDNSDiscovery("8.8.8.8", "192.168.1.100")
+
+	// Capture log output
+	var logBuf strings.Builder
+	oldOutput := log.Writer()
+	log.SetOutput(&logBuf)
+	defer log.SetOutput(oldOutput)
+
+	msg := "Test log message"
+	d.throttledLog(msg)
+	d.throttledLog(msg)
+	d.throttledLog(msg)
+
+	count := strings.Count(logBuf.String(), msg)
+	if count != 1 {
+		t.Errorf("Expected log message to appear once due to throttling, but appeared %d times", count)
+	}
+
+	// Advance time by 11 seconds to bypass throttling
+	d.lastLogMu.Lock()
+	d.lastLog[msg] = time.Now().Add(-11 * time.Second)
+	d.lastLogMu.Unlock()
+
+	d.throttledLog(msg)
+	count = strings.Count(logBuf.String(), msg)
+	if count != 2 {
+		t.Errorf("Expected log message to appear twice after advancing time, but appeared %d times", count)
+	}
+}
+
+func TestDNSDiscovery_LoopPrevention(t *testing.T) {
+	serviceIP := "192.168.1.100"
+	bindAddr := "127.0.0.1:53"
+	upstreamDNS := "127.0.0.1:53"
+	d := NewDNSDiscovery(upstreamDNS, serviceIP)
+	d.bindAddr = bindAddr
+
+	// Capture log output to avoid panic if it's being throttled/logged
+	var logBuf strings.Builder
+	oldOutput := log.Writer()
+	log.SetOutput(&logBuf)
+	defer log.SetOutput(oldOutput)
+
+	m := new(dns.Msg)
+	m.SetQuestion("google.com.", dns.TypeA)
+
+	rw := &mockResponseWriter{}
+	d.forward(rw, m)
+
+	if rw.msg == nil {
+		t.Fatal("Expected a response message")
+	}
+
+	if rw.msg.Rcode != dns.RcodeServerFailure {
+		t.Errorf("Expected RcodeServerFailure (2), got %d", rw.msg.Rcode)
+	}
+}
+
+func TestDNSDiscovery_EmptyUpstream(t *testing.T) {
+	serviceIP := "192.168.1.100"
+	upstreamDNS := "" // Empty upstream
+	d := NewDNSDiscovery(upstreamDNS, serviceIP)
+	d.bindAddr = ":53"
+
+	m := new(dns.Msg)
+	m.SetQuestion("google.com.", dns.TypeA)
+
+	rw := &mockResponseWriter{}
+	d.ServeDNS(rw, m)
+
+	if rw.msg == nil {
+		t.Fatal("Expected a response message, got nil")
+	}
+
+	if rw.msg.Rcode != dns.RcodeServerFailure {
+		t.Errorf("Expected RcodeServerFailure (2) for empty upstream, got %d", rw.msg.Rcode)
+	}
+}
+
+func TestDNSDiscovery_ForwardTimeout(t *testing.T) {
+	serviceIP := "192.168.1.100"
+	// Use an IP that is unroutable or doesn't exist on the network to ensure timeout
+	upstreamDNS := "192.0.2.1:53" // TEST-NET-1, usually non-routable
+	d := NewDNSDiscovery(upstreamDNS, serviceIP)
+
+	m := new(dns.Msg)
+	m.SetQuestion("google.com.", dns.TypeA)
+
+	rw := &mockResponseWriter{}
+
+	start := time.Now()
+	d.forward(rw, m)
+	duration := time.Since(start)
+
+	if duration < 2*time.Second {
+		t.Errorf("Expected forward to take at least 2 seconds (timeout), but took %v", duration)
+	}
+
+	if rw.msg == nil || rw.msg.Rcode != dns.RcodeServerFailure {
+		t.Errorf("Expected RcodeServerFailure after timeout")
+	}
+}