Merge branch 'ipset-subs' of git://github.com/dsheets/AdGuardHome into dsheets-ipset-subs

dnsforward/ipset.go

@@ -89,10 +89,49 @@ func (c *ipsetCtx) getIP(rr dns.RR) net.IP {
 	}
 }
 
-// Add IP addresses of the specified in configuration domain names to an ipset list
+// Find the ipsets for a given host (accounting for subdomain wildcards)
-func (c *ipsetCtx) process(ctx *dnsContext) int {
+func (c *ipsetCtx) getIpsetNames(host string) ([]string, bool) {
+	var ipsetNames []string
+	var found bool
+
+	// search for matching ipset hosts starting with most specific subdomain
+	i := 0
+	for i != -1 {
+		host = host[i:]
+
+		ipsetNames, found = c.ipsetList[host]
+		if found {
+			break
+		}
+
+		// move slice up to the parent domain
+		i = strings.Index(host, ".")
+		if i != -1 {
+			i++
+		}
+	}
+
+	return ipsetNames, found
+}
+
+func addToIpset(host string, ipsetName string, ipStr string) {
+	code, out, err := util.RunCommand("ipset", "add", ipsetName, ipStr)
+	if err != nil {
+		log.Info("IPSET: %s(%s) -> %s: %s", host, ipStr, ipsetName, err)
+		return
+	}
+	if code != 0 {
+		log.Info("IPSET: ipset add:  code:%d  output:'%s'", code, out)
+		return
+	}
+	log.Debug("IPSET: added %s(%s) -> %s", host, ipStr, ipsetName)
+}
+
+// Compute which addresses to add to which ipsets for a particular DNS query response
+// Call addMember for each (host, ipset, ip) triple
+func (c *ipsetCtx) processMembers(ctx *dnsContext, addMember func(string, string, string)) int {
 	req := ctx.proxyCtx.Req
-	if !(req.Question[0].Qtype == dns.TypeA ||
+	if req == nil || !(req.Question[0].Qtype == dns.TypeA ||
 		req.Question[0].Qtype == dns.TypeAAAA) ||
 		!ctx.responseFromUpstream {
 		return resultDone
@@ -101,33 +140,31 @@ func (c *ipsetCtx) process(ctx *dnsContext) int {
 	host := req.Question[0].Name
 	host = strings.TrimSuffix(host, ".")
 	host = strings.ToLower(host)
-	ipsetNames, found := c.ipsetList[host]
+	ipsetNames, found := c.getIpsetNames(host)
 	if !found {
 		return resultDone
 	}
 
 	log.Debug("IPSET: found ipsets %v for host %s", ipsetNames, host)
 
-	for _, it := range ctx.proxyCtx.Res.Answer {
+	if ctx.proxyCtx.Res != nil {
-		ip := c.getIP(it)
+		for _, it := range ctx.proxyCtx.Res.Answer {
-		if ip == nil {
+			ip := c.getIP(it)
-			continue
+			if ip == nil {
-		}
+				continue
+			}
 
-		ipStr := ip.String()
+			ipStr := ip.String()
-		for _, name := range ipsetNames {
+			for _, name := range ipsetNames {
-			code, out, err := util.RunCommand("ipset", "add", name, ipStr)
+				addMember(host, name, ipStr)
-			if err != nil {
-				log.Info("IPSET: %s(%s) -> %s: %s", host, ipStr, name, err)
-				continue
 			}
-			if code != 0 {
-				log.Info("IPSET: ipset add:  code:%d  output:'%s'", code, out)
-				continue
-			}
-			log.Debug("IPSET: added %s(%s) -> %s", host, ipStr, name)
 		}
 	}
 
 	return resultDone
 }
+
+// Add IP addresses of the specified in configuration domain names to an ipset list
+func (c *ipsetCtx) process(ctx *dnsContext) int {
+	return c.processMembers(ctx, addToIpset)
+}

dnsforward/ipset_test.go

@@ -1,6 +1,7 @@
 package dnsforward
 
 import (
+	"net"
 	"testing"
 
 	"github.com/AdguardTeam/dnsproxy/proxy"
@@ -8,14 +9,92 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
-func TestIPSET(t *testing.T) {
+var s Server
-	s := Server{}
+var c ipsetCtx
-	s.conf.IPSETList = append(s.conf.IPSETList, "HOST.com/name")
+var ctx *dnsContext
-	s.conf.IPSETList = append(s.conf.IPSETList, "host2.com,host3.com/name23")
+
-	s.conf.IPSETList = append(s.conf.IPSETList, "host4.com/name4,name41")
+type Binding struct {
-	c := ipsetCtx{}
+	host  string
+	ipset string
+	ipStr string
+}
+
+var b map[Binding]int
+
+func setup() {
+	s = Server{}
+	s.conf.IPSETList = []string{
+		"HOST.com/name",
+		"host2.com,host3.com/name23",
+		"host4.com/name4,name41",
+		"sub.host4.com/subhost4",
+	}
+
+	c = ipsetCtx{}
 	c.init(s.conf.IPSETList)
 
+	ctx = &dnsContext{
+		srv: &s,
+	}
+	ctx.responseFromUpstream = true
+	ctx.proxyCtx = &proxy.DNSContext{}
+
+	b = make(map[Binding]int)
+}
+
+func makeReq(fqdn string, qtype uint16) *dns.Msg {
+	return &dns.Msg{
+		Question: []dns.Question{
+			{
+				Name:  fqdn,
+				Qtype: qtype,
+			},
+		},
+	}
+}
+
+func makeReqA(fqdn string) *dns.Msg {
+	return makeReq(fqdn, dns.TypeA)
+}
+
+func makeReqAAAA(fqdn string) *dns.Msg {
+	return makeReq(fqdn, dns.TypeAAAA)
+}
+
+func makeA(fqdn string, ip net.IP) *dns.A {
+	return &dns.A{
+		Hdr: dns.RR_Header{Name: fqdn, Rrtype: dns.TypeA, Class: dns.ClassINET, Ttl: 0},
+		A:   ip,
+	}
+}
+
+func makeAAAA(fqdn string, ip net.IP) *dns.AAAA {
+	return &dns.AAAA{
+		Hdr:  dns.RR_Header{Name: fqdn, Rrtype: dns.TypeAAAA, Class: dns.ClassINET, Ttl: 0},
+		AAAA: ip,
+	}
+}
+
+func makeCNAME(fqdn string, cnameFqdn string) *dns.CNAME {
+	return &dns.CNAME{
+		Hdr:    dns.RR_Header{Name: fqdn, Rrtype: dns.TypeCNAME, Class: dns.ClassINET, Ttl: 0},
+		Target: cnameFqdn,
+	}
+}
+
+func addToBindings(host string, ipset string, ipStr string) {
+	binding := Binding{host, ipset, ipStr}
+	count := b[binding]
+	b[binding] = count + 1
+}
+
+func doProcess(t *testing.T) {
+	assert.Equal(t, resultDone, c.processMembers(ctx, addToBindings))
+}
+
+func TestIpsetParsing(t *testing.T) {
+	setup()
+
 	assert.Equal(t, "name", c.ipsetList["host.com"][0])
 	assert.Equal(t, "name23", c.ipsetList["host2.com"][0])
 	assert.Equal(t, "name23", c.ipsetList["host3.com"][0])
@@ -24,18 +103,97 @@ func TestIPSET(t *testing.T) {
 
 	_, ok := c.ipsetList["host0.com"]
 	assert.False(t, ok)
+}
 
-	ctx := &dnsContext{
+func TestIpsetNoQuestion(t *testing.T) {
-		srv: &s,
+	setup()
-	}
+
-	ctx.proxyCtx = &proxy.DNSContext{}
+	doProcess(t)
-	ctx.proxyCtx.Req = &dns.Msg{
+	assert.Equal(t, 0, len(b))
-		Question: []dns.Question{
+}
-			{
+
-				Name:  "host.com.",
+func TestIpsetNoAnswer(t *testing.T) {
-				Qtype: dns.TypeA,
+	setup()
-			},
+
+	ctx.proxyCtx.Req = makeReqA("HOST4.COM.")
+
+	doProcess(t)
+	assert.Equal(t, 0, len(b))
+}
+
+func TestIpsetCache(t *testing.T) {
+	setup()
+
+	ctx.proxyCtx.Req = makeReqA("HOST4.COM.")
+	ctx.proxyCtx.Res = &dns.Msg{
+		Answer: []dns.RR{
+			makeA("HOST4.COM.", net.IPv4(127, 0, 0, 1)),
+			makeAAAA("HOST4.COM.", net.IPv6loopback),
 		},
 	}
-	assert.Equal(t, resultDone, c.process(ctx))
+
+	doProcess(t)
+
+	assert.Equal(t, 1, b[Binding{"host4.com", "name4", "127.0.0.1"}])
+	assert.Equal(t, 1, b[Binding{"host4.com", "name41", "127.0.0.1"}])
+	assert.Equal(t, 1, b[Binding{"host4.com", "name4", net.IPv6loopback.String()}])
+	assert.Equal(t, 1, b[Binding{"host4.com", "name41", net.IPv6loopback.String()}])
+	assert.Equal(t, 4, len(b))
+
+	doProcess(t)
+
+	assert.Equal(t, 1, b[Binding{"host4.com", "name4", "127.0.0.1"}])
+	assert.Equal(t, 1, b[Binding{"host4.com", "name41", "127.0.0.1"}])
+	assert.Equal(t, 1, b[Binding{"host4.com", "name4", net.IPv6loopback.String()}])
+	assert.Equal(t, 1, b[Binding{"host4.com", "name41", net.IPv6loopback.String()}])
+	assert.Equal(t, 4, len(b))
+}
+
+func TestIpsetSubdomainOverride(t *testing.T) {
+	setup()
+
+	ctx.proxyCtx.Req = makeReqA("sub.host4.com.")
+	ctx.proxyCtx.Res = &dns.Msg{
+		Answer: []dns.RR{
+			makeA("sub.host4.com.", net.IPv4(127, 0, 0, 1)),
+		},
+	}
+
+	doProcess(t)
+
+	assert.Equal(t, 1, b[Binding{"sub.host4.com", "subhost4", "127.0.0.1"}])
+	assert.Equal(t, 1, len(b))
+}
+
+func TestIpsetSubdomainWildcard(t *testing.T) {
+	setup()
+
+	ctx.proxyCtx.Req = makeReqA("sub.host.com.")
+	ctx.proxyCtx.Res = &dns.Msg{
+		Answer: []dns.RR{
+			makeA("sub.host.com.", net.IPv4(127, 0, 0, 1)),
+		},
+	}
+
+	doProcess(t)
+
+	assert.Equal(t, 1, b[Binding{"sub.host.com", "name", "127.0.0.1"}])
+	assert.Equal(t, 1, len(b))
+}
+
+func TestIpsetCnameThirdParty(t *testing.T) {
+	setup()
+
+	ctx.proxyCtx.Req = makeReqA("host.com.")
+	ctx.proxyCtx.Res = &dns.Msg{
+		Answer: []dns.RR{
+			makeCNAME("host.com.", "foo.bar.baz.elb.amazonaws.com."),
+			makeA("foo.bar.baz.elb.amazonaws.com.", net.IPv4(8, 8, 8, 8)),
+		},
+	}
+
+	doProcess(t)
+
+	assert.Equal(t, 1, b[Binding{"host.com", "name", "8.8.8.8"}])
+	assert.Equal(t, 1, len(b))
 }