package smallwebwaf_test

import (
	"bufio"
	"context"
	"encoding/pem"
	"io"
	"maps"
	"net"
	"net/http"
	"net/http/httptest"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"sync/atomic"
	"syscall"
	"testing"
	"time"

	"sneak.berlin/go/smallwebwaf/internal/servewaf"
	"sneak.berlin/go/smallwebwaf/internal/smallwebwaf"
)

const (
	// wafBackend is the setting that names servewaf to rproxy.
	wafBackend = "SWRP_WAF_BACKEND"
	// rproxyListenAddr is the setting that says where rproxy listens.
	rproxyListenAddr = "SWRP_LISTEN_ADDR"
	// rproxyUpstreamURL is the setting that says where the app is.
	rproxyUpstreamURL = "SWRP_UPSTREAM_URL"
	// rproxyTimeout is the setting that says how long rproxy waits for servewaf.
	rproxyTimeout = "SWRP_TIMEOUT_MS"
	// bannedProbe is a path the default rule file bans a client for.
	bannedProbe = "/.env"
	// otherClient is a client the tests do not ban.
	otherClient = "198.51.100.7"
)

// runRProxy runs `smallwebwaf rproxy` with the settings in env until ctx is
// done, and returns its exit status.
func runRProxy(ctx context.Context, env map[string]string, out *output) int {
	return smallwebwaf.RunRProxy(ctx, smallwebwaf.Params{
		Version: testVersion,
		LookupEnv: func(name string) (string, bool) {
			value, ok := env[name]

			return value, ok
		},
		Stdout: out,
	})
}

// startRProxy runs rproxy with servewaf at backend and the app at appURL, and
// returns its URL and what it writes. It believes X-Forwarded-For from the
// tests. It stops as the test ends, and that it stopped in order is checked.
func startRProxy(t *testing.T, backend, appURL string) (string, *output) {
	t.Helper()

	ctx, stop := context.WithCancel(t.Context())
	out := &output{}
	exited := make(chan int, 1)

	go func() {
		exited <- runRProxy(ctx, map[string]string{
			rproxyListenAddr:       localhost + ":0",
			"SWRP_UPSTREAM_URL":    appURL,
			"SWRP_TRUSTED_PROXIES": localhost + "/32",
			wafBackend:             backend,
		}, out)
	}()

	t.Cleanup(func() {
		stop()

		select {
		case status := <-exited:
			if status != 0 {
				t.Errorf("rproxy's exit status %d, want 0; output:\n%s", status, out.text())
			}
		case <-time.After(waitLimit):
			t.Error("rproxy still running after being told to stop")
		}

		out.line(t, "msg", "stopped")
	})

	addr, _ := out.line(t, "msg", "starting")["address"].(string)

	return "http://" + addr, out
}

// startServeWAF runs servewaf with the settings in env, and returns its address
// and URL, what it writes, and what stops it, which checks that it stopped in
// order.
func startServeWAF(t *testing.T, env map[string]string) (string, *output, func()) {
	t.Helper()

	ctx, stop := context.WithCancel(t.Context())
	out := &output{}
	exited := make(chan int, 1)

	go func() {
		exited <- run(ctx, env, out)
	}()

	addr, _ := out.line(t, "msg", "starting")["address"].(string)

	return addr, out, func() {
		stop()

		select {
		case status := <-exited:
			if status != 0 {
				t.Fatalf("servewaf's exit status %d, want 0; output:\n%s", status, out.text())
			}
		case <-time.After(waitLimit):
			t.Fatal("servewaf still running after being told to stop")
		}
	}
}

// serveWAFSettings are servewaf's settings with the default rule file, to listen
// at any free port.
func serveWAFSettings(t *testing.T) map[string]string {
	t.Helper()

	return map[string]string{
		listenAddr: localhost + ":0",
		stateDir:   t.TempDir(),
		rulesDir:   filepath.Join("..", "..", "share", "rules.d"),
	}
}

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

	addr, serveWAFOut, stop := startServeWAF(t, serveWAFSettings(t))
	defer stop()

	backend := "http://" + addr
	first, firstOut := startRProxy(t, backend, startApp(t))
	second, _ := startRProxy(t, backend, startApp(t))

	wantStatus(t, first+"/", placed, http.StatusOK)
	wantStatus(t, second+"/", placed, http.StatusOK)

	// The default rule file bans a probe for this.
	wantStatus(t, first+bannedProbe, placed, http.StatusForbidden)
	wantStatus(t, second+"/", placed, http.StatusForbidden)
	wantStatus(t, second+"/", otherClient, http.StatusOK)

	// One request's three lines carry one id, rproxy's with servewaf's answer.
	line := firstOut.line(t, "path", bannedProbe)
	if line["disposition"] != servewaf.DispositionDeny || line["client_ip"] != placed {
		t.Errorf("rproxy's line is %v", line)
	}

	id, _ := line["request_id"].(string)
	decision := serveWAFLine(t, serveWAFOut, id, "decision")
	report := serveWAFLine(t, serveWAFOut, id, "report")

	if decision["host"] != line["host"] || decision["client_ip"] != placed ||
		report["status"] != float64(http.StatusForbidden) {
		t.Errorf("servewaf's lines are %v and %v", decision, report)
	}
}

