package pipeline

import (
	"context"
	"errors"
	"net"
	"net/http"
	"net/http/httptest"
	"net/netip"
	"net/url"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
)

func TestCDNDialPinsAndChecksEveryAnswer(t *testing.T) {
	t.Setenv("EVOBGP_CDN_ALLOW_PRIVATE", "0")
	answers := [][]net.IP{
		{net.ParseIP("203.0.113.10")},
		{net.ParseIP("127.0.0.1")},
	}
	lookups := 0
	var dialed []string
	lookup := func(context.Context, string) ([]net.IP, error) {
		answer := answers[lookups]
		lookups++
		return answer, nil
	}
	dial := func(_ context.Context, _, address string) (net.Conn, error) {
		dialed = append(dialed, address)
		return nil, nil
	}
	if _, err := cdnDial(context.Background(), "tcp", "cdn.example:443", lookup, dial); err != nil {
		t.Fatal(err)
	}
	if _, err := cdnDial(context.Background(), "tcp", "cdn.example:443", lookup, dial); err == nil {
		t.Fatal("changed DNS answer must be rejected before dial")
	}
	if lookups != 2 || len(dialed) != 1 || dialed[0] != "203.0.113.10:443" {
		t.Fatalf("lookups=%d dialed=%v", lookups, dialed)
	}
}

func TestCDNDialRejectsMixedAndMappedPrivateAnswers(t *testing.T) {
	t.Setenv("EVOBGP_CDN_ALLOW_PRIVATE", "0")
	tests := []struct {
		name string
		ips  []net.IP
	}{
		{"mixed", []net.IP{net.ParseIP("203.0.113.10"), net.ParseIP("10.0.0.5")}},
		{"mapped loopback", []net.IP{net.ParseIP("::ffff:127.0.0.1")}},
		{"mapped private", []net.IP{net.ParseIP("::ffff:10.0.0.1")}},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			called := false
			_, err := cdnDial(context.Background(), "tcp", "cdn.example:443", func(context.Context, string) ([]net.IP, error) {
				return tc.ips, nil
			}, func(context.Context, string, string) (net.Conn, error) {
				called = true
				return nil, nil
			})
			if err == nil || called {
				t.Fatalf("want rejection before dial, got err=%v called=%v", err, called)
			}
		})
	}
	if !isBlockedCDNIP(netip.MustParseAddr("::ffff:169.254.169.254")) {
		t.Fatal("mapped metadata address must be blocked")
	}
}

func TestCDNDialFallsBackOnlyToCheckedIP(t *testing.T) {
	t.Setenv("EVOBGP_CDN_ALLOW_PRIVATE", "0")
	var dialed []string
	_, err := cdnDial(context.Background(), "tcp", "cdn.example:443", func(context.Context, string) ([]net.IP, error) {
		return []net.IP{net.ParseIP("2001:db8::1"), net.ParseIP("203.0.113.10")}, nil
	}, func(_ context.Context, _, address string) (net.Conn, error) {
		dialed = append(dialed, address)
		if len(dialed) == 1 {
			return nil, errors.New("unreachable")
		}
		return nil, nil
	})
	if err != nil || len(dialed) != 2 || dialed[0] != "[2001:db8::1]:443" || dialed[1] != "203.0.113.10:443" {
		t.Fatalf("err=%v dialed=%v", err, dialed)
	}
}

func TestCDNDialStopsOnCancellation(t *testing.T) {
	t.Setenv("EVOBGP_CDN_ALLOW_PRIVATE", "0")
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	called := 0
	_, err := cdnDial(ctx, "tcp", "cdn.example:443", func(context.Context, string) ([]net.IP, error) {
		return []net.IP{net.ParseIP("203.0.113.10"), net.ParseIP("203.0.113.11")}, nil
	}, func(context.Context, string, string) (net.Conn, error) {
		called++
		cancel()
		return nil, errors.New("unreachable")
	})
	if !errors.Is(err, context.Canceled) || called != 1 {
		t.Fatalf("err=%v dial attempts=%d", err, called)
	}
}

