package rproxy_test

import (
	"bytes"
	"io"
	"log/slog"
	"net"
	"net/http"
	"net/http/httptest"
	"strconv"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	"sneak.berlin/go/smallwebwaf/internal/config"
	"sneak.berlin/go/smallwebwaf/internal/proxy"
	"sneak.berlin/go/smallwebwaf/internal/requestlog"
	"sneak.berlin/go/smallwebwaf/internal/rproxy"
)

// post is the start of a POST whose headers end with what follows it.
const post = "POST / HTTP/1.1\r\nHost: site.example\r\n"

// announcedLength is the header that announces a body of size bytes, and the
// end of the headers.
func announcedLength(size int) string {
	return "Content-Length: " + strconv.Itoa(size) + "\r\n\r\n"
}

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

	cfg, err := config.RProxyFromEnvironment(func(name string) (string, bool) {
		if name == "SWRP_WAF_BACKEND" {
			return nowhere, true
		}

		return "", false
	})
	if err != nil {
		t.Fatalf("settings: %v", err)
	}

	server := rproxy.New(rproxy.Params{
		Config:     cfg,
		RequestLog: io.Discard,
		ProcessLog: requestlog.NewProcessLogger(io.Discard, "test", slog.LevelInfo),
	})

	if server.Addr != ":8080" || server.MaxHeaderBytes != 28<<10 ||
		server.IdleTimeout != 2*time.Minute || server.ReadHeaderTimeout != time.Minute {
		t.Errorf("server listens on %q with header limit %d, idle time %s and "+
			"header timeout %s", server.Addr, server.MaxHeaderBytes,
			server.IdleTimeout, server.ReadHeaderTimeout)
	}
}

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

	for _, tc := range []struct {
		name, limit string
		size, asked int
	}{
		{"with the limit within the first part", "1K", 1025, 1024},
		{"with the limit past the first part", "128K", 129 << 10, questionBodyMaxBytes},
	} {
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()

			servewaf := startServeWAF(t, allow)
			app := startApp(t, func(http.ResponseWriter, *http.Request) {
				t.Error("the app was asked")
			})
			addr, _ := startRProxy(t, servewaf.URL, app.URL, map[string]string{
				maxRequestSize: tc.limit,
			})

			got := visit(t, newRequest(t, http.MethodPost, addr, "/",
				bytes.NewReader(make([]byte, tc.size))))
			if got.status != http.StatusRequestEntityTooLarge {
				t.Errorf("answered %+v, want 413", got)
			}

			if asked := servewaf.question(t); len(asked.body) != tc.asked {
				t.Errorf("the question carries %d bytes, want %d", len(asked.body), tc.asked)
			}

			report := servewaf.report(t)
			if report.Status != http.StatusRequestEntityTooLarge {
				t.Errorf("report %+v, want status 413", report)
			}
		})
	}
}

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

	servewaf := startServeWAF(t, deny)
	app := startApp(t, greet)
	addr, _ := startRProxy(t, servewaf.URL, app.URL, map[string]string{
		maxRequestSize: "1K",
	})

	got := visit(t, newRequest(t, http.MethodPost, addr, "/",
		bytes.NewReader(make([]byte, 1025))))
	if got.status != http.StatusForbidden {
		t.Errorf("answered %+v, want servewaf's 403", got)
	}
}

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

	for _, tc := range []struct {
		name, limit string
		size        int
		// mayReachApp is true when the part of the body the limit lets through
		// goes on to the app before the limit is met. The app may still not have
		// started on the request when it is cut off.
		mayReachApp bool
	}{
		{"within the first part", "1K", 2 << 10, false},
		{"past the first part", "128K", 129 << 10, true},
	} {
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()

			var asked atomic.Bool

			servewaf := startServeWAF(t, allow)
			app := startApp(t, func(w http.ResponseWriter, r *http.Request) {
				asked.Store(true)
				greet(w, r)
			})
			addr, _ := startRProxy(t, servewaf.URL, app.URL, map[string]string{
				maxRequestSize: tc.limit,
			})

			conn := dial(t, addr)
			send(t, conn, post+"Transfer-Encoding: chunked\r\n\r\n"+
				strconv.FormatInt(int64(tc.size), 16)+"\r\n"+strings.Repeat("x", tc.size)+
				"\r\n0\r\n\r\n")

			if status := readStatus(t, conn); status != http.StatusRequestEntityTooLarge {
				t.Errorf("status %d, want 413", status)
			}

			if asked.Load() && !tc.mayReachApp {
				t.Error("the app was asked")
			}
		})
	}
}

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

	large := strings.Repeat("x", 2<<10)

	t.Run("announced, is refused", func(t *testing.T) {
		t.Parallel()

		servewaf := startServeWAF(t, allow)
		app := startApp(t, func(w http.ResponseWriter, _ *http.Request) {
			w.Header().Set(contentLengthHeader, strconv.Itoa(len(large)))
			_, _ = io.WriteString(w, large)
		})
		addr, _ := startRProxy(t, servewaf.URL, app.URL, map[string]string{
			maxResponseSize: "1K",
		})

		got := get(t, addr, "/")
		if got.status != http.StatusBadGateway {
			t.Errorf("answered %d, want 502", got.status)
		}

		if report := servewaf.report(t); report.Status != http.StatusBadGateway {
			t.Errorf("report %+v, want status 502", report)
		}
	})

	t.Run("not announced, is cut off", func(t *testing.T) {
		t.Parallel()

		servewaf := startServeWAF(t, allow)
		app := startApp(t, func(w http.ResponseWriter, _ *http.Request) {
			_, _ = io.WriteString(w, large[:1])
			_ = http.NewResponseController(w).Flush()
			_, _ = io.WriteString(w, large[1:])
		})
		addr, _ := startRProxy(t, servewaf.URL, app.URL, map[string]string{
			maxResponseSize: "1K",
		})

		conn := dial(t, addr)
		send(t, conn, "GET / HTTP/1.1\r\nHost: site.example\r\n\r\n")

		text, err := io.ReadAll(conn)
		if err != nil || bytes.Contains(text, []byte(large)) {
			t.Errorf("read %d bytes (%v), want the answer cut off", len(text), err)
		}

		report := servewaf.report(t)
		if report.Status != http.StatusOK || report.ResponseBytes > 1<<10 {
			t.Errorf("report %+v, want status 200 and at most 1K of the answer", report)
		}
	})
}

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

	for _, tc := range []struct {
		name, rest string
	}{
		{"with its length announced", announcedLength(100) + "some of it"},
		{"with no length announced", "Transfer-Encoding: chunked\r\n\r\na\r\nsome of it\r\n"},
	} {
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()

			servewaf := startServeWAF(t, allow)
			app := startApp(t, greet)
			addr, out := startRProxy(t, servewaf.URL, app.URL, map[string]string{
				maxRequestTime: "500",
			})

			conn := dial(t, addr)
			send(t, conn, post+tc.rest)

			if status := readStatus(t, conn); status != http.StatusRequestTimeout {
				t.Errorf("status %d, want 408", status)
			}

			line := out.requestLine(t)
			if line["disposition"] != nil || line["request_bytes"] != 10.0 {
				t.Errorf("line %v, want no disposition and 10 bytes", line)
			}

			if len(servewaf.questions) != 0 {
				t.Error("servewaf was asked")
			}

			servewaf.wantNoReport(t)
		})
	}
}

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

	servewaf := startServeWAF(t, allow)
	app := startApp(t, greet)
	addr, out := startRProxy(t, servewaf.URL, app.URL, nil)

	conn := dial(t, addr)
	send(t, conn, post+announcedLength(100)+"some of it")
	_ = conn.Close()

	line := out.requestLine(t)
	if line["status"] != 0.0 || line["disposition"] != nil {
		t.Errorf("line %v, want no answer and no disposition", line)
	}

	if len(servewaf.questions) != 0 {
		t.Error("servewaf was asked")
	}
}

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

	for _, tc := range []struct {
		name, request string
		stopsSending  bool
	}{
		{
			"a malformed chunk", post + "Transfer-Encoding: chunked\r\n\r\nnot a chunk\r\n",
			false,
		},
		{"a visitor that stops sending", post + announcedLength(100) + "some of it", true},
	} {
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()

			servewaf := startServeWAF(t, allow)
			app := startApp(t, greet)
			addr, out := startRProxy(t, servewaf.URL, app.URL, nil)

			conn := dial(t, addr)
			send(t, conn, tc.request)

			if tc.stopsSending {
				tcpConn, ok := conn.(*net.TCPConn)
				if !ok {
					t.Fatalf("the connection is a %T", conn)
				}

				err := tcpConn.CloseWrite()
				if err != nil {
					t.Fatalf("stop sending: %v", err)
				}
			}

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

			line := out.requestLine(t)
			if line["status"] != 0.0 || line["disposition"] != nil {
				t.Errorf("line %v, want no answer and no disposition", line)
			}

			if len(servewaf.questions) != 0 {
				t.Error("servewaf was asked")
			}
		})
	}
}

