From 2d3d6d2f1355f21a77f8ef6cf45a5bc531c8c935 Mon Sep 17 00:00:00 2001
From: Notealot <714804968@qq.com>
Date: Sun, 24 Oct 2021 08:34:03 +0800
Subject: [PATCH] Provide custom options of TrustedPlatform for another CDN
 services (#2906)

* refine TrustedPlatform and docs

* refactor for switch

* refactor switch to if statement
---
 README.md       | 29 +++++++++++++++++++++++++++++
 context.go      | 14 +++++---------
 context_test.go | 14 +++++++++++++-
 gin.go          |  4 ++--
 4 files changed, 49 insertions(+), 12 deletions(-)

diff --git a/README.md b/README.md
index 0fa78ec..cad746d 100644
--- a/README.md
+++ b/README.md
@@ -78,6 +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)
   - [Testing](#testing)
   - [Users](#users)
 
@@ -2236,6 +2237,34 @@ func main() {
 }
 ```
 
+**Notice:** If you are using a CDN service, you can set the `Engine.TrustedPlatform`
+to skip TrustedProxies check, it has a higher priority than TrustedProxies. 
+Look at the example below:
+```go
+import (
+	"fmt"
+
+	"github.com/gin-gonic/gin"
+)
+
+func main() {
+
+	router := gin.Default()
+	// Use predefined header gin.PlatformXXX
+	router.TrustedPlatform = gin.PlatformGoogleAppEngine
+	// Or set your own trusted request header for another trusted proxy service
+	// Don't set it to any suspect request header, it's unsafe
+	router.TrustedPlatform = "X-CDN-IP"
+
+	router.GET("/", func(c *gin.Context) {
+		// If you set TrustedPlatform, ClientIP() will resolve the
+		// corresponding header and return IP directly
+		fmt.Printf("ClientIP: %s\n", c.ClientIP())
+	})
+	router.Run()
+}
+```
+
 ## Testing
 
 The `net/http/httptest` package is preferable way for HTTP testing.
diff --git a/context.go b/context.go
index bc2c38e..58f38c8 100644
--- a/context.go
+++ b/context.go
@@ -735,20 +735,16 @@ func (c *Context) ShouldBindBodyWith(obj interface{}, bb binding.BindingBody) (e
 	return bb.BindBody(body, obj)
 }
 
-// ClientIP implements a best effort algorithm to return the real client IP.
+// ClientIP implements one best effort algorithm to return the real client IP.
 // It called c.RemoteIP() under the hood, to check if the remote IP is a trusted proxy or not.
 // If it is it will then try to parse the headers defined in Engine.RemoteIPHeaders (defaulting to [X-Forwarded-For, X-Real-Ip]).
 // If the headers are not syntactically valid OR the remote IP does not correspond to a trusted proxy,
 // the remote IP (coming form Request.RemoteAddr) is returned.
 func (c *Context) ClientIP() string {
-	// Check if we're running on a trusted platform
-	switch c.engine.TrustedPlatform {
-	case PlatformGoogleAppEngine:
-		if addr := c.requestHeader("X-Appengine-Remote-Addr"); addr != "" {
-			return addr
-		}
-	case PlatformCloudflare:
-		if addr := c.requestHeader("CF-Connecting-IP"); addr != "" {
+	// Check if we're running on a trusted platform, continue running backwards if error
+	if c.engine.TrustedPlatform != "" {
+		// Developers can define their own header of Trusted Platform or use predefined constants
+		if addr := c.requestHeader(c.engine.TrustedPlatform); addr != "" {
 			return addr
 		}
 	}
diff --git a/context_test.go b/context_test.go
index e9fe88f..c286c0f 100644
--- a/context_test.go
+++ b/context_test.go
@@ -1458,8 +1458,20 @@ func TestContextClientIP(t *testing.T) {
 	c.engine.TrustedPlatform = PlatformGoogleAppEngine
 	assert.Equal(t, "50.50.50.50", c.ClientIP())
 
-	// Test the legacy flag
+	// Use custom TrustedPlatform header
+	c.engine.TrustedPlatform = "X-CDN-IP"
+	c.Request.Header.Set("X-CDN-IP", "80.80.80.80")
+	assert.Equal(t, "80.80.80.80", c.ClientIP())
+	// wrong header
+	c.engine.TrustedPlatform = "X-Wrong-Header"
+	assert.Equal(t, "40.40.40.40", c.ClientIP())
+
+	c.Request.Header.Del("X-CDN-IP")
+	// TrustedPlatform is empty
 	c.engine.TrustedPlatform = ""
+	assert.Equal(t, "40.40.40.40", c.ClientIP())
+
+	// Test the legacy flag
 	c.engine.AppEngine = true
 	assert.Equal(t, "50.50.50.50", c.ClientIP())
 	c.engine.AppEngine = false
diff --git a/gin.go b/gin.go
index 1774717..1d9c9fe 100644
--- a/gin.go
+++ b/gin.go
@@ -59,10 +59,10 @@ type RoutesInfo []RouteInfo
 const (
 	// When running on Google App Engine. Trust X-Appengine-Remote-Addr
 	// for determining the client's IP
-	PlatformGoogleAppEngine = "google-app-engine"
+	PlatformGoogleAppEngine = "X-Appengine-Remote-Addr"
 	// When using Cloudflare's CDN. Trust CF-Connecting-IP for determining
 	// the client's IP
-	PlatformCloudflare = "cloudflare"
+	PlatformCloudflare = "CF-Connecting-IP"
 )
 
 // Engine is the framework's instance, it contains the muxer, middleware and configuration settings.