func TestCDNClientRejectsUnsafeRedirectsAndProxy(t *testing.T) {
	t.Setenv("EVOBGP_CDN_ALLOW_PRIVATE", "0")
	t.Setenv("HTTPS_PROXY", "http://127.0.0.1:8080")
	client := NewCDNHTTPClient()
	tr, ok := client.Transport.(*cdnRoundTripper)
	if !ok || tr.transport.Proxy != nil {
		t.Fatal("CDN transport must dial directly")
	}
	for _, raw := range []string{
		"https://127.0.0.1/secret",
		"https://[::ffff:127.0.0.1]/secret",
		"http://cdn.example/unsafe",
		"https://user:password@cdn.example/secret",
	} {
		t.Run(raw, func(t *testing.T) {
			u, err := url.Parse(raw)
			if err != nil {
				t.Fatal(err)
			}
			if err := client.CheckRedirect(&http.Request{URL: u}, nil); err == nil {
				t.Fatal("unsafe redirect accepted")
			}
		})
	}
}

type cdnTestRoundTripper func(*http.Request) (*http.Response, error)

func (f cdnTestRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) { return f(req) }

func TestCDNClientRejectsOpaqueTransport(t *testing.T) {
	_, err := secureCDNClient(&http.Client{Transport: cdnTestRoundTripper(func(*http.Request) (*http.Response, error) {
		return nil, errors.New("must not be called")
	})})
	if err == nil {
		t.Fatal("opaque transport could bypass checked dial")
	}
}

func TestCDNClientRejectsPrivateTLSRedirectBeforeDial(t *testing.T) {
	t.Setenv("EVOBGP_CDN_ALLOW_PRIVATE", "0")
	var hits atomic.Int32
	var secretHits atomic.Int32
	srv := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		hits.Add(1)
		if r.URL.Path == "/secret" {
			secretHits.Add(1)
		}
		http.Redirect(w, r, "https://private.example.com/secret", http.StatusFound)
	}))
	defer srv.Close()

	client := NewCDNHTTPClient()
	tr := client.Transport.(*cdnRoundTripper).transport
	tr.TLSClientConfig = srv.Client().Transport.(*http.Transport).TLSClientConfig.Clone()
	actual := srv.Listener.Addr().String()
	var dialed []string
	var dialedMu sync.Mutex
	dialer := &net.Dialer{}
	tr.DialContext = func(ctx context.Context, network, address string) (net.Conn, error) {
		return cdnDial(ctx, network, address, func(_ context.Context, host string) ([]net.IP, error) {
			if host == "private.example.com" {
				return []net.IP{net.ParseIP("127.0.0.1")}, nil
			}
			return []net.IP{net.ParseIP("203.0.113.10")}, nil
		}, func(ctx context.Context, network, checkedAddress string) (net.Conn, error) {
			dialedMu.Lock()
			dialed = append(dialed, checkedAddress)
			dialedMu.Unlock()
			return dialer.DialContext(ctx, network, actual)
		})
	}
	_, port, err := net.SplitHostPort(actual)
	if err != nil {
		t.Fatal(err)
	}
	resp, err := client.Get("https://cdn.example.com:" + port + "/redirect")
	if resp != nil {
		_ = resp.Body.Close()
	}
	if err == nil || !strings.Contains(err.Error(), "blocked address") {
		t.Fatalf("private TLS redirect was not rejected: %v", err)
	}
	dialedMu.Lock()
	dialedSnapshot := append([]string(nil), dialed...)
	dialedMu.Unlock()
	if hits.Load() != 1 || secretHits.Load() != 0 || len(dialedSnapshot) != 1 || dialedSnapshot[0] != "203.0.113.10:"+port {
		t.Fatalf("hits=%d secretHits=%d dialed=%v", hits.Load(), secretHits.Load(), dialedSnapshot)
	}
}