Add support for private OCI registries (#990)

Signed-off-by: Radoslav Dimitrov <[email protected]>

Author: rdimitrov

cmd/thv/app/run.go

@@ -9,6 +9,7 @@ import (
 
 	"github.com/stacklok/toolhive/pkg/config"
 	"github.com/stacklok/toolhive/pkg/container"
+	"github.com/stacklok/toolhive/pkg/container/runtime"
 	"github.com/stacklok/toolhive/pkg/logger"
 	"github.com/stacklok/toolhive/pkg/permissions"
 	"github.com/stacklok/toolhive/pkg/process"
@@ -285,7 +286,7 @@ func runCmdFunc(cmd *cobra.Command, args []string) error {
 	// If we are running in detached mode, or the CLI is wrapped by the K8s operator,
 	// we use the DetachedEnvVarValidator.
 	var envVarValidator runner.EnvVarValidator
-	if process.IsDetached() || container.IsKubernetesRuntime() {
+	if process.IsDetached() || runtime.IsKubernetesRuntime() {
 		envVarValidator = &runner.DetachedEnvVarValidator{}
 	} else {
 		envVarValidator = &runner.CLIEnvVarValidator{}
@@ -297,7 +298,7 @@ func runCmdFunc(cmd *cobra.Command, args []string) error {
 	// Only pull image if we are not running in Kubernetes mode.
 	// This split will go away if we implement a separate command or binary
 	// for running MCP servers in Kubernetes.
-	if !container.IsKubernetesRuntime() {
+	if !runtime.IsKubernetesRuntime() {
 		// Take the MCP server we were supplied and either fetch the image, or
 		// build it from a protocol scheme. If the server URI refers to an image
 		// in our trusted registry, we will also fetch the image metadata.

cmd/thv/main.go

@@ -6,7 +6,7 @@ import (
 
 	"github.com/stacklok/toolhive/cmd/thv/app"
 	"github.com/stacklok/toolhive/pkg/client"
-	"github.com/stacklok/toolhive/pkg/container"
+	"github.com/stacklok/toolhive/pkg/container/runtime"
 	"github.com/stacklok/toolhive/pkg/logger"
 )
 
@@ -19,7 +19,7 @@ func main() {
 	client.CheckAndPerformAutoDiscoveryMigration()
 
 	// Skip update check for completion command or if we are running in kubernetes
-	if err := app.NewRootCmd(!app.IsCompletionCommand(os.Args) && !container.IsKubernetesRuntime()).Execute(); err != nil {
+	if err := app.NewRootCmd(!app.IsCompletionCommand(os.Args) && !runtime.IsKubernetesRuntime()).Execute(); err != nil {
 		os.Exit(1)
 	}
 }

pkg/container/docker/client.go

@@ -56,7 +56,7 @@ func NewClient(ctx context.Context) (*Client, error) {
 		return nil, err // there is already enough context in the error.
 	}
 
-	imageManager := images.NewDockerImageManager(dockerClient)
+	imageManager := images.NewRegistryImageManager(dockerClient)
 
 	c := &Client{
 		runtimeType:  runtimeType,

pkg/container/factory.go

@@ -4,7 +4,6 @@ package container
 
 import (
 	"context"
-	"os"
 
 	"github.com/stacklok/toolhive/pkg/container/docker"
 	"github.com/stacklok/toolhive/pkg/container/kubernetes"
@@ -21,7 +20,7 @@ func NewFactory() *Factory {
 
 // Create creates a container runtime
 func (*Factory) Create(ctx context.Context) (runtime.Runtime, error) {
-	if !IsKubernetesRuntime() {
+	if !runtime.IsKubernetesRuntime() {
 		client, err := docker.NewClient(ctx)
 		if err != nil {
 			return nil, err
@@ -41,9 +40,3 @@ func (*Factory) Create(ctx context.Context) (runtime.Runtime, error) {
 func NewMonitor(rt runtime.Runtime, containerID, containerName string) runtime.Monitor {
 	return docker.NewMonitor(rt, containerID, containerName)
 }
-
-// IsKubernetesRuntime returns true if the runtime is Kubernetes
-// isn't the best way to do this, but for now it's good enough
-func IsKubernetesRuntime() bool {
-	return os.Getenv("KUBERNETES_SERVICE_HOST") != ""
-}

pkg/container/images/docker.go

@@ -49,6 +49,30 @@ func (d *DockerImageManager) ImageExists(ctx context.Context, imageName string)
 
 // BuildImage builds a Docker image from a Dockerfile in the specified context directory
 func (d *DockerImageManager) BuildImage(ctx context.Context, contextDir, imageName string) error {
+	return buildDockerImage(ctx, d.client, contextDir, imageName)
+}
+
+// PullImage pulls an image from a registry
+func (d *DockerImageManager) PullImage(ctx context.Context, imageName string) error {
+	logger.Infof("Pulling image: %s", imageName)
+
+	// Pull the image
+	reader, err := d.client.ImagePull(ctx, imageName, dockerimage.PullOptions{})
+	if err != nil {
+		return fmt.Errorf("failed to pull image: %v", err)
+	}
+	defer reader.Close()
+
+	// Parse and filter the pull output
+	if err := parsePullOutput(reader, os.Stdout); err != nil {
+		return fmt.Errorf("failed to process pull output: %v", err)
+	}
+
+	return nil
+}
+
+// buildDockerImage builds a Docker image using the Docker client API
+func buildDockerImage(ctx context.Context, dockerClient *client.Client, contextDir, imageName string) error {
 	logger.Infof("Building image %s from context directory %s", imageName, contextDir)
 
 	// Create a tar archive of the context directory
@@ -76,7 +100,7 @@ func (d *DockerImageManager) BuildImage(ctx context.Context, contextDir, imageNa
 		Remove:     true,
 	}
 
-	response, err := d.client.ImageBuild(ctx, tarFile, buildOptions)
+	response, err := dockerClient.ImageBuild(ctx, tarFile, buildOptions)
 	if err != nil {
 		return fmt.Errorf("failed to build image: %v", err)
 	}
@@ -90,25 +114,6 @@ func (d *DockerImageManager) BuildImage(ctx context.Context, contextDir, imageNa
 	return nil
 }
 
-// PullImage pulls an image from a registry
-func (d *DockerImageManager) PullImage(ctx context.Context, imageName string) error {
-	logger.Infof("Pulling image: %s", imageName)
-
-	// Pull the image
-	reader, err := d.client.ImagePull(ctx, imageName, dockerimage.PullOptions{})
-	if err != nil {
-		return fmt.Errorf("failed to pull image: %v", err)
-	}
-	defer reader.Close()
-
-	// Parse and filter the pull output
-	if err := parsePullOutput(reader, os.Stdout); err != nil {
-		return fmt.Errorf("failed to process pull output: %v", err)
-	}
-
-	return nil
-}
-
 // createTarFromDir creates a tar archive from a directory
 func createTarFromDir(srcDir string, writer io.Writer) error {
 	// Create a new tar writer

pkg/container/images/image.go

@@ -5,6 +5,7 @@ import (
 	"context"
 
 	"github.com/stacklok/toolhive/pkg/container/docker/sdk"
+	"github.com/stacklok/toolhive/pkg/container/runtime"
 	"github.com/stacklok/toolhive/pkg/logger"
 )
 
@@ -26,15 +27,20 @@ type ImageManager interface {
 // NewImageManager creates an instance of ImageManager appropriate
 // for the current environment, or returns an error if it is not supported.
 func NewImageManager(ctx context.Context) ImageManager {
-	// ASSUMPTION: This only works for the Docker runtime.
-	// Otherwise, return a no-op implementation.
+	// Check if we are running in a Kubernetes environment
+	if runtime.IsKubernetesRuntime() {
+		logger.Debug("running in Kubernetes environment, using no-op image manager")
+		return &NoopImageManager{}
+	}
+
+	// Check if we are running in a Docker or compatible environment
 	dockerClient, _, _, err := sdk.NewDockerClient(ctx)
 	if err != nil {
 		logger.Debug("no docker runtime found, using no-op image manager")
 		return &NoopImageManager{}
 	}
 
-	return NewDockerImageManager(dockerClient)
+	return NewRegistryImageManager(dockerClient)
 }
 
 // NoopImageManager is a no-op implementation of ImageManager.

pkg/container/images/keychain.go

@@ -0,0 +1,73 @@
+package images
+
+import (
+	"os"
+	"strings"
+
+	"github.com/google/go-containerregistry/pkg/authn"
+)
+
+// envKeychain implements a keychain that reads credentials from environment variables
+type envKeychain struct{}
+
+// Resolve implements the authn.Keychain interface
+func (*envKeychain) Resolve(target authn.Resource) (authn.Authenticator, error) {
+	registry := target.RegistryStr()
+
+	// Try registry-specific environment variables first
+	// Format: REGISTRY_<NORMALIZED_REGISTRY_NAME>_USERNAME/PASSWORD, i.e., REGISTRY_DOCKER_IO_USERNAME
+	normalizedRegistry := strings.ToUpper(strings.ReplaceAll(registry, ".", "_"))
+	normalizedRegistry = strings.ReplaceAll(normalizedRegistry, "-", "_")
+
+	username := os.Getenv("REGISTRY_" + normalizedRegistry + "_USERNAME")
+	password := os.Getenv("REGISTRY_" + normalizedRegistry + "_PASSWORD")
+
+	// If registry-specific vars not found, try generic one REGISTRY_USERNAME/PASSWORD
+	if username == "" || password == "" {
+		username = os.Getenv("REGISTRY_USERNAME")
+		password = os.Getenv("REGISTRY_PASSWORD")
+	}
+
+	if username != "" && password != "" {
+		return &authn.Basic{
+			Username: username,
+			Password: password,
+		}, nil
+	}
+
+	return authn.Anonymous, nil
+}
+
+// compositeKeychain combines multiple keychains and tries them in order
+type compositeKeychain struct {
+	keychains []authn.Keychain
+}
+
+// Resolve implements the authn.Keychain interface
+func (c *compositeKeychain) Resolve(target authn.Resource) (authn.Authenticator, error) {
+	for _, keychain := range c.keychains {
+		auth, err := keychain.Resolve(target)
+		if err != nil {
+			continue
+		}
+
+		// Check if we got actual credentials (not anonymous)
+		if auth != authn.Anonymous {
+			return auth, nil
+		}
+	}
+
+	// If no keychain provided credentials, return anonymous
+	return authn.Anonymous, nil
+}
+
+// NewCompositeKeychain creates a keychain that tries environment variables first,
+// then falls back to the default keychain
+func NewCompositeKeychain() authn.Keychain {
+	return &compositeKeychain{
+		keychains: []authn.Keychain{
+			&envKeychain{},        // Try environment variables first
+			authn.DefaultKeychain, // Then try default keychain (Docker config, etc.)
+		},
+	}
+}

pkg/container/images/registry.go

@@ -0,0 +1,127 @@
+package images
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"runtime"
+
+	"github.com/docker/docker/client"
+	"github.com/google/go-containerregistry/pkg/authn"
+	"github.com/google/go-containerregistry/pkg/name"
+	"github.com/google/go-containerregistry/pkg/v1"
+	"github.com/google/go-containerregistry/pkg/v1/daemon"
+	"github.com/google/go-containerregistry/pkg/v1/remote"
+
+	"github.com/stacklok/toolhive/pkg/logger"
+)
+
+// RegistryImageManager implements the ImageManager interface using go-containerregistry
+// for direct registry operations without requiring a Docker daemon.
+// However, for building images from Dockerfiles, it still uses the Docker client.
+type RegistryImageManager struct {
+	keychain     authn.Keychain
+	platform     *v1.Platform
+	dockerClient *client.Client
+}
+
+// NewRegistryImageManager creates a new RegistryImageManager instance
+func NewRegistryImageManager(dockerClient *client.Client) *RegistryImageManager {
+	return &RegistryImageManager{
+		keychain:     NewCompositeKeychain(), // Use composite keychain (env vars + default)
+		platform:     getDefaultPlatform(),   // Use a default platform based on host architecture
+		dockerClient: dockerClient,           // Used solely for building images from Dockerfiles
+	}
+}
+
+// getDefaultPlatform returns the default platform for containers
+// Uses host architecture
+func getDefaultPlatform() *v1.Platform {
+	return &v1.Platform{
+		Architecture: runtime.GOARCH,
+		OS:           "linux", // TODO: Should we support Windows too?
+	}
+}
+
+// ImageExists checks if an image exists locally in the daemon or remotely in the registry
+func (*RegistryImageManager) ImageExists(_ context.Context, imageName string) (bool, error) {
+	// Parse the image reference
+	ref, err := name.ParseReference(imageName)
+	if err != nil {
+		return false, fmt.Errorf("failed to parse image reference %q: %w", imageName, err)
+	}
+
+	// First check if image exists locally in daemon
+	if _, err := daemon.Image(ref); err != nil {
+		// Image does not exist locally
+		return false, nil
+	}
+	// Image exists locally
+	return true, nil
+}
+
+// PullImage pulls an image from a registry and saves it to the local daemon
+func (r *RegistryImageManager) PullImage(ctx context.Context, imageName string) error {
+	logger.Infof("Pulling image: %s", imageName)
+
+	// Parse the image reference
+	ref, err := name.ParseReference(imageName)
+	if err != nil {
+		return fmt.Errorf("failed to parse image reference %q: %w", imageName, err)
+	}
+
+	// Configure remote options
+	remoteOpts := []remote.Option{
+		remote.WithAuthFromKeychain(r.keychain),
+		remote.WithContext(ctx),
+	}
+
+	if r.platform != nil {
+		remoteOpts = append(remoteOpts, remote.WithPlatform(*r.platform))
+	}
+
+	// Pull the image from the registry
+	img, err := remote.Image(ref, remoteOpts...)
+	if err != nil {
+		return fmt.Errorf("failed to pull image from registry: %w", err)
+	}
+
+	// Convert reference to tag for daemon.Write
+	tag, ok := ref.(name.Tag)
+	if !ok {
+		// If it's not a tag, try to convert to tag
+		tag, err = name.NewTag(ref.String())
+		if err != nil {
+			return fmt.Errorf("failed to convert reference to tag: %w", err)
+		}
+	}
+
+	// Save the image to the local daemon
+	response, err := daemon.Write(tag, img)
+	if err != nil {
+		return fmt.Errorf("failed to write image to daemon: %w", err)
+	}
+
+	// Display success message
+	fmt.Fprintf(os.Stdout, "Successfully pulled %s\n", imageName)
+	logger.Infof("Pull complete for image: %s, response: %s", imageName, response)
+
+	return nil
+}
+
+// BuildImage builds a Docker image from a Dockerfile in the specified context directory
+func (r *RegistryImageManager) BuildImage(ctx context.Context, contextDir, imageName string) error {
+	return buildDockerImage(ctx, r.dockerClient, contextDir, imageName)
+}
+
+// WithKeychain sets the keychain for authentication
+func (r *RegistryImageManager) WithKeychain(keychain authn.Keychain) *RegistryImageManager {
+	r.keychain = keychain
+	return r
+}
+
+// WithPlatform sets the platform for the RegistryImageManager
+func (r *RegistryImageManager) WithPlatform(platform *v1.Platform) *RegistryImageManager {
+	r.platform = platform
+	return r
+}

pkg/container/runtime/types.go

@@ -5,6 +5,7 @@ package runtime
 import (
 	"context"
 	"io"
+	"os"
 	"time"
 
 	"github.com/stacklok/toolhive/pkg/permissions"
@@ -219,3 +220,9 @@ type Mount struct {
 	// ReadOnly indicates if the mount is read-only
 	ReadOnly bool
 }
+
+// IsKubernetesRuntime returns true if the runtime is Kubernetes
+// isn't the best way to do this, but for now it's good enough
+func IsKubernetesRuntime() bool {
+	return os.Getenv("KUBERNETES_SERVICE_HOST") != ""
+}