smallwebwaf/internal/sww/discriminator_test.go

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