diff --git a/pkg/client/connector/sshcli.go b/pkg/client/connector/sshcli.go index af591c88..b641b516 100644 --- a/pkg/client/connector/sshcli.go +++ b/pkg/client/connector/sshcli.go @@ -5,11 +5,12 @@ import ( "fmt" "net" "os" + "os/exec" "path/filepath" "strconv" + "strings" "github.com/docker/cli/cli/connhelper/commandconn" - "github.com/psviderski/uncloud/internal/machine" "golang.org/x/net/proxy" "google.golang.org/grpc" "google.golang.org/grpc/credentials/insecure" @@ -61,17 +62,30 @@ func controlSocketPath() string { } func (c *SSHCLIConnector) Connect(ctx context.Context) (*grpc.ClientConn, error) { - // Create gRPC client with a dialer that spawns a new SSH connection on demand. - // Each dial attempt runs `ssh ... uncloudd dial-stdio`, reusing the control socket if available. + // Validate SSH connectivity by running a no-op command on the remote machine. This also + // establishes the control socket (ControlMaster=auto) so subsequent connections reuse it. + probeArgs := append(c.buildSSHArgs(), "true") + probe := exec.CommandContext(ctx, "ssh", probeArgs...) + if output, err := probe.CombinedOutput(); err != nil { + return nil, fmt.Errorf("SSH connection to '%s': %w: %s", + c.config.Destination(), err, strings.TrimSpace(string(output))) + } + + // Create gRPC client with a dialer that spawns new SSH connections on demand, + // reusing the control socket established above. grpcConn, err := grpc.NewClient( "passthrough:///", // Dummy target since we're using a custom dialer. grpc.WithTransportCredentials(insecure.NewCredentials()), grpc.WithDefaultServiceConfig(defaultServiceConfig), grpc.WithContextDialer(func(ctx context.Context, _ string) (net.Conn, error) { - args := c.buildSSHArgs() - conn, err := commandconn.New(ctx, "ssh", args...) + dialArgs := append(c.buildSSHArgs(), "uncloudd", "dial-stdio") + if c.config.SockPath != "" { + dialArgs = append(dialArgs, "--socket", c.config.SockPath) + } + + conn, err := commandconn.New(ctx, "ssh", dialArgs...) if err != nil { - return nil, fmt.Errorf("SSH connection to %s: %w", c.config.Destination(), err) + return nil, fmt.Errorf("SSH connection to '%s': %w", c.config.Destination(), err) } return conn, nil }), @@ -83,8 +97,9 @@ func (c *SSHCLIConnector) Connect(ctx context.Context) (*grpc.ClientConn, error) return grpcConn, nil } -// buildSSHArgs constructs the SSH command arguments to run `uncloudd dial-stdio` on the remote machine reusing -// the established connection via control socket. +// buildSSHArgs constructs the SSH command arguments with connection options and destination. The options +// include control socket settings for connection reuse if necessary. +// The remote command is not included and should be appended by the caller. func (c *SSHCLIConnector) buildSSHArgs() []string { var args []string @@ -120,14 +135,6 @@ func (c *SSHCLIConnector) buildSSHArgs() []string { // Add [user@]host destination. args = append(args, c.config.Destination()) - // Add remote command: uncloudd dial-stdio - args = append(args, "uncloudd", "dial-stdio") - - // Add socket path if non-default. - if c.config.SockPath != "" && c.config.SockPath != machine.DefaultUncloudSockPath { - args = append(args, "--socket", c.config.SockPath) - } - return args } @@ -137,10 +144,22 @@ func (c *SSHCLIConnector) Dialer() (proxy.ContextDialer, error) { return nil, fmt.Errorf("SSH connector not configured") } - return &sshCLIDialer{ - config: c.config, - controlSockPath: c.controlSockPath, - }, nil + return c, nil +} + +// DialContext establishes a connection to the target address through an SSH tunnel using -W flag. +func (c *SSHCLIConnector) DialContext(ctx context.Context, network, address string) (net.Conn, error) { + if network != "tcp" { + return nil, fmt.Errorf("unsupported network type: %s", network) + } + + args := append(c.buildSSHArgs(), "-W", address) + conn, err := commandconn.New(ctx, "ssh", args...) + if err != nil { + return nil, fmt.Errorf("SSH connection to '%s' for dialing '%s': %w", c.config.Destination(), address, err) + } + + return conn, nil } func (c *SSHCLIConnector) Close() error { @@ -148,64 +167,3 @@ func (c *SSHCLIConnector) Close() error { // The SSH control socket may persist for connection reuse across CLI invocations. return nil } - -// 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 { - 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 specified. - if d.config.Port != 0 { - args = append(args, "-p", strconv.Itoa(d.config.Port)) - } - - // Add identity file if specified. - if d.config.KeyPath != "" { - args = append(args, "-i", d.config.KeyPath) - } - - // Add -W flag for stdin/stdout forwarding to target address. - args = append(args, "-W", address) - - // Add [user@]host destination. - args = append(args, d.config.Destination()) - - return args -} - -// DialContext establishes a connection to the target address through an SSH tunnel using -W flag. -func (d *sshCLIDialer) DialContext(ctx context.Context, network, address string) (net.Conn, error) { - // Only support TCP connections. - if network != "tcp" { - return nil, fmt.Errorf("unsupported network type: %s", network) - } - - // Build SSH command arguments. - args := d.buildDialArgs(address) - - // Create connection using docker's commandconn. - conn, err := commandconn.New(ctx, "ssh", args...) - if err != nil { - return nil, fmt.Errorf("SSH connection to %s for dialing %s: %w", d.config.Destination(), address, err) - } - - return conn, nil -} diff --git a/pkg/client/connector/sshcli_test.go b/pkg/client/connector/sshcli_test.go index d913c77d..86e0b463 100644 --- a/pkg/client/connector/sshcli_test.go +++ b/pkg/client/connector/sshcli_test.go @@ -4,7 +4,6 @@ import ( "strings" "testing" - "github.com/psviderski/uncloud/internal/machine" "github.com/stretchr/testify/assert" ) @@ -24,7 +23,7 @@ func TestSSHCLIConnector_buildSSHArgs(t *testing.T) { Host: "example.com", }, 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"}, + expected: []string{"-o", "ControlMaster=auto", "-o", "ControlPath=/tmp/test.sock", "-o", "ControlPersist=10m", "-o", "ConnectTimeout=5", "-T", "root@example.com"}, }, { name: "basic connection without control socket", @@ -33,7 +32,7 @@ func TestSSHCLIConnector_buildSSHArgs(t *testing.T) { Host: "example.com", }, controlSockPath: "", - expected: []string{"-o", "ConnectTimeout=5", "-T", "root@example.com", "uncloudd", "dial-stdio"}, + expected: []string{"-o", "ConnectTimeout=5", "-T", "root@example.com"}, }, { name: "with custom port", @@ -43,7 +42,7 @@ func TestSSHCLIConnector_buildSSHArgs(t *testing.T) { Port: 2222, }, 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"}, + expected: []string{"-o", "ControlMaster=auto", "-o", "ControlPath=/tmp/test.sock", "-o", "ControlPersist=10m", "-o", "ConnectTimeout=5", "-T", "-p", "2222", "root@example.com"}, }, { name: "with identity file", @@ -53,27 +52,7 @@ func TestSSHCLIConnector_buildSSHArgs(t *testing.T) { KeyPath: "/path/to/key", }, 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", - config: SSHConnectorConfig{ - User: "root", - Host: "example.com", - SockPath: "/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)", - config: SSHConnectorConfig{ - User: "root", - Host: "example.com", - SockPath: machine.DefaultUncloudSockPath, - }, - 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"}, + expected: []string{"-o", "ControlMaster=auto", "-o", "ControlPath=/tmp/test.sock", "-o", "ControlPersist=10m", "-o", "ConnectTimeout=5", "-T", "-i", "/path/to/key", "root@example.com"}, }, { name: "all options combined", @@ -85,7 +64,7 @@ func TestSSHCLIConnector_buildSSHArgs(t *testing.T) { SockPath: "/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"}, + 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"}, }, { name: "port 0 not included", @@ -95,7 +74,7 @@ func TestSSHCLIConnector_buildSSHArgs(t *testing.T) { Port: 0, }, 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"}, + expected: []string{"-o", "ControlMaster=auto", "-o", "ControlPath=/tmp/test.sock", "-o", "ControlPersist=10m", "-o", "ConnectTimeout=5", "-T", "root@example.com"}, }, { name: "port 22 included when explicit", @@ -105,7 +84,7 @@ func TestSSHCLIConnector_buildSSHArgs(t *testing.T) { 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"}, + expected: []string{"-o", "ControlMaster=auto", "-o", "ControlPath=/tmp/test.sock", "-o", "ControlPersist=10m", "-o", "ConnectTimeout=5", "-T", "-p", "22", "root@example.com"}, }, } @@ -120,94 +99,6 @@ func TestSSHCLIConnector_buildSSHArgs(t *testing.T) { } } -func TestSSHCLIDialer_buildDialArgs(t *testing.T) { - t.Parallel() - - tests := []struct { - name string - config SSHConnectorConfig - controlSockPath string - address string - expected []string - }{ - { - name: "basic connection without control socket", - config: SSHConnectorConfig{ - User: "root", - Host: "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: "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, - }, - 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 and control socket", - config: SSHConnectorConfig{ - User: "root", - Host: "example.com", - Port: 22, - KeyPath: "/home/user/.ssh/id_rsa", - }, - 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 and control socket", - config: SSHConnectorConfig{ - User: "root", - Host: "example.com", - Port: 2222, - KeyPath: "/home/user/.ssh/id_rsa", - }, - 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: "port 0 not included", - config: SSHConnectorConfig{ - User: "root", - Host: "example.com", - Port: 0, - }, - 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"}, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - t.Parallel() - - 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().