Repository: cloudwego/hertz
Pull request #936 (2023-09-06 11:49:36)

fix(client): conn leak in stream mode when resp is no content(204)
#### 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.
修复在流式读响应模式下如果响应体是无body（204状态码）类型连接泄露的问题

#### (Optional) More detailed description for this PR(en: English/zh: Chinese).
<!--
Provide more detailed info for review(e.g., it's recommended to provide perf data if this is a perf type PR).
-->
en:
The closing logic timing of the streaming and non-streaming connections is different. In Non-streaming scene the conn of resp will be closed/released after the parsing of resp is completed. The streaming scene needs to operate the connection due to the user's post-processing logic, so it is controlled through a callback form (callback supports active and Passive gc finalizer triggers), but the scene of 204 just fell in the middle of these two processing logics, so the callback was not hung up, and it was not closed directly.
zh(optional):
流式和非流式的连接关闭逻辑时机不同，非流式解完resp就关闭了，流式场景由于用户后置逻辑需要操作连接，因此是通过一个callback的形态来控制的（callback支持主动和被动gc finalizer触发），但是204这个场景正好掉到了这两种处理逻辑中间，导致没有挂上callback，也没有直接关闭。

#### (Optional) Which issue(s) this PR fixes:
<!--
Automatically closes linked issue when PR is merged.
Eg: `Fixes #<issue number>`, or `Fixes (paste link of issue)`.
-->

#### (Optional) The PR that updates user documentation:
<!--
If the current PR requires user awareness at the usage level, please submit a PR to update user docs. [User docs repo](https://github.com/cloudwego/cloudwego.github.io)
-->


Patch:
diff --git a/pkg/protocol/http1/client.go b/pkg/protocol/http1/client.go
index c2b7cf5..09fbc10 100644
--- a/pkg/protocol/http1/client.go
+++ b/pkg/protocol/http1/client.go
@@ -72,9 +72,12 @@ import (
 	respI "github.com/cloudwego/hertz/pkg/protocol/http1/resp"
 )
 
-var errConnectionClosed = errs.NewPublic("the server closed connection before returning the first response byte. " +
-	"Make sure the server returns 'Connection: close' response header before closing the connection")
-var errTimeout = errs.New(errs.ErrTimeout, errs.ErrorTypePublic, "host client")
+var (
+	errConnectionClosed = errs.NewPublic("the server closed connection before returning the first response byte. " +
+		"Make sure the server returns 'Connection: close' response header before closing the connection")
+
+	errTimeout = errs.New(errs.ErrTimeout, errs.ErrorTypePublic, "host client")
+)
 
 // HostClient balances http requests among hosts listed in Addr.
 //
@@ -699,7 +702,8 @@ func (c *HostClient) doNonNilReqResp(req *protocol.Request, resp *protocol.Respo
 
 	shouldCloseConn = resetConnection || req.ConnectionClose() || resp.ConnectionClose()
 
-	if c.ResponseBodyStream {
+	// In stream mode, we still can close/release the connection immediately if there is no content on the wire.
+	if c.ResponseBodyStream && resp.BodyStream() != protocol.NoResponseBody {
 		return false, err
 	}
 
diff --git a/pkg/protocol/request.go b/pkg/protocol/request.go
index ef70f92..d047318 100644
--- a/pkg/protocol/request.go
+++ b/pkg/protocol/request.go
@@ -76,7 +76,6 @@ var (
 // NoBody is an io.ReadCloser with no bytes. Read always returns EOF
 // and Close always returns nil. It can be used in an outgoing client
 // request to explicitly signal that a request has zero bytes.
-// An alternative, however, is to simply set Request.Body to nil.
 var NoBody = noBody{}
 
 type noBody struct{}
diff --git a/pkg/protocol/response.go b/pkg/protocol/response.go
index 016af92..8beb385 100644
--- a/pkg/protocol/response.go
+++ b/pkg/protocol/response.go
@@ -54,7 +54,13 @@ import (
 	"github.com/cloudwego/hertz/pkg/network"
 )
 
-var responsePool sync.Pool
+var (
+	responsePool sync.Pool
+	// NoResponseBody is an io.ReadCloser with no bytes. Read always returns EOF
+	// and Close always returns nil. It can be used in an ingoing client
+	// response to explicitly signal that a response has zero bytes.
+	NoResponseBody = noBody{}
+)
 
 // Response represents HTTP response.
 //
@@ -334,6 +340,9 @@ func (resp *Response) SetBody(body []byte) {
 }
 
 func (resp *Response) BodyStream() io.Reader {
+	if resp.bodyStream == nil {
+		resp.bodyStream = NoResponseBody
+	}
 	return resp.bodyStream
 }
 


Test patch:
diff --git a/pkg/protocol/http1/client_test.go b/pkg/protocol/http1/client_test.go
index 7bc389b..7c45d14 100644
--- a/pkg/protocol/http1/client_test.go
+++ b/pkg/protocol/http1/client_test.go
@@ -536,3 +536,43 @@ func TestConnNotRetry(t *testing.T) {
 	assert.True(t, logbuf.String() == "")
 	protocol.ReleaseResponse(resp)
 }
+
+type countCloseConn struct {
+	network.Conn
+	isClose bool
+}
+
+func (c *countCloseConn) Close() error {
+	c.isClose = true
+	return nil
+}
+
+func newCountCloseConn(s string) *countCloseConn {
+	return &countCloseConn{
+		Conn: mock.NewConn(s),
+	}
+}
+
+func TestStreamNoContent(t *testing.T) {
+	conn := newCountCloseConn("HTTP/1.1 204 Foo Bar\r\nContent-Type: aab\r\nTrailer: Foo\r\nContent-Encoding: deflate\r\nTransfer-Encoding: chunked\r\n\r\n0\r\nFoo: bar\r\n\r\nHTTP/1.2")
+
+	c := &HostClient{
+		ClientOptions: &ClientOptions{
+			Dialer: newSlowConnDialer(func(network, addr string) (network.Conn, error) {
+				return conn, nil
+			}),
+		},
+		Addr: "foobar",
+	}
+
+	c.ResponseBodyStream = true
+
+	req := protocol.AcquireRequest()
+	req.SetRequestURI("http://foobar/baz")
+	req.Header.SetConnectionClose(true)
+	resp := protocol.AcquireResponse()
+
+	c.Do(context.Background(), req, resp)
+
+	assert.True(t, conn.isClose)
+}