traefik/pkg/ip/strategy_test.go

package ip

import (
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestRemoteAddrStrategy_GetIP(t *testing.T) {
	testCases := []struct {
		desc     string
		expected string
	}{
		{
			desc:     "Use RemoteAddr",
			expected: "192.0.2.1",
		},
	}

	for _, test := range testCases {
		t.Run(test.desc, func(t *testing.T) {
			t.Parallel()

			strategy := RemoteAddrStrategy{}
			req := httptest.NewRequest(http.MethodGet, "http://127.0.0.1", nil)
			actual := strategy.GetIP(req)
			assert.Equal(t, test.expected, actual)
		})
	}
}

func TestDepthStrategy_GetIP(t *testing.T) {
	testCases := []struct {
		desc          string
		depth         int
		xForwardedFor string
		expected      string
	}{
		{
			desc:          "Use depth",
			depth:         3,
			xForwardedFor: "10.0.0.4,10.0.0.3,10.0.0.2,10.0.0.1",
			expected:      "10.0.0.3",
		},
		{
			desc:          "Use nonexistent depth in XForwardedFor",
			depth:         2,
			xForwardedFor: "",
			expected:      "",
		},
		{
			desc:          "Use depth that match the first IP in XForwardedFor",
			depth:         2,
			xForwardedFor: "10.0.0.2,10.0.0.1",
			expected:      "10.0.0.2",
		},
	}

	for _, test := range testCases {
		t.Run(test.desc, func(t *testing.T) {
			t.Parallel()

			strategy := DepthStrategy{Depth: test.depth}
			req := httptest.NewRequest(http.MethodGet, "http://127.0.0.1", nil)
			req.Header.Set(xForwardedFor, test.xForwardedFor)
			actual := strategy.GetIP(req)
			assert.Equal(t, test.expected, actual)
		})
	}
}

func TestTrustedIPsStrategy_GetIP(t *testing.T) {
	testCases := []struct {
		desc          string
		trustedIPs    []string
		xForwardedFor string
		expected      string
		useRemote     bool
	}{
		{
			desc:          "Trust all IPs",
			trustedIPs:    []string{"10.0.0.4", "10.0.0.3", "10.0.0.2", "10.0.0.1"},
			xForwardedFor: "10.0.0.4,10.0.0.3,10.0.0.2,10.0.0.1",
			expected:      "",
		},
		{
			desc:          "Do not trust all IPs",
			trustedIPs:    []string{"10.0.0.2", "10.0.0.1"},
			xForwardedFor: "10.0.0.4,10.0.0.3,10.0.0.2,10.0.0.1",
			expected:      "10.0.0.3",
		},
		{
			desc:          "Do not trust all IPs with CIDR",
			trustedIPs:    []string{"10.0.0.1/24"},
			xForwardedFor: "127.0.0.1,10.0.0.4,10.0.0.3,10.0.0.2,10.0.0.1",
			expected:      "127.0.0.1",
		},
		{
			desc:          "Trust all IPs with CIDR",
			trustedIPs:    []string{"10.0.0.1/24"},
			xForwardedFor: "10.0.0.4,10.0.0.3,10.0.0.2,10.0.0.1",
			expected:      "",
		},
	}

	for _, test := range testCases {
		t.Run(test.desc, func(t *testing.T) {
			t.Parallel()

			checker, err := NewChecker(test.trustedIPs)
			require.NoError(t, err)

			strategy := PoolStrategy{Checker: checker}
			req := httptest.NewRequest(http.MethodGet, "http://127.0.0.1", nil)
			req.Header.Set(xForwardedFor, test.xForwardedFor)
			actual := strategy.GetIP(req)
			assert.Equal(t, test.expected, actual)
		})
	}
}