From 7274dcf622388ee960468c9c956400d98dc6d3d4 Mon Sep 17 00:00:00 2001
From: nadoo <287492+nadoo@users.noreply.github.com>
Date: Mon, 6 Aug 2018 08:41:54 +0800
Subject: [PATCH] dns: set timout only when multiple upstream dns servers set

---
 dns/client.go  | 5 ++++-
 dns/message.go | 2 +-
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/dns/client.go b/dns/client.go
index bdf612c..a8b4ecd 100644
--- a/dns/client.go
+++ b/dns/client.go
@@ -118,7 +118,10 @@ func (c *Client) exchange(qname string, reqBytes []byte, preferTCP bool) (server
 	for _, server = range servers {
 		rc, err = dialer.Dial(network, server)
 		// TODO: support timeout setting for different upstream server
-		rc.SetDeadline(time.Now().Add(time.Duration(3) * time.Second))
+		if len(servers) > 1 {
+			rc.SetDeadline(time.Now().Add(time.Duration(3) * time.Second))
+		}
+
 		if err != nil {
 			log.F("[dns] failed to connect to server %v: %v", server, err)
 			continue
diff --git a/dns/message.go b/dns/message.go
index d8613f3..e6c4b6b 100644
--- a/dns/message.go
+++ b/dns/message.go
@@ -119,7 +119,7 @@ func (m *Message) Marshal() ([]byte, error) {
 
 // UnmarshalMessage unmarshals []bytes to Message
 func UnmarshalMessage(b []byte) (*Message, error) {
-	if len(b) <= 2+HeaderLen {
+	if len(b) < HeaderLen {
 		return nil, errors.New("UnmarshalMessage: not enough data")
 	}