From b27f153426a438ac10ab274b00674d267276bd6a Mon Sep 17 00:00:00 2001
From: nadoo <287492+nadoo@users.noreply.github.com>
Date: Mon, 6 Aug 2018 08:03:07 +0800
Subject: [PATCH] dns: check bytes length in unmarshal message

---
 dns/client.go  | 1 +
 dns/message.go | 8 ++++++++
 2 files changed, 9 insertions(+)

diff --git a/dns/client.go b/dns/client.go
index 617665f..ede53b9 100644
--- a/dns/client.go
+++ b/dns/client.go
@@ -134,6 +134,7 @@ func (c *Client) exchange(qname string, reqBytes []byte, preferTCP bool) (server
 		if err == nil {
 			break
 		}
+
 		log.F("[dns] failed to exchange with server %v: %v", server, err)
 	}
 
diff --git a/dns/message.go b/dns/message.go
index c583c2c..d8613f3 100644
--- a/dns/message.go
+++ b/dns/message.go
@@ -119,6 +119,10 @@ func (m *Message) Marshal() ([]byte, error) {
 
 // UnmarshalMessage unmarshals []bytes to Message
 func UnmarshalMessage(b []byte) (*Message, error) {
+	if len(b) <= 2+HeaderLen {
+		return nil, errors.New("UnmarshalMessage: not enough data")
+	}
+
 	m := &Message{unMarshaled: b}
 	err := UnmarshalHeader(b[:HeaderLen], &m.Header)
 	if err != nil {
@@ -280,6 +284,10 @@ func (m *Message) UnmarshalQuestion(b []byte, q *Question) (n int, err error) {
 		return 0, errors.New("unmarshal question must not be nil")
 	}
 
+	if len(b) <= 5 {
+		return 0, errors.New("UnmarshalQuestion: not enough data")
+	}
+
 	domain, idx, err := m.UnmarshalDomain(b)
 	if err != nil {
 		return 0, err