Rename the program to sww and its deciding daemon to discriminator (closes #218) (#221)

internal/sww/discriminator_test.go

@@ -0,0 +1,204 @@
+package sww_test
+
+import (
+	"bytes"
+	"context"
+	"crypto/rand"
+	"io"
+	"net/http"
+	"os"
+	"os/exec"
+	"path/filepath"
+	"strings"
+	"syscall"
+	"testing"
+	"time"
+
+	"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()
+
+	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+"="+discriminatorCommand,
+		listenAddr+"="+localhost+":0",
+		stateDir+"="+t.TempDir(),
+		rulesDir+"="+t.TempDir(),
+		lookupSource+"=off")
+
+	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)
+	ask(t, addr, "/", http.StatusOK, discriminator.DispositionAllow)
+
+	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())
+	}
+}
+
+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)
+	})
+}