internal/http/request.go - apisix-go-plugin-runner - Git at Google

 /*
  * 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 (
 	"encoding/binary"
 	"net"
 	"net/http"
 	"net/url"
 	"reflect"
 	"sync"

 	"github.com/api7/ext-plugin-proto/go/A6"
 	ei "github.com/api7/ext-plugin-proto/go/A6/ExtraInfo"
 	hrc "github.com/api7/ext-plugin-proto/go/A6/HTTPReqCall"
 	flatbuffers "github.com/google/flatbuffers/go"

 	"github.com/apache/apisix-go-plugin-runner/internal/util"
 	"github.com/apache/apisix-go-plugin-runner/pkg/common"
 	pkgHTTP "github.com/apache/apisix-go-plugin-runner/pkg/http"
 	"github.com/apache/apisix-go-plugin-runner/pkg/log"
 )

 type Request struct {
 	// the root of the flatbuffers HTTPReqCall Request msg
 	r *hrc.Req

 	conn            net.Conn
 	extraInfoHeader []byte

 	path []byte

 	hdr    *Header
 	rawHdr http.Header

 	args    url.Values
 	rawArgs url.Values

 	vars map[string][]byte
 }

 func (r *Request) ConfToken() uint32 {
 	return r.r.ConfToken()
 }

 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) 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 cloneUrlValues(oldV url.Values) url.Values {
 	nv := 0
 	for _, vv := range oldV {
 		nv += len(vv)
 	}
 	sv := make([]string, nv)
 	newV := make(url.Values, len(oldV))
 	for k, vv := range oldV {
 		n := copy(sv, vv)
 		newV[k] = sv[:n:n]
 		sv = sv[n:]
 	}
 	return newV
 }

 func (r *Request) Args() url.Values {
 	if r.args == nil {
 		args := url.Values{}
 		size := r.r.ArgsLength()
 		obj := A6.TextEntry{}
 		for i := 0; i < size; i++ {
 			if r.r.Args(&obj, i) {
 				args.Add(string(obj.Name()), string(obj.Value()))
 			}
 		}
 		r.args = args
 		r.rawArgs = cloneUrlValues(args)
 	}
 	return r.args
 }

 func (r *Request) Var(name string) ([]byte, error) {
 	if r.vars == nil {
 		r.vars = map[string][]byte{}
 	}

 	var v []byte
 	var found bool

 	if v, found = r.vars[name]; !found {
 		var err error

 		builder := util.GetBuilder()
 		varName := builder.CreateString(name)
 		ei.VarStart(builder)
 		ei.VarAddName(builder, varName)
 		varInfo := ei.VarEnd(builder)
 		v, err = r.askExtraInfo(builder, ei.InfoVar, varInfo)
 		util.PutBuilder(builder)

 		if err != nil {
 			return nil, err
 		}

 		r.vars[name] = v
 	}
 	return v, nil
 }

 func (r *Request) Reset() {
 	r.path = nil
 	r.hdr = nil
 	r.args = nil

 	r.vars = nil
 	r.conn = nil

 	// Keep the fields below
 	// r.extraInfoHeader = nil
 }

 func (r *Request) FetchChanges(id uint32, builder *flatbuffers.Builder) bool {
 	if r.path == nil && r.hdr == nil && r.args == nil {
 		return false
 	}

 	var path flatbuffers.UOffsetT
 	if r.path != nil {
 		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)
 	}

 	var argsVec flatbuffers.UOffsetT
 	if r.args != nil {
 		args := []flatbuffers.UOffsetT{}
 		oldArgs := r.rawArgs
 		newArgs := r.args
 		for n := range oldArgs {
 			if _, ok := newArgs[n]; !ok {
 				// deleted
 				name := builder.CreateString(n)
 				A6.TextEntryStart(builder)
 				A6.TextEntryAddName(builder, name)
 				te := A6.TextEntryEnd(builder)
 				args = append(args, te)
 			}
 		}
 		for n, v := range newArgs {
 			if raw, ok := oldArgs[n]; !ok || !reflect.DeepEqual(raw, v) {
 				// set / add
 				for _, vv := range v {
 					name := builder.CreateString(n)
 					value := builder.CreateString(vv)
 					A6.TextEntryStart(builder)
 					A6.TextEntryAddName(builder, name)
 					A6.TextEntryAddValue(builder, value)
 					te := A6.TextEntryEnd(builder)
 					args = append(args, te)
 				}
 			}
 		}
 		size := len(args)
 		hrc.RewriteStartArgsVector(builder, size)
 		for i := size - 1; i >= 0; i-- {
 			te := args[i]
 			builder.PrependUOffsetT(te)
 		}
 		argsVec = builder.EndVector(size)
 	}

 	hrc.RewriteStart(builder)
 	if path > 0 {
 		hrc.RewriteAddPath(builder, path)
 	}
 	if hdrVec > 0 {
 		hrc.RewriteAddHeaders(builder, hdrVec)
 	}
 	if argsVec > 0 {
 		hrc.RewriteAddArgs(builder, argsVec)
 	}
 	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 (r *Request) BindConn(c net.Conn) {
 	r.conn = c
 }

 func (r *Request) askExtraInfo(builder *flatbuffers.Builder,
 	infoType ei.Info, info flatbuffers.UOffsetT) ([]byte, error) {

 	ei.ReqStart(builder)
 	ei.ReqAddInfoType(builder, infoType)
 	ei.ReqAddInfo(builder, info)
 	eiRes := ei.ReqEnd(builder)
 	builder.Finish(eiRes)

 	c := r.conn
 	if len(r.extraInfoHeader) == 0 {
 		r.extraInfoHeader = make([]byte, util.HeaderLen)
 	}
 	header := r.extraInfoHeader

 	out := builder.FinishedBytes()
 	size := len(out)
 	binary.BigEndian.PutUint32(header, uint32(size))
 	header[0] = util.RPCExtraInfo

 	n, err := c.Write(header)
 	if err != nil {
 		util.WriteErr(n, err)
 		return nil, common.ErrConnClosed
 	}

 	n, err = c.Write(out)
 	if err != nil {
 		util.WriteErr(n, err)
 		return nil, common.ErrConnClosed
 	}

 	n, err = c.Read(header)
 	if util.ReadErr(n, err, util.HeaderLen) {
 		return nil, common.ErrConnClosed
 	}

 	ty := header[0]
 	header[0] = 0
 	length := binary.BigEndian.Uint32(header)

 	log.Infof("receive rpc type: %d data length: %d", ty, length)

 	buf := make([]byte, length)
 	n, err = c.Read(buf)
 	if util.ReadErr(n, err, int(length)) {
 		return nil, common.ErrConnClosed
 	}

 	resp := ei.GetRootAsResp(buf, 0)
 	res := resp.ResultBytes()
 	return res, nil
 }

 var reqPool = sync.Pool{
 	New: func() interface{} {
 		return &Request{}
 	},
 }

 func CreateRequest(buf []byte) *Request {
 	req := reqPool.Get().(*Request)
 	req.r = hrc.GetRootAsReq(buf, 0)
 	return req
 }

 func ReuseRequest(r *Request) {
 	r.Reset()
 	reqPool.Put(r)
 }

 type Header struct {
 	http.Header
 }

 func newHeader() *Header {
 	return &Header{
 		Header: http.Header{},
 	}
 }

 func (h *Header) View() http.Header {
 	return h.Header
 }
	/*
	* 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 (
	"encoding/binary"
	"net"
	"net/http"
	"net/url"
	"reflect"
	"sync"

	"github.com/api7/ext-plugin-proto/go/A6"
	ei "github.com/api7/ext-plugin-proto/go/A6/ExtraInfo"
	hrc "github.com/api7/ext-plugin-proto/go/A6/HTTPReqCall"
	flatbuffers "github.com/google/flatbuffers/go"

	"github.com/apache/apisix-go-plugin-runner/internal/util"
	"github.com/apache/apisix-go-plugin-runner/pkg/common"
	pkgHTTP "github.com/apache/apisix-go-plugin-runner/pkg/http"
	"github.com/apache/apisix-go-plugin-runner/pkg/log"
	)

	type Request struct {
	// the root of the flatbuffers HTTPReqCall Request msg
	r *hrc.Req

	conn net.Conn
	extraInfoHeader []byte

	path []byte

	hdr *Header
	rawHdr http.Header

	args url.Values
	rawArgs url.Values

	vars map[string][]byte
	}

	func (r *Request) ConfToken() uint32 {
	return r.r.ConfToken()
	}

	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) 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 cloneUrlValues(oldV url.Values) url.Values {
	nv := 0
	for _, vv := range oldV {
	nv += len(vv)
	}
	sv := make([]string, nv)
	newV := make(url.Values, len(oldV))
	for k, vv := range oldV {
	n := copy(sv, vv)
	newV[k] = sv[:n:n]
	sv = sv[n:]
	}
	return newV
	}

	func (r *Request) Args() url.Values {
	if r.args == nil {
	args := url.Values{}
	size := r.r.ArgsLength()
	obj := A6.TextEntry{}
	for i := 0; i < size; i++ {
	if r.r.Args(&obj, i) {
	args.Add(string(obj.Name()), string(obj.Value()))
	}
	}
	r.args = args
	r.rawArgs = cloneUrlValues(args)
	}
	return r.args
	}

	func (r *Request) Var(name string) ([]byte, error) {
	if r.vars == nil {
	r.vars = map[string][]byte{}
	}

	var v []byte
	var found bool

	if v, found = r.vars[name]; !found {
	var err error

	builder := util.GetBuilder()
	varName := builder.CreateString(name)
	ei.VarStart(builder)
	ei.VarAddName(builder, varName)
	varInfo := ei.VarEnd(builder)
	v, err = r.askExtraInfo(builder, ei.InfoVar, varInfo)
	util.PutBuilder(builder)

	if err != nil {
	return nil, err
	}

	r.vars[name] = v
	}
	return v, nil
	}

	func (r *Request) Reset() {
	r.path = nil
	r.hdr = nil
	r.args = nil

	r.vars = nil
	r.conn = nil

	// Keep the fields below
	// r.extraInfoHeader = nil
	}

	func (r Request) FetchChanges(id uint32, builder flatbuffers.Builder) bool {
	if r.path == nil && r.hdr == nil && r.args == nil {
	return false
	}

	var path flatbuffers.UOffsetT
	if r.path != nil {
	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)
	}

	var argsVec flatbuffers.UOffsetT
	if r.args != nil {
	args := []flatbuffers.UOffsetT{}
	oldArgs := r.rawArgs
	newArgs := r.args
	for n := range oldArgs {
	if _, ok := newArgs[n]; !ok {
	// deleted
	name := builder.CreateString(n)
	A6.TextEntryStart(builder)
	A6.TextEntryAddName(builder, name)
	te := A6.TextEntryEnd(builder)
	args = append(args, te)
	}
	}
	for n, v := range newArgs {
	if raw, ok := oldArgs[n]; !ok \|\| !reflect.DeepEqual(raw, v) {
	// set / add
	for _, vv := range v {
	name := builder.CreateString(n)
	value := builder.CreateString(vv)
	A6.TextEntryStart(builder)
	A6.TextEntryAddName(builder, name)
	A6.TextEntryAddValue(builder, value)
	te := A6.TextEntryEnd(builder)
	args = append(args, te)
	}
	}
	}
	size := len(args)
	hrc.RewriteStartArgsVector(builder, size)
	for i := size - 1; i >= 0; i-- {
	te := args[i]
	builder.PrependUOffsetT(te)
	}
	argsVec = builder.EndVector(size)
	}

	hrc.RewriteStart(builder)
	if path > 0 {
	hrc.RewriteAddPath(builder, path)
	}
	if hdrVec > 0 {
	hrc.RewriteAddHeaders(builder, hdrVec)
	}
	if argsVec > 0 {
	hrc.RewriteAddArgs(builder, argsVec)
	}
	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 (r *Request) BindConn(c net.Conn) {
	r.conn = c
	}

	func (r Request) askExtraInfo(builder flatbuffers.Builder,
	infoType ei.Info, info flatbuffers.UOffsetT) ([]byte, error) {

	ei.ReqStart(builder)
	ei.ReqAddInfoType(builder, infoType)
	ei.ReqAddInfo(builder, info)
	eiRes := ei.ReqEnd(builder)
	builder.Finish(eiRes)

	c := r.conn
	if len(r.extraInfoHeader) == 0 {
	r.extraInfoHeader = make([]byte, util.HeaderLen)
	}
	header := r.extraInfoHeader

	out := builder.FinishedBytes()
	size := len(out)
	binary.BigEndian.PutUint32(header, uint32(size))
	header[0] = util.RPCExtraInfo

	n, err := c.Write(header)
	if err != nil {
	util.WriteErr(n, err)
	return nil, common.ErrConnClosed
	}

	n, err = c.Write(out)
	if err != nil {
	util.WriteErr(n, err)
	return nil, common.ErrConnClosed
	}

	n, err = c.Read(header)
	if util.ReadErr(n, err, util.HeaderLen) {
	return nil, common.ErrConnClosed
	}

	ty := header[0]
	header[0] = 0
	length := binary.BigEndian.Uint32(header)

	log.Infof("receive rpc type: %d data length: %d", ty, length)

	buf := make([]byte, length)
	n, err = c.Read(buf)
	if util.ReadErr(n, err, int(length)) {
	return nil, common.ErrConnClosed
	}

	resp := ei.GetRootAsResp(buf, 0)
	res := resp.ResultBytes()
	return res, nil
	}

	var reqPool = sync.Pool{
	New: func() interface{} {
	return &Request{}
	},
	}

	func CreateRequest(buf []byte) *Request {
	req := reqPool.Get().(*Request)
	req.r = hrc.GetRootAsReq(buf, 0)
	return req
	}

	func ReuseRequest(r *Request) {
	r.Reset()
	reqPool.Put(r)
	}

	type Header struct {
	http.Header
	}

	func newHeader() *Header {
	return &Header{
	Header: http.Header{},
	}
	}

	func (h *Header) View() http.Header {
	return h.Header
	}