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