fixes #2016 - make IP() and IPs() more reliable

ctx.go

@@ -649,33 +649,77 @@ func (c *Ctx) Port() string {
 }
 
 // IP returns the remote IP address of the request.
+// If ProxyHeader is configured, it will parse that header and return the first valid IP address
 // Please use Config.EnableTrustedProxyCheck to prevent header spoofing, in case when your app is behind the proxy.
 func (c *Ctx) IP() string {
 	if c.IsProxyTrusted() && len(c.app.config.ProxyHeader) > 0 {
-		return c.Get(c.app.config.ProxyHeader)
+		return c.extractIPFromHeader(c.app.config.ProxyHeader)
 	}
 
 	return c.fasthttp.RemoteIP().String()
 }
 
-// IPs returns an string slice of IP addresses specified in the X-Forwarded-For request header.
-func (c *Ctx) IPs() (ips []string) {
-	header := c.fasthttp.Request.Header.Peek(HeaderXForwardedFor)
-	if len(header) == 0 {
-		return
-	}
-	ips = make([]string, bytes.Count(header, []byte(","))+1)
-	var commaPos, i int
+// extractValidIPs will return a slice of strings that represent valid IP addresses
+// in the input string. The order is maintained. The separator is a comma
+func extractValidIPs(input string) (validIPs []string) {
+
+	// try to gather IPs in the input with minimal allocations to improve performance
+	ips := make([]string, bytes.Count([]byte(input), []byte(","))+1)
+	var commaPos, i, validCount int
 	for {
-		commaPos = bytes.IndexByte(header, ',')
+		commaPos = bytes.IndexByte([]byte(input), ',')
 		if commaPos != -1 {
-			ips[i] = utils.Trim(c.app.getString(header[:commaPos]), ' ')
-			header, i = header[commaPos+1:], i+1
+			if net.ParseIP(utils.Trim(input[:commaPos], ' ')) != nil {
+				ips[i] = utils.Trim(input[:commaPos], ' ')
+				validCount++
+			}
+			input, i = input[commaPos+1:], i+1
 		} else {
-			ips[i] = utils.Trim(c.app.getString(header), ' ')
-			return
+			if net.ParseIP(utils.Trim(input, ' ')) != nil {
+				ips[i] = utils.Trim(input, ' ')
+				validCount++
+			}
+			break
+		}
+	}
+
+	// filter out any invalid IP(s) that we found
+	if len(ips) == validCount {
+		validIPs = ips
+	} else {
+		validIPs = make([]string, validCount)
+		var validIndex int
+		for n := range ips {
+			if ips[n] != "" {
+				validIPs[validIndex] = ips[n]
+				validIndex++
+			}
 		}
 	}
+	return
+}
+
+// extractIPFromHeader will attempt to pull the real client IP from the given header
+// currently it will return the first valid IP address in header
+func (c *Ctx) extractIPFromHeader(header string) string {
+	// extract only valid IPs from the header's value
+	validIps := extractValidIPs(c.Get(header))
+
+	// since X-Forwarded-For has no RFC, it's really up to the proxy to decide whether to append
+	// or prepend IPs to this list. For example, the AWS ALB will prepend but the F5 BIG-IP will append ;(
+	// for now lets just go with the first value in the list...
+	if len(validIps) > 0 {
+		return validIps[0]
+	}
+
+	// return the IP from the stack if we could not find any valid Ips
+	return c.fasthttp.RemoteIP().String()
+}
+
+// IPs returns a string slice of IP addresses specified in the X-Forwarded-For request header.
+// Only valid IP addresses are returned
+func (c *Ctx) IPs() (ips []string) {
+	return extractValidIPs(c.Get(HeaderXForwardedFor))
 }
 
 // Is returns the matching content type,

ctx_test.go

@@ -1099,19 +1099,51 @@ func Test_Ctx_PortInHandler(t *testing.T) {
 // go test -run Test_Ctx_IP
 func Test_Ctx_IP(t *testing.T) {
 	t.Parallel()
+
 	app := New()
 	c := app.AcquireCtx(&fasthttp.RequestCtx{})
 	defer app.ReleaseCtx(c)
+
+	// default behaviour will return the remote IP from the stack
+	utils.AssertEqual(t, "0.0.0.0", c.IP())
+
+	// X-Forwarded-For is set, but it is ignored because proxyHeader is not set
+	c.Request().Header.Set(HeaderXForwardedFor, "0.0.0.1")
 	utils.AssertEqual(t, "0.0.0.0", c.IP())
 }
 
 // go test -run Test_Ctx_IP_ProxyHeader
 func Test_Ctx_IP_ProxyHeader(t *testing.T) {
 	t.Parallel()
-	app := New(Config{ProxyHeader: "Real-Ip"})
-	c := app.AcquireCtx(&fasthttp.RequestCtx{})
-	defer app.ReleaseCtx(c)
-	utils.AssertEqual(t, "", c.IP())
+
+	// make sure that the same behaviour exists for different proxy header names
+	proxyHeaderNames := []string{"Real-Ip", HeaderXForwardedFor}
+
+	for _, proxyHeaderName := range proxyHeaderNames {
+		app := New(Config{ProxyHeader: proxyHeaderName})
+		c := app.AcquireCtx(&fasthttp.RequestCtx{})
+
+		// when proxy header is enabled and the value is a valid IP, we return it
+		c.Request().Header.Set(proxyHeaderName, "0.0.0.1")
+		utils.AssertEqual(t, "0.0.0.1", c.IP())
+
+		// when proxy header is enabled and the value is a list of IPs, we return the first valid IP
+		c.Request().Header.Set(proxyHeaderName, "0.0.0.1, 0.0.0.2")
+		utils.AssertEqual(t, "0.0.0.1", c.IP())
+
+		c.Request().Header.Set(proxyHeaderName, "invalid, 0.0.0.2, 0.0.0.3")
+		utils.AssertEqual(t, "0.0.0.2", c.IP())
+
+		// when proxy header is enabled but the value is empty, we will ignore the header
+		c.Request().Header.Set(proxyHeaderName, "")
+		utils.AssertEqual(t, "0.0.0.0", c.IP())
+
+		// when proxy header is enabled but the value is not an IP, we will ignore the header
+		c.Request().Header.Set(proxyHeaderName, "not-valid-ip")
+		utils.AssertEqual(t, "0.0.0.0", c.IP())
+
+		app.ReleaseCtx(c)
+	}
 }
 
 // go test -run Test_Ctx_IP_UntrustedProxy
@@ -1140,14 +1172,32 @@ func Test_Ctx_IPs(t *testing.T) {
 	app := New()
 	c := app.AcquireCtx(&fasthttp.RequestCtx{})
 	defer app.ReleaseCtx(c)
+
+	// normal happy path test case
 	c.Request().Header.Set(HeaderXForwardedFor, "127.0.0.1, 127.0.0.2, 127.0.0.3")
 	utils.AssertEqual(t, []string{"127.0.0.1", "127.0.0.2", "127.0.0.3"}, c.IPs())
 
+	// inconsistent space formatting
 	c.Request().Header.Set(HeaderXForwardedFor, "127.0.0.1,127.0.0.2  ,127.0.0.3")
 	utils.AssertEqual(t, []string{"127.0.0.1", "127.0.0.2", "127.0.0.3"}, c.IPs())
 
+	// invalid IPs are in the header
+	c.Request().Header.Set(HeaderXForwardedFor, "invalid, 127.0.0.1, 127.0.0.2")
+	utils.AssertEqual(t, []string{"127.0.0.1", "127.0.0.2"}, c.IPs())
+	c.Request().Header.Set(HeaderXForwardedFor, "127.0.0.1, invalid, 127.0.0.2")
+	utils.AssertEqual(t, []string{"127.0.0.1", "127.0.0.2"}, c.IPs())
+
+	// ensure that the ordering of IPs in the header is maintained
+	c.Request().Header.Set(HeaderXForwardedFor, "127.0.0.3, 127.0.0.1, 127.0.0.2")
+	utils.AssertEqual(t, []string{"127.0.0.3", "127.0.0.1", "127.0.0.2"}, c.IPs())
+
+	// empty header
 	c.Request().Header.Set(HeaderXForwardedFor, "")
 	utils.AssertEqual(t, 0, len(c.IPs()))
+
+	// missing header
+	c.Request()
+	utils.AssertEqual(t, 0, len(c.IPs()))
 }
 
 // go test -v -run=^$ -bench=Benchmark_Ctx_IPs -benchmem -count=4
@@ -1165,6 +1215,34 @@ func Benchmark_Ctx_IPs(b *testing.B) {
 	utils.AssertEqual(b, []string{"127.0.0.1", "127.0.0.1", "127.0.0.1"}, res)
 }
 
+func Benchmark_Ctx_IP_With_ProxyHeader(b *testing.B) {
+	app := New(Config{ProxyHeader: HeaderXForwardedFor})
+	c := app.AcquireCtx(&fasthttp.RequestCtx{})
+	defer app.ReleaseCtx(c)
+	c.Request().Header.Set(HeaderXForwardedFor, "127.0.0.1")
+	var res string
+	b.ReportAllocs()
+	b.ResetTimer()
+	for n := 0; n < b.N; n++ {
+		res = c.IP()
+	}
+	utils.AssertEqual(b, "127.0.0.1", res)
+}
+
+func Benchmark_Ctx_IP(b *testing.B) {
+	app := New()
+	c := app.AcquireCtx(&fasthttp.RequestCtx{})
+	defer app.ReleaseCtx(c)
+	c.Request()
+	var res string
+	b.ReportAllocs()
+	b.ResetTimer()
+	for n := 0; n < b.N; n++ {
+		res = c.IP()
+	}
+	utils.AssertEqual(b, "0.0.0.0", res)
+}
+
 // go test -run Test_Ctx_Is
 func Test_Ctx_Is(t *testing.T) {
 	t.Parallel()