Serve available pubkeys lazily

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**
 

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",

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

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 "<type> <key>".
+	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")
+}

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 <type> <key> <host>. we want to drop the host
-	return bytes.Join(bytes.Fields(pubkey)[0:2], []byte(" ")), nil
+	// pubkey is in the form <type> <key> <host>. 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)