diff --git a/internal/cli/config/connection.go b/internal/cli/config/connection.go index 2051b95c..3b736b99 100644 --- a/internal/cli/config/connection.go +++ b/internal/cli/config/connection.go @@ -30,7 +30,7 @@ type MachineConnection struct { MachineID string `yaml:"machine_id,omitempty"` } -func (c MachineConnection) String() string { +func (c *MachineConnection) String() string { if c.SSH != "" { return "ssh://" + string(c.SSH) } else if c.SSHCLI != "" { @@ -69,36 +69,37 @@ func (c *MachineConnection) Validate() error { } // SSHDestination represents an SSH destination string in the canonical form of "user@host:port". -// The default user "root" and port 22 can be omitted. +// Empty user or port components are omitted. type SSHDestination string +// NewSSHDestination constructs an SSHDestination from user, host, and port components. +// If user is empty, it is omitted. +// If port is 0, it is omitted. func NewSSHDestination(user, host string, port int) SSHDestination { dst := host - if port != 0 && port != DefaultSSHPort { + if port != 0 { dst = net.JoinHostPort(host, strconv.Itoa(port)) } - if user == "" { - user = DefaultSSHUser + if user != "" { + dst = fmt.Sprintf("%s@%s", user, dst) } - dst = user + "@" + dst + return SSHDestination(dst) } +// Parse parses the SSH destination string into user, host, and port components. +// If user is not specified, it returns an empty string. +// If port is not specified, it returns 0. func (d SSHDestination) Parse() (user string, host string, port int, err error) { host = string(d) if strings.Contains(host, "@") { user, host, _ = strings.Cut(host, "@") } - if user == "" { - user = DefaultSSHUser - } h, p, sErr := net.SplitHostPort(host) if sErr == nil { host = h port, err = strconv.Atoi(p) } - if port == 0 { - port = DefaultSSHPort - } + return } diff --git a/internal/sshexec/ssh.go b/internal/sshexec/ssh.go index 435f8880..442b4acf 100644 --- a/internal/sshexec/ssh.go +++ b/internal/sshexec/ssh.go @@ -4,6 +4,7 @@ import ( "fmt" "net" "os" + osuser "os/user" "strconv" "time" @@ -13,7 +14,17 @@ import ( ) func Connect(user, host string, port int, sshKeyPath string) (*ssh.Client, error) { + // Use the current OS user if no user is specified to be make it consistent with ssh CLI behavior. + if user == "" { + if u, err := osuser.Current(); err == nil { + user = u.Username + } + } + if port == 0 { + port = 22 + } addr := net.JoinHostPort(host, strconv.Itoa(port)) + // Try to connect using SSH agent only. agentAuth, agentClose, agentErr := sshAgentAuth() if agentErr == nil { diff --git a/pkg/client/connector/ssh.go b/pkg/client/connector/ssh.go index b1dec958..dfd06595 100644 --- a/pkg/client/connector/ssh.go +++ b/pkg/client/connector/ssh.go @@ -24,6 +24,16 @@ type SSHConnectorConfig struct { SockPath string } +// Destination returns the destination string that can be passed to ssh CLI in the format [user@]host. +func (cfg *SSHConnectorConfig) Destination() string { + dst := cfg.Host + if cfg.User != "" { + dst = fmt.Sprintf("%s@%s", cfg.User, dst) + } + + return dst +} + // SSHConnector establishes a connection to the machine API through an SSH tunnel to the machine. type SSHConnector struct { config SSHConnectorConfig @@ -48,7 +58,11 @@ func (c *SSHConnector) Connect(ctx context.Context) (*grpc.ClientConn, error) { var err error c.client, err = sshexec.Connect(c.config.User, c.config.Host, c.config.Port, c.config.KeyPath) if err != nil { - return nil, fmt.Errorf("SSH login to %s@%s:%d: %w", c.config.User, c.config.Host, c.config.Port, err) + dst := c.config.Destination() + if c.config.Port != 0 { + dst = fmt.Sprintf("%s:%d", dst, c.config.Port) + } + return nil, fmt.Errorf("SSH login to %s: %w", dst, err) } } diff --git a/pkg/client/connector/sshcli.go b/pkg/client/connector/sshcli.go index bbd4831b..41afa772 100644 --- a/pkg/client/connector/sshcli.go +++ b/pkg/client/connector/sshcli.go @@ -4,6 +4,8 @@ import ( "context" "fmt" "net" + "os" + "path/filepath" "strconv" "github.com/docker/cli/cli/connhelper/commandconn" @@ -18,26 +20,72 @@ import ( type SSHCLIConnector struct { config SSHConnectorConfig conn net.Conn + // Path to SSH control socket for connection reuse. + controlSockPath string } func NewSSHCLIConnector(cfg *SSHConnectorConfig) *SSHCLIConnector { - return &SSHCLIConnector{config: *cfg} + return &SSHCLIConnector{ + config: *cfg, + controlSockPath: controlSocketPath(), + } +} + +// controlSocketPath returns a unique control socket path for the SSH connection. +// Returns an empty string if unable to find or create a suitable path. +func controlSocketPath() string { + // %C is expanded by `ssh` to a hash of user, local and remote hostnames, port, and the contents + // of the ProxyJump option. This ensures that shared connections are uniquely identified. + sockName := fmt.Sprintf("uc_control_%%C.sock") + + // Prefer XDG_RUNTIME_DIR if set, fall back to ~/.ssh if it exists. + if dir := os.Getenv("XDG_RUNTIME_DIR"); dir != "" { + return filepath.Join(dir, sockName) + } + if home, err := os.UserHomeDir(); err == nil { + sshDir := filepath.Join(home, ".ssh") + if fi, sErr := os.Stat(sshDir); sErr == nil && fi.IsDir() { + return filepath.Join(sshDir, sockName) + } + } + + // Last resort: create a subdirectory in temp with restricted permissions. + tmpDir := filepath.Join(os.TempDir(), fmt.Sprintf("uncloud-%d", os.Getuid())) + path := filepath.Join(tmpDir, sockName) + if len(path)-2+40 < 104 { // 40 chars for %C hash, 104 is typical UNIX socket path limit + if err := os.MkdirAll(tmpDir, 0o700); err == nil { + return path + } + } + + return "" } // sshCLIDialer implements proxy.ContextDialer by spawning SSH processes with -W flag. type sshCLIDialer struct { config SSHConnectorConfig + // Shared control socket path from SSHCLIConnector for connection reuse. + controlSockPath string } // buildDialArgs constructs SSH command arguments for -W flag dialing. func (d *sshCLIDialer) buildDialArgs(address string) []string { - args := []string{} + var args []string + + if d.controlSockPath != "" { + // Try to reuse the existing control connection without initiating a new one. + // Falls back to direct connection if the control socket is not available. + args = append(args, "-o", "ControlMaster=no") + args = append(args, "-o", "ControlPath="+d.controlSockPath) + } // Add connection timeout to fail fast when node is down. args = append(args, "-o", "ConnectTimeout=5") + // Disable pseudo-terminal allocation to prevent SSH from executing as a login shell. + args = append(args, "-T") - // Add port if non-standard. - if d.config.Port != 0 && d.config.Port != 22 { + // Add port if specified. + if d.config.Port != 0 { args = append(args, "-p", strconv.Itoa(d.config.Port)) } @@ -49,8 +97,8 @@ func (d *sshCLIDialer) buildDialArgs(address string) []string { // Add -W flag for stdin/stdout forwarding to target address. args = append(args, "-W", address) - // Add user@host. - args = append(args, d.config.User+"@"+d.config.Host) + // Add [user@]host destination. + args = append(args, d.config.Destination()) return args } @@ -65,28 +113,28 @@ func (d *sshCLIDialer) DialContext(ctx context.Context, network, address string) // Build SSH command arguments. args := d.buildDialArgs(address) - // Create connection using commandconn. + // Create connection using docker's commandconn. conn, err := commandconn.New(ctx, "ssh", args...) if err != nil { - return nil, fmt.Errorf("SSH connection to %s@%s for dialing %s: %w", d.config.User, d.config.Host, address, err) + return nil, fmt.Errorf("SSH connection to %s for dialing %s: %w", d.config.Destination(), address, err) } return conn, nil } func (c *SSHCLIConnector) Connect(ctx context.Context) (*grpc.ClientConn, error) { - // Build SSH command arguments. - args := c.buildSSHArgs() + if c.conn == nil { + args := c.buildSSHArgs() - // Create connection using commandconn. - conn, err := commandconn.New(ctx, "ssh", args...) - if err != nil { - return nil, fmt.Errorf("SSH connection to %s@%s: %w", c.config.User, c.config.Host, err) + // Create connection using docker's commandconn. + conn, err := commandconn.New(ctx, "ssh", args...) + if err != nil { + return nil, fmt.Errorf("SSH connection to %s: %w", c.config.Destination(), err) + } + c.conn = conn } - c.conn = conn - // Create gRPC client over the connection. - // Use a custom dialer that returns our existing connection. + // Create gRPC client over the connection. Use a custom dialer that returns our existing connection. grpcConn, err := grpc.NewClient( "passthrough:///", // Dummy target since we're using a custom dialer. grpc.WithTransportCredentials(insecure.NewCredentials()), @@ -103,15 +151,32 @@ func (c *SSHCLIConnector) Connect(ctx context.Context) (*grpc.ClientConn, error) return grpcConn, nil } -// buildSSHArgs constructs the SSH command arguments. +// buildSSHArgs constructs the SSH command arguments to run `uncloudd dial-stdio` on the remote machine reusing +// the established connection via control socket. func (c *SSHCLIConnector) buildSSHArgs() []string { - args := []string{} + var args []string + + // Add control socket options for connection reuse if available. + if c.controlSockPath != "" { + args = append(args, "-o", "ControlMaster=auto") + args = append(args, "-o", "ControlPath="+c.controlSockPath) + + // Keep the established connection alive for a short duration after the last session closes to allow reuse. + controlPersist := "10m" + // Override the default duration with the UNCLOUD_SSH_CONTROL_PERSIST env variable. + if v := os.Getenv("UNCLOUD_SSH_CONTROL_PERSIST"); v != "" { + controlPersist = v + } + args = append(args, "-o", "ControlPersist="+controlPersist) + } // Add connection timeout to fail fast when node is down. args = append(args, "-o", "ConnectTimeout=5") + // Disable pseudo-terminal allocation to prevent SSH from executing as a login shell. + args = append(args, "-T") - // Add port if non-standard. - if c.config.Port != 0 && c.config.Port != 22 { + // Add port if specified. + if c.config.Port != 0 { args = append(args, "-p", strconv.Itoa(c.config.Port)) } @@ -120,8 +185,8 @@ func (c *SSHCLIConnector) buildSSHArgs() []string { args = append(args, "-i", c.config.KeyPath) } - // Add user@host. - args = append(args, c.config.User+"@"+c.config.Host) + // Add [user@]host destination. + args = append(args, c.config.Destination()) // Add remote command: uncloudd dial-stdio args = append(args, "uncloudd", "dial-stdio") @@ -141,7 +206,8 @@ func (c *SSHCLIConnector) Dialer() (proxy.ContextDialer, error) { } return &sshCLIDialer{ - config: c.config, + config: c.config, + controlSockPath: c.controlSockPath, }, nil } diff --git a/pkg/client/connector/sshcli_test.go b/pkg/client/connector/sshcli_test.go index a3da9df9..d913c77d 100644 --- a/pkg/client/connector/sshcli_test.go +++ b/pkg/client/connector/sshcli_test.go @@ -1,6 +1,7 @@ package connector import ( + "strings" "testing" "github.com/psviderski/uncloud/internal/machine" @@ -11,17 +12,28 @@ func TestSSHCLIConnector_buildSSHArgs(t *testing.T) { t.Parallel() tests := []struct { - name string - config SSHConnectorConfig - expected []string + name string + config SSHConnectorConfig + controlSockPath string + expected []string }{ { - name: "basic connection", + name: "basic connection with control socket", config: SSHConnectorConfig{ User: "root", Host: "example.com", }, - expected: []string{"-o", "ConnectTimeout=5", "root@example.com", "uncloudd", "dial-stdio"}, + controlSockPath: "/tmp/test.sock", + expected: []string{"-o", "ControlMaster=auto", "-o", "ControlPath=/tmp/test.sock", "-o", "ControlPersist=10m", "-o", "ConnectTimeout=5", "-T", "root@example.com", "uncloudd", "dial-stdio"}, + }, + { + name: "basic connection without control socket", + config: SSHConnectorConfig{ + User: "root", + Host: "example.com", + }, + controlSockPath: "", + expected: []string{"-o", "ConnectTimeout=5", "-T", "root@example.com", "uncloudd", "dial-stdio"}, }, { name: "with custom port", @@ -30,7 +42,8 @@ func TestSSHCLIConnector_buildSSHArgs(t *testing.T) { Host: "example.com", Port: 2222, }, - expected: []string{"-o", "ConnectTimeout=5", "-p", "2222", "root@example.com", "uncloudd", "dial-stdio"}, + controlSockPath: "/tmp/test.sock", + expected: []string{"-o", "ControlMaster=auto", "-o", "ControlPath=/tmp/test.sock", "-o", "ControlPersist=10m", "-o", "ConnectTimeout=5", "-T", "-p", "2222", "root@example.com", "uncloudd", "dial-stdio"}, }, { name: "with identity file", @@ -39,7 +52,8 @@ func TestSSHCLIConnector_buildSSHArgs(t *testing.T) { Host: "example.com", KeyPath: "/path/to/key", }, - expected: []string{"-o", "ConnectTimeout=5", "-i", "/path/to/key", "root@example.com", "uncloudd", "dial-stdio"}, + controlSockPath: "/tmp/test.sock", + expected: []string{"-o", "ControlMaster=auto", "-o", "ControlPath=/tmp/test.sock", "-o", "ControlPersist=10m", "-o", "ConnectTimeout=5", "-T", "-i", "/path/to/key", "root@example.com", "uncloudd", "dial-stdio"}, }, { name: "with custom socket path", @@ -48,7 +62,8 @@ func TestSSHCLIConnector_buildSSHArgs(t *testing.T) { Host: "example.com", SockPath: "/custom/path/uncloud.sock", }, - expected: []string{"-o", "ConnectTimeout=5", "root@example.com", "uncloudd", "dial-stdio", "--socket", "/custom/path/uncloud.sock"}, + controlSockPath: "/tmp/test.sock", + expected: []string{"-o", "ControlMaster=auto", "-o", "ControlPath=/tmp/test.sock", "-o", "ControlPersist=10m", "-o", "ConnectTimeout=5", "-T", "root@example.com", "uncloudd", "dial-stdio", "--socket", "/custom/path/uncloud.sock"}, }, { name: "with default socket path (not included)", @@ -57,7 +72,8 @@ func TestSSHCLIConnector_buildSSHArgs(t *testing.T) { Host: "example.com", SockPath: machine.DefaultUncloudSockPath, }, - expected: []string{"-o", "ConnectTimeout=5", "root@example.com", "uncloudd", "dial-stdio"}, + controlSockPath: "/tmp/test.sock", + expected: []string{"-o", "ControlMaster=auto", "-o", "ControlPath=/tmp/test.sock", "-o", "ControlPersist=10m", "-o", "ConnectTimeout=5", "-T", "root@example.com", "uncloudd", "dial-stdio"}, }, { name: "all options combined", @@ -68,16 +84,28 @@ func TestSSHCLIConnector_buildSSHArgs(t *testing.T) { KeyPath: "/path/to/key", SockPath: "/custom/path/uncloud.sock", }, - expected: []string{"-o", "ConnectTimeout=5", "-p", "2222", "-i", "/path/to/key", "root@example.com", "uncloudd", "dial-stdio", "--socket", "/custom/path/uncloud.sock"}, + controlSockPath: "/tmp/test.sock", + expected: []string{"-o", "ControlMaster=auto", "-o", "ControlPath=/tmp/test.sock", "-o", "ControlPersist=10m", "-o", "ConnectTimeout=5", "-T", "-p", "2222", "-i", "/path/to/key", "root@example.com", "uncloudd", "dial-stdio", "--socket", "/custom/path/uncloud.sock"}, }, { - name: "port 22 not included (default)", + name: "port 0 not included", config: SSHConnectorConfig{ User: "root", Host: "example.com", Port: 0, }, - expected: []string{"-o", "ConnectTimeout=5", "root@example.com", "uncloudd", "dial-stdio"}, + controlSockPath: "/tmp/test.sock", + expected: []string{"-o", "ControlMaster=auto", "-o", "ControlPath=/tmp/test.sock", "-o", "ControlPersist=10m", "-o", "ConnectTimeout=5", "-T", "root@example.com", "uncloudd", "dial-stdio"}, + }, + { + name: "port 22 included when explicit", + config: SSHConnectorConfig{ + User: "root", + Host: "example.com", + Port: 22, + }, + controlSockPath: "/tmp/test.sock", + expected: []string{"-o", "ControlMaster=auto", "-o", "ControlPath=/tmp/test.sock", "-o", "ControlPersist=10m", "-o", "ConnectTimeout=5", "-T", "-p", "22", "root@example.com", "uncloudd", "dial-stdio"}, }, } @@ -85,7 +113,7 @@ func TestSSHCLIConnector_buildSSHArgs(t *testing.T) { t.Run(tt.name, func(t *testing.T) { t.Parallel() - c := &SSHCLIConnector{config: tt.config} + c := &SSHCLIConnector{config: tt.config, controlSockPath: tt.controlSockPath} got := c.buildSSHArgs() assert.Equal(t, tt.expected, got) }) @@ -96,61 +124,77 @@ func TestSSHCLIDialer_buildDialArgs(t *testing.T) { t.Parallel() tests := []struct { - name string - config SSHConnectorConfig - address string - expected []string + name string + config SSHConnectorConfig + controlSockPath string + address string + expected []string }{ { - name: "basic connection", + name: "basic connection without control socket", config: SSHConnectorConfig{ User: "root", Host: "example.com", }, - address: "10.210.1.1:5000", - expected: []string{"-o", "ConnectTimeout=5", "-W", "10.210.1.1:5000", "root@example.com"}, + controlSockPath: "", + address: "10.210.1.1:5000", + expected: []string{"-o", "ConnectTimeout=5", "-T", "-W", "10.210.1.1:5000", "root@example.com"}, }, { - name: "custom port", + name: "basic connection with control socket", + config: SSHConnectorConfig{ + User: "root", + Host: "example.com", + }, + controlSockPath: "/tmp/test.sock", + address: "10.210.1.1:5000", + expected: []string{"-o", "ControlMaster=no", "-o", "ControlPath=/tmp/test.sock", "-o", "ConnectTimeout=5", "-T", "-W", "10.210.1.1:5000", "root@example.com"}, + }, + { + name: "custom port with control socket", config: SSHConnectorConfig{ User: "root", Host: "example.com", Port: 2222, }, - address: "10.210.1.1:5000", - expected: []string{"-o", "ConnectTimeout=5", "-p", "2222", "-W", "10.210.1.1:5000", "root@example.com"}, + controlSockPath: "/tmp/test.sock", + address: "10.210.1.1:5000", + expected: []string{"-o", "ControlMaster=no", "-o", "ControlPath=/tmp/test.sock", "-o", "ConnectTimeout=5", "-T", "-p", "2222", "-W", "10.210.1.1:5000", "root@example.com"}, }, { - name: "with identity file", + name: "with identity file and control socket", config: SSHConnectorConfig{ User: "root", Host: "example.com", Port: 22, KeyPath: "/home/user/.ssh/id_rsa", }, - address: "10.210.1.1:5000", - expected: []string{"-o", "ConnectTimeout=5", "-i", "/home/user/.ssh/id_rsa", "-W", "10.210.1.1:5000", "root@example.com"}, + controlSockPath: "/tmp/test.sock", + address: "10.210.1.1:5000", + expected: []string{"-o", "ControlMaster=no", "-o", "ControlPath=/tmp/test.sock", "-o", "ConnectTimeout=5", "-T", "-p", "22", "-i", "/home/user/.ssh/id_rsa", "-W", "10.210.1.1:5000", "root@example.com"}, }, { - name: "custom port with identity file", + name: "custom port with identity file and control socket", config: SSHConnectorConfig{ User: "root", Host: "example.com", Port: 2222, KeyPath: "/home/user/.ssh/id_rsa", }, - address: "10.210.1.1:5000", - expected: []string{"-o", "ConnectTimeout=5", "-p", "2222", "-i", "/home/user/.ssh/id_rsa", "-W", "10.210.1.1:5000", "root@example.com"}, + controlSockPath: "/tmp/test.sock", + address: "10.210.1.1:5000", + expected: []string{"-o", "ControlMaster=no", "-o", "ControlPath=/tmp/test.sock", "-o", "ConnectTimeout=5", "-T", "-p", "2222", "-i", "/home/user/.ssh/id_rsa", "-W", "10.210.1.1:5000", "root@example.com"}, }, { - name: "zero port defaults to 22", + name: "port 0 not included", config: SSHConnectorConfig{ User: "root", Host: "example.com", Port: 0, }, - address: "10.210.1.1:5000", - expected: []string{"-o", "ConnectTimeout=5", "-W", "10.210.1.1:5000", "root@example.com"}, + controlSockPath: "/tmp/test.sock", + address: "10.210.1.1:5000", + expected: []string{"-o", "ControlMaster=no", "-o", "ControlPath=/tmp/test.sock", "-o", "ConnectTimeout=5", "-T", "-W", "10.210.1.1:5000", "root@example.com"}, }, } @@ -158,9 +202,28 @@ func TestSSHCLIDialer_buildDialArgs(t *testing.T) { t.Run(tt.name, func(t *testing.T) { t.Parallel() - d := &sshCLIDialer{config: tt.config} + d := &sshCLIDialer{config: tt.config, controlSockPath: tt.controlSockPath} got := d.buildDialArgs(tt.address) assert.Equal(t, tt.expected, got) }) } } + +func TestControlSocketPath(t *testing.T) { + // Note: Cannot use t.Parallel() because a subtest uses t.Setenv(). + + path1 := controlSocketPath() + path2 := controlSocketPath() + + assert.Equal(t, path1, path2) + assert.True(t, strings.HasSuffix(path1, ".sock")) + assert.Contains(t, path1, "%C") + + t.Run("uses XDG_RUNTIME_DIR when set", func(t *testing.T) { + runDir := "/user/runtime/dir" + t.Setenv("XDG_RUNTIME_DIR", runDir) + + path := controlSocketPath() + assert.True(t, strings.HasPrefix(path, runDir)) + }) +}