// serveWAFLine waits for servewaf's line with request id and servewaf_line kind.
func serveWAFLine(t *testing.T, out *output, id, kind string) map[string]any {
	t.Helper()

	deadline := time.Now().Add(waitLimit)
	for time.Now().Before(deadline) {
		line := out.findLine(t, func(fields map[string]any) bool {
			return fields["request_id"] == id && fields["servewaf_line"] == kind
		})
		if line != nil {
			return line
		}

		time.Sleep(pollInterval)
	}

	t.Fatalf("no %s line for %s in:\n%s", kind, id, out.text())

	return nil
}

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

	env := serveWAFSettings(t)
	addr, _, stop := startServeWAF(t, env)
	site, out := startRProxy(t, "http://"+addr, startApp(t))

	// A visitor who keeps asking throughout, and is never refused.
	visitorDone := make(chan struct{})
	visitorStop := make(chan struct{})

	var refused atomic.Int64

	go func() {
		defer close(visitorDone)

		for {
			select {
			case <-visitorStop:
				return
			case <-time.After(pollInterval):
			}

			if !answersNormally(t.Context(), site+"/", otherClient) {
				refused.Add(1)
			}
		}
	}()

	wantStatus(t, site+bannedProbe, placed, http.StatusForbidden)
	wantStatus(t, site+"/", placed, http.StatusForbidden)
	stop()

	// With servewaf down, the site answers, and rproxy says the default did it.
	wantStatus(t, site+"/after-stop", placed, http.StatusOK)

	line := out.line(t, "path", "/after-stop")
	if line["disposition"] != servewaf.DispositionAllow ||
		line["default_disposition"] != true {
		t.Errorf("rproxy's line is %v", line)
	}

	// Back at its address, with its state files, it applies the ban again.
	restarted := maps.Clone(env)
	restarted[listenAddr] = addr
	_, _, stop = startServeWAF(t, restarted)

	defer stop()

	wantStatus(t, site+"/", placed, http.StatusForbidden)

	close(visitorStop)
	<-visitorDone

	if refused.Load() != 0 {
		t.Errorf("the site refused its visitor %d times", refused.Load())
	}
}

// answersNormally reports whether a request to url from the client at from, as
// X-Forwarded-For names it, gets the app's answer.
func answersNormally(ctx context.Context, url, from string) bool {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, http.NoBody)
	if err != nil {
		return false
	}

	req.Header.Set("X-Forwarded-For", from)

	transport := &http.Transport{}
	defer transport.CloseIdleConnections()

	res, err := transport.RoundTrip(req)
	if err != nil {
		return false
	}

	body, err := io.ReadAll(res.Body)
	_ = res.Body.Close()

	return err == nil && res.StatusCode == http.StatusOK && string(body) == greeting
}

