[apisix-go-plugin-runner] 14/22: feat: extend pkg Request

[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 d8647f9bf36785c976fd0400abab488ead73212d
Author: spacewander <sp...@gmail.com>
AuthorDate: Tue May 25 14:58:51 2021 +0800

    feat: extend pkg Request
---
 internal/http/request.go      |  54 +++++++++++++++++++-
 internal/http/request_test.go | 113 ++++++++++++++++++++++++++++++++++++++++++
 internal/plugin/plugin.go     |   2 +-
 pkg/http/http.go              |  16 +++++-
 4 files changed, 180 insertions(+), 5 deletions(-)

diff --git a/internal/http/request.go b/internal/http/request.go
index 596b0c4..570c245 100644
--- a/internal/http/request.go
+++ b/internal/http/request.go
@@ -15,22 +15,72 @@
 package http
 
 import (
+	"net"
+
 	hrc "github.com/api7/ext-plugin-proto/go/A6/HTTPReqCall"
+	flatbuffers "github.com/google/flatbuffers/go"
 )
 
 type Request struct {
 	// the root of the flatbuffers HTTPReqCall Request msg
 	r *hrc.Req
+
+	path []byte
 }
 
-func (r Request) ConfToken() uint32 {
+func (r *Request) ConfToken() uint32 {
 	return r.r.ConfToken()
 }
 
-func (r Request) Id() uint32 {
+func (r *Request) ID() uint32 {
 	return r.r.Id()
 }
 
+func (r *Request) SrcIP() net.IP {
+	return r.r.SrcIpBytes()
+}
+
+func (r *Request) Method() string {
+	return r.r.Method().String()
+}
+
+func (r *Request) Path() []byte {
+	if r.path == nil {
+		return r.r.Path()
+	}
+	return r.path
+}
+
+func (r *Request) SetPath(path []byte) {
+	r.path = path
+}
+
+func (r *Request) FetchChanges(id uint32, builder *flatbuffers.Builder) bool {
+	if r.path == nil {
+		return false
+	}
+
+	var path flatbuffers.UOffsetT
+	if r.path != nil {
+		path = builder.CreateByteString(r.path)
+	}
+
+	hrc.RewriteStart(builder)
+	if path > 0 {
+		hrc.RewriteAddPath(builder, path)
+	}
+	rewrite := hrc.RewriteEnd(builder)
+
+	hrc.RespStart(builder)
+	hrc.RespAddId(builder, id)
+	hrc.RespAddActionType(builder, hrc.ActionRewrite)
+	hrc.RespAddAction(builder, rewrite)
+	res := hrc.RespEnd(builder)
+	builder.Finish(res)
+
+	return true
+}
+
 func CreateRequest(buf []byte) *Request {
 	req := &Request{
 		r: hrc.GetRootAsReq(buf, 0),
diff --git a/internal/http/request_test.go b/internal/http/request_test.go
new file mode 100644
index 0000000..667138e
--- /dev/null
+++ b/internal/http/request_test.go
@@ -0,0 +1,113 @@
+// 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"
+	"testing"
+
+	"github.com/apache/apisix-go-plugin-runner/internal/util"
+	"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"
+)
+
+func getRewriteAction(t *testing.T, b *flatbuffers.Builder) *hrc.Rewrite {
+	buf := b.FinishedBytes()
+	res := hrc.GetRootAsResp(buf, 0)
+	tab := &flatbuffers.Table{}
+	if res.Action(tab) {
+		assert.Equal(t, hrc.ActionRewrite, res.ActionType())
+		rewrite := &hrc.Rewrite{}
+		rewrite.Init(tab.Bytes, tab.Pos)
+		return rewrite
+	}
+	return nil
+}
+
+type reqOpt struct {
+	srcIP  []byte
+	method A6.Method
+	path   string
+}
+
+func buildReq(opt reqOpt) []byte {
+	builder := flatbuffers.NewBuilder(1024)
+
+	var ip flatbuffers.UOffsetT
+	if len(opt.srcIP) > 0 {
+		ip = builder.CreateByteVector(opt.srcIP)
+	}
+
+	var path flatbuffers.UOffsetT
+	if opt.path != "" {
+		path = builder.CreateString(opt.path)
+	}
+
+	hrc.ReqStart(builder)
+	hrc.ReqAddId(builder, 233)
+	hrc.ReqAddConfToken(builder, 1)
+	if ip > 0 {
+		hrc.ReqAddSrcIp(builder, ip)
+	}
+	if opt.method != 0 {
+		hrc.ReqAddMethod(builder, opt.method)
+	}
+	if path > 0 {
+		hrc.ReqAddPath(builder, path)
+	}
+	r := hrc.ReqEnd(builder)
+	builder.Finish(r)
+	return builder.FinishedBytes()
+}
+
+func TestSrcIp(t *testing.T) {
+	for _, ip := range []net.IP{
+		net.IPv4(127, 0, 0, 1),
+		net.IPv4(127, 2, 3, 1),
+		net.ParseIP("2001:db8::68"),
+		net.ParseIP("::12"),
+	} {
+		out := buildReq(reqOpt{srcIP: ip})
+		r := CreateRequest(out)
+		assert.Equal(t, ip, r.SrcIP())
+	}
+}
+
+func TestMethod(t *testing.T) {
+	for _, m := range []A6.Method{
+		A6.MethodGET,
+		A6.MethodPATCH,
+	} {
+		out := buildReq(reqOpt{method: m})
+		r := CreateRequest(out)
+		assert.Equal(t, m.String(), r.Method())
+	}
+}
+
+func TestPath(t *testing.T) {
+	out := buildReq(reqOpt{path: "/apisix"})
+	r := CreateRequest(out)
+	assert.Equal(t, "/apisix", string(r.Path()))
+
+	r.SetPath([]byte("/go"))
+	assert.Equal(t, "/go", string(r.Path()))
+
+	builder := util.GetBuilder()
+	assert.True(t, r.FetchChanges(1, builder))
+	rewrite := getRewriteAction(t, builder)
+	assert.Equal(t, "/go", string(rewrite.Path()))
+}
diff --git a/internal/plugin/plugin.go b/internal/plugin/plugin.go
index e4d9ef6..3ca4159 100644
--- a/internal/plugin/plugin.go
+++ b/internal/plugin/plugin.go
@@ -57,7 +57,7 @@ func HTTPReqCall(buf []byte) (*flatbuffers.Builder, error) {
 		return nil, err
 	}
 
-	id := req.Id()
+	id := req.ID()
 	builder := reportAction(id, req, resp)
 	return builder, nil
 }
diff --git a/pkg/http/http.go b/pkg/http/http.go
index de83956..ebae8ac 100644
--- a/pkg/http/http.go
+++ b/pkg/http/http.go
@@ -14,6 +14,8 @@
 // limitations under the License.
 package http
 
+import "net"
+
 // Request represents the HTTP request received by APISIX.
 // We don't use net/http's Request because it doesn't suit our purpose.
 // Take `Request.Header` as an example:
@@ -24,6 +26,16 @@ 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
+	// ID returns the request id
+	ID() uint32
+	// SrcIP returns the client's IP
+	SrcIP() net.IP
+	// Method returns the HTTP method (GET, POST, PUT, etc.)
+	Method() string
+	// Path returns the path part of the client's URI (without query string and the other parts)
+	// It won't be equal to the one in the Request-Line sent by the client if it has
+	// been rewritten by APISIX
+	Path() []byte
+	// SetPath is the setter for Path
+	SetPath([]byte)
 }