smallwebwaf/internal/smallwebwaf/rproxy_test.go

package smallwebwaf_test

import (
	"context"
	"io"
	"maps"
	"net"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync/atomic"
	"syscall"
	"testing"
	"time"

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

const (
	// wafBackend is the setting that names servewaf to rproxy.
	wafBackend = "SWRP_WAF_BACKEND"
	// 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{
			"SWRP_LISTEN_ADDR":     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 <- runAs(ctx, true, 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 addr.
func serveWAFSettings(t *testing.T, addr string) map[string]string {
	t.Helper()

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

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

	addr, serveWAFOut, stop := startServeWAF(t, serveWAFSettings(t, localhost+":0"))
	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"] != proxy.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, localhost+":0")
	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"] != proxy.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
}

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

	for _, tc := range []struct {
		name, setting string
		env           map[string]string
	}{
		{"no backend", wafBackend, nil},
		{
			"another invalid", "SWRP_TIMEOUT_MS",
			map[string]string{wafBackend: "http://servewaf:8080", "SWRP_TIMEOUT_MS": "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",
		"SWRP_LISTEN_ADDR": 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()

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

	ctx, cancel := context.WithTimeout(t.Context(), 3*waitLimit)
	defer 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",
		// Nothing listens there; the health check does not ask.
		wafBackend+"=http://127.0.0.1:9")

	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)

	req, err := http.NewRequestWithContext(t.Context(), http.MethodGet,
		"http://"+addr+proxy.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())
	}
}