feat: support request body rewrite #127 (#151)

.github/workflows/lint.yml

@@ -32,9 +32,9 @@ jobs:
     steps:
       - uses: actions/checkout@v2
       - name: setup go
-        uses: actions/setup-go@v1
+        uses: actions/setup-go@v4
         with:
-          go-version: '1.15'
+          go-version: '1.17'
 
       - name: Download golangci-lint
         run: curl -sSfL https://raw.githubusercontent.com/golangci/golangci-lint/master/install.sh | sh -s -- -b $(go env GOPATH)/bin v1.39.0

.gitignore

@@ -7,3 +7,6 @@ release/
 coverage.txt
 logs/
 *.svg
+.vscode
+go.work
+go.work.sum

cmd/go-runner/plugins/limit_req.go

@@ -5,7 +5,7 @@
 // (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
+//	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,

cmd/go-runner/plugins/request_body_rewrite.go

@@ -0,0 +1,64 @@
+/*
+ * 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 plugins
+
+import (
+	"encoding/json"
+	"net/http"
+
+	pkgHTTP "github.com/apache/apisix-go-plugin-runner/pkg/http"
+	"github.com/apache/apisix-go-plugin-runner/pkg/log"
+	"github.com/apache/apisix-go-plugin-runner/pkg/plugin"
+)
+
+const requestBodyRewriteName = "request-body-rewrite"
+
+func init() {
+	if err := plugin.RegisterPlugin(&RequestBodyRewrite{}); err != nil {
+		log.Fatalf("failed to register plugin %s: %s", requestBodyRewriteName, err.Error())
+	}
+}
+
+type RequestBodyRewrite struct {
+	plugin.DefaultPlugin
+}
+
+type RequestBodyRewriteConfig struct {
+	NewBody string `json:"new_body"`
+}
+
+func (*RequestBodyRewrite) Name() string {
+	return requestBodyRewriteName
+}
+
+func (p *RequestBodyRewrite) ParseConf(in []byte) (interface{}, error) {
+	conf := RequestBodyRewriteConfig{}
+	err := json.Unmarshal(in, &conf)
+	if err != nil {
+		log.Errorf("failed to parse config for plugin %s: %s", p.Name(), err.Error())
+	}
+	return conf, err
+}
+
+func (*RequestBodyRewrite) RequestFilter(conf interface{}, _ http.ResponseWriter, r pkgHTTP.Request) {
+	newBody := conf.(RequestBodyRewriteConfig).NewBody
+	if newBody == "" {
+		return
+	}
+	r.SetBody([]byte(newBody))
+}

cmd/go-runner/plugins/request_body_rewrite_test.go

@@ -0,0 +1,132 @@
+/*
+ * 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 plugins
+
+import (
+	"context"
+	"net"
+	"net/http"
+	"net/url"
+	"testing"
+
+	pkgHTTP "github.com/apache/apisix-go-plugin-runner/pkg/http"
+	"github.com/stretchr/testify/require"
+)
+
+func TestRequestBodyRewrite_ParseConf(t *testing.T) {
+	testCases := []struct {
+		name    string
+		in      []byte
+		expect  string
+		wantErr bool
+	}{
+		{
+			"happy path",
+			[]byte(`{"new_body":"hello"}`),
+			"hello",
+			false,
+		},
+		{
+			"empty conf",
+			[]byte(``),
+			"",
+			true,
+		},
+		{
+			"empty body",
+			[]byte(`{"new_body":""}`),
+			"",
+			false,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			p := new(RequestBodyRewrite)
+			conf, err := p.ParseConf(tc.in)
+			if tc.wantErr {
+				require.Error(t, err)
+			} else {
+				require.NoError(t, err)
+			}
+			require.Equal(t, tc.expect, conf.(RequestBodyRewriteConfig).NewBody)
+		})
+	}
+}
+
+func TestRequestBodyRewrite_RequestFilter(t *testing.T) {
+	req := &mockHTTPRequest{body: []byte("hello")}
+	p := new(RequestBodyRewrite)
+	conf, err := p.ParseConf([]byte(`{"new_body":"See ya"}`))
+	require.NoError(t, err)
+	p.RequestFilter(conf, nil, req)
+	require.Equal(t, []byte("See ya"), req.body)
+}
+
+// mockHTTPRequest implements pkgHTTP.Request
+type mockHTTPRequest struct {
+	body []byte
+}
+
+func (r *mockHTTPRequest) SetBody(body []byte) {
+	r.body = body
+}
+
+func (*mockHTTPRequest) Args() url.Values {
+	panic("unimplemented")
+}
+
+func (*mockHTTPRequest) Body() ([]byte, error) {
+	panic("unimplemented")
+}
+
+func (*mockHTTPRequest) Context() context.Context {
+	panic("unimplemented")
+}
+
+func (*mockHTTPRequest) Header() pkgHTTP.Header {
+	panic("unimplemented")
+}
+
+func (*mockHTTPRequest) ID() uint32 {
+	panic("unimplemented")
+}
+
+func (*mockHTTPRequest) Method() string {
+	panic("unimplemented")
+}
+
+func (*mockHTTPRequest) Path() []byte {
+	panic("unimplemented")
+}
+
+func (*mockHTTPRequest) RespHeader() http.Header {
+	panic("unimplemented")
+}
+
+func (*mockHTTPRequest) SetPath([]byte) {
+	panic("unimplemented")
+}
+
+func (*mockHTTPRequest) SrcIP() net.IP {
+	panic("unimplemented")
+}
+
+func (*mockHTTPRequest) Var(string) ([]byte, error) {
+	panic("unimplemented")
+}

go.mod

@@ -4,7 +4,7 @@ go 1.15
 
 require (
 	github.com/ReneKroon/ttlcache/v2 v2.4.0
-	github.com/api7/ext-plugin-proto v0.6.0
+	github.com/api7/ext-plugin-proto v0.6.1
 	github.com/google/flatbuffers v2.0.0+incompatible
 	github.com/spf13/cobra v1.2.1
 	github.com/stretchr/testify v1.7.0

go.sum

@@ -43,8 +43,8 @@ github.com/ReneKroon/ttlcache/v2 v2.4.0 h1:KywGhjik+ZFTDXMNLiPECSzmdx2yNvAlDNKES
 github.com/ReneKroon/ttlcache/v2 v2.4.0/go.mod h1:zbo6Pv/28e21Z8CzzqgYRArQYGYtjONRxaAKGxzQvG4=
 github.com/alvaroloes/enumer v1.1.2/go.mod h1:FxrjvuXoDAx9isTJrv4c+T410zFi0DtXIT0m65DJ+Wo=
 github.com/antihax/optional v1.0.0/go.mod h1:uupD/76wgC+ih3iEmQUL+0Ugr19nfwCT1kdvxnR2qWY=
-github.com/api7/ext-plugin-proto v0.6.0 h1:xmgcKwWRiM9EpBIs1wYJ7Ife/YnLl4IL2NEy4417g60=
-github.com/api7/ext-plugin-proto v0.6.0/go.mod h1:8dbdAgCESeqwZ0IXirbjLbshEntmdrAX3uet+LW3jVU=
+github.com/api7/ext-plugin-proto v0.6.1 h1:eQN0oHacL97ezVGWVmsRigt+ClcpgjipUq0rmW8BG4g=
+github.com/api7/ext-plugin-proto v0.6.1/go.mod h1:8dbdAgCESeqwZ0IXirbjLbshEntmdrAX3uet+LW3jVU=
 github.com/armon/circbuf v0.0.0-20150827004946-bbbad097214e/go.mod h1:3U/XgcO3hCbHZ8TKRvWD2dDTCfh9M9ya+I9JpbB7O8o=
 github.com/armon/go-metrics v0.0.0-20180917152333-f0300d1749da/go.mod h1:Q73ZrmVTwzkszR9V5SSuryQ31EELlFMUz1kKyl939pY=
 github.com/armon/go-radix v0.0.0-20180808171621-7fddfc383310/go.mod h1:ufUuZ+zHj4x4TnLV4JWEpy2hxWSpsRywHrMgIH9cCH8=

internal/http/request.go

@@ -189,6 +189,10 @@ func (r *Request) Body() ([]byte, error) {
 	return v, nil
 }
 
+func (r *Request) SetBody(body []byte) {
+	r.body = body
+}
+
 func (r *Request) Reset() {
 	defer r.cancel()
 	r.path = nil
@@ -205,7 +209,7 @@ func (r *Request) Reset() {
 }
 
 func (r *Request) FetchChanges(id uint32, builder *flatbuffers.Builder) bool {
-	if r.path == nil && r.hdr == nil && r.args == nil && r.respHdr == nil {
+	if !r.hasChanges() {
 		return false
 	}
 
@@ -214,6 +218,11 @@ func (r *Request) FetchChanges(id uint32, builder *flatbuffers.Builder) bool {
 		path = builder.CreateByteString(r.path)
 	}
 
+	var body flatbuffers.UOffsetT
+	if r.body != nil {
+		body = builder.CreateByteVector(r.body)
+	}
+
 	var hdrVec, respHdrVec flatbuffers.UOffsetT
 	if r.hdr != nil {
 		hdrs := []flatbuffers.UOffsetT{}
@@ -314,6 +323,9 @@ func (r *Request) FetchChanges(id uint32, builder *flatbuffers.Builder) bool {
 	if path > 0 {
 		hrc.RewriteAddPath(builder, path)
 	}
+	if body > 0 {
+		hrc.RewriteAddBody(builder, body)
+	}
 	if hdrVec > 0 {
 		hrc.RewriteAddHeaders(builder, hdrVec)
 	}
@@ -346,6 +358,11 @@ func (r *Request) Context() context.Context {
 	return context.Background()
 }
 
+func (r *Request) hasChanges() bool {
+	return r.path != nil || r.hdr != nil ||
+		r.args != nil || r.respHdr != nil || r.body != nil
+}
+
 func (r *Request) askExtraInfo(builder *flatbuffers.Builder,
 	infoType ei.Info, info flatbuffers.UOffsetT) ([]byte, error) {
 

internal/http/request_test.go

@@ -502,6 +502,17 @@ func TestBody(t *testing.T) {
 	}()
 
 	v, err := r.Body()
-	assert.Nil(t, err)
+	assert.NoError(t, err)
 	assert.Equal(t, "Hello, Go Runner", string(v))
+
+	const newBody = "Hello, Rust Runner"
+	r.SetBody([]byte(newBody))
+	v, err = r.Body()
+	assert.NoError(t, err)
+	assert.Equal(t, []byte(newBody), v)
+
+	builder := util.GetBuilder()
+	assert.True(t, r.FetchChanges(1, builder))
+	rewrite := getRewriteAction(t, builder)
+	assert.Equal(t, []byte(newBody), rewrite.BodyBytes())
 }

pkg/http/http.go

@@ -67,6 +67,9 @@ type Request interface {
 	// pkg/common.ErrConnClosed type is returned.
 	Body() ([]byte, error)
 
+	// SetBody rewrites the original request body
+	SetBody([]byte)
+
 	// Context returns the request's context.
 	//
 	// The returned context is always non-nil; it defaults to the

tests/e2e/plugins/plugins_request_body_rewrite_test.go

@@ -0,0 +1,69 @@
+/*
+ * 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 plugins_test
+
+import (
+	"net/http"
+
+	"github.com/apache/apisix-go-plugin-runner/tests/e2e/tools"
+	"github.com/gavv/httpexpect/v2"
+	"github.com/onsi/ginkgo"
+	"github.com/onsi/ginkgo/extensions/table"
+)
+
+var _ = ginkgo.Describe("RequestBodyRewrite Plugin", func() {
+	table.DescribeTable("tries to test request body rewrite feature",
+		func(tc tools.HttpTestCase) {
+			tools.RunTestCase(tc)
+		},
+		table.Entry("config APISIX", tools.HttpTestCase{
+			Object: tools.GetA6CPExpect(),
+			Method: http.MethodPut,
+			Path:   "/apisix/admin/routes/1",
+			Body: `{
+				"uri":"/echo",
+				"plugins":{
+					"ext-plugin-pre-req":{
+						"conf":[
+							{
+								"name":"request-body-rewrite",
+								"value":"{\"new_body\":\"request body rewrite\"}"
+							}
+						]
+					}
+				},
+				"upstream":{
+					"nodes":{
+						"web:8888":1
+					},
+					"type":"roundrobin"
+				}
+			}`,
+			Headers:           map[string]string{"X-API-KEY": tools.GetAdminToken()},
+			ExpectStatusRange: httpexpect.Status2xx,
+		}),
+		table.Entry("should rewrite request body", tools.HttpTestCase{
+			Object:       tools.GetA6DPExpect(),
+			Method:       http.MethodGet,
+			Path:         "/echo",
+			Body:         "hello hello world world",
+			ExpectBody:   "request body rewrite",
+			ExpectStatus: http.StatusOK,
+		}),
+	)
+})