proxy: imp code

AdguardTeam · May 12, 2023 · aa08e6d · aa08e6d
1 parent cf32f74
commit aa08e6d
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 15 deletions.
diff --git a/proxy/cache.go b/proxy/cache.go
@@ -225,8 +225,10 @@ func (c *cache) getWithSubnet(req *dns.Msg, n *net.IPNet) (ci *cacheItem, expire
 	}
 
 	var data []byte
-	for mask, _ := n.Mask.Size(); mask >= 0 && data == nil; mask-- {
-		k = msgToKeyWithSubnet(req, n.IP, mask)
+	for mask, bits := n.Mask.Size(); mask >= 0 && data == nil; mask-- {
+		ecsIP := n.IP.Mask(net.CIDRMask(mask, bits))
+		k = msgToKeyWithSubnet(req, ecsIP, mask)
+
 		data = c.itemsWithSubnet.Get(k)
 	}
 
@@ -485,15 +487,15 @@ func msgToKey(m *dns.Msg) (b []byte) {
 // 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)
+	keyLen := 1 + 2*packedMsgLenSz + 1 + len(q.Name)
 	ipLen := len(ecsIP)
 	masked := mask != 0
 	if masked {
-		cap += ipLen
+		keyLen += ipLen
 	}
 
 	// Initialize the slice.
-	key = make([]byte, cap)
+	key = make([]byte, keyLen)
 	k := 0
 
 	// Put DO.
@@ -516,8 +518,7 @@ func msgToKeyWithSubnet(m *dns.Msg, ecsIP net.IP, mask int) (key []byte) {
 	key[k] = uint8(mask)
 	k++
 	if masked {
-		ip := ecsIP.Mask(net.CIDRMask(mask, ipLen*8))
-		k += copy(key[k:], ip)
+		k += copy(key[k:], ecsIP)
 	}
 	copy(key[k:], strings.ToLower(q.Name))
 

diff --git a/proxy/cache_test.go b/proxy/cache_test.go
@@ -705,20 +705,28 @@ func TestSubnet(t *testing.T) {
 func TestCache_getWithSubnet_mask(t *testing.T) {
 	const testFQDN = "example.com."
 
-	c := newCache(testCacheSize, true, true)
 	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, net.IP{4, 4, 4, 4})},
+		Answer: []dns.RR{newRR(t, testFQDN, dns.TypeA, 300, ansIP)},
 	}).SetReply(req)
 
-	// Store mask that contains testIP.
+	// Cache IP network that contains the testIP.
 	c.setWithSubnet(
 		resp,
 		upstreamWithAddr,
-		&net.IPNet{IP: net.IP{176, 112, 176, 0}, Mask: net.CIDRMask(20, netutil.IPv4BitLen)},
+		&net.IPNet{IP: cachedIP, Mask: cidrMask},
 	)
 
 	t.Run("mask_matched", func(t *testing.T) {
@@ -727,14 +735,14 @@ func TestCache_getWithSubnet_mask(t *testing.T) {
 			Mask: net.CIDRMask(24, netutil.IPv4BitLen),
 		})
 		assert.False(t, expired)
-		assert.Equal(t, msgToKeyWithSubnet(req, testIP, 20), 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 := testutil.RequireTypeAssert[*dns.A](t, ci.m.Answer[0])
-		assert.True(t, a.A.Equal(net.IP{4, 4, 4, 4}))
+		assert.True(t, a.A.Equal(ansIP))
 	})
 
 	t.Run("no_mask_matched", func(t *testing.T) {

diff --git a/proxy/proxy_cache.go b/proxy/proxy_cache.go
@@ -85,8 +85,7 @@ func (p *Proxy) cacheResp(d *DNSContext) {
 		// valid for all addresses that fall within that range.
 		//
 		// See RFC 7871 Section 7.3.1.
-		// TODO(e.burkov): Should be scope < reqOnes?
-		if scope < ones {
+		if scope < reqOnes {
 			ecs.Mask = net.CIDRMask(scope, bits)
 			ecs.IP = ecs.IP.Mask(ecs.Mask)
 		}