// sendRaw sends request, as it is, to rproxy at site, and returns the status of
// the answer.
func sendRaw(t *testing.T, site, request string) int {
	t.Helper()

	conn, err := (&net.Dialer{}).DialContext(t.Context(), "tcp",
		strings.TrimPrefix(site, "http://"))
	if err != nil {
		t.Fatalf("dial: %v", err)
	}

	defer func() {
		_ = conn.Close()
	}()

	err = conn.SetDeadline(time.Now().Add(waitLimit))
	if err != nil {
		t.Fatalf("set the deadline: %v", err)
	}

	_, err = io.WriteString(conn, request)
	if err != nil {
		t.Fatalf("send: %v", err)
	}

	res, err := http.ReadResponse(bufio.NewReader(conn), nil)
	if err != nil {
		t.Fatalf("read the answer: %v", err)
	}

	_ = res.Body.Close()

	return res.StatusCode
}

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

	addr, serveWAFOut, stop := startServeWAF(t, serveWAFSettings(t))
	defer stop()

	site, _ := startRProxy(t, "http://"+addr, startApp(t))

	// Read as a further request, this body would ban otherClient on every site.
	const smuggledID = "smuggled-request"

	smuggled := "GET " + bannedProbe + " HTTP/1.1\r\nHost: site.example\r\n" +
		servewaf.ClientHeader + ": " + otherClient + "\r\n" +
		"X-Request-ID: " + smuggledID + "\r\n\r\n"

	status := sendRaw(t, site, "POST /upload HTTP/1.1\r\nHost: site.example\r\n"+
		"X-Forwarded-For: "+placed+"\r\nTransfer-Encoding: chunked\r\n\r\n"+
		strconv.FormatInt(int64(len(smuggled)), 16)+"\r\n"+smuggled+"\r\n0\r\n\r\n")
	if status != http.StatusOK {
		t.Errorf("status %d, want the app's 200", status)
	}

	// Read as a further request, it would act straight after the answer.
	for range 20 {
		wantStatus(t, site+"/", otherClient, http.StatusOK)
		time.Sleep(pollInterval)
	}

	line := serveWAFOut.findLine(t, func(fields map[string]any) bool {
		return fields["request_id"] == smuggledID
	})
	if line != nil {
		t.Errorf("servewaf read the body as a request: %v", line)
	}
}

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

	addr, _, stop := startServeWAF(t, serveWAFSettings(t))
	defer stop()

	site, out := startRProxy(t, "http://"+addr, startApp(t))

	wantStatus(t, site+bannedProbe, placed, http.StatusForbidden)

	status := sendRaw(t, site, "GET x:admin HTTP/1.1\r\nHost: site.example\r\n"+
		"X-Forwarded-For: "+placed+"\r\n\r\n")
	if status != http.StatusBadRequest {
		t.Errorf("status %d, want 400", status)
	}

	line := out.line(t, "path", "admin")
	if line["disposition"] != nil || line["default_disposition"] != nil {
		t.Errorf("rproxy's line is %v, want no disposition", line)
	}
}

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

	const zoned = "2001:db8::1%eth0"

	addr, _, stop := startServeWAF(t, serveWAFSettings(t))
	defer stop()

	site, out := startRProxy(t, "http://"+addr, startApp(t))

	wantStatus(t, site+bannedProbe, zoned, http.StatusForbidden)
	wantStatus(t, site+"/", zoned, http.StatusForbidden)
	wantStatus(t, site+"/", "2001:db8::1", http.StatusForbidden)

	line := out.line(t, "path", bannedProbe)
	if line["client_ip"] != "2001:db8::1" || line["default_disposition"] != nil {
		t.Errorf("rproxy's line is %v", line)
	}
}

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

	const status = 799

	app := httptest.NewServer(http.HandlerFunc(
		func(w http.ResponseWriter, _ *http.Request) {
			w.WriteHeader(status)
		}))
	t.Cleanup(app.Close)

	addr, serveWAFOut, stop := startServeWAF(t, serveWAFSettings(t))
	defer stop()

	site, out := startRProxy(t, "http://"+addr, app.URL)

	wantStatus(t, site+"/", placed, status)

	line := out.line(t, "path", "/")
	if line["status"] != float64(status) {
		t.Errorf("rproxy's line is %v, want the status the visitor got", line)
	}

	id, _ := line["request_id"].(string)

	report := serveWAFLine(t, serveWAFOut, id, "report")
	if report["status"] != 599.0 {
		t.Errorf("servewaf's line is %v, want status 599", report)
	}
}

