From b16f5fb610586805e89be17ed9c6fd453287a748 Mon Sep 17 00:00:00 2001
From: Erez Rokah <erezrokah@users.noreply.github.com>
Date: Fri, 29 Dec 2023 11:48:10 +0100
Subject: [PATCH] feat: Add Support for pulling images from CloudQuery Docker
 registry (#190)

---
 go.mod                  |  2 +-
 managedplugin/docker.go | 34 ++++++++++++++++++++++++++++++----
 managedplugin/plugin.go |  2 +-
 3 files changed, 32 insertions(+), 6 deletions(-)

diff --git a/go.mod b/go.mod
index 41b7369..8b3b040 100644
--- a/go.mod
+++ b/go.mod
@@ -6,6 +6,7 @@ require (
 	github.com/apache/arrow/go/v15 v15.0.0-20231223155039-ec41209ea02b
 	github.com/avast/retry-go/v4 v4.5.1
 	github.com/cloudquery/cloudquery-api-go v1.6.3
+	github.com/distribution/reference v0.5.0
 	github.com/docker/docker v24.0.7+incompatible
 	github.com/docker/go-connections v0.4.0
 	github.com/ghodss/yaml v1.0.0
@@ -37,7 +38,6 @@ require (
 	github.com/chenzhuoyu/iasm v0.9.1 // indirect
 	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/deepmap/oapi-codegen v1.15.0 // indirect
-	github.com/distribution/reference v0.5.0 // indirect
 	github.com/docker/distribution v2.8.3+incompatible // indirect
 	github.com/docker/go-units v0.5.0 // indirect
 	github.com/fatih/structs v1.1.0 // indirect
diff --git a/managedplugin/docker.go b/managedplugin/docker.go
index 6b68652..58ab32d 100644
--- a/managedplugin/docker.go
+++ b/managedplugin/docker.go
@@ -5,9 +5,12 @@ import (
 	"encoding/json"
 	"fmt"
 	"io"
+	"strings"
 
+	"github.com/distribution/reference"
 	"github.com/docker/docker/api/types"
 	"github.com/docker/docker/api/types/filters"
+	"github.com/docker/docker/api/types/registry"
 	"github.com/docker/docker/client"
 	"github.com/schollz/progressbar/v3"
 )
@@ -81,14 +84,37 @@ func isDockerImageAvailable(ctx context.Context, imageName string) (bool, error)
 	return len(images) > 0, nil
 }
 
-func pullDockerImage(ctx context.Context, imageName string) error {
-	cli, err := client.NewClientWithOpts(client.FromEnv)
+func pullDockerImage(ctx context.Context, imageName string, authToken string) error {
+	// Pull the image
+	additionalHeaders := make(map[string]string)
+	opts := types.ImagePullOptions{}
+	if authToken != "" && strings.HasPrefix(imageName, "registry.cloudquery.io") {
+		namedRef, err := reference.ParseNormalizedNamed(imageName)
+		if err != nil {
+			return fmt.Errorf("failed to parse Docker image tag: %v", err)
+		}
+		nameWithTag, ok := namedRef.(reference.NamedTagged)
+		if !ok {
+			return fmt.Errorf("failed to parse Docker image tag: %v", err)
+		}
+		additionalHeaders["X-Meta-Plugin-Version"] = nameWithTag.Tag()
+		authConfig := registry.AuthConfig{
+			Username: "managedplugin",
+			Password: authToken,
+		}
+		encodedAuth, err := registry.EncodeAuthConfig(authConfig)
+		if err != nil {
+			return fmt.Errorf("failed to encode Docker auth config: %v", err)
+		}
+		opts.RegistryAuth = encodedAuth
+	}
+
+	cli, err := client.NewClientWithOpts(client.FromEnv, client.WithHTTPHeaders(additionalHeaders))
 	if err != nil {
 		return fmt.Errorf("failed to create Docker client: %v", err)
 	}
 
-	// Pull the image
-	out, err := cli.ImagePull(ctx, imageName, types.ImagePullOptions{})
+	out, err := cli.ImagePull(ctx, imageName, opts)
 	if err != nil {
 		return fmt.Errorf("failed to pull Docker image: %v", err)
 	}
diff --git a/managedplugin/plugin.go b/managedplugin/plugin.go
index 84a3ab1..22427b5 100644
--- a/managedplugin/plugin.go
+++ b/managedplugin/plugin.go
@@ -173,7 +173,7 @@ func (c *Client) downloadPlugin(ctx context.Context, typ PluginType) error {
 		if imageAvailable, err := isDockerImageAvailable(ctx, c.config.Path); err != nil {
 			return err
 		} else if !imageAvailable {
-			return pullDockerImage(ctx, c.config.Path)
+			return pullDockerImage(ctx, c.config.Path, c.authToken)
 		}
 		return nil
 	case RegistryCloudQuery: