diff --git a/README.md b/README.md index 9f341f5..d78a2d5 100644 --- a/README.md +++ b/README.md @@ -27,7 +27,9 @@ make build-httpserver go run ./cmd/httpserver/main.go [--listen-addr=127.0.0.1:8080] [--ssh-pubkey-file=/etc/ssh/ssh_host_ed25519_key.pub] [--ssh-pubkey-file=/path/to/container_key.pub] ``` -You can specify multiple `--ssh-pubkey-file` flags to serve multiple public keys. The server will serve all pubkeys at the `/pubkey` endpoint, separated by newlines. +You can specify multiple `--ssh-pubkey-file` flags to serve multiple public keys. The server serves all currently-available pubkeys at the `/pubkey` endpoint, separated by newlines. + +Pubkey files are read lazily on each request, so a key that only becomes available after the server starts (for example a key generated once an encrypted disk is unlocked) is served as soon as it appears, with no restart. A file that is missing or not yet readable is simply skipped. If no key is available yet, `/pubkey` responds with `503 Service Unavailable`. **Install dev dependencies** diff --git a/cmd/httpserver/main.go b/cmd/httpserver/main.go index e12f828..5b392bf 100644 --- a/cmd/httpserver/main.go +++ b/cmd/httpserver/main.go @@ -17,7 +17,7 @@ var flags []cli.Flag = []cli.Flag{ &cli.StringSliceFlag{ Name: "ssh-pubkey-file", Value: cli.NewStringSlice("/etc/ssh/ssh_host_ed25519_key.pub"), - Usage: "path to file containing pubkey to serve (can be specified multiple times)", + Usage: "path to file containing pubkey to serve (read on each request, missing files are skipped, can be specified multiple times)", }, &cli.StringFlag{ Name: "listen-addr", diff --git a/httpserver/handler.go b/httpserver/handler.go index 018e10b..c5d9472 100644 --- a/httpserver/handler.go +++ b/httpserver/handler.go @@ -1,12 +1,18 @@ package httpserver import ( + "bytes" "net/http" "time" "github.com/flashbots/ssh-pubkey-server/metrics" ) +// handleGetPubkey serves every currently-available pubkey, newline-joined. +// Files are read on each request, and any path that is missing or not yet +// readable is skipped — so the response grows as keys become available (e.g. a +// key behind an encrypted disk that is unlocked later) without a restart. When +// no key is available yet, it responds 503. func (s *Server) handleGetPubkey(w http.ResponseWriter, r *http.Request) { m := s.metricsSrv.Float64Histogram( "request_duration_api", @@ -18,8 +24,24 @@ func (s *Server) handleGetPubkey(w http.ResponseWriter, r *http.Request) { m.Record(r.Context(), float64(time.Since(start).Microseconds())) }(time.Now()) - _, err := w.Write(s.sshPubkeys) - if err != nil { + var keys [][]byte + for _, path := range s.cfg.SSHPubkeyPaths { + pubkey, err := readAndFormatPubkey(path) + if err != nil { + s.log.Debug("pubkey not available, skipping", "path", path, "err", err) + continue + } + if pubkey != nil { + keys = append(keys, pubkey) + } + } + + if len(keys) == 0 { + w.WriteHeader(http.StatusServiceUnavailable) + return + } + + if _, err := w.Write(bytes.Join(keys, []byte("\n"))); err != nil { s.log.Error("could not serve pubkey", "err", err) } } diff --git a/httpserver/handler_test.go b/httpserver/handler_test.go index 429e401..9e47316 100644 --- a/httpserver/handler_test.go +++ b/httpserver/handler_test.go @@ -5,6 +5,8 @@ import ( "log/slog" "net/http" "net/http/httptest" + "os" + "path/filepath" "testing" "time" @@ -110,3 +112,58 @@ func Test_Handlers_Healthcheck_Drain_Undrain(t *testing.T) { 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") +} diff --git a/httpserver/server.go b/httpserver/server.go index de9cae5..81dbd46 100644 --- a/httpserver/server.go +++ b/httpserver/server.go @@ -36,12 +36,12 @@ type Server struct { isReady atomic.Bool log *slog.Logger - sshPubkeys []byte - srv *http.Server metricsSrv *metrics.MetricsServer } +var errMalformedPubkey = errors.New("malformed pubkey: want at least 2 fields") + func readAndFormatPubkey(path string) ([]byte, error) { if path == "" { return nil, nil @@ -52,8 +52,14 @@ func readAndFormatPubkey(path string) ([]byte, error) { return nil, err } - // pubkey is in the form . we want to drop the host - return bytes.Join(bytes.Fields(pubkey)[0:2], []byte(" ")), nil + // pubkey is in the form . we want to drop the host. + // A file may be empty or half-written (e.g. while the container is writing + // its host key), so guard against fewer than two fields rather than panic. + fields := bytes.Fields(pubkey) + if len(fields) < 2 { + return nil, errMalformedPubkey + } + return bytes.Join(fields[0:2], []byte(" ")), nil } func New(cfg *HTTPServerConfig) (srv *Server, err error) { @@ -62,29 +68,16 @@ func New(cfg *HTTPServerConfig) (srv *Server, err error) { return nil, err } - var pubkeys [][]byte - - // Read all specified pubkey files - for _, path := range cfg.SSHPubkeyPaths { - if pubkey, err := readAndFormatPubkey(path); err != nil { - return nil, err - } else if pubkey != nil { - pubkeys = append(pubkeys, pubkey) - } - } - - combinedPubkeys := bytes.Join(pubkeys, []byte("\n")) - srv = &Server{ cfg: cfg, log: cfg.Log, - sshPubkeys: combinedPubkeys, srv: nil, metricsSrv: metricsSrv, } srv.isReady.Swap(true) mux := chi.NewRouter() + mux.With(srv.httpLogger).Get("/pubkey", srv.handleGetPubkey) // Never serve at `/` (root) path mux.With(srv.httpLogger).Get("/livez", srv.handleLivenessCheck) mux.With(srv.httpLogger).Get("/readyz", srv.handleReadinessCheck)