Skip to content

Commit

Permalink
Pull request: proxy: fix edns cache
Browse files Browse the repository at this point in the history 
Merge in GO/dnsproxy from 329-edns-addr-cache to master

Squashed commit of the following:

commit 6b19f15
Author: Dimitry Kolyshev <dkolyshev@adguard.com>
Date:   Wed May 17 10:34:50 2023 +0300

    proxy: imp code

commit f3f8215
Author: Dimitry Kolyshev <dkolyshev@adguard.com>
Date:   Tue May 16 15:43:43 2023 +0300

    proxy: imp code

commit 9352a3d
Author: Dimitry Kolyshev <dkolyshev@adguard.com>
Date:   Tue May 16 15:29:11 2023 +0300

    proxy: imp code

commit ae8b196
Author: Dimitry Kolyshev <dkolyshev@adguard.com>
Date:   Tue May 16 13:45:42 2023 +0300

    proxy: imp code

commit 3a6eea1
Author: Dimitry Kolyshev <dkolyshev@adguard.com>
Date:   Tue May 16 11:01:14 2023 +0300

    proxy: imp code

commit 281db6d
Author: Dimitry Kolyshev <dkolyshev@adguard.com>
Date:   Mon May 15 15:11:36 2023 +0300

    proxy: edns cache bits

commit b5b1ab3
Author: Dimitry Kolyshev <dkolyshev@adguard.com>
Date:   Mon May 15 12:12:55 2023 +0300

    Revert "proxy: edns cache bits"

    This reverts commit 33e1218.

commit 33e1218
Author: Dimitry Kolyshev <dkolyshev@adguard.com>
Date:   Mon May 15 11:51:39 2023 +0300

    proxy: edns cache bits

commit e8b4b2c
Author: Dimitry Kolyshev <dkolyshev@adguard.com>
Date:   Fri May 12 12:41:09 2023 +0300

    proxy: imp tests

commit 0e15a81
Author: Dimitry Kolyshev <dkolyshev@adguard.com>
Date:   Fri May 12 12:37:36 2023 +0300

    proxy: imp tests

commit 9f34110
Author: Dimitry Kolyshev <dkolyshev@adguard.com>
Date:   Fri May 12 12:24:02 2023 +0300

    proxy: imp code

commit aa08e6d
Author: Dimitry Kolyshev <dkolyshev@adguard.com>
Date:   Fri May 12 11:56:47 2023 +0300

    proxy: imp code

commit cf32f74
Author: Dimitry Kolyshev <dkolyshev@adguard.com>
Date:   Fri May 12 11:15:09 2023 +0300

    proxy: imp tests

commit d73664e
Author: Dimitry Kolyshev <dkolyshev@adguard.com>
Date:   Thu May 11 16:43:43 2023 +0300

    proxy: fix edns cache
  • Loading branch information
@Mizzick
Mizzick committed May 17, 2023
1 parent 4f1c527 commit 936bd45
Show file tree
Hide file tree
Showing 3 changed files with 138 additions and 88 deletions.
80 changes: 57 additions & 23 deletions proxy/cache.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,9 @@ import (
"github.com/AdguardTeam/dnsproxy/upstream"
glcache "github.com/AdguardTeam/golibs/cache"
"github.com/AdguardTeam/golibs/log"
"github.com/AdguardTeam/golibs/mathutil"
"github.com/miekg/dns"
"golang.org/x/exp/slices"
)

// defaultCacheSize is the size of cache in bytes by default.
Expand Down Expand Up @@ -210,8 +212,8 @@ func (c *cache) get(req *dns.Msg) (ci *cacheItem, expired bool, key []byte) {
return ci, expired, key
}

// getWithSubnet returns cached item for the req if it's found by n. expired is
// true if the item's TTL is expired. k is the resulting key for req. It's
// getWithSubnet returns cached item for the req if it's found by n. expired
// is true if the item's TTL is expired. k is the resulting key for req. It's
// returned to avoid recalculating it afterwards.
//
// Note that a slow longest-prefix-match algorithm is used, so cache searches
Expand All @@ -224,9 +226,38 @@ func (c *cache) getWithSubnet(req *dns.Msg, n *net.IPNet) (ci *cacheItem, expire
return nil, false, nil
}

