package proxy_test

import (
	"fmt"
	"maps"
	"net"
	"net/http"
	"net/http/httptest"
	"net/netip"
	"slices"
	"strconv"
	"strings"
	"testing"
	"testing/synctest"
	"time"

	"sneak.berlin/go/smallwebwaf/internal/bans"
	"sneak.berlin/go/smallwebwaf/internal/proxy"
	"sneak.berlin/go/smallwebwaf/internal/ratelimit"
	"sneak.berlin/go/smallwebwaf/internal/reputation"
	"sneak.berlin/go/smallwebwaf/internal/requestlog"
)

const (
	// answerTimeLimit is the default of rproxy's SWRP_TIMEOUT_MS, the time
	// servewaf has to answer.
	answerTimeLimit = 2 * time.Second
	// defaultLookupWait is the default of SWWAF_LOOKUP_TIMEOUT.
	defaultLookupWait = 800 * time.Millisecond
	// The header limits SPEC.md gives.
	cookieLimit      = 50
	headerFieldLimit = 100
	// cookieHeader is the header the cookie limit counts.
	cookieHeader = "Cookie"
	// tooLarge is how servewaf refuses a request over the header limits.
	tooLarge = http.StatusRequestHeaderFieldsTooLarge
	// silentAbuseIPDBURL is where AbuseIPDB never answers: at the stand-in
	// that unansweredGeoJSURL names too.
	silentAbuseIPDBURL = "unanswered://abuseipdb/api/v2/check"
	// wholeLimit is the percentage of a rate limit that lowers it not at all.
	wholeLimit = 100
	// clientsDecision is the CrowdSec decision list holding client alone.
	clientsDecision = `[{"duration": "4h0m0s", "scenario": "crowdsecurity/ssh-bf", ` +
		`"scope": "Ip", "type": "ban", "value": "` + client + `"}]`
)

// cookies returns a Cookie header's value with count cookies.
func cookies(count int) string {
	parts := make([]string, 0, count)
	for i := range count {
		parts = append(parts, "c"+strconv.Itoa(i)+"=1")
	}

	return strings.Join(parts, "; ")
}

// fillHeaders adds headers to a question until the checks see total header
// fields: all but ClientHeader and PeerHeader, which servewaf removes. Go's
// client sends nothing else but Host, which the limit does not count, once
// User-Agent is set.
func fillHeaders(req *http.Request, total int) {
	req.Header.Set("User-Agent", "test")

	for i := len(req.Header) - 2; i < total; i++ {
		req.Header.Set("X-Filler-"+strconv.Itoa(i), "1")
	}
}

func TestSixteenThousandCookiesAreRefusedWellWithinTheTimeLimit(t *testing.T) {
	t.Parallel()

	s := startServeWAF(t, "", map[string]string{wafMode: block})
	req := s.question(siteA, client, "/")
	req.Header.Set(cookieHeader, strings.Repeat("a;", 16000))

	began := time.Now()
	got := do(t, req)

	if took := time.Since(began); took >= answerTimeLimit/2 {
		t.Errorf("answered after %s, want well within %s", took, answerTimeLimit)
	}

	wantDisposition(t, got, tooLarge, proxy.DispositionDeny)

	line := s.out.requestLine(t)
	wantLine(t, line, tooLarge, requestlog.ActionTooLarge)

	if line.DurationWAF != nil {
		t.Errorf("the Core Rule Set read the request, for %f ms", *line.DurationWAF)
	}
}

