[apisix-go-plugin-runner] 13/22: feat: record response

[sp...@apache.org]

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

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

commit 38309c734030f838e90b4660df3376d4622047ec
Author: spacewander <sp...@gmail.com>
AuthorDate: Mon May 24 13:04:12 2021 +0800

    feat: record response
---
 internal/http/request.go       |  4 ++
 internal/http/response.go      | 92 ++++++++++++++++++++++++++++++++++++++---
 internal/http/response_test.go | 93 ++++++++++++++++++++++++++++++++++++++++++
 internal/plugin/plugin.go      | 11 ++++-
 internal/plugin/plugin_test.go |  2 +
 pkg/http/http.go               |  2 +
 6 files changed, 196 insertions(+), 8 deletions(-)

diff --git a/internal/http/request.go b/internal/http/request.go
index 4c96aa2..596b0c4 100644
--- a/internal/http/request.go
+++ b/internal/http/request.go
@@ -27,6 +27,10 @@ func (r Request) ConfToken() uint32 {
 	return r.r.ConfToken()
 }
 
+func (r Request) Id() uint32 {
+	return r.r.Id()
+}
+
 func CreateRequest(buf []byte) *Request {
 	req := &Request{
 		r: hrc.GetRootAsReq(buf, 0),
diff --git a/internal/http/response.go b/internal/http/response.go
index 46fbeeb..ef28b9a 100644
--- a/internal/http/response.go
+++ b/internal/http/response.go
@@ -14,24 +14,104 @@
 // limitations under the License.
 package http
 
-import "net/http"
+import (
+	"bytes"
+	"net/http"
+
+	"github.com/api7/ext-plugin-proto/go/A6"
+	hrc "github.com/api7/ext-plugin-proto/go/A6/HTTPReqCall"
+	flatbuffers "github.com/google/flatbuffers/go"
+)
 
 type Response struct {
+	hdr  http.Header
+	body *bytes.Buffer
 	code int
 }
 
-func (r Response) Header() http.Header {
-	return nil
+func (r *Response) Header() http.Header {
+	r.hdr = http.Header{}
+	return r.hdr
 }
 
-func (r Response) Write([]byte) (int, error) {
-	return 0, nil
+func (r *Response) Write(b []byte) (int, error) {
+	if r.body == nil {
+		r.body = &bytes.Buffer{}
+	}
+
+	return r.body.Write(b)
 }
 
-func (r Response) WriteHeader(statusCode int) {
+func (r *Response) WriteHeader(statusCode int) {
+	if r.code != 0 {
+		// official WriteHeader can't override written status
+		// keep the same behavior
+		return
+	}
 	r.code = statusCode
 }
 
+func (r *Response) FetchChanges(id uint32, builder *flatbuffers.Builder) bool {
+	if r.body == nil && r.code == 0 && len(r.hdr) == 0 {
+		return false
+	}
+
+	hdrLen := len(r.hdr)
+	var hdrVec flatbuffers.UOffsetT
+	if hdrLen > 0 {
+		hdrs := []flatbuffers.UOffsetT{}
+		for n, arr := range r.hdr {
+			for _, v := range arr {
+				name := builder.CreateString(n)
+				value := builder.CreateString(v)
+				A6.TextEntryStart(builder)
+				A6.TextEntryAddName(builder, name)
+				A6.TextEntryAddValue(builder, value)
+				te := A6.TextEntryEnd(builder)
+				hdrs = append(hdrs, te)
+			}
+		}
+		size := len(hdrs)
+		hrc.StopStartHeadersVector(builder, size)
+		for i := size - 1; i >= 0; i-- {
+			te := hdrs[i]
+			builder.PrependUOffsetT(te)
+		}
+		hdrVec = builder.EndVector(size)
+	}
+
+	var bodyVec flatbuffers.UOffsetT
+	if r.body != nil {
+		b := r.body.Bytes()
+		if len(b) > 0 {
+			bodyVec = builder.CreateByteVector(b)
+		}
+	}
+
+	hrc.StopStart(builder)
+	if r.code == 0 {
+		hrc.StopAddStatus(builder, 200)
+	} else {
+		hrc.StopAddStatus(builder, uint16(r.code))
+	}
+	if hdrLen > 0 {
+		hrc.StopAddHeaders(builder, hdrVec)
+	}
+	if r.body != nil {
+		hrc.StopAddBody(builder, bodyVec)
+	}
+	stop := hrc.StopEnd(builder)
+
+	hrc.RespStart(builder)
+	hrc.RespAddId(builder, id)
+	hrc.RespAddActionType(builder, hrc.ActionStop)
+	hrc.RespAddAction(builder, stop)
+	res := hrc.RespEnd(builder)
+	builder.Finish(res)
+
+	return true
+}
+
 func CreateResponse() *Response {
 	return &Response{}
 }
diff --git a/internal/http/response_test.go b/internal/http/response_test.go
new file mode 100644
index 0000000..5eddd4c
--- /dev/null
+++ b/internal/http/response_test.go
@@ -0,0 +1,93 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements.  See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License.  You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+package http
+
+import (
+	"net/http"
+	"testing"
+
+	"github.com/api7/ext-plugin-proto/go/A6"
+	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"
+)
+
+func getStopAction(t *testing.T, b *flatbuffers.Builder) *hrc.Stop {
+	buf := b.FinishedBytes()
+	res := hrc.GetRootAsResp(buf, 0)
+	tab := &flatbuffers.Table{}
+	if res.Action(tab) {
+		assert.Equal(t, hrc.ActionStop, res.ActionType())
+		stop := &hrc.Stop{}
+		stop.Init(tab.Bytes, tab.Pos)
+		return stop
+	}
+	return nil
+}
+
+func TestFetchChanges(t *testing.T) {
+	r := CreateResponse()
+	r.Write([]byte("hello"))
+	h := r.Header()
+	h.Set("foo", "bar")
+	h.Add("foo", "baz")
+	h.Add("cat", "dog")
+	r.Write([]byte(" world"))
+	assert.Equal(t, "dog", h.Get("cat"))
+	builder := util.GetBuilder()
+	assert.True(t, r.FetchChanges(1, builder))
+
+	stop := getStopAction(t, builder)
+	assert.Equal(t, uint16(200), stop.Status())
+	assert.Equal(t, []byte("hello world"), stop.BodyBytes())
+
+	res := http.Header{}
+	assert.Equal(t, 3, stop.HeadersLength())
+	for i := 0; i < stop.HeadersLength(); i++ {
+		e := &A6.TextEntry{}
+		stop.Headers(e, i)
+		res.Add(string(e.Name()), string(e.Value()))
+	}
+	assert.Equal(t, h, res)
+}
+
+func TestFetchChangesEmptyResponse(t *testing.T) {
+	r := CreateResponse()
+	builder := util.GetBuilder()
+	assert.False(t, r.FetchChanges(1, builder))
+}
+
+func TestFetchChangesStatusOnly(t *testing.T) {
+	r := CreateResponse()
+	r.WriteHeader(400)
+	builder := util.GetBuilder()
+	assert.True(t, r.FetchChanges(1, builder))
+
+	stop := getStopAction(t, builder)
+	assert.Equal(t, uint16(400), stop.Status())
+}
+
+func TestWriteHeaderTwice(t *testing.T) {
+	r := CreateResponse()
+	r.WriteHeader(400)
+	r.WriteHeader(503)
+	builder := util.GetBuilder()
+	assert.True(t, r.FetchChanges(1, builder))
+
+	stop := getStopAction(t, builder)
+	assert.Equal(t, uint16(400), stop.Status())
+}
diff --git a/internal/plugin/plugin.go b/internal/plugin/plugin.go
index ba07c8a..e4d9ef6 100644
--- a/internal/plugin/plugin.go
+++ b/internal/plugin/plugin.go
@@ -29,9 +29,15 @@ func handle(conf RuleConf, w http.ResponseWriter, r pkgHTTP.Request) error {
 	return nil
 }
 
-func reportAction(req *inHTTP.Request, resp *inHTTP.Response) *flatbuffers.Builder {
+func reportAction(id uint32, req *inHTTP.Request, resp *inHTTP.Response) *flatbuffers.Builder {
 	builder := util.GetBuilder()
+
+	if resp != nil && resp.FetchChanges(id, builder) {
+		return builder
+	}
+
 	hrc.RespStart(builder)
+	hrc.RespAddId(builder, id)
 	res := hrc.RespEnd(builder)
 	builder.Finish(res)
 	return builder
@@ -51,6 +57,7 @@ func HTTPReqCall(buf []byte) (*flatbuffers.Builder, error) {
 		return nil, err
 	}
 
-	builder := reportAction(req, resp)
+	id := req.Id()
+	builder := reportAction(id, req, resp)
 	return builder, nil
 }
diff --git a/internal/plugin/plugin_test.go b/internal/plugin/plugin_test.go
index 5433cca..c433f3f 100644
--- a/internal/plugin/plugin_test.go
+++ b/internal/plugin/plugin_test.go
@@ -29,6 +29,7 @@ func TestHTTPReqCall(t *testing.T) {
 
 	builder := flatbuffers.NewBuilder(1024)
 	hrc.ReqStart(builder)
+	hrc.ReqAddId(builder, 233)
 	hrc.ReqAddConfToken(builder, 1)
 	r := hrc.ReqEnd(builder)
 	builder.Finish(r)
@@ -39,5 +40,6 @@ func TestHTTPReqCall(t *testing.T) {
 
 	out = b.FinishedBytes()
 	resp := hrc.GetRootAsResp(out, 0)
+	assert.Equal(t, uint32(233), resp.Id())
 	assert.Equal(t, hrc.ActionNONE, resp.ActionType())
 }
diff --git a/pkg/http/http.go b/pkg/http/http.go
index 54edde1..de83956 100644
--- a/pkg/http/http.go
+++ b/pkg/http/http.go
@@ -24,4 +24,6 @@ package http
 // So the server must parse all the headers, ...". The official API is suboptimal, which
 // is even worse in our case as it is not a real HTTP server.
 type Request interface {
+	// Id returns the request id
+	Id() uint32
 }