Support both ECDSA + RSA for SSH Server (#65)

Author: Kaiede

remote_shell_service.go

@@ -2,6 +2,8 @@ package main
 
 import (
 	"bufio"
+	"crypto/ecdsa"
+	"crypto/elliptic"
 	"crypto/rand"
 	"crypto/rsa"
 	"crypto/subtle"
@@ -17,11 +19,17 @@ import (
 	"github.com/gliderlabs/ssh"
 	"github.com/google/uuid"
 	"go.uber.org/zap"
+	gossh "golang.org/x/crypto/ssh"
 	"golang.org/x/term"
 )
 
 type ConsoleTarget int32
 
+const (
+	RSAKeyType string = "RSA PRIVATE KEY"
+	ECKeyType         = "EC PRIVATE KEY"
+)
+
 const (
 	stdOutTarget ConsoleTarget = 0
 	stdErrTarget ConsoleTarget = 1
@@ -194,55 +202,232 @@ func consoleInRoutine(stdIn io.Reader, console *Console, logger *zap.Logger) {
 	}
 }
 
-func ensureHostKey(logger *zap.Logger) (string, error) {
+const (
+	// Current filename, hides on Linux systems.
+	HostKeyFilename string = ".hostKey.pem"
+
+	// Old filename, not hidden.
+	OldHostKeyFilename = "hostKey.pem"
+)
+
+// Use the hidden form first, but fallback to the non-hidden one if it already exists.
+func pickHostKeyPath(homeDir string) string {
+	defaultKeyfilePath := filepath.Join(homeDir, HostKeyFilename)
+	_, err := os.Stat(defaultKeyfilePath)
+	if !os.IsNotExist(err) {
+		return defaultKeyfilePath
+	}
+
+	fallbackKeyfilePath := filepath.Join(homeDir, OldHostKeyFilename)
+	_, err = os.Stat(fallbackKeyfilePath)
+	if !os.IsNotExist(err) {
+		return fallbackKeyfilePath
+	}
+
+	return defaultKeyfilePath
+}
+
+// Exists to clean up the non-hidden key file if it still exists
+func cleanupOldHostKey() error {
 	homeDir, err := os.UserHomeDir()
 	if err != nil {
-		return "", err
+		return err
 	}
 
-	keyfilePath := filepath.Join(homeDir, "hostKey.pem")
+	keyfilePath := filepath.Join(homeDir, OldHostKeyFilename)
 	_, err = os.Stat(keyfilePath)
 	if os.IsNotExist(err) {
-		logger.Info("Generating host key for remote shell server.")
-		hostKey, err := rsa.GenerateKey(rand.Reader, 4096)
+		return nil
+	}
+
+	err = os.Remove(keyfilePath)
+	if err != nil {
+		return err
+	}
+
+	_, err = os.Stat(keyfilePath)
+	if !os.IsNotExist(err) {
+		return err
+	}
+
+	return nil
+}
+
+type hostKeys struct {
+	rsaKey *rsa.PrivateKey
+	ecKey  *ecdsa.PrivateKey
+}
+
+func populateKeys(keys *hostKeys, logger *zap.Logger) (bool, error) {
+	didAdd := false
+	if keys.ecKey == nil {
+		logger.Info("Generating ECDSA SSH Host Key")
+		ellipticKey, err := ecdsa.GenerateKey(elliptic.P384(), rand.Reader)
 		if err != nil {
-			return keyfilePath, err
+			return didAdd, err
 		}
 
-		err = hostKey.Validate()
+		keys.ecKey = ellipticKey
+		didAdd = true
+	}
+
+	if keys.rsaKey == nil {
+		logger.Info("Generating RSA SSH Host Key")
+		rsaKey, err := rsa.GenerateKey(rand.Reader, 4096)
 		if err != nil {
-			return keyfilePath, err
+			return didAdd, err
 		}
 
-		hostDER := x509.MarshalPKCS1PrivateKey(hostKey)
-		hostBlock := pem.Block{
-			Type:    "RSA PRIVATE KEY",
-			Headers: nil,
-			Bytes:   hostDER,
+		keys.rsaKey = rsaKey
+		didAdd = true
+	}
+
+	return didAdd, nil
+}
+
+func writeKeys(hostKeyPath string, keys *hostKeys, logger *zap.Logger) error {
+	keysFile, err := os.OpenFile(hostKeyPath, os.O_CREATE+os.O_WRONLY+os.O_TRUNC, 0600)
+	if err != nil {
+		return err
+	}
+
+	defer keysFile.Close()
+
+	logger.Info(fmt.Sprintf("Writing Host Keys to %s.", hostKeyPath))
+	if keys.ecKey != nil {
+		ecDER, err := x509.MarshalECPrivateKey(keys.ecKey)
+		if err != nil {
+			return err
+		}
+
+		ecBlock := pem.Block{
+			Type:  ECKeyType,
+			Bytes: ecDER,
 		}
-		hostPEM := pem.EncodeToMemory(&hostBlock)
 
-		err = os.WriteFile(keyfilePath, hostPEM, 0600)
-		return keyfilePath, err
+		pem.Encode(keysFile, &ecBlock)
 	}
 
-	return keyfilePath, err
+	if keys.rsaKey != nil {
+		rsaDER := x509.MarshalPKCS1PrivateKey(keys.rsaKey)
+		rsaBlock := pem.Block{
+			Type:  RSAKeyType,
+			Bytes: rsaDER,
+		}
+
+		pem.Encode(keysFile, &rsaBlock)
+	}
+
+	return nil
+}
+
+func readKeys(hostKeyPath string) (*hostKeys, error) {
+	bytes, err := os.ReadFile(hostKeyPath)
+	if err != nil {
+		return nil, err
+	}
+
+	var keys hostKeys
+	for len(bytes) > 0 {
+		pemBlock, next := pem.Decode(bytes)
+		if pemBlock == nil {
+			break
+		}
+
+		switch pemBlock.Type {
+		case RSAKeyType:
+			rsaKey, err := x509.ParsePKCS1PrivateKey(pemBlock.Bytes)
+			if err != nil {
+				return &keys, err
+			}
+			keys.rsaKey = rsaKey
+		case ECKeyType:
+			ecKey, err := x509.ParseECPrivateKey(pemBlock.Bytes)
+			if err != nil {
+				return &keys, err
+			}
+			keys.ecKey = ecKey
+		}
+
+		bytes = next
+	}
+
+	return &keys, nil
+}
+
+func ensureHostKeys(logger *zap.Logger) (*hostKeys, error) {
+	homeDir, err := os.UserHomeDir()
+	if err != nil {
+		return nil, err
+	}
+
+	keyfilePath := pickHostKeyPath(homeDir)
+	defaultKeyfilePath := filepath.Join(homeDir, HostKeyFilename)
+	fileChanged := keyfilePath != defaultKeyfilePath
+	_, err = os.Stat(keyfilePath)
+	if os.IsNotExist(err) {
+		logger.Info("Generating host keys for remote shell server.")
+		var hostKeys hostKeys
+		addedKeys, err := populateKeys(&hostKeys, logger)
+
+		if (fileChanged || addedKeys) && err == nil {
+			writeKeys(defaultKeyfilePath, &hostKeys, logger)
+		}
+		return &hostKeys, err
+	} else {
+		logger.Info(fmt.Sprintf("Reading host keys for remote shell from %s.", keyfilePath))
+		hostKeys, err := readKeys(keyfilePath)
+		if err != nil {
+			return nil, err
+		}
+
+		// Populate missing keys (older files only have RSA)
+		addedKeys, err := populateKeys(hostKeys, logger)
+
+		if (fileChanged || addedKeys) && err == nil {
+			writeKeys(defaultKeyfilePath, hostKeys, logger)
+		}
+		return hostKeys, err
+	}
+}
+
+func twinKeys(keys *hostKeys) ssh.Option {
+	return func(srv *ssh.Server) error {
+		rsaSigner, err := gossh.NewSignerFromKey(keys.rsaKey)
+		if err != nil {
+			return err
+		}
+		srv.AddHostKey(rsaSigner)
+
+		ecSigner, err := gossh.NewSignerFromKey(keys.ecKey)
+		if err != nil {
+			return err
+		}
+		srv.AddHostKey(ecSigner)
+
+		return nil
+	}
 }
 
 func runRemoteShellServer(console *Console, logger *zap.Logger) {
 	logger.Info("Starting remote shell server on 2222...")
 	ssh.Handle(func(s ssh.Session) { handleSession(s, console, logger) })
 
-	hostKeyPath, err := ensureHostKey(logger)
+	hostKeys, err := ensureHostKeys(logger)
 	if err != nil {
-		logger.Error("Unable to ensure host key exists", zap.Error(err))
+		logger.Error("Unable to ensure host keys exist", zap.Error(err))
 		return
 	}
 
+	err = cleanupOldHostKey()
+	if err != nil {
+		logger.Warn("Unable to remote old host key file", zap.Error(err))
+	}
+
 	log.Fatal(ssh.ListenAndServe(
 		":2222",
 		nil,
-		ssh.HostKeyFile(hostKeyPath),
+		twinKeys(hostKeys),
 		ssh.PasswordAuth(func(ctx ssh.Context, password string) bool { return passwordHandler(ctx, password, logger) }),
 	))
 }