From 9c1dd2ba6bf1b04c9d86befd01c8d39fa0172693 Mon Sep 17 00:00:00 2001 From: Matteo Pace Date: Sat, 23 Dec 2023 21:03:41 +0100 Subject: [PATCH] fix: body callbacks return the whole buffered body size (#418) --- examples/http_body/main.go | 23 +++++++++-------------- examples/http_body/main_test.go | 4 ++-- proxywasm/proxytest/http.go | 4 ++-- proxywasm/proxytest/http_test.go | 6 ++++-- 4 files changed, 17 insertions(+), 20 deletions(-) diff --git a/examples/http_body/main.go b/examples/http_body/main.go index 1f76ec89..c8a07a81 100644 --- a/examples/http_body/main.go +++ b/examples/http_body/main.go @@ -69,16 +69,14 @@ type setBodyContext struct { // Embed the default root http context here, // so that we don't need to reimplement all the methods. types.DefaultHttpContext - modifyResponse bool - totalRequestBodySize int - totalResponseBodySize int - bufferOperation string + modifyResponse bool + bufferOperation string } // Override types.DefaultHttpContext. func (ctx *setBodyContext) OnHttpRequestHeaders(numHeaders int, endOfStream bool) types.Action { mode, err := proxywasm.GetHttpRequestHeader("buffer-replace-at") - if mode == "response" { + if err == nil && mode == "response" { ctx.modifyResponse = true } @@ -112,13 +110,12 @@ func (ctx *setBodyContext) OnHttpRequestBody(bodySize int, endOfStream bool) typ return types.ActionContinue } - ctx.totalRequestBodySize += bodySize if !endOfStream { // Wait until we see the entire body to replace. return types.ActionPause } - - originalBody, err := proxywasm.GetHttpRequestBody(0, ctx.totalRequestBodySize) + // Being the body never been sent upstream so far, bodySize is the total size of the body received. + originalBody, err := proxywasm.GetHttpRequestBody(0, bodySize) if err != nil { proxywasm.LogErrorf("failed to get request body: %v", err) return types.ActionContinue @@ -160,13 +157,12 @@ func (ctx *setBodyContext) OnHttpResponseBody(bodySize int, endOfStream bool) ty return types.ActionContinue } - ctx.totalResponseBodySize += bodySize if !endOfStream { // Wait until we see the entire body to replace. return types.ActionPause } - originalBody, err := proxywasm.GetHttpResponseBody(0, ctx.totalResponseBodySize) + originalBody, err := proxywasm.GetHttpResponseBody(0, bodySize) if err != nil { proxywasm.LogErrorf("failed to get response body: %v", err) return types.ActionContinue @@ -189,7 +185,7 @@ func (ctx *setBodyContext) OnHttpResponseBody(bodySize int, endOfStream bool) ty } type echoBodyContext struct { - // mbed the default plugin context + // Embed the default plugin context // so that you don't need to reimplement all the methods by yourself. types.DefaultHttpContext totalRequestBodySize int @@ -197,14 +193,13 @@ type echoBodyContext struct { // Override types.DefaultHttpContext. func (ctx *echoBodyContext) OnHttpRequestBody(bodySize int, endOfStream bool) types.Action { - ctx.totalRequestBodySize += bodySize + ctx.totalRequestBodySize = bodySize if !endOfStream { // Wait until we see the entire body to replace. return types.ActionPause } - // Send the request body as the response body. - body, _ := proxywasm.GetHttpRequestBody(0, ctx.totalRequestBodySize) + body, _ := proxywasm.GetHttpRequestBody(0, bodySize) if err := proxywasm.SendHttpResponse(200, nil, body, -1); err != nil { panic(err) } diff --git a/examples/http_body/main_test.go b/examples/http_body/main_test.go index 8493b162..e0674566 100644 --- a/examples/http_body/main_test.go +++ b/examples/http_body/main_test.go @@ -264,7 +264,7 @@ func TestEchoBodyContext_OnHttpRequestBody(t *testing.T) { // Create http context. id := host.InitializeHttpContext() - for _, frame := range []string{"frame1...", "frame2..."} { + for _, frame := range []string{"frame1...", "frame2...", "frame3..."} { // Call OnRequestHeaders without "content-length" action := host.CallOnRequestBody(id, []byte(frame), false /* end of stream */) @@ -282,7 +282,7 @@ func TestEchoBodyContext_OnHttpRequestBody(t *testing.T) { localResponse := host.GetSentLocalResponse(id) require.NotNil(t, localResponse) require.Equal(t, uint32(200), localResponse.StatusCode) - require.Equal(t, "frame1...frame2...", string(localResponse.Data)) + require.Equal(t, "frame1...frame2...frame3...", string(localResponse.Data)) }) }) } diff --git a/proxywasm/proxytest/http.go b/proxywasm/proxytest/http.go index 843eb181..49c66391 100644 --- a/proxywasm/proxytest/http.go +++ b/proxywasm/proxytest/http.go @@ -416,7 +416,7 @@ func (h *httpHostEmulator) CallOnRequestBody(contextID uint32, body []byte, endO cs.requestBody = append(cs.requestBodyBuffer, body...) cs.action = internal.ProxyOnRequestBody(contextID, - len(body), endOfStream) + len(cs.requestBody), endOfStream) if cs.action == types.ActionPause { // Buffering requested cs.requestBodyBuffer = cs.requestBody @@ -435,7 +435,7 @@ func (h *httpHostEmulator) CallOnResponseBody(contextID uint32, body []byte, end cs.responseBody = append(cs.responseBodyBuffer, body...) cs.action = internal.ProxyOnResponseBody(contextID, - len(body), endOfStream) + len(cs.responseBody), endOfStream) if cs.action == types.ActionPause { // Buffering requested cs.responseBodyBuffer = cs.responseBody diff --git a/proxywasm/proxytest/http_test.go b/proxywasm/proxytest/http_test.go index 791c2a52..74b92893 100644 --- a/proxywasm/proxytest/http_test.go +++ b/proxywasm/proxytest/http_test.go @@ -85,13 +85,15 @@ func TestBodyBuffering(t *testing.T) { name: "buffered", buffered: true, action: types.ActionPause, - logged: "11111", + // The first chunk has been buffered, therefore it will be retrieved when calling GetHttpRequestBody at the end of stream. + logged: "1111122222", }, { name: "unbuffered", buffered: false, action: types.ActionContinue, - logged: "22222", + // The first chunk has not been buffered, therefore it will not be retrieved when calling GetHttpRequestBody at the end of stream. + logged: "22222", }, }