From f5741f61ac25188c01f1bd3398ce403f8e2b2e51 Mon Sep 17 00:00:00 2001
From: eddc005 <eddc005@proton.me>
Date: Sun, 2 Jun 2024 23:35:00 +0100
Subject: [PATCH] feat: add multiple addresses support for DNS modifier

---
 modifier/interface.go |  2 +-
 modifier/udp/dns.go   | 58 ++++++++++++++++++++++++++++++-------------
 ruleset/expr.go       |  4 +--
 3 files changed, 44 insertions(+), 20 deletions(-)

diff --git a/modifier/interface.go b/modifier/interface.go
index 0340a87..5a8a17c 100644
--- a/modifier/interface.go
+++ b/modifier/interface.go
@@ -4,7 +4,7 @@ type Modifier interface {
 	// Name returns the name of the modifier.
 	Name() string
 	// New returns a new modifier instance.
-	New(args map[string]interface{}) (Instance, error)
+	New(args map[string][]interface{}) (Instance, error)
 }
 
 type Instance interface{}
diff --git a/modifier/udp/dns.go b/modifier/udp/dns.go
index afab276..fe3f763 100644
--- a/modifier/udp/dns.go
+++ b/modifier/udp/dns.go
@@ -1,7 +1,10 @@
 package udp
 
 import (
+	"encoding/binary"
 	"errors"
+	"hash/fnv"
+	"math/rand/v2"
 	"net"
 
 	"github.com/apernet/OpenGFW/modifier"
@@ -24,23 +27,28 @@ func (m *DNSModifier) Name() string {
 	return "dns"
 }
 
-func (m *DNSModifier) New(args map[string]interface{}) (modifier.Instance, error) {
+func (m *DNSModifier) New(args map[string][]interface{}) (modifier.Instance, error) {
 	i := &dnsModifierInstance{}
-	aStr, ok := args["a"].(string)
-	if ok {
-		a := net.ParseIP(aStr).To4()
-		if a == nil {
-			return nil, &modifier.ErrInvalidArgs{Err: errInvalidIP}
+	i.seed = rand.Uint32()
+	for _, arg := range args["a"] {
+		aStr, ok := arg.(string)
+		if ok {
+			a := net.ParseIP(aStr).To4()
+			if a == nil {
+				return nil, &modifier.ErrInvalidArgs{Err: errInvalidIP}
+			}
+			i.A = append(i.A, a)
 		}
-		i.A = a
 	}
-	aaaaStr, ok := args["aaaa"].(string)
-	if ok {
-		aaaa := net.ParseIP(aaaaStr).To16()
-		if aaaa == nil {
-			return nil, &modifier.ErrInvalidArgs{Err: errInvalidIP}
+	for _, arg := range args["aaaa"] {
+		aaaaStr, ok := arg.(string)
+		if ok {
+			aaaa := net.ParseIP(aaaaStr).To16()
+			if aaaa == nil {
+				return nil, &modifier.ErrInvalidArgs{Err: errInvalidIP}
+			}
+			i.AAAA = append(i.AAAA, aaaa)
 		}
-		i.AAAA = aaaa
 	}
 	return i, nil
 }
@@ -48,8 +56,9 @@ func (m *DNSModifier) New(args map[string]interface{}) (modifier.Instance, error
 var _ modifier.UDPModifierInstance = (*dnsModifierInstance)(nil)
 
 type dnsModifierInstance struct {
-	A    net.IP
-	AAAA net.IP
+	A    []net.IP
+	AAAA []net.IP
+	seed uint32
 }
 
 func (i *dnsModifierInstance) Process(data []byte) ([]byte, error) {
@@ -64,26 +73,41 @@ func (i *dnsModifierInstance) Process(data []byte) ([]byte, error) {
 	if len(dns.Questions) == 0 {
 		return nil, &modifier.ErrInvalidPacket{Err: errEmptyDNSQuestion}
 	}
+
+	// Hash the query name so that DNS response is fixed for a given query.
+	// Use a random seed to avoid determinism.
+	hashStringToIndex := func(b []byte, sliceLength int, seed uint32) int {
+        h := fnv.New32a()
+		seedBytes := make([]byte, 4)
+        binary.LittleEndian.PutUint32(seedBytes, seed)
+		h.Write(seedBytes)
+        h.Write(b)
+        hashValue := h.Sum32()
+        return int(hashValue % uint32(sliceLength))
+    }
+
 	// In practice, most if not all DNS clients only send one question
 	// per packet, so we don't care about the rest for now.
 	q := dns.Questions[0]
 	switch q.Type {
 	case layers.DNSTypeA:
 		if i.A != nil {
+			idx := hashStringToIndex(q.Name, len(i.A), i.seed)
 			dns.Answers = []layers.DNSResourceRecord{{
 				Name:  q.Name,
 				Type:  layers.DNSTypeA,
 				Class: layers.DNSClassIN,
-				IP:    i.A,
+				IP:    i.A[idx],
 			}}
 		}
 	case layers.DNSTypeAAAA:
 		if i.AAAA != nil {
+			idx := hashStringToIndex(q.Name, len(i.AAAA), i.seed)
 			dns.Answers = []layers.DNSResourceRecord{{
 				Name:  q.Name,
 				Type:  layers.DNSTypeAAAA,
 				Class: layers.DNSClassIN,
-				IP:    i.AAAA,
+				IP:    i.AAAA[idx],
 			}}
 		}
 	}
diff --git a/ruleset/expr.go b/ruleset/expr.go
index 373d0d3..b1e4fe6 100644
--- a/ruleset/expr.go
+++ b/ruleset/expr.go
@@ -32,8 +32,8 @@ type ExprRule struct {
 }
 
 type ModifierEntry struct {
-	Name string                 `yaml:"name"`
-	Args map[string]interface{} `yaml:"args"`
+	Name string                   `yaml:"name"`
+	Args map[string][]interface{} `yaml:"args"`
 }
 
 func ExprRulesFromYAML(file string) ([]ExprRule, error) {