Add support for ipv6 subnet in ipStrategy

pkg/ip/strategy.go

@@ -3,6 +3,7 @@ package ip
 import (
 	"net"
 	"net/http"
+	"net/netip"
 	"strings"
 )
 
@@ -16,7 +17,10 @@ type Strategy interface {
 }
 
 // RemoteAddrStrategy a strategy that always return the remote address.
-type RemoteAddrStrategy struct{}
+type RemoteAddrStrategy struct {
+	// IPv6Subnet instructs the strategy to return the first IP of the subnet where IP belongs.
+	IPv6Subnet *int
+}
 
 // GetIP returns the selected IP.
 func (s *RemoteAddrStrategy) GetIP(req *http.Request) string {
@@ -24,15 +28,22 @@ func (s *RemoteAddrStrategy) GetIP(req *http.Request) string {
 	if err != nil {
 		return req.RemoteAddr
 	}
+
+	if s.IPv6Subnet != nil {
+		return getIPv6SubnetIP(ip, *s.IPv6Subnet)
+	}
+
 	return ip
 }
 
 // DepthStrategy a strategy based on the depth inside the X-Forwarded-For from right to left.
 type DepthStrategy struct {
 	Depth int
+	// IPv6Subnet instructs the strategy to return the first IP of the subnet where IP belongs.
+	IPv6Subnet *int
 }
 
-// GetIP return the selected IP.
+// GetIP returns the selected IP.
 func (s *DepthStrategy) GetIP(req *http.Request) string {
 	xff := req.Header.Get(xForwardedFor)
 	xffs := strings.Split(xff, ",")
@@ -40,7 +51,14 @@ func (s *DepthStrategy) GetIP(req *http.Request) string {
 	if len(xffs) < s.Depth {
 		return ""
 	}
-	return strings.TrimSpace(xffs[len(xffs)-s.Depth])
+
+	ip := strings.TrimSpace(xffs[len(xffs)-s.Depth])
+
+	if s.IPv6Subnet != nil {
+		return getIPv6SubnetIP(ip, *s.IPv6Subnet)
+	}
+
+	return ip
 }
 
 // PoolStrategy is a strategy based on an IP Checker.
@@ -72,3 +90,23 @@ func (s *PoolStrategy) GetIP(req *http.Request) string {
 
 	return ""
 }
+
+// getIPv6SubnetIP returns the IPv6 subnet IP.
+// It returns the original IP when it is not an IPv6, or if parsing the IP has failed with an error.
+func getIPv6SubnetIP(ip string, ipv6Subnet int) string {
+	addr, err := netip.ParseAddr(ip)
+	if err != nil {
+		return ip
+	}
+
+	if !addr.Is6() {
+		return ip
+	}
+
+	prefix, err := addr.Prefix(ipv6Subnet)
+	if err != nil {
+		return ip
+	}
+
+	return prefix.Addr().String()
+}

pkg/ip/strategy_test.go

@@ -9,23 +9,81 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
+const (
+	ipv6Basic            = "::abcd:ffff:c0a8:1"
+	ipv6BracketsPort     = "[::abcd:ffff:c0a8:1]:80"
+	ipv6BracketsZonePort = "[::abcd:ffff:c0a8:1%1]:80"
+)
+
 func TestRemoteAddrStrategy_GetIP(t *testing.T) {
 	testCases := []struct {
-		desc     string
-		expected string
+		desc       string
+		expected   string
+		remoteAddr string
+		ipv6Subnet *int
 	}{
+		// Valid IP format
 		{
-			desc:     "Use RemoteAddr",
+			desc:     "Use RemoteAddr, ipv4",
 			expected: "192.0.2.1",
 		},
+		{
+			desc:       "Use RemoteAddr, ipv6 brackets with port, no IPv6 subnet",
+			remoteAddr: ipv6BracketsPort,
+			expected:   "::abcd:ffff:c0a8:1",
+		},
+		{
+			desc:       "Use RemoteAddr, ipv6 brackets with zone and port, no IPv6 subnet",
+			remoteAddr: ipv6BracketsZonePort,
+			expected:   "::abcd:ffff:c0a8:1%1",
+		},
+
+		// Invalid IPv6 format
+		{
+			desc:       "Use RemoteAddr, ipv6 basic, missing brackets, no IPv6 subnet",
+			remoteAddr: ipv6Basic,
+			expected:   ipv6Basic,
+		},
+
+		// Valid IP format with subnet
+		{
+			desc:       "Use RemoteAddr, ipv4, ignore subnet",
+			expected:   "192.0.2.1",
+			ipv6Subnet: intPtr(24),
+		},
+		{
+			desc:       "Use RemoteAddr, ipv6 brackets with port, subnet",
+			remoteAddr: ipv6BracketsPort,
+			expected:   "::abcd:0:0:0",
+			ipv6Subnet: intPtr(80),
+		},
+		{
+			desc:       "Use RemoteAddr, ipv6 brackets with zone and port, subnet",
+			remoteAddr: ipv6BracketsZonePort,
+			expected:   "::abcd:0:0:0",
+			ipv6Subnet: intPtr(80),
+		},
+
+		// Valid IP, invalid subnet
+		{
+			desc:       "Use RemoteAddr, ipv6 brackets with port, invalid subnet",
+			remoteAddr: ipv6BracketsPort,
+			expected:   "::abcd:ffff:c0a8:1",
+			ipv6Subnet: intPtr(500),
+		},
 	}
 
 	for _, test := range testCases {
 		t.Run(test.desc, func(t *testing.T) {
 			t.Parallel()
 
-			strategy := RemoteAddrStrategy{}
+			strategy := RemoteAddrStrategy{
+				IPv6Subnet: test.ipv6Subnet,
+			}
 			req := httptest.NewRequest(http.MethodGet, "http://127.0.0.1", nil)
+			if test.remoteAddr != "" {
+				req.RemoteAddr = test.remoteAddr
+			}
 			actual := strategy.GetIP(req)
 			assert.Equal(t, test.expected, actual)
 		})
@@ -38,6 +96,7 @@ func TestDepthStrategy_GetIP(t *testing.T) {
 		depth         int
 		xForwardedFor string
 		expected      string
+		ipv6Subnet    *int
 	}{
 		{
 			desc:          "Use depth",
@@ -57,13 +116,30 @@ func TestDepthStrategy_GetIP(t *testing.T) {
 			xForwardedFor: "10.0.0.2,10.0.0.1",
 			expected:      "10.0.0.2",
 		},
+		{
+			desc:          "Use depth with IPv4 subnet",
+			depth:         2,
+			xForwardedFor: "10.0.0.3,10.0.0.2,10.0.0.1",
+			expected:      "10.0.0.2",
+			ipv6Subnet:    intPtr(80),
+		},
+		{
+			desc:          "Use depth with IPv6 subnet",
+			depth:         2,
+			xForwardedFor: "10.0.0.3," + ipv6Basic + ",10.0.0.1",
+			expected:      "::abcd:0:0:0",
+			ipv6Subnet:    intPtr(80),
+		},
 	}
 
 	for _, test := range testCases {
 		t.Run(test.desc, func(t *testing.T) {
 			t.Parallel()
 
-			strategy := DepthStrategy{Depth: test.depth}
+			strategy := DepthStrategy{
+				Depth:      test.depth,
+				IPv6Subnet: test.ipv6Subnet,
+			}
 			req := httptest.NewRequest(http.MethodGet, "http://127.0.0.1", nil)
 			req.Header.Set(xForwardedFor, test.xForwardedFor)
 			actual := strategy.GetIP(req)
@@ -121,3 +197,7 @@ func TestTrustedIPsStrategy_GetIP(t *testing.T) {
 		})
 	}
 }
+
+func intPtr(value int) *int {
+	return &value
+}