Repository: megaease/easegress
Pull request #502 (2022-02-10 09:11:56)

add close option for HTTPContext.Request().SetBody method
The request body in HTTPContext can be `io.Reader` or `io.ReadCloser`, but when we call `SetBody` method in `HTTPContext.Request`, it only replace the reader and not close previous reader.

This pr fix this issue by adding a flag to SetBody method. For now, except the case of PrimarySecondaryReader, all other cases will close previous reader when set new reader.


Patch:
diff --git a/pkg/context/httpcontext.go b/pkg/context/httpcontext.go
index efdaf6e8..dc5198a9 100644
--- a/pkg/context/httpcontext.go
+++ b/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
 
diff --git a/pkg/context/httprequest.go b/pkg/context/httprequest.go
index ea93169b..01e28554 100644
--- a/pkg/context/httprequest.go
+++ b/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 {
diff --git a/pkg/context/httptemplate.go b/pkg/context/httptemplate.go
index a0f4c88a..5b1d6a18 100644
--- a/pkg/context/httptemplate.go
+++ b/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
 }
diff --git a/pkg/filter/headertojson/headertojson.go b/pkg/filter/headertojson/headertojson.go
index 8a745b00..23c34b41 100644
--- a/pkg/filter/headertojson/headertojson.go
+++ b/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 ""
 }
diff --git a/pkg/filter/proxy/proxy.go b/pkg/filter/proxy/proxy.go
index b10fbbbd..ebe525b2 100644
--- a/pkg/filter/proxy/proxy.go
+++ b/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)
diff --git a/pkg/filter/remotefilter/remotefilter.go b/pkg/filter/remotefilter/remotefilter.go
index 6e8152d4..4af1e6c0 100644
--- a/pkg/filter/remotefilter/remotefilter.go
+++ b/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
diff --git a/pkg/filter/requestadaptor/requestadaptor.go b/pkg/filter/requestadaptor/requestadaptor.go
index 43f17776..3dde4a75 100644
--- a/pkg/filter/requestadaptor/requestadaptor.go
+++ b/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 ""
diff --git a/pkg/filter/retryer/retryer.go b/pkg/filter/retryer/retryer.go
index 988f9a86..949887ed 100644
--- a/pkg/filter/retryer/retryer.go
+++ b/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("")
 
diff --git a/pkg/util/callbackreader/callbackreader.go b/pkg/util/callbackreader/callbackreader.go
index acc9ee5e..69d452cf 100644
--- a/pkg/util/callbackreader/callbackreader.go
+++ b/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
+}


Test patch:
diff --git a/pkg/context/contexttest/request.go b/pkg/context/contexttest/request.go
index aac1a914..0b4a2dc5 100644
--- a/pkg/context/contexttest/request.go
+++ b/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)
 	}
 }