smallwebwaf/internal/smallwebwaf/healthcheck_test.go

package smallwebwaf_test

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

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

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

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

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

	out := &output{}
	exited := make(chan int, 1)
	settings := map[string]string{
		listenAddr:  localhost + ":0",
		upstreamURL: app.URL,
		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.
	env := map[string]string{listenAddr: ":" + port, upstreamURL: app.URL}

	wantHealthCheck(t, env, 0, "")

	// The health check reads those two settings alone, here given as
	// files: a removed or invalid token file, or an invalid value of
	// another setting, does not fail it.
	for _, other := range []struct{ name, value string }{
		{"SWWAF_METRICS_TOKEN_FILE", filepath.Join(t.TempDir(), "removed")},
		{"SWWAF_METRICS_TOKEN_FILE", writeFile(t, "too short\n")},
		{"SWWAF_MODE", "neither"},
	} {
		wantHealthCheck(t, map[string]string{
			listenAddr + "_FILE":  writeFile(t, ":"+port+"\n"),
			upstreamURL + "_FILE": writeFile(t, app.URL+"\n"),
			other.name:            other.value,
		}, 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, map[string]string{
		listenAddr:  ":" + port,
		upstreamURL: "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, map[string]string{listenAddr: ":0"}, 1,
		"unhealthy: ask smallwebwaf: ")
	wantHealthCheck(t, map[string]string{listenAddr: "8080"}, 1,
		"unhealthy: invalid setting: SWWAF_LISTEN_ADDR: ")

	stop()

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

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

	var stderr bytes.Buffer

	noSettings := func(string) (string, bool) {
		return "", false
	}

	got := smallwebwaf.HealthCheck(t.Context(), []string{"now"}, noSettings, &stderr)

	want := "smallwebwaf 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 := smallwebwaf.HealthCheck(t.Context(), nil, func(name string) (string, bool) {
		value, ok := env[name]

		return value, ok
	}, &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
}