add close option for HTTPContext.Request().SetBody method

pkg/context/contexttest/request.go

@@ -44,7 +44,7 @@ type MockedHTTPRequest struct {
 	MockedCookies     func() []*http.Cookie
 	MockedAddCookie   func(cookie *http.Cookie)
 	MockedBody        func() io.Reader
-	MockedSetBody     func(io.Reader)
+	MockedSetBody     func(io.Reader, bool)
 	MockedStd         func() *http.Request
 	MockedSize        func() uint64
 }
@@ -189,9 +189,9 @@ func (r *MockedHTTPRequest) Body() io.Reader {
 }
 
 // SetBody mocks the SetBody function of HTTPRequest
-func (r *MockedHTTPRequest) SetBody(body io.Reader) {
+func (r *MockedHTTPRequest) SetBody(body io.Reader, closePreviousReader bool) {
 	if r.MockedSetBody != nil {
-		r.MockedSetBody(body)
+		r.MockedSetBody(body, closePreviousReader)
 	}
 }
 

pkg/context/httpcontext.go

@@ -100,7 +100,7 @@ type (
 		AddCookie(cookie *http.Cookie)
 
 		Body() io.Reader
-		SetBody(io.Reader)
+		SetBody(io.Reader, bool)
 
 		Std() *http.Request
 

pkg/context/httprequest.go

@@ -155,8 +155,8 @@ func (r *httpRequest) Body() io.Reader {
 	return r.body
 }
 
-func (r *httpRequest) SetBody(reader io.Reader) {
-	r.body = callbackreader.New(reader)
+func (r *httpRequest) SetBody(reader io.Reader, closePreviousReader bool) {
+	r.body.SetReader(reader, closePreviousReader)
 }
 
 func (r *httpRequest) Size() uint64 {

pkg/context/httptemplate.go

@@ -332,7 +332,7 @@ func saveReqBody(e *HTTPTemplate, filterName string, ctx HTTPContext) error {
 	}
 	e.Engine.SetDict(fmt.Sprintf(filterReqBody, filterName), string(bodyBuff.Bytes()))
 	// Reset back into Request's Body
-	ctx.Request().SetBody(bodyBuff)
+	ctx.Request().SetBody(bodyBuff, true)
 
 	return nil
 }

pkg/filter/headertojson/headertojson.go

@@ -203,6 +203,6 @@ func (h *HeaderToJSON) handle(ctx context.HTTPContext) string {
 	if err != nil {
 		return resultJSONEncodeDecodeErr
 	}
-	ctx.Request().SetBody(bytes.NewReader(bodyBytes))
+	ctx.Request().SetBody(bytes.NewReader(bodyBytes), true)
 	return ""
 }

pkg/filter/proxy/proxy.go

@@ -318,7 +318,7 @@ func (b *Proxy) Handle(ctx context.HTTPContext) (result string) {
 func (b *Proxy) handle(ctx context.HTTPContext) (result string) {
 	if b.mirrorPool != nil && b.mirrorPool.filter.Filter(ctx) {
 		primaryBody, secondaryBody := newPrimarySecondaryReader(ctx.Request().Body())
-		ctx.Request().SetBody(primaryBody)
+		ctx.Request().SetBody(primaryBody, false)
 
 		wg := &sync.WaitGroup{}
 		wg.Add(1)

pkg/filter/remotefilter/remotefilter.go

@@ -309,7 +309,7 @@ func (rf *RemoteFilter) unmarshalHTTPContext(buff []byte, ctx context.HTTPContex
 	r.SetPath(re.Path)
 	r.SetQuery(re.Query)
 	r.Header().Reset(re.Header)
-	r.SetBody(bytes.NewReader(re.Body))
+	r.SetBody(bytes.NewReader(re.Body), true)
 
 	if we == nil {
 		return

pkg/filter/requestadaptor/requestadaptor.go

@@ -151,10 +151,10 @@ func (ra *RequestAdaptor) handle(ctx context.HTTPContext) string {
 				logger.Errorf("BUG request render body failed, template %s, err %v",
 					ra.spec.Body, err)
 			} else {
-				ctx.Request().SetBody(bytes.NewReader([]byte(body)))
+				ctx.Request().SetBody(bytes.NewReader([]byte(body)), true)
 			}
 		} else {
-			ctx.Request().SetBody(bytes.NewReader([]byte(ra.spec.Body)))
+			ctx.Request().SetBody(bytes.NewReader([]byte(ra.spec.Body)), true)
 		}
 		ctx.Request().Header().Del("Content-Encoding")
 	}
@@ -203,7 +203,7 @@ func (ra *RequestAdaptor) processCompress(ctx context.HTTPContext) string {
 	}
 	gw.Close()
 
-	ctx.Request().SetBody(&buf)
+	ctx.Request().SetBody(&buf, true)
 	ctx.Request().Header().Set("Content-Encoding", "gzip")
 	return ""
 }
@@ -220,7 +220,7 @@ func (ra *RequestAdaptor) processDecompress(ctx context.HTTPContext) string {
 		if err != nil {
 			return resultDecompressFail
 		}
-		ctx.Request().SetBody(bytes.NewReader(data))
+		ctx.Request().SetBody(bytes.NewReader(data), true)
 		ctx.Request().Header().Del("Content-Encoding")
 	}
 	return ""

pkg/filter/retryer/retryer.go

@@ -179,7 +179,7 @@ func (r *Retryer) handle(ctx context.HTTPContext, u *URLRule) string {
 	data, _ := io.ReadAll(ctx.Request().Body())
 	for {
 		attempt++
-		ctx.Request().SetBody(bytes.NewReader(data))
+		ctx.Request().SetBody(bytes.NewReader(data), true)
 
 		result := ctx.CallNextHandler("")
 

pkg/util/callbackreader/callbackreader.go

@@ -88,3 +88,15 @@ func (cr *CallbackReader) Close() error {
 
 	return nil
 }
+
+// SetReader replace previous reader with new reader. If closePreviousReader set to true, it will close
+// previous reader.
+func (cr *CallbackReader) SetReader(reader io.Reader, closePreviousReader bool) {
+	if closePreviousReader {
+		closer, ok := cr.reader.(io.Closer)
+		if ok {
+			closer.Close()
+		}
+	}
+	cr.reader = reader
+}