Repository: cloudwego/hertz
Pull request #934 (2023-09-03 22:50:03)

optimize: adjusting the judgement logic of isTrustedProxy
#### What type of PR is this?

optimize

<!--
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
-->

#### 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.
- [x] 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.
- 调整 isTrustedProxy 的判断逻辑，现在的 clientIP 的默认逻辑中 isTrustedProxy 正常情况下一定返回的是 false，导致即便用户的 header 中存在 xr和xl 也不会反回，而是返回 remoteaddress。
- gin 的默认实现是宽松的，isTrustedProxy 会默认返回 true，而不是 false。



#### (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)
-->

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/cloudwego/hertz/pull/934?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.03%`** :warning:
> Comparison is base [(`74cedc7`)](https://app.codecov.io/gh/cloudwego/hertz/commit/74cedc74ad2804f9eb88a70ad2565893cb260edf?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego) 81.95% compared to head [(`c535316`)](https://app.codecov.io/gh/cloudwego/hertz/pull/934?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego) 81.93%.
> Report is 3 commits behind head on develop.

> :exclamation: Current head c535316 differs from pull request most recent head 906ec47. Consider uploading reports for the commit 906ec47 to get more accurate results

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


```diff
@@             Coverage Diff             @@
##           develop     #934      +/-   ##
===========================================
- Coverage    81.95%   81.93%   -0.03%     
===========================================
  Files           97       97              
  Lines         9623     9633      +10     
