Merge pull request #21 from cedws/identityfile

feat: add ssh_key_file field to connections config

internal/cli/cli.go

@@ -4,15 +4,17 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"github.com/charmbracelet/huh"
-	"google.golang.org/protobuf/types/known/emptypb"
 	"net/netip"
 	"uncloud/internal/cli/client"
 	"uncloud/internal/cli/client/connector"
 	"uncloud/internal/cli/config"
+	"uncloud/internal/fs"
 	"uncloud/internal/machine"
 	"uncloud/internal/machine/api/pb"
 	"uncloud/internal/sshexec"
+
+	"github.com/charmbracelet/huh"
+	"google.golang.org/protobuf/types/known/emptypb"
 )
 
 const defaultClusterName = "default"
@@ -89,10 +91,14 @@ func (cli *CLI) ConnectCluster(ctx context.Context, clusterName string) (*client
 		if err != nil {
 			return nil, fmt.Errorf("parse SSH connection %q: %w", conn.SSH, err)
 		}
+
+		keyPath := fs.ExpandHomeDir(conn.SSHKeyFile)
+
 		sshConfig := &connector.SSHConnectorConfig{
 			User:    user,
 			Host:    host,
 			Port:    port,
+			KeyPath: keyPath,
 		}
 		return client.New(ctx, connector.NewSSHConnector(sshConfig))
 	} else if conn.TCP.IsValid() {
@@ -157,9 +163,11 @@ func (cli *CLI) initRemoteMachine(
 			return fmt.Errorf("set current cluster: %w", err)
 		}
 	}
+
 	// Save the machine's SSH connection details in the cluster config.
 	connCfg := config.MachineConnection{
 		SSH:        config.NewSSHDestination(remoteMachine.User, remoteMachine.Host, remoteMachine.Port),
+		SSHKeyFile: remoteMachine.KeyPath,
 	}
 	cli.config.Clusters[clusterName].Connections = append(cli.config.Clusters[clusterName].Connections, connCfg)
 	if err = cli.config.Save(); err != nil {

internal/cli/config/connection.go

@@ -18,6 +18,7 @@ type MachineConnection struct {
 	TCP        netip.AddrPort `toml:"tcp,omitempty"`
 	Host       string         `toml:"host,omitempty"`
 	PublicKey  secret.Secret  `toml:"public_key,omitempty"`
+	SSHKeyFile string         `toml:"ssh_key_file,omitempty"`
 }
 
 // SSHDestination represents an SSH destination string in the canonical form of "user@host:port".

internal/fs/fs.go

@@ -5,8 +5,23 @@ import (
 	"os"
 	"os/user"
 	"strconv"
+	"strings"
 )
 
+func ExpandHomeDir(path string) string {
+	if len(path) == 0 {
+		return path
+	}
+	if path[0] == '~' {
+		home, err := os.UserHomeDir()
+		if err != nil {
+			return path
+		}
+		return strings.Replace(path, "~", home, 1)
+	}
+	return path
+}
+
 // LookupUIDGID returns the user and group IDs for the given username.
 func LookupUIDGID(username string) (uid, gid int, err error) {
 	usr, err := user.Lookup(username)

internal/fs/fs_test.go

@@ -0,0 +1,22 @@
+package fs
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestExpandHomeDir(t *testing.T) {
+	t.Run("empty", func(t *testing.T) {
+		assert.Equal(t, "", ExpandHomeDir(""))
+	})
+
+	t.Run("no home", func(t *testing.T) {
+		assert.Equal(t, "/path", ExpandHomeDir("/path"))
+	})
+
+	t.Run("home", func(t *testing.T) {
+		t.Setenv("HOME", "/home/user")
+		assert.Equal(t, "/home/user/path", ExpandHomeDir("~/path"))
+	})
+}

internal/sshexec/ssh.go

@@ -2,12 +2,13 @@ package sshexec
 
 import (
 	"fmt"
-	"golang.org/x/crypto/ssh"
-	"golang.org/x/crypto/ssh/agent"
 	"net"
 	"os"
 	"strconv"
 	"time"
+
+	"golang.org/x/crypto/ssh"
+	"golang.org/x/crypto/ssh/agent"
 )
 
 func Connect(user, host string, port int, sshKeyPath string) (*ssh.Client, error) {
@@ -34,6 +35,9 @@ func Connect(user, host string, port int, sshKeyPath string) (*ssh.Client, error
 	}
 
 	keyAuth, err := privateKeyAuth(sshKeyPath)
+	if err != nil {
+		return nil, err
+	}
 	config := &ssh.ClientConfig{
 		User:            user,
 		Auth:            []ssh.AuthMethod{keyAuth},