Support multiple keys (#5)

* Support multiple keys

* Fix add to known hosts script

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] [--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.
+
 **Install dev dependencies**
 
 ```bash

cmd/cli/add_to_known_hosts.sh

@@ -8,5 +8,11 @@ if [ $1 = "-h" ]; then
 	exit 0;
 fi
 
-pubkey=`curl -s $1/pubkey`
-ssh-keyscan -H "$2" 2>/dev/null | grep "${pubkey}"
+pubkeys=`curl -s $1/pubkey`
+host_keys=`ssh-keyscan -H "$2" 2>/dev/null; ssh-keyscan -H -p 10022 "$2" 2>/dev/null`
+
+echo "$pubkeys" | while IFS= read -r pubkey; do
+	if [ -n "$pubkey" ]; then
+		echo "$host_keys" | grep "${pubkey}"
+	fi
+done

cmd/httpserver/main.go

@@ -14,10 +14,10 @@ import (
 )
 
 var flags []cli.Flag = []cli.Flag{
-	&cli.StringFlag{
+	&cli.StringSliceFlag{
 		Name:  "ssh-pubkey-file",
-		Value: "/etc/ssh/ssh_host_ed25519_key.pub",
-		Usage: "path to file containing pubkey to serve",
+		Value: cli.NewStringSlice("/etc/ssh/ssh_host_ed25519_key.pub"),
+		Usage: "path to file containing pubkey to serve (can be specified multiple times)",
 	},
 	&cli.StringFlag{
 		Name:  "listen-addr",
@@ -99,7 +99,7 @@ func main() {
 				ReadTimeout:              60 * time.Second,
 				WriteTimeout:             30 * time.Second,
 
-				SSHPubkeyPath: cCtx.String("ssh-pubkey-file"),
+				SSHPubkeyPaths: cCtx.StringSlice("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,10 @@ 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(),
+		SSHPubkeyPaths: []string{"./test_key.pub", "./test_key.pub"},
 	})
 	require.NoError(t, err)
 
@@ -45,7 +45,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, expectedOutput, data)
 	}
 
 	{ // Check health

httpserver/server.go

@@ -28,37 +28,57 @@ type HTTPServerConfig struct {
 	ReadTimeout              time.Duration
 	WriteTimeout             time.Duration
 
-	SSHPubkeyPath string
+	SSHPubkeyPaths []string
 }
 
 type Server struct {
 	cfg     *HTTPServerConfig
 	isReady atomic.Bool
 	log     *slog.Logger
 
-	sshPubkey []byte
+	sshPubkeys []byte
 
 	srv        *http.Server
 	metricsSrv *metrics.MetricsServer
 }
 
-func New(cfg *HTTPServerConfig) (srv *Server, err error) {
-	metricsSrv, err := metrics.New(common.PackageName, cfg.MetricsAddr)
+func readAndFormatPubkey(path string) ([]byte, error) {
+	if path == "" {
+		return nil, nil
+	}
+
+	pubkey, err := os.ReadFile(path)
 	if err != nil {
 		return nil, err
 	}
 
-	sshPubkey, err := os.ReadFile(cfg.SSHPubkeyPath)
+	// 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
 	}
-	// pubkey is in the form <type> <key> <host>. we want to drop the host
-	sshPubkey = bytes.Join(bytes.Fields(sshPubkey)[0:2], []byte(" "))
+
+	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,
-		sshPubkey:  sshPubkey,
+		sshPubkeys: combinedPubkeys,
 		srv:        nil,
 		metricsSrv: metricsSrv,
 	}