[apisix-go-plugin-runner] 16/22: feat: add Request.Header

[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 65cbe80b2aba18644d8e485dc2f3a4f7814af36b
Author: spacewander <sp...@gmail.com>
AuthorDate: Wed May 26 10:41:05 2021 +0800

    feat: add Request.Header
---
 internal/http/request.go      | 78 ++++++++++++++++++++++++++++++++++++++++-
 internal/http/request_test.go | 81 +++++++++++++++++++++++++++++++++++++++++--
 pkg/http/http.go              | 26 +++++++++++++-
 3 files changed, 180 insertions(+), 5 deletions(-)

diff --git a/internal/http/request.go b/internal/http/request.go
index 570c245..11c9ef4 100644
--- a/internal/http/request.go
+++ b/internal/http/request.go
@@ -16,7 +16,10 @@ package http
 
 import (
 	"net"
+	"net/http"
 
+	pkgHTTP "github.com/apache/apisix-go-plugin-runner/pkg/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"
 )
@@ -26,6 +29,9 @@ type Request struct {
 	r *hrc.Req
 
 	path []byte
+
+	hdr    *Header
+	rawHdr http.Header
 }
 
 func (r *Request) ConfToken() uint32 {
@@ -55,8 +61,25 @@ func (r *Request) SetPath(path []byte) {
 	r.path = path
 }
 
+func (r *Request) Header() pkgHTTP.Header {
+	if r.hdr == nil {
+		hdr := newHeader()
+		hh := hdr.View()
+		size := r.r.HeadersLength()
+		obj := A6.TextEntry{}
+		for i := 0; i < size; i++ {
+			if r.r.Headers(&obj, i) {
+				hh.Add(string(obj.Name()), string(obj.Value()))
+			}
+		}
+		r.hdr = hdr
+		r.rawHdr = hdr.Clone()
+	}
+	return r.hdr
+}
+
 func (r *Request) FetchChanges(id uint32, builder *flatbuffers.Builder) bool {
-	if r.path == nil {
+	if r.path == nil && r.hdr == nil {
 		return false
 	}
 
@@ -65,10 +88,49 @@ func (r *Request) FetchChanges(id uint32, builder *flatbuffers.Builder) bool {
 		path = builder.CreateByteString(r.path)
 	}
 
+	var hdrVec flatbuffers.UOffsetT
+	if r.hdr != nil {
+		hdrs := []flatbuffers.UOffsetT{}
+		oldHdr := r.rawHdr
+		newHdr := r.hdr.View()
+		for n := range oldHdr {
+			if _, ok := newHdr[n]; !ok {
+				// deleted
+				name := builder.CreateString(n)
+				A6.TextEntryStart(builder)
+				A6.TextEntryAddName(builder, name)
+				te := A6.TextEntryEnd(builder)
+				hdrs = append(hdrs, te)
+			}
+		}
+		for n, v := range newHdr {
+			if raw, ok := oldHdr[n]; !ok || raw[0] != v[0] {
+				// set
+				name := builder.CreateString(n)
+				value := builder.CreateString(v[0])
+				A6.TextEntryStart(builder)
+				A6.TextEntryAddName(builder, name)
+				A6.TextEntryAddValue(builder, value)
+				te := A6.TextEntryEnd(builder)
+				hdrs = append(hdrs, te)
+			}
+		}
+		size := len(hdrs)
+		hrc.RewriteStartHeadersVector(builder, size)
+		for i := size - 1; i >= 0; i-- {
+			te := hdrs[i]
+			builder.PrependUOffsetT(te)
+		}
+		hdrVec = builder.EndVector(size)
+	}
+
 	hrc.RewriteStart(builder)
 	if path > 0 {
 		hrc.RewriteAddPath(builder, path)
 	}
+	if hdrVec > 0 {
+		hrc.RewriteAddHeaders(builder, hdrVec)
+	}
 	rewrite := hrc.RewriteEnd(builder)
 
 	hrc.RespStart(builder)
@@ -87,3 +149,17 @@ func CreateRequest(buf []byte) *Request {
 	}
 	return req
 }
+
+type Header struct {
+	http.Header
+}
+
+func newHeader() *Header {
+	return &Header{
+		Header: http.Header{},
+	}
+}
+
+func (h *Header) View() http.Header {
+	return h.Header
+}
diff --git a/internal/http/request_test.go b/internal/http/request_test.go
index 667138e..54fed1e 100644
--- a/internal/http/request_test.go
+++ b/internal/http/request_test.go
@@ -16,6 +16,7 @@ package http
 
 import (
 	"net"
+	"net/http"
 	"testing"
 
 	"github.com/apache/apisix-go-plugin-runner/internal/util"
@@ -38,10 +39,16 @@ func getRewriteAction(t *testing.T, b *flatbuffers.Builder) *hrc.Rewrite {
 	return nil
 }
 
+type pair struct {
+	name  string
+	value string
+}
+
 type reqOpt struct {
-	srcIP  []byte
-	method A6.Method
-	path   string
+	srcIP   []byte
+	method  A6.Method
+	path    string
+	headers []pair
 }
 
 func buildReq(opt reqOpt) []byte {
@@ -57,6 +64,28 @@ func buildReq(opt reqOpt) []byte {
 		path = builder.CreateString(opt.path)
 	}
 
+	hdrLen := len(opt.headers)
+	var hdrVec flatbuffers.UOffsetT
+	if hdrLen > 0 {
+		hdrs := []flatbuffers.UOffsetT{}
+		for _, v := range opt.headers {
+			name := builder.CreateString(v.name)
+			value := builder.CreateString(v.value)
+			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)
+	}
+
 	hrc.ReqStart(builder)
 	hrc.ReqAddId(builder, 233)
 	hrc.ReqAddConfToken(builder, 1)
@@ -69,6 +98,9 @@ func buildReq(opt reqOpt) []byte {
 	if path > 0 {
 		hrc.ReqAddPath(builder, path)
 	}
+	if hdrVec > 0 {
+		hrc.ReqAddHeaders(builder, hdrVec)
+	}
 	r := hrc.ReqEnd(builder)
 	builder.Finish(r)
 	return builder.FinishedBytes()
@@ -111,3 +143,46 @@ func TestPath(t *testing.T) {
 	rewrite := getRewriteAction(t, builder)
 	assert.Equal(t, "/go", string(rewrite.Path()))
 }
+
+func TestHeader(t *testing.T) {
+	out := buildReq(reqOpt{headers: []pair{
+		{"k", "v"},
+		{"cache-control", "no-cache"},
+		{"cache-control", "no-store"},
+		{"cat", "dog"},
+	}})
+	r := CreateRequest(out)
+	hdr := r.Header()
+	assert.Equal(t, "v", hdr.Get("k"))
+	assert.Equal(t, "no-cache", hdr.Get("Cache-Control"))
+	assert.Equal(t, "no-cache", hdr.Get("cache-control"))
+
+	hdr.Del("empty")
+	hdr.Del("k")
+	assert.Equal(t, "", hdr.Get("k"))
+
+	hdr.Set("cache-control", "max-age=10s")
+	assert.Equal(t, "max-age=10s", hdr.Get("Cache-Control"))
+	hdr.Del("cache-Control")
+	assert.Equal(t, "", hdr.Get("cache-control"))
+
+	hdr.Set("k", "v2")
+	hdr.Del("cat")
+
+	builder := util.GetBuilder()
+	assert.True(t, r.FetchChanges(1, builder))
+	rewrite := getRewriteAction(t, builder)
+	assert.Equal(t, 3, rewrite.HeadersLength())
+
+	exp := http.Header{}
+	exp.Set("Cache-Control", "")
+	exp.Set("cat", "")
+	exp.Set("k", "v2")
+	res := http.Header{}
+	for i := 0; i < rewrite.HeadersLength(); i++ {
+		e := &A6.TextEntry{}
+		rewrite.Headers(e, i)
+		res.Add(string(e.Name()), string(e.Value()))
+	}
+	assert.Equal(t, exp, res)
+}
diff --git a/pkg/http/http.go b/pkg/http/http.go
index ebae8ac..cbbceeb 100644
--- a/pkg/http/http.go
+++ b/pkg/http/http.go
@@ -14,7 +14,10 @@
 // limitations under the License.
 package http
 
-import "net"
+import (
+	"net"
+	"net/http"
+)
 
 // Request represents the HTTP request received by APISIX.
 // We don't use net/http's Request because it doesn't suit our purpose.
@@ -38,4 +41,25 @@ type Request interface {
 	Path() []byte
 	// SetPath is the setter for Path
 	SetPath([]byte)
+	// Header returns the HTTP headers
+	Header() Header
+}
+
+// Header is like http.Header, but only implements the subset of its methods
+type Header interface {
+	// Set sets the header entries associated with key to the single element value.
+	// It replaces any existing values associated with key.
+	// The key is case insensitive
+	Set(key, value string)
+	// Del deletes the values associated with key. The key is case insensitive
+	Del(key string)
+	// Get gets the first value associated with the given key.
+	// If there are no values associated with the key, Get returns "".
+	// It is case insensitive
+	Get(key string) string
+	// View returns the internal structure. It is expected for read operations. Any write operation
+	// won't be recorded
+	View() http.Header
+
+	// TODO: support Add
 }