Repository: cloudwego/hertz
Pull request #815 (2023-06-14 08:47:57)

fix: close abnormal conn when release
#### What type of PR is this?
fix

#### Check the PR title.
<!--
The description of the title will be attached in Release Notes, 
so please describe it from user-oriented, what this PR does / why we need it.
Please check your PR title with the below requirements:
-->
- [x] This PR title match the format: `<type>(optional scope): <description>`.
- [x] The description of this PR title is user-oriented and clear enough for others to understand.
- [ ] Attach the PR updating the user documentation if the current PR requires user awareness at the usage level. [User docs repo](https://github.com/cloudwego/cloudwego.github.io).

#### (Optional) Translate the PR title into Chinese.
关闭 release body stream 阶段出现异常的链接


#### (Optional) More detail description for this PR(en: English/zh: Chinese).
<!--
Provide more detailed info for review. If it is a perf type PR, perf data is suggested to give.
-->
en: close connection if error happened when release the body stream, otherwise possible dirty data will affect subsequent requests for this connection.
zh(optional): 释放body流时如果出错则关闭连接，否则，可能的脏数据会影响这个连接后续的请求。

#### (Optional) Which issue(s) this PR fixes:
None

#### (Optional) The PR that updates user documentation:
None


codecov[bot]: ## [Codecov](https://app.codecov.io/gh/cloudwego/hertz/pull/815?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego) Report
Patch coverage: **`80.00`**% and project coverage change: **`+0.04`** :tada:
> Comparison is base [(`cdcafeb`)](https://app.codecov.io/gh/cloudwego/hertz/commit/cdcafeb7183a81cd6537a4bdde5af534340946b1?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego) 75.63% compared to head [(`6e6db88`)](https://app.codecov.io/gh/cloudwego/hertz/pull/815?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego) 75.67%.

> :exclamation: Current head 6e6db88 differs from pull request most recent head 0aa5981. Consider uploading reports for the commit 0aa5981 to get more accurate results

<details><summary>Additional details and impacted files</summary>


```diff
@@             Coverage Diff             @@
##           develop     #815      +/-   ##
===========================================
+ Coverage    75.63%   75.67%   +0.04%     
===========================================
  Files           97       97              
  Lines         9480     9484       +4     
===========================================
+ Hits          7170     7177       +7     
+ Misses        1859     1856       -3     
  Partials       451      451              
```


| [Impacted Files](https://app.codecov.io/gh/cloudwego/hertz/pull/815?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego) | Coverage Δ | |
|---|---|---|
| [pkg/protocol/http1/ext/stream.go](https://app.codecov.io/gh/cloudwego/hertz/pull/815?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego#diff-cGtnL3Byb3RvY29sL2h0dHAxL2V4dC9zdHJlYW0uZ28=) | `14.37% <ø> (ø)` | |
| [pkg/protocol/http1/client.go](https://app.codecov.io/gh/cloudwego/hertz/pull/815?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego#diff-cGtnL3Byb3RvY29sL2h0dHAxL2NsaWVudC5nbw==) | `58.31% <50.00%> (ø)` | |
| [pkg/protocol/http1/resp/response.go](https://app.codecov.io/gh/cloudwego/hertz/pull/815?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego#diff-cGtnL3Byb3RvY29sL2h0dHAxL3Jlc3AvcmVzcG9uc2UuZ28=) | `71.94% <87.50%> (+0.83%)` | :arrow_up: |

... and [1 file with indirect coverage changes](https://app.codecov.io/gh/cloudwego/hertz/pull/815/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego)


</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/cloudwego/hertz/pull/815?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego).   
:loudspeaker: Do you have feedback about the report comment? [Let us know in this issue](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego).


Patch:
diff --git a/pkg/protocol/http1/client.go b/pkg/protocol/http1/client.go
index 353ecc8..13f07a1 100644
--- a/pkg/protocol/http1/client.go
+++ b/pkg/protocol/http1/client.go
@@ -628,8 +628,8 @@ func (c *HostClient) doNonNilReqResp(req *protocol.Request, resp *protocol.Respo
 	if !c.ResponseBodyStream {
 		err = respI.ReadHeaderAndLimitBody(resp, zr, c.MaxResponseBodySize)
 	} else {
-		err = respI.ReadBodyStream(resp, zr, c.MaxResponseBodySize, func() error {
-			if shouldCloseConn {
+		err = respI.ReadBodyStream(resp, zr, c.MaxResponseBodySize, func(shouldClose bool) error {
+			if shouldCloseConn || shouldClose {
 				c.closeConn(cc)
 			} else {
 				c.releaseConn(cc)
diff --git a/pkg/protocol/http1/ext/stream.go b/pkg/protocol/http1/ext/stream.go
index a09e1e3..3ff25d4 100644
--- a/pkg/protocol/http1/ext/stream.go
+++ b/pkg/protocol/http1/ext/stream.go
@@ -309,6 +309,7 @@ func (rs *bodyStream) skipRest() error {
 }
 
 // ReleaseBodyStream releases the body stream.
+// Error of skipRest may be returned if there is one.
 //
 // NOTE: Be careful to use this method unless you know what it's for.
 func ReleaseBodyStream(requestReader io.Reader) (err error) {
diff --git a/pkg/protocol/http1/resp/response.go b/pkg/protocol/http1/resp/response.go
index 8b39b0c..0366336 100644
--- a/pkg/protocol/http1/resp/response.go
+++ b/pkg/protocol/http1/resp/response.go
@@ -50,6 +50,7 @@ import (
 
 	"github.com/cloudwego/hertz/pkg/common/bytebufferpool"
 	errs "github.com/cloudwego/hertz/pkg/common/errors"
+	"github.com/cloudwego/hertz/pkg/common/hlog"
 	"github.com/cloudwego/hertz/pkg/network"
 	"github.com/cloudwego/hertz/pkg/protocol"
 	"github.com/cloudwego/hertz/pkg/protocol/consts"
@@ -128,14 +129,21 @@ func ReadHeaderAndLimitBody(resp *protocol.Response, r network.Reader, maxBodySi
 
 type clientRespStream struct {
 	r             io.Reader
-	closeCallback func() error
+	closeCallback func(shouldClose bool) error
 }
 
 func (c *clientRespStream) Close() (err error) {
 	runtime.SetFinalizer(c, nil)
-	ext.ReleaseBodyStream(c.r)
+	// If error happened in release, the connection may be in abnormal state.
+	// Close it in the callback in order to avoid other unexpected problems.
+	err = ext.ReleaseBodyStream(c.r)
+	shouldClose := false
+	if err != nil {
+		shouldClose = true
+		hlog.Warnf("connection will be closed instead of recycled because an error occurred during the stream body release: %s", err.Error())
+	}
 	if c.closeCallback != nil {
-		err = c.closeCallback()
+		err = c.closeCallback(shouldClose)
 	}
 	c.reset()
 	return
@@ -157,7 +165,7 @@ var clientRespStreamPool = sync.Pool{
 	},
 }
 
-func convertClientRespStream(bs io.Reader, fn func() error) *clientRespStream {
+func convertClientRespStream(bs io.Reader, fn func(shouldClose bool) error) *clientRespStream {
 	clientStream := clientRespStreamPool.Get().(*clientRespStream)
 	clientStream.r = bs
 	clientStream.closeCallback = fn
@@ -166,7 +174,7 @@ func convertClientRespStream(bs io.Reader, fn func() error) *clientRespStream {
 }
 
 // ReadBodyStream reads response body in stream
-func ReadBodyStream(resp *protocol.Response, r network.Reader, maxBodySize int, closeCallBack func() error) error {
+func ReadBodyStream(resp *protocol.Response, r network.Reader, maxBodySize int, closeCallBack func(shouldClose bool) error) error {
 	resp.ResetBody()
 	err := ReadHeader(&resp.Header, r)
 	if err != nil {


Test patch:
diff --git a/pkg/common/test/mock/network.go b/pkg/common/test/mock/network.go
index 1038eb7..3d47475 100644
--- a/pkg/common/test/mock/network.go
+++ b/pkg/common/test/mock/network.go
@@ -161,6 +161,22 @@ func NewSlowReadConn(source string) *SlowReadConn {
 	return &SlowReadConn{Conn: NewConn(source)}
 }
 
+type ErrorReadConn struct {
+	*Conn
+	errorToReturn error
+}
+
+func NewErrorReadConn(err error) *ErrorReadConn {
+	return &ErrorReadConn{
+		Conn:          NewConn(""),
+		errorToReturn: err,
+	}
+}
+
+func (er *ErrorReadConn) Peek(n int) ([]byte, error) {
+	return nil, er.errorToReturn
+}
+
 type SlowWriteConn struct {
 	*Conn
 	writeTimeout time.Duration
diff --git a/pkg/protocol/http1/resp/response_test.go b/pkg/protocol/http1/resp/response_test.go
index 7178971..cff0a57 100644
--- a/pkg/protocol/http1/resp/response_test.go
+++ b/pkg/protocol/http1/resp/response_test.go
@@ -57,6 +57,7 @@ import (
 	"github.com/cloudwego/hertz/pkg/common/utils"
 	"github.com/cloudwego/hertz/pkg/protocol"
 	"github.com/cloudwego/hertz/pkg/protocol/consts"
+	"github.com/cloudwego/hertz/pkg/protocol/http1/ext"
 	"github.com/cloudwego/netpoll"
 )
 
@@ -614,6 +615,27 @@ func testResponseBodyStreamWithTrailer(t *testing.T, body []byte, disableNormali
 	}
 }
 
+func TestResponseReadBodyStreamBadReader(t *testing.T) {
+	t.Parallel()
+
+	resp := protocol.AcquireResponse()
+
+	errReader := mock.NewErrorReadConn(errors.New("test error"))
+
+	bodyBuf := resp.BodyBuffer()
+	bodyBuf.Reset()
+
+	bodyStream := ext.AcquireBodyStream(bodyBuf, errReader, resp.Header.Trailer(), 100)
+	resp.ConstructBodyStream(bodyBuf, convertClientRespStream(bodyStream, func(shouldClose bool) error {
+		assert.True(t, shouldClose)
+		return nil
+	}))
+
+	stBody := resp.BodyStream()
+	closer, _ := stBody.(io.Closer)
+	closer.Close()
+}
+
 func TestSetResponseBodyStreamFixedSize(t *testing.T) {
 	t.Parallel()