From 5d33eedfdb2d574e09b0ade8a556c930a16e59f7 Mon Sep 17 00:00:00 2001
From: Pierre Fenoll <pierrefenoll@gmail.com>
Date: Tue, 3 May 2022 21:53:35 +0200
Subject: [PATCH] connhelper/ssh: allow passing socket path

Signed-off-by: Pierre Fenoll <pierrefenoll@gmail.com>
---
 client/connhelper/ssh/ssh.go      | 30 +++++++++++++++++-------------
 client/connhelper/ssh/ssh_test.go | 14 ++++++++------
 2 files changed, 25 insertions(+), 19 deletions(-)

diff --git a/client/connhelper/ssh/ssh.go b/client/connhelper/ssh/ssh.go
index 6bd4cd8d2..e3666b572 100644
--- a/client/connhelper/ssh/ssh.go
+++ b/client/connhelper/ssh/ssh.go
@@ -23,25 +23,31 @@ func Helper(u *url.URL) (*connhelper.ConnectionHelper, error) {
 	}
 	return &connhelper.ConnectionHelper{
 		ContextDialer: func(ctx context.Context, addr string) (net.Conn, error) {
-			ctxFlags := []string{}
+			args := []string{}
 			if sp.User != "" {
-				ctxFlags = append(ctxFlags, "-l", sp.User)
+				args = append(args, "-l", sp.User)
 			}
 			if sp.Port != "" {
-				ctxFlags = append(ctxFlags, "-p", sp.Port)
+				args = append(args, "-p", sp.Port)
 			}
-			ctxFlags = append(ctxFlags, "--", sp.Host)
+			args = append(args, "--", sp.Host)
+			args = append(args, "buildctl")
+			if socket := sp.Socket; socket != "" {
+				args = append(args, "--addr", "unix://"+socket)
+			}
+			args = append(args, "dial-stdio")
 			// using background context because context remains active for the duration of the process, after dial has completed
-			return commandconn.New(context.Background(), "ssh", append(ctxFlags, []string{"buildctl", "dial-stdio"}...)...)
+			return commandconn.New(context.Background(), "ssh", args...)
 		},
 	}, nil
 }
 
 // Spec
 type Spec struct {
-	User string
-	Host string
-	Port string
+	User   string
+	Host   string
+	Port   string
+	Socket string
 }
 
 // SpecFromURL creates Spec from URL.
@@ -49,8 +55,9 @@ type Spec struct {
 // Only <host> part is mandatory.
 func SpecFromURL(u *url.URL) (*Spec, error) {
 	sp := Spec{
-		Host: u.Hostname(),
-		Port: u.Port(),
+		Host:   u.Hostname(),
+		Port:   u.Port(),
+		Socket: u.Path,
 	}
 	if user := u.User; user != nil {
 		sp.User = user.Username()
@@ -61,9 +68,6 @@ func SpecFromURL(u *url.URL) (*Spec, error) {
 	if sp.Host == "" {
 		return nil, errors.Errorf("no host specified")
 	}
-	if u.Path != "" {
-		return nil, errors.Errorf("extra path after the host: %q", u.Path)
-	}
 	if u.RawQuery != "" {
 		return nil, errors.Errorf("extra query after the host: %q", u.RawQuery)
 	}
diff --git a/client/connhelper/ssh/ssh_test.go b/client/connhelper/ssh/ssh_test.go
index 52b7f6f9d..e4afb31cf 100644
--- a/client/connhelper/ssh/ssh_test.go
+++ b/client/connhelper/ssh/ssh_test.go
@@ -12,14 +12,16 @@ func TestSpecFromURL(t *testing.T) {
 		"ssh://foo": {
 			Host: "foo",
 		},
-		"ssh://me@foo:10022": {
-			User: "me", Host: "foo", Port: "10022",
+		"ssh://me@foo:10022/s/o/c/k/e/t.sock": {
+			User: "me", Host: "foo", Port: "10022", Socket: "/s/o/c/k/e/t.sock",
 		},
 		"ssh://me:passw0rd@foo": nil,
-		"ssh://foo/bar":         nil,
-		"ssh://foo?bar":         nil,
-		"ssh://foo#bar":         nil,
-		"ssh://":                nil,
+		"ssh://foo/bar": {
+			Host: "foo", Socket: "/bar",
+		},
+		"ssh://foo?bar": nil,
+		"ssh://foo#bar": nil,
+		"ssh://":        nil,
 	}
 	for s, expected := range cases {
 		u, err := url.Parse(s)