driver(docker): opt to set additional headers to the client

Signed-off-by: CrazyMax <[email protected]>
docker · Oct 11, 2023 · 3a80666 · 3a80666
1 parent c5398fa
commit 3a80666
Show file tree

Hide file tree

Showing 5 changed files with 35 additions and 18 deletions.
diff --git a/driver/docker-container/driver.go b/driver/docker-container/driver.go
@@ -359,7 +359,7 @@ func (d *Driver) Rm(ctx context.Context, force, rmVolume, rmDaemon bool) error {
 	return nil
 }
 
-func (d *Driver) Client(ctx context.Context) (*client.Client, error) {
+func (d *Driver) Client(ctx context.Context, copts ...driver.ClientOption) (*client.Client, error) {
 	_, conn, err := d.exec(ctx, []string{"buildctl", "dial-stdio"})
 	if err != nil {
 		return nil, err
@@ -390,7 +390,7 @@ func (d *Driver) Factory() driver.Factory {
 	return d.factory
 }
 
-func (d *Driver) Features(ctx context.Context) map[driver.Feature]bool {
+func (d *Driver) Features(ctx context.Context, copts ...driver.ClientOption) map[driver.Feature]bool {
 	return map[driver.Feature]bool{
 		driver.OCIExporter:    true,
 		driver.DockerExporter: true,
@@ -399,7 +399,7 @@ func (d *Driver) Features(ctx context.Context) map[driver.Feature]bool {
 	}
 }
 
-func (d *Driver) HostGatewayIP(ctx context.Context) (net.IP, error) {
+func (d *Driver) HostGatewayIP(ctx context.Context, copts ...driver.ClientOption) (net.IP, error) {
 	return nil, errors.New("host-gateway is not supported by the docker-container driver")
 }
 

diff --git a/driver/docker/driver.go b/driver/docker/driver.go
@@ -54,10 +54,15 @@ func (d *Driver) Rm(ctx context.Context, force, rmVolume, rmDaemon bool) error {
 	return nil
 }
 
-func (d *Driver) Client(ctx context.Context) (*client.Client, error) {
+func (d *Driver) Client(ctx context.Context, copts ...driver.ClientOption) (*client.Client, error) {
+	co := driver.ClientOptions{}
+	for _, opt := range copts {
+		opt(&co)
+	}
+
 	opts := []client.ClientOpt{
 		client.WithContextDialer(func(context.Context, string) (net.Conn, error) {
-			return d.DockerAPI.DialHijack(ctx, "/grpc", "h2c", nil)
+			return d.DockerAPI.DialHijack(ctx, "/grpc", "h2c", co.DialMeta)
 		}), client.WithSessionDialer(func(ctx context.Context, proto string, meta map[string][]string) (net.Conn, error) {
 			return d.DockerAPI.DialHijack(ctx, "/session", proto, meta)
 		}),
@@ -78,10 +83,10 @@ type features struct {
 	list map[driver.Feature]bool
 }
 
-func (d *Driver) Features(ctx context.Context) map[driver.Feature]bool {
+func (d *Driver) Features(ctx context.Context, copts ...driver.ClientOption) map[driver.Feature]bool {
 	d.features.once.Do(func() {
 		var useContainerdSnapshotter bool
-		if c, err := d.Client(ctx); err == nil {
+		if c, err := d.Client(ctx, copts...); err == nil {
 			workers, _ := c.ListWorkers(ctx)
 			for _, w := range workers {
 				if _, ok := w.Labels["org.mobyproject.buildkit.worker.snapshotter"]; ok {
@@ -106,9 +111,9 @@ type hostGateway struct {
 	err  error
 }
 
-func (d *Driver) HostGatewayIP(ctx context.Context) (net.IP, error) {
+func (d *Driver) HostGatewayIP(ctx context.Context, copts ...driver.ClientOption) (net.IP, error) {
 	d.hostGateway.once.Do(func() {
-		c, err := d.Client(ctx)
+		c, err := d.Client(ctx, copts...)
 		if err != nil {
 			d.hostGateway.err = err
 			return

diff --git a/driver/driver.go b/driver/driver.go
@@ -57,9 +57,9 @@ type Driver interface {
 	Version(context.Context) (string, error)
 	Stop(ctx context.Context, force bool) error
 	Rm(ctx context.Context, force, rmVolume, rmDaemon bool) error
-	Client(ctx context.Context) (*client.Client, error)
-	Features(ctx context.Context) map[Feature]bool
-	HostGatewayIP(ctx context.Context) (net.IP, error)
+	Client(ctx context.Context, copts ...ClientOption) (*client.Client, error)
+	Features(ctx context.Context, copts ...ClientOption) map[Feature]bool
+	HostGatewayIP(ctx context.Context, copts ...ClientOption) (net.IP, error)
 	IsMobyDriver() bool
 	Config() InitConfig
 }
@@ -91,3 +91,15 @@ func Boot(ctx, clientContext context.Context, d *DriverHandle, pw progress.Write
 		return c, nil
 	}
 }
+
+type ClientOption func(*ClientOptions)
+
+type ClientOptions struct {
+	DialMeta map[string][]string
+}
+
+func WithDialMeta(dialMeta map[string][]string) ClientOption {
+	return func(o *ClientOptions) {
+		o.DialMeta = dialMeta
+	}
+}
diff --git a/driver/kubernetes/driver.go b/driver/kubernetes/driver.go
@@ -189,7 +189,7 @@ func (d *Driver) Rm(ctx context.Context, force, rmVolume, rmDaemon bool) error {
 	return nil
 }
 
-func (d *Driver) Client(ctx context.Context) (*client.Client, error) {
+func (d *Driver) Client(ctx context.Context, copts ...driver.ClientOption) (*client.Client, error) {
 	restClient := d.clientset.CoreV1().RESTClient()
 	restClientConfig, err := d.KubeClientConfig.ClientConfig()
 	if err != nil {
@@ -228,7 +228,7 @@ func (d *Driver) Factory() driver.Factory {
 	return d.factory
 }
 
-func (d *Driver) Features(ctx context.Context) map[driver.Feature]bool {
+func (d *Driver) Features(ctx context.Context, copts ...driver.ClientOption) map[driver.Feature]bool {
 	return map[driver.Feature]bool{
 		driver.OCIExporter:    true,
 		driver.DockerExporter: d.DockerAPI != nil,
@@ -237,6 +237,6 @@ func (d *Driver) Features(ctx context.Context) map[driver.Feature]bool {
 	}
 }
 
-func (d *Driver) HostGatewayIP(ctx context.Context) (net.IP, error) {
+func (d *Driver) HostGatewayIP(ctx context.Context, copts ...driver.ClientOption) (net.IP, error) {
 	return nil, errors.New("host-gateway is not supported by the kubernetes driver")
 }
diff --git a/driver/remote/driver.go b/driver/remote/driver.go
@@ -63,7 +63,7 @@ func (d *Driver) Rm(ctx context.Context, force, rmVolume, rmDaemon bool) error {
 	return nil
 }
 
-func (d *Driver) Client(ctx context.Context) (*client.Client, error) {
+func (d *Driver) Client(ctx context.Context, copts ...driver.ClientOption) (*client.Client, error) {
 	opts := []client.ClientOpt{}
 
 	exp, err := detect.Exporter()
@@ -84,7 +84,7 @@ func (d *Driver) Client(ctx context.Context) (*client.Client, error) {
 	return client.New(ctx, d.InitConfig.EndpointAddr, opts...)
 }
 
-func (d *Driver) Features(ctx context.Context) map[driver.Feature]bool {
+func (d *Driver) Features(ctx context.Context, copts ...driver.ClientOption) map[driver.Feature]bool {
 	return map[driver.Feature]bool{
 		driver.OCIExporter:    true,
 		driver.DockerExporter: true,
@@ -93,7 +93,7 @@ func (d *Driver) Features(ctx context.Context) map[driver.Feature]bool {
 	}
 }
 
-func (d *Driver) HostGatewayIP(ctx context.Context) (net.IP, error) {
+func (d *Driver) HostGatewayIP(ctx context.Context, copts ...driver.ClientOption) (net.IP, error) {
 	return nil, errors.New("host-gateway is not supported by the remote driver")
 }