smallwebwaf/internal/discriminator/proxy_test.go

package discriminator_test

import (
	"bufio"
	"bytes"
	"encoding/json"
	"io"
	"log/slog"
	"maps"
	"net"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/spf13/viper"

	"sneak.berlin/go/smallwebwaf/internal/alerts"
	"sneak.berlin/go/smallwebwaf/internal/config"
	"sneak.berlin/go/smallwebwaf/internal/discriminator"
	"sneak.berlin/go/smallwebwaf/internal/lookup"
	"sneak.berlin/go/smallwebwaf/internal/requestlog"
	"sneak.berlin/go/smallwebwaf/internal/rproxy"
	"sneak.berlin/go/smallwebwaf/internal/rules"
)

const (
	// waitLimit bounds how long a test waits for what should happen.
	waitLimit = 10 * time.Second
	// pollInterval is how often a test looks for a log line.
	pollInterval = 10 * time.Millisecond
	// localhost is where every test server listens, and so the address rproxy
	// sees each test's requests come from.
	localhost = "127.0.0.1"
	// requestType is the type that marks a request log line.
	requestType = "request"
	// protocol is the protocol of every question.
	protocol = "HTTP/1.1"
	// trustLocalhost trusts the address every test connects from, and a
	// network for proxies in front of it.
	trustLocalhost = localhost + "/32,10.0.0.0/8"
	// appHost is the host every test asks for.
	appHost = "app.example"
	// client is the client's address, as a proxy names it.
	client = "203.0.113.9"
	// forwardedFor is the header that lists the client and its proxies,
	// and forwardedProto the one that gives the scheme the client used.
	forwardedFor   = "X-Forwarded-For"
	forwardedProto = "X-Forwarded-Proto"
	// secure is the scheme a client reached traefik with, and plain the one
	// rproxy serves.
	secure = "https"
	plain  = "http"
)

// The settings the tests set.
const (
	mode                 = "SWW_WAF_MODE"
	trustedProxies       = "SWW_RPROXY_TRUSTED_PROXIES"
	ipv6GroupPrefix      = "SWW_WAF_IPV6_GROUP_PREFIX"
	maxTrackedClients    = "SWW_WAF_MAX_TRACKED_CLIENTS"
	allowNets            = "SWW_WAF_ALLOW_NETS"
	rateLimitExemptNets  = "SWW_WAF_RATE_LIMIT_EXEMPT_NETS"
	denyNets             = "SWW_WAF_DENY_NETS"
	rateLimitPerMinute   = "SWW_WAF_RATE_LIMIT_PER_MINUTE"
	rateLimitPerDay      = "SWW_WAF_RATE_LIMIT_PER_DAY"
	rateLimitExemptPaths = "SWW_WAF_RATE_LIMIT_EXEMPT_PATHS"
	lookupSource         = "SWW_WAF_LOOKUP_SOURCE"
	lookupDBPath         = "SWW_WAF_LOOKUP_DB_PATH"
	lookupTimeout        = "SWW_WAF_LOOKUP_TIMEOUT"
	addLookupHeaders     = "SWW_WAF_ADD_LOOKUP_HEADERS"
	deniedCountries      = "SWW_WAF_DENIED_COUNTRIES"
	allowedCountries     = "SWW_WAF_EXCLUSIVELY_ALLOWED_COUNTRIES"
	banResponse          = "SWW_WAF_BAN_RESPONSE"
	limitBanDuration     = "SWW_WAF_LIMIT_BAN_DURATION"
	limitBanRepeatWindow = "SWW_WAF_LIMIT_BAN_REPEAT_WINDOW"
	maxBanDuration       = "SWW_WAF_MAX_BAN_DURATION"
	maxBans              = "SWW_WAF_MAX_BANS"
	banScopeV4Prefix     = "SWW_WAF_BAN_SCOPE_V4_PREFIX"
	instanceName         = "SWW_WAF_INSTANCE_NAME"
	logRequestHeaders    = "SWW_WAF_LOG_REQUEST_HEADERS"
	attackBanDuration    = "SWW_WAF_ATTACK_BAN_DURATION"
	rulesDir             = "SWW_WAF_RULES_DIR"
	wafMode              = "SWW_WAF_CORE_RULE_SET_MODE"
)

// off is the value that switches a setting off.
const off = "off"

// on is the value that switches a setting that takes true or false on.
const on = "true"

// output collects what sww writes on stdout.
type output struct {
	mu  sync.Mutex
	buf bytes.Buffer
}

