Repository: cloudwego/hertz
Pull request #530 (2023-01-10 18:19:11)

optimize(hz): refuse all new operation if there already exists a project
#### What type of PR is this?

optimize

#### 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.

- 项目存在时拒绝使用任何 new 操作


#### Which issue(s) this PR fixes:

Fixes https://github.com/cloudwego/hertz/issues/529

<!--
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/530?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego) Report
Base: **69.75**% // Head: **69.87**% // Increases project coverage by **`+0.12%`** :tada:
> Coverage data is based on head [(`7752d13`)](https://codecov.io/gh/cloudwego/hertz/pull/530?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego) compared to base [(`d469803`)](https://codecov.io/gh/cloudwego/hertz/commit/d4698036dfd4a1700d403713c8e6495f54676307?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego).
> Patch coverage: 100.00% of modified lines in pull request are covered.

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


```diff
@@             Coverage Diff             @@
##           develop     #530      +/-   ##
===========================================
+ Coverage    69.75%   69.87%   +0.12%     
===========================================
  Files           93       93              
  Lines         8878     8887       +9     
===========================================
+ Hits          6193     6210      +17     
+ Misses        2315     2307       -8     
  Partials       370      370              
```


| [Impacted Files](https://codecov.io/gh/cloudwego/hertz/pull/530?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego) | Coverage Δ | |
|---|---|---|
| [pkg/protocol/header.go](https://codecov.io/gh/cloudwego/hertz/pull/530?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego#diff-cGtnL3Byb3RvY29sL2hlYWRlci5nbw==) | `75.59% <100.00%> (+0.65%)` | :arrow_up: |
| [pkg/protocol/cookie.go](https://codecov.io/gh/cloudwego/hertz/pull/530?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego#diff-cGtnL3Byb3RvY29sL2Nvb2tpZS5nbw==) | `88.47% <0.00%> (+2.05%)` | :arrow_up: |

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/530?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).

li-jin-gou: cc @welkeyever @Duslia @wzekin 

Patch:
diff --git a/cmd/hz/app/app.go b/cmd/hz/app/app.go
index 2069d55..abf7c07 100644
--- a/cmd/hz/app/app.go
+++ b/cmd/hz/app/app.go
@@ -44,6 +44,15 @@ func New(c *cli.Context) error {
 	setLogVerbose(args.Verbose)
 	logs.Debugf("args: %#v\n", args)
 
+	exist, err := util.PathExist(filepath.Join(args.OutDir, meta.ManifestFile))
+	if err != nil {
+		return cli.Exit(err, meta.LoadError)
+	}
+
+	if exist && !args.ForceNew {
+		return cli.Exit(fmt.Errorf("the current is already a hertz project, if you want to regenerate it you can specify \"-force\""), meta.LoadError)
+	}
+
 	err = GenerateLayout(args)
 	if err != nil {
 		return cli.Exit(err, meta.GenerateLayoutError)
@@ -163,6 +172,7 @@ func Init() *cli.App {
 	thriftPluginsFlag := cli.StringSliceFlag{Name: "thrift-plugins", Usage: "Specify plugins for the thriftgo. ({plugin_name}:{options})"}
 	protoPluginsFlag := cli.StringSliceFlag{Name: "protoc-plugins", Usage: "Specify plugins for the protoc. ({plugin_name}:{options}:{out_dir})"}
 	noRecurseFlag := cli.BoolFlag{Name: "no_recurse", Usage: "Generate master model only.", Destination: &globalArgs.NoRecurse}
+	forceNewFlag := cli.BoolFlag{Name: "force", Aliases: []string{"f"}, Usage: "Force new a project, which will overwrite the generated files", Destination: &globalArgs.ForceNew}
 
 	jsonEnumStrFlag := cli.BoolFlag{Name: "json_enumstr", Usage: "Use string instead of num for json enums when idl is thrift.", Destination: &globalArgs.JSONEnumStr}
 	unsetOmitemptyFlag := cli.BoolFlag{Name: "unset_omitempty", Usage: "Remove 'omitempty' tag for generated struct.", Destination: &globalArgs.UnsetOmitempty}
@@ -205,6 +215,7 @@ func Init() *cli.App {
 				&protoOptionsFlag,
 				&optPkgFlag,
 				&noRecurseFlag,
+				&forceNewFlag,
 
 				&jsonEnumStrFlag,
 				&unsetOmitemptyFlag,
diff --git a/cmd/hz/config/argument.go b/cmd/hz/config/argument.go
index 253b2e2..4ed71f5 100644
--- a/cmd/hz/config/argument.go
+++ b/cmd/hz/config/argument.go
@@ -64,6 +64,7 @@ type Argument struct {
 	Excludes             []string
 	NoRecurse            bool
 	HandlerByMethod      bool
+	ForceNew             bool
 
 	CustomizeLayout  string
 	CustomizePackage string
diff --git a/pkg/app/context.go b/pkg/app/context.go
index 157674c..e86cc4b 100644
--- a/pkg/app/context.go
+++ b/pkg/app/context.go
@@ -78,22 +78,72 @@ type Handler interface {
 
 type ClientIP func(ctx *RequestContext) string
 
-var defaultClientIP = func(ctx *RequestContext) string {
-	RemoteIPHeaders := []string{"X-Real-IP", "X-Forwarded-For"}
-	for _, headerName := range RemoteIPHeaders {
-		ip := ctx.Request.Header.Get(headerName)
-		if ip != "" {
-			return ip
+type ClientIPOptions struct {
+	RemoteIPHeaders []string
+	TrustedProxies  map[string]bool
+}
+
+var defaultClientIPOptions = ClientIPOptions{
+	RemoteIPHeaders: []string{"X-Real-IP", "X-Forwarded-For"},
+	TrustedProxies: map[string]bool{
+		"0.0.0.0": true,
+	},
+}
+
+// ClientIPWithOption used to generate custom ClientIP function and set by engine.SetClientIPFunc
+func ClientIPWithOption(opts ClientIPOptions) ClientIP {
+	return func(ctx *RequestContext) string {
+		RemoteIPHeaders := opts.RemoteIPHeaders
+		TrustedProxies := opts.TrustedProxies
+
+		remoteIP, _, err := net.SplitHostPort(strings.TrimSpace(ctx.RemoteAddr().String()))
+		if err != nil {
+			return ""
+		}
+		trusted := isTrustedProxy(TrustedProxies, remoteIP)
+
+		if trusted {
+			for _, headerName := range RemoteIPHeaders {
+				ip, valid := validateHeader(TrustedProxies, ctx.Request.Header.Get(headerName))
+				if valid {
+					return ip
+				}
+			}
 		}
+
+		return remoteIP
 	}
+}
 
-	if ip, _, err := net.SplitHostPort(strings.TrimSpace(ctx.RemoteAddr().String())); err == nil {
-		return ip
+// isTrustedProxy will check whether the IP address is included in the trusted list according to TrustedProxies
+func isTrustedProxy(trustedProxies map[string]bool, remoteIP string) bool {
+	return trustedProxies[remoteIP]
+}
+
+// validateHeader will parse X-Real-IP and X-Forwarded-For header and return the Initial client IP address or an untrusted IP address
+func validateHeader(trustedProxies map[string]bool, header string) (clientIP string, valid bool) {
+	if header == "" {
+		return "", false
 	}
+	items := strings.Split(header, ",")
+	for i := len(items) - 1; i >= 0; i-- {
+		ipStr := strings.TrimSpace(items[i])
+		ip := net.ParseIP(ipStr)
+		if ip == nil {
+			break
+		}
 
-	return ""
+		// X-Forwarded-For is appended by proxy
+		// Check IPs in reverse order and stop when find untrusted proxy
+		if (i == 0) || (!isTrustedProxy(trustedProxies, ipStr)) {
+			return ipStr, true
+		}
+	}
+	return "", false
 }
 
+var defaultClientIP = ClientIPWithOption(defaultClientIPOptions)
+
 // SetClientIPFunc sets ClientIP function implementation to get ClientIP.
 // Deprecated: Use engine.SetClientIPFunc instead of SetClientIPFunc
 func SetClientIPFunc(fn ClientIP) {
diff --git a/pkg/protocol/http1/ext/common.go b/pkg/protocol/http1/ext/common.go
index 7d8498c..94526b7 100644
--- a/pkg/protocol/http1/ext/common.go
+++ b/pkg/protocol/http1/ext/common.go
@@ -387,6 +387,70 @@ func stripSpace(b []byte) []byte {
 	return b
 }
 
+func SkipTrailer(r network.Reader) error {
+	n := 1
+	for {
+		err := trySkipTrailer(r, n)
+		if err == nil {
+			return nil
+		}
+		if !errors.Is(err, errs.ErrNeedMore) {
+			return err
+		}
+		// No more data available on the wire, try block peek(by netpoll)
+		if n == r.Len() {
+			n++
+
+			continue
+		}
+		n = r.Len()
+	}
+}
+
+func trySkipTrailer(r network.Reader, n int) error {
+	b, err := r.Peek(n)
+	if len(b) == 0 {
+		// Return ErrTimeout on any timeout.
+		if err != nil && strings.Contains(err.Error(), "timeout") {
+			return errs.New(errs.ErrTimeout, errs.ErrorTypePublic, "read response header")
+		}
+
+		if n == 1 || err == io.EOF {
+			return io.EOF
+		}
+
+		return errs.NewPublicf("error when reading request trailer: %w", err)
+	}
+	b = MustPeekBuffered(r)
+	headersLen, errParse := skipTrailer(b)
+	if errParse != nil {
+		if err == io.EOF {
+			return err
+		}
+		return HeaderError("response", err, errParse, b)
+	}
+	MustDiscard(r, headersLen)
+	return nil
+}
+
+func skipTrailer(buf []byte) (int, error) {
+	skip := 0
+	strCRLFLen := len(bytestr.StrCRLF)
+	for {
+		index := bytes.Index(buf, bytestr.StrCRLF)
+		if index == -1 {
+			return 0, errs.ErrNeedMore
+		}
+
+		buf = buf[index+strCRLFLen:]
+		skip += index + strCRLFLen
+
+		if index == 0 {
+			return skip, nil
+		}
+	}
+}
+
 func ReadTrailer(t *protocol.Trailer, r network.Reader) error {
 	n := 1
 	for {
diff --git a/pkg/protocol/http1/ext/stream.go b/pkg/protocol/http1/ext/stream.go
index 2dd203d..a09e1e3 100644
--- a/pkg/protocol/http1/ext/stream.go
+++ b/pkg/protocol/http1/ext/stream.go
@@ -46,6 +46,7 @@ import (
 	"io"
 	"sync"
 
+	"github.com/cloudwego/hertz/internal/bytestr"
 	"github.com/cloudwego/hertz/pkg/common/bytebufferpool"
 	errs "github.com/cloudwego/hertz/pkg/common/errors"
 	"github.com/cloudwego/hertz/pkg/common/utils"
@@ -73,6 +74,8 @@ type bodyStream struct {
 	offset          int
 	contentLength   int
 	chunkLeft       int
+	// whether the chunk has reached the EOF
+	chunkEOF bool
 }
 
 func ReadBodyWithStreaming(zr network.Reader, contentLength, maxBodySize int, dst []byte) (b []byte, err error) {
@@ -114,6 +117,7 @@ func AcquireBodyStream(b *bytebufferpool.ByteBuffer, r network.Reader, t *protoc
 	rs.reader = r
 	rs.contentLength = contentLength
 	rs.trailer = t
+	rs.chunkEOF = false
 
 	return rs
 }
@@ -125,6 +129,10 @@ func (rs *bodyStream) Read(p []byte) (int, error) {
 		}
 	}()
 	if rs.contentLength == -1 {
+		if rs.chunkEOF {
+			return 0, io.EOF
+		}
+
 		if rs.chunkLeft == 0 {
 			chunkSize, err := utils.ParseChunkSize(rs.reader)
 			if err != nil {
@@ -133,6 +141,7 @@ func (rs *bodyStream) Read(p []byte) (int, error) {
 			if chunkSize == 0 {
 				err = ReadTrailer(rs.trailer, rs.reader)
 				if err == nil {
+					rs.chunkEOF = true
 					err = io.EOF
 				}
 				return 0, err
@@ -226,6 +235,45 @@ func (rs *bodyStream) skipRest() error {
 	if rs.prefetchedBytes == nil {
 		return nil
 	}
+
+	// the request is chunked encoding
+	if rs.contentLength == -1 {
+		if rs.chunkEOF {
+			return nil
+		}
+
+		strCRLFLen := len(bytestr.StrCRLF)
+		for {
+			chunkSize, err := utils.ParseChunkSize(rs.reader)
+			if err != nil {
+				return err
+			}
+
+			if chunkSize == 0 {
+				rs.chunkEOF = true
+				return SkipTrailer(rs.reader)
+			}
+
+			err = rs.reader.Skip(chunkSize)
+			if err != nil {
+				return err
+			}
+
+			crlf, err := rs.reader.Peek(strCRLFLen)
+			if err != nil {
+				return err
+			}
+
+			if !bytes.Equal(crlf, bytestr.StrCRLF) {
+				return errBrokenChunk
+			}
+
+			err = rs.reader.Skip(strCRLFLen)
+			if err != nil {
+				return err
+			}
+		}
+	}
 	// max value of pSize is 8193, it's safe.
 	pSize := int(rs.prefetchedBytes.Size())
 	if rs.contentLength <= pSize || rs.offset == rs.contentLength {


Test patch:
diff --git a/pkg/app/context_test.go b/pkg/app/context_test.go
index 4e62a94..31d6da0 100644
--- a/pkg/app/context_test.go
+++ b/pkg/app/context_test.go
@@ -673,18 +673,48 @@ func TestContextContentType(t *testing.T) {
 func TestClientIp(t *testing.T) {
 	c := NewContext(0)
 	c.conn = mock.NewConn("")
-	// Case 1
-	c.Request.Header.Set("X-Forwarded-For", "126.0.0.2")
+	// 0.0.0.0 simulates a trusted proxy server
+	c.Request.Header.Set("X-Forwarded-For", "  126.0.0.2, 0.0.0.0 ")
 	val := c.ClientIP()
 	if val != "126.0.0.2" {
 		t.Fatalf("unexpected %v. Expecting %v", val, "126.0.0.2")
 	}
-	// Case 2
+	// no proxy server
+	c = NewContext(0)
+	c.conn = mock.NewConn("")
 	c.Request.Header.Set("X-Real-Ip", "126.0.0.1")
 	val = c.ClientIP()
 	if val != "126.0.0.1" {
 		t.Fatalf("unexpected %v. Expecting %v", val, "126.0.0.1")
 	}
+	// custom RemoteIPHeaders and TrustedProxies
+	opts := ClientIPOptions{
+		RemoteIPHeaders: []string{"X-Forwarded-For", "X-Real-IP"},
+		TrustedProxies: map[string]bool{
+			"0.0.0.0": true,
+		},
+	}
+	c = NewContext(0)
+	c.SetClientIPFunc(ClientIPWithOption(opts))
+	c.conn = mock.NewConn("")
+	c.Request.Header.Set("X-Forwarded-For", "  126.0.0.2, 0.0.0.0 ")
+	val = c.ClientIP()
+	if val != "126.0.0.2" {
+		t.Fatalf("unexpected %v. Expecting %v", val, "126.0.0.2")
+	}
+	// no trusted proxy server
+	opts = ClientIPOptions{
+		RemoteIPHeaders: []string{"X-Forwarded-For", "X-Real-IP"},
+		TrustedProxies:  nil,
+	}
+	c = NewContext(0)
+	c.SetClientIPFunc(ClientIPWithOption(opts))
+	c.conn = mock.NewConn("")
+	c.Request.Header.Set("X-Forwarded-For", "  126.0.0.2, 0.0.0.0 ")
+	val = c.ClientIP()
+	if val != "0.0.0.0" {
+		t.Fatalf("unexpected %v. Expecting %v", val, "0.0.0.0")
+	}
 }
 
 func TestSetClientIPFunc(t *testing.T) {
diff --git a/pkg/protocol/http1/ext/stream_test.go b/pkg/protocol/http1/ext/stream_test.go
new file mode 100644
index 0000000..57cad37
--- /dev/null
+++ b/pkg/protocol/http1/ext/stream_test.go
@@ -0,0 +1,90 @@
+/*
+ * Copyright 2022 CloudWeGo Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package ext
+
+import (
+	"fmt"
+	"io"
+	"testing"
+
+	"github.com/cloudwego/hertz/pkg/common/bytebufferpool"
+	"github.com/cloudwego/hertz/pkg/common/test/assert"
+	"github.com/cloudwego/hertz/pkg/common/test/mock"
+	"github.com/cloudwego/hertz/pkg/protocol"
+)
+
+func createChunkedBody(body, rest []byte, trailer map[string]string, hasTrailer bool) []byte {
+	var b []byte
+	chunkSize := 1
+	for len(body) > 0 {
+		if chunkSize > len(body) {
+			chunkSize = len(body)
+		}
+		b = append(b, []byte(fmt.Sprintf("%x\r\n", chunkSize))...)
+		b = append(b, body[:chunkSize]...)
+		b = append(b, []byte("\r\n")...)
+		body = body[chunkSize:]
+		chunkSize++
+	}
+	if hasTrailer {
+		b = append(b, "0\r\n"...)
+		for k, v := range trailer {
+			b = append(b, k...)
+			b = append(b, ": "...)
+			b = append(b, v...)
+			b = append(b, "\r\n"...)
+		}
+		b = append(b, "\r\n"...)
+	}
+	return append(b, rest...)
+}
+
+func testChunkedSkipRest(t *testing.T, data, rest string) {
+	var pool bytebufferpool.Pool
+	reader := mock.NewZeroCopyReader(data)
+
+	bs := AcquireBodyStream(pool.Get(), reader, &protocol.Trailer{}, -1)
+	err := bs.(*bodyStream).skipRest()
+	assert.Nil(t, err)
+
+	rest_data, err := io.ReadAll(reader)
+	assert.Nil(t, err)
+	assert.DeepEqual(t, rest, string(rest_data))
+}
+
+func testChunkedSkipRestWithBodySize(t *testing.T, bodySize int) {
+	body := mock.CreateFixedBody(bodySize)
+	rest := mock.CreateFixedBody(bodySize)
+	data := createChunkedBody(body, rest, map[string]string{"foo": "bar"}, true)
+
+	testChunkedSkipRest(t, string(data), string(rest))
+}
+
+func TestChunkedSkipRest(t *testing.T) {
+	t.Parallel()
+
+	testChunkedSkipRest(t, "0\r\n\r\n", "")
+	testChunkedSkipRest(t, "0\r\n\r\nHTTP/1.1 / POST", "HTTP/1.1 / POST")
+	testChunkedSkipRest(t, "0\r\nHertz: test\r\nfoo: bar\r\n\r\nHTTP/1.1 / POST", "HTTP/1.1 / POST")
+
+	testChunkedSkipRestWithBodySize(t, 5)
+
+	// medium-size body
+	testChunkedSkipRestWithBodySize(t, 43488)
+
+	// big body
+	testChunkedSkipRestWithBodySize(t, 3*1024*1024)
+}