fix: a race when reusing flatbuffers.Builder (#35)
diff --git a/internal/server/server.go b/internal/server/server.go
index 890997f..720df4d 100644
--- a/internal/server/server.go
+++ b/internal/server/server.go
@@ -45,18 +45,14 @@
dealRPCTest func(buf []byte) (*flatbuffers.Builder, error)
)
-func generateErrorReport(err error) (ty byte, out []byte) {
+func generateErrorReport(err error) *flatbuffers.Builder {
if err == ttlcache.ErrNotFound {
log.Warnf("%s", err)
} else {
log.Errorf("%s", err)
}
- ty = util.RPCError
- bd := ReportError(err)
- out = bd.FinishedBytes()
- util.PutBuilder(bd)
- return
+ return ReportError(err)
}
func recoverPanic() {
@@ -65,7 +61,7 @@
}
}
-func dispatchRPC(ty byte, in []byte, conn net.Conn) (byte, []byte) {
+func dispatchRPC(ty byte, in []byte, conn net.Conn) (*flatbuffers.Builder, error) {
var err error
var bd *flatbuffers.Builder
switch ty {
@@ -78,21 +74,18 @@
default:
err = UnknownType{ty}
}
+ return bd, err
+}
- var out []byte
- if err != nil {
- ty, out = generateErrorReport(err)
- } else {
- out = bd.FinishedBytes()
+func checkIfDataTooLarge(bd *flatbuffers.Builder) *flatbuffers.Builder {
+ out := bd.FinishedBytes()
+ size := len(out)
+ if size > util.MaxDataSize {
+ err := fmt.Errorf("the max length of data is %d but got %d", util.MaxDataSize, size)
util.PutBuilder(bd)
- size := len(out)
- if size > util.MaxDataSize {
- err = fmt.Errorf("the max length of data is %d but got %d", util.MaxDataSize, size)
- ty, out = generateErrorReport(err)
- }
+ bd = generateErrorReport(err)
}
-
- return ty, out
+ return bd
}
func handleConn(c net.Conn) {
@@ -122,8 +115,16 @@
break
}
- ty, out := dispatchRPC(ty, buf, c)
+ bd, err := dispatchRPC(ty, buf, c)
+ if err != nil {
+ util.PutBuilder(bd)
+ bd = generateErrorReport(err)
+ } else {
+ bd = checkIfDataTooLarge(bd)
+ }
+
+ out := bd.FinishedBytes()
size := len(out)
binary.BigEndian.PutUint32(header, uint32(size))
header[0] = ty
@@ -139,6 +140,8 @@
util.WriteErr(n, err)
break
}
+
+ util.PutBuilder(bd)
}
}
diff --git a/internal/server/server_test.go b/internal/server/server_test.go
index 74bf35c..606c2cb 100644
--- a/internal/server/server_test.go
+++ b/internal/server/server_test.go
@@ -27,7 +27,6 @@
"time"
hrc "github.com/api7/ext-plugin-proto/go/A6/HTTPReqCall"
- flatbuffers "github.com/google/flatbuffers/go"
"github.com/stretchr/testify/assert"
"github.com/apache/apisix-go-plugin-runner/internal/util"
@@ -53,28 +52,23 @@
}
func TestDispatchRPC_UnknownType(t *testing.T) {
- ty, _ := dispatchRPC(126, []byte(""), nil)
- assert.Equal(t, byte(util.RPCError), ty)
+ _, err := dispatchRPC(126, []byte(""), nil)
+ assert.Equal(t, UnknownType{126}, err)
}
func TestDispatchRPC_OutTooLarge(t *testing.T) {
- dealRPCTest = func(buf []byte) (*flatbuffers.Builder, error) {
- builder := util.GetBuilder()
- bodyVec := builder.CreateByteVector(make([]byte, util.MaxDataSize+1))
- hrc.StopStart(builder)
- hrc.StopAddBody(builder, bodyVec)
- stop := hrc.StopEnd(builder)
+ builder := util.GetBuilder()
+ bodyVec := builder.CreateByteVector(make([]byte, util.MaxDataSize+1))
+ hrc.StopStart(builder)
+ hrc.StopAddBody(builder, bodyVec)
+ stop := hrc.StopEnd(builder)
- hrc.RespStart(builder)
- hrc.RespAddId(builder, 1)
- hrc.RespAddActionType(builder, hrc.ActionStop)
- hrc.RespAddAction(builder, stop)
- res := hrc.RespEnd(builder)
- builder.Finish(res)
- return builder, nil
- }
- ty, _ := dispatchRPC(util.RPCTest, []byte(""), nil)
- assert.Equal(t, byte(util.RPCError), ty)
+ hrc.RespStart(builder)
+ hrc.RespAddId(builder, 1)
+ hrc.RespAddActionType(builder, hrc.ActionStop)
+ hrc.RespAddAction(builder, stop)
+ res := hrc.RespEnd(builder)
+ builder.Finish(res)
}
func TestRun(t *testing.T) {
