From abe6de223df2324f89384f2d05627aee2ec55b3a Mon Sep 17 00:00:00 2001
From: Miek Gieben <miek@miek.nl>
Date: Sat, 5 Apr 2014 08:06:49 +0100
Subject: [PATCH] Add some more checks

---
 msg.go        | 14 ++++++++++----
 parse_test.go | 15 +++++++++------
 2 files changed, 19 insertions(+), 10 deletions(-)

diff --git a/msg.go b/msg.go
index 0cd9813c..f537f242 100644
--- a/msg.go
+++ b/msg.go
@@ -976,7 +976,7 @@ func unpackStructValue(val reflect.Value, msg []byte, off int) (off1 int, err er
 				if off == rdend {
 					break
 				}
-				if off+net.IPv6len > rdend {
+				if off+net.IPv6len > rdend || off+net.IPv6len > lenmsg {
 					return lenmsg, &Error{err: "overflow unpacking aaaa"}
 				}
 				fv.Set(reflect.ValueOf(net.IP{msg[off], msg[off+1], msg[off+2], msg[off+3], msg[off+4],
@@ -988,6 +988,9 @@ func unpackStructValue(val reflect.Value, msg []byte, off int) (off1 int, err er
 				serv := make([]uint16, 0)
 				j := 0
 				for off < rdend {
+					if off+1 > lenmsg {
+						return lenmsg, &Error{err: "overflow unpacking wks"}
+					}
 					b := msg[off]
 					// Check the bits one by one, and set the type
 					if b&0x80 == 0x80 {
@@ -1023,7 +1026,7 @@ func unpackStructValue(val reflect.Value, msg []byte, off int) (off1 int, err er
 					break
 				}
 				// Rest of the record is the type bitmap
-				if off+2 > rdend {
+				if off+2 > rdend || off+2 > lenmsg {
 					return lenmsg, &Error{err: "overflow unpacking nsecx"}
 				}
 				nsec := make([]uint16, 0)
@@ -1037,15 +1040,18 @@ func unpackStructValue(val reflect.Value, msg []byte, off int) (off1 int, err er
 						// A length window of zero is strange. If there
 						// the window should not have been specified. Bail out
 						// println("dns: length == 0 when unpacking NSEC")
-						return lenmsg, ErrRdata
+						return lenmsg, &Error{err: "overflow unpacking nsecx"}
 					}
 					if length > 32 {
-						return lenmsg, ErrRdata
+						return lenmsg, &Error{err: "overflow unpacking nsecx"}
 					}
 
 					// Walk the bytes in the window - and check the bit settings...
 					off += 2
 					for j := 0; j < length; j++ {
+						if off+j+1 > lenmsg {
+							return lenmsg, &Error{err: "overflow unpacking nsecx"}
+						}
 						b := msg[off+j]
 						// Check the bits one by one, and set the type
 						if b&0x80 == 0x80 {
diff --git a/parse_test.go b/parse_test.go
index aebd5117..a826d540 100644
--- a/parse_test.go
+++ b/parse_test.go
@@ -1115,11 +1115,14 @@ func TestTxtLong(t *testing.T) {
 	}
 }
 
-func TestMalformedPacket1(t *testing.T) {
-	packet := "00441553000000010000000000010563646e6a730a636c6f7564666c61726503636f6d0363646e0a636c6f7564666c617265036e657400001c00010000291000000080000000"
-	data, _ := hex.DecodeString(packet)
+func TestMalformedPackets(t *testing.T) {
+	var packets = []string{
+		"0021641c000000010000000000000b757361706f6f6c70726f7303636f6d0000100001",
+	}
 
-	// This crashes godns
-	var msg Msg
-	msg.Unpack(data)
+	for _, packet := range packets {
+		data, _ := hex.DecodeString(packet)
+		var msg Msg
+		msg.Unpack(data)
+	}
 }