===========================================
+ Hits          7887     7893       +6     
- Misses        1262     1264       +2     
- Partials       474      476       +2     
```


| [Files Changed](https://app.codecov.io/gh/cloudwego/hertz/pull/934?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego) | Coverage Δ | |
|---|---|---|
| [pkg/protocol/request.go](https://app.codecov.io/gh/cloudwego/hertz/pull/934?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego#diff-cGtnL3Byb3RvY29sL3JlcXVlc3QuZ28=) | `82.86% <ø> (ø)` | |
| [pkg/protocol/response.go](https://app.codecov.io/gh/cloudwego/hertz/pull/934?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego#diff-cGtnL3Byb3RvY29sL3Jlc3BvbnNlLmdv) | `84.30% <0.00%> (-1.00%)` | :arrow_down: |
| [pkg/app/context.go](https://app.codecov.io/gh/cloudwego/hertz/pull/934?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego#diff-cGtnL2FwcC9jb250ZXh0Lmdv) | `95.43% <88.23%> (-0.54%)` | :arrow_down: |
| [pkg/protocol/http1/client.go](https://app.codecov.io/gh/cloudwego/hertz/pull/934?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego#diff-cGtnL3Byb3RvY29sL2h0dHAxL2NsaWVudC5nbw==) | `73.60% <100.00%> (+0.17%)` | :arrow_up: |


</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/cloudwego/hertz/pull/934?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego).   
:loudspeaker: Have feedback on the report? [Share it here](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: PATL @welkeyever  @Duslia 

Patch:
diff --git a/pkg/app/context.go b/pkg/app/context.go
index 559b84d..9a6e454 100644
--- a/pkg/app/context.go
+++ b/pkg/app/context.go
@@ -80,48 +80,72 @@ type ClientIP func(ctx *RequestContext) string
 
 type ClientIPOptions struct {
 	RemoteIPHeaders []string
-	TrustedProxies  map[string]bool
+	TrustedCIDRs    []*net.IPNet
 }
 
-var defaultClientIPOptions = ClientIPOptions{
-	RemoteIPHeaders: []string{"X-Real-IP", "X-Forwarded-For"},
-	TrustedProxies: map[string]bool{
-		"0.0.0.0": true,
+var defaultTrustedCIDRs = []*net.IPNet{
+	{ // 0.0.0.0/0 (IPv4)
+		IP:   net.IP{0x0, 0x0, 0x0, 0x0},
+		Mask: net.IPMask{0x0, 0x0, 0x0, 0x0},
+	},
+	{ // ::/0 (IPv6)
+		IP:   net.IP{0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0},
+		Mask: net.IPMask{0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0},
 	},
 }
 
+var defaultClientIPOptions = ClientIPOptions{
+	RemoteIPHeaders: []string{"X-Forwarded-For", "X-Real-IP"},
+	TrustedCIDRs:    defaultTrustedCIDRs,
+}
+
 // 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
+		TrustedCIDRs := opts.TrustedCIDRs
 
-		remoteIP, _, err := net.SplitHostPort(strings.TrimSpace(ctx.RemoteAddr().String()))
+		remoteIPStr, _, err := net.SplitHostPort(strings.TrimSpace(ctx.RemoteAddr().String()))
 		if err != nil {
 			return ""
 		}
-		trusted := isTrustedProxy(TrustedProxies, remoteIP)
+
+		remoteIP := net.ParseIP(remoteIPStr)
+		if remoteIP == nil {
+			return ""
+		}
+
+		trusted := isTrustedProxy(TrustedCIDRs, remoteIP)
 
 		if trusted {
 			for _, headerName := range RemoteIPHeaders {
-				ip, valid := validateHeader(TrustedProxies, ctx.Request.Header.Get(headerName))
+				ip, valid := validateHeader(TrustedCIDRs, ctx.Request.Header.Get(headerName))
 				if valid {
 					return ip
 				}
 			}
 		}
 
-		return remoteIP
+		return remoteIPStr
 	}
 }
 
-// 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]
+// isTrustedProxy will check whether the IP address is included in the trusted list according to trustedCIDRs
+func isTrustedProxy(trustedCIDRs []*net.IPNet, remoteIP net.IP) bool {
+	if trustedCIDRs == nil {
+		return false
+	}
+
+	for _, cidr := range trustedCIDRs {
+		if cidr.Contains(remoteIP) {
+			return true
+		}
+	}
+	return false
 }
 
 // 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) {
+func validateHeader(trustedCIDRs []*net.IPNet, header string) (clientIP string, valid bool) {
 	if header == "" {
 		return "", false
 	}
@@ -135,7 +159,7 @@ func validateHeader(trustedProxies map[string]bool, header string) (clientIP str
 
 		// X-Forwarded-For is appended by proxy
 		// Check IPs in reverse order and stop when find untrusted proxy
-		if (i == 0) || (!isTrustedProxy(trustedProxies, ipStr)) {
+		if (i == 0) || (!isTrustedProxy(trustedCIDRs, ip)) {
 			return ipStr, true
 		}
 	}


Test patch:
diff --git a/pkg/app/context_test.go b/pkg/app/context_test.go
index 7211e39..88cd263 100644
--- a/pkg/app/context_test.go
+++ b/pkg/app/context_test.go
@@ -24,6 +24,7 @@ import (
 	"fmt"
 	"html/template"
 	"io/ioutil"
+	"net"
 	"os"
 	"reflect"
 	"strings"
@@ -802,51 +803,66 @@ func TestContextContentType(t *testing.T) {
 	assert.DeepEqual(t, consts.MIMEApplicationJSONUTF8, bytesconv.B2s(c.ContentType()))
 }
 
-func TestClientIp(t *testing.T) {
+type MockIpConn struct {
+	*mock.Conn
+	RemoteIp string
+	Port     int
+}
+
+func (c *MockIpConn) RemoteAddr() net.Addr {
+	return &net.UDPAddr{
+		IP:   net.ParseIP(c.RemoteIp),
+		Port: c.Port,
+	}
+}
+
+func newContextClientIPTest() *RequestContext {
 	c := NewContext(0)
-	c.conn = mock.NewConn("")
-	// 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")
-	}
-	// 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
+	c.conn = &MockIpConn{
+		Conn:     mock.NewConn(""),
+		RemoteIp: "127.0.0.1",
+		Port:     8080,
+	}
+	c.Request.Header.Set("X-Real-IP", " 10.10.10.10  ")
+	c.Request.Header.Set("X-Forwarded-For", "  20.20.20.20, 30.30.30.30")
+	return c
+}
+
+func TestClientIp(t *testing.T) {
+	c := newContextClientIPTest()
+	// default X-Forwarded-For and X-Real-IP behaviour
+	assert.DeepEqual(t, "20.20.20.20", c.ClientIP())
+
+	c.Request.Header.DelBytes([]byte("X-Forwarded-For"))
+	assert.DeepEqual(t, "10.10.10.10", c.ClientIP())
+
+	c.Request.Header.Set("X-Forwarded-For", "30.30.30.30  ")
+	assert.DeepEqual(t, "30.30.30.30", c.ClientIP())
+
+	// No trusted CIDRS
+	c = newContextClientIPTest()
 	opts := ClientIPOptions{
 		RemoteIPHeaders: []string{"X-Forwarded-For", "X-Real-IP"},
-		TrustedProxies: map[string]bool{
-			"0.0.0.0": true,
-		},
+		TrustedCIDRs:    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 != "126.0.0.2" {
-		t.Fatalf("unexpected %v. Expecting %v", val, "126.0.0.2")
-	}
-	// no trusted proxy server
+	assert.DeepEqual(t, "127.0.0.1", c.ClientIP())
+
+	_, cidr, _ := net.ParseCIDR("30.30.30.30/32")
 	opts = ClientIPOptions{
 		RemoteIPHeaders: []string{"X-Forwarded-For", "X-Real-IP"},
-		TrustedProxies:  nil,
+		TrustedCIDRs:    []*net.IPNet{cidr},
 	}
-	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")
+	assert.DeepEqual(t, "127.0.0.1", c.ClientIP())
+
+	_, cidr, _ = net.ParseCIDR("127.0.0.1/32")
+	opts = ClientIPOptions{
+		RemoteIPHeaders: []string{"X-Forwarded-For", "X-Real-IP"},
+		TrustedCIDRs:    []*net.IPNet{cidr},
 	}
+	c.SetClientIPFunc(ClientIPWithOption(opts))
+	assert.DeepEqual(t, "30.30.30.30", c.ClientIP())
 }
 
 func TestSetClientIPFunc(t *testing.T) {
diff --git a/pkg/common/utils/ioutil_test.go b/pkg/common/utils/ioutil_test.go
index c477180..e9a573b 100644
--- a/pkg/common/utils/ioutil_test.go
+++ b/pkg/common/utils/ioutil_test.go
@@ -106,6 +106,23 @@ func TestIoutilCopyBuffer(t *testing.T) {
 	assert.DeepEqual(t, written, srcLen)
 	assert.DeepEqual(t, err, nil)
 	assert.DeepEqual(t, []byte(str), writeBuffer.Bytes())
+
+	// Test when no data is readable
+	writeBuffer.Reset()
+	emptySrc := bytes.NewBufferString("")
+	written, err = CopyBuffer(dst, emptySrc, buf)
+	assert.DeepEqual(t, written, int64(0))
+	assert.Nil(t, err)
+	assert.DeepEqual(t, []byte(""), writeBuffer.Bytes())
+
+	// Test a LimitedReader
+	writeBuffer.Reset()
+	limit := int64(5)
+	limitedSrc := io.LimitedReader{R: bytes.NewBufferString(str), N: limit}
+	written, err = CopyBuffer(dst, &limitedSrc, buf)
+	assert.DeepEqual(t, written, limit)
+	assert.Nil(t, err)
+	assert.DeepEqual(t, []byte(str[:limit]), writeBuffer.Bytes())
 }
 
 func TestIoutilCopyBufferWithIoWriter(t *testing.T) {
@@ -198,7 +215,7 @@ func TestIoutilCopyBufferWithNilBufferAndIoLimitedReader(t *testing.T) {
 
 func TestIoutilCopyZeroAlloc(t *testing.T) {
 	var writeBuffer bytes.Buffer
-	str := string("hertz is very good!!!")
+	str := "hertz is very good!!!"
 	src := bytes.NewBufferString(str)
 	dst := network.NewWriter(&writeBuffer)
 	srcLen := int64(src.Len())
@@ -207,4 +224,49 @@ func TestIoutilCopyZeroAlloc(t *testing.T) {
 	assert.DeepEqual(t, written, srcLen)
 	assert.DeepEqual(t, err, nil)
 	assert.DeepEqual(t, []byte(str), writeBuffer.Bytes())
+
+	// Test when no data is readable
+	writeBuffer.Reset()
+	emptySrc := bytes.NewBufferString("")
+	written, err = CopyZeroAlloc(dst, emptySrc)
+	assert.DeepEqual(t, written, int64(0))
+	assert.Nil(t, err)
+	assert.DeepEqual(t, []byte(""), writeBuffer.Bytes())
+}
+
+func TestIoutilCopyBufferWithEmptyBuffer(t *testing.T) {
+	var writeBuffer bytes.Buffer
+	str := "hertz is very good!!!"
+	src := bytes.NewBufferString(str)
+	dst := network.NewWriter(&writeBuffer)
+	// Use a non-empty buffer of length 0
+	emptyBuf := make([]byte, 0)
+	func() {
+		defer func() {
+			if r := recover(); r != nil {
+				assert.DeepEqual(t, "empty buffer in io.CopyBuffer", r)
+			}
+		}()
+
+		written, err := CopyBuffer(dst, src, emptyBuf)
+		assert.Nil(t, err)
+		assert.DeepEqual(t, written, int64(len(str)))
+		assert.DeepEqual(t, []byte(str), writeBuffer.Bytes())
+	}()
+}
+
+func TestIoutilCopyBufferWithLimitedReader(t *testing.T) {
+	var writeBuffer bytes.Buffer
+	str := "hertz is very good!!!"
+	src := bytes.NewBufferString(str)
+	limit := int64(5)
+	limitedSrc := io.LimitedReader{R: src, N: limit}
+	dst := network.NewWriter(&writeBuffer)
+	var buf []byte
+
+	// Test LimitedReader status
+	written, err := CopyBuffer(dst, &limitedSrc, buf)
+	assert.Nil(t, err)
+	assert.DeepEqual(t, written, limit)
+	assert.DeepEqual(t, []byte(str[:limit]), writeBuffer.Bytes())
 }
diff --git a/pkg/common/utils/network_test.go b/pkg/common/utils/network_test.go
index 232e35a..78be77d 100644
--- a/pkg/common/utils/network_test.go
+++ b/pkg/common/utils/network_test.go
@@ -41,3 +41,14 @@ func TestTLSRecordHeaderLooksLikeHTTP(t *testing.T) {
 		assert.DeepEqual(t, expectedResult, TLSRecordHeaderLooksLikeHTTP(value))
 	}
 }
+
+func TestLocalIP(t *testing.T) {
+	// Mock the localIP variable for testing purposes.
+	localIP = "192.168.0.1"
+
+	// Ensure that LocalIP() returns the expected local IP.
+	expectedIP := "192.168.0.1"
+	if got := LocalIP(); got != expectedIP {
+		assert.DeepEqual(t, got, expectedIP)
+	}
+}
diff --git a/pkg/common/utils/path_test.go b/pkg/common/utils/path_test.go
index 6ebb843..b98a8c7 100644
--- a/pkg/common/utils/path_test.go
+++ b/pkg/common/utils/path_test.go
@@ -82,6 +82,11 @@ func TestPathCleanPath(t *testing.T) {
 	expectedMultiSlashPath := "/Foo/Bar/go/src/github.com/cloudwego"
 	cleanMultiSlashPath := CleanPath(multiSlashPath)
 	assert.DeepEqual(t, expectedMultiSlashPath, cleanMultiSlashPath)
+
+	inputPath := "/Foo/Bar/go/src/github.com/cloudwego/hertz/pkg/common/utils/path_test.go/."
+	expectedPath := "/Foo/Bar/go/src/github.com/cloudwego/hertz/pkg/common/utils/path_test.go/"
+	cleanedPath := CleanPath(inputPath)
+	assert.DeepEqual(t, expectedPath, cleanedPath)
 }
 
 // The Function AddMissingPort can only add the missed port, don't consider the other error case.