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

This commit is contained in:
Dreamacro
2022-08-13 13:07:35 +08:00
parent 5940f62794
commit 3946d771e5
5 changed files with 129 additions and 43 deletions

View File

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