Feature: sync missing resolver logic from premium, but still net.IP on opensource

dns/client.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"crypto/tls"
 	"fmt"
+	"math/rand"
 	"net"
 	"strings"
 
@@ -36,9 +37,13 @@ func (c *client) ExchangeContext(ctx context.Context, m *D.Msg) (*D.Msg, error)
 			return nil, fmt.Errorf("dns %s not a valid ip", c.host)
 		}
 	} else {
-		if ip, err = resolver.ResolveIPWithResolver(c.host, c.r); err != nil {
+		ips, err := resolver.LookupIPWithResolver(ctx, c.host, c.r)
+		if err != nil {
 			return nil, fmt.Errorf("use default dns resolve failed: %w", err)
+		} else if len(ips) == 0 {
+			return nil, fmt.Errorf("%w: %s", resolver.ErrIPNotFound, c.host)
 		}
+		ip = ips[rand.Intn(len(ips))]
 	}
 
 	network := "udp"

dns/doh.go

@@ -4,7 +4,9 @@ import (
 	"bytes"
 	"context"
 	"crypto/tls"
+	"fmt"
 	"io"
+	"math/rand"
 	"net"
 	"net/http"
 
@@ -91,10 +93,13 @@ func newDoHClient(url, iface string, r *Resolver) *dohClient {
 					return nil, err
 				}
 
-				ip, err := resolver.ResolveIPWithResolver(host, r)
+				ips, err := resolver.LookupIPWithResolver(ctx, host, r)
 				if err != nil {
 					return nil, err
+				} else if len(ips) == 0 {
+					return nil, fmt.Errorf("%w: %s", resolver.ErrIPNotFound, host)
 				}
+				ip := ips[rand.Intn(len(ips))]
 
 				options := []dialer.Option{}
 				if iface != "" {

dns/resolver.go

@@ -42,19 +42,23 @@ type Resolver struct {
 	policy                *trie.DomainTrie
 }
 
-// ResolveIP request with TypeA and TypeAAAA, priority return TypeA
-func (r *Resolver) ResolveIP(host string) (ip net.IP, err error) {
-	ch := make(chan net.IP, 1)
+// LookupIP request with TypeA and TypeAAAA, priority return TypeA
+func (r *Resolver) LookupIP(ctx context.Context, host string) (ip []net.IP, err error) {
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
+	ch := make(chan []net.IP, 1)
+
 	go func() {
 		defer close(ch)
-		ip, err := r.resolveIP(host, D.TypeAAAA)
+		ip, err := r.lookupIP(ctx, host, D.TypeAAAA)
 		if err != nil {
 			return
 		}
 		ch <- ip
 	}()
 
-	ip, err = r.resolveIP(host, D.TypeA)
+	ip, err = r.lookupIP(ctx, host, D.TypeA)
 	if err == nil {
 		return
 	}
@@ -67,14 +71,47 @@ func (r *Resolver) ResolveIP(host string) (ip net.IP, err error) {
 	return ip, nil
 }
 
+// ResolveIP request with TypeA and TypeAAAA, priority return TypeA
+func (r *Resolver) ResolveIP(host string) (ip net.IP, err error) {
+	ips, err := r.LookupIP(context.Background(), host)
+	if err != nil {
+		return nil, err
+	} else if len(ips) == 0 {
+		return nil, fmt.Errorf("%w: %s", resolver.ErrIPNotFound, host)
+	}
+	return ips[rand.Intn(len(ips))], nil
+}
+
+// LookupIPv4 request with TypeA
+func (r *Resolver) LookupIPv4(ctx context.Context, host string) ([]net.IP, error) {
+	return r.lookupIP(ctx, host, D.TypeA)
+}
+
 // ResolveIPv4 request with TypeA
 func (r *Resolver) ResolveIPv4(host string) (ip net.IP, err error) {
-	return r.resolveIP(host, D.TypeA)
+	ips, err := r.lookupIP(context.Background(), host, D.TypeA)
+	if err != nil {
+		return nil, err
+	} else if len(ips) == 0 {
+		return nil, fmt.Errorf("%w: %s", resolver.ErrIPNotFound, host)
+	}
+	return ips[rand.Intn(len(ips))], nil
+}
+
+// LookupIPv6 request with TypeAAAA
+func (r *Resolver) LookupIPv6(ctx context.Context, host string) ([]net.IP, error) {
+	return r.lookupIP(ctx, host, D.TypeAAAA)
 }
 
 // ResolveIPv6 request with TypeAAAA
 func (r *Resolver) ResolveIPv6(host string) (ip net.IP, err error) {
-	return r.resolveIP(host, D.TypeAAAA)
+	ips, err := r.lookupIP(context.Background(), host, D.TypeAAAA)
+	if err != nil {
+		return nil, err
+	} else if len(ips) == 0 {
+		return nil, fmt.Errorf("%w: %s", resolver.ErrIPNotFound, host)
+	}
+	return ips[rand.Intn(len(ips))], nil
 }
 
 func (r *Resolver) shouldIPFallback(ip net.IP) bool {
@@ -253,14 +290,15 @@ func (r *Resolver) ipExchange(ctx context.Context, m *D.Msg) (msg *D.Msg, err er
 	return
 }
 
-func (r *Resolver) resolveIP(host string, dnsType uint16) (ip net.IP, err error) {
-	ip = net.ParseIP(host)
+func (r *Resolver) lookupIP(ctx context.Context, host string, dnsType uint16) ([]net.IP, error) {
+	ip := net.ParseIP(host)
 	if ip != nil {
-		isIPv4 := ip.To4() != nil
+		ip4 := ip.To4()
+		isIPv4 := ip4 != nil
 		if dnsType == D.TypeAAAA && !isIPv4 {
-			return ip, nil
+			return []net.IP{ip}, nil
 		} else if dnsType == D.TypeA && isIPv4 {
-			return ip, nil
+			return []net.IP{ip4}, nil
 		} else {
 			return nil, resolver.ErrIPVersion
 		}
@@ -275,13 +313,10 @@ func (r *Resolver) resolveIP(host string, dnsType uint16) (ip net.IP, err error)
 	}
 
 	ips := msgToIP(msg)
-	ipLength := len(ips)
-	if ipLength == 0 {
+	if len(ips) == 0 {
 		return nil, resolver.ErrIPNotFound
 	}
-
-	ip = ips[rand.Intn(ipLength)]
-	return
+	return ips, nil
 }
 
 func (r *Resolver) msgToDomain(msg *D.Msg) string {