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

This commit is contained in:
Andrey Meshkov
2020-10-08 14:59:09 +03:00
2 changed files with 232 additions and 37 deletions

View File

@@ -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
func (c *ipsetCtx) process(ctx *dnsContext) int {
// Find the ipsets for a given host (accounting for subdomain wildcards)
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 {
ip := c.getIP(it)
if ip == nil {
continue
}
if ctx.proxyCtx.Res != nil {
for _, it := range ctx.proxyCtx.Res.Answer {
ip := c.getIP(it)
if ip == nil {
continue
}
ipStr := ip.String()
for _, name := range ipsetNames {
code, out, err := util.RunCommand("ipset", "add", name, ipStr)
if err != nil {
log.Info("IPSET: %s(%s) -> %s: %s", host, ipStr, name, err)
continue
ipStr := ip.String()
for _, name := range ipsetNames {
addMember(host, name, ipStr)
}
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)
}

View File

@@ -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) {
s := Server{}
s.conf.IPSETList = append(s.conf.IPSETList, "HOST.com/name")
s.conf.IPSETList = append(s.conf.IPSETList, "host2.com,host3.com/name23")
s.conf.IPSETList = append(s.conf.IPSETList, "host4.com/name4,name41")
c := ipsetCtx{}
var s Server
var c ipsetCtx
var ctx *dnsContext
type Binding struct {
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{
srv: &s,
}
ctx.proxyCtx = &proxy.DNSContext{}
ctx.proxyCtx.Req = &dns.Msg{
Question: []dns.Question{
{
Name: "host.com.",
Qtype: dns.TypeA,
},
func TestIpsetNoQuestion(t *testing.T) {
setup()
doProcess(t)
assert.Equal(t, 0, len(b))
}
func TestIpsetNoAnswer(t *testing.T) {
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))
}