package proxy

import (
	"bufio"
	"context"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/netip"
	"net/url"
	"slices"
	"strings"
	"testing"
)

func TestIsPrivate(t *testing.T) {
	for addr, want := range map[string]bool{
		"127.0.0.1":        true,
		"10.1.2.3":         true,
		"169.254.169.254":  true,
		"100.64.0.1":       true,
		"192.168.1.1":      true,
		"0.0.0.0":          true,
		"::1":              true,
		"::ffff:127.0.0.1": true,
		"fe80::1":          true,
		"fd00::1":          true,
		"192.0.0.8":        true,
		"198.18.0.1":       true,
		"240.0.0.1":        true,
		"255.255.255.255":  true,
		// IPv6 addresses carrying a private IPv4 one.
		"64:ff9b::a9fe:a9fe": true, // NAT64, 169.254.169.254
		"64:ff9b::808:808":   false,
		"2002:7f00:1::":      true, // 6to4, 127.0.0.1
		"2002:808:808::":     false,
		"::a00:1":            true, // IPv4-compatible 10.0.0.1
		"fec0::1":            true,
		"8.8.8.8":            false,
		"2606:4700::1111":    false,
	} {
		if got := IsPrivate(netip.MustParseAddr(addr)); got != want {
			t.Errorf("IsPrivate(%s) = %v", addr, got)
		}
	}
}

// start runs a proxy that allows only the hosts given, and an upstream
// HTTP server; it returns their addresses.
func start(t *testing.T, allowPrivate bool, allowed ...string) (proxyAddr string, upstreamPort int) {
	up, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	upSrv := &http.Server{Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { fmt.Fprint(w, "upstream") })}
	go upSrv.Serve(up)
	t.Cleanup(func() { upSrv.Close() })

	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	port := up.Addr().(*net.TCPAddr).Port
	s := &Server{
		Check:        func(_ context.Context, host string) (bool, string) { return slices.Contains(allowed, host), "" },
		Ports:        []int{port},
		AllowPrivate: allowPrivate,
	}
	go s.Serve(l)
	t.Cleanup(s.Close)
	return l.Addr().String(), port
}

func get(t *testing.T, proxyAddr, target string) (int, string) {
	u, _ := url.Parse("http://" + proxyAddr)
	c := &http.Client{Transport: &http.Transport{Proxy: http.ProxyURL(u)}}
	resp, err := c.Get(target)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()
	body, _ := io.ReadAll(resp.Body)
	return resp.StatusCode, string(body)
}

func TestHTTP(t *testing.T) {
	addr, port := start(t, true, "localhost")
	if code, body := get(t, addr, fmt.Sprintf("http://localhost:%d/", port)); code != 200 || body != "upstream" {
		t.Errorf("allowed host: %d %q", code, body)
	}
	if code, body := get(t, addr, fmt.Sprintf("http://127.0.0.1:%d/", port)); code != 403 || !strings.HasPrefix(body, Marker) {
		t.Errorf("unlisted host: %d %q", code, body)
	}
	if code, _ := get(t, addr, "http://localhost:1/"); code != 403 {
		t.Errorf("unlisted port: %d", code)
	}
}

func TestPrivateRefused(t *testing.T) {
	addr, port := start(t, false, "localhost")
	if code, body := get(t, addr, fmt.Sprintf("http://localhost:%d/", port)); code != 403 || !strings.Contains(body, "private") {
		t.Errorf("private address: %d %q", code, body)
	}
}

func TestConnect(t *testing.T) {
	addr, port := start(t, true, "localhost")
	conn, err := net.Dial("tcp", addr)
	if err != nil {
		t.Fatal(err)
	}
	defer conn.Close()
	fmt.Fprintf(conn, "CONNECT localhost:%d HTTP/1.1\r\nHost: localhost:%d\r\n\r\n", port, port)
	br := bufio.NewReader(conn)
	status, err := br.ReadString('\n')
	if err != nil || !strings.Contains(status, "200") {
		t.Fatalf("CONNECT: %q %v", status, err)
	}
	br.ReadString('\n') // the blank line ending the response
	fmt.Fprintf(conn, "GET / HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\n\r\n")
	rest, _ := io.ReadAll(br)
	if !strings.Contains(string(rest), "upstream") {
		t.Errorf("tunnel: %q", rest)
	}
}

func TestNormalHost(t *testing.T) {
	for in, want := range map[string]string{
		"Example.COM.":    "example.com",
		"a-b_c.example":   "a-b_c.example",
		"127.0.0.1":       "127.0.0.1",
		"::FFFF:1.2.3.4":  "::ffff:1.2.3.4",
		"evil\x1b[2J.com": "",
		"ex ample.com":    "",
		"a..b":            "",
		"":                "",
		"bücher.de":       "",
		"\u212aey.com":    "", // the Kelvin sign, not k
	} {
		got, ok := normalHost(in)
		if ok != (want != "") || got != want {
			t.Errorf("normalHost(%q) = %q, %v; want %q", in, got, ok, want)
		}
	}
}

// A plain HTTP request reaches the host approved under that name, never
// under a Host header of the client's own (which a CDN would route by).
// net/http's server takes the host of an absolute URI over the header, as
// RFC 9112 says; this keeps it so.
func TestHostHeaderIsTheApprovedHost(t *testing.T) {
	up, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	upSrv := &http.Server{Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { fmt.Fprint(w, "host="+r.Host) })}
	go upSrv.Serve(up)
	t.Cleanup(func() { upSrv.Close() })
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	port := up.Addr().(*net.TCPAddr).Port
	s := &Server{Check: func(_ context.Context, host string) (bool, string) { return host == "127.0.0.1", "" }, Ports: []int{port}, AllowPrivate: true}
	go s.Serve(l)
	t.Cleanup(s.Close)

	c, err := net.Dial("tcp", l.Addr().String())
	if err != nil {
		t.Fatal(err)
	}
	defer c.Close()
	fmt.Fprintf(c, "GET http://127.0.0.1:%d/ HTTP/1.1\r\nHost: fronted.example\r\nConnection: close\r\n\r\n", port)
	resp, err := http.ReadResponse(bufio.NewReader(c), nil)
	if err != nil {
		t.Fatal(err)
	}
	body, _ := io.ReadAll(resp.Body)
	if want := fmt.Sprintf("host=127.0.0.1:%d", port); string(body) != want {
		t.Errorf("upstream saw %q, want %q", body, want)
	}
}