smallwebwaf/internal/sww/healthcheck_test.go

package sww_test

import (
	"bytes"
	"context"
	"net"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

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

const (
	// commandSetting says which command a container runs, and which the health
	// check checks.
	commandSetting = "SWW_COMMAND"
	// discriminatorCommand is the command that runs discriminator.
	discriminatorCommand = "discriminator"
)

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

	ctx, stop := context.WithCancel(t.Context())
	defer stop()

	out := &output{}
	exited := make(chan int, 1)
	settings := map[string]string{
		listenAddr: localhost + ":0",
		stateDir:   t.TempDir(),
		rulesDir:   t.TempDir(),
	}

	go func() {
		exited <- run(ctx, settings, out)
	}()

	addr, _ := out.line(t, "msg", "starting")["address"].(string)
	_, port, _ := net.SplitHostPort(addr)

	// The container's settings: an address to listen on with an empty
	// host part, which the health check asks at 127.0.0.1.
	wantHealthCheck(t, map[string]string{
		commandSetting: discriminatorCommand, listenAddr: ":" + port,
	}, 0, "")

	// The health check reads that setting alone, here given as a file: a removed
	// or invalid token file, or an invalid value of another setting, does not
	// fail it.
	for _, other := range []struct{ name, value string }{
		{"SWW_WAF_METRICS_TOKEN_FILE", filepath.Join(t.TempDir(), "removed")},
		{"SWW_WAF_METRICS_TOKEN_FILE", writeFile(t, "too short\n")},
		{"SWW_WAF_MODE", "neither"},
	} {
		wantHealthCheck(t, map[string]string{
			commandSetting:       discriminatorCommand,
			listenAddr + "_FILE": writeFile(t, ":"+port+"\n"),
			other.name:           other.value,
		}, 0, "")
	}

	// Listening on port 0 takes any free port, so nothing listens on port 0
	// itself.
	wantHealthCheck(t, map[string]string{
		commandSetting: discriminatorCommand, listenAddr: ":0",
	}, 1, "unhealthy: ask discriminator: ")
	wantHealthCheck(t, map[string]string{
		commandSetting: discriminatorCommand, listenAddr: "8080",
	}, 1, "unhealthy: invalid setting: SWW_WAF_LISTEN_ADDR: ")

	stop()

	select {
	case <-exited:
	case <-time.After(waitLimit):
		t.Fatal("still running after being told to stop")
	}
}

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

	app := httptest.NewServer(http.NotFoundHandler())
	defer app.Close()

	// Nothing listens at the backend. The health check asks rproxy and the app,
	// whether or not discriminator answers.
	site, _ := startRProxy(t, "http://127.0.0.1:9", app.URL)
	_, port, _ := net.SplitHostPort(strings.TrimPrefix(site, "http://"))
	settings := func(more ...string) map[string]string {
		env := map[string]string{rproxyListenAddr: ":" + port, rproxyUpstreamURL: app.URL}
		for i := 0; i < len(more); i += 2 {
			env[more[i]] = more[i+1]
		}

		return env
	}

	wantHealthCheck(t, settings(), 0, "")
	wantHealthCheck(t, settings(commandSetting, ""), 0, "")
	wantHealthCheck(t, settings(commandSetting, "rproxy"), 0, "")

	// SWW_COMMAND cannot be given as a file.
	wantHealthCheck(t, settings(commandSetting+"_FILE", writeFile(t, "discriminator\n")),
		0, "")

	// Another setting does not fail it, even an invalid one, and neither does
	// SWW_RPROXY_WAF_BACKEND being unset.
	wantHealthCheck(t, settings("SWW_RPROXY_TIMEOUT_MS", "soon"), 0, "")

	// Another test can take a closed server's port at once, so the app's
	// address here is a multicast address, to which TCP cannot connect.
	wantHealthCheck(t, settings(rproxyUpstreamURL, "http://224.0.0.1"), 1,
		"unhealthy: connect to the app: ")

	// Listening on port 0 takes any free port, so nothing listens on port 0
	// itself.
	wantHealthCheck(t, settings(rproxyListenAddr, ":0"), 1, "unhealthy: ask rproxy: ")
	wantHealthCheck(t, settings(rproxyListenAddr, "8080"), 1,
		"unhealthy: invalid setting: SWW_RPROXY_LISTEN_ADDR: ")

	// A setting of rproxy's can be given as a file too, here the listen
	// address, which the health check reads.
	wantHealthCheck(t, map[string]string{
		rproxyListenAddr + "_FILE": writeFile(t, ":"+port+"\n"),
		rproxyUpstreamURL:          app.URL,
	}, 0, "")
}

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

	wantHealthCheck(t, map[string]string{commandSetting: "healthcheck"}, 1,
		"unhealthy: invalid setting: SWW_COMMAND: "+
			"\"healthcheck\" is not rproxy or discriminator\n")
}

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

	var stderr bytes.Buffer

	got := sww.HealthCheck(t.Context(), []string{"now"}, viperFrom(nil), &stderr)

	want := "sww healthcheck: unexpected argument \"now\"\n"
	if got != 1 || stderr.String() != want {
		t.Errorf("health check returned %d and wrote %q, want 1 and %q",
			got, stderr.String(), want)
	}
}

// wantHealthCheck runs the health check with the settings in env, and
// checks its exit status and the start of what it writes to stderr,
// which is nothing when message is empty.
func wantHealthCheck(t *testing.T, env map[string]string, status int, message string) {
	t.Helper()

	var stderr bytes.Buffer

	got := sww.HealthCheck(t.Context(), nil, viperFrom(env), &stderr)

	wrote := stderr.String()
	if got != status || !strings.HasPrefix(wrote, message) ||
		(message == "" && wrote != "") {
		t.Errorf("health check returned %d and wrote %q, want %d and %q",
			got, wrote, status, message)
	}
}

// writeFile writes contents to a file in a directory of its own, removed
// when the test ends, and returns the file's path.
func writeFile(t *testing.T, contents string) string {
	t.Helper()

	path := filepath.Join(t.TempDir(), "setting")

	err := os.WriteFile(path, []byte(contents), 0o600)
	if err != nil {
		t.Fatalf("write %s: %v", path, err)
	}

	return path
}