// startRProxyProcess runs `smallwebwaf rproxy` in a process of its own, with the
// settings in env, each as NAME=value. It returns where rproxy listens, the
// process, and what rproxy writes.
func startRProxyProcess(t *testing.T, env ...string) (string, *exec.Cmd, *output) {
	t.Helper()

	binary, err := os.Executable()
	if err != nil {
		t.Fatalf("find the test binary: %v", err)
	}

	ctx, cancel := context.WithTimeout(t.Context(), 3*waitLimit)
	t.Cleanup(cancel)

	out := &output{}
	cmd := exec.CommandContext(ctx, binary) //nolint:gosec // this test binary
	cmd.Stdout = out

	cmd.Env = append(os.Environ(),
		commandForMain+"=rproxy",
		"SWRP_LISTEN_ADDR="+localhost+":0")
	cmd.Env = append(cmd.Env, env...)

	err = cmd.Start()
	if err != nil {
		t.Fatalf("start: %v", err)
	}

	t.Cleanup(func() {
		_ = cmd.Process.Kill()
		_ = cmd.Wait()
	})

	addr, _ := out.line(t, "msg", "starting")["address"].(string)

	return addr, cmd, out
}

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

	protocols := make(chan string, 10)
	backend := httptest.NewUnstartedServer(http.HandlerFunc(
		func(w http.ResponseWriter, r *http.Request) {
			if r.Header.Get(servewaf.ClientHeader) != "" {
				protocols <- r.Proto

				w.Header().Set(servewaf.DispositionHeader, servewaf.DispositionDeny)
				w.WriteHeader(http.StatusForbidden)
			}
		}))
	backend.EnableHTTP2 = true
	backend.StartTLS()
	t.Cleanup(backend.Close)

	// rproxy has no setting for the certificates it trusts, so it is told of this
	// one as Go's client is told of any, in a process of its own.
	certificates := filepath.Join(t.TempDir(), "servewaf.pem")

	err := os.WriteFile(certificates, pem.EncodeToMemory(&pem.Block{
		Type: "CERTIFICATE", Bytes: backend.Certificate().Raw,
	}), 0o600)
	if err != nil {
		t.Fatalf("write the certificate: %v", err)
	}

	addr, _, _ := startRProxyProcess(t,
		wafBackend+"="+backend.URL,
		"SWRP_UPSTREAM_URL="+startApp(t),
		"SSL_CERT_FILE="+certificates)

	for name, headers := range map[string]string{
		"no such header":      "",
		"Upgrade":             "Upgrade: anything\r\n",
		"Connection":          "Connection: anything\r\n",
		"a WebSocket upgrade": "Connection: Upgrade\r\nUpgrade: websocket\r\n",
	} {
		status := sendRaw(t, addr, "GET / HTTP/1.1\r\nHost: site.example\r\n"+headers+"\r\n")
		if status != http.StatusForbidden {
			t.Errorf("with %s: status %d, want servewaf's 403", name, status)
		}

		select {
		case protocol := <-protocols:
			if protocol != "HTTP/1.1" {
				t.Errorf("with %s: asked over %s, want HTTP/1.1", name, protocol)
			}
		case <-time.After(waitLimit):
			t.Errorf("with %s: servewaf was not asked", name)
		}
	}
}

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

	hostname, err := os.Hostname()
	if err != nil {
		t.Fatalf("the host's name: %v", err)
	}

	// Nothing listens at the backend's address; no request asks it.
	_, out := startRProxy(t, "http://127.0.0.1:9", startApp(t))

	line := out.line(t, "msg", "starting")
	if line["instance"] != hostname || line["level"] != "INFO" {
		t.Errorf("rproxy's line is %v, want instance %s at INFO", line, hostname)
	}
}

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

	reportGotThere := make(chan struct{})
	release := make(chan struct{})
	backend := httptest.NewServer(http.HandlerFunc(
		func(w http.ResponseWriter, r *http.Request) {
			if r.Header.Get(servewaf.ClientHeader) != "" {
				w.Header().Set(servewaf.DispositionHeader, servewaf.DispositionAllow)

				return
			}

			close(reportGotThere)
			<-release
			w.WriteHeader(http.StatusNoContent)
		}))
	t.Cleanup(backend.Close)

	ctx, stop := context.WithCancel(t.Context())
	defer stop()

	env := map[string]string{
		rproxyListenAddr:    localhost + ":0",
		"SWRP_UPSTREAM_URL": startApp(t),
		rproxyTimeout:       "60000",
		wafBackend:          backend.URL,
	}
	out := &output{}
	exited := make(chan int, 1)

	go func() {
		exited <- runRProxy(ctx, env, out)
	}()

	addr, _ := out.line(t, "msg", "starting")["address"].(string)

	wantStatus(t, "http://"+addr+"/", placed, http.StatusOK)

	select {
	case <-reportGotThere:
	case <-time.After(waitLimit):
		t.Fatal("no report")
	}

	stop()

	select {
	case <-exited:
		t.Error("rproxy stopped with a report on its way")
		close(release)

		return
	case <-time.After(10 * pollInterval):
	}

	close(release)

	select {
	case status := <-exited:
		if status != 0 {
			t.Errorf("rproxy's exit status %d, want 0; output:\n%s", status, out.text())
		}
	case <-time.After(waitLimit):
		t.Error("rproxy still running, the report having been answered")
	}
}

