package httpserver import ( "io" "log/slog" "net/http" "net/http/httptest" "os" "path/filepath" "testing" "time" "github.com/flashbots/ssh-pubkey-server/common" "github.com/stretchr/testify/require" ) func getTestLogger() *slog.Logger { return common.SetupLogger(&common.LoggingOpts{ Debug: true, JSON: false, Service: "test", Version: "test", }) } func Test_Handlers_Healthcheck_Drain_Undrain(t *testing.T) { const ( latency = 200 * time.Millisecond listenAddr = ":8080" ) //nolint: exhaustruct s, err := New(&HTTPServerConfig{ DrainDuration: latency, ListenAddr: listenAddr, Log: getTestLogger(), SSHPubkeyPaths: []string{"./test_key.pub", "./test_key.pub"}, }) require.NoError(t, err) { // Check pubkey req := httptest.NewRequest(http.MethodGet, "http://localhost"+listenAddr+"/pubkey", nil) //nolint:goconst,nolintlint w := httptest.NewRecorder() s.handleGetPubkey(w, req) resp := w.Result() defer resp.Body.Close() data, err := io.ReadAll(resp.Body) require.NoError(t, err) require.Equal(t, http.StatusOK, resp.StatusCode) expectedKey := []byte("ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIFGGVd5nQewq0hETk2Tr/P7OZxTW/4aftdfh9/cAe7FC") expectedOutput := append(expectedKey, '\n') expectedOutput = append(expectedOutput, expectedKey...) require.Equal(t, expectedOutput, data) } { // Check health req := httptest.NewRequest(http.MethodGet, "http://localhost"+listenAddr+"/readyz", nil) //nolint:goconst,nolintlint w := httptest.NewRecorder() s.handleReadinessCheck(w, req) resp := w.Result() defer resp.Body.Close() _, err := io.ReadAll(resp.Body) require.NoError(t, err) require.Equal(t, http.StatusOK, resp.StatusCode, "Healthcheck must return `Ok` before draining") } { // Drain req := httptest.NewRequest(http.MethodGet, "http://localhost"+listenAddr+"/drain", nil) w := httptest.NewRecorder() start := time.Now() s.handleDrain(w, req) duration := time.Since(start) resp := w.Result() defer resp.Body.Close() _, err := io.ReadAll(resp.Body) require.NoError(t, err) require.Equal(t, http.StatusOK, resp.StatusCode, "Must return `Ok` for calls to `/drain`") require.GreaterOrEqual(t, duration, latency, "Must wait long enough during draining") } { // Check health req := httptest.NewRequest(http.MethodGet, "http://localhost"+listenAddr+"/readyz", nil) w := httptest.NewRecorder() s.handleReadinessCheck(w, req) resp := w.Result() defer resp.Body.Close() _, err := io.ReadAll(resp.Body) require.NoError(t, err) require.Equal(t, http.StatusServiceUnavailable, resp.StatusCode, "Healthcheck must return `Service Unavailable` after draining") } { // Undrain req := httptest.NewRequest(http.MethodGet, "http://localhost"+listenAddr+"/undrain", nil) w := httptest.NewRecorder() s.handleUndrain(w, req) resp := w.Result() defer resp.Body.Close() _, err := io.ReadAll(resp.Body) require.NoError(t, err) require.Equal(t, http.StatusOK, resp.StatusCode, "Must return `Ok` for calls to `/undrain`") time.Sleep(latency) } { // Check health req := httptest.NewRequest(http.MethodGet, "http://localhost"+listenAddr+"/readyz", nil) w := httptest.NewRecorder() s.handleReadinessCheck(w, req) resp := w.Result() defer resp.Body.Close() _, err := io.ReadAll(resp.Body) require.NoError(t, err) require.Equal(t, http.StatusOK, resp.StatusCode, "Healthcheck must return `Ok` after undraining") } } func Test_Handlers_Pubkey_LazyAvailability(t *testing.T) { dir := t.TempDir() pathA := filepath.Join(dir, "a.pub") pathB := filepath.Join(dir, "b.pub") // The host field is dropped by readAndFormatPubkey, so the served output is // just " ". require.NoError(t, os.WriteFile(pathA, []byte("ssh-ed25519 AAAAKEYA comment"), 0o600)) expectedA := []byte("ssh-ed25519 AAAAKEYA") expectedB := []byte("ssh-ed25519 AAAAKEYB") //nolint: exhaustruct s, err := New(&HTTPServerConfig{ ListenAddr: ":8080", Log: getTestLogger(), SSHPubkeyPaths: []string{pathA, pathB}, // pathB does not exist yet }) require.NoError(t, err) get := func() (int, []byte) { req := httptest.NewRequest(http.MethodGet, "http://localhost/pubkey", nil) w := httptest.NewRecorder() s.handleGetPubkey(w, req) resp := w.Result() defer resp.Body.Close() body, readErr := io.ReadAll(resp.Body) require.NoError(t, readErr) return resp.StatusCode, body } // Only the first key exists yet: /pubkey serves the available subset. code, body := get() require.Equal(t, http.StatusOK, code) require.Equal(t, expectedA, body, "/pubkey must return only the available key") // The second key appears later (e.g. after the disk is unlocked): served // with no restart, thanks to per-request reads. require.NoError(t, os.WriteFile(pathB, []byte("ssh-ed25519 AAAAKEYB comment"), 0o600)) code, body = get() require.Equal(t, http.StatusOK, code) require.Equal(t, []byte(string(expectedA)+"\n"+string(expectedB)), body, "/pubkey must return both keys once available") // A half-written (empty) file is skipped rather than panicking. require.NoError(t, os.WriteFile(pathA, []byte(""), 0o600)) code, body = get() require.Equal(t, http.StatusOK, code) require.Equal(t, expectedB, body, "/pubkey must skip an empty/malformed key file") // When no key is available yet, /pubkey reports not-ready. require.NoError(t, os.Remove(pathA)) require.NoError(t, os.Remove(pathB)) code, _ = get() require.Equal(t, http.StatusServiceUnavailable, code, "/pubkey must return 503 when no key is available") }