Repository: cloudwego/hertz
Pull request #515 (2023-01-05 05:48:49)

fix: client cannot parse stream body correctly when transfer-encoding is 'identity'
#### What type of PR is this?
<!--
Add one of the following kinds:

build: Changes that affect the build system or external dependencies (example scopes: gulp, broccoli, npm)
ci: Changes to our CI configuration files and scripts (example scopes: Travis, Circle, BrowserStack, SauceLabs)
docs: Documentation only changes
feat: A new feature
optimize: A new optimization
fix: A bug fix
perf: A code change that improves performance
refactor: A code change that neither fixes a bug nor adds a feature
style: Changes that do not affect the meaning of the code (white space, formatting, missing semi-colons, etc)
test: Adding missing tests or correcting existing tests
chore: Changes to the build process or auxiliary tools and libraries such as documentation generation
-->
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.


#### (Optional) Translate the PR title into Chinese.
Hertz client 当 response 的 transfer-encoding 为 identity 时，不能正确解析stream body

#### (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:
zh(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)`.
-->


codecov[bot]: # [Codecov](https://codecov.io/gh/cloudwego/hertz/pull/515?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego) Report
Base: **69.34**% // Head: **69.72**% // Increases project coverage by **`+0.37%`** :tada:
> Coverage data is based on head [(`aed75d7`)](https://codecov.io/gh/cloudwego/hertz/pull/515?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego) compared to base [(`81332e6`)](https://codecov.io/gh/cloudwego/hertz/commit/81332e6a707d26302400b66373223078bc05f9b0?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego).
> Patch coverage: 0.00% of modified lines in pull request are covered.

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


```diff
@@             Coverage Diff             @@
##           develop     #515      +/-   ##
===========================================
+ Coverage    69.34%   69.72%   +0.37%     
===========================================
  Files           93       93              
  Lines         8874     8878       +4     