// sendAndKeepOpen sends request to rproxy at addr, and returns the connection. A
// read from it ends with an error once waitLimit is up.
func sendAndKeepOpen(t *testing.T, addr, request string) net.Conn {
	t.Helper()

	conn, err := (&net.Dialer{}).DialContext(t.Context(), "tcp", addr)
	if err != nil {
		t.Fatalf("dial: %v", err)
	}

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

	err = conn.SetDeadline(time.Now().Add(waitLimit))
	if err != nil {
		t.Fatalf("set the deadline: %v", err)
	}

	_, err = io.WriteString(conn, request)
	if err != nil {
		t.Fatalf("send: %v", err)
	}

	return conn
}

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

	const (
		cutOffAfter     = 5 * time.Second
		upstreamSetting = "SWRP_UPSTREAM_URL"
	)

	asked := make(chan struct{})
	release := make(chan struct{})
	app := httptest.NewServer(http.HandlerFunc(
		func(http.ResponseWriter, *http.Request) {
			close(asked)
			<-release
		}))
	t.Cleanup(app.Close)
	t.Cleanup(func() { close(release) })

	backend := httptest.NewServer(http.HandlerFunc(
		func(w http.ResponseWriter, _ *http.Request) {
			w.Header().Set(servewaf.DispositionHeader, servewaf.DispositionAllow)
		}))
	t.Cleanup(backend.Close)

	ctx, stop := context.WithCancel(t.Context())
	defer stop()

	env := map[string]string{
		rproxyListenAddr: localhost + ":0",
		upstreamSetting:  app.URL,
		wafBackend:       backend.URL,
	}
	out := &output{}
	exited := make(chan int, 1)

	go func() {
		exited <- runRProxy(ctx, env, out)
	}()

	addr, _ := out.line(t, "msg", "starting")["address"].(string)
	conn := sendAndKeepOpen(t, addr, "GET / HTTP/1.1\r\nHost: site.example\r\n\r\n")

	select {
	case <-asked:
	case <-time.After(waitLimit):
		t.Fatal("the app was not asked")
	}

	stopped := time.Now()

	stop()

	text, err := io.ReadAll(conn)
	if err != nil || len(text) != 0 {
		t.Errorf("read %q (%v), want the connection cut off", text, err)
	}

	if elapsed := time.Since(stopped); elapsed < cutOffAfter {
		t.Errorf("cut off after %s, want %s", elapsed, cutOffAfter)
	}

	select {
	case status := <-exited:
		if status != 0 {
			t.Errorf("rproxy's exit status %d, want 0; output:\n%s", status, out.text())
		}
	case <-time.After(waitLimit):
		t.Error("rproxy still running after the cut-off")
	}

	out.line(t, "msg", "requests still in progress were cut off")
}

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

	for _, tc := range []struct {
		name, setting string
		env           map[string]string
	}{
		{"no backend", wafBackend, nil},
		{
			"another invalid", rproxyTimeout,
			map[string]string{wafBackend: "http://servewaf:8080", rproxyTimeout: "soon"},
		},
	} {
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()

			out := &output{}

			if status := runRProxy(t.Context(), tc.env, out); status != 1 {
				t.Errorf("exit status %d, want 1", status)
			}

			line := out.line(t, "msg", "invalid setting")
			message, _ := line["error"].(string)

			if !strings.HasPrefix(message, tc.setting+": ") {
				t.Errorf("start refused with %v, want it to name %s", line, tc.setting)
			}
		})
	}
}

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

	taken, err := (&net.ListenConfig{}).Listen(t.Context(), "tcp", localhost+":0")
	if err != nil {
		t.Fatalf("listen: %v", err)
	}

	defer func() {
		_ = taken.Close()
	}()

	out := &output{}

	status := runRProxy(t.Context(), map[string]string{
		wafBackend:       "http://servewaf:8080",
		rproxyListenAddr: taken.Addr().String(),
	}, out)
	if status != 1 {
		t.Errorf("exit status %d, want 1", status)
	}

	out.line(t, "msg", "cannot listen on SWRP_LISTEN_ADDR")
}

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

	// Nothing listens there; the health check does not ask.
	addr, cmd, out := startRProxyProcess(t, wafBackend+"=http://127.0.0.1:9")

	req, err := http.NewRequestWithContext(t.Context(), http.MethodGet,
		"http://"+addr+servewaf.HealthPath, http.NoBody)
	if err != nil {
		t.Fatalf("new request: %v", err)
	}

	transport := &http.Transport{}
	defer transport.CloseIdleConnections()

	res, err := transport.RoundTrip(req)
	if err != nil {
		t.Fatalf("health check: %v", err)
	}

	body, err := io.ReadAll(res.Body)
	_ = res.Body.Close()

	if err != nil || string(body) != "ok\n" {
		t.Errorf("health check answered %q (%v)", body, err)
	}

	err = cmd.Process.Signal(syscall.SIGTERM)
	if err != nil {
		t.Fatalf("signal: %v", err)
	}

	err = cmd.Wait()
	if err != nil {
		t.Errorf("exit: %v; output:\n%s", err, out.text())
	}
}