// Write adds lines sww writes.
func (o *output) Write(p []byte) (int, error) {
	o.mu.Lock()
	defer o.mu.Unlock()

	return o.buf.Write(p)
}

// text returns everything written so far.
func (o *output) text() string {
	o.mu.Lock()
	defer o.mu.Unlock()

	return o.buf.String()
}

// lines returns every line written so far, decoded.
func (o *output) lines(t *testing.T) []map[string]any {
	t.Helper()
	o.mu.Lock()
	defer o.mu.Unlock()

	var lines []map[string]any

	for text := range strings.Lines(o.buf.String()) {
		var line map[string]any

		err := json.Unmarshal([]byte(text), &line)
		if err != nil {
			t.Fatalf("output line %q is not JSON: %v", text, err)
		}

		lines = append(lines, line)
	}

	return lines
}

// logLine is a request log line, as typed fields and as the JSON object
// it was written as.
type logLine struct {
	requestlog.Line

	fields map[string]any
}

// requestLines waits for count request log lines and returns them.
func (o *output) requestLines(t *testing.T, count int) []logLine {
	t.Helper()

	deadline := time.Now().Add(waitLimit)
	for time.Now().Before(deadline) {
		var found []logLine

		for _, fields := range o.lines(t) {
			if fields["type"] == requestType {
				found = append(found, decodeLine(t, fields))
			}
		}

		if len(found) >= count {
			return found
		}

		time.Sleep(pollInterval)
	}

	t.Fatalf("fewer than %d request log lines after %s", count, waitLimit)

	return nil
}

// requestLine waits for the request log line of a test's one request.
func (o *output) requestLine(t *testing.T) logLine {
	t.Helper()

	return o.requestLines(t, 1)[0]
}

// endedLines waits for count requests that went through rproxy, or to discriminator's
// own endpoints, to have ended, and returns the last log line of each, in the
// order discriminator took them. A question's last line is the one its report writes.
func (o *output) endedLines(t *testing.T, count int) []logLine {
	t.Helper()

	deadline := time.Now().Add(waitLimit)
	for time.Now().Before(deadline) {
		found := o.lastLines(t)

		ended := len(found) >= count
		for _, line := range found[:min(count, len(found))] {
			if line.DiscriminatorLine == requestlog.DiscriminatorLineDecision {
				ended = false
			}
		}

		if ended {
			return found
		}

		time.Sleep(pollInterval)
	}

	t.Fatalf("fewer than %d requests have ended after %s", count, waitLimit)

	return nil
}

// lastLines returns the last request log line written so far of each request,
// in the order discriminator took them.
func (o *output) lastLines(t *testing.T) []logLine {
	t.Helper()

	var found []logLine

	places := map[string]int{}

	for _, fields := range o.lines(t) {
		if fields["type"] != requestType {
			continue
		}

		line := decodeLine(t, fields)

		place, seen := places[line.RequestID]
		if seen {
			found[place] = line

			continue
		}

		places[line.RequestID] = len(found)
		found = append(found, line)
	}

	return found
}

// endedLine waits for a test's one request to have ended, and returns its last
// log line.
func (o *output) endedLine(t *testing.T) logLine {
	t.Helper()

	return o.endedLines(t, 1)[0]
}

// decodeLine reads a request log line's fields into a logLine.
func decodeLine(t *testing.T, fields map[string]any) logLine {
	t.Helper()

	encoded, err := json.Marshal(fields)
	if err != nil {
		t.Fatalf("encode %v: %v", fields, err)
	}

	line := logLine{fields: fields}

	err = json.Unmarshal(encoded, &line.Line)
	if err != nil {
		t.Fatalf("decode %s: %v", encoded, err)
	}

	return line
}

// startApp starts app as the app rproxy passes requests to.
func startApp(t *testing.T, app http.HandlerFunc) *httptest.Server {
	t.Helper()

	server := httptest.NewServer(app)
	t.Cleanup(server.Close)

	return server
}

// startProxy starts discriminator, and rproxy in front of the app at appURL, each with
// the settings in env on top of the defaults, and returns where rproxy listens
// and what discriminator writes.
func startProxy(t *testing.T, appURL string, env map[string]string) (string, *output) {
	t.Helper()

	return startProxyWithGeoJS(t, appURL, "", env)
}

// startProxyWithGeoJS is startProxy with clients' AS numbers and
// countries looked up at geojsURL.
func startProxyWithGeoJS(
	t *testing.T, appURL, geojsURL string, env map[string]string,
) (string, *output) {
	t.Helper()

	addr, out, _ := startProxyWithClock(t, appURL, geojsURL, time.Now, env)

	return addr, out
}