func TestEachHeaderLimitRefusesARequestOverIt(t *testing.T) {
	t.Parallel()

	repeated := make([]string, headerFieldLimit)
	for i := range repeated {
		repeated[i] = "1"
	}

	for _, tc := range []struct {
		name string
		make func(*http.Request)
		// refused is true for a request over a limit, and false for one at it.
		refused bool
	}{
		{"cookies at the limit", func(req *http.Request) {
			req.Header.Set(cookieHeader, cookies(cookieLimit))
		}, false},
		{"cookies over the limit", func(req *http.Request) {
			req.Header.Set(cookieHeader, cookies(cookieLimit+1))
		}, true},
		{"empty parts are no cookies", func(req *http.Request) {
			req.Header.Set(cookieHeader, cookies(cookieLimit)+"; ;;")
		}, false},
		{"cookies over the limit across two headers", func(req *http.Request) {
			req.Header[cookieHeader] = []string{cookies(cookieLimit), "last=1"}
		}, true},
		{"cookies without a value count", func(req *http.Request) {
			req.Header.Set(cookieHeader, strings.Repeat("a;", cookieLimit+1))
		}, true},
		{"header fields at the limit", func(req *http.Request) {
			fillHeaders(req, headerFieldLimit)
		}, false},
		{"header fields over the limit", func(req *http.Request) {
			fillHeaders(req, headerFieldLimit+1)
		}, true},
		{"a header sent many times counts each time", func(req *http.Request) {
			req.Header["X-Filler"] = repeated
		}, true},
	} {
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()

			s := startServeWAF(t, "", map[string]string{wafMode: block})
			req := s.question(siteA, client, "/")
			tc.make(req)

			got := do(t, req)
			line := s.out.requestLine(t)

			if tc.refused {
				wantDisposition(t, got, tooLarge, proxy.DispositionDeny)
				wantLine(t, line, tooLarge, requestlog.ActionTooLarge)

				return
			}

			wantDisposition(t, got, http.StatusOK, proxy.DispositionAllow)
			wantLine(t, line, http.StatusOK, requestlog.ActionForward)

			if line.DurationWAF == nil {
				t.Error("the Core Rule Set did not read a request within the limits")
			}
		})
	}
}

func TestHeaderLimitsHoldWhateverTheCoreRuleSetSettingsSay(t *testing.T) {
	t.Parallel()

	for _, env := range []map[string]string{
		{wafMode: off},
		{wafMode: "detect"},
		{wafMode: block, "SWWAF_WAF_EXEMPT_PATHS": "/"},
	} {
		t.Run(fmt.Sprint(env), func(t *testing.T) {
			t.Parallel()

			s := startServeWAF(t, "", env)
			req := s.question(siteA, client, "/")
			req.Header.Set(cookieHeader, cookies(cookieLimit+1))

			wantDisposition(t, do(t, req), tooLarge, proxy.DispositionDeny)
		})
	}
}

func TestClientInTheAllowNetsSkipsTheHeaderLimits(t *testing.T) {
	t.Parallel()

	s := startServeWAF(t, "", map[string]string{wafMode: block, allowNets: client})
	req := s.question(siteA, client, "/")
	req.Header.Set(cookieHeader, cookies(cookieLimit+1))

	wantDisposition(t, do(t, req), http.StatusOK, proxy.DispositionAllow)
}

func TestObserveModeAllowsARequestOverTheHeaderLimitsUnread(t *testing.T) {
	t.Parallel()

	s := startServeWAF(t, "", map[string]string{mode: observe, wafMode: block})
	req := s.question(siteA, client, "/")
	req.Header.Set(cookieHeader, strings.Repeat("a;", 16000))

	began := time.Now()
	got := do(t, req)

	if took := time.Since(began); took >= answerTimeLimit/2 {
		t.Errorf("answered after %s, want well within %s", took, answerTimeLimit)
	}

	wantDisposition(t, got, http.StatusOK, proxy.DispositionAllow)

	line := s.out.requestLine(t)
	if line.WouldAction != requestlog.ActionTooLarge || line.DurationWAF != nil {
		t.Errorf("would_action %q and duration_waf %v, want %q and none",
			line.WouldAction, line.DurationWAF, requestlog.ActionTooLarge)
	}
}

func TestRequestsOverTheHeaderLimitsCountTowardTheRateLimits(t *testing.T) {
	t.Parallel()

	s := startServeWAF(t, "", map[string]string{rateLimitPerMinute: "1"})
	ask := func() answer {
		req := s.question(siteA, client, "/")
		req.Header.Set(cookieHeader, cookies(cookieLimit+1))

		return do(t, req)
	}

	wantDisposition(t, ask(), tooLarge, proxy.DispositionDeny)
	wantDisposition(t, ask(), http.StatusForbidden, proxy.DispositionDeny)

	lines := s.out.requestLines(t, 2)
	wantLine(t, lines[1], http.StatusForbidden, requestlog.ActionRateLimited)
}

func TestWithoutACommandTheHeaderLimitsDoNotApply(t *testing.T) {
	t.Parallel()

	app := startApp(t, func(http.ResponseWriter, *http.Request) {})
	addr, out := startProxy(t, app.URL, nil)

	req := newRequest(t, http.MethodGet, addr, "/", http.NoBody)
	req.Header.Set(cookieHeader, cookies(cookieLimit+1))
	fillHeaders(req, headerFieldLimit+1)

	wantStatus(t, do(t, req), http.StatusOK)
	wantLine(t, out.requestLine(t), http.StatusOK, requestlog.ActionForward)
}

