From 88932f9167a530ef89b474c47b09cb538decfe5c Mon Sep 17 00:00:00 2001
From: zealic <zealic@gmail.com>
Date: Wed, 25 May 2016 11:22:09 +0800
Subject: [PATCH] Add dns-timeout support.

---
 acme/dns_challenge.go | 8 ++++++--
 cli.go                | 4 ++++
 cli_handlers.go       | 4 ++++
 3 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/acme/dns_challenge.go b/acme/dns_challenge.go
index b77e916a..2f45e2a9 100644
--- a/acme/dns_challenge.go
+++ b/acme/dns_challenge.go
@@ -26,6 +26,9 @@ var RecursiveNameservers = []string{
 	"google-public-dns-b.google.com:53",
 }
 
+// DNSTimeout is used to override the default DNS timeout of 10 seconds.
+var DNSTimeout = 10 * time.Second
+
 // DNS01Record returns a DNS record which will fulfill the `dns-01` challenge
 func DNS01Record(domain, keyAuth string) (fqdn string, value string, ttl int) {
 	keyAuthShaBytes := sha256.Sum256([]byte(keyAuth))
@@ -161,10 +164,11 @@ func dnsQuery(fqdn string, rtype uint16, nameservers []string, recursive bool) (
 	// Will retry the request based on the number of servers (n+1)
 	for i := 1; i <= len(nameservers)+1; i++ {
 		ns := nameservers[i%len(nameservers)]
-		in, err = dns.Exchange(m, ns)
+		udp := &dns.Client{Net: "udp", Timeout: DNSTimeout}
+		in, _, err = udp.Exchange(m, ns)
 
 		if err == dns.ErrTruncated {
-			tcp := &dns.Client{Net: "tcp"}
+			tcp := &dns.Client{Net: "tcp", Timeout: DNSTimeout}
 			// If the TCP request suceeds, the err will reset to nil
 			in, _, err = tcp.Exchange(m, ns)
 		}
diff --git a/cli.go b/cli.go
index 438557fe..947ed345 100644
--- a/cli.go
+++ b/cli.go
@@ -150,6 +150,10 @@ func main() {
 			Name:  "http-timeout",
 			Usage: "Set the HTTP timeout value to a specific value in seconds. The default is 10 seconds.",
 		},
+		cli.IntFlag{
+			Name:  "dns-timeout",
+			Usage: "Set the DNS timeout value to a specific value in seconds. The default is 10 seconds.",
+		},
 	}
 
 	err = app.Run(os.Args)
diff --git a/cli_handlers.go b/cli_handlers.go
index 6a2fe56e..06d534c4 100644
--- a/cli_handlers.go
+++ b/cli_handlers.go
@@ -37,6 +37,10 @@ func setup(c *cli.Context) (*Configuration, *Account, *acme.Client) {
 		acme.HTTPTimeout = time.Duration(c.GlobalInt("http-timeout")) * time.Second
 	}
 
+	if c.GlobalIsSet("dns-timeout") {
+		acme.DNSTimeout = time.Duration(c.GlobalInt("dns-timeout")) * time.Second
+	}
+
 	err := checkFolder(c.GlobalString("path"))
 	if err != nil {
 		logger().Fatalf("Could not check/create path: %s", err.Error())