From ac48b8bee6e7870050b3a665c4a772b4049db26c Mon Sep 17 00:00:00 2001
From: Connor Edwards <38229097+cedws@users.noreply.github.com>
Date: Fri, 7 Feb 2025 21:58:04 +0000
Subject: [PATCH] feat: add identity_file field to config connections config

---
 internal/cli/cli.go               | 16 +++++++++++-----
 internal/cli/config/connection.go |  9 +++++----
 internal/fs/fs.go                 | 13 +++++++++++++
 internal/fs/fs_test.go            | 22 ++++++++++++++++++++++
 internal/sshexec/ssh.go           |  8 ++++++--
 5 files changed, 57 insertions(+), 11 deletions(-)
 create mode 100644 internal/fs/fs_test.go

diff --git a/internal/cli/cli.go b/internal/cli/cli.go
index 25b6c62d..904b0b5b 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.IdentityFile)
+
 		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() {
diff --git a/internal/cli/config/connection.go b/internal/cli/config/connection.go
index 8a4158df..736f4a3f 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"`
+	IdentityFile string         `toml:"identity_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..6794a068 100644
--- a/internal/fs/fs.go
+++ b/internal/fs/fs.go
@@ -5,8 +5,21 @@ import (
 	"os"
 	"os/user"
 	"strconv"
+	"strings"
 )
 
+func ExpandHomeDir(path string) string {
+	if len(path) == 0 {
+		return path
+	}
+	if path[0] == '~' {
+		// TODO: Improve compat with other OSes
+		path = strings.Replace(path, "~", "${HOME}", 1)
+		return os.ExpandEnv(path)
+	}
+	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},