===========================================
+ Hits          6154     6190      +36     
+ Misses        2355     2318      -37     
- Partials       365      370       +5     
```


| [Impacted Files](https://codecov.io/gh/cloudwego/hertz/pull/515?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego) | Coverage Δ | |
|---|---|---|
| [pkg/network/netpoll/connection.go](https://codecov.io/gh/cloudwego/hertz/pull/515?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego#diff-cGtnL25ldHdvcmsvbmV0cG9sbC9jb25uZWN0aW9uLmdv) | `73.33% <0.00%> (-7.16%)` | :arrow_down: |
| [pkg/protocol/http1/ext/stream.go](https://codecov.io/gh/cloudwego/hertz/pull/515?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego#diff-cGtnL3Byb3RvY29sL2h0dHAxL2V4dC9zdHJlYW0uZ28=) | `0.00% <0.00%> (ø)` | |
| [pkg/app/fs.go](https://codecov.io/gh/cloudwego/hertz/pull/515?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego#diff-cGtnL2FwcC9mcy5nbw==) | `68.65% <0.00%> (ø)` | |
| [pkg/app/context.go](https://codecov.io/gh/cloudwego/hertz/pull/515?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego#diff-cGtnL2FwcC9jb250ZXh0Lmdv) | `84.30% <0.00%> (ø)` | |
| [pkg/protocol/header.go](https://codecov.io/gh/cloudwego/hertz/pull/515?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego#diff-cGtnL3Byb3RvY29sL2hlYWRlci5nbw==) | `74.93% <0.00%> (ø)` | |
| [pkg/protocol/request.go](https://codecov.io/gh/cloudwego/hertz/pull/515?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego#diff-cGtnL3Byb3RvY29sL3JlcXVlc3QuZ28=) | `83.24% <0.00%> (ø)` | |
| [pkg/app/client/client.go](https://codecov.io/gh/cloudwego/hertz/pull/515?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego#diff-cGtnL2FwcC9jbGllbnQvY2xpZW50Lmdv) | `74.39% <0.00%> (ø)` | |
| [pkg/common/ut/request.go](https://codecov.io/gh/cloudwego/hertz/pull/515?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego#diff-cGtnL2NvbW1vbi91dC9yZXF1ZXN0Lmdv) | `93.10% <0.00%> (ø)` | |
| [pkg/route/routergroup.go](https://codecov.io/gh/cloudwego/hertz/pull/515?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego#diff-cGtnL3JvdXRlL3JvdXRlcmdyb3VwLmdv) | `77.06% <0.00%> (ø)` | |
| [pkg/common/errors/errors.go](https://codecov.io/gh/cloudwego/hertz/pull/515?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego#diff-cGtnL2NvbW1vbi9lcnJvcnMvZXJyb3JzLmdv) | `87.95% <0.00%> (ø)` | |
| ... and [3 more](https://codecov.io/gh/cloudwego/hertz/pull/515?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego) | |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego)

</details>

[:umbrella: View full report at Codecov](https://codecov.io/gh/cloudwego/hertz/pull/515?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/cmd/hz/app/app.go b/cmd/hz/app/app.go
index 4261bd4..a66f624 100644
--- a/cmd/hz/app/app.go
+++ b/cmd/hz/app/app.go
@@ -141,7 +141,7 @@ func Init() *cli.App {
 	outDirFlag := cli.StringFlag{Name: "out_dir", Usage: "Specify the project path.", Destination: &globalArgs.OutDir}
 	handlerDirFlag := cli.StringFlag{Name: "handler_dir", Usage: "Specify the handler path.", Destination: &globalArgs.HandlerDir}
 	modelDirFlag := cli.StringFlag{Name: "model_dir", Usage: "Specify the model path.", Destination: &globalArgs.ModelDir}
-	clientDirFlag := cli.StringFlag{Name: "client_dir", Usage: "Specify the client path. If not specified, no client code is generated.", Destination: &globalArgs.ClientDir}
+	clientDirFlag := cli.StringFlag{Name: "client_dir", Usage: "Specify the client path. If not specified, IDL generated path is used for 'client' command; no client code is generated for 'new' command", Destination: &globalArgs.ClientDir}
 
 	optPkgFlag := cli.StringSliceFlag{Name: "option_package", Aliases: []string{"P"}, Usage: "Specify the package path. ({include_path}={import_path})"}
 	includesFlag := cli.StringSliceFlag{Name: "proto_path", Aliases: []string{"I"}, Usage: "Add an IDL search path for includes. (Valid only if idl is protobuf)"}
@@ -263,6 +263,7 @@ func Init() *cli.App {
 				&idlFlag,
 				&moduleFlag,
 				&modelDirFlag,
+				&clientDirFlag,
 
 				&includesFlag,
 				&thriftOptionsFlag,
diff --git a/cmd/hz/generator/client.go b/cmd/hz/generator/client.go
index 7ad2b1c..37b2cac 100644
--- a/cmd/hz/generator/client.go
+++ b/cmd/hz/generator/client.go
@@ -44,7 +44,7 @@ type ClientFile struct {
 func (pkgGen *HttpPackageGenerator) genClient(pkg *HttpPackage, clientDir string) error {
 	for _, s := range pkg.Services {
 		cliDir := util.SubDir(clientDir, util.ToSnakeCase(s.Name))
-		hertzClientPath := filepath.Join(cliDir, "hertz_client.go")
+		hertzClientPath := filepath.Join(cliDir, hertzClientTplName)
 		isExist, err := util.PathExist(hertzClientPath)
 		if err != nil {
 			return err
diff --git a/cmd/hz/generator/package.go b/cmd/hz/generator/package.go
index fa8a6a5..0476271 100644
--- a/cmd/hz/generator/package.go
+++ b/cmd/hz/generator/package.go
@@ -141,9 +141,12 @@ func (pkgGen *HttpPackageGenerator) Generate(pkg *HttpPackage) error {
 	}
 
 	if pkgGen.CmdType == meta.CmdClient {
+		// default client dir
 		clientDir := pkgGen.IdlClientDir
-		// support multiple service, so disable "clientDir/hertz"
-		// clientDir = util.SubDir(clientDir, "hertz")
+		// user specify client dir
+		if len(pkgGen.ClientDir) != 0 {
+			clientDir = pkgGen.ClientDir
+		}
 		if err := pkgGen.genClient(pkg, clientDir); err != nil {
 			return err
 		}
diff --git a/cmd/hz/generator/package_tpl.go b/cmd/hz/generator/package_tpl.go
index ac2197b..3eba61e 100644
--- a/cmd/hz/generator/package_tpl.go
+++ b/cmd/hz/generator/package_tpl.go
@@ -289,9 +289,12 @@ type use interface {
 
 // Definition of global data and types.
 type ResponseResultDecider func(statusCode int, rawResponse *protocol.Response) (isError bool)
-type bindRequestBodyFunc func(c *cli, r *request) (contentType string, body io.Reader, err error)
-type beforeRequestFunc func(*cli, *request) error
-type afterResponseFunc func(*cli, *response) error
+
+type (
+	bindRequestBodyFunc func(c *cli, r *request) (contentType string, body io.Reader, err error)
+	beforeRequestFunc   func(*cli, *request) error
+	afterResponseFunc   func(*cli, *response) error
+)
 
 var (
 	hdrContentTypeKey     = http.CanonicalHeaderKey("Content-Type")
@@ -906,11 +909,13 @@ func (s *{{$.ServiceName}}Client) {{$MethodInfo.Name}}(context context.Context,
 		setRequestOption(reqOpt...).
 		setResult(httpResp).
 		execute("{{$MethodInfo.HTTPMethod}}", "{{$MethodInfo.Path}}")
-	if err == nil {
-		resp = httpResp
+	if err != nil {
+		return nil, nil, err
 	}
+    
+	resp = httpResp
 	rawResponse = ret.rawResponse
-	return resp, rawResponse, err
+	return resp, rawResponse, nil
 }
 {{end}}
 
diff --git a/pkg/app/server/option.go b/pkg/app/server/option.go
index 7395775..bf44829 100644
--- a/pkg/app/server/option.go
+++ b/pkg/app/server/option.go
@@ -333,7 +333,7 @@ func WithDisablePrintRoute(b bool) config.Option {
 
 // WithOnAccept sets the callback function when a new connection is accepted but cannot
 // receive data in netpoll. In go net, it will be called before converting tls connection
-func WithOnAccept(fn func(conn network.Conn) context.Context) config.Option {
+func WithOnAccept(fn func(conn net.Conn) context.Context) config.Option {
 	return config.Option{F: func(o *config.Options) {
 		o.OnAccept = fn
 	}}
diff --git a/pkg/common/config/option.go b/pkg/common/config/option.go
index 3ba3d38..03c84a0 100644
--- a/pkg/common/config/option.go
+++ b/pkg/common/config/option.go
@@ -83,7 +83,7 @@ type Options struct {
 	// If you'd like to check whether the peer IP is in the blacklist, you can use OnAccept.
 	// In go net, OnAccept is executed after connection accepted but before establishing
 	// tls connection. OnConnect is executed after establishing tls connection.
-	OnAccept  func(conn network.Conn) context.Context
+	OnAccept  func(conn net.Conn) context.Context
 	OnConnect func(ctx context.Context, conn network.Conn) context.Context
 
 	// Registry is used for service registry.
diff --git a/pkg/network/netpoll/connection.go b/pkg/network/netpoll/connection.go
index 33d2875..683f9c3 100644
--- a/pkg/network/netpoll/connection.go
+++ b/pkg/network/netpoll/connection.go
@@ -49,6 +49,12 @@ func (c *Conn) Peek(n int) (b []byte, err error) {
 	return
 }
 
+func (c *Conn) Read(p []byte) (int, error) {
+	n, err := c.Conn.Read(p)
+	err = normalizeErr(err)
+	return n, err
+}
+
 func (c *Conn) Skip(n int) error {
 	return c.Conn.Skip(n)
 }
diff --git a/pkg/network/netpoll/transport.go b/pkg/network/netpoll/transport.go
index da28bf2..c213060 100644
--- a/pkg/network/netpoll/transport.go
+++ b/pkg/network/netpoll/transport.go
@@ -41,7 +41,7 @@ type transporter struct {
 	listener         net.Listener
 	eventLoop        netpoll.EventLoop
 	listenConfig     *net.ListenConfig
-	OnAccept         func(conn network.Conn) context.Context
+	OnAccept         func(conn net.Conn) context.Context
 	OnConnect        func(ctx context.Context, conn network.Conn) context.Context
 }
 
diff --git a/pkg/network/standard/transport.go b/pkg/network/standard/transport.go
index b6131b2..2605b2e 100644
--- a/pkg/network/standard/transport.go
+++ b/pkg/network/standard/transport.go
@@ -46,7 +46,7 @@ type transport struct {
 	tls              *tls.Config
 	listenConfig     *net.ListenConfig
 	lock             sync.Mutex
-	OnAccept         func(conn network.Conn) context.Context
+	OnAccept         func(conn net.Conn) context.Context
 	OnConnect        func(ctx context.Context, conn network.Conn) context.Context
 }
 
@@ -73,7 +73,7 @@ func (t *transport) serve() (err error) {
 		}
 
 		if t.OnAccept != nil {
-			ctx = t.OnAccept(c)
+			ctx = t.OnAccept(conn)
 		}
 
 		if t.tls != nil {
diff --git a/pkg/protocol/http1/ext/stream.go b/pkg/protocol/http1/ext/stream.go
index 268a137..d72e995 100644
--- a/pkg/protocol/http1/ext/stream.go
+++ b/pkg/protocol/http1/ext/stream.go
@@ -211,7 +211,7 @@ func (rs *bodyStream) Read(p []byte) (int, error) {
 	if err != nil {
 		// the data on stream may be incomplete
 		if err == io.EOF {
-			if rs.offset != rs.contentLength {
+			if rs.offset != rs.contentLength && rs.contentLength != -2 {
 				err = io.ErrUnexpectedEOF
 			}
 			// ensure that skipRest works fine


Test patch:
diff --git a/pkg/app/server/hertz_test.go b/pkg/app/server/hertz_test.go
index c5c1c8e..de6a593 100644
--- a/pkg/app/server/hertz_test.go
+++ b/pkg/app/server/hertz_test.go
@@ -803,7 +803,7 @@ func TestOnprepare(t *testing.T) {
 	assert.DeepEqual(t, "the server closed connection before returning the first response byte. Make sure the server returns 'Connection: close' response header before closing the connection", err.Error())
 
 	h = New(
-		WithOnAccept(func(conn network.Conn) context.Context {
+		WithOnAccept(func(conn net.Conn) context.Context {
 			conn.Close()
 			return context.Background()
 		}),
@@ -817,4 +817,18 @@ func TestOnprepare(t *testing.T) {
 	if err == nil {
 		t.Fatalf("err should not be nil")
 	}
+
+	h = New(
+		WithOnAccept(func(conn net.Conn) context.Context {
+			assert.DeepEqual(t, conn.LocalAddr().String(), "127.0.0.1:9231")
+			return context.Background()
+		}),
+		WithHostPorts("localhost:9231"),
+		WithTransport(standard.NewTransporter))
+	h.GET("/ping", func(ctx context.Context, c *app.RequestContext) {
+		c.JSON(consts.StatusOK, utils.H{"ping": "pong"})
+	})
+	go h.Spin()
+	time.Sleep(time.Second)
+	c.Get(context.Background(), nil, "http://127.0.0.1:9231/ping")
 }
diff --git a/pkg/protocol/http1/resp/response_test.go b/pkg/protocol/http1/resp/response_test.go
index 1b4eefb..cf45f65 100644
--- a/pkg/protocol/http1/resp/response_test.go
+++ b/pkg/protocol/http1/resp/response_test.go
@@ -521,3 +521,67 @@ func testResponseReadLimitBodySuccess(t *testing.T, s string, maxBodySize int) {
 		t.Fatalf("unexpected error: %s. s=%q, maxBodySize=%d", err, s, maxBodySize)
 	}
 }
+
+func testResponseReadBodyStreamSuccess(t *testing.T, resp *protocol.Response, response string, expectedStatusCode, expectedContentLength int,
+	expectedContentType, expectedBody, expectedTrailer string,
+) {
+	zr := mock.NewZeroCopyReader(response)
+	err := ReadBodyStream(resp, zr, 0, nil)
+	if err != nil {
+		t.Fatalf("Unexpected error: %s", err)
+	}
+	assert.True(t, resp.IsBodyStream())
+
+	body, err := io.ReadAll(resp.BodyStream())
+	if err != nil && err != io.EOF {
+		t.Fatalf("Unexpected error: %s", err)
+	}
+	verifyResponseHeader(t, &resp.Header, expectedStatusCode, expectedContentLength, expectedContentType, "")
+	if !bytes.Equal(body, []byte(expectedBody)) {
+		t.Fatalf("Unexpected body %q. Expected %q", resp.Body(), []byte(expectedBody))
+	}
+	assert.VerifyTrailer(t, zr, expectedTrailer)
+}
+
+func TestResponseReadBodyStream(t *testing.T) {
+	t.Parallel()
+
+	resp := &protocol.Response{}
+
+	// usual response
+	testResponseReadBodyStreamSuccess(t, resp, "HTTP/1.1 200 OK\r\nContent-Length: 10\r\nContent-Type: foo/bar\r\n\r\n0123456789",
+		consts.StatusOK, 10, "foo/bar", "0123456789", "")
+
+	// zero response
+	testResponseReadBodyStreamSuccess(t, resp, "HTTP/1.1 500 OK\r\nContent-Length: 0\r\nContent-Type: foo/bar\r\n\r\n",
+		consts.StatusInternalServerError, 0, "foo/bar", "", "")
+
+	// response with trailer
+	testResponseReadBodyStreamSuccess(t, resp, "HTTP/1.1 300 OK\r\nContent-Length: 5\r\nContent-Type: bar\r\n\r\n56789aaa",
+		consts.StatusMultipleChoices, 5, "bar", "56789", "aaa")
+
+	// no content-length ('identity' transfer-encoding)
+	testResponseReadBodyStreamSuccess(t, resp, "HTTP/1.1 200 OK\r\nContent-Type: foobar\r\n\r\nzxxc",
+		consts.StatusOK, -2, "foobar", "zxxc", "")
+
+	// explicitly stated 'Transfer-Encoding: identity'
+	testResponseReadBodyStreamSuccess(t, resp, "HTTP/1.1 234 ss\r\nContent-Type: xxx\r\n\r\nxag",
+		234, -2, "xxx", "xag", "")
+
+	// big 'identity' response
+	body := string(mock.CreateFixedBody(100500))
+	testResponseReadBodyStreamSuccess(t, resp, "HTTP/1.1 200 OK\r\nContent-Type: aa\r\n\r\n"+body,
+		consts.StatusOK, -2, "aa", body, "")
+
+	// chunked response
+	testResponseReadBodyStreamSuccess(t, resp, "HTTP/1.1 200 OK\r\nContent-Type: text/html\r\nTransfer-Encoding: chunked\r\n\r\n4\r\nqwer\r\n2\r\nty\r\n0\r\n\r\nzzzzz",
+		consts.StatusOK, -1, "text/html", "qwerty", "zzzzz")
+
+	// chunked response with non-chunked Transfer-Encoding.
+	testResponseReadBodyStreamSuccess(t, resp, "HTTP/1.1 230 OK\r\nContent-Type: text\r\nTransfer-Encoding: aaabbb\r\n\r\n2\r\ner\r\n2\r\nty\r\n0\r\n\r\nwe",
+		230, -1, "text", "erty", "we")
+
+	// zero chunked response
+	testResponseReadBodyStreamSuccess(t, resp, "HTTP/1.1 200 OK\r\nContent-Type: text/html\r\nTransfer-Encoding: chunked\r\n\r\n0\r\n\r\nzzz",
+		consts.StatusOK, -1, "text/html", "", "zzz")
+}