proxy: edns cache bits

proxy/cache.go

@@ -226,9 +226,7 @@ func (c *cache) getWithSubnet(req *dns.Msg, n *net.IPNet) (ci *cacheItem, expire
 
 	var data []byte
 	for mask, bits := n.Mask.Size(); mask >= 0 && data == nil; mask-- {
-		ecsIP := n.IP.Mask(net.CIDRMask(mask, bits))
-		k = msgToKeyWithSubnet(req, ecsIP, mask)
-
+		k = msgToKeyWithSubnet(req, n.IP, mask, bits)
 		data = c.itemsWithSubnet.Get(k)
 	}
 
@@ -287,8 +285,8 @@ func (c *cache) setWithSubnet(m *dns.Msg, u upstream.Upstream, subnet *net.IPNet
 		return
 	}
 
-	pref, _ := subnet.Mask.Size()
-	key := msgToKeyWithSubnet(m, subnet.IP.Mask(subnet.Mask), pref)
+	pref, bits := subnet.Mask.Size()
+	key := msgToKeyWithSubnet(m, subnet.IP, pref, bits)
 	packed := item.pack()
 
 	c.itemsWithSubnetLock.Lock()
@@ -483,9 +481,8 @@ func msgToKey(m *dns.Msg) (b []byte) {
 }
 
 // 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) {
+// mask, client's IP address and question's name of m.
+func msgToKeyWithSubnet(m *dns.Msg, ecsIP net.IP, mask, bits int) (key []byte) {
 	q := m.Question[0]
 	keyLen := 1 + 2*packedMsgLenSz + 1 + len(q.Name)
 	ipLen := len(ecsIP)
@@ -518,7 +515,13 @@ func msgToKeyWithSubnet(m *dns.Msg, ecsIP net.IP, mask int) (key []byte) {
 	key[k] = uint8(mask)
 	k++
 	if masked {
-		k += copy(key[k:], ecsIP)
+		// Apply mask on the bits level to reduce allocations.
+		am := net.CIDRMask(mask, bits)
+		n := len(ecsIP)
+		for i := 0; i < n; i++ {
+			key[k] = ecsIP[i] & am[i]
+			k++
+		}
 	}
 	copy(key[k:], strings.ToLower(q.Name))
 

proxy/cache_test.go

@@ -608,7 +608,7 @@ func TestCache_getWithSubnet(t *testing.T) {
 	t.Run("different_ip", func(t *testing.T) {
 		ci, expired, key := c.getWithSubnet(req, &net.IPNet{IP: ip2234, Mask: mask24})
 		assert.False(t, expired)
-		assert.Equal(t, msgToKeyWithSubnet(req, ip2234, 0), key)
+		assert.Equal(t, msgToKeyWithSubnet(req, ip2234, 0, netutil.IPv4BitLen), key)
 		assert.Nil(t, ci)
 	})
 
@@ -627,7 +627,7 @@ func TestCache_getWithSubnet(t *testing.T) {
 	t.Run("with_subnet_1", func(t *testing.T) {
 		ci, expired, key := c.getWithSubnet(req, &net.IPNet{IP: ip1234, Mask: mask24})
 		assert.False(t, expired)
-		assert.Equal(t, msgToKeyWithSubnet(req, ip1234.Mask(mask16), 16), key)
+		assert.Equal(t, msgToKeyWithSubnet(req, ip1234, 16, netutil.IPv4BitLen), key)
 
 		require.NotNil(t, ci)
 		require.NotNil(t, ci.m)
@@ -640,7 +640,7 @@ func TestCache_getWithSubnet(t *testing.T) {
 	t.Run("with_subnet_2", func(t *testing.T) {
 		ci, expired, key := c.getWithSubnet(req, &net.IPNet{IP: ip2234, Mask: mask24})
 		assert.False(t, expired)
-		assert.Equal(t, msgToKeyWithSubnet(req, ip2234.Mask(mask16), 16), key)
+		assert.Equal(t, msgToKeyWithSubnet(req, ip2234, 16, netutil.IPv4BitLen), key)
 
 		require.NotNil(t, ci)
 		require.NotNil(t, ci.m)
@@ -653,7 +653,7 @@ func TestCache_getWithSubnet(t *testing.T) {
 	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)
+		assert.Equal(t, msgToKeyWithSubnet(req, ip1234, 0, netutil.IPv4BitLen), key)
 
 		require.NotNil(t, ci)
 		require.NotNil(t, ci.m)

proxy/proxy_test.go

@@ -1028,7 +1028,7 @@ func TestECSProxyCacheMinMaxTTL(t *testing.T) {
 	})
 	assert.False(t, expired)
 
-	assert.Equal(t, key, msgToKeyWithSubnet(d.Req, clientIP, 24))
+	assert.Equal(t, key, msgToKeyWithSubnet(d.Req, clientIP, 24, netutil.IPv4BitLen))
 	assert.True(t, ci.m.Answer[0].Header().Ttl == prx.CacheMinTTL)
 
 	// 2nd request
@@ -1055,7 +1055,7 @@ func TestECSProxyCacheMinMaxTTL(t *testing.T) {
 		Mask: net.CIDRMask(24, netutil.IPv4BitLen),
 	})
 	assert.False(t, expired)
-	assert.Equal(t, key, msgToKeyWithSubnet(d.Req, clientIP, 24))
+	assert.Equal(t, key, msgToKeyWithSubnet(d.Req, clientIP, 24, netutil.IPv4BitLen))
 	assert.True(t, ci.m.Answer[0].Header().Ttl == prx.CacheMaxTTL)
 }