philipp/uncloud
1
0
Fork 0
mirror of https://github.com/psviderski/uncloud.git synced 2026-08-26 11:03:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

chore: make ssh+cli connections reuse one SSH connection via control socket. Fix image push

Browse Source
This commit is contained in:
Pasha Sviderski
2026-01-28 16:49:08 +10:00
parent e0a63a3f49
commit c91a964513
5 changed files with 225 additions and 70 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/cli/config/connection.go
+13 -12
View File
@@ -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
}
internal/sshexec/ssh.go
+11
View File
@@ -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 {
pkg/client/connector/ssh.go
+15 -1
View File
@@ -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)
}
}
pkg/client/connector/sshcli.go
+90 -24
View File
@@ -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
}
pkg/client/connector/sshcli_test.go
+96 -33
View File
@@ -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))
})
}