Fix buffer allocation issues

Fix a number or issues with how buffers for connections are processed including:
* Fix a number of uses of len(..) instead of cap(..) eliminating unnecessary buffer reallocation and invalid assignment of buffers smaller than defaultBufSize, which could cause a panic.
* Eliminate redundant size test in takeBuffer.
* Change buffer takeXXX functions to return an error to make it explicit that they can fail and fix missing error check in handleAuthResult.
* Always initialise nullMask in writeExecutePacket.

Also:
* Fix some typo's and unnecessary UTF-8 characters in comments.
* Add benchmarks for buffer and connection creation to validate memory allocations.
* Add myself / company to AUTHORS.

Author: stevenh

AUTHORS

@@ -73,6 +73,7 @@ Shuode Li <elemount at qq.com>
 Soroush Pour <me at soroushjp.com>
 Stan Putrya <root.vagner at gmail.com>
 Stanley Gunawan <gunawan.stanley at gmail.com>
+Steven Hartland <steven.hartland at multiplay.co.uk>
 Thomas Wodarek <wodarekwebpage at gmail.com>
 Tom Jenkinson <tom at tjenkinson.me>
 Xiangyu Hu <xiangyu.hu at outlook.com>
@@ -90,3 +91,4 @@ Keybase Inc.
 Percona LLC
 Pivotal Inc.
 Stripe Inc.
+Multiplay Ltd.

auth.go

@@ -360,13 +360,15 @@ func (mc *mysqlConn) handleAuthResult(oldAuthData []byte, plugin string) error {
 					pubKey := mc.cfg.pubKey
 					if pubKey == nil {
 						// request public key from server
-						data := mc.buf.takeSmallBuffer(4 + 1)
+						data, err := mc.buf.takeSmallBuffer(4 + 1)
+						if err != nil {
+							return err
+						}
 						data[4] = cachingSha2PasswordRequestPublicKey
 						mc.writePacket(data)
 
 						// parse public key
-						data, err := mc.readPacket()
-						if err != nil {
+						if data, err = mc.readPacket(); err != nil {
 							return err
 						}
 

benchmark_test.go

@@ -317,3 +317,30 @@ func BenchmarkExecContext(b *testing.B) {
 		})
 	}
 }
+
+var buf buffer
+
+func BenchmarkNewBuffer(b *testing.B) {
+	b.ReportAllocs()
+
+	var bu buffer
+	for i := 0; i < b.N; i++ {
+		bu = newBuffer(nil)
+	}
+
+	buf = bu
+}
+
+var con *mysqlConn
+
+func BenchmarkNewConn(b *testing.B) {
+	b.ReportAllocs()
+
+	var c *mysqlConn
+	for i := 0; i < b.N; i++ {
+		c = &mysqlConn{}
+		c.buf = newBuffer(nil)
+	}
+
+	con = c
+}

buffer.go

@@ -29,6 +29,7 @@ type buffer struct {
 	timeout time.Duration
 }
 
