From fbe907722f4c6b0a23c3f32f5355fd8e5a0f968f Mon Sep 17 00:00:00 2001 From: alexhulbert Date: Fri, 23 May 2025 19:20:33 -0400 Subject: [PATCH] Add cli option for passing a second key --- README.md | 4 +++- cmd/httpserver/main.go | 12 ++++++++--- httpserver/handler.go | 2 +- httpserver/handler_test.go | 14 ++++++++----- httpserver/server.go | 41 +++++++++++++++++++++++++++++++------- 5 files changed, 56 insertions(+), 17 deletions(-) diff --git a/README.md b/README.md index ad25475..c462c61 100644 --- a/README.md +++ b/README.md @@ -24,9 +24,11 @@ make build-httpserver **Run pubkey server** ```bash -go run ./cmd/httpserver/main.go [--listen-addr=127.0.0.1:8080] [--ssh-pubkey-file=/etc/ssh/ssh_host_ed25519_key.pub] +go run ./cmd/httpserver/main.go [--listen-addr=127.0.0.1:8080] [--host-ssh-pubkey-file=/etc/ssh/ssh_host_ed25519_key.pub] [--container-ssh-pubkey-file=/path/to/container_key.pub] ``` +The server will serve both pubkeys (if provided) at the `/pubkey` endpoint, separated by a newline. The `--container-ssh-pubkey-file` flag is optional. + **Install dev dependencies** ```bash diff --git a/cmd/httpserver/main.go b/cmd/httpserver/main.go index b46b5d8..3ecde25 100644 --- a/cmd/httpserver/main.go +++ b/cmd/httpserver/main.go @@ -15,9 +15,14 @@ import ( var flags []cli.Flag = []cli.Flag{ &cli.StringFlag{ - Name: "ssh-pubkey-file", + Name: "host-ssh-pubkey-file", Value: "/etc/ssh/ssh_host_ed25519_key.pub", - Usage: "path to file containing pubkey to serve", + Usage: "path to file containing host pubkey to serve", + }, + &cli.StringFlag{ + Name: "container-ssh-pubkey-file", + Value: "", + Usage: "path to file containing container pubkey to serve", }, &cli.StringFlag{ Name: "listen-addr", @@ -99,7 +104,8 @@ func main() { ReadTimeout: 60 * time.Second, WriteTimeout: 30 * time.Second, - SSHPubkeyPath: cCtx.String("ssh-pubkey-file"), + HostSSHPubkeyPath: cCtx.String("host-ssh-pubkey-file"), + ContainerSSHPubkeyPath: cCtx.String("container-ssh-pubkey-file"), } srv, err := httpserver.New(cfg) diff --git a/httpserver/handler.go b/httpserver/handler.go index d2f88b4..018e10b 100644 --- a/httpserver/handler.go +++ b/httpserver/handler.go @@ -18,7 +18,7 @@ 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.sshPubkey) + _, err := w.Write(s.sshPubkeys) if err != nil { s.log.Error("could not serve pubkey", "err", err) } diff --git a/httpserver/handler_test.go b/httpserver/handler_test.go index 147d63c..a7c23ff 100644 --- a/httpserver/handler_test.go +++ b/httpserver/handler_test.go @@ -29,10 +29,11 @@ func Test_Handlers_Healthcheck_Drain_Undrain(t *testing.T) { //nolint: exhaustruct s, err := New(&HTTPServerConfig{ - DrainDuration: latency, - ListenAddr: listenAddr, - Log: getTestLogger(), - SSHPubkeyPath: "./test_key.pub", + DrainDuration: latency, + ListenAddr: listenAddr, + Log: getTestLogger(), + HostSSHPubkeyPath: "./test_key.pub", + ContainerSSHPubkeyPath: "./test_key.pub", }) require.NoError(t, err) @@ -45,7 +46,10 @@ func Test_Handlers_Healthcheck_Drain_Undrain(t *testing.T) { data, err := io.ReadAll(resp.Body) require.NoError(t, err) require.Equal(t, http.StatusOK, resp.StatusCode) - require.Equal(t, data, []byte("ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIFGGVd5nQewq0hETk2Tr/P7OZxTW/4aftdfh9/cAe7FC")) + expectedKey := []byte("ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIFGGVd5nQewq0hETk2Tr/P7OZxTW/4aftdfh9/cAe7FC") + expectedOutput := append(expectedKey, '\n') + expectedOutput = append(expectedOutput, expectedKey...) + require.Equal(t, data, expectedOutput) } { // Check health diff --git a/httpserver/server.go b/httpserver/server.go index 6af944a..6fdb485 100644 --- a/httpserver/server.go +++ b/httpserver/server.go @@ -28,7 +28,8 @@ type HTTPServerConfig struct { ReadTimeout time.Duration WriteTimeout time.Duration - SSHPubkeyPath string + HostSSHPubkeyPath string + ContainerSSHPubkeyPath string } type Server struct { @@ -36,29 +37,55 @@ type Server struct { isReady atomic.Bool log *slog.Logger - sshPubkey []byte + sshPubkeys []byte srv *http.Server metricsSrv *metrics.MetricsServer } +func readAndFormatPubkey(path string) ([]byte, error) { + if path == "" { + return nil, nil + } + + pubkey, err := os.ReadFile(path) + if err != nil { + return nil, err + } + + // pubkey is in the form . we want to drop the host + return bytes.Join(bytes.Fields(pubkey)[0:2], []byte(" ")), nil +} + func New(cfg *HTTPServerConfig) (srv *Server, err error) { metricsSrv, err := metrics.New(common.PackageName, cfg.MetricsAddr) if err != nil { return nil, err } - sshPubkey, err := os.ReadFile(cfg.SSHPubkeyPath) - if err != nil { + var pubkeys [][]byte + + // Read host pubkey + if hostPubkey, err := readAndFormatPubkey(cfg.HostSSHPubkeyPath); err != nil { return nil, err + } else if hostPubkey != nil { + pubkeys = append(pubkeys, hostPubkey) } - // pubkey is in the form . we want to drop the host - sshPubkey = bytes.Join(bytes.Fields(sshPubkey)[0:2], []byte(" ")) + + // Read container pubkey + if containerPubkey, err := readAndFormatPubkey(cfg.ContainerSSHPubkeyPath); err != nil { + return nil, err + } else if containerPubkey != nil { + pubkeys = append(pubkeys, containerPubkey) + } + + // Combine pubkeys with newline + combinedPubkeys := bytes.Join(pubkeys, []byte("\n")) srv = &Server{ cfg: cfg, log: cfg.Log, - sshPubkey: sshPubkey, + sshPubkeys: combinedPubkeys, srv: nil, metricsSrv: metricsSrv, }