From 4fe85da03eafb942061d01ea9466f5a0e2f0f4a7 Mon Sep 17 00:00:00 2001
From: BaiZe1998 <1157467179@qq.com>
Date: Sun, 20 Nov 2022 20:01:33 +0800
Subject: [PATCH] fix: ClientIP handling is unsafe

---
 pkg/app/context.go      | 68 +++++++++++++++++++++++++++++++++++------
 pkg/app/context_test.go | 36 ++++++++++++++++++++--
 2 files changed, 92 insertions(+), 12 deletions(-)

diff --git a/pkg/app/context.go b/pkg/app/context.go
index 157674ce2..e86cc4b1a 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/app/context_test.go b/pkg/app/context_test.go
index fc3c147d8..d408013db 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) {