From 0b56a6001211eb81f6e9f0e7b78b3601d4914db1 Mon Sep 17 00:00:00 2001
From: Notealot <714804968@qq.com>
Date: Mon, 29 Nov 2021 11:14:02 +0800
Subject: [PATCH] Refactor and fix

---
 README.md       |  4 ++--
 context.go      | 52 ++++++++-----------------------------------------
 context_test.go | 10 +++++++++-
 gin.go          | 41 +++++++++++++++++++++++++++++++++++++-
 4 files changed, 59 insertions(+), 48 deletions(-)

diff --git a/README.md b/README.md
index cad746d6..203fe79d 100644
--- a/README.md
+++ b/README.md
@@ -78,7 +78,7 @@ Gin is a web framework written in Go (Golang). It features a martini-like API wi
     - [http2 server push](#http2-server-push)
     - [Define format for the log of routes](#define-format-for-the-log-of-routes)
     - [Set and get a cookie](#set-and-get-a-cookie)
-  - [Don't trust all proxies](#don't-trust-all-proxies)
+  - [Don't trust all proxies](#do-not-trust-all-proxies)
   - [Testing](#testing)
   - [Users](#users)
 
@@ -2197,7 +2197,7 @@ func main() {
 }
 ```
 
-## Don't trust all proxies
+## Do not trust all proxies
 
 Gin lets you specify which headers to hold the real client IP (if any),
 as well as specifying which proxies (or direct clients) you trust to
diff --git a/context.go b/context.go
index 2aa91e11..8146be37 100644
--- a/context.go
+++ b/context.go
@@ -757,10 +757,14 @@ func (c *Context) ClientIP() string {
 		}
 	}
 
-	remoteIP, trusted := c.RemoteIP()
+	// It also checks if the remoteIP is a trusted proxy or not.
+	// In order to perform this validation, it will see if the IP is contained within at least one of the CIDR blocks
+	// defined by Engine.SetTrustedProxies()
+	remoteIP := net.ParseIP(c.RemoteIP())
 	if remoteIP == nil {
 		return ""
 	}
+	trusted := c.engine.isTrustedProxy(remoteIP)
 
 	if trusted && c.engine.ForwardedByClientIP && c.engine.RemoteIPHeaders != nil {
 		for _, headerName := range c.engine.RemoteIPHeaders {
@@ -773,53 +777,13 @@ func (c *Context) ClientIP() string {
 	return remoteIP.String()
 }
 
-func (e *Engine) isTrustedProxy(ip net.IP) bool {
-	if e.trustedCIDRs != nil {
-		for _, cidr := range e.trustedCIDRs {
-			if cidr.Contains(ip) {
-				return true
-			}
-		}
-	}
-	return false
-}
-
 // RemoteIP parses the IP from Request.RemoteAddr, normalizes and returns the IP (without the port).
-// It also checks if the remoteIP is a trusted proxy or not.
-// In order to perform this validation, it will see if the IP is contained within at least one of the CIDR blocks
-// defined by Engine.SetTrustedProxies()
-func (c *Context) RemoteIP() (net.IP, bool) {
+func (c *Context) RemoteIP() string {
 	ip, _, err := net.SplitHostPort(strings.TrimSpace(c.Request.RemoteAddr))
 	if err != nil {
-		return nil, false
+		return ""
 	}
-	remoteIP := net.ParseIP(ip)
-	if remoteIP == nil {
-		return nil, false
-	}
-
-	return remoteIP, c.engine.isTrustedProxy(remoteIP)
-}
-
-func (e *Engine) validateHeader(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 {
-			return "", false
-		}
-
-		// X-Forwarded-For is appended by proxy
-		// Check IPs in reverse order and stop when find untrusted proxy
-		if (i == 0) || (!e.isTrustedProxy(ip)) {
-			return ipStr, true
-		}
-	}
-	return
+	return ip
 }
 
 // ContentType returns the Content-Type header of the request.
diff --git a/context_test.go b/context_test.go
index c286c0f4..4d002c23 100644
--- a/context_test.go
+++ b/context_test.go
@@ -12,6 +12,7 @@ import (
 	"html/template"
 	"io"
 	"mime/multipart"
+	"net"
 	"net/http"
 	"net/http/httptest"
 	"os"
@@ -1404,6 +1405,11 @@ func TestContextClientIP(t *testing.T) {
 	// Tests exercising the TrustedProxies functionality
 	resetContextForClientIPTests(c)
 
+	// IPv6 support
+	c.Request.RemoteAddr = "[::1]:12345"
+	assert.Equal(t, "20.20.20.20", c.ClientIP())
+
+	resetContextForClientIPTests(c)
 	// No trusted proxies
 	_ = c.engine.SetTrustedProxies([]string{})
 	c.engine.RemoteIPHeaders = []string{"X-Forwarded-For"}
@@ -1500,6 +1506,7 @@ func resetContextForClientIPTests(c *Context) {
 	c.Request.Header.Set("CF-Connecting-IP", "60.60.60.60")
 	c.Request.RemoteAddr = "  40.40.40.40:42123 "
 	c.engine.TrustedPlatform = ""
+	c.engine.trustedCIDRs = defaultTrustedCIDRs
 	c.engine.AppEngine = false
 }
 
@@ -2051,7 +2058,8 @@ func TestRemoteIPFail(t *testing.T) {
 	c, _ := CreateTestContext(httptest.NewRecorder())
 	c.Request, _ = http.NewRequest("POST", "/", nil)
 	c.Request.RemoteAddr = "[:::]:80"
-	ip, trust := c.RemoteIP()
+	ip := net.ParseIP(c.RemoteIP())
+	trust := c.engine.isTrustedProxy(ip)
 	assert.Nil(t, ip)
 	assert.False(t, trust)
 }
diff --git a/gin.go b/gin.go
index 1d9c9fee..42f281f3 100644
--- a/gin.go
+++ b/gin.go
@@ -28,7 +28,13 @@ var (
 
 var defaultPlatform string
 
-var defaultTrustedCIDRs = []*net.IPNet{{IP: net.IP{0x0, 0x0, 0x0, 0x0}, Mask: net.IPMask{0x0, 0x0, 0x0, 0x0}}} // 0.0.0.0/0
+var defaultTrustedCIDRs = []*net.IPNet{
+	{IP: net.IP{0x0, 0x0, 0x0, 0x0}, Mask: net.IPMask{0x0, 0x0, 0x0, 0x0}}, // 0.0.0.0/0
+	{ // ::/0
+		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},
+	},
+}
 
 // HandlerFunc defines the handler used by gin middleware as return value.
 type HandlerFunc func(*Context)
@@ -411,6 +417,39 @@ func (engine *Engine) parseTrustedProxies() error {
 	return err
 }
 
+// isTrustedProxy will check whether the IP address is included in the trusted list according to Engine.trustedCIDRs
+func (engine *Engine) isTrustedProxy(ip net.IP) bool {
+	if engine.trustedCIDRs != nil {
+		for _, cidr := range engine.trustedCIDRs {
+			if cidr.Contains(ip) {
+				return true
+			}
+		}
+	}
+	return false
+}
+
+// validateHeader will parse X-Forwarded-For header and return the trusted client IP address
+func (engine *Engine) validateHeader(header string) (clientIP string, valid bool) {
+	if header != "" {
+		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
+			}
+
+			// X-Forwarded-For is appended by proxy
+			// Check IPs in reverse order and stop when find untrusted proxy
+			if (i == 0) || (!engine.isTrustedProxy(ip)) {
+				return ipStr, true
+			}
+		}
+	}
+	return "", false
+}
+
 // parseIP parse a string representation of an IP and returns a net.IP with the
 // minimum byte representation or nil if input is invalid.
 func parseIP(ip string) net.IP {