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/dispatcher/default.go

@@ -193,10 +193,15 @@ func shouldOverride(ctx context.Context, result SniffResult, request session.Sni
 		if strings.HasPrefix(protocolString, p) {
 			return true
 		}
-		if fakeDNSEngine != nil && protocolString != "bittorrent" && p == "fakedns" &&
-			destination.Address.Family().IsIP() && fakeDNSEngine.GetFakeIPRange().Contains(destination.Address.IP()) {
-			newError("Using sniffer ", protocolString, " since the fake DNS missed").WriteToLog(session.ExportIDToError(ctx))
-			return true
+		if fkr0, ok := fakeDNSEngine.(dns.FakeDNSEngineRev0); ok && protocolString != "bittorrent" && p == "fakedns" &&
+			destination.Address.Family().IsIP() && fkr0.IsIPInIPPool(destination.Address) {
+				newError("Using sniffer ", protocolString, " since the fake DNS missed").WriteToLog(session.ExportIDToError(ctx))
+				return true
+		}
+		if resultSubset, ok := result.(SnifferIsProtoSubsetOf); ok {
+			if resultSubset.IsProtoSubsetOf(p) {
+				return true
+			}
 		}
 	}
 

app/dispatcher/fakednssniffer.go

@@ -2,6 +2,7 @@ package dispatcher
 
 import (
 	"context"
+	"strings"
 
 	"github.com/xtls/xray-core/common"
 	"github.com/xtls/xray-core/common/net"
@@ -32,6 +33,15 @@ func newFakeDNSSniffer(ctx context.Context) (protocolSnifferWithMetadata, error)
 				return &fakeDNSSniffResult{domainName: domainFromFakeDNS}, nil
 			}
 		}
+
+		if ipAddressInRangeValueI := ctx.Value(ipAddressInRange); ipAddressInRangeValueI != nil {
+			ipAddressInRangeValue := ipAddressInRangeValueI.(*ipAddressInRangeOpt)
+			if fkr0, ok := fakeDNSEngine.(dns.FakeDNSEngineRev0); ok {
+				inPool := fkr0.IsIPInIPPool(Target.Address)
+				ipAddressInRangeValue.addressInRange = &inPool
+			}
+		}
+
 		return nil, common.ErrNoClue
 	}, metadataSniffer: true}, nil
 }
@@ -47,3 +57,61 @@ func (fakeDNSSniffResult) Protocol() string {
 func (f fakeDNSSniffResult) Domain() string {
 	return f.domainName
 }
+
+type fakeDNSExtraOpts int
+
+const ipAddressInRange fakeDNSExtraOpts = 1
+
+type ipAddressInRangeOpt struct {
+	addressInRange *bool
+}
+
+type DNSThenOthersSniffResult struct {
+	domainName           string
+	protocolOriginalName string
+}
+
+func (f DNSThenOthersSniffResult) IsProtoSubsetOf(protocolName string) bool {
+	return strings.HasPrefix(protocolName, f.protocolOriginalName)
+}
+
+func (DNSThenOthersSniffResult) Protocol() string {
+	return "fakedns+others"
+}
+
+func (f DNSThenOthersSniffResult) Domain() string {
+	return f.domainName
+}
+
+func newFakeDNSThenOthers(ctx context.Context, fakeDNSSniffer protocolSnifferWithMetadata, others []protocolSnifferWithMetadata) (
+	protocolSnifferWithMetadata, error) { // nolint: unparam
+	// ctx may be used in the future
+	_ = ctx
+	return protocolSnifferWithMetadata{
+		protocolSniffer: func(ctx context.Context, bytes []byte) (SniffResult, error) {
+			ipAddressInRangeValue := &ipAddressInRangeOpt{}
+			ctx = context.WithValue(ctx, ipAddressInRange, ipAddressInRangeValue)
+			result, err := fakeDNSSniffer.protocolSniffer(ctx, bytes)
+			if err == nil {
+				return result, nil
+			}
+			if ipAddressInRangeValue.addressInRange != nil {
+				if *ipAddressInRangeValue.addressInRange {
+					for _, v := range others {
+						if v.metadataSniffer || bytes != nil {
+							if result, err := v.protocolSniffer(ctx, bytes); err == nil {
+								return DNSThenOthersSniffResult{domainName: result.Domain(), protocolOriginalName: result.Protocol()}, nil
+							}
+						}
+					}
+					return nil, common.ErrNoClue
+				}
+				newError("ip address not in fake dns range, return as is").AtDebug().WriteToLog()
+				return nil, common.ErrNoClue
+			}
+			newError("fake dns sniffer did not set address in range option, assume false.").AtWarning().WriteToLog()
+			return nil, common.ErrNoClue
+		},
+		metadataSniffer: false,
+	}, nil
+}

app/dispatcher/sniffer.go

@@ -37,7 +37,12 @@ func NewSniffer(ctx context.Context) *Sniffer {
 		},
 	}
 	if sniffer, err := newFakeDNSSniffer(ctx); err == nil {
+		others := ret.sniffer
 		ret.sniffer = append(ret.sniffer, sniffer)
+		fakeDNSThenOthers, err := newFakeDNSThenOthers(ctx, sniffer, others)
+		if err == nil {
+			ret.sniffer = append([]protocolSnifferWithMetadata{fakeDNSThenOthers}, ret.sniffer...)
+		}
 	}
 	return ret
 }
@@ -121,3 +126,7 @@ func (c compositeResult) ProtocolForDomainResult() string {
 type SnifferResultComposite interface {
 	ProtocolForDomainResult() string
 }
+
+type SnifferIsProtoSubsetOf interface {
+	IsProtoSubsetOf(protocolName string) bool
+}