var data []byte
for mask, _ := n.Mask.Size(); mask >= 0 && data == nil; mask-- {
k = msgToKeyWithSubnet(req, n.IP, mask)
ecsIP := n.IP.Mask(n.Mask)
ipLen := len(ecsIP)
m, _ := n.Mask.Size()

k = msgToKeyWithSubnet(req, ecsIP, m)
data := c.itemsWithSubnet.Get(k)

// In order to reduce allocations we apply mask on bits level. As the key
// k has ecsIP in bytes slice representation, each iteration we can just
// clear one bit in the end of it by applying the bitmask.
for bitmask := ^byte(0); m >= 0 && data == nil; m-- {
// Set mask identification byte in the key.
k[keyMaskIndex] = byte(m)

// In case mask is zero, the key doesn't have IP in it.
if m == 0 {
k = slices.Delete(k, keyIPIndex, keyIPIndex+ipLen)
data = c.itemsWithSubnet.Get(k)

continue
}

// Shift or renew bitmask.
if m%8 == 0 {
bitmask = ^byte(0)
} else {
bitmask <<= 1
}

// Clear the last non-zero bit in the byte of the IP address.
k[keyIPIndex+m/8] &= bitmask

data = c.itemsWithSubnet.Get(k)
}

Expand Down Expand Up @@ -286,7 +317,7 @@ func (c *cache) setWithSubnet(m *dns.Msg, u upstream.Upstream, subnet *net.IPNet
}

pref, _ := subnet.Mask.Size()
key := msgToKeyWithSubnet(m, subnet.IP, pref)
key := msgToKeyWithSubnet(m, subnet.IP.Mask(subnet.Mask), pref)
packed := item.pack()

c.itemsWithSubnetLock.Lock()
Expand Down Expand Up @@ -480,44 +511,47 @@ func msgToKey(m *dns.Msg) (b []byte) {
return b
}

const (
// keyMaskIndex is the index of the byte with mask ones value.
keyMaskIndex = 1 + 2*packedMsgLenSz

// keyIPIndex is the start index of the IP address in the key.
keyIPIndex = keyMaskIndex + 1
)

// msgToKeyWithSubnet constructs the cache key from DO bit, type, class, subnet
// mask, client's IP address and question's name of m. ecsIP is expected to be
// masked already.
func msgToKeyWithSubnet(m *dns.Msg, ecsIP net.IP, mask int) (key []byte) {
q := m.Question[0]
cap := 1 + 2*packedMsgLenSz + 1 + len(q.Name)
ipLen := len(ecsIP)
keyLen := keyIPIndex + len(q.Name)
masked := mask != 0
if masked {
cap += ipLen
keyLen += len(ecsIP)
}

// Initialize the slice.
key = make([]byte, cap)
k := 0
key = make([]byte, keyLen)

// Put DO.
if opt := m.IsEdns0(); opt != nil && opt.Do() {
key[k] = 1
} else {
key[k] = 0
}
k++
opt := m.IsEdns0()
key[0] = mathutil.BoolToNumber[byte](opt != nil && opt.Do())

// Put Qtype.
//
// TODO(d.kolyshev): We should put Qtype in key[1:].
binary.BigEndian.PutUint16(key[:], q.Qtype)
k += packedMsgLenSz

// Put Qclass.
binary.BigEndian.PutUint16(key[k:], q.Qclass)
k += packedMsgLenSz
binary.BigEndian.PutUint16(key[1+packedMsgLenSz:], q.Qclass)

// Add mask.
key[k] = uint8(mask)
k++
key[keyMaskIndex] = uint8(mask)
k := keyIPIndex
if masked {
k += copy(key[k:], ecsIP)
k += copy(key[keyIPIndex:], ecsIP)
}

copy(key[k:], strings.ToLower(q.Name))

return key
Expand Down
144 changes: 80 additions & 64 deletions proxy/cache_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -583,122 +583,138 @@ func setAndGetCache(t *testing.T, c *cache, g *sync.WaitGroup, host, ip string)
}, 1100*time.Millisecond, 100*time.Millisecond, "cache for %s should already be removed", host)
}

