deps: vendor utility dependencies

Signed-off-by: Jacob Howard <[email protected]>

Author: xenoscopic

pkg/inference/backends/llamacpp/download.go

@@ -13,7 +13,7 @@ import (
 	"runtime"
 	"strings"
 
-	"github.com/docker/model-runner/pkg/dockerhub"
+	"github.com/docker/model-runner/pkg/internal/dockerhub"
 	"github.com/docker/model-runner/pkg/paths"
 )
 

pkg/inference/models/manager.go

@@ -11,9 +11,7 @@ import (
 	"github.com/docker/model-distribution/pkg/distribution"
 	"github.com/docker/model-distribution/pkg/types"
 	"github.com/docker/model-runner/pkg/inference"
-	"github.com/docker/model-runner/pkg/ipc"
 	"github.com/docker/model-runner/pkg/logger"
-	"github.com/docker/model-runner/pkg/paths"
 )
 
 const (
@@ -36,24 +34,13 @@ type Manager struct {
 }
 
 // NewManager creates a new model's manager.
-func NewManager(log logger.ComponentLogger, transport http.RoundTripper) *Manager {
-	// Create the distribution client
-	distributionClient, err := distribution.NewClient(
-		distribution.WithStoreRootPath(paths.DockerHome("models")),
-		distribution.WithLogger(log.WithField("component", "model-distribution")),
-		distribution.WithTransport(transport),
-		distribution.WithUserAgent(ipc.UserAgent),
-	)
-	if err != nil {
-		log.Errorf("Failed to create distribution client: %v", err)
-		// Continue without distribution client
-	}
+func NewManager(log logger.ComponentLogger, client *distribution.Client) *Manager {
 	// Create the manager.
 	m := &Manager{
 		log:                log,
 		pullTokens:         make(chan struct{}, maximumConcurrentModelPulls),
 		router:             http.NewServeMux(),
-		distributionClient: distributionClient,
+		distributionClient: client,
 	}
 
 	// Register routes.

pkg/internal/archive/archive.go

@@ -0,0 +1,42 @@
+package archive
+
+import (
+	"fmt"
+	"path/filepath"
+	"strings"
+)
+
+// CheckRelative returns an error if the filename path escapes dir.
+// This is used to protect against path traversal attacks when extracting archives.
+// It also rejects absolute filename paths.
+func CheckRelative(dir, filename string) (string, error) {
+	if filepath.IsAbs(filename) {
+		return "", fmt.Errorf("archive path has absolute path: %q", filename)
+	}
+	target := filepath.Join(dir, filename)
+	if resolved, err := filepath.EvalSymlinks(target); err == nil {
+		target = resolved
+		if resolved, err = filepath.EvalSymlinks(dir); err == nil {
+			dir = resolved
+		}
+	}
+	rel, err := filepath.Rel(dir, target)
+	if err != nil {
+		return "", err
+	}
+	if strings.HasPrefix(rel, "..") {
+		return "", fmt.Errorf("archive file %q escapes %q", target, dir)
+	}
+	return target, nil
+}
+
+// CheckSymlink returns an error if the link path escapes dir.
+// This is used to protect against path traversal attacks when extracting archives.
+// It also rejects absolute linkname paths.
+func CheckSymlink(dir, name, linkname string) error {
+	if filepath.IsAbs(linkname) {
+		return fmt.Errorf("archive path has absolute link: %q", linkname)
+	}
+	_, err := CheckRelative(dir, filepath.Join(filepath.Dir(name), linkname))
+	return err
+}

pkg/internal/dockerhub/download.go

@@ -0,0 +1,136 @@
+package dockerhub
+
+import (
+	"context"
+	"encoding/base64"
+	"errors"
+	"fmt"
+	"log"
+	"os"
+	"path/filepath"
+	"strings"
+	"time"
+
+	"github.com/containerd/containerd/content"
+	"github.com/containerd/containerd/content/local"
+	"github.com/containerd/containerd/images"
+	"github.com/containerd/containerd/images/archive"
+	"github.com/containerd/containerd/remotes"
+	"github.com/containerd/containerd/remotes/docker"
+	"github.com/containerd/platforms"
+	"github.com/docker/model-runner/pkg/internal/jsonutil"
+	v1 "github.com/opencontainers/image-spec/specs-go/v1"
+	"github.com/sirupsen/logrus"
+)
+
+func PullPlatform(ctx context.Context, image, destination, requiredOs, requiredArch string) error {
+	if err := os.MkdirAll(filepath.Dir(destination), 0o755); err != nil {
+		return fmt.Errorf("creating destination directory %s: %w", filepath.Dir(destination), err)
+	}
+	output, err := os.Create(destination)
+	if err != nil {
+		return fmt.Errorf("creating destination file %s: %w", destination, err)
+	}
+	tmpDir, err := os.MkdirTemp("", "docker-pull")
+	if err != nil {
+		return fmt.Errorf("creating temp directory: %w", err)
+	}
+	defer os.RemoveAll(tmpDir)
+	store, err := local.NewStore(tmpDir)
+	if err != nil {
+		return fmt.Errorf("creating new content store: %w", err)
+	}
+	desc, err := retry(ctx, 10, 1*time.Second, func() (*v1.Descriptor, error) { return fetch(ctx, store, image, requiredOs, requiredArch) })
+	if err != nil {
+		return fmt.Errorf("fetching image: %w", err)
+	}
+	return archive.Export(ctx, store, output, archive.WithManifest(*desc, image), archive.WithSkipMissing(store))
+}
+
+func retry(ctx context.Context, attempts int, sleep time.Duration, f func() (*v1.Descriptor, error)) (*v1.Descriptor, error) {
+	var err error
+	var result *v1.Descriptor
+	for i := 0; i < attempts; i++ {
+		if i > 0 {
+			log.Printf("retry %d after error: %v\n", i, err)
+			select {
+			case <-ctx.Done():
+				return nil, ctx.Err()
+			case <-time.After(sleep):
+			}
+		}
+		result, err = f()
+		if err == nil {
+			return result, nil
+		}
+	}
+	return nil, fmt.Errorf("after %d attempts, last error: %s", attempts, err)
+}
+
+func fetch(ctx context.Context, store content.Store, ref, requiredOs, requiredArch string) (*v1.Descriptor, error) {
+	resolver := docker.NewResolver(docker.ResolverOptions{
+		Hosts: docker.ConfigureDefaultRegistries(
+			docker.WithAuthorizer(
+				docker.NewDockerAuthorizer(
+					docker.WithAuthCreds(dockerCredentials)))),
+	})
+	name, desc, err := resolver.Resolve(ctx, ref)
+	if err != nil {
+		return nil, err
+	}
+	fetcher, err := resolver.Fetcher(ctx, name)
+	if err != nil {
+		return nil, err
+	}
+
+	childrenHandler := images.ChildrenHandler(store)
+	if requiredOs != "" && requiredArch != "" {
+		requiredPlatform := platforms.Only(v1.Platform{OS: requiredOs, Architecture: requiredArch})
+		childrenHandler = images.LimitManifests(images.FilterPlatforms(images.ChildrenHandler(store), requiredPlatform), requiredPlatform, 1)
+	}
+	h := images.Handlers(remotes.FetchHandler(store, fetcher), childrenHandler)
+	if err := images.Dispatch(ctx, h, nil, desc); err != nil {
+		return nil, err
+	}
+	return &desc, nil
+}
+
+func dockerCredentials(host string) (string, string, error) {
+	hubUsername, hubPassword := os.Getenv("DOCKER_HUB_USER"), os.Getenv("DOCKER_HUB_PASSWORD")
+	if hubUsername != "" && hubPassword != "" {
+		return hubUsername, hubPassword, nil
+	}
+	logrus.WithField("host", host).Debug("checking for registry auth config")
+	home, err := os.UserHomeDir()
+	if err != nil {
+		return "", "", err
+	}
+	credentialConfig := filepath.Join(home, ".docker", "config.json")
+	cfg := struct {
+		Auths map[string]struct {
+			Auth string
+		}
+	}{}
+	if err := jsonutil.ReadFile(credentialConfig, &cfg); err != nil {
+		if errors.Is(err, os.ErrNotExist) {
+			return "", "", nil
+		}
+		return "", "", err
+	}
+	for h, r := range cfg.Auths {
+		if h == host {
+			creds, err := base64.StdEncoding.DecodeString(r.Auth)
+			if err != nil {
+				return "", "", err
+			}
+			parts := strings.SplitN(string(creds), ":", 2)
+			if len(parts) != 2 {
+				logrus.Debugf("skipping not user/password auth for registry %s: %s", host, parts[0])
+				return "", "", nil
+			}
+			logrus.Debugf("using auth for registry %s: user=%s", host, parts[0])
+			return parts[0], parts[1], nil
+		}
+	}
+	return "", "", nil
+}