// newServeWAFInBubble makes a servewaf for a synctest bubble: GeoJS never
// answers, a setting needs its answer, and everything is at its default but
// for env. No request to it may use the network.
func newServeWAFInBubble(t *testing.T, env map[string]string) (*proxy.Server, *output) {
	t.Helper()

	settings := map[string]string{deniedCountries: "KP"}
	maps.Copy(settings, env)

	server, out, _ := newServerWith(t, proxy.NewServeWAF, "http://app.invalid",
		unansweredGeoJSURL, time.Now, settings)

	return server, out
}

// askDirectly hands a question about a request from client straight to server,
// as change leaves it, and returns the answer's status and how long it took.
func askDirectly(
	t *testing.T, server *proxy.Server, change func(*http.Request),
) (int, time.Duration) {
	t.Helper()

	req := httptest.NewRequestWithContext(t.Context(), http.MethodGet, "/", http.NoBody)
	req.RemoteAddr = net.JoinHostPort(rproxyPeer, "1234")
	req.Header.Set(proxy.ClientHeader, client)
	req.Header.Set(requestIDHeader, "the-id")
	change(req)

	recorder := httptest.NewRecorder()
	began := time.Now()

	server.Handler.ServeHTTP(recorder, req)

	return recorder.Code, time.Since(began)
}

// banClient bans client at now.
func banClient(_ *testing.T, server *proxy.Server, now time.Time) {
	server.Ledger.BanForLimit(netip.MustParsePrefix(client+"/32"), now, bans.Notes{})
}

// countClient counts one request from client at now for the rate limits.
func countClient(_ *testing.T, server *proxy.Server, now time.Time) {
	server.Limiter.Count(netip.MustParsePrefix(client+"/32"), now, wholeLimit)
}

// listClient puts client on the blocklist at dropURL.
func listClient(t *testing.T, server *proxy.Server, now time.Time) {
	t.Helper()

	err := server.Lists.Load([]reputation.List{{
		URL: dropURL, Tried: now, Fetched: now, Lines: []string{client},
	}})
	if err != nil {
		t.Fatalf("load the blocklist: %v", err)
	}
}

// decideOnClient puts client on the CrowdSec decision list.
func decideOnClient(t *testing.T, server *proxy.Server, now time.Time) {
	t.Helper()
	loadDecisions(t, server, now, clientsDecision)
}

// leaveAsItIs prepares nothing.
func leaveAsItIs(*testing.T, *proxy.Server, time.Time) {}

func TestARefusedRequestDoesNotWaitForASlowGeoJS(t *testing.T) {
	t.Parallel()

	const forbidden = http.StatusForbidden

	for _, tc := range []struct {
		name string
		env  map[string]string
		// prepare gets the server ready, at now, before the question.
		prepare func(t *testing.T, server *proxy.Server, now time.Time)
		// cookie is the request's Cookie header.
		cookie string
		status int
		action string
	}{
		{
			name: "SWWAF_DENY_NETS", env: map[string]string{denyNets: client},
			prepare: leaveAsItIs, status: forbidden, action: requestlog.ActionDenied,
		},
		{
			name:    "a ban",
			prepare: banClient, status: forbidden, action: requestlog.ActionBanned,
		},
		{
			name: "the rate limits", env: map[string]string{rateLimitPerMinute: "1"},
			prepare: countClient, status: forbidden, action: requestlog.ActionRateLimited,
		},
		{
			name:    "the rate limits, with a biased threshold",
			env:     map[string]string{rateLimitPerMinute: "1", unknownLimitPercent: "50"},
			prepare: countClient, status: forbidden, action: requestlog.ActionRateLimited,
		},
		{
			name: "a blocklist", env: map[string]string{blocklistURLs: dropURL},
			prepare: listClient, status: forbidden, action: requestlog.ActionDenied,
		},
		{
			name:    "the CrowdSec decision list",
			env:     map[string]string{crowdSecURL: lapi, crowdSecKey: bouncerKey},
			prepare: decideOnClient, status: forbidden, action: requestlog.ActionBanned,
		},
		{
			name: "the header limits", cookie: cookies(cookieLimit + 1),
			prepare: leaveAsItIs, status: tooLarge, action: requestlog.ActionTooLarge,
		},
	} {
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()

			// In a synctest bubble, as TestRequestWaitsAsLongAsTheLookupTimeoutSays
			// says, so that a request that does not wait takes no time at all.
			synctest.Test(t, func(t *testing.T) {
				server, out := newServeWAFInBubble(t, tc.env)
				tc.prepare(t, server, time.Now())

				status, took := askDirectly(t, server, func(req *http.Request) {
					req.Header.Set(cookieHeader, tc.cookie)
				})

				if status != tc.status || took != 0 {
					t.Errorf("answered %d after %s, want %d at once; output:\n%s",
						status, took, tc.status, out.text())
				}

				wantLine(t, out.requestLine(t), tc.status, tc.action)

				// GeoJS's request is abandoned before the bubble ends.
				time.Sleep(defaultLookupWait)
			})
		})
	}
}

