package egress

import (
	"context"
	"errors"
	"net"
	"sync/atomic"
	"testing"

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

func TestGuardedDialerAllowsOnlyConfiguredPrivateCIDRs(t *testing.T) {
	_, allowed, err := net.ParseCIDR("10.20.0.0/16")
	require.NoError(t, err)

	testCases := []struct {
		name       string
		ips        []string
		wantDial   bool
		wantDenied bool
	}{
		{name: "configured private", ips: []string{"10.20.1.2"}, wantDial: true},
		{name: "unconfigured private", ips: []string{"10.21.1.2"}, wantDenied: true},
		{
			name:       "configured policy also confines public addresses",
			ips:        []string{"8.8.8.8"},
			wantDenied: true,
		},
		{
			name:       "mixed answer fails closed",
			ips:        []string{"10.20.1.2", "10.21.1.2"},
			wantDenied: true,
		},
	}
	for _, testCase := range testCases {
		t.Run(testCase.name, func(t *testing.T) {
			var dialCalls atomic.Int32
			resolved := make([]net.IP, 0, len(testCase.ips))
			for _, raw := range testCase.ips {
				resolved = append(resolved, net.ParseIP(raw))
			}
			dialer := newGuardedDialer(Config{
				TargetNetworks: []*net.IPNet{allowed},
				Resolver:       staticResolver{ips: resolved},
				DialContext: func(context.Context, string, string) (net.Conn, error) {
					dialCalls.Add(1)
					return nil, errors.New("dial attempted")
				},
			})
			_, err := dialer.DialContext(
				context.Background(), "tcp", "jenkins.example:443")
			if testCase.wantDenied {
				require.ErrorIs(t, err, ErrDestinationDenied)
				require.Zero(t, dialCalls.Load())
				return
			}
			require.ErrorContains(t, err, "dial attempted")
			require.Equal(t, int32(1), dialCalls.Load())
		})
	}
}

func TestGuardedDialerNeverAllowsHardDeniedAddresses(t *testing.T) {
	for _, rawIP := range []string{
		"0.0.0.0",
		"127.0.0.1",
		"169.254.169.254",
		"224.0.0.1",
		"100.64.0.1",
		"100.100.100.200",
		"168.63.129.16",
		"::",
		"::1",
		"fe80::1",
		"ff00::1",
		"fd00:ec2::254",
	} {
		t.Run(rawIP, func(t *testing.T) {
			var dialCalls atomic.Int32
			dialer := newGuardedDialer(Config{
				AllowPrivate: true,
				Resolver: staticResolver{
					ips: []net.IP{net.ParseIP(rawIP)},
				},
				DialContext: func(context.Context, string, string) (net.Conn, error) {
					dialCalls.Add(1)
					return nil, errors.New("must not dial")
				},
			})
			_, err := dialer.DialContext(
				context.Background(), "tcp", "jenkins.example:443")
			require.ErrorIs(t, err, ErrDestinationDenied)
			require.Zero(t, dialCalls.Load())
		})
	}
}

func TestGuardedDialerAllowsPrivateAddressesOnlyInConnectionMode(t *testing.T) {
	for _, rawIP := range []string{
		"8.8.8.8",
		"10.20.1.2",
		"172.20.1.2",
		"192.168.1.2",
		"fd12:3456::10",
	} {
		t.Run(rawIP, func(t *testing.T) {
			var dialCalls atomic.Int32
			dialer := newGuardedDialer(Config{
				AllowPrivate: true,
				Resolver: staticResolver{
					ips: []net.IP{net.ParseIP(rawIP)},
				},
				DialContext: func(context.Context, string, string) (net.Conn, error) {
					dialCalls.Add(1)
					return nil, errors.New("dial attempted")
				},
			})
			_, err := dialer.DialContext(
				context.Background(), "tcp", "jenkins.example:443")
			require.ErrorContains(t, err, "dial attempted")
			require.NotErrorIs(t, err, ErrDestinationDenied)
			require.Equal(t, int32(1), dialCalls.Load())
		})
	}
}

func TestGuardedDialerConnectionModeFailsClosedOnMixedHardDeniedAnswer(t *testing.T) {
	var dialCalls atomic.Int32
	dialer := newGuardedDialer(Config{
		AllowPrivate: true,
		Resolver: staticResolver{
			ips: []net.IP{
				net.ParseIP("10.20.1.2"),
				net.ParseIP("169.254.169.254"),
			},
		},
		DialContext: func(context.Context, string, string) (net.Conn, error) {
			dialCalls.Add(1)
			return nil, errors.New("must not dial")
		},
	})
	_, err := dialer.DialContext(
		context.Background(), "tcp", "jenkins.example:443")
	require.ErrorIs(t, err, ErrDestinationDenied)
	require.Zero(t, dialCalls.Load())
}

func TestGuardedDialerConnectionModeResolvesDNSOnEveryDial(t *testing.T) {
	var resolveCalls atomic.Int32
	var dialed []string
	dialer := newGuardedDialer(Config{
		AllowPrivate: true,
		Resolver: resolverFunc(func(
			context.Context,
			string,
			string,
		) ([]net.IP, error) {
			if resolveCalls.Add(1) == 1 {
				return []net.IP{net.ParseIP("10.20.1.2")}, nil
			}
			return []net.IP{net.ParseIP("192.168.1.2")}, nil
		}),
		DialContext: func(_ context.Context, _, address string) (net.Conn, error) {
			dialed = append(dialed, address)
			return nil, errors.New("dial attempted")
		},
	})

	for range 2 {
		_, err := dialer.DialContext(
			context.Background(), "tcp", "jenkins.example:443")
		require.ErrorContains(t, err, "dial attempted")
	}
	require.Equal(t, int32(2), resolveCalls.Load())
	require.Equal(t, []string{"10.20.1.2:443", "192.168.1.2:443"}, dialed)
}

func TestGuardedDialerClonesConfiguredCIDRs(t *testing.T) {
	_, allowed, err := net.ParseCIDR("10.20.0.0/16")
	require.NoError(t, err)
	dialer := newGuardedDialer(Config{
		TargetNetworks: []*net.IPNet{allowed},
		Resolver: staticResolver{
			ips: []net.IP{net.ParseIP("10.20.1.2")},
		},
		DialContext: func(context.Context, string, string) (net.Conn, error) {
			return nil, errors.New("dial attempted")
		},
	})

	allowed.IP[0] = 11
	_, err = dialer.DialContext(
		context.Background(), "tcp", "jenkins.example:443")
	require.ErrorContains(t, err, "dial attempted")
	require.NotErrorIs(t, err, ErrDestinationDenied)
}

type resolverFunc func(context.Context, string, string) ([]net.IP, error)

func (resolve resolverFunc) LookupIP(
	ctx context.Context,
	network string,
	host string,
) ([]net.IP, error) {
	return resolve(ctx, network, host)
}