From 66189202346bb8d49267b116cbb553e26f5db96f Mon Sep 17 00:00:00 2001
From: Justin Bastress <jbastress@censys.io>
Date: Mon, 1 Oct 2018 11:08:26 -0400
Subject: [PATCH] add some tighter bounds checking in MSSQL scanner, and if
 there is an uncaught panic, log the body that caused it

---
 modules/mssql/connection.go | 20 ++++++++++++++++++--
 utility.go                  | 17 +++++++++++++++++
 2 files changed, 35 insertions(+), 2 deletions(-)

diff --git a/modules/mssql/connection.go b/modules/mssql/connection.go
index 58e5e83..7cce7d4 100644
--- a/modules/mssql/connection.go
+++ b/modules/mssql/connection.go
@@ -11,7 +11,7 @@ import (
 	"strings"
 	"time"
 
-	logrus "github.com/sirupsen/logrus"
+	"github.com/sirupsen/logrus"
 	"github.com/zmap/zgrab2"
 )
 
@@ -401,6 +401,9 @@ func (options PreloginOptions) Encode() ([]byte, error) {
 	for _, ik := range sortedKeys {
 		k := PreloginOptionToken(ik)
 		v := options[k]
+		if len(cursor) < 5 {
+			return nil, fmt.Errorf("encode: size mismatch (options.Size()=%d)", options.Size())
+		}
 		cursor[0] = byte(k)
 		if offset > 0xffff {
 			return nil, ErrTooLarge
@@ -411,6 +414,9 @@ func (options PreloginOptions) Encode() ([]byte, error) {
 		offset += len(v)
 		cursor = cursor[5:]
 	}
+	if len(cursor) < 1 {
+		return nil, fmt.Errorf("encode: size mismatch (options.Size()=%d, len(sortedKeys)=%d)", options.Size(), len(sortedKeys))
+	}
 	// Write the terminator after the last PL_OPTION header
 	// (and just before the first value)
 	cursor[0] = 0xff
@@ -421,6 +427,9 @@ func (options PreloginOptions) Encode() ([]byte, error) {
 // returned in rest.
 // If body can't be decoded as a PRELOGIN body, returns nil, nil, ErrInvalidData
 func decodePreloginOptions(body []byte) (result *PreloginOptions, rest []byte, err error) {
+	if len(body) < 1 {
+		return nil, nil, ErrInvalidData
+	}
 	cursor := body[:]
 	options := make(PreloginOptions)
 	max := 0
@@ -506,7 +515,13 @@ func (options PreloginOptions) MarshalJSON() ([]byte, error) {
 
 	fedAuthRequired, hasFedAuthRequired := opts[PreloginFedAuthRequired]
 	if hasFedAuthRequired {
-		aux.FedAuthRequired = &fedAuthRequired[0]
+		temp := uint8(0)
+		if len(fedAuthRequired) > 0 {
+			temp = fedAuthRequired[0]
+		} else {
+			logrus.Debugf("fedAuthRequired was present but empty (options=%#v)", options)
+		}
+		aux.FedAuthRequired = &temp
 	}
 
 	nonce, hasNonce := opts[PreloginNonce]
@@ -632,6 +647,7 @@ func (connection *Connection) readPreloginPacket() (*TDSPacket, *PreloginOptions
 	if packet.Type != TDSPacketTypeTabularResult {
 		return packet, nil, &zgrab2.ScanError{Status: zgrab2.SCAN_APPLICATION_ERROR, Err: err}
 	}
+	defer zgrab2.LogPanic("Error decoding Prelogin packet %#v", packet.Body)
 	plOptions, rest, err := decodePreloginOptions(packet.Body)
 	if err != nil {
 		return packet, nil, err
diff --git a/utility.go b/utility.go
index 8ac1e05..44fd536 100644
--- a/utility.go
+++ b/utility.go
@@ -10,6 +10,8 @@ import (
 	"time"
 
 	"github.com/zmap/zflags"
+	"github.com/sirupsen/logrus"
+	"runtime/debug"
 )
 
 var parser *flags.Parser
@@ -209,3 +211,18 @@ func IsTimeoutError(err error) bool {
 
 	return false
 }
+
+// LogPanic is intended to be called from within defer -- if there was no panic, it returns without
+// doing anything. Otherwise, it logs the stacktrace, the panic error, and the provided message
+// before re-raising the original panic.
+// Example:
+//     defer zgrab2.LogPanic("Error decoding body '%x'", body)
+func LogPanic(format string, args...interface{}) {
+	err := recover()
+	if err == nil {
+		return
+	}
+	logrus.Errorf("Uncaught panic at %s: %v", string(debug.Stack()), err)
+	logrus.Errorf(format, args...)
+	panic(err)
+}