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 <[email protected]>
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
 }

Reply via email to