consume context in local nameserver.

app/dns/nameserver.go

@@ -2,31 +2,34 @@ package dns
 
 import (
 	"context"
-	"time"
 
 	"v2ray.com/core/common/net"
 )
 
-var (
-	multiQuestionDNS = map[net.Address]bool{
-		net.IPAddress([]byte{8, 8, 8, 8}): true,
-		net.IPAddress([]byte{8, 8, 4, 4}): true,
-		net.IPAddress([]byte{9, 9, 9, 9}): true,
-	}
-)
-
-type ARecord struct {
-	IPs    []net.IP
-	Expire time.Time
-}
-
 type NameServer interface {
 	QueryIP(ctx context.Context, domain string) ([]net.IP, error)
 }
 
 type LocalNameServer struct {
+	resolver net.Resolver
 }
 
-func (*LocalNameServer) QueryIP(ctx context.Context, domain string) ([]net.IP, error) {
-	return net.LookupIP(domain)
+func (s *LocalNameServer) QueryIP(ctx context.Context, domain string) ([]net.IP, error) {
+	ipAddr, err := s.resolver.LookupIPAddr(ctx, domain)
+	if err != nil {
+		return nil, err
+	}
+	var ips []net.IP
+	for _, addr := range ipAddr {
+		ips = append(ips, addr.IP)
+	}
+	return ips, nil
+}
+
+func NewLocalNameServer() *LocalNameServer {
+	return &LocalNameServer{
+		resolver: net.Resolver{
+			PreferGo: true,
+		},
+	}
 }

app/dns/nameserver_test.go

@@ -0,0 +1,21 @@
+package dns_test
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	. "v2ray.com/core/app/dns"
+	. "v2ray.com/ext/assert"
+)
+
+func TestLocalNameServer(t *testing.T) {
+	assert := With(t)
+
+	s := NewLocalNameServer()
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second*2)
+	ips, err := s.QueryIP(ctx, "google.com")
+	cancel()
+	assert(err, IsNil)
+	assert(len(ips), GreaterThan, 0)
+}

app/dns/udpns.go

@@ -16,6 +16,14 @@ import (
 	"v2ray.com/core/transport/internet/udp"
 )
 
+var (
+	multiQuestionDNS = map[net.Address]bool{
+		net.IPAddress([]byte{8, 8, 8, 8}): true,
+		net.IPAddress([]byte{8, 8, 4, 4}): true,
+		net.IPAddress([]byte{9, 9, 9, 9}): true,
+	}
+)
+
 type IPRecord struct {
 	IP     net.IP
 	Expire time.Time

common/net/system.go

@@ -51,3 +51,5 @@ type TCPListener = net.TCPListener
 type UnixListener = net.UnixListener
 
 var ResolveUnixAddr = net.ResolveUnixAddr
+
+type Resolver = net.Resolver