dockerimg: update SSHTheater to use Ed25519 keys

Co-Authored-By: sketch <hello@sketch.dev>
Change-ID: sa8eb476bc38e0479k
diff --git a/dockerimg/ssh_theater.go b/dockerimg/ssh_theater.go
index 1b060cc..91d3c21 100644
--- a/dockerimg/ssh_theater.go
+++ b/dockerimg/ssh_theater.go
@@ -3,9 +3,8 @@
 import (
 	"bufio"
 	"bytes"
+	"crypto/ed25519"
 	"crypto/rand"
-	"crypto/rsa"
-	"crypto/x509"
 	"encoding/pem"
 	"fmt"
 	"io/fs"
@@ -19,7 +18,7 @@
 	"golang.org/x/crypto/ssh/knownhosts"
 )
 
-const keyBitSize = 2048
+// Ed25519 has a fixed key size, no bit size constant needed
 
 // SSHTheater does the necessary key pair generation, known_hosts updates, ssh_config file updates etc steps
 // so that ssh can connect to a locally running sketch container to other local processes like vscode without
@@ -33,6 +32,8 @@
 // Include $HOME/.config/sketch/ssh_config
 //
 // where $HOME is your home directory.
+//
+// SSHTheater uses Ed25519 keys for improved security and performance.
 type SSHTheater struct {
 	cntrName string
 	sshHost  string
@@ -58,6 +59,7 @@
 // these files do not already exist.  These key pair files are not deleted by #Cleanup,
 // so they can be re-used across invocations of sketch. This means every sketch container
 // that runs on this host will use the same ssh server identity.
+// The system uses Ed25519 keys for better security and performance.
 //
 // If this doesn't return an error, you should be able to run "ssh <cntrName>"
 // in a terminal on your host machine to open a shell into the container without having
@@ -224,15 +226,25 @@
 	return hosts
 }
 
