package main

import (
	"context"
	"errors"
	"net"
	"net/netip"
	"testing"
)

func testPolicy(answer func(host string) []netip.Addr) *Policy {
	policy := NewPolicy()
	policy.lookup = func(ctx context.Context, host string) ([]netip.Addr, error) {
		return answer(host), nil
	}
	return policy
}

func publicAnswer(string) []netip.Addr {
	return []netip.Addr{netip.MustParseAddr("93.184.216.34")}
}

func codeOf(err error) string {
	var policyErr *PolicyError
	if errors.As(err, &policyErr) {
		return policyErr.Code
	}
	if err == nil {
		return ""
	}
	return "other:" + err.Error()
}

func TestCheckTable(t *testing.T) {
	policy := testPolicy(publicAnswer)
	cases := []struct {
		name   string
		url    string
		method string
		mode   string
		want   string
	}{
		{"apex", "https://reddit.com/r/go", "GET", "fetch", ""},
		{"subdomain", "https://old.reddit.com/r/go.json", "get", "fetch", ""},
		{"short domain", "https://redd.it/abc", "GET", "any", ""},
		{"explicit 443", "https://www.reddit.com:443/", "GET", "fetch", ""},
		{"plain http 80", "http://reddit.com/", "GET", "fetch", ""},
		{"trailing dot", "https://reddit.com./", "GET", "fetch", ""},
		{"upper case host", "https://WWW.REDDIT.COM/", "GET", "fetch", ""},
		{"prefix attack", "https://evilreddit.com/", "GET", "fetch", RefuseHost},
		{"suffix attack", "https://reddit.com.evil.net/", "GET", "fetch", RefuseHost},
		{"hyphen attack", "https://not-reddit.com/", "GET", "fetch", RefuseHost},
		{"other host", "https://example.com/", "GET", "fetch", RefuseHost},
		{"port 8080", "https://reddit.com:8080/", "GET", "fetch", RefusePort},
		{"port 22", "http://reddit.com:22/", "GET", "fetch", RefusePort},
		{"port 0", "https://reddit.com:0/", "GET", "fetch", RefusePort},
		{"ipv4 literal", "http://127.0.0.1/", "GET", "fetch", RefuseHost},
		{"ipv4 literal 80", "http://93.184.216.34/", "GET", "fetch", RefuseHost},
		{"ipv6 literal", "http://[::1]/", "GET", "fetch", RefuseHost},
		{"credentials", "https://user:pw@reddit.com/", "GET", "fetch", RefuseBadURL},
		{"scheme ftp", "ftp://reddit.com/", "GET", "fetch", RefuseBadURL},
		{"scheme file", "file:///etc/passwd", "GET", "fetch", RefuseBadURL},
		{"empty", "", "GET", "fetch", RefuseBadURL},
		{"post", "https://reddit.com/", "POST", "fetch", RefuseMethod},
		{"delete", "https://reddit.com/", "DELETE", "fetch", RefuseMethod},
		{"bad port text", "https://reddit.com:abc/", "GET", "fetch", RefuseBadURL},
	}
	for _, item := range cases {
		_, err := policy.Check(item.url, item.method, item.mode)
		if got := codeOf(err); got != item.want {
			t.Errorf("%s: got %q want %q", item.name, got, item.want)
		}
	}
	long := "https://reddit.com/" + string(make([]byte, MaxURLChars))
	if _, err := policy.Check(long, "GET", "fetch"); codeOf(err) != RefuseBadURL {
		t.Errorf("long url not refused: %v", err)
	}
}

func TestCheckHostPortModes(t *testing.T) {
	policy := testPolicy(publicAnswer)
	if _, err := policy.CheckHostPort("reddit.com", 443, "tunnel"); err != nil {
		t.Fatal(err)
	}
	if _, err := policy.CheckHostPort("reddit.com", 25, "tunnel"); codeOf(err) != RefusePort {
		t.Fatalf("port: %v", err)
	}
	if _, err := policy.CheckHostPort("evilreddit.com", 443, "tunnel"); codeOf(err) != RefuseHost {
		t.Fatalf("host: %v", err)
	}
}

func TestPublicAddressTable(t *testing.T) {
	cases := map[string]bool{
		"8.8.8.8": true, "93.184.216.34": true, "2606:4700:4700::1111": true,
		"10.0.0.1": false, "10.255.255.255": false, "127.0.0.1": false, "127.1.2.3": false,
		"169.254.169.254": false, "172.16.0.1": false, "172.31.255.255": false, "172.32.0.1": true,
		"192.168.1.1": false, "100.64.0.1": false, "100.127.255.255": false, "100.128.0.1": true,
		"0.0.0.0": false, "0.1.2.3": false, "224.0.0.1": false, "239.255.255.250": false,
		"255.255.255.255": false, "240.0.0.1": false, "198.18.0.1": false, "192.0.2.1": false,
		"::": false, "::1": false, "fe80::1": false, "fc00::1": false, "fd12:3456::1": false,
		"ff02::1": false, "::ffff:127.0.0.1": false, "::ffff:10.0.0.1": false, "::ffff:8.8.8.8": true,
		"2001:db8::1": false, "64:ff9b::7f00:1": false, "2002:7f00:1::1": false, "2001::1": false,
	}
	for text, want := range cases {
		if got := IsPublicAddr(netip.MustParseAddr(text)); got != want {
			t.Errorf("%s: got %v want %v", text, got, want)
		}
	}
	if IsPublicAddr(netip.Addr{}) {
		t.Error("invalid address counted as public")
	}
	if IsPublicAddr(netip.MustParseAddr("fe80::1%eth0")) {
		t.Error("zoned address counted as public")
	}
}

