From 79d29bb41eaade12d6fa56da3c38492e7d432f9c Mon Sep 17 00:00:00 2001
From: Maxime Suret <11944422+msuret@users.noreply.github.com>
Date: Sat, 9 May 2020 15:43:17 +0200
Subject: [PATCH] serve sftp: add support for multiple host keys by repeating
 --key flag

---
 cmd/serve/sftp/server.go | 46 +++++++++++++++++++++-------------------
 cmd/serve/sftp/sftp.go   | 14 ++++++------
 2 files changed, 31 insertions(+), 29 deletions(-)

diff --git a/cmd/serve/sftp/server.go b/cmd/serve/sftp/server.go
index bd87483c5..2899f2f55 100644
--- a/cmd/serve/sftp/server.go
+++ b/cmd/serve/sftp/server.go
@@ -206,34 +206,36 @@ func (s *server) serve() (err error) {
 	}
 
 	// Load the private key, from the cache if not explicitly configured
-	keyPath := s.opt.Key
+	keyPaths := s.opt.HostKeys
 	cachePath := filepath.Join(config.CacheDir, "serve-sftp")
-	if keyPath == "" {
-		keyPath = filepath.Join(cachePath, "id_rsa")
+	if len(keyPaths) == 0 {
+		keyPaths = []string{filepath.Join(cachePath, "id_rsa")}
 	}
-	private, err := loadPrivateKey(keyPath)
-	if err != nil && s.opt.Key == "" {
-		fs.Debugf(nil, "Failed to load %q: %v", keyPath, err)
-		// If loading a cached key failed, make the keys and retry
-		err = os.MkdirAll(cachePath, 0700)
-		if err != nil {
-			return errors.Wrap(err, "failed to create cache path")
+	for _, keyPath := range keyPaths {
+		private, err := loadPrivateKey(keyPath)
+		if err != nil && len(s.opt.HostKeys) == 0 {
+			fs.Debugf(nil, "Failed to load %q: %v", keyPath, err)
+			// If loading a cached key failed, make the keys and retry
+			err = os.MkdirAll(cachePath, 0700)
+			if err != nil {
+				return errors.Wrap(err, "failed to create cache path")
+			}
+			const bits = 2048
+			fs.Logf(nil, "Generating %d bit key pair at %q", bits, keyPath)
+			err = makeSSHKeyPair(bits, keyPath+".pub", keyPath)
+			if err != nil {
+				return errors.Wrap(err, "failed to create SSH key pair")
+			}
+			// reload the new keys
+			private, err = loadPrivateKey(keyPath)
 		}
-		const bits = 2048
-		fs.Logf(nil, "Generating %d bit key pair at %q", bits, keyPath)
-		err = makeSSHKeyPair(bits, keyPath+".pub", keyPath)
 		if err != nil {
-			return errors.Wrap(err, "failed to create SSH key pair")
+			return err
 		}
-		// reload the new keys
-		private, err = loadPrivateKey(keyPath)
-	}
-	if err != nil {
-		return err
-	}
-	fs.Debugf(nil, "Loaded private key from %q", keyPath)
+		fs.Debugf(nil, "Loaded private key from %q", keyPath)
 
-	s.config.AddHostKey(private)
+		s.config.AddHostKey(private)
+	}
 
 	// Once a ServerConfig has been configured, connections can be
 	// accepted.
diff --git a/cmd/serve/sftp/sftp.go b/cmd/serve/sftp/sftp.go
index 5cc27c91b..c442db261 100644
--- a/cmd/serve/sftp/sftp.go
+++ b/cmd/serve/sftp/sftp.go
@@ -19,12 +19,12 @@ import (
 
 // Options contains options for the http Server
 type Options struct {
-	ListenAddr     string // Port to listen on
-	Key            string // Path to private host key
-	AuthorizedKeys string // Path to authorized keys file
-	User           string // single username
-	Pass           string // password for user
-	NoAuth         bool   // allow no authentication on connections
+	ListenAddr     string   // Port to listen on
+	HostKeys       []string // Paths to private host keys
+	AuthorizedKeys string   // Path to authorized keys file
+	User           string   // single username
+	Pass           string   // password for user
+	NoAuth         bool     // allow no authentication on connections
 }
 
 // DefaultOpt is the default values used for Options
@@ -40,7 +40,7 @@ var Opt = DefaultOpt
 func AddFlags(flagSet *pflag.FlagSet, Opt *Options) {
 	rc.AddOption("sftp", &Opt)
 	flags.StringVarP(flagSet, &Opt.ListenAddr, "addr", "", Opt.ListenAddr, "IPaddress:Port or :Port to bind server to.")
-	flags.StringVarP(flagSet, &Opt.Key, "key", "", Opt.Key, "SSH private host key file (leave blank to auto generate)")
+	flags.StringArrayVarP(flagSet, &Opt.HostKeys, "key", "", Opt.HostKeys, "SSH private host key file (Can be multi-valued, leave blank to auto generate)")
 	flags.StringVarP(flagSet, &Opt.AuthorizedKeys, "authorized-keys", "", Opt.AuthorizedKeys, "Authorized keys file")
 	flags.StringVarP(flagSet, &Opt.User, "user", "", Opt.User, "User name for authentication.")
 	flags.StringVarP(flagSet, &Opt.Pass, "pass", "", Opt.Pass, "Password for authentication.")