Error while expanding buffer in dechunker

Serious issue that could cause corrupt data in records.
Fixes #183

Author: Peter Wilhelmsson

neo4j/internal/bolt/bolt3server_test.go

@@ -91,7 +91,7 @@ func (s *bolt3server) waitForHello() {
 }
 
 func (s *bolt3server) receiveMsg() *testStruct {
-	buf, err := dechunkMessage(s.conn, []byte{})
+	_, buf, err := dechunkMessage(s.conn, []byte{})
 	if err != nil {
 		panic(err)
 	}

neo4j/internal/bolt/bolt4server_test.go

@@ -94,7 +94,7 @@ func (s *bolt4server) waitForHello() map[string]interface{} {
 }
 
 func (s *bolt4server) receiveMsg() *testStruct {
-	buf, err := dechunkMessage(s.conn, []byte{})
+	_, buf, err := dechunkMessage(s.conn, []byte{})
 	if err != nil {
 		panic(err)
 	}

neo4j/internal/bolt/chunker_test.go

@@ -81,7 +81,7 @@ func TestChunker(ot *testing.T) {
 
 	receiveAndAssertMessage := func(t *testing.T, rd io.Reader, expected []byte) {
 		t.Helper()
-		msg, err := dechunkMessage(rd, []byte{})
+		_, msg, err := dechunkMessage(rd, []byte{})
 		AssertNoError(t, err)
 		assertSlices(t, msg, expected)
 	}

neo4j/internal/bolt/dechunker.go

@@ -24,32 +24,37 @@ import (
 	"io"
 )
 
-func dechunkMessage(rd io.Reader, msgBuf []byte) ([]byte, error) {
+// dechunkMessage takes a buffer to be reused and returns the reusable buffer
+// (might have been reallocated to handle growth), the message buffer and
+// error.
+func dechunkMessage(rd io.Reader, msgBuf []byte) ([]byte, []byte, error) {
 	sizeBuf := []byte{0x00, 0x00}
 	off := 0
 
 	for {
 		_, err := io.ReadFull(rd, sizeBuf)
 		if err != nil {
-			return msgBuf, err
+			return msgBuf, nil, err
 		}
 		chunkSize := int(binary.BigEndian.Uint16(sizeBuf))
 		if chunkSize == 0 {
 			if off > 0 {
-				return msgBuf[:off], nil
+				return msgBuf, msgBuf[:off], nil
 			}
 			// Got a nop chunk
 			continue
 		}
 
 		// Need to expand buffer
 		if (off + chunkSize) > cap(msgBuf) {
-			msgBuf = append(msgBuf, make([]byte, off+chunkSize)...)
+			newMsgBuf := make([]byte, (off+chunkSize)+4096)
+			copy(newMsgBuf, msgBuf)
+			msgBuf = newMsgBuf
 		}
 		// Read the chunk into buffer
 		_, err = io.ReadFull(rd, msgBuf[off:(off+chunkSize)])
 		if err != nil {
-			return msgBuf, err
+			return msgBuf, nil, err
 		}
 		off += chunkSize
 	}

neo4j/internal/bolt/dechunker_test.go

@@ -39,13 +39,16 @@ func TestDechunker(t *testing.T) {
 		{size: 1021, max: 0x7},
 		{size: 0xffff78, max: 0x30},
 		{size: 3, max: 0xffff},
+		{size: 1021, max: 0xff90},
+		{size: 0xffff78, max: 0xff90},
 	}
-	for _, msg := range messages {
+	for msgi, msg := range messages {
 		// Prepare message
 		str := &bytes.Buffer{}
 		total := msg.size
 		this := uint16(0)
 
+		b := byte(0)
 		for total > 0 {
 			// Write size
 			if total > uint32(msg.max) {
@@ -60,15 +63,30 @@ func TestDechunker(t *testing.T) {
 
 			// Write data
 			buf = make([]byte, int(this))
+			for i := range buf {
+				buf[i] = b
+				b++
+			}
 			str.Write(buf)
 		}
 		// Write end of mesage marker
 		str.Write([]byte{0x00, 0x00})
 
 		// Dechunk the message
-		buf, err = dechunkMessage(str, buf)
+		var msgBuf []byte
+		buf, msgBuf, err = dechunkMessage(str, buf)
 		AssertNoError(t, err)
-		AssertLen(t, buf, int(msg.size))
+		AssertLen(t, msgBuf, int(msg.size))
+		// Check content of buffer
+		b = 0
+		for i := range msgBuf {
+			if msgBuf[i] != b {
+				t.Errorf("Wrong content in buffer at %d, %d vs %d (message %d)", i, msgBuf[i], b, msgi)
+				return
+			}
+			b++
+		}
+
 		// Check that buffer increases or stays put
 		if cap(buf) < prevCap {
 			t.Errorf("Underlying buffer should be reused")

neo4j/internal/bolt/hydratedehydrate_test.go

@@ -47,7 +47,7 @@ func TestDehydrateHydrate(ot *testing.T) {
 		buf := &bytes.Buffer{}
 		out.send(buf)
 
-		byts, err := dechunkMessage(buf, []byte{})
+		_, byts, err := dechunkMessage(buf, []byte{})
 		if err != nil {
 			t.Fatal(err)
 		}

neo4j/internal/bolt/incoming.go

@@ -23,18 +23,17 @@ import (
 )
 
 type incoming struct {
-	buf []byte
+	buf []byte // Reused buffer
 	hyd hydrator
 }
 
 func (i *incoming) next(rd io.Reader) (interface{}, error) {
 	// Get next message from transport layer
-	buf, err := dechunkMessage(rd, i.buf)
+	var err error
+	var msg []byte
+	i.buf, msg, err = dechunkMessage(rd, i.buf)
 	if err != nil {
 		return nil, err
 	}
-	// Reuse buffer for next dechunk
-	i.buf = buf
-
-	return i.hyd.hydrate(buf)
+	return i.hyd.hydrate(msg)
 }

neo4j/internal/bolt/outgoing_test.go

@@ -99,7 +99,7 @@ func TestOutgoing(ot *testing.T) {
 		out.send(buf)
 
 		// Dechunk it
-		byts, err := dechunkMessage(buf, []byte{})
+		_, byts, err := dechunkMessage(buf, []byte{})
 		if err != nil {
 			t.Fatal(err)
 		}