func TestResolveRefusesAnyPrivateAnswer(t *testing.T) {
	cases := map[string][]string{
		"loopback":    {"127.0.0.1"},
		"private":     {"10.1.2.3"},
		"mixed":       {"93.184.216.34", "192.168.0.1"},
		"v6 loopback": {"::1"},
		"v6 mapped":   {"::ffff:10.0.0.1"},
		"cgnat":       {"100.64.5.5"},
	}
	for name, list := range cases {
		addrs := []netip.Addr{}
		for _, item := range list {
			addrs = append(addrs, netip.MustParseAddr(item))
		}
		policy := testPolicy(func(string) []netip.Addr { return addrs })
		if _, err := policy.Resolve(context.Background(), "reddit.com"); codeOf(err) != RefuseAddress {
			t.Errorf("%s: %v", name, err)
		}
	}
	empty := testPolicy(func(string) []netip.Addr { return nil })
	if _, err := empty.Resolve(context.Background(), "reddit.com"); codeOf(err) != FailDNS {
		t.Errorf("empty answer: %v", err)
	}
}

func TestDialConnectsToValidatedIPAndDefeatsRebinding(t *testing.T) {
	calls := 0
	policy := testPolicy(func(string) []netip.Addr {
		calls++
		if calls == 1 {
			return []netip.Addr{netip.MustParseAddr("93.184.216.34")}
		}
		return []netip.Addr{netip.MustParseAddr("127.0.0.1")}
	})
	var dialed []string
	policy.dial = func(ctx context.Context, network, address string) (net.Conn, error) {
		dialed = append(dialed, address)
		client, server := net.Pipe()
		server.Close()
		return client, nil
	}
	if _, err := policy.DialContext(context.Background(), "tcp", "www.reddit.com:443"); err != nil {
		t.Fatal(err)
	}
	if len(dialed) != 1 || dialed[0] != "93.184.216.34:443" {
		t.Fatalf("dialed %v", dialed)
	}
	if _, err := policy.DialContext(context.Background(), "tcp", "www.reddit.com:443"); codeOf(err) != RefuseAddress {
		t.Fatalf("rebinding not refused: %v", err)
	}
	if len(dialed) != 1 {
		t.Fatalf("private address was dialed: %v", dialed)
	}
	if _, err := policy.DialContext(context.Background(), "tcp", "evil.com:443"); codeOf(err) != RefuseHost {
		t.Fatalf("host: %v", err)
	}
	if _, err := policy.DialContext(context.Background(), "tcp", "reddit.com:22"); codeOf(err) != RefusePort {
		t.Fatalf("port: %v", err)
	}
}

func TestRedirectLeavingAllowlistRefused(t *testing.T) {
	policy := testPolicy(publicAnswer)
	cases := map[string]string{
		"https://old.reddit.com/x":   "",
		"https://redd.it/y":          "",
		"https://evil.com/":          RefuseRedirect,
		"https://evilreddit.com/":    RefuseRedirect,
		"http://reddit.com:8080/":    RefuseRedirect,
		"http://127.0.0.1/":          RefuseRedirect,
		"ftp://reddit.com/":          RefuseRedirect,
		"https://reddit.com.evil.io": RefuseRedirect,
	}
	for location, want := range cases {
		_, err := policy.CheckRedirect(location, "GET")
		if got := codeOf(err); got != want {
			t.Errorf("%s: got %q want %q", location, got, want)
		}
	}
	if _, err := policy.CheckRedirect("https://reddit.com/", "POST"); codeOf(err) != RefuseRedirect {
		t.Error("method change through redirect allowed")
	}
}

func TestRefineOnlyNarrows(t *testing.T) {
	policy := testPolicy(publicAnswer)
	rejected := policy.Refine([]Destination{
		{Suffix: "old.reddit.com", Methods: []string{"GET", "POST"}, Ports: []int{80, 443, 8080}, Modes: []string{"fetch", "any"}},
		{Suffix: "redd.it", Methods: []string{"GET"}, Ports: []int{443}, Modes: []string{"tunnel", "fetch"}},
		{Suffix: "evil.com", Methods: []string{"GET"}, Ports: []int{443}, Modes: []string{"fetch"}},
		{Suffix: "evilreddit.com", Methods: []string{"GET"}, Ports: []int{443}, Modes: []string{"fetch"}},
		{Suffix: "old.reddit.com", Methods: []string{"GET"}, Ports: []int{443}, Modes: []string{"fetch"}},
	})
	if len(rejected) != 3 {
		t.Fatalf("rejected %v", rejected)
	}
	if _, err := policy.Check("https://old.reddit.com/", "GET", "fetch"); err != nil {
		t.Fatal(err)
	}
	if _, err := policy.Check("https://old.reddit.com/", "POST", "fetch"); codeOf(err) != RefuseMethod {
		t.Errorf("post widened: %v", err)
	}
	if _, err := policy.Check("http://old.reddit.com:8080/", "GET", "fetch"); codeOf(err) != RefusePort {
		t.Errorf("port widened: %v", err)
	}
	if _, err := policy.Check("https://www.reddit.com/", "GET", "fetch"); codeOf(err) != RefuseHost {
		t.Errorf("dropped destination still allowed: %v", err)
	}
	if _, err := policy.Check("http://redd.it/a", "GET", "fetch"); codeOf(err) != RefusePort {
		t.Errorf("redd.it port 80 should be gone: %v", err)
	}
	policy.Refine(nil)
	if _, err := policy.Check("https://redd.it/a", "GET", "fetch"); codeOf(err) != RefuseHost {
		t.Errorf("empty refinement must allow nothing: %v", err)
	}
}