From 1eff9310f8da42250528bdf0aabcb278d082dae5 Mon Sep 17 00:00:00 2001
From: Justin Chadwell <me@jedevc.com>
Date: Tue, 17 May 2022 16:12:07 +0100
Subject: [PATCH] remote: add additional connhelpers to buildx

Signed-off-by: Justin Chadwell <me@jedevc.com>
---
 commands/util.go               |  9 +++------
 driver/remote/factory.go       | 15 +++++++++------
 driver/remote/util/endpoint.go | 26 ++++++++++++++++++++++++++
 3 files changed, 38 insertions(+), 12 deletions(-)
 create mode 100644 driver/remote/util/endpoint.go

diff --git a/commands/util.go b/commands/util.go
index 9d90098e..d37c196d 100644
--- a/commands/util.go
+++ b/commands/util.go
@@ -9,6 +9,7 @@ import (
 	"github.com/docker/buildx/build"
 	"github.com/docker/buildx/driver"
 	ctxkube "github.com/docker/buildx/driver/kubernetes/context"
+	remoteutil "github.com/docker/buildx/driver/remote/util"
 	"github.com/docker/buildx/store"
 	"github.com/docker/buildx/store/storeutil"
 	"github.com/docker/buildx/util/platformutil"
@@ -45,12 +46,8 @@ func validateEndpoint(dockerCli command.Cli, ep string) (string, error) {
 
 // validateBuildkitEndpoint validates that endpoint is a valid buildkit host
 func validateBuildkitEndpoint(ep string) (string, error) {
-	endpoint, err := url.Parse(ep)
-	if err != nil {
-		return "", errors.Wrapf(err, "failed to parse endpoint %s", ep)
-	}
-	if endpoint.Scheme != "tcp" && endpoint.Scheme != "unix" {
-		return "", errors.Errorf("unrecognized url scheme %s", endpoint.Scheme)
+	if err := remoteutil.IsValidEndpoint(ep); err != nil {
+		return "", err
 	}
 	return ep, nil
 }
diff --git a/driver/remote/factory.go b/driver/remote/factory.go
index cf256618..83ce16e6 100644
--- a/driver/remote/factory.go
+++ b/driver/remote/factory.go
@@ -4,10 +4,15 @@ import (
 	"context"
 	"net/url"
 	"path/filepath"
-	"regexp"
 	"strings"
 
+	// import connhelpers for special url schemes
+	_ "github.com/moby/buildkit/client/connhelper/dockercontainer"
+	_ "github.com/moby/buildkit/client/connhelper/kubepod"
+	_ "github.com/moby/buildkit/client/connhelper/ssh"
+
 	"github.com/docker/buildx/driver"
+	util "github.com/docker/buildx/driver/remote/util"
 	dockerclient "github.com/docker/docker/client"
 	"github.com/pkg/errors"
 )
@@ -15,8 +20,6 @@ import (
 const prioritySupported = 20
 const priorityUnsupported = 90
 
-var schemeRegexp = regexp.MustCompile("^(tcp|unix)://")
-
 func init() {
 	driver.Register(&factory{})
 }
@@ -33,10 +36,10 @@ func (*factory) Usage() string {
 }
 
 func (*factory) Priority(ctx context.Context, endpoint string, api dockerclient.APIClient) int {
-	if schemeRegexp.MatchString(endpoint) {
-		return prioritySupported
+	if util.IsValidEndpoint(endpoint) != nil {
+		return priorityUnsupported
 	}
-	return priorityUnsupported
+	return prioritySupported
 }
 
 func (f *factory) New(ctx context.Context, cfg driver.InitConfig) (driver.Driver, error) {
diff --git a/driver/remote/util/endpoint.go b/driver/remote/util/endpoint.go
new file mode 100644
index 00000000..3b535f52
--- /dev/null
+++ b/driver/remote/util/endpoint.go
@@ -0,0 +1,26 @@
+package remote
+
+import (
+	"net/url"
+
+	"github.com/pkg/errors"
+)
+
+var schemes = map[string]struct{}{
+	"tcp":              {},
+	"unix":             {},
+	"ssh":              {},
+	"docker-container": {},
+	"kube-pod":         {},
+}
+
+func IsValidEndpoint(ep string) error {
+	endpoint, err := url.Parse(ep)
+	if err != nil {
+		return errors.Wrapf(err, "failed to parse endpoint %s", ep)
+	}
+	if _, ok := schemes[endpoint.Scheme]; !ok {
+		return errors.Errorf("unrecognized url scheme %s", endpoint.Scheme)
+	}
+	return nil
+}