func TestSubnet(t *testing.T) {
c := newCache(testCacheSize, true, false)
func TestCache_getWithSubnet(t *testing.T) {
const testFQDN = "example.com."

ip1234, ip2234, ip3234 := net.IP{1, 2, 3, 4}, net.IP{2, 2, 3, 4}, net.IP{3, 2, 3, 4}
req := (&dns.Msg{}).SetQuestion("example.com.", dns.TypeA)
req := (&dns.Msg{}).SetQuestion(testFQDN, dns.TypeA)
mask16 := net.CIDRMask(16, netutil.IPv4BitLen)
mask24 := net.CIDRMask(24, netutil.IPv4BitLen)

c := newCache(testCacheSize, true, false)

t.Run("empty", func(t *testing.T) {
ci, expired, _ := c.getWithSubnet(req, &net.IPNet{
IP: ip1234,
Mask: net.CIDRMask(24, netutil.IPv4BitLen),
})
ci, expired, _ := c.getWithSubnet(req, &net.IPNet{IP: ip1234, Mask: mask24})
assert.Nil(t, ci)
assert.False(t, expired)
})

// Add a response with subnet.
resp := (&dns.Msg{
MsgHdr: dns.MsgHdr{
Response: true,
},
Answer: []dns.RR{newRR(t, "example.com.", dns.TypeA, 1, net.IP{1, 1, 1, 1})},
}).SetQuestion("example.com.", dns.TypeA)
c.setWithSubnet(
resp,
upstreamWithAddr,
&net.IPNet{IP: ip1234, Mask: net.CIDRMask(16, netutil.IPv4BitLen)},
)
Answer: []dns.RR{newRR(t, testFQDN, dns.TypeA, 1, net.IP{1, 1, 1, 1})},
}).SetReply(req)
c.setWithSubnet(resp, upstreamWithAddr, &net.IPNet{IP: ip1234, Mask: mask16})

t.Run("different_ip", func(t *testing.T) {
ci, expired, key := c.getWithSubnet(req, &net.IPNet{
IP: ip2234,
Mask: net.CIDRMask(24, netutil.IPv4BitLen),
})
ci, expired, key := c.getWithSubnet(req, &net.IPNet{IP: ip2234, Mask: mask24})
assert.False(t, expired)
assert.Equal(t, msgToKeyWithSubnet(req, ip2234, 0), key)

require.Nil(t, ci)
assert.Nil(t, ci)
})

// Add a response entry with subnet #2.
resp = (&dns.Msg{
MsgHdr: dns.MsgHdr{
Response: true,
},
Answer: []dns.RR{newRR(t, "example.com.", dns.TypeA, 1, net.IP{2, 2, 2, 2})},
}).SetQuestion("example.com.", dns.TypeA)
c.setWithSubnet(
resp,
upstreamWithAddr,
&net.IPNet{IP: ip2234, Mask: net.CIDRMask(16, netutil.IPv4BitLen)},
)
Answer: []dns.RR{newRR(t, testFQDN, dns.TypeA, 1, net.IP{2, 2, 2, 2})},
}).SetReply(req)
c.setWithSubnet(resp, upstreamWithAddr, &net.IPNet{IP: ip2234, Mask: mask16})

// Add a response entry without subnet.
resp = (&dns.Msg{
MsgHdr: dns.MsgHdr{
Response: true,
},
Answer: []dns.RR{newRR(t, "example.com.", dns.TypeA, 1, net.IP{3, 3, 3, 3})},
}).SetQuestion("example.com.", dns.TypeA)
c.setWithSubnet(
resp,
upstreamWithAddr,
&net.IPNet{IP: nil, Mask: nil},
)
Answer: []dns.RR{newRR(t, testFQDN, dns.TypeA, 1, net.IP{3, 3, 3, 3})},
}).SetReply(req)
c.setWithSubnet(resp, upstreamWithAddr, &net.IPNet{IP: nil, Mask: nil})

t.Run("with_subnet_1", func(t *testing.T) {
ci, expired, key := c.getWithSubnet(req, &net.IPNet{
IP: ip1234,
Mask: net.CIDRMask(24, netutil.IPv4BitLen),
})
ci, expired, key := c.getWithSubnet(req, &net.IPNet{IP: ip1234, Mask: mask24})
assert.False(t, expired)
assert.Equal(t, msgToKeyWithSubnet(req, ip1234, 16), key)
assert.Equal(t, msgToKeyWithSubnet(req, ip1234.Mask(mask16), 16), key)

require.NotNil(t, ci)
require.NotNil(t, ci.m)
require.NotEmpty(t, ci.m.Answer)

a, ok := ci.m.Answer[0].(*dns.A)
require.True(t, ok)

a := testutil.RequireTypeAssert[*dns.A](t, ci.m.Answer[0])
assert.True(t, a.A.Equal(net.IP{1, 1, 1, 1}))
})

