From 8d0c8ad4381cbf23e1e6653561f7f14a26544d34 Mon Sep 17 00:00:00 2001
From: Simon Zolin <s.zolin@adguard.com>
Date: Tue, 18 Aug 2020 12:36:52 +0300
Subject: [PATCH] + DNS: resolve host names to IP addresses leased by AGH DHCP
 server

Close #1956

Squashed commit of the following:

commit 21f11632c871e9c17faa77f9cd6a7aa836559779
Author: Simon Zolin <s.zolin@adguard.com>
Date:   Mon Aug 17 19:54:24 2020 +0300

    + DNS: resolve host names to IP addresses leased by AGH DHCP server
---
 dnsforward/dnsforward.go |  3 ++
 dnsforward/handle_dns.go | 65 +++++++++++++++++++++++++++++++++++++++-
 2 files changed, 67 insertions(+), 1 deletion(-)

diff --git a/dnsforward/dnsforward.go b/dnsforward/dnsforward.go
index 95db8967..e73824b0 100644
--- a/dnsforward/dnsforward.go
+++ b/dnsforward/dnsforward.go
@@ -51,6 +51,9 @@ type Server struct {
 	stats      stats.Stats
 	access     *accessCtx
 
+	tableHostToIP     map[string]net.IP // "hostname -> IP" table for internal addresses (DHCP)
+	tableHostToIPLock sync.Mutex
+
 	tablePTR     map[string]string // "IP -> hostname" table for reverse lookup
 	tablePTRLock sync.Mutex
 
diff --git a/dnsforward/handle_dns.go b/dnsforward/handle_dns.go
index 462f3750..6860dcf3 100644
--- a/dnsforward/handle_dns.go
+++ b/dnsforward/handle_dns.go
@@ -1,6 +1,7 @@
 package dnsforward
 
 import (
+	"net"
 	"strings"
 	"time"
 
@@ -42,6 +43,7 @@ func (s *Server) handleDNSRequest(_ *proxy.Proxy, d *proxy.DNSContext) error {
 	type modProcessFunc func(ctx *dnsContext) int
 	mods := []modProcessFunc{
 		processInitial,
+		processInternalHosts,
 		processInternalIPAddrs,
 		processFilteringBeforeRequest,
 		processUpstream,
@@ -102,20 +104,81 @@ func (s *Server) onDHCPLeaseChanged(flags int) {
 		return
 	}
 
+	hostToIP := make(map[string]net.IP)
 	m := make(map[string]string)
+
 	ll := s.dhcpServer.Leases(dhcpd.LeasesAll)
+
 	for _, l := range ll {
 		if len(l.Hostname) == 0 {
 			continue
 		}
+
 		m[l.IP.String()] = l.Hostname
+
+		ip := make(net.IP, 4)
+		copy(ip, l.IP.To4())
+		hostToIP[l.Hostname] = ip
 	}
-	log.Debug("DNS: added %d PTR entries from DHCP", len(m))
+
+	log.Debug("DNS: added %d A/PTR entries from DHCP", len(m))
+
+	s.tableHostToIPLock.Lock()
+	s.tableHostToIP = hostToIP
+	s.tableHostToIPLock.Unlock()
+
 	s.tablePTRLock.Lock()
 	s.tablePTR = m
 	s.tablePTRLock.Unlock()
 }
 
+// Respond to A requests if the target host name is associated with a lease from our DHCP server
+func processInternalHosts(ctx *dnsContext) int {
+	s := ctx.srv
+	req := ctx.proxyCtx.Req
+	if !(req.Question[0].Qtype == dns.TypeA || req.Question[0].Qtype == dns.TypeAAAA) {
+		return resultDone
+	}
+
+	host := req.Question[0].Name
+	host = strings.ToLower(host)
+	if !strings.HasSuffix(host, ".lan.") {
+		return resultDone
+	}
+	host = strings.TrimSuffix(host, ".lan.")
+
+	s.tableHostToIPLock.Lock()
+	if s.tableHostToIP == nil {
+		s.tableHostToIPLock.Unlock()
+		return resultDone
+	}
+	ip, ok := s.tableHostToIP[host]
+	s.tableHostToIPLock.Unlock()
+	if !ok {
+		return resultDone
+	}
+
+	log.Debug("DNS: internal record: %s -> %s", req.Question[0].Name, ip.String())
+
+	resp := s.makeResponse(req)
+
+	if req.Question[0].Qtype == dns.TypeA {
+		a := &dns.A{}
+		a.Hdr = dns.RR_Header{
+			Name:   req.Question[0].Name,
+			Rrtype: dns.TypeA,
+			Ttl:    s.conf.BlockedResponseTTL,
+			Class:  dns.ClassINET,
+		}
+		a.A = make([]byte, 4)
+		copy(a.A, ip)
+		resp.Answer = append(resp.Answer, a)
+	}
+
+	ctx.proxyCtx.Res = resp
+	return resultDone
+}
+
 // Respond to PTR requests if the target IP address is leased by our DHCP server
 func processInternalIPAddrs(ctx *dnsContext) int {
 	s := ctx.srv