// deadlineRecorder notes the time limits rproxy sets on reading the visitor's body
// and on sending its answer.
type deadlineRecorder struct {
	*httptest.ResponseRecorder

	readDeadline  time.Time
	writeDeadline time.Time
}

// SetReadDeadline notes deadline.
func (w *deadlineRecorder) SetReadDeadline(deadline time.Time) error {
	w.readDeadline = deadline

	return nil
}

// SetWriteDeadline notes deadline.
func (w *deadlineRecorder) SetWriteDeadline(deadline time.Time) error {
	w.writeDeadline = deadline

	return nil
}

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

	const (
		limit = 5 * time.Minute
		// announced is over what Go's server reads of a body left unread, so that
		// the app's own server answers without waiting for the rest.
		announced = 100 * questionBodyMaxBytes
	)

	servewaf := startServeWAF(t, allow)
	app := startApp(t, func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusNoContent)
	})
	_, _, server := startServer(t, servewaf.URL, app.URL, map[string]string{
		maxRequestTime: strconv.FormatInt(limit.Milliseconds(), 10),
	})

	// The visitor sends the first part and a little more, and then nothing.
	stalled, stop := io.Pipe()

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

	req := httptest.NewRequestWithContext(t.Context(), http.MethodPost, "/",
		io.MultiReader(bytes.NewReader(make([]byte, questionBodyMaxBytes+1)), stalled))
	req.ContentLength = announced
	req.Header.Set(contentLengthHeader, strconv.Itoa(announced))

	w := &deadlineRecorder{ResponseRecorder: httptest.NewRecorder()}

	before := time.Now()

	server.Handler.ServeHTTP(w, req)

	if w.Code != http.StatusNoContent {
		t.Errorf("status %d, want the app's 204", w.Code)
	}

	if w.readDeadline.Before(before.Add(limit)) ||
		w.readDeadline.After(time.Now().Add(limit)) {
		t.Errorf("the reading of the body ends at %s, want %s from the request",
			w.readDeadline, limit)
	}
}

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

	const limit = 5 * time.Minute

	for _, tc := range []struct {
		name, target string
		status       int
	}{
		{"servewaf's deny", "/", http.StatusForbidden},
		{"rproxy's own 400", "x:admin", http.StatusBadRequest},
	} {
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()

			servewaf := startServeWAF(t, deny)
			app := startApp(t, greet)
			_, _, server := startServer(t, servewaf.URL, app.URL, map[string]string{
				maxResponseTime: strconv.FormatInt(limit.Milliseconds(), 10),
			})

			req := httptest.NewRequestWithContext(t.Context(), http.MethodGet, tc.target,
				http.NoBody)
			w := &deadlineRecorder{ResponseRecorder: httptest.NewRecorder()}

			before := time.Now()

			server.Handler.ServeHTTP(w, req)

			if w.Code != tc.status {
				t.Errorf("status %d, want %d", w.Code, tc.status)
			}

			if w.writeDeadline.Before(before.Add(limit)) ||
				w.writeDeadline.After(time.Now().Add(limit)) {
				t.Errorf("the time limit on sending ends at %s, want %s from the answer",
					w.writeDeadline, limit)
			}
		})
	}
}

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

	servewaf := startServeWAF(t, allow)
	app := startApp(t, greet)
	addr, _ := startRProxy(t, servewaf.URL, app.URL, map[string]string{
		maxRequestTime: "3000",
	})

	conn := dial(t, addr)
	send(t, conn, post+announcedLength(2*questionBodyMaxBytes)+
		strings.Repeat("x", questionBodyMaxBytes+1))

	if status := readStatus(t, conn); status != http.StatusRequestTimeout {
		t.Errorf("status %d, want 408", status)
	}

	if asked := servewaf.question(t); len(asked.body) != questionBodyMaxBytes {
		t.Errorf("the question carries %d bytes", len(asked.body))
	}

	if report := servewaf.report(t); report.Status != http.StatusRequestTimeout {
		t.Errorf("report %+v, want status 408", report)
	}
}

