From a21e4a7deb2d60b40ea4ebcedf6e22a56dcb4d04 Mon Sep 17 00:00:00 2001
From: Loyalsoldier <10487845+Loyalsoldier@users.noreply.github.com>
Date: Thu, 22 Apr 2021 08:12:52 +0800
Subject: [PATCH] Fix: FakeDNS IPv6 & refine JSON unmarshal (#925)

---
 infra/conf/fakedns.go | 115 ++++++++++++++++++++++++++++++------------
 1 file changed, 82 insertions(+), 33 deletions(-)

diff --git a/infra/conf/fakedns.go b/infra/conf/fakedns.go
index ab6aaa0be..91698c2cf 100644
--- a/infra/conf/fakedns.go
+++ b/infra/conf/fakedns.go
@@ -1,78 +1,127 @@
 package conf
 
 import (
-	"github.com/golang/protobuf/proto"
+	"encoding/json"
+	"strings"
 
 	"github.com/v2fly/v2ray-core/v4/app/dns/fakedns"
 )
 
 type FakeDNSPoolElementConfig struct {
 	IPPool  string `json:"ipPool"`
-	LruSize int64  `json:"poolSize"`
+	LRUSize int64  `json:"poolSize"`
 }
 
 type FakeDNSConfig struct {
-	IPPool  string                      `json:"ipPool"`
-	LruSize int64                       `json:"poolSize"`
-	Pools   *[]FakeDNSPoolElementConfig `json:"pools,omitempty"`
+	pool  *FakeDNSPoolElementConfig
+	pools []*FakeDNSPoolElementConfig
 }
 
-func (f FakeDNSConfig) Build() (proto.Message, error) {
-	if f.Pools != nil {
-		fakeDNSPool := &fakedns.FakeDnsPoolMulti{}
-		for _, v := range *f.Pools {
-			fakeDNSPool.Pools = append(fakeDNSPool.Pools, &fakedns.FakeDnsPool{IpPool: v.IPPool, LruSize: v.LruSize})
-		}
-		return fakeDNSPool, nil
+// UnmarshalJSON implements encoding/json.Unmarshaler.UnmarshalJSON
+func (f *FakeDNSConfig) UnmarshalJSON(data []byte) error {
+	var pool FakeDNSPoolElementConfig
+	var pools []*FakeDNSPoolElementConfig
+	switch {
+	case json.Unmarshal(data, &pool) == nil:
+		f.pool = &pool
+	case json.Unmarshal(data, &pools) == nil:
+		f.pools = pools
+	default:
+		return newError("invalid fakedns config")
+	}
+	return nil
+}
+
+func (f *FakeDNSConfig) Build() (*fakedns.FakeDnsPoolMulti, error) {
+	fakeDNSPool := fakedns.FakeDnsPoolMulti{}
+
+	if f.pool != nil {
+		fakeDNSPool.Pools = append(fakeDNSPool.Pools, &fakedns.FakeDnsPool{
+			IpPool:  f.pool.IPPool,
+			LruSize: f.pool.LRUSize,
+		})
+		return &fakeDNSPool, nil
 	}
 
-	return &fakedns.FakeDnsPool{
-		IpPool:  f.IPPool,
-		LruSize: f.LruSize,
-	}, nil
+	if f.pools != nil {
+		for _, v := range f.pools {
+			fakeDNSPool.Pools = append(fakeDNSPool.Pools, &fakedns.FakeDnsPool{IpPool: v.IPPool, LruSize: v.LRUSize})
+		}
+		return &fakeDNSPool, nil
+	}
+
+	return nil, newError("no valid FakeDNS config")
 }
 
 type FakeDNSPostProcessingStage struct{}
 
-func (FakeDNSPostProcessingStage) Process(conf *Config) error {
-	var fakeDNSInUse bool
+func (FakeDNSPostProcessingStage) Process(config *Config) error {
+	fakeDNSInUse := false
+	isIPv4Enable, isIPv6Enable := true, true
 
-	if conf.DNSConfig != nil {
-		for _, v := range conf.DNSConfig.Servers {
-			if v.Address.Family().IsDomain() {
-				if v.Address.Domain() == "fakedns" {
-					fakeDNSInUse = true
-				}
+	if config.DNSConfig != nil {
+		for _, v := range config.DNSConfig.Servers {
+			if v.Address.Family().IsDomain() && strings.EqualFold(v.Address.Domain(), "fakedns") {
+				fakeDNSInUse = true
 			}
 		}
+
+		switch strings.ToLower(config.DNSConfig.QueryStrategy) {
+		case "useip4", "useipv4", "use_ip4", "use_ipv4", "use_ip_v4", "use-ip4", "use-ipv4", "use-ip-v4":
+			isIPv4Enable, isIPv6Enable = true, false
+		case "useip6", "useipv6", "use_ip6", "use_ipv6", "use_ip_v6", "use-ip6", "use-ipv6", "use-ip-v6":
+			isIPv4Enable, isIPv6Enable = false, true
+		}
 	}
 
 	if fakeDNSInUse {
-		if conf.FakeDNS == nil {
-			// Add a Fake DNS Config if there is none
-			conf.FakeDNS = &FakeDNSConfig{
-				IPPool:  "198.18.0.0/15",
-				LruSize: 65535,
+		// Add a Fake DNS Config if there is none
+		if config.FakeDNS == nil {
+			config.FakeDNS = &FakeDNSConfig{}
+			switch {
+			case isIPv4Enable && isIPv6Enable:
+				config.FakeDNS.pools = []*FakeDNSPoolElementConfig{
+					{
+						IPPool:  "198.18.0.0/15",
+						LRUSize: 32768,
+					},
+					{
+						IPPool:  "fc00::/18",
+						LRUSize: 32768,
+					},
+				}
+			case !isIPv4Enable && isIPv6Enable:
+				config.FakeDNS.pool = &FakeDNSPoolElementConfig{
+					IPPool:  "fc00::/18",
+					LRUSize: 65535,
+				}
+			case isIPv4Enable && !isIPv6Enable:
+				config.FakeDNS.pool = &FakeDNSPoolElementConfig{
+					IPPool:  "198.18.0.0/15",
+					LRUSize: 65535,
+				}
 			}
 		}
+
 		found := false
 		// Check if there is a Outbound with necessary sniffer on
 		var inbounds []InboundDetourConfig
 
-		if len(conf.InboundConfigs) > 0 {
-			inbounds = append(inbounds, conf.InboundConfigs...)
+		if len(config.InboundConfigs) > 0 {
+			inbounds = append(inbounds, config.InboundConfigs...)
 		}
 		for _, v := range inbounds {
 			if v.SniffingConfig != nil && v.SniffingConfig.Enabled && v.SniffingConfig.DestOverride != nil {
 				for _, dov := range *v.SniffingConfig.DestOverride {
-					if dov == "fakedns" {
+					if strings.EqualFold(dov, "fakedns") || strings.EqualFold(dov, "fakedns+others") {
 						found = true
+						break
 					}
 				}
 			}
 		}
 		if !found {
-			newError("Defined Fake DNS but haven't enabled fake dns sniffing at any inbound.").AtWarning().WriteToLog()
+			newError("Defined FakeDNS but haven't enabled FakeDNS destOverride at any inbound.").AtWarning().WriteToLog()
 		}
 	}