fix: body callbacks return the whole buffered body size (#418)

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,22 +185,21 @@ 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
 }
 
 // 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)
 	}

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))
 		})
 	})
 }

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

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",
 		},
 	}