From 110f029b5fbeb0a5e96f9188addbfb81d4fa677c Mon Sep 17 00:00:00 2001
From: CrazyMax <crazy-max@users.noreply.github.com>
Date: Thu, 28 Sep 2023 16:40:09 +0200
Subject: [PATCH] driver(docker): allow attaching additional headers to the
 client

Signed-off-by: CrazyMax <crazy-max@users.noreply.github.com>
---
 driver/docker-container/driver.go |  2 +-
 driver/docker/driver.go           |  9 +++++++--
 driver/driver.go                  | 14 +++++++++++++-
 driver/kubernetes/driver.go       |  2 +-
 driver/manager.go                 |  8 ++++----
 driver/remote/driver.go           |  2 +-
 6 files changed, 27 insertions(+), 10 deletions(-)

diff --git a/driver/docker-container/driver.go b/driver/docker-container/driver.go
index d42d233a94f7..75e2f553898b 100644
--- 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
diff --git a/driver/docker/driver.go b/driver/docker/driver.go
index e31175c33b88..748a7f3c28aa 100644
--- a/driver/docker/driver.go
+++ b/driver/docker/driver.go
@@ -51,10 +51,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.Meta)
 		}), client.WithSessionDialer(func(ctx context.Context, proto string, meta map[string][]string) (net.Conn, error) {
 			return d.DockerAPI.DialHijack(ctx, "/session", proto, meta)
 		}),
diff --git a/driver/driver.go b/driver/driver.go
index 8642a543cb8c..137a6fb2b2bc 100644
--- a/driver/driver.go
+++ b/driver/driver.go
@@ -58,7 +58,7 @@ 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)
+	Client(ctx context.Context, opts ...ClientOption) (*client.Client, error)
 	Features(ctx context.Context) map[Feature]bool
 	IsMobyDriver() bool
 	Config() InitConfig
@@ -110,3 +110,15 @@ func historyAPISupported(ctx context.Context, c *client.Client) bool {
 		}
 	}
 }
+
+type ClientOption func(*ClientOptions)
+
+type ClientOptions struct {
+	Meta map[string][]string
+}
+
+func WithMeta(meta map[string][]string) ClientOption {
+	return func(o *ClientOptions) {
+		o.Meta = meta
+	}
+}
diff --git a/driver/kubernetes/driver.go b/driver/kubernetes/driver.go
index 55c7853da0cb..5a29856ddb11 100644
--- 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 {
diff --git a/driver/manager.go b/driver/manager.go
index b4ec318f396c..eddddc8c2d51 100644
--- a/driver/manager.go
+++ b/driver/manager.go
@@ -159,9 +159,9 @@ type DriverHandle struct {
 	hostGatewayIPErr        error
 }
 
-func (d *DriverHandle) Client(ctx context.Context) (*client.Client, error) {
+func (d *DriverHandle) Client(ctx context.Context, copts ...ClientOption) (*client.Client, error) {
 	d.once.Do(func() {
-		d.client, d.err = d.Driver.Client(ctx)
+		d.client, d.err = d.Driver.Client(ctx, copts...)
 	})
 	return d.client, d.err
 }
@@ -173,9 +173,9 @@ func (d *DriverHandle) Features(ctx context.Context) map[Feature]bool {
 	return d.features
 }
 
-func (d *DriverHandle) HistoryAPISupported(ctx context.Context) bool {
+func (d *DriverHandle) HistoryAPISupported(ctx context.Context, copts ...ClientOption) bool {
 	d.historyAPISupportedOnce.Do(func() {
-		if c, err := d.Client(ctx); err == nil {
+		if c, err := d.Client(ctx, copts...); err == nil {
 			d.historyAPISupported = historyAPISupported(ctx, c)
 		}
 	})
diff --git a/driver/remote/driver.go b/driver/remote/driver.go
index c12cf995b58e..8b7eb8add47b 100644
--- a/driver/remote/driver.go
+++ b/driver/remote/driver.go
@@ -61,7 +61,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()