[apisix-go-plugin-runner] branch master updated: fix: set the correct rpc response type when meet an error (#102)

[so...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

soulbird pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/apisix-go-plugin-runner.git


The following commit(s) were added to refs/heads/master by this push:
     new a3a31cd  fix: set the correct rpc response type when meet an error (#102)
a3a31cd is described below

commit a3a31cd4fb8bdc0cbb5828b6a6a3a2bf80621c62
Author: 刘奕聪 <im...@gmail.com>
AuthorDate: Tue Sep 27 11:51:59 2022 +0800

    fix: set the correct rpc response type when meet an error (#102)
    
    Co-authored-by: yakovliu <ya...@tencent.com>
---
 internal/server/server.go      | 39 +++++++++++++++++++++++----------------
 internal/server/server_test.go | 19 ++++++++++++++++++-
 2 files changed, 41 insertions(+), 17 deletions(-)

diff --git a/internal/server/server.go b/internal/server/server.go
index c28ddf4..5ac291f 100644
--- a/internal/server/server.go
+++ b/internal/server/server.go
@@ -73,33 +73,40 @@ func recoverPanic() {
 	}
 }
 
-func dispatchRPC(ty byte, in []byte, conn net.Conn) *flatbuffers.Builder {
+func dispatchRPC(ty byte, in []byte, conn net.Conn) (*flatbuffers.Builder, byte) {
 	var err error
 	var bd *flatbuffers.Builder
 	hl, ok := typeHandlerMap[ty]
 	if !ok {
-		err = UnknownType{ty}
-	} else {
-		bd, err = hl(in, conn)
+		log.Warnf("unknown rpc type: %d", ty)
+		return generateErrorReport(UnknownType{ty}), util.RPCError
 	}
 
+	bd, err = hl(in, conn)
 	if err != nil {
-		bd = generateErrorReport(err)
-	} else {
-		bd = checkIfDataTooLarge(bd)
+		return generateErrorReport(err), util.RPCError
+	}
+
+	replaced, ok := checkIfDataTooLarge(bd)
+	if !ok {
+		return replaced, util.RPCError
 	}
-	return bd
+
+	return bd, ty
 }
 
-func checkIfDataTooLarge(bd *flatbuffers.Builder) *flatbuffers.Builder {
+func checkIfDataTooLarge(bd *flatbuffers.Builder) (*flatbuffers.Builder, bool) {
 	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)
-		bd = generateErrorReport(err)
+	if size < util.MaxDataSize {
+		return bd, true
 	}
-	return bd
+
+	err := fmt.Errorf("the max length of data is %d but got %d", util.MaxDataSize, size)
+	util.PutBuilder(bd)
+	bd = generateErrorReport(err)
+
+	return bd, false
 }
 
 func handleConn(c net.Conn) {
@@ -129,11 +136,11 @@ func handleConn(c net.Conn) {
 			break
 		}
 
-		bd := dispatchRPC(ty, buf, c)
+		bd, respTy := dispatchRPC(ty, buf, c)
 		out := bd.FinishedBytes()
 		size := len(out)
 		binary.BigEndian.PutUint32(header, uint32(size))
-		header[0] = ty
+		header[0] = respTy
 
 		n, err = c.Write(header)
 		if err != nil {
diff --git a/internal/server/server_test.go b/internal/server/server_test.go
index c698cb7..b854973 100644
--- a/internal/server/server_test.go
+++ b/internal/server/server_test.go
@@ -29,6 +29,7 @@ import (
 	hrc "github.com/api7/ext-plugin-proto/go/A6/HTTPReqCall"
 	"github.com/stretchr/testify/assert"
 
+	"github.com/apache/apisix-go-plugin-runner/internal/plugin"
 	"github.com/apache/apisix-go-plugin-runner/internal/util"
 )
 
@@ -52,10 +53,22 @@ func TestGetConfCacheTTL(t *testing.T) {
 }
 
 func TestDispatchRPC_UnknownType(t *testing.T) {
-	bd := dispatchRPC(126, []byte(""), nil)
+	bd, ty := dispatchRPC(126, []byte(""), nil)
 	err := UnknownType{126}
 	expectBd := ReportError(err)
 	assert.Equal(t, expectBd.FinishedBytes(), bd.FinishedBytes())
+	assert.Equal(t, ty, byte(util.RPCError))
+}
+
+func TestDispatchRPC_KnownType(t *testing.T) {
+	bd := util.GetBuilder()
+	hrc.ReqStart(bd)
+	hrc.ReqAddConfToken(bd, 1)
+	r := hrc.ReqEnd(bd)
+	bd.Finish(r)
+
+	_, ty := dispatchRPC(util.RPCHTTPReqCall, bd.FinishedBytes(), nil)
+	assert.Equal(t, ty, byte(util.RPCError))
 }
 
 func TestDispatchRPC_OutTooLarge(t *testing.T) {
@@ -117,3 +130,7 @@ func TestRun(t *testing.T) {
 	_, err = os.Stat(path)
 	assert.True(t, os.IsNotExist(err))
 }
+
+func init() {
+	plugin.InitConfCache(time.Second)
+}