smallwebwaf/internal/sww/discriminator_test.go

package sww_test

import (
	"bytes"
	"context"
	"crypto/rand"
	"io"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"
	"syscall"
	"testing"
	"time"
	// The time zone database, so that the process TestMainRunsDiscriminator
	// starts finds the zone TZ names on any host.
	_ "time/tzdata"

	"sneak.berlin/go/smallwebwaf/internal/discriminator"
	"sneak.berlin/go/smallwebwaf/internal/sww"
)

// commandForMain, while set in the environment, has this test binary run Main
// with the command it gives, in place of the tests, so that a test can run
// Main in a process of its own.
const commandForMain = "SWW_TEST_COMMAND"

func TestMain(m *testing.M) {
	command, set := os.LookupEnv(commandForMain)
	if !set {
		os.Exit(m.Run())
	}

	os.Args = strings.Fields("sww " + command)

	os.Exit(sww.Main(testVersion))
}

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

	// Kiritimati is 14 hours ahead of UTC; every time sww gives stays in UTC.
	const zone = "Pacific/Kiritimati"

	_, err := time.LoadLocation(zone)
	if err != nil {
		t.Fatalf("load %s: %v", zone, err)
	}

	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()

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

	cmd.Env = append(os.Environ(),
		"TZ="+zone,
		commandForMain+"="+discriminatorCommand,
		listenAddr+"="+localhost+":0",
		stateDir+"="+dir,
		rulesDir+"="+t.TempDir(),
		lookupSource+"=off")

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

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

	starting := out.line(t, "msg", "starting")
	addr, _ := starting["address"].(string)
	ask(t, addr, "/", http.StatusOK, discriminator.DispositionAllow)

	// The files are written as sww stops.
	err = cmd.Process.Signal(syscall.SIGTERM)
	if err != nil {
		t.Fatalf("signal: %v", err)
	}

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

	wantUTC(t, "the process line", starting["time"])
	wantUTC(t, "the request line", out.line(t, "type", "request")["time"])

	path := filepath.Join(dir, "clients.json")

	clients, err := os.ReadFile(path) //nolint:gosec // a file in the test's directory
	if err != nil {
		t.Fatalf("read clients.json: %v", err)
	}

	stamp := regexp.MustCompile(`"\d{4}-\d\d-\d\dT[^"]*"`)

	times := stamp.FindAllString(string(clients), -1)
	if len(times) == 0 {
		t.Errorf("clients.json holds no time:\n%s", clients)
	}

	for _, stored := range times {
		wantUTC(t, "clients.json", strings.Trim(stored, `"`))
	}
}

// wantUTC checks that value, a time read from the log line or file that
// where names, is in UTC.
func wantUTC(t *testing.T, where string, value any) {
	t.Helper()

	text, _ := value.(string)
	if !strings.HasSuffix(text, "Z") {
		t.Errorf("%s gives the time %v, want one in UTC", where, value)
	}
}

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

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

	for _, command := range []string{"", "serve"} {
		t.Run("command "+command, func(t *testing.T) {
			t.Parallel()

			ctx, cancel := context.WithTimeout(t.Context(), waitLimit)
			defer cancel()

			var stdout, stderr bytes.Buffer

			cmd := exec.CommandContext(ctx, binary)
			cmd.Stdout, cmd.Stderr = &stdout, &stderr

			cmd.Env = append(os.Environ(), commandForMain+"="+command)

			err := cmd.Run()

			if code := cmd.ProcessState.ExitCode(); code != 2 {
				t.Errorf("exit status %d (%v), want 2", code, err)
			}

			for _, name := range []string{"rproxy", discriminatorCommand, "healthcheck"} {
				if !strings.Contains(stderr.String(), name) {
					t.Errorf("stderr %q does not name %s", stderr.String(), name)
				}
			}

			if stdout.Len() != 0 {
				t.Errorf("stdout %q, want nothing", stdout.String())
			}
		})
	}
}

// ask sends discriminator at addr rproxy's question about a request for path from
// placed, checks the answer's status and disposition, and closes it.
func ask(t *testing.T, addr, path string, status int, disposition string) {
	t.Helper()

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

	req.Header.Set(discriminator.ClientHeader, placed)
	req.Header.Set("X-Request-ID", rand.Text())

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

	res, err := (&http.Client{Transport: transport}).Do(req)
	if err != nil {
		t.Fatalf("question: %v", err)
	}

	_, _ = io.Copy(io.Discard, res.Body)
	_ = res.Body.Close()

	if res.StatusCode != status ||
		res.Header.Get(discriminator.DispositionHeader) != disposition {
		t.Errorf("answered %d with disposition %q, want %d with %q", res.StatusCode,
			res.Header.Get(discriminator.DispositionHeader), status, disposition)
	}
}

// runDiscriminator runs `sww discriminator` with the settings in env, has use
// send it requests at the address it listens on, stops it as SIGTERM does, and
// checks that it stopped in order.
func runDiscriminator(t *testing.T, env map[string]string, use func(addr string)) {
	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)
	use(addr)
	stop()

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

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

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

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

	runDiscriminator(t, env, func(addr string) {
		ask(t, addr, "/", http.StatusOK, discriminator.DispositionAllow)
		// The default rule file bans a probe for this.
		ask(t, addr, "/.env", http.StatusForbidden, discriminator.DispositionDeny)
	})

	runDiscriminator(t, env, func(addr string) {
		ask(t, addr, "/", http.StatusForbidden, discriminator.DispositionDeny)
	})
}