Support size-prefixed buffers and add tests for size-prefixed messages (#6232)

go/builder.go

@@ -22,6 +22,7 @@ type Builder struct {
 }
 
 const fileIdentifierLength = 4
+const sizePrefixLength = 4
 
 // NewBuilder initializes a Builder of size `initial_size`.
 // The internal buffer is grown as needed.
@@ -580,11 +581,53 @@ func (b *Builder) FinishWithFileIdentifier(rootTable UOffsetT, fid []byte) {
 	b.Finish(rootTable)
 }
 
+// FinishSizePrefixed finalizes a buffer, pointing to the given `rootTable`.
+// The buffer is prefixed with the size of the buffer, excluding the size
+// of the prefix itself.
+func (b *Builder) FinishSizePrefixed(rootTable UOffsetT) {
+	b.finish(rootTable, true)
+}
+
+// FinishSizePrefixedWithFileIdentifier finalizes a buffer, pointing to the given `rootTable`
+// and applies a file identifier. The buffer is prefixed with the size of the buffer,
+// excluding the size of the prefix itself.
+func (b *Builder) FinishSizePrefixedWithFileIdentifier(rootTable UOffsetT, fid []byte) {
+	if fid == nil || len(fid) != fileIdentifierLength {
+		panic("incorrect file identifier length")
+	}
+	// In order to add a file identifier and size prefix to the flatbuffer message,
+	// we need to prepare an alignment, a size prefix length, and file identifier length
+	b.Prep(b.minalign, SizeInt32+fileIdentifierLength+sizePrefixLength)
+	for i := fileIdentifierLength - 1; i >= 0; i-- {
+		// place the file identifier
+		b.PlaceByte(fid[i])
+	}
+	// finish
+	b.finish(rootTable, true)
+}
+
 // Finish finalizes a buffer, pointing to the given `rootTable`.
 func (b *Builder) Finish(rootTable UOffsetT) {
+	b.finish(rootTable, false)
+}
+
+// finish finalizes a buffer, pointing to the given `rootTable`
+// with an optional size prefix.
+func (b *Builder) finish(rootTable UOffsetT, sizePrefix bool) {
 	b.assertNotNested()
-	b.Prep(b.minalign, SizeUOffsetT)
+
+	if sizePrefix {
+		b.Prep(b.minalign, SizeUOffsetT+sizePrefixLength)
+	} else {
+		b.Prep(b.minalign, SizeUOffsetT)
+	}
+
 	b.PrependUOffsetT(rootTable)
+
+	if sizePrefix {
+		b.PlaceUint32(uint32(b.Offset()))
+	}
+
 	b.finished = true
 }
 

go/lib.go

@@ -11,3 +11,15 @@ func GetRootAs(buf []byte, offset UOffsetT, fb FlatBuffer) {
 	n := GetUOffsetT(buf[offset:])
 	fb.Init(buf, n+offset)
 }
+
+// GetSizePrefixedRootAs is a generic helper to initialize a FlatBuffer with the provided size-prefixed buffer
+// bytes and its data offset
+func GetSizePrefixedRootAs(buf []byte, offset UOffsetT, fb FlatBuffer) {
+	n := GetUOffsetT(buf[offset+sizePrefixLength:])
+	fb.Init(buf, n+offset+sizePrefixLength)
+}
+
+// GetSizePrefix reads the size from a size-prefixed flatbuffer
+func GetSizePrefix(buf []byte, offset UOffsetT) uint32 {
+	return GetUint32(buf[offset:])
+}