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

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
+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))
})
}