// startAppThatTakesNothing starts an app that accepts connections and never
// reads from them, and returns its URL.
func startAppThatTakesNothing(t *testing.T) string {
	t.Helper()

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

	stopped := make(chan struct{})

	t.Cleanup(func() {
		close(stopped)

		_ = listener.Close()
	})

	go func() {
		for {
			conn, err := listener.Accept()
			if err != nil {
				return
			}

			go func() {
				<-stopped

				_ = conn.Close()
			}()
		}
	}()

	return "http://" + listener.Addr().String()
}

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

	// More than the connections hold while nobody reads, written a little at a
	// time, so that rproxy soon waits for the app.
	const (
		bodySize  = 64 << 20
		writeSize = 32 << 10
	)

	servewaf := startServeWAF(t, allow)
	addr, _ := startRProxy(t, servewaf.URL, startAppThatTakesNothing(t),
		map[string]string{maxRequestTime: "3000"})

	conn := dial(t, addr)
	send(t, conn, post+announcedLength(bodySize))

	go func() {
		part := make([]byte, writeSize)

		for range bodySize / writeSize {
			_, err := conn.Write(part)
			if err != nil {
				return
			}
		}
	}()

	if status := readStatus(t, conn); status != http.StatusGatewayTimeout {
		t.Errorf("status %d, want 504", status)
	}
}

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

	release := make(chan struct{})
	servewaf := startServeWAF(t, allow)
	app := startApp(t, func(http.ResponseWriter, *http.Request) {
		<-release
	})

	t.Cleanup(func() { close(release) })

	addr, _ := startRProxy(t, servewaf.URL, app.URL, map[string]string{
		maxResponseTime: "300",
	})

	got := get(t, addr, "/")
	if got.status != http.StatusGatewayTimeout {
		t.Errorf("answered %d, want 504", got.status)
	}

	if report := servewaf.report(t); report.Status != http.StatusGatewayTimeout {
		t.Errorf("report %+v, want status 504", report)
	}
}

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

	release := make(chan struct{})
	servewaf := startServeWAF(t, allow)
	app := startApp(t, func(w http.ResponseWriter, _ *http.Request) {
		_, _ = io.WriteString(w, "the start")
		_ = http.NewResponseController(w).Flush()

		<-release
	})

	t.Cleanup(func() { close(release) })

	addr, _ := startRProxy(t, servewaf.URL, app.URL, map[string]string{
		maxResponseTime: "300",
	})

	conn := dial(t, addr)
	send(t, conn, "GET / HTTP/1.1\r\nHost: site.example\r\n\r\n")

	// dial's deadline ends this read with an error if the connection stays open.
	_, err := io.ReadAll(conn)
	if err != nil {
		t.Errorf("read: %v, want the connection closed", err)
	}

	if report := servewaf.report(t); report.Status != http.StatusOK {
		t.Errorf("report %+v, want the status the visitor got", report)
	}
}

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

	// More than the connections hold while nobody reads.
	const (
		bodySize  = 64 << 20
		writeSize = 32 << 10
	)

	servewaf := startServeWAF(t, allow)
	app := startApp(t, func(w http.ResponseWriter, _ *http.Request) {
		part := make([]byte, writeSize)

		for range bodySize / writeSize {
			_, err := w.Write(part)
			if err != nil {
				return
			}
		}
	})
	addr, out := startRProxy(t, servewaf.URL, app.URL, map[string]string{
		maxResponseTime: "500",
	})

	// The visitor asks, and never reads the answer.
	conn := dial(t, addr)
	send(t, conn, "GET / HTTP/1.1\r\nHost: site.example\r\n\r\n")

	line := out.requestLine(t)

	responseBytes, _ := line["response_bytes"].(float64)
	if line["status"] != float64(http.StatusOK) || responseBytes >= bodySize {
		t.Errorf("line %v, want status 200 and the answer cut off", line)
	}

	report := servewaf.report(t)
	if report.Status != http.StatusOK || report.ResponseBytes >= bodySize {
		t.Errorf("report %+v, want status 200 and the answer cut off", report)
	}
}

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

	servewaf := startServeWAF(t, allow)
	addr, _ := startRProxy(t, servewaf.URL, nowhere, nil)

	got := get(t, addr, "/")
	if got.status != http.StatusBadGateway {
		t.Errorf("answered %d, want 502", got.status)
	}

	if report := servewaf.report(t); report.Status != http.StatusBadGateway {
		t.Errorf("report %+v, want status 502", report)
	}
}

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

	servewaf := startServeWAF(t, allow)
	app := startApp(t, func(w http.ResponseWriter, _ *http.Request) {
		_, _ = io.WriteString(w, greeting)
	})
	addr, _ := startRProxy(t, servewaf.URL, app.URL, map[string]string{
		maxRequestTime: "3000",
	})

	// Far more announced than sent, so that Go's server answers for the app
	// without waiting for the rest.
	conn := dial(t, addr)
	send(t, conn, post+announcedLength(100*questionBodyMaxBytes)+
		strings.Repeat("x", questionBodyMaxBytes))

	if status := readStatus(t, conn); status != http.StatusOK {
		t.Errorf("status %d, want the app's 200", status)
	}
}

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

	const answer = "the end of a long answer"

	servewaf := startServeWAF(t, allow)
	app := startApp(t, func(w http.ResponseWriter, r *http.Request) {
		_, _ = io.Copy(io.Discard, r.Body)
		_ = http.NewResponseController(w).Flush()

		time.Sleep(2100 * time.Millisecond)

		_, _ = io.WriteString(w, answer)
	})
	addr, _ := startRProxy(t, servewaf.URL, app.URL, map[string]string{
		maxRequestTime: "2000",
	})

	got := visit(t, newRequest(t, http.MethodPost, addr, "/", strings.NewReader("a body")))
	if got.status != http.StatusOK || got.body != answer {
		t.Errorf("answered %d %q, want the app's whole answer", got.status, got.body)
	}
}

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

	for _, tc := range []struct {
		name, start string
		status      int
	}{
		{"a deny", post, http.StatusForbidden},
		{
			"the health check",
			"GET " + proxy.HealthPath + " HTTP/1.1\r\nHost: site.example\r\n",
			http.StatusOK,
		},
	} {
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()

			servewaf := startServeWAF(t, deny)
			app := startApp(t, greet)
			addr, _ := startRProxy(t, servewaf.URL, app.URL, map[string]string{
				maxRequestTime: "500",
			})

			// Go's server waits for a little of the body left before it sends an answer,
			// and the time for sending the request ends that wait.
			conn := dial(t, addr)
			send(t, conn, tc.start+announcedLength(questionBodyMaxBytes+100)+
				strings.Repeat("x", questionBodyMaxBytes))

			if status := readStatus(t, conn); status != tc.status {
				t.Errorf("status %d, want %d", status, tc.status)
			}
		})
	}
}