+// newBuffer allocates and returns a new buffer.
 func newBuffer(nc net.Conn) buffer {
 	var b [defaultBufSize]byte
 	return buffer{
@@ -49,11 +50,13 @@ func (b *buffer) fill(need int) error {
 	// grow buffer if necessary
 	// TODO: let the buffer shrink again at some point
 	//       Maybe keep the org buf slice and swap back?
-	if need > len(b.buf) {
+	if need > cap(b.buf) {
 		// Round up to the next multiple of the default size
 		newBuf := make([]byte, ((need/defaultBufSize)+1)*defaultBufSize)
 		copy(newBuf, b.buf)
 		b.buf = newBuf
+	} else if need > len(b.buf) {
+		b.buf = b.buf[:need]
 	}
 
 	b.idx = 0
@@ -109,39 +112,39 @@ func (b *buffer) readNext(need int) ([]byte, error) {
 // If possible, a slice from the existing buffer is returned.
 // Otherwise a bigger buffer is made.
 // Only one buffer (total) can be used at a time.
-func (b *buffer) takeBuffer(length int) []byte {
+func (b *buffer) takeBuffer(length int) ([]byte, error) {
 	if b.length > 0 {
-		return nil
+		return nil, ErrBusyBuffer
 	}
 
 	// test (cheap) general case first
-	if length <= defaultBufSize || length <= cap(b.buf) {
-		return b.buf[:length]
+	if length <= cap(b.buf) {
+		return b.buf[:length], nil
 	}
 
 	if length < maxPacketSize {
 		b.buf = make([]byte, length)
-		return b.buf
+		return b.buf, nil
 	}
-	return make([]byte, length)
+	return make([]byte, length), nil
 }
 
 // shortcut which can be used if the requested buffer is guaranteed to be
 // smaller than defaultBufSize
 // Only one buffer (total) can be used at a time.
-func (b *buffer) takeSmallBuffer(length int) []byte {
+func (b *buffer) takeSmallBuffer(length int) ([]byte, error) {
 	if b.length > 0 {
-		return nil
+		return nil, ErrBusyBuffer
 	}
-	return b.buf[:length]
+	return b.buf[:length], nil
 }
 
 // takeCompleteBuffer returns the complete existing buffer.
 // This can be used if the necessary buffer size is unknown.
 // Only one buffer (total) can be used at a time.
-func (b *buffer) takeCompleteBuffer() []byte {
+func (b *buffer) takeCompleteBuffer() ([]byte, error) {
 	if b.length > 0 {
-		return nil
+		return nil, ErrBusyBuffer
 	}
-	return b.buf
+	return b.buf, nil
 }

connection.go

@@ -182,10 +182,10 @@ func (mc *mysqlConn) interpolateParams(query string, args []driver.Value) (strin
 		return "", driver.ErrSkip
 	}
 
-	buf := mc.buf.takeCompleteBuffer()
-	if buf == nil {
+	buf, err := mc.buf.takeCompleteBuffer()
+	if err != nil {
 		// can not take the buffer. Something must be wrong with the connection
-		errLog.Print(ErrBusyBuffer)
+		errLog.Print(err)
 		return "", ErrInvalidConn
 	}
 	buf = buf[:0]

driver.go

@@ -50,7 +50,7 @@ func RegisterDial(net string, dial DialFunc) {
 
 // Open new Connection.
 // See https://github.com/go-sql-driver/mysql#dsn-data-source-name for how
-// the DSN string is formated
+// the DSN string is formatted
 func (d MySQLDriver) Open(dsn string) (driver.Conn, error) {
 	var err error
 

packets.go

@@ -51,7 +51,7 @@ func (mc *mysqlConn) readPacket() ([]byte, error) {
 		mc.sequence++
 
 		// packets with length 0 terminate a previous packet which is a
-		// multiple of (2^24)−1 bytes long
+		// multiple of (2^24)-1 bytes long
 		if pktLen == 0 {
 			// there was no previous packet
 			if prevData == nil {
@@ -288,10 +288,10 @@ func (mc *mysqlConn) writeHandshakeResponsePacket(authResp []byte, addNUL bool,
 	}
 
 	// Calculate packet length and get buffer with that size
-	data := mc.buf.takeSmallBuffer(pktLen + 4)
-	if data == nil {
+	data, err := mc.buf.takeSmallBuffer(pktLen + 4)
+	if err != nil {
 		// cannot take the buffer. Something must be wrong with the connection
-		errLog.Print(ErrBusyBuffer)
+		errLog.Print(err)
 		return errBadConnNoWrite
 	}
 
@@ -375,10 +375,10 @@ func (mc *mysqlConn) writeAuthSwitchPacket(authData []byte, addNUL bool) error {
 	if addNUL {
 		pktLen++
 	}
-	data := mc.buf.takeSmallBuffer(pktLen)
-	if data == nil {
+	data, err := mc.buf.takeSmallBuffer(pktLen)
+	if err != nil {
 		// cannot take the buffer. Something must be wrong with the connection
-		errLog.Print(ErrBusyBuffer)
+		errLog.Print(err)
 		return errBadConnNoWrite
 	}
 
@@ -399,10 +399,10 @@ func (mc *mysqlConn) writeCommandPacket(command byte) error {
 	// Reset Packet Sequence
 	mc.sequence = 0
 
-	data := mc.buf.takeSmallBuffer(4 + 1)
-	if data == nil {
+	data, err := mc.buf.takeSmallBuffer(4 + 1)
+	if err != nil {
 		// cannot take the buffer. Something must be wrong with the connection
-		errLog.Print(ErrBusyBuffer)
+		errLog.Print(err)
 		return errBadConnNoWrite
 	}
 
@@ -418,10 +418,10 @@ func (mc *mysqlConn) writeCommandPacketStr(command byte, arg string) error {
 	mc.sequence = 0
 
 	pktLen := 1 + len(arg)
-	data := mc.buf.takeBuffer(pktLen + 4)
-	if data == nil {
+	data, err := mc.buf.takeBuffer(pktLen + 4)
+	if err != nil {
 		// cannot take the buffer. Something must be wrong with the connection
-		errLog.Print(ErrBusyBuffer)
+		errLog.Print(err)
 		return errBadConnNoWrite
 	}
 
@@ -439,10 +439,10 @@ func (mc *mysqlConn) writeCommandPacketUint32(command byte, arg uint32) error {
 	// Reset Packet Sequence
 	mc.sequence = 0
 
-	data := mc.buf.takeSmallBuffer(4 + 1 + 4)
-	if data == nil {
+	data, err := mc.buf.takeSmallBuffer(4 + 1 + 4)
+	if err != nil {
 		// cannot take the buffer. Something must be wrong with the connection
-		errLog.Print(ErrBusyBuffer)
+		errLog.Print(err)
 		return errBadConnNoWrite
 	}
 
@@ -895,7 +895,7 @@ func (stmt *mysqlStmt) writeExecutePacket(args []driver.Value) error {
 	const minPktLen = 4 + 1 + 4 + 1 + 4
 	mc := stmt.mc
 
-	// Determine threshould dynamically to avoid packet size shortage.
+	// Determine threshold dynamically to avoid packet size shortage.
 	longDataSize := mc.maxAllowedPacket / (stmt.paramCount + 1)
 	if longDataSize < 64 {
 		longDataSize = 64
@@ -905,15 +905,16 @@ func (stmt *mysqlStmt) writeExecutePacket(args []driver.Value) error {
 	mc.sequence = 0
 
 	var data []byte
+	var err error
 
 	if len(args) == 0 {
-		data = mc.buf.takeBuffer(minPktLen)
+		data, err = mc.buf.takeBuffer(minPktLen)
 	} else {
-		data = mc.buf.takeCompleteBuffer()
+		data, err = mc.buf.takeCompleteBuffer()
 	}
-	if data == nil {
+	if err != nil {
 		// cannot take the buffer. Something must be wrong with the connection
-		errLog.Print(ErrBusyBuffer)
+		errLog.Print(err)
 		return errBadConnNoWrite
 	}
 
@@ -939,23 +940,25 @@ func (stmt *mysqlStmt) writeExecutePacket(args []driver.Value) error {
 		pos := minPktLen
 
 		var nullMask []byte
-		if maskLen, typesLen := (len(args)+7)/8, 1+2*len(args); pos+maskLen+typesLen >= len(data) {
+		maskLen, typesLen := (len(args)+7)/8, 1+2*len(args)
+		l := pos + maskLen + typesLen
+		if l > cap(data) {
 			// buffer has to be extended but we don't know by how much so
 			// we depend on append after all data with known sizes fit.
 			// We stop at that because we deal with a lot of columns here
 			// which makes the required allocation size hard to guess.
 			tmp := make([]byte, pos+maskLen+typesLen)
 			copy(tmp[:pos], data[:pos])
 			data = tmp
-			nullMask = data[pos : pos+maskLen]
-			pos += maskLen
-		} else {
-			nullMask = data[pos : pos+maskLen]
-			for i := 0; i < maskLen; i++ {
-				nullMask[i] = 0
-			}
-			pos += maskLen
+		} else if l > len(data) {
+			data = data[:l]
+		}
+
+		nullMask = data[pos : pos+maskLen]
+		for i := range nullMask {
+			nullMask[i] = 0
 		}
+		pos += maskLen
 
 		// newParameterBoundFlag 1 [1 byte]
 		data[pos] = 0x01