package main

import (
	"context"
	"net"
	"net/http"
	"net/http/httptest"
	"net/netip"
	"strings"
	"sync"
	"testing"
	"time"
)

type fetchRig struct {
	policy *Policy
	f      *Fetcher
	mu     sync.Mutex
	dialed []string
}

func newFetchRig(t *testing.T, handler http.Handler) *fetchRig {
	t.Helper()
	server := httptest.NewServer(handler)
	t.Cleanup(server.Close)
	rig := &fetchRig{}
	rig.policy = testPolicy(publicAnswer)
	rig.policy.dial = func(ctx context.Context, network, address string) (net.Conn, error) {
		rig.mu.Lock()
		rig.dialed = append(rig.dialed, address)
		rig.mu.Unlock()
		var dialer net.Dialer
		return dialer.DialContext(ctx, "tcp", server.Listener.Addr().String())
	}
	rig.f = NewFetcher(rig.policy)
	return rig
}

func job(url string) Frame {
	return Frame{Type: "job", JobID: "j1", Method: "GET", URL: url, Timeout: 5}
}

func TestFetchSuccessStripsIdentityAndHopByHop(t *testing.T) {
	var seen http.Header
	rig := newFetchRig(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		seen = r.Header.Clone()
		w.Header().Set("Set-Cookie", "session=abc")
		w.Header().Set("X-Thing", "kept")
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"ok":true}`))
	}))
	frame := job("http://www.reddit.com/r/go.json")
	frame.Headers = map[string]any{
		"Cookie": "a=b", "Authorization": "Bearer x", "Proxy-Authorization": "y", "Connection": "close",
		"Accept": "application/json", "X-Bad": "a\r\nb", "Host": "evil.com", "X-Number": 5,
	}
	result, err := rig.f.Do(context.Background(), frame, 1000, 1<<20)
	if err != nil {
		t.Fatal(err)
	}
	if result.Status != 200 || string(result.Body) != `{"ok":true}` || result.Truncated {
		t.Fatalf("result %+v", result)
	}
	for _, name := range []string{"Cookie", "Authorization", "Proxy-Authorization", "X-Bad", "X-Number"} {
		if seen.Get(name) != "" {
			t.Errorf("%s was forwarded", name)
		}
	}
	if seen.Get("Accept") != "application/json" {
		t.Error("allowed header dropped")
	}
	if _, ok := result.Headers["set-cookie"]; ok {
		t.Error("set-cookie reported")
	}
	if result.Headers["x-thing"] != "kept" {
		t.Errorf("headers %v", result.Headers)
	}
	if len(rig.dialed) != 1 || rig.dialed[0] != "93.184.216.34:80" {
		t.Errorf("dialed %v", rig.dialed)
	}
}

func TestFetchTruncatesAtBodyCapAndBudget(t *testing.T) {
	rig := newFetchRig(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		_, _ = w.Write([]byte(strings.Repeat("x", 5000)))
	}))
	result, err := rig.f.Do(context.Background(), job("http://reddit.com/a"), 1000, 1<<20)
	if err != nil || len(result.Body) != 1000 || !result.Truncated {
		t.Fatalf("cap: %v %+v", err, result)
	}
	result, err = rig.f.Do(context.Background(), job("http://reddit.com/a"), 1000, 300)
	if err != nil || len(result.Body) != 300 || !result.Truncated {
		t.Fatalf("budget: %v %+v", err, result)
	}
	result, err = rig.f.Do(context.Background(), job("http://reddit.com/a"), 10000, 1<<20)
	if err != nil || len(result.Body) != 5000 || result.Truncated {
		t.Fatalf("whole: %v %+v", err, result)
	}
}

func TestFetchRefusals(t *testing.T) {
	rig := newFetchRig(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		_, _ = w.Write([]byte("ok"))
	}))
	cases := []struct {
		name   string
		frame  Frame
		prep   func()
		expect string
	}{
		{"host", job("http://evilreddit.com/"), nil, RefuseHost},
		{"port", job("http://reddit.com:8080/"), nil, RefusePort},
		{"method", Frame{JobID: "j", Method: "POST", URL: "http://reddit.com/", Timeout: 5}, nil, RefuseMethod},
		{"literal", job("http://10.0.0.1/"), nil, RefuseHost},
		{"private resolve", job("http://reddit.com/"), func() {
			rig.policy.lookup = func(context.Context, string) ([]netip.Addr, error) {
				return []netip.Addr{netip.MustParseAddr("192.168.1.10")}, nil
			}
		}, RefuseAddress},
	}
	for _, item := range cases {
		if item.prep != nil {
			item.prep()
		}
		before := len(rig.dialed)
		_, err := rig.f.Do(context.Background(), item.frame, 1000, 1<<20)
		if got := ErrorCode(err); got != item.expect {
			t.Errorf("%s: got %q (%v)", item.name, got, err)
		}
		if len(rig.dialed) != before {
			t.Errorf("%s: a connection was attempted", item.name)
		}
	}
}

func TestFetchRedirects(t *testing.T) {
	var hits []string
	rig := newFetchRig(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		hits = append(hits, r.Host+r.URL.Path)
		switch r.URL.Path {
		case "/to-allowed":
			http.Redirect(w, r, "http://old.reddit.com/final", http.StatusFound)
		case "/to-evil":
			http.Redirect(w, r, "http://evil.example/steal", http.StatusFound)
		case "/to-private":
			http.Redirect(w, r, "http://127.0.0.1/admin", http.StatusFound)
		case "/to-port":
			http.Redirect(w, r, "http://reddit.com:8080/x", http.StatusFound)
		case "/loop":
			http.Redirect(w, r, "/loop", http.StatusFound)
		default:
			_, _ = w.Write([]byte("final"))
		}
	}))
	result, err := rig.f.Do(context.Background(), job("http://reddit.com/to-allowed"), 1000, 1<<20)
	if err != nil || string(result.Body) != "final" {
		t.Fatalf("allowed redirect: %v", err)
	}
	for _, path := range []string{"/to-evil", "/to-private", "/to-port", "/loop"} {
		_, err := rig.f.Do(context.Background(), job("http://reddit.com"+path), 1000, 1<<20)
		if ErrorCode(err) != RefuseRedirect {
			t.Errorf("%s: got %q (%v)", path, ErrorCode(err), err)
		}
	}
	for _, hit := range hits {
		if strings.HasPrefix(hit, "evil.example") || strings.HasPrefix(hit, "127.0.0.1") {
			t.Errorf("request reached %s", hit)
		}
	}
}

func TestFetchTimeout(t *testing.T) {
	rig := newFetchRig(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		select {
		case <-r.Context().Done():
		case <-time.After(3 * time.Second):
		}
	}))
	frame := job("http://reddit.com/slow")
	frame.Timeout = 1
	started := time.Now()
	_, err := rig.f.Do(context.Background(), frame, 1000, 1<<20)
	if ErrorCode(err) != RefuseTimeout {
		t.Fatalf("got %q (%v)", ErrorCode(err), err)
	}
	if time.Since(started) > 2500*time.Millisecond {
		t.Fatal("timeout was not enforced")
	}
}

func TestClampTimeout(t *testing.T) {
	cases := map[int]time.Duration{0: 25 * time.Second, -3: 25 * time.Second, 7: 7 * time.Second, 500: 60 * time.Second}
	for in, want := range cases {
		if got := clampTimeout(in); got != want {
			t.Errorf("%d: %v", in, got)
		}
	}
}