Modify DNS resolver option (#400)

raviqqe · Jul 17, 2024 · 985b1a3 · 985b1a3
1 parent 743a4e0
commit 985b1a3
Show file tree

Hide file tree

Showing 6 changed files with 19 additions and 21 deletions.
diff --git a/.snapshots/TestHelp b/.snapshots/TestHelp
@@ -22,6 +22,7 @@ Application Options:
                                             sitemap.xml (deprecated)
       --header=<header>...                  Custom headers
   -f, --ignore-fragments                    Ignore URL fragments
+      --dns-resolver=<address>              Custom DNS resolver
       --format=[text|json|junit]            Output format (default: text)
       --json                                Output results in JSON (deprecated)
       --experimental-verbose-json           Include successful results in JSON

diff --git a/arguments.go b/arguments.go
@@ -16,14 +16,14 @@ type arguments struct {
 	MaxConnections         int      `short:"c" long:"max-connections" value-name:"<count>" default:"512" description:"Maximum number of HTTP connections"`
 	MaxConnectionsPerHost  int      `long:"max-connections-per-host" value-name:"<count>" default:"512" description:"Maximum number of HTTP connections per host"`
 	MaxResponseBodySize    int      `long:"max-response-body-size" value-name:"<size>" default:"10000000" description:"Maximum response body size to read"`
-	CustomDnsAddr          string   `long:"custom-dns-addr" description:"Custom DNS server to be used for requests"`
 	RawExcludedPatterns    []string `short:"e" long:"exclude" value-name:"<pattern>..." description:"Exclude URLs matched with given regular expressions"`
 	RawIncludedPatterns    []string `short:"i" long:"include" value-name:"<pattern>..." description:"Include URLs matched with given regular expressions"`
 	FollowRobotsTxt        bool     `long:"follow-robots-txt" description:"Follow robots.txt when scraping pages"`
 	FollowSitemapXML       bool     `long:"follow-sitemap-xml" description:"Scrape only pages listed in sitemap.xml (deprecated)"`
 	RawHeaders             []string `long:"header" value-name:"<header>..." description:"Custom headers"`
 	// TODO Remove a short option.
 	IgnoreFragments bool   `short:"f" long:"ignore-fragments" description:"Ignore URL fragments"`
+	DnsResolver     string `long:"dns-resolver" value-name:"<address>" description:"Custom DNS resolver"`
 	Format          string `long:"format" description:"Output format" default:"text" choice:"text" choice:"json" choice:"junit"`
 	// TODO Remove this option.
 	JSONOutput bool `long:"json" description:"Output results in JSON (deprecated)"`

diff --git a/command.go b/command.go
@@ -56,7 +56,7 @@ func (c *command) runWithError(ss []string) (bool, error) {
 						SkipTLSVerification:   args.SkipTLSVerification,
 						Timeout:               time.Duration(args.Timeout) * time.Second,
 						Header:                args.Header,
-						CustomDnsAddr:         args.CustomDnsAddr,
+						DnsResolver:           args.DnsResolver,
 					},
 				),
 				args.RateLimit,

diff --git a/fasthttp_http_client_factory.go b/fasthttp_http_client_factory.go
@@ -17,28 +17,25 @@ func newFasthttpHttpClientFactory() *fasthttpHttpClientFactory {
 }
 
 func (*fasthttpHttpClientFactory) Create(o httpClientOptions) httpClient {
-	d := func(addr string) (net.Conn, error) {
-		return fasthttp.DialTimeout(addr, tcpTimeout)
+	d := func(address string) (net.Conn, error) {
+		return fasthttp.DialTimeout(address, tcpTimeout)
 	}
 
 	if o.Proxy != "" {
 		d = fasthttpproxy.FasthttpHTTPDialerTimeout(o.Proxy, tcpTimeout)
-	}
-
-	if o.CustomDnsAddr != "" {
-		d = func(addr string) (net.Conn, error) {
-			dialer := fasthttp.TCPDialer{
-				Concurrency: 1000,
-				Resolver: &net.Resolver{
-					PreferGo:     true,
-					StrictErrors: false,
-					Dial: func(ctx context.Context, network string, address string) (net.Conn, error) {
-						internalDialer := net.Dialer{}
-						return internalDialer.DialContext(ctx, "udp", o.CustomDnsAddr)
-					},
+	} else if o.DnsResolver != "" {
+		nd := &net.Dialer{}
+		td := &fasthttp.TCPDialer{
+			Concurrency: concurrency,
+			Resolver: &net.Resolver{
+				Dial: func(ctx context.Context, network, address string) (net.Conn, error) {
+					return nd.DialContext(ctx, "udp", o.DnsResolver)
 				},
-			}
-			return dialer.DialTimeout(addr, tcpTimeout)
+			},
+		}
+
+		d = func(address string) (net.Conn, error) {
+			return td.DialTimeout(address, tcpTimeout)
 		}
 	}
 

diff --git a/fasthttp_http_client_factory_test.go b/fasthttp_http_client_factory_test.go
@@ -11,5 +11,5 @@ func TestFasthttpHttpClientFactoryCreateWithProxy(t *testing.T) {
 }
 
 func TestFasthttpHttpClientFactoryCreateWithCustomDnsResolver(t *testing.T) {
-	newFasthttpHttpClientFactory().Create(httpClientOptions{CustomDnsAddr: "1.1.1.1:53"})
+	newFasthttpHttpClientFactory().Create(httpClientOptions{DnsResolver: "1.1.1.1:53"})
 }
diff --git a/http_client_factory.go b/http_client_factory.go
@@ -13,7 +13,7 @@ type httpClientOptions struct {
 	SkipTLSVerification bool
 	Timeout             time.Duration
 	Header              http.Header
-	CustomDnsAddr       string
+	DnsResolver         string
 }
 
 type httpClientFactory interface {