Fakedns improvements (#731)

Co-authored-by: Shelikhoo <xiaokangwang@outlook.com>
Co-authored-by: sixg0000d <sixg0000d@gmail.com>
Co-authored-by: Loyalsoldier <10487845+Loyalsoldier@users.noreply.github.com>

app/dns/fakedns/fake.go

@@ -20,12 +20,30 @@ type Holder struct {
 	config *FakeDnsPool
 }
 
+func (fkdns *Holder) IsIPInIPPool(ip net.Address) bool {
+	if ip.Family().IsDomain() {
+		return false
+	}
+	return fkdns.ipRange.Contains(ip.IP())
+}
+
+func (fkdns *Holder) GetFakeIPForDomain3(domain string, ipv4, ipv6 bool) []net.Address {
+	isIPv6 := fkdns.ipRange.IP.To4() == nil
+	if (isIPv6 && ipv6) || (!isIPv6 && ipv4) {
+		return fkdns.GetFakeIPForDomain(domain)
+	}
+	return []net.Address{}
+}
+
 func (*Holder) Type() interface{} {
 	return (*dns.FakeDNSEngine)(nil)
 }
 
 func (fkdns *Holder) Start() error {
-	return fkdns.initializeFromConfig()
+	if fkdns.config != nil && fkdns.config.IpPool != "" && fkdns.config.LruSize != 0 {
+		return fkdns.initializeFromConfig()
+	}
+	return newError("invalid fakeDNS setting")
 }
 
 func (fkdns *Holder) Close() error {
@@ -41,7 +59,7 @@ func NewFakeDNSHolder() (*Holder, error) {
 	if fkdns, err = NewFakeDNSHolderConfigOnly(nil); err != nil {
 		return nil, newError("Unable to create Fake Dns Engine").Base(err).AtError()
 	}
-	err = fkdns.initialize(dns.FakeIPPool, 65535)
+	err = fkdns.initialize(dns.FakeIPv4Pool, 65535)
 	if err != nil {
 		return nil, err
 	}
@@ -117,9 +135,92 @@ func (fkdns *Holder) GetDomainFromFakeDNS(ip net.Address) string {
 	return ""
 }
 
-// GetFakeIPRange return fake IP range from configuration
-func (fkdns *Holder) GetFakeIPRange() *gonet.IPNet {
-	return fkdns.ipRange
+type HolderMulti struct {
+	holders []*Holder
+
+	config *FakeDnsPoolMulti
+}
+
+func (h *HolderMulti) IsIPInIPPool(ip net.Address) bool {
+	if ip.Family().IsDomain() {
+		return false
+	}
+	for _, v := range h.holders {
+		if v.IsIPInIPPool(ip) {
+			return true
+		}
+	}
+	return false
+}
+
+func (h *HolderMulti) GetFakeIPForDomain3(domain string, ipv4, ipv6 bool) []net.Address {
+	var ret []net.Address
+	for _, v := range h.holders {
+		ret = append(ret, v.GetFakeIPForDomain3(domain, ipv4, ipv6)...)
+	}
+	return ret
+}
+
+func (h *HolderMulti) GetFakeIPForDomain(domain string) []net.Address {
+	var ret []net.Address
+	for _, v := range h.holders {
+		ret = append(ret, v.GetFakeIPForDomain(domain)...)
+	}
+	return ret
+}
+
+func (h *HolderMulti) GetDomainFromFakeDNS(ip net.Address) string {
+	for _, v := range h.holders {
+		if domain := v.GetDomainFromFakeDNS(ip); domain != "" {
+			return domain
+		}
+	}
+	return ""
+}
+
+func (h *HolderMulti) Type() interface{} {
+	return (*dns.FakeDNSEngine)(nil)
+}
+
+func (h *HolderMulti) Start() error {
+	for _, v := range h.holders {
+		if v.config != nil && v.config.IpPool != "" && v.config.LruSize != 0 {
+			if err := v.Start(); err != nil {
+				return newError("Cannot start all fake dns pools").Base(err)
+			}
+		} else {
+			return newError("invalid fakeDNS setting")
+		}
+	}
+	return nil
+}
+
+func (h *HolderMulti) Close() error {
+	for _, v := range h.holders {
+		if err := v.Close(); err != nil {
+			return newError("Cannot close all fake dns pools").Base(err)
+		}
+	}
+	return nil
+}
+
+func (h *HolderMulti) createHolderGroups() error {
+	for _, v := range h.config.Pools {
+		holder, err := NewFakeDNSHolderConfigOnly(v)
+		if err != nil {
+			return err
+		}
+		h.holders = append(h.holders, holder)
+	}
+	return nil
+}
+
+func NewFakeDNSHolderMulti(conf *FakeDnsPoolMulti) (*HolderMulti, error) {
+	holderMulti := &HolderMulti{nil, conf}
+	if err := holderMulti.createHolderGroups(); err != nil {
+		return nil, err
+	}
+	return holderMulti, nil
 }
 
 func init() {
@@ -131,4 +232,13 @@ func init() {
 		}
 		return f, nil
 	}))
+
+	common.Must(common.RegisterConfig((*FakeDnsPoolMulti)(nil), func(ctx context.Context, config interface{}) (interface{}, error) {
+		var f *HolderMulti
+		var err error
+		if f, err = NewFakeDNSHolderMulti(config.(*FakeDnsPoolMulti)); err != nil {
+			return nil, err
+		}
+		return f, nil
+	}))
 }