package httpz

import (
	"net/netip"
	"testing"

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

func p[T any](x T) *T {
	return &x
}

func TestSafeTransport_Control(t *testing.T) {
	control := NewSafeNetDialerControl(&SafeNetControlConfig{})

	testcases := []struct {
		name          string
		network       string
		address       string
		expectedError *string
	}{
		{
			name:          "invalid network type",
			network:       "udp4",
			expectedError: p("invalid network. Wanted tcp4 or tcp6 but got \"udp4\""),
		},
		{
			name:          "invalid address",
			network:       "tcp4",
			address:       "foo",
			expectedError: p("invalid address. Wanted a valid host:port pair but got \"foo\""),
		},
		{
			name:          "invalid port",
			network:       "tcp4",
			address:       "0.0.0.0:foo",
			expectedError: p("invalid address. Wanted a valid host:port pair but got \"0.0.0.0:foo\""),
		},
		{
			name:          "invalid port, negative number",
			network:       "tcp4",
			address:       "0.0.0.0:-1",
			expectedError: p("invalid address. Wanted a valid host:port pair but got \"0.0.0.0:-1\""),
		},
		{
			name:          "invalid ip address",
			network:       "tcp4",
			address:       "0.0.0:80",
			expectedError: p("invalid address. Wanted a valid host:port pair but got \"0.0.0:80\""),
		},
		{
			name:          "invalid ip address, is unresolved hostname",
			network:       "tcp4",
			address:       "example.com:80",
			expectedError: p("invalid address. Wanted a valid host:port pair but got \"example.com:80\""),
		},
		{
			name:          "invalid port",
			network:       "tcp4",
			address:       "0.0.0.0:22",
			expectedError: p("invalid port. wanted one of [80 443] but got 22"),
		},
		{
			name:          "not allowed ip address, unspecified IPv4",
			network:       "tcp4",
			address:       "0.0.0.0:80",
			expectedError: p("not allowed IP address. IP address \"0.0.0.0\" is not allowed"),
		},
		{
			name:          "not allowed ip address, unspecified IPv6",
			network:       "tcp6",
			address:       "[::]:80",
			expectedError: p("not allowed IP address. IP address \"::\" is not allowed"),
		},
		{
			name:          "not allowed ip address, loopback IPv4 (RFC 1122)",
			network:       "tcp4",
			address:       "127.0.0.1:80",
			expectedError: p("not allowed IP address. IP address \"127.0.0.1\" is not allowed"),
		},
		{
			name:          "not allowed ip address, loopback IPv6 (RFC 4291)",
			network:       "tcp6",
			address:       "[::1]:80",
			expectedError: p("not allowed IP address. IP address \"::1\" is not allowed"),
		},
		{
			name:          "not allowed ip address, link-local IPv4 (RFC 3927)",
			network:       "tcp4",
			address:       "169.254.0.1:80",
			expectedError: p("not allowed IP address. IP address \"169.254.0.1\" is not allowed"),
		},
		{
			name:          "not allowed ip address, link-local IPv6 (RFC 4291)",
			network:       "tcp6",
			address:       "[fe80::1]:80",
			expectedError: p("not allowed IP address. IP address \"fe80::1\" is not allowed"),
		},
		{
			name:          "not allowed ip address, multicast IPv4 (RFC 1112)",
			network:       "tcp4",
			address:       "224.0.0.1:80",
			expectedError: p("not allowed IP address. IP address \"224.0.0.1\" is not allowed"),
		},
		{
			name:          "not allowed ip address, multicast IPv6 (RFC 4291)",
			network:       "tcp6",
			address:       "[ff02::1]:80",
			expectedError: p("not allowed IP address. IP address \"ff02::1\" is not allowed"),
		},
		{
			name:          "not allowed ip address, private IPv4 10.0.0.0/8 (first) (RFC 1918)",
			network:       "tcp4",
			address:       "10.0.0.1:80",
			expectedError: p("not allowed IP address. IP address \"10.0.0.1\" is not allowed"),
		},
		{
			name:          "not allowed ip address, private IPv4 10.0.0.0/8 (last) (RFC 1918)",
			network:       "tcp4",
			address:       "10.255.255.255:80",
			expectedError: p("not allowed IP address. IP address \"10.255.255.255\" is not allowed"),
		},
		{
			name:          "not allowed ip address, private IPv4 172.16.0.0/12 (first) (RFC 1918)",
			network:       "tcp4",
			address:       "172.16.0.1:80",
			expectedError: p("not allowed IP address. IP address \"172.16.0.1\" is not allowed"),
		},
		{
			name:          "not allowed ip address, private IPv4 172.16.0.0/12 (last) (RFC 1918)",
			network:       "tcp4",
			address:       "172.31.255.255:80",
			expectedError: p("not allowed IP address. IP address \"172.31.255.255\" is not allowed"),
		},
		{
			name:          "not allowed ip address, private IPv4 192.168.0.0/16 (first) (RFC 1918)",
			network:       "tcp4",
			address:       "192.168.0.1:80",
			expectedError: p("not allowed IP address. IP address \"192.168.0.1\" is not allowed"),
		},
		{
			name:          "not allowed ip address, private IPv4 192.168.0.0/16 (last) (RFC 1918)",
			network:       "tcp4",
			address:       "192.168.255.255:80",
			expectedError: p("not allowed IP address. IP address \"192.168.255.255\" is not allowed"),
		},
		{
			name:          "not allowed ip address, private (Unique Local Addresses) IPv6 fc00::/7 (first) (RFC 4193)",
			network:       "tcp6",
			address:       "[fc00::1]:80",
			expectedError: p("not allowed IP address. IP address \"fc00::1\" is not allowed"),
		},
		{
			name:          "not allowed ip address, private (Unique Local Addresses) IPv6 fc00::/7 (first) (RFC 4193)",
			network:       "tcp6",
			address:       "[fdff:ffff:ffff:ffff:ffff:ffff:ffff:ffff]:80",
			expectedError: p("not allowed IP address. IP address \"fdff:ffff:ffff:ffff:ffff:ffff:ffff:ffff\" is not allowed"),
		},
		{
			name:          "not allowed ip address, IPv4-mapped IPv6 address (RFC 4291)", // They are actually mapped in the Go stdlib when parsing the IP.
			network:       "tcp6",
			address:       "[::ffff:169.254.169.254]:80", // it's a link-local address
			expectedError: p("not allowed IP address. IP address \"::ffff:169.254.169.254\" is not allowed"),
		},
		{
			name:    "allowed public IPv4 address, random address",
			network: "tcp4",
			address: "142.250.181.238:443", // random address google.com resolved to, dig +short google.com A
		},
		{
			name:    "allowed public IPv6 address, random address",
			network: "tcp4",
			address: "[2a00:1450:4001:82f::200e]:443", // random address google.com resolved to, dig +short google.com AAAA
		},
	}

	for _, tc := range testcases {
		t.Run(tc.name, func(t *testing.T) {
			err := control(tc.network, tc.address, nil)
			if tc.expectedError == nil {
				assert.NoError(t, err)
			} else {
				assert.EqualError(t, err, *tc.expectedError)
			}
		})
	}
}

func TestSafeTransport_Allow(t *testing.T) {
	testcases := []struct {
		name           string
		allowedIPs     []netip.Addr
		allowedIPCIDRs []netip.Prefix
		network        string
		address        string
	}{
		{
			name:       "allow loopback interface by explicit IP",
			allowedIPs: []netip.Addr{netip.MustParseAddr("127.0.0.1")},
			network:    "tcp4",
			address:    "127.0.0.1:80",
		},
		{
			name:           "allow private IPv4 range by CIDR",
			allowedIPCIDRs: []netip.Prefix{netip.MustParsePrefix("10.0.0.0/8")},
			network:        "tcp4",
			address:        "10.10.10.10:80",
		},
	}

	for _, tc := range testcases {
		t.Run(tc.name, func(t *testing.T) {
			control := NewSafeNetDialerControl(&SafeNetControlConfig{
				AllowedIPs:     tc.allowedIPs,
				AllowedIPCIDRs: tc.allowedIPCIDRs,
			})

			err := control(tc.network, tc.address, nil)
			assert.NoError(t, err)
		})
	}
}

func TestSafeTransport_Block(t *testing.T) {
	testcases := []struct {
		name           string
		blockedIPs     []netip.Addr
		blockedIPCIDRs []netip.Prefix
		network        string
		address        string
		expectedError  string
	}{
		{
			name:          "block single public IPv4 address",
			blockedIPs:    []netip.Addr{netip.MustParseAddr("142.250.181.238")}, // random address google.com resolved to, dig +short google.com A
			network:       "tcp4",
			address:       "142.250.181.238:443",
			expectedError: "not allowed IP address. IP address \"142.250.181.238\" is blocked",
		},
		{
			name:           "block public IPv4 CIDR",
			blockedIPCIDRs: []netip.Prefix{netip.MustParsePrefix("142.250.0.0/16")},
			network:        "tcp4",
			address:        "142.250.181.238:443",
			expectedError:  "not allowed IP address. IP address \"142.250.181.238\" is blocked",
		},
	}

	for _, tc := range testcases {
		t.Run(tc.name, func(t *testing.T) {
			control := NewSafeNetDialerControl(&SafeNetControlConfig{
				BlockedIPs:     tc.blockedIPs,
				BlockedIPCIDRs: tc.blockedIPCIDRs,
			})

			err := control(tc.network, tc.address, nil)
			assert.EqualError(t, err, tc.expectedError)
		})
	}
}

func TestSafeTransport_BlockBeforeAllow(t *testing.T) {
	testcases := []struct {
		name           string
		allowedIPs     []netip.Addr
		allowedIPCIDRs []netip.Prefix
		blockedIPs     []netip.Addr
		blockedIPCIDRs []netip.Prefix
		network        string
		address        string
		expectedError  string
	}{
		{
			name:          "block single public IPv4 address before allow",
			allowedIPs:    []netip.Addr{netip.MustParseAddr("142.250.181.238")}, // random address google.com resolved to, dig +short google.com A
			blockedIPs:    []netip.Addr{netip.MustParseAddr("142.250.181.238")}, // random address google.com resolved to, dig +short google.com A
			network:       "tcp4",
			address:       "142.250.181.238:443",
			expectedError: "not allowed IP address. IP address \"142.250.181.238\" is blocked",
		},
		{
			name:           "block public IPv4 CIDR before allow",
			allowedIPCIDRs: []netip.Prefix{netip.MustParsePrefix("142.250.0.0/16")},
			blockedIPCIDRs: []netip.Prefix{netip.MustParsePrefix("142.250.0.0/16")},
			network:        "tcp4",
			address:        "142.250.181.238:443",
			expectedError:  "not allowed IP address. IP address \"142.250.181.238\" is blocked",
		},
	}

	for _, tc := range testcases {
		t.Run(tc.name, func(t *testing.T) {
			control := NewSafeNetDialerControl(&SafeNetControlConfig{
				AllowedIPs:     tc.allowedIPs,
				AllowedIPCIDRs: tc.allowedIPCIDRs,
				BlockedIPs:     tc.blockedIPs,
				BlockedIPCIDRs: tc.blockedIPCIDRs,
			})

			err := control(tc.network, tc.address, nil)
			assert.EqualError(t, err, tc.expectedError)
		})
	}
}