func TestAQuestionWaitsForASlowGeoJSOnlyAsLongAsTheDefaultSays(t *testing.T) {
	t.Parallel()

	for _, tc := range []struct {
		name   string
		env    map[string]string
		status int
		action string
	}{
		{"nothing refuses it", nil, http.StatusOK, requestlog.ActionForward},
		{
			"a country list needs the answer", map[string]string{allowedCountries: "DE"},
			http.StatusForbidden, requestlog.ActionCountryDenied,
		},
	} {
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()

			synctest.Test(t, func(t *testing.T) {
				server, out := newServeWAFInBubble(t, tc.env)
				status, took := askDirectly(t, server, func(*http.Request) {})

				if status != tc.status || took != defaultLookupWait {
					t.Errorf("answered %d after %s, want %d after %s; output:\n%s",
						status, took, tc.status, defaultLookupWait, out.text())
				}

				wantLine(t, out.requestLine(t), tc.status, tc.action)
			})
		})
	}
}

func TestServeWAFAnswersWithinTheTimeLimitWhileEveryOutsideSourceHangs(t *testing.T) {
	t.Parallel()

	held, askedGeoJS, _ := startHeldGeoJS(t)

	// A resolver that takes every query and answers none.
	resolver, err := (&net.ListenConfig{}).ListenPacket(t.Context(), "udp", localhost+":0")
	if err != nil {
		t.Fatalf("listen: %v", err)
	}

	t.Cleanup(func() { _ = resolver.Close() })

	server, out, _ := newServerWith(t, func(params proxy.Params) *proxy.Server {
		params.AbuseIPDBURL = silentAbuseIPDBURL

		return proxy.NewServeWAF(params)
	}, "http://app.invalid", held, time.Now, map[string]string{
		deniedCountries: "KP", wafMode: block,
		dnsblZones: dnsblZone, dnsblResolver: resolver.LocalAddr().String(),
		abuseIPDBKey:  accountKey,
		blocklistURLs: held + "/drop.txt", asnLimitPercentURL: held + "/asn.txt",
		crowdSecURL: held, crowdSecKey: bouncerKey,
	})

	// AbuseIPDB is asked only about a client with an offence.
	group := netip.MustParsePrefix(client + "/32")
	server.Limiter.AddToHistory(group, time.Now(), ratelimit.Request{Attack: true})

	go server.Lists.Run(t.Context())

	status, took := askDirectly(t, server, func(req *http.Request) {
		req.Header.Set("User-Agent", "test")
		req.Header.Set("Accept", "*/*")
		req.Host = siteA
	})

	if status != http.StatusOK || took >= answerTimeLimit {
		t.Errorf("answered %d after %s, want %d within %s; output:\n%s",
			status, took, http.StatusOK, answerTimeLimit, out.text())
	}

	// Each outside source was asked, and none has answered. The stand-in for
	// GeoJS is the lists' too, which name no address.
	if !slices.Contains(askedGeoJS(), client) {
		t.Errorf("GeoJS was asked about %q, not %s", askedGeoJS(), client)
	}

	if queries := server.DNSBL.Queries(dnsblZone); queries != 1 {
		t.Errorf("%d queries to the DNS blocklist, want 1", queries)
	}

	if checked := server.AbuseIPDB.Checked(); checked != 1 {
		t.Errorf("%d checks with AbuseIPDB, want 1", checked)
	}
}