Refactor local DNS cache partitioning

dns/client.go

@@ -2,6 +2,7 @@ package dns
 
 import (
 	"context"
+	"encoding/binary"
 	"errors"
 	"hash/fnv"
 	"net"
@@ -18,6 +19,7 @@ import (
 	"github.com/sagernet/sing/common/task"
 	"github.com/sagernet/sing/contrab/freelru"
 	"github.com/sagernet/sing/contrab/maphash"
+	"github.com/sagernet/sing/service"
 
 	"github.com/miekg/dns"
 )
@@ -43,6 +45,7 @@ type Client struct {
 	initRDRCFunc      func() adapter.RDRCStore
 	dnsCache          adapter.DNSCacheStore
 	initDNSCacheFunc  func() adapter.DNSCacheStore
+	networkManager    adapter.NetworkManager
 	logger            logger.ContextLogger
 	cache             *freelru.Cache[dnsCacheKey, *dns.Msg]
 	cacheLock         compatible.Map[dnsCacheKey, chan struct{}]
@@ -104,53 +107,56 @@ func (k dnsCacheKey) persistentName() string {
 }
 
 func (c *Client) newCacheKey(transport adapter.DNSTransport, question dns.Question, message *dns.Msg, options adapter.DNSQueryOptions) dnsCacheKey {
-	cacheKey := dnsCacheKey{Question: question, transportTag: transport.Tag(), clientSubnet: c.effectiveClientSubnet(message, options)}
-	environmentTransport, withEnvironment := transport.(adapter.DNSTransportWithEnvironment)
-	if withEnvironment {
-		cacheKey.environment = environmentHash(environmentTransport.Environment())
+	clientSubnet := options.ClientSubnet
+	if !clientSubnet.IsValid() {
+		clientSubnet = c.clientSubnet
+	}
+	if !clientSubnet.IsValid() {
+		clientSubnet = clientSubnetFromMessage(message)
+	}
+	return dnsCacheKey{
+		Question:     question,
+		transportTag: transport.Tag(),
+		clientSubnet: clientSubnet,
+		environment:  c.environmentHash(transport),
 	}
-	return cacheKey
 }
 
 func (c *Client) finishCacheKey(transport adapter.DNSTransport, key dnsCacheKey) (dnsCacheKey, bool) {
-	environmentTransport, withEnvironment := transport.(adapter.DNSTransportWithEnvironment)
-	if !withEnvironment {
-		return key, true
-	}
-	environment := environmentHash(environmentTransport.Environment())
-	if environment == key.environment {
-		return key, true
-	}
-	if key.environment == 0 {
+	environment := c.environmentHash(transport)
+	if environment == key.environment || key.environment == 0 {
 		key.environment = environment
 		return key, true
 	}
 	return key, false
 }
 
-func environmentHash(environment []string) uint64 {
-	if len(environment) == 0 {
+func (c *Client) environmentHash(transport adapter.DNSTransport) uint64 {
+	environmentTransport, withEnvironment := transport.(adapter.DNSTransportWithEnvironment)
+	if !withEnvironment {
 		return 0
 	}
+	var networkEnvironment uint64
+	if c.networkManager != nil {
+		networkEnvironment = c.networkManager.NetworkEnvironment()
+	}
+	environment := environmentTransport.Environment()
+	if len(environment) == 0 {
+		return networkEnvironment
+	}
 	digest := fnv.New64a()
 	for _, entry := range environment {
 		digest.Write([]byte(entry))
 		digest.Write([]byte{0})
 	}
+	var hashBytes [8]byte
+	binary.BigEndian.PutUint64(hashBytes[:], networkEnvironment)
+	digest.Write(hashBytes[:])
 	return digest.Sum64()
 }
 
-func (c *Client) effectiveClientSubnet(message *dns.Msg, options adapter.DNSQueryOptions) netip.Prefix {
-	if options.ClientSubnet.IsValid() {
-		return options.ClientSubnet
-	}
-	if c.clientSubnet.IsValid() {
-		return c.clientSubnet
-	}
-	return clientSubnetFromMessage(message)
-}
-
 func (c *Client) Start() {
+	c.networkManager = service.FromContext[adapter.NetworkManager](c.ctx)
 	if c.initRDRCFunc != nil {
 		c.rdrc = c.initRDRCFunc()
 	}