diff --git a/internal/cli/cli.go b/internal/cli/cli.go index 25b6c62d..7accb247 100644 --- a/internal/cli/cli.go +++ b/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, + 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), + 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 { diff --git a/internal/cli/config/connection.go b/internal/cli/config/connection.go index 8a4158df..97273872 100644 --- a/internal/cli/config/connection.go +++ b/internal/cli/config/connection.go @@ -14,10 +14,11 @@ const ( ) type MachineConnection struct { - SSH SSHDestination `toml:"ssh,omitempty"` - TCP netip.AddrPort `toml:"tcp,omitempty"` - Host string `toml:"host,omitempty"` - PublicKey secret.Secret `toml:"public_key,omitempty"` + SSH SSHDestination `toml:"ssh,omitempty"` + 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". diff --git a/internal/fs/fs.go b/internal/fs/fs.go index 61330db7..86e7a33d 100644 --- a/internal/fs/fs.go +++ b/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) diff --git a/internal/fs/fs_test.go b/internal/fs/fs_test.go new file mode 100644 index 00000000..5abac425 --- /dev/null +++ b/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")) + }) +} diff --git a/internal/sshexec/ssh.go b/internal/sshexec/ssh.go index 8e741b63..11ee94f8 100644 --- a/internal/sshexec/ssh.go +++ b/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},