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:])
+}

src/idl_gen_go.cpp

@@ -255,17 +255,28 @@ class GoGenerator : public BaseGenerator {
   void NewRootTypeFromBuffer(const StructDef &struct_def,
                              std::string *code_ptr) {
     std::string &code = *code_ptr;
+    std::string size_prefix[] = { "", "SizePrefixed" };
 
-    code += "func GetRootAs";
+    for (int i = 0; i < 2; i++) {
-    code += struct_def.name;
+      code += "func Get" + size_prefix[i] + "RootAs";
-    code += "(buf []byte, offset flatbuffers.UOffsetT) ";
+      code += struct_def.name;
-    code += "*" + struct_def.name + "";
+      code += "(buf []byte, offset flatbuffers.UOffsetT) ";
-    code += " {\n";
+      code += "*" + struct_def.name + "";
-    code += "\tn := flatbuffers.GetUOffsetT(buf[offset:])\n";
+      code += " {\n";
-    code += "\tx := &" + struct_def.name + "{}\n";
+      if (i == 0) {
-    code += "\tx.Init(buf, n+offset)\n";
+        code += "\tn := flatbuffers.GetUOffsetT(buf[offset:])\n";
-    code += "\treturn x\n";
+      } else {
-    code += "}\n\n";
+        code += "\tn := flatbuffers.GetUOffsetT(buf[offset+flatbuffers.SizeUint32:])\n";
+      }
+      code += "\tx := &" + struct_def.name + "{}\n";
+      if (i == 0) {
+        code += "\tx.Init(buf, n+offset)\n";
+      } else {
+        code += "\tx.Init(buf, n+offset+flatbuffers.SizeUint32)\n";
+      }
+      code += "\treturn x\n";
+      code += "}\n\n";
+    }
   }
 
   // Initialize an existing object with other data, to avoid an allocation.

tests/MyGame/Example/Monster.go

@@ -442,6 +442,13 @@ func GetRootAsMonster(buf []byte, offset flatbuffers.UOffsetT) *Monster {
 	return x
 }
 
+func GetSizePrefixedRootAsMonster(buf []byte, offset flatbuffers.UOffsetT) *Monster {
+	n := flatbuffers.GetUOffsetT(buf[offset+flatbuffers.SizeUint32:])
+	x := &Monster{}
+	x.Init(buf, n+offset+flatbuffers.SizeUint32)
+	return x
+}
+
 func (rcv *Monster) Init(buf []byte, i flatbuffers.UOffsetT) {
 	rcv._tab.Bytes = buf
 	rcv._tab.Pos = i

tests/MyGame/Example/Referrable.go

@@ -39,6 +39,13 @@ func GetRootAsReferrable(buf []byte, offset flatbuffers.UOffsetT) *Referrable {
 	return x
 }
 
+func GetSizePrefixedRootAsReferrable(buf []byte, offset flatbuffers.UOffsetT) *Referrable {
+	n := flatbuffers.GetUOffsetT(buf[offset+flatbuffers.SizeUint32:])
+	x := &Referrable{}
+	x.Init(buf, n+offset+flatbuffers.SizeUint32)
+	return x
+}
+
 func (rcv *Referrable) Init(buf []byte, i flatbuffers.UOffsetT) {
 	rcv._tab.Bytes = buf
 	rcv._tab.Pos = i

tests/MyGame/Example/Stat.go

@@ -46,6 +46,13 @@ func GetRootAsStat(buf []byte, offset flatbuffers.UOffsetT) *Stat {
 	return x
 }
 
+func GetSizePrefixedRootAsStat(buf []byte, offset flatbuffers.UOffsetT) *Stat {
+	n := flatbuffers.GetUOffsetT(buf[offset+flatbuffers.SizeUint32:])
+	x := &Stat{}
+	x.Init(buf, n+offset+flatbuffers.SizeUint32)
+	return x
+}
+
 func (rcv *Stat) Init(buf []byte, i flatbuffers.UOffsetT) {
 	rcv._tab.Bytes = buf
 	rcv._tab.Pos = i

tests/MyGame/Example/TestSimpleTableWithEnum.go

@@ -39,6 +39,13 @@ func GetRootAsTestSimpleTableWithEnum(buf []byte, offset flatbuffers.UOffsetT) *
 	return x
 }
 
+func GetSizePrefixedRootAsTestSimpleTableWithEnum(buf []byte, offset flatbuffers.UOffsetT) *TestSimpleTableWithEnum {
+	n := flatbuffers.GetUOffsetT(buf[offset+flatbuffers.SizeUint32:])
+	x := &TestSimpleTableWithEnum{}
+	x.Init(buf, n+offset+flatbuffers.SizeUint32)
+	return x
+}
+
 func (rcv *TestSimpleTableWithEnum) Init(buf []byte, i flatbuffers.UOffsetT) {
 	rcv._tab.Bytes = buf
 	rcv._tab.Pos = i

tests/MyGame/Example/TypeAliases.go

@@ -98,6 +98,13 @@ func GetRootAsTypeAliases(buf []byte, offset flatbuffers.UOffsetT) *TypeAliases
 	return x
 }
 
+func GetSizePrefixedRootAsTypeAliases(buf []byte, offset flatbuffers.UOffsetT) *TypeAliases {
+	n := flatbuffers.GetUOffsetT(buf[offset+flatbuffers.SizeUint32:])
+	x := &TypeAliases{}
+	x.Init(buf, n+offset+flatbuffers.SizeUint32)
+	return x
+}
+
 func (rcv *TypeAliases) Init(buf []byte, i flatbuffers.UOffsetT) {
 	rcv._tab.Bytes = buf
 	rcv._tab.Pos = i

tests/MyGame/Example2/Monster.go

@@ -36,6 +36,13 @@ func GetRootAsMonster(buf []byte, offset flatbuffers.UOffsetT) *Monster {
 	return x
 }
 
+func GetSizePrefixedRootAsMonster(buf []byte, offset flatbuffers.UOffsetT) *Monster {
+	n := flatbuffers.GetUOffsetT(buf[offset+flatbuffers.SizeUint32:])
+	x := &Monster{}
+	x.Init(buf, n+offset+flatbuffers.SizeUint32)
+	return x
+}
+
 func (rcv *Monster) Init(buf []byte, i flatbuffers.UOffsetT) {
 	rcv._tab.Bytes = buf
 	rcv._tab.Pos = i

tests/MyGame/InParentNamespace.go

@@ -36,6 +36,13 @@ func GetRootAsInParentNamespace(buf []byte, offset flatbuffers.UOffsetT) *InPare
 	return x
 }
 
+func GetSizePrefixedRootAsInParentNamespace(buf []byte, offset flatbuffers.UOffsetT) *InParentNamespace {
+	n := flatbuffers.GetUOffsetT(buf[offset+flatbuffers.SizeUint32:])
+	x := &InParentNamespace{}
+	x.Init(buf, n+offset+flatbuffers.SizeUint32)
+	return x
+}
+
 func (rcv *InParentNamespace) Init(buf []byte, i flatbuffers.UOffsetT) {
 	rcv._tab.Bytes = buf
 	rcv._tab.Pos = i

tests/go_test.go

@@ -87,13 +87,13 @@ func TestAll(t *testing.T) {
 
 	// Verify that using the generated Go code builds a buffer without
 	// returning errors:
-	generated, off := CheckGeneratedBuild(t.Fatalf)
+	generated, off := CheckGeneratedBuild(false, t.Fatalf)
 
 	// Verify that the buffer generated by Go code is readable by the
 	// generated Go code:
-	CheckReadBuffer(generated, off, t.Fatalf)
+	CheckReadBuffer(generated, off, false, t.Fatalf)
-	CheckMutateBuffer(generated, off, t.Fatalf)
+	CheckMutateBuffer(generated, off, false, t.Fatalf)
-	CheckObjectAPI(generated, off, t.Fatalf)
+	CheckObjectAPI(generated, off, false, t.Fatalf)
 
 	// Verify that the buffer generated by C++ code is readable by the
 	// generated Go code:
@@ -101,9 +101,9 @@ func TestAll(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	CheckReadBuffer(monsterDataCpp, 0, t.Fatalf)
+	CheckReadBuffer(monsterDataCpp, 0, false, t.Fatalf)
-	CheckMutateBuffer(monsterDataCpp, 0, t.Fatalf)
+	CheckMutateBuffer(monsterDataCpp, 0, false, t.Fatalf)
-	CheckObjectAPI(monsterDataCpp, 0, t.Fatalf)
+	CheckObjectAPI(monsterDataCpp, 0, false, t.Fatalf)
 
 	// Verify that vtables are deduplicated when written:
 	CheckVtableDeduplication(t.Fatalf)
@@ -127,6 +127,9 @@ func TestAll(t *testing.T) {
 	// Check a parent namespace import
 	CheckParentNamespace(t.Fatalf)
 
+	// Check size-prefixed flatbuffers
+	CheckSizePrefixedBuffer(t.Fatalf)
+
 	// If the filename of the FlatBuffers file generated by the Java test
 	// is given, check that Go code can read it, and that Go code
 	// generates an identical buffer when used to create the example data:
@@ -135,7 +138,7 @@ func TestAll(t *testing.T) {
 		if err != nil {
 			t.Fatal(err)
 		}
-		CheckReadBuffer(monsterDataJava, 0, t.Fatalf)
+		CheckReadBuffer(monsterDataJava, 0, false, t.Fatalf)
 		CheckByteEquality(generated[off:], monsterDataJava, t.Fatalf)
 	}
 
@@ -153,11 +156,19 @@ func TestAll(t *testing.T) {
 
 // CheckReadBuffer checks that the given buffer is evaluated correctly
 // as the example Monster.
-func CheckReadBuffer(buf []byte, offset flatbuffers.UOffsetT, fail func(string, ...interface{})) {
+func CheckReadBuffer(buf []byte, offset flatbuffers.UOffsetT, sizePrefix bool, fail func(string, ...interface{})) {
 	// try the two ways of generating a monster
-	monster1 := example.GetRootAsMonster(buf, offset)
+	var monster1 *example.Monster
 	monster2 := &example.Monster{}
-	flatbuffers.GetRootAs(buf, offset, monster2)
+
+	if sizePrefix {
+		monster1 = example.GetSizePrefixedRootAsMonster(buf, offset)
+		flatbuffers.GetSizePrefixedRootAs(buf, offset, monster2)
+	} else {
+		monster1 = example.GetRootAsMonster(buf, offset)
+		flatbuffers.GetRootAs(buf, offset, monster2)
+	}
+
 	for _, monster := range []*example.Monster{monster1, monster2} {
 		if got := monster.Hp(); 80 != got {
 			fail(FailString("hp", 80, got))
@@ -320,13 +331,18 @@ func CheckReadBuffer(buf []byte, offset flatbuffers.UOffsetT, fail func(string,
 
 // CheckMutateBuffer checks that the given buffer can be mutated correctly
 // as the example Monster. Only available scalar values are mutated.
-func CheckMutateBuffer(org []byte, offset flatbuffers.UOffsetT, fail func(string, ...interface{})) {
+func CheckMutateBuffer(org []byte, offset flatbuffers.UOffsetT, sizePrefix bool, fail func(string, ...interface{})) {
 	// make a copy to mutate
 	buf := make([]byte, len(org))
 	copy(buf, org)
 
 	// load monster data from the buffer
-	monster := example.GetRootAsMonster(buf, offset)
+	var monster *example.Monster
+	if sizePrefix {
+		monster = example.GetSizePrefixedRootAsMonster(buf, offset)
+	} else {
+		monster = example.GetRootAsMonster(buf, offset)
+	}
 
 	// test case struct
 	type testcase struct {
@@ -409,7 +425,12 @@ func CheckMutateBuffer(org []byte, offset flatbuffers.UOffsetT, fail func(string
 
 	// To make sure the buffer has changed accordingly
 	// Read data from the buffer and verify all fields
-	monster = example.GetRootAsMonster(buf, offset)
+	if sizePrefix {
+		monster = example.GetSizePrefixedRootAsMonster(buf, offset)
+	} else {
+		monster = example.GetRootAsMonster(buf, offset)
+	}
+
 	for _, t := range testForMutatedValues {
 		if !t.testfn() {
 			fail("field '" + t.field + "' doesn't have the expected mutated value")
@@ -429,7 +450,12 @@ func CheckMutateBuffer(org []byte, offset flatbuffers.UOffsetT, fail func(string
 	// back to their original values and compare buffers.
 	// This test is done to make sure mutations do not do
 	// any unnecessary changes to the buffer.
-	monster = example.GetRootAsMonster(buf, offset)
+	if sizePrefix {
+		monster = example.GetSizePrefixedRootAsMonster(buf, offset)
+	} else {
+		monster = example.GetRootAsMonster(buf, offset)
+	}
+
 	monster.MutateHp(80)
 	monster.MutateTestbool(true)
 	monster.Pos(nil).MutateX(1.0)
@@ -453,8 +479,14 @@ func CheckMutateBuffer(org []byte, offset flatbuffers.UOffsetT, fail func(string
 	}
 }
 
-func CheckObjectAPI(buf []byte, offset flatbuffers.UOffsetT, fail func(string, ...interface{})) {
+func CheckObjectAPI(buf []byte, offset flatbuffers.UOffsetT, sizePrefix bool, fail func(string, ...interface{})) {
-	monster := example.GetRootAsMonster(buf, offset).UnPack()
+	var monster *example.MonsterT
+
+	if sizePrefix {
+		monster = example.GetSizePrefixedRootAsMonster(buf, offset).UnPack()
+	} else {
+		monster = example.GetRootAsMonster(buf, offset).UnPack()
+	}
 
 	if got := monster.Hp; 80 != got {
 		fail(FailString("hp", 80, got))
@@ -1183,7 +1215,7 @@ func CheckGetRootAsForNonRootTable(fail func(string, ...interface{})) {
 }
 
 // CheckGeneratedBuild uses generated code to build the example Monster.
-func CheckGeneratedBuild(fail func(string, ...interface{})) ([]byte, flatbuffers.UOffsetT) {
+func CheckGeneratedBuild(sizePrefix bool, fail func(string, ...interface{})) ([]byte, flatbuffers.UOffsetT) {
 	b := flatbuffers.NewBuilder(0)
 	str := b.CreateString("MyMonster")
 	test1 := b.CreateString("test1")
@@ -1227,7 +1259,11 @@ func CheckGeneratedBuild(fail func(string, ...interface{})) ([]byte, flatbuffers
 	example.MonsterAddTestarrayofstring(b, testArrayOfString)
 	mon := example.MonsterEnd(b)
 
-	b.Finish(mon)
+	if sizePrefix {
+		b.FinishSizePrefixed(mon)
+	} else {
+		b.Finish(mon)
+	}
 
 	return b.Bytes, b.Head()
 }
@@ -1626,6 +1662,27 @@ func CheckParentNamespace(fail func(string, ...interface{})) {
 	}
 }
 
+func CheckSizePrefixedBuffer(fail func(string, ...interface{})) {
+	// Generate a size-prefixed flatbuffer
+	generated, off := CheckGeneratedBuild(true, fail)
+
+	// Check that the size prefix is the size of monsterdata_go_wire.mon minus 4
+	size := flatbuffers.GetSizePrefix(generated, off)
+	if size != 220 {
+		fail("mismatch between size prefix and expected size")
+	}
+
+	// Check that the buffer can be used as expected
+	CheckReadBuffer(generated, off, true, fail)
+	CheckMutateBuffer(generated, off, true, fail)
+	CheckObjectAPI(generated, off, true, fail)
+
+	// Write generated bfufer out to a file
+	if err := ioutil.WriteFile(outData+".sp", generated[off:], os.FileMode(0644)); err != nil {
+		fail("failed to write file: %s", err)
+	}
+}
+
 // Include simple random number generator to ensure results will be the
 // same cross platform.
 // http://en.wikipedia.org/wiki/Park%E2%80%93Miller_random_number_generator
@@ -1835,7 +1892,7 @@ func BenchmarkVtableDeduplication(b *testing.B) {
 // BenchmarkParseGold measures the speed of parsing the 'gold' data
 // used throughout this test suite.
 func BenchmarkParseGold(b *testing.B) {
-	buf, offset := CheckGeneratedBuild(b.Fatalf)
+	buf, offset := CheckGeneratedBuild(false, b.Fatalf)
 	monster := example.GetRootAsMonster(buf, offset)
 
 	// use these to prevent allocations:
@@ -1897,7 +1954,7 @@ func BenchmarkParseGold(b *testing.B) {
 
 // BenchmarkBuildGold uses generated code to build the example Monster.
 func BenchmarkBuildGold(b *testing.B) {
-	buf, offset := CheckGeneratedBuild(b.Fatalf)
+	buf, offset := CheckGeneratedBuild(false, b.Fatalf)
 	bytes_length := int64(len(buf[offset:]))
 
 	reuse_str := "MyMonster"

tests/namespace_test/NamespaceA/NamespaceB/TableInNestedNS.go

@@ -39,6 +39,13 @@ func GetRootAsTableInNestedNS(buf []byte, offset flatbuffers.UOffsetT) *TableInN
 	return x
 }
 
+func GetSizePrefixedRootAsTableInNestedNS(buf []byte, offset flatbuffers.UOffsetT) *TableInNestedNS {
+	n := flatbuffers.GetUOffsetT(buf[offset+flatbuffers.SizeUint32:])
+	x := &TableInNestedNS{}
+	x.Init(buf, n+offset+flatbuffers.SizeUint32)
+	return x
+}
+
 func (rcv *TableInNestedNS) Init(buf []byte, i flatbuffers.UOffsetT) {
 	rcv._tab.Bytes = buf
 	rcv._tab.Pos = i

tests/namespace_test/NamespaceA/SecondTableInA.go

@@ -42,6 +42,13 @@ func GetRootAsSecondTableInA(buf []byte, offset flatbuffers.UOffsetT) *SecondTab
 	return x
 }
 
+func GetSizePrefixedRootAsSecondTableInA(buf []byte, offset flatbuffers.UOffsetT) *SecondTableInA {
+	n := flatbuffers.GetUOffsetT(buf[offset+flatbuffers.SizeUint32:])
+	x := &SecondTableInA{}
+	x.Init(buf, n+offset+flatbuffers.SizeUint32)
+	return x
+}
+
 func (rcv *SecondTableInA) Init(buf []byte, i flatbuffers.UOffsetT) {
 	rcv._tab.Bytes = buf
 	rcv._tab.Pos = i

tests/namespace_test/NamespaceA/TableInFirstNS.go

@@ -49,6 +49,13 @@ func GetRootAsTableInFirstNS(buf []byte, offset flatbuffers.UOffsetT) *TableInFi
 	return x
 }
 
+func GetSizePrefixedRootAsTableInFirstNS(buf []byte, offset flatbuffers.UOffsetT) *TableInFirstNS {
+	n := flatbuffers.GetUOffsetT(buf[offset+flatbuffers.SizeUint32:])
+	x := &TableInFirstNS{}
+	x.Init(buf, n+offset+flatbuffers.SizeUint32)
+	return x
+}
+
 func (rcv *TableInFirstNS) Init(buf []byte, i flatbuffers.UOffsetT) {
 	rcv._tab.Bytes = buf
 	rcv._tab.Pos = i

tests/namespace_test/NamespaceC/TableInC.go

@@ -46,6 +46,13 @@ func GetRootAsTableInC(buf []byte, offset flatbuffers.UOffsetT) *TableInC {
 	return x
 }
 
+func GetSizePrefixedRootAsTableInC(buf []byte, offset flatbuffers.UOffsetT) *TableInC {
+	n := flatbuffers.GetUOffsetT(buf[offset+flatbuffers.SizeUint32:])
+	x := &TableInC{}
+	x.Init(buf, n+offset+flatbuffers.SizeUint32)
+	return x
+}
+
 func (rcv *TableInC) Init(buf []byte, i flatbuffers.UOffsetT) {
 	rcv._tab.Bytes = buf
 	rcv._tab.Pos = i