-func encodePrivateKeyToPEM(privateKey *rsa.PrivateKey) []byte {
+func encodePrivateKeyToPEM(privateKey ed25519.PrivateKey) []byte {
 	pemBlock := &pem.Block{
-		Type:  "RSA PRIVATE KEY",
-		Bytes: x509.MarshalPKCS1PrivateKey(privateKey),
+		Type:  "OPENSSH PRIVATE KEY",
+		Bytes: MarshalED25519PrivateKey(privateKey),
 	}
 	pemBytes := pem.EncodeToMemory(pemBlock)
 	return pemBytes
 }
 
+// MarshalED25519PrivateKey encodes an Ed25519 private key in the OpenSSH private key format
+func MarshalED25519PrivateKey(key ed25519.PrivateKey) []byte {
+	// Marshal the private key using the SSH library
+	pkBytes, err := ssh.MarshalPrivateKey(key, "")
+	if err != nil {
+		panic(fmt.Sprintf("failed to marshal private key: %v", err))
+	}
+	return pem.EncodeToMemory(pkBytes)
+}
+
 func (c *SSHTheater) writeKeyToFile(keyBytes []byte, filename string) error {
 	err := c.fs.WriteFile(filename, keyBytes, 0o600)
 	return err
@@ -243,14 +255,14 @@
 		return nil, nil
 	}
 
-	privateKey, err := c.kg.GeneratePrivateKey(keyBitSize)
+	privateKey, publicKey, err := c.kg.GenerateKeyPair()
 	if err != nil {
-		return nil, fmt.Errorf("error generating private key: %w", err)
+		return nil, fmt.Errorf("error generating key pair: %w", err)
 	}
 
-	publicRsaKey, err := c.kg.GeneratePublicKey(&privateKey.PublicKey)
+	sshPublicKey, err := c.kg.ConvertToSSHPublicKey(publicKey)
 	if err != nil {
-		return nil, fmt.Errorf("error generating public key: %w", err)
+		return nil, fmt.Errorf("error converting to SSH public key: %w", err)
 	}
 
 	privateKeyPEM := encodePrivateKeyToPEM(privateKey)
@@ -259,13 +271,13 @@
 	if err != nil {
 		return nil, fmt.Errorf("error writing private key to file %w", err)
 	}
-	pubKeyBytes := ssh.MarshalAuthorizedKey(publicRsaKey)
+	pubKeyBytes := ssh.MarshalAuthorizedKey(sshPublicKey)
 
 	err = c.writeKeyToFile([]byte(pubKeyBytes), idPath+".pub")
 	if err != nil {
 		return nil, fmt.Errorf("error writing public key to file %w", err)
 	}
-	return publicRsaKey, nil
+	return sshPublicKey, nil
 }
 
 func (c *SSHTheater) addSketchHostMatchIfMissing(cfg *ssh_config.Config) error {
@@ -558,19 +570,20 @@
 
 // KeyGenerator represents an interface for generating SSH keys for testability
 type KeyGenerator interface {
-	GeneratePrivateKey(bitSize int) (*rsa.PrivateKey, error)
-	GeneratePublicKey(privateKey *rsa.PublicKey) (ssh.PublicKey, error)
+	GenerateKeyPair() (ed25519.PrivateKey, ed25519.PublicKey, error)
+	ConvertToSSHPublicKey(publicKey ed25519.PublicKey) (ssh.PublicKey, error)
 }
 
 // RealKeyGenerator is the default implementation of KeyGenerator
 type RealKeyGenerator struct{}
 
-func (kg *RealKeyGenerator) GeneratePrivateKey(bitSize int) (*rsa.PrivateKey, error) {
-	return rsa.GenerateKey(rand.Reader, bitSize)
+func (kg *RealKeyGenerator) GenerateKeyPair() (ed25519.PrivateKey, ed25519.PublicKey, error) {
+	publicKey, privateKey, err := ed25519.GenerateKey(rand.Reader)
+	return privateKey, publicKey, err
 }
 
-func (kg *RealKeyGenerator) GeneratePublicKey(privateKey *rsa.PublicKey) (ssh.PublicKey, error) {
-	return ssh.NewPublicKey(privateKey)
+func (kg *RealKeyGenerator) ConvertToSSHPublicKey(publicKey ed25519.PublicKey) (ssh.PublicKey, error) {
+	return ssh.NewPublicKey(publicKey)
 }
 
 // CheckSSHReachability checks if the user's SSH config includes the Sketch SSH config file
diff --git a/dockerimg/ssh_theater_test.go b/dockerimg/ssh_theater_test.go
index a4e41a4..3ca80a7 100644
--- a/dockerimg/ssh_theater_test.go
+++ b/dockerimg/ssh_theater_test.go
@@ -3,8 +3,8 @@
 import (
 	"bufio"
 	"bytes"
+	"crypto/ed25519"
 	"crypto/rand"
-	"crypto/rsa"
 	"fmt"
 	"io/fs"
 	"os"
@@ -186,50 +186,52 @@
 
 // MockKeyGenerator implements KeyGenerator interface for testing
 type MockKeyGenerator struct {
-	privateKey *rsa.PrivateKey
-	publicKey  ssh.PublicKey
-	FailOn     map[string]error
+	privateKey   ed25519.PrivateKey
+	publicKey    ed25519.PublicKey
+	sshPublicKey ssh.PublicKey
+	FailOn       map[string]error
 }
 
-func NewMockKeyGenerator(privateKey *rsa.PrivateKey, publicKey ssh.PublicKey) *MockKeyGenerator {
+func NewMockKeyGenerator(privateKey ed25519.PrivateKey, publicKey ed25519.PublicKey, sshPublicKey ssh.PublicKey) *MockKeyGenerator {
 	return &MockKeyGenerator{
-		privateKey: privateKey,
-		publicKey:  publicKey,
-		FailOn:     make(map[string]error),
+		privateKey:   privateKey,
+		publicKey:    publicKey,
+		sshPublicKey: sshPublicKey,
+		FailOn:       make(map[string]error),
 	}
 }
 
-func (m *MockKeyGenerator) GeneratePrivateKey(bitSize int) (*rsa.PrivateKey, error) {
-	if err, ok := m.FailOn["GeneratePrivateKey"]; ok {
-		return nil, err
+func (m *MockKeyGenerator) GenerateKeyPair() (ed25519.PrivateKey, ed25519.PublicKey, error) {
+	if err, ok := m.FailOn["GenerateKeyPair"]; ok {
+		return nil, nil, err
 	}
-	return m.privateKey, nil
+	return m.privateKey, m.publicKey, nil
 }
 
-func (m *MockKeyGenerator) GeneratePublicKey(privateKey *rsa.PublicKey) (ssh.PublicKey, error) {
-	if err, ok := m.FailOn["GeneratePublicKey"]; ok {
+func (m *MockKeyGenerator) ConvertToSSHPublicKey(publicKey ed25519.PublicKey) (ssh.PublicKey, error) {
+	if err, ok := m.FailOn["ConvertToSSHPublicKey"]; ok {
 		return nil, err
 	}
-	return m.publicKey, nil
+	return m.sshPublicKey, nil
 }
 
 // setupMocks sets up common mocks for testing
-func setupMocks(t *testing.T) (*MockFileSystem, *MockKeyGenerator, *rsa.PrivateKey) {
-	// Generate a real private key using real random
-	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
+func setupMocks(t *testing.T) (*MockFileSystem, *MockKeyGenerator, ed25519.PrivateKey) {
+	// Generate a real Ed25519 key pair
+	publicKey, privateKey, err := ed25519.GenerateKey(rand.Reader)
 	if err != nil {
-		t.Fatalf("Failed to generate test private key: %v", err)
+		t.Fatalf("Failed to generate test key pair: %v", err)
 	}
 
-	// Generate a test public key
-	publicKey, err := ssh.NewPublicKey(&privateKey.PublicKey)
+	// Generate a test SSH public key
+	sshPublicKey, err := ssh.NewPublicKey(publicKey)
 	if err != nil {
-		t.Fatalf("Failed to generate test public key: %v", err)
+		t.Fatalf("Failed to generate test SSH public key: %v", err)
 	}
 
 	// Create mocks
 	mockFS := NewMockFileSystem()
-	mockKG := NewMockKeyGenerator(privateKey, publicKey)
+	mockKG := NewMockKeyGenerator(privateKey, publicKey, sshPublicKey)
 
 	return mockFS, mockKG, privateKey
 }
@@ -314,7 +316,7 @@
 
 	// Verify public key content format
 	pubKeyContent, _ := mockFS.ReadFile(pubKeyPath)
-	if !bytes.HasPrefix(pubKeyContent, []byte("ssh-rsa ")) {
+	if !bytes.HasPrefix(pubKeyContent, []byte("ssh-ed25519 ")) {
 		t.Errorf("Public key does not have expected format, got: %s", pubKeyContent)
 	}
 }
@@ -471,9 +473,9 @@
 	ssh, mockFS, _ := setupTestSSHTheater(t)
 
 	// Setup server public key
-	privateKey, _ := ssh.kg.GeneratePrivateKey(2048)
-	publicKey, _ := ssh.kg.GeneratePublicKey(&privateKey.PublicKey)
-	ssh.serverPublicKey = publicKey
+	_, publicKey, _ := ssh.kg.GenerateKeyPair()
+	sshPublicKey, _ := ssh.kg.ConvertToSSHPublicKey(publicKey)
+	ssh.serverPublicKey = sshPublicKey
 
 	// Create host line to be removed
 	hostLine := "[localhost]:2222 ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQ..."
@@ -524,14 +526,14 @@
 	knownHostsPath := filepath.Join(tempDir, "known_hosts")
 	serverIdentityPath := filepath.Join(tempDir, "server_identity")
 
-	// Create private key for server key
-	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
+	// Create keys for server key
+	publicKey, _, err := ed25519.GenerateKey(rand.Reader)
 	if err != nil {
-		t.Fatalf("Failed to generate private key: %v", err)
+		t.Fatalf("Failed to generate key pair: %v", err)
 	}
-	publicKey, err := ssh.NewPublicKey(&privateKey.PublicKey)
+	sshPublicKey, err := ssh.NewPublicKey(publicKey)
 	if err != nil {
-		t.Fatalf("Failed to generate public key: %v", err)
+		t.Fatalf("Failed to generate SSH public key: %v", err)
 	}
 
 	// Initialize files
@@ -551,7 +553,7 @@
 		userIdentityPath:   userIdentityPath,
 		knownHostsPath:     knownHostsPath,
 		serverIdentityPath: serverIdentityPath,
-		serverPublicKey:    publicKey,
+		serverPublicKey:    sshPublicKey,
 		fs:                 &RealFileSystem{},
 		kg:                 &RealKeyGenerator{},
 	}
@@ -670,7 +672,7 @@
 	// Test directory creation failure
 	mockFS := NewMockFileSystem()
 	mockFS.FailOn["MkdirAll"] = fmt.Errorf("mock mkdir error")
-	mockKG := NewMockKeyGenerator(nil, nil)
+	mockKG := NewMockKeyGenerator(nil, nil, nil)
 
 	// Set HOME environment variable for the test
 	oldHome := os.Getenv("HOME")
@@ -685,8 +687,8 @@
 
 	// Test key generation failure
 	mockFS = NewMockFileSystem()
-	mockKG = NewMockKeyGenerator(nil, nil)
-	mockKG.FailOn["GeneratePrivateKey"] = fmt.Errorf("mock key generation error")
+	mockKG = NewMockKeyGenerator(nil, nil, nil)
+	mockKG.FailOn["GenerateKeyPair"] = fmt.Errorf("mock key generation error")
 
 	_, err = newSSHTheatherWithDeps("test-container", "localhost", "2222", mockFS, mockKG)
 	if err == nil || !strings.Contains(err.Error(), "key generation error") {