Add cli option for passing a second key

Author: alexhulbert

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

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)

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

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

httpserver/server.go

@@ -28,37 +28,64 @@ type HTTPServerConfig struct {
 	ReadTimeout              time.Duration
 	WriteTimeout             time.Duration
 
-	SSHPubkeyPath string
+	HostSSHPubkeyPath      string
+	ContainerSSHPubkeyPath string
 }
 
 type Server struct {
 	cfg     *HTTPServerConfig
 	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 <type> <key> <host>. 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 <type> <key> <host>. 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,
 	}