t.Run("with_subnet_2", func(t *testing.T) {
ci, expired, key := c.getWithSubnet(req, &net.IPNet{
IP: ip2234,
Mask: net.CIDRMask(24, netutil.IPv4BitLen),
})
ci, expired, key := c.getWithSubnet(req, &net.IPNet{IP: ip2234, Mask: mask24})
assert.False(t, expired)
assert.Equal(t, msgToKeyWithSubnet(req, ip2234, 16), key)
assert.Equal(t, msgToKeyWithSubnet(req, ip2234.Mask(mask16), 16), key)

require.NotNil(t, ci)
require.NotNil(t, ci.m)
require.NotEmpty(t, ci.m.Answer)

a, ok := ci.m.Answer[0].(*dns.A)
require.True(t, ok)

a := testutil.RequireTypeAssert[*dns.A](t, ci.m.Answer[0])
assert.True(t, a.A.Equal(net.IP{2, 2, 2, 2}))
})

t.Run("with_subnet_3", func(t *testing.T) {
ci, expired, key := c.getWithSubnet(req, &net.IPNet{IP: ip3234, Mask: mask24})
assert.False(t, expired)
assert.Equal(t, msgToKeyWithSubnet(req, ip1234, 0), key)

require.NotNil(t, ci)
require.NotNil(t, ci.m)
require.NotEmpty(t, ci.m.Answer)

a := testutil.RequireTypeAssert[*dns.A](t, ci.m.Answer[0])
assert.True(t, a.A.Equal(net.IP{3, 3, 3, 3}))
})
}

func TestCache_getWithSubnet_mask(t *testing.T) {
const testFQDN = "example.com."

testIP := net.IP{176, 112, 191, 0}
noMatchIP := net.IP{177, 112, 191, 0}

// cachedIP/cidrMask network contains the testIP.
const cidrMaskOnes = 20
cidrMask := net.CIDRMask(cidrMaskOnes, netutil.IPv4BitLen)
cachedIP := net.IP{176, 112, 176, 0}

ansIP := net.IP{4, 4, 4, 4}

c := newCache(testCacheSize, true, true)

req := (&dns.Msg{}).SetQuestion(testFQDN, dns.TypeA)
resp := (&dns.Msg{
Answer: []dns.RR{newRR(t, testFQDN, dns.TypeA, 300, ansIP)},
}).SetReply(req)

// Cache IP network that contains the testIP.
c.setWithSubnet(
resp,
upstreamWithAddr,
&net.IPNet{IP: cachedIP, Mask: cidrMask},
)

t.Run("mask_matched", func(t *testing.T) {
ci, expired, key := c.getWithSubnet(req, &net.IPNet{
IP: ip3234,
IP: testIP,
Mask: net.CIDRMask(24, netutil.IPv4BitLen),
})
assert.False(t, expired)
assert.Equal(t, msgToKeyWithSubnet(req, ip1234, 0), key)
assert.Equal(t, msgToKeyWithSubnet(req, testIP.Mask(cidrMask), cidrMaskOnes), key)

require.NotNil(t, ci)
require.NotNil(t, ci.m)
require.NotEmpty(t, ci.m.Answer)

a, ok := ci.m.Answer[0].(*dns.A)
require.True(t, ok)
a := testutil.RequireTypeAssert[*dns.A](t, ci.m.Answer[0])
assert.True(t, a.A.Equal(ansIP))
})

assert.True(t, a.A.Equal(net.IP{3, 3, 3, 3}))
t.Run("no_mask_matched", func(t *testing.T) {
ci, expired, key := c.getWithSubnet(req, &net.IPNet{
IP: noMatchIP,
Mask: net.CIDRMask(24, netutil.IPv4BitLen),
})
assert.False(t, expired)
assert.Equal(t, msgToKeyWithSubnet(req, noMatchIP, 0), key)
assert.Nil(t, ci)
})
}

Expand Down
2 changes: 1 addition & 1 deletion proxy/proxy_cache.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -85,7 +85,7 @@ func (p *Proxy) cacheResp(d *DNSContext) {
// valid for all addresses that fall within that range.
//
// See RFC 7871 Section 7.3.1.
if scope < ones {
if scope < reqOnes {
ecs.Mask = net.CIDRMask(scope, bits)
ecs.IP = ecs.IP.Mask(ecs.Mask)
}
Expand Down

0 comments on commit 936bd45

Please sign in to comment.