GODRIVER-3809 Fix *streamingByteSrc.readSlice(). (#2326)

Author: qingyang-hu

bson/buffered_byte_src.go

@@ -68,7 +68,7 @@ func (b *bufferedByteSrc) discard(n int) (int, error) {
 	return n, nil
 }
 
-// readSlice scans buf[offset:] for the first occurrence of delim, returns
+// readSlice reads buf[offset:] for the first occurrence of delim, returning
 // buf[offset:idx+1], and advances offset past it; errors if delim not found.
 func (b *bufferedByteSrc) readSlice(delim byte) ([]byte, error) {
 	// Ensure we don't read past the end of the buffer.

bson/streaming_byte_src.go

@@ -56,15 +56,43 @@ func (s *streamingByteSrc) discard(n int) (int, error) {
 	return m, err
 }
 
-// readSlice scans buf[offset:] for the first occurrence of delim, returns
-// buf[offset:idx+1], and advances offset past it; errors if delim not found.
+// readSlice reads until the first occurrence of delim, returning a slice
+// containing the data up to and including the delimiter, and advances offset
+// past it; errors if delim not found.
 func (s *streamingByteSrc) readSlice(delim byte) ([]byte, error) {
-	data, err := s.br.ReadSlice(delim)
-	if err != nil {
-		return nil, err
+	var full [][]byte
+	var frag []byte
+	var err error
+	var n int
+
+	for {
+		if l := len(frag); l > 0 {
+			// Make a copy of the fragment to accumulate full buffers.
+			buf := make([]byte, l)
+			copy(buf, frag)
+			full = append(full, buf)
+		}
+		frag, err = s.br.ReadSlice(delim)
+		n += len(frag)
+		if err != bufio.ErrBufferFull {
+			break
+		}
+	}
+	s.offset += int64(n)
+
+	// If ReadSlice is only called once, we can return the fragment directly.
+	if len(full) == 0 {
+		return frag, err
+	}
+
+	// Allocate new buffer to hold the full buffers and the fragment.
+	buf := make([]byte, n)
+	n = 0
+	for i := range full {
+		n += copy(buf[n:], full[i])
 	}
-	s.offset += int64(len(data))
-	return data, nil
+	copy(buf[n:], frag)
+	return buf, err
 }
 
 // pos returns the current read position in the buffer.

bson/unmarshaling_cases_test.go

@@ -7,7 +7,9 @@
 package bson
 
 import (
+	"fmt"
 	"reflect"
+	"strings"
 )
 
 type unmarshalingTestCase struct {
@@ -58,6 +60,15 @@ func unmarshalingTestCases() []unmarshalingTestCase {
 		Foo []byte
 	}
 
+	longKey := strings.Repeat("k", 16_000_000)
+	tLongKey := reflect.StructOf([]reflect.StructField{
+		{
+			Name: "Foo",
+			Type: reflect.TypeOf(false),
+			Tag:  reflect.StructTag(fmt.Sprintf(`bson:"%s"`, longKey)),
+		},
+	})
+
 	return []unmarshalingTestCase{
 		{
 			name: "small struct",
@@ -235,6 +246,16 @@ func unmarshalingTestCases() []unmarshalingTestCase {
 				{Key: "b_ptr_tracker", Value: MaxKey{}},
 			}),
 		},
+		{
+			name:  "long key",
+			sType: tLongKey,
+			want: func() any {
+				vLongKey := reflect.New(tLongKey)
+				vLongKey.Elem().Field(0).SetBool(true)
+				return vLongKey.Interface()
+			}(),
+			data: docToBytes(D{{longKey, true}}),
+		},
 	}
 }