// startProxyWithClock is startProxyWithGeoJS with requests counted and
// bans made by the time now tells, and returns discriminator's server as well, whose
// Addr is where discriminator listens. Unless env sets SWW_WAF_RULES_DIR, it is an
// empty directory, of no rules, and unless it sets SWW_WAF_INSTANCE_NAME, that is app,
// the label instance of every metric.
func startProxyWithClock(
	t *testing.T, appURL, geojsURL string, now func() time.Time,
	env map[string]string,
) (string, *output, *discriminator.Server) {
	t.Helper()

	addr, out, server, _ := startProxyWithAlerts(t, appURL, geojsURL, now, env)

	return addr, out, server
}

// startProxyWithAlerts is startProxyWithClock, and returns the queue of
// the alerts discriminator raises as well, as newDiscriminator makes them.
func startProxyWithAlerts(
	t *testing.T, appURL, geojsURL string, now func() time.Time,
	env map[string]string,
) (string, *output, *discriminator.Server, *alerts.Queue) {
	t.Helper()

	server, out, alertQueue := newDiscriminator(t, geojsURL, now, env)
	serve(t, server)

	return startRProxy(t, server.Addr, appURL, env), out, server, alertQueue
}

// serve serves server on a port of its own until the test ends, and sets its
// Addr to where it listens.
func serve(t *testing.T, server *discriminator.Server) {
	t.Helper()

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

	server.Addr = listener.Addr().String()

	go func() {
		_ = server.Serve(listener)
	}()

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

// viperFrom returns a viper Set with each of settings, as the viper
// config.NewViper returns one bound to a process's environment.
func viperFrom(settings map[string]string) *viper.Viper {
	v := viper.New()
	for name, value := range settings {
		v.Set(name, value)
	}

	return v
}

// startRProxy starts rproxy in front of the app at appURL, asking the discriminator
// at discriminatorAddr, with the settings in env on top of the defaults, and returns
// where it listens. What it writes is dropped. SWW_RPROXY_TIMEOUT_MS is a minute unless
// env sets it, so that on a busy host SWW_RPROXY_DEFAULT_DISPOSITION does not decide.
func startRProxy(
	t *testing.T, discriminatorAddr, appURL string, env map[string]string,
) string {
	t.Helper()

	settings := map[string]string{
		"SWW_RPROXY_WAF_BACKEND":  "http://" + discriminatorAddr,
		"SWW_RPROXY_UPSTREAM_URL": appURL,
		"SWW_RPROXY_TIMEOUT_MS":   "60000",
	}
	maps.Copy(settings, env)

	cfg, err := config.RProxyFromEnvironment(viperFrom(settings))
	if err != nil {
		t.Fatalf("settings %v: %v", settings, err)
	}

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

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

	go func() {
		_ = server.Serve(listener)
	}()

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

	return listener.Addr().String()
}

// newDiscriminator makes discriminator's server, without starting it, and
// returns it, what it writes, and the queue of the alerts it raises, as the
// settings in env make it. No alert is sent from the queue: they wait in it,
// for the test to look at.
// With no geojsURL, there is no stand-in for GeoJS to look clients up at, and
// SWW_WAF_LOOKUP_SOURCE is off unless env sets it. While it is file, the lookup
// database SWW_WAF_LOOKUP_DB_PATH names is read. Clients are checked with AbuseIPDB
// at abuseIPDBURL while env sets SWW_WAF_ABUSEIPDB_KEY.
// SWW_WAF_CORE_RULE_SET_MODE is off unless env sets it, so that only the tests of the
// Core Rule Set have their requests inspected by it.
func newDiscriminator(
	t *testing.T, geojsURL string, now func() time.Time, env map[string]string,
) (*discriminator.Server, *output, *alerts.Queue) {
	t.Helper()

	return newDiscriminatorWith(t, discriminator.New, geojsURL, now, env)
}

// newDiscriminatorWith is newDiscriminator for the server build makes.
func newDiscriminatorWith(
	t *testing.T, build func(discriminator.Params) *discriminator.Server,
	geojsURL string, now func() time.Time, env map[string]string,
) (*discriminator.Server, *output, *alerts.Queue) {
	t.Helper()

	settings := map[string]string{
		rulesDir: t.TempDir(), instanceName: "app", wafMode: off,
	}
	if geojsURL == "" {
		settings[lookupSource] = off
	}

	maps.Copy(settings, env)

	cfg, err := config.FromEnvironment(viperFrom(settings))
	if err != nil {
		t.Fatalf("settings %v: %v", settings, err)
	}

	out := &output{}
	processLog := requestlog.NewProcessLogger(out, cfg.InstanceName, cfg.LogLevel)

	ruleFiles, err := rules.Load(rules.Params{
		Dir: cfg.RulesDir, Enabled: cfg.RulesEnabled, ProcessLog: processLog,
	})
	if err != nil {
		t.Fatalf("rule files: %v", err)
	}

	alertQueue := alerts.New(alerts.Params{
		WebhookURL: cfg.AlertWebhookURL,
		Events:     cfg.AlertEvents,
		Cooldown:   cfg.AlertCooldown,
		MaxPerHour: cfg.AlertMaxPerHour,
		Instance:   cfg.InstanceName,
		Now:        now,
		ProcessLog: processLog,
	})

	var lookupFile *lookup.File

	if cfg.LookupSource == fileSource {
		lookupFile, err = lookup.OpenFile(lookup.FileParams{
			Path: cfg.LookupDBPath, Now: now, ProcessLog: processLog, Alerts: alertQueue,
		})
		if err != nil {
			t.Fatalf("lookup database: %v", err)
		}
	}

	coreRuleSet, err := discriminator.NewCoreRuleSet(cfg)
	if err != nil {
		t.Fatalf("Core Rule Set: %v", err)
	}

	server := build(discriminator.Params{
		Config:       cfg,
		RequestLog:   out,
		ProcessLog:   processLog,
		GeoJSURL:     geojsURL,
		AbuseIPDBURL: abuseIPDBURL,
		LookupFile:   lookupFile,
		CoreRuleSet:  coreRuleSet,
		Now:          now,
		Rules:        ruleFiles,
		Alerts:       alertQueue,
	})

	return server, out, alertQueue
}

// newClient returns an HTTP client that sends requests as they are made,
// with no compression of its own.
func newClient(t *testing.T) *http.Client {
	t.Helper()

	transport := &http.Transport{DisableCompression: true}
	t.Cleanup(transport.CloseIdleConnections)

	return &http.Client{Transport: transport}
}

// answer is a response as a test reads it: the status, the headers, as
// much of the body as arrived, and the error that ended the reading, nil
// when the whole body arrived.
type answer struct {
	status int
	header http.Header
	body   []byte
	err    error
}

// readAnswer reads all of res, and closes its body.
func readAnswer(res *http.Response) answer {
	body, err := io.ReadAll(res.Body)
	_ = res.Body.Close()

	return answer{status: res.StatusCode, header: res.Header, body: body, err: err}
}

// newRequest makes a request for path to addr, where rproxy or discriminator listens.
func newRequest(t *testing.T, method, addr, path string, body io.Reader) *http.Request {
	t.Helper()

	req, err := http.NewRequestWithContext(t.Context(), method, "http://"+addr+path, body)
	if err != nil {
		t.Fatalf("new request: %v", err)
	}

	return req
}

// do sends req and reads the answer.
func do(t *testing.T, req *http.Request) answer {
	t.Helper()

	res, err := newClient(t).Do(req)
	if err != nil {
		t.Fatalf("%s %s: %v", req.Method, req.URL.Path, err)
	}

	return readAnswer(res)
}

// get sends a GET request for path to addr.
func get(t *testing.T, addr, path string) answer {
	t.Helper()

	return do(t, newRequest(t, http.MethodGet, addr, path, http.NoBody))
}

// dial opens a connection to addr, for requests the HTTP client cannot make,
// such as one that stops sending halfway.
func dial(t *testing.T, addr string) net.Conn {
	t.Helper()

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

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

	return conn
}

// send writes text to conn.
func send(t *testing.T, conn net.Conn, text string) {
	t.Helper()

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

// readResponse reads the answer to a request sent on conn.
func readResponse(t *testing.T, conn net.Conn) answer {
	t.Helper()

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

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

	return readAnswer(res)
}

// wantLine checks the request log line's status and action.
func wantLine(t *testing.T, line logLine, status int, action string) {
	t.Helper()

	if line.Status != status || line.Action != action {
		t.Errorf("log line has status %d and action %q, want %d and %q",
			line.Status, line.Action, status, action)
	}
}

// wantStatus checks an answer's status.
func wantStatus(t *testing.T, got answer, status int) {
	t.Helper()

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