add support for reading from file; re-enable args a and aaaa

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{}

modifier/udp/dns.go

@@ -1,11 +1,14 @@
 package udp
 
 import (
+	"bufio"
 	"encoding/binary"
 	"errors"
+	"fmt"
 	"hash/fnv"
 	"math/rand/v2"
 	"net"
+	"os"
 
 	"github.com/apernet/OpenGFW/modifier"
 
@@ -17,37 +20,102 @@ var _ modifier.Modifier = (*DNSModifier)(nil)
 
 var (
 	errInvalidIP           = errors.New("invalid ip")
+	errInvalidIPList       = errors.New("invalid ip list")
+	errInvalidIpListFile   = errors.New("unable to open or parse ip list file")
 	errNotValidDNSResponse = errors.New("not a valid dns response")
 	errEmptyDNSQuestion    = errors.New("empty dns question")
 )
 
+func fmtErrInvalidIP(ip string) error {
+	return fmt.Errorf("invalid ip: %s", ip)
+}
+
+func fmtErrInvalidIpListFile(filePath string) error {
+	return fmt.Errorf("unable to open or parse ip list file: %s", filePath)
+}
+
 type DNSModifier struct{}
 
 func (m *DNSModifier) Name() string {
 	return "dns"
 }
 
-func (m *DNSModifier) New(args map[string][]interface{}) (modifier.Instance, error) {
+func (m *DNSModifier) parseIpEntry(entry interface{}, i *dnsModifierInstance) error {
-	i := &dnsModifierInstance{}
+	entryStr, ok := entry.(string)
-	i.seed = rand.Uint32()
+	if !ok {
-	for _, arg := range args["a"] {
+		return &modifier.ErrInvalidArgs{Err: errInvalidIP}
-		aStr, ok := arg.(string)
+	}
-		if ok {
+
-			a := net.ParseIP(aStr).To4()
+	ip := net.ParseIP(entryStr)
-			if a == nil {
+	if ip == nil {
-				return nil, &modifier.ErrInvalidArgs{Err: errInvalidIP}
+		return &modifier.ErrInvalidArgs{Err: fmtErrInvalidIP(entryStr)}
-			}
+	}
-			i.A = append(i.A, a)
+	if ip4 := ip.To4(); ip4 != nil {
+		i.A = append(i.A, ip4)
+	} else {
+		i.AAAA = append(i.AAAA, ip)
+	}
+
+	return nil
+}
+
+func (m *DNSModifier) parseIpList(list []interface{}, i *dnsModifierInstance) error {
+	for _, entry := range list {
+		if err := m.parseIpEntry(entry, i); err != nil {
+			return err
 		}
 	}
-	for _, arg := range args["aaaa"] {
+	return nil
-		aaaaStr, ok := arg.(string)
+}
-		if ok {
+
-			aaaa := net.ParseIP(aaaaStr).To16()
+func (m *DNSModifier) parseIpListFile(filePath string, i *dnsModifierInstance) error {
-			if aaaa == nil {
+	file, err := os.Open(filePath)
-				return nil, &modifier.ErrInvalidArgs{Err: errInvalidIP}
+	if err != nil {
+		return &modifier.ErrInvalidArgs{Err: fmtErrInvalidIpListFile(filePath)}
+	}
+	defer file.Close()
+
+	scanner := bufio.NewScanner(file)
+	for scanner.Scan() {
+		line := scanner.Text()
+		if err := m.parseIpEntry(line, i); err != nil {
+			return err
+		}
+	}
+
+	if err := scanner.Err(); err != nil {
+		return &modifier.ErrInvalidArgs{Err: fmtErrInvalidIpListFile(filePath)}
+	}
+
+	return nil
+}
+
+func (m *DNSModifier) New(args map[string]interface{}) (modifier.Instance, error) {
+	i := &dnsModifierInstance{}
+	i.seed = rand.Uint32()
+
+	for key, value := range args {
+		switch key {
+		case "a", "aaaa":
+			if err := m.parseIpEntry(value, i); err != nil {
+				return nil, err
+			}
+		case "list":
+			if list, ok := value.([]interface{}); ok {
+				if err := m.parseIpList(list, i); err != nil {
+					return nil, err
+				}
+			} else {
+				return nil, &modifier.ErrInvalidArgs{Err: errInvalidIPList}
+			}
+		case "file":
+			if filePath, ok := value.(string); ok {
+				if err := m.parseIpListFile(filePath, i); err != nil {
+					return nil, err
+				}
+			} else {
+				return nil, &modifier.ErrInvalidArgs{Err: errInvalidIpListFile}
 			}
-			i.AAAA = append(i.AAAA, aaaa)
 		}
 	}
 	return i, nil
@@ -77,14 +145,14 @@ func (i *dnsModifierInstance) Process(data []byte) ([]byte, error) {
 	// 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()
+		h := fnv.New32a()
 		seedBytes := make([]byte, 4)
-        binary.LittleEndian.PutUint32(seedBytes, seed)
+		binary.LittleEndian.PutUint32(seedBytes, seed)
 		h.Write(seedBytes)
-        h.Write(b)
+		h.Write(b)
-        hashValue := h.Sum32()
+		hashValue := h.Sum32()
-        return int(hashValue % uint32(sliceLength))
+		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.

ruleset/expr.go

@@ -32,8 +32,8 @@ type ExprRule struct {
 }
 
 type ModifierEntry struct {
-	Name string                   `yaml:"name"`
+	Name string                 `yaml:"name"`
-	Args map[string][]interface{} `yaml:"args"`
+	Args map[string]interface{} `yaml:"args"`
 }
 
 func ExprRulesFromYAML(file string) ([]ExprRule, error) {