Mirror PyPI packages

docs/configuration.md

@@ -1506,6 +1506,8 @@ The `proxy mirror` command pre-populates the cache from various sources. It acce
 
 Packages are downloaded from the configured upstreams: npm, cargo, gem, golang, hex, pub, maven and nuget from `upstream.npm`, `upstream.cargo_download`, `upstream.gem`, `upstream.go`, `upstream.hex`, `upstream.pub`, `upstream.maven` and `upstream.nuget`. npm packages in `upstream.npm_routes` come from their route and are cached under the route's keys, so clients get them from the cache as they would after an install. Mirror API jobs do the same.
 
+A PyPI version is mirrored with every file `upstream.pypi` lists for it: the sdist and each wheel, checked against the listed sha256 and cached under the name and version the download handler reads from the filename. Add the `file_name` qualifier to mirror one file, for example `pkg:pypi/zope.interface@7.2?file_name=zope.interface-7.2-cp312-cp312-macosx_11_0_arm64.whl`.
+
 | Flag | Default | Description |
 |------|---------|-------------|
 | `--sbom` | | Path to CycloneDX or SPDX SBOM file |

internal/handler/pypi.go

@@ -649,13 +649,19 @@ var archiveExtensions = []string{".tar.gz", ".tar.bz2", ".tar.xz", ".tar.Z", ".t
 // bdist_msi formats, which share a filename layout.
 var windowsInstallerExtensions = []string{".exe", ".msi"}
 
-// parseFilename extracts package name and version from a PyPI filename.
-// Handles wheels, sdists and legacy bdist formats:
+func (h *PyPIHandler) parseFilename(filename string) (name, version string) {
+	return PyPIFilenameIdentity(filename)
+}
+
+// PyPIFilenameIdentity extracts the package name and version a PyPI file is
+// cached under from its filename, so proxy mirror stores files where the
+// download handler looks for them. Handles wheels, sdists and legacy bdist
+// formats:
 // - requests-2.31.0-py3-none-any.whl
 // - requests-2.31.0.tar.gz
 // - numpy-1.8.0-py2.7-macosx-10.9-x86_64.egg
 // - numpy-1.8.0.win32-py2.7.exe
-func (h *PyPIHandler) parseFilename(filename string) (name, version string) {
+func PyPIFilenameIdentity(filename string) (name, version string) {
 	// PEP 658/714 core-metadata sidecars are the distribution filename plus
 	// ".metadata"; they describe the same name and version. Without this, pip's
 	// metadata-only fetches fall back to a hash-derived package identifier.

internal/mirror/mirror.go

@@ -3,6 +3,7 @@ package mirror
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"log/slog"
 	"sync"
@@ -226,6 +227,10 @@ func (m *Mirror) mirrorOne(ctx context.Context, pv PackageVersion, tracker *prog
 			"ecosystem", pv.Ecosystem, "name", pv.Name, "version", pv.Version, "reason", m.RefusePackages)
 		return
 	}
+	if pv.Ecosystem == "pypi" {
+		m.mirrorFiles(ctx, pv, tracker)
+		return
+	}
 	result, err := m.proxy.GetOrFetchArtifact(ctx, pv.Ecosystem, pv.Name, pv.Version, "")
 	if err != nil {
 		tracker.failed.Add(1)
@@ -251,3 +256,70 @@ func (m *Mirror) mirrorOne(ctx context.Context, pv PackageVersion, tracker *prog
 			"size", result.Artifact.Size)
 	}
 }
+
+// mirrorFiles mirrors every file of a PyPI version, or only the one named by
+// the file_name qualifier. Each file is cached under the name and version
+// the download handler reads from its filename, so clients are served it.
+func (m *Mirror) mirrorFiles(ctx context.Context, pv PackageVersion, tracker *progressTracker) {
+	fail := func(err error) {
+		tracker.failed.Add(1)
+		tracker.addError(pv.Ecosystem, pv.Name, pv.Version, err.Error())
+		m.logger.Warn("mirror failed",
+			"ecosystem", pv.Ecosystem, "name", pv.Name, "version", pv.Version, "error", err)
+	}
+	files, err := m.fileLister().ListFiles(ctx, pv.Ecosystem, pv.Name, pv.Version)
+	if err != nil {
+		fail(fmt.Errorf("listing files: %w", err))
+		return
+	}
+	wantFile := pv.Qualifiers["file_name"]
+	var matched, fetched int
+	var size int64
+	for _, file := range files {
+		if wantFile != "" && file.Filename != wantFile {
+			continue
+		}
+		matched++
+		name, version := handler.PyPIFilenameIdentity(file.Filename)
+		if name == "" {
+			name, version = pv.Name, pv.Version
+		}
+		result, err := m.proxy.GetOrFetchArtifactFromURLWithIntegrity(ctx, pv.Ecosystem, name, version, file.Filename, file.URL, file.Integrity)
+		if err != nil {
+			fail(fmt.Errorf("%s: %w", file.Filename, err))
+			return
+		}
+		if result.Reader != nil {
+			_ = result.Reader.Close()
+		}
+		if !result.Cached {
+			fetched++
+			size += result.Artifact.Size
+		}
+	}
+	switch {
+	case matched == 0 && wantFile != "":
+		fail(fmt.Errorf("no file named %q", wantFile))
+	case matched == 0:
+		fail(errors.New("version has no files"))
+	case fetched == 0:
+		tracker.skipped.Add(1)
+		m.logger.Debug("already cached",
+			"ecosystem", pv.Ecosystem, "name", pv.Name, "version", pv.Version)
+	default:
+		tracker.completed.Add(1)
+		tracker.bytes.Add(size)
+		m.logger.Info("mirrored",
+			"ecosystem", pv.Ecosystem, "name", pv.Name, "version", pv.Version,
+			"files", fetched, "size", size)
+	}
+}
+
+// fileLister lists files from where Versions lists versions, so both come
+// from the configured upstream.
+func (m *Mirror) fileLister() FileLister { //nolint:ireturn // any lister will do
+	if lister, ok := m.Versions.(FileLister); ok {
+		return lister
+	}
+	return RegistryVersions{}
+}

internal/mirror/source.go

@@ -14,7 +14,7 @@ type PackageVersion struct {
 	Name      string
 	Version   string
 	// Qualifiers are the PURL qualifiers. Images use repository_url, tag and
-	// arch.
+	// arch; PyPI uses file_name to mirror one file of the version.
 	Qualifiers map[string]string
 	// Platforms limit which platforms of an image index are mirrored.
 	Platforms []string
@@ -78,6 +78,8 @@ func (s *PURLSource) Enumerate(ctx context.Context, fn func(PackageVersion) erro
 			if isImageEcosystem(ecosystem) {
 				pv.Qualifiers = p.Qualifiers.Map()
 				pv.Platforms = s.Platforms
+			} else if fileName := p.Qualifier("file_name"); fileName != "" && ecosystem == "pypi" {
+				pv.Qualifiers = map[string]string{"file_name": fileName}
 			}
 			if err := fn(pv); err != nil {
 				return err

internal/mirror/versions.go

@@ -2,6 +2,7 @@ package mirror
 
 import (
 	"context"
+	"fmt"
 
 	"github.com/git-pkgs/purl"
 	"github.com/git-pkgs/registries"
@@ -22,8 +23,40 @@ type RegistryVersions struct {
 	URL func(purlType, name string) (string, error)
 }
 
+// FileLister lists the files of one package version, for ecosystems that
+// publish several files per version.
+type FileLister interface {
+	ListFiles(ctx context.Context, ecosystem, name, version string) ([]registries.Artifact, error)
+}
+
 // ListVersions implements VersionLister.
 func (r RegistryVersions) ListVersions(ctx context.Context, ecosystem, name string) ([]string, error) {
+	versions, err := r.fetchVersions(ctx, ecosystem, name)
+	if err != nil {
+		return nil, err
+	}
+	result := make([]string, len(versions))
+	for i, v := range versions {
+		result[i] = v.Number
+	}
+	return result, nil
+}
+
+// ListFiles implements FileLister.
+func (r RegistryVersions) ListFiles(ctx context.Context, ecosystem, name, version string) ([]registries.Artifact, error) {
+	versions, err := r.fetchVersions(ctx, ecosystem, name)
+	if err != nil {
+		return nil, err
+	}
+	for _, v := range versions {
+		if v.Number == version {
+			return v.Artifacts, nil
+		}
+	}
+	return nil, fmt.Errorf("%w: %s %s@%s", registries.ErrNotFound, ecosystem, name, version)
+}
+
+func (r RegistryVersions) fetchVersions(ctx context.Context, ecosystem, name string) ([]registries.Version, error) {
 	purlType := purl.EcosystemToPURLType(ecosystem)
 	var baseURL string
 	if r.URL != nil {
@@ -40,13 +73,5 @@ func (r RegistryVersions) ListVersions(ctx context.Context, ecosystem, name stri
 	if err != nil {
 		return nil, err
 	}
-	versions, err := reg.FetchVersions(ctx, name)
-	if err != nil {
-		return nil, err
-	}
-	result := make([]string, len(versions))
-	for i, v := range versions {
-		result[i] = v.Number
-	}
-	return result, nil
+	return reg.FetchVersions(ctx, name)
 }

internal/server/mirror_fetch_test.go

@@ -52,7 +52,9 @@ func mirrorWithConfig(t *testing.T, cfg *config.Config, purls ...string) (*mirro
 	}
 	logger := slog.New(slog.NewTextHandler(io.Discard, nil))
 	proxy, _ := NewMirrorProxy(cfg, db, store, logger)
-	progress, err := mirror.New(proxy, db, store, logger, 1).Run(context.Background(), &mirror.PURLSource{PURLs: purls})
+	m := mirror.New(proxy, db, store, logger, 1)
+	m.Versions = NewMirrorVersions(cfg.Upstream, proxy.HTTPClient)
+	progress, err := m.Run(context.Background(), &mirror.PURLSource{PURLs: purls})
 	if err != nil {
 		t.Fatalf("Run() error = %v", err)
 	}

internal/server/mirror_pypi_test.go

@@ -0,0 +1,140 @@
+package server
+
+import (
+	"crypto/sha256"
+	"encoding/hex"
+	"encoding/json"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"sync"
+	"testing"
+
+	"github.com/git-pkgs/proxy/internal/config"
+	"github.com/git-pkgs/proxy/internal/handler"
+)
+
+const (
+	mirrorWheelPath = "/packages/ab/cd/example_pkg-1.0-py3-none-any.whl"
+	mirrorSdistPath = "/packages/ef/01/example.pkg-1.0.tar.gz"
+)
+
+// pypiUpstream serves the JSON API for example-pkg 1.0 with a wheel and an
+// sdist, whose filenames spell the name differently, and counts file
+// downloads by path.
+type pypiUpstream struct {
+	*httptest.Server
+	mu    sync.Mutex
+	files map[string]string
+	hits  map[string]int
+}
+
+func newPyPIUpstream(t *testing.T, digests map[string]string) *pypiUpstream {
+	t.Helper()
+	u := &pypiUpstream{
+		files: map[string]string{mirrorWheelPath: "wheel", mirrorSdistPath: "sdist"},
+		hits:  map[string]int{},
+	}
+	u.Server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.URL.Path == "/pypi/example-pkg/json" {
+			var files []map[string]any
+			for _, path := range []string{mirrorWheelPath, mirrorSdistPath} {
+				sum := sha256.Sum256([]byte(u.files[path]))
+				digest := hex.EncodeToString(sum[:])
+				if d, ok := digests[path]; ok {
+					digest = d
+				}
+				files = append(files, map[string]any{
+					"url":     u.URL + path,
+					"digests": map[string]string{"sha256": digest},
+					"size":    len(u.files[path]),
+				})
+			}
+			_ = json.NewEncoder(w).Encode(map[string]any{
+				"info":     map[string]any{"name": "example-pkg"},
+				"releases": map[string]any{"1.0": files},
+			})
+			return
+		}
+		body, ok := u.files[r.URL.Path]
+		if !ok {
+			http.NotFound(w, r)
+			return
+		}
+		u.mu.Lock()
+		u.hits[r.URL.Path]++
+		u.mu.Unlock()
+		_, _ = io.WriteString(w, body)
+	}))
+	t.Cleanup(u.Close)
+	return u
+}
+
+func (u *pypiUpstream) hitsFor(path string) int {
+	u.mu.Lock()
+	defer u.mu.Unlock()
+	return u.hits[path]
+}
+
+func pypiMirrorConfig(upstream *pypiUpstream) *config.Config {
+	cfg := config.Default()
+	cfg.Upstream.PyPI = upstream.URL
+	cfg.Upstream.PyPIDownload = upstream.URL
+	return cfg
+}
+
+func TestMirrorPyPIFetchesEveryFile(t *testing.T) {
+	upstream := newPyPIUpstream(t, nil)
+	cfg := pypiMirrorConfig(upstream)
+
+	progress, proxy := mirrorWithConfig(t, cfg, "pkg:pypi/example-pkg@1.0")
+
+	if progress.Completed != 1 || progress.Failed != 0 {
+		t.Fatalf("completed = %d, failed = %d, errors = %v", progress.Completed, progress.Failed, progress.Errors)
+	}
+	for _, path := range []string{mirrorWheelPath, mirrorSdistPath} {
+		if got := upstream.hitsFor(path); got != 1 {
+			t.Errorf("%s fetched %d times, want 1", path, got)
+		}
+	}
+
+	// pip gets both files from the cache once the upstream is down.
+	upstream.Close()
+	pypi := handler.NewPyPIHandlerWithUpstreams(proxy, "http://proxy.test", cfg.Upstream.PyPI, cfg.Upstream.PyPIDownload)
+	for path, want := range map[string]string{mirrorWheelPath: "wheel", mirrorSdistPath: "sdist"} {
+		rec := httptest.NewRecorder()
+		pypi.Routes().ServeHTTP(rec, httptest.NewRequest(http.MethodGet, path, nil))
+		if rec.Code != http.StatusOK || rec.Body.String() != want {
+			t.Errorf("%s from cache: status = %d, body = %q, want 200 and %q", path, rec.Code, rec.Body, want)
+		}
+	}
+}
+
+func TestMirrorPyPIFileNameQualifier(t *testing.T) {
+	upstream := newPyPIUpstream(t, nil)
+
+	progress, _ := mirrorWithConfig(t, pypiMirrorConfig(upstream), "pkg:pypi/example-pkg@1.0?file_name=example.pkg-1.0.tar.gz")
+
+	if progress.Completed != 1 || upstream.hitsFor(mirrorSdistPath) != 1 {
+		t.Errorf("completed = %d, sdist hits = %d, errors = %v; want the sdist", progress.Completed, upstream.hitsFor(mirrorSdistPath), progress.Errors)
+	}
+	if upstream.hitsFor(mirrorWheelPath) != 0 {
+		t.Error("the wheel was mirrored although file_name names the sdist")
+	}
+
+	progress, _ = mirrorWithConfig(t, pypiMirrorConfig(upstream), "pkg:pypi/example-pkg@1.0?file_name=missing-1.0.tar.gz")
+	if progress.Failed != 1 || len(progress.Errors) != 1 || !strings.Contains(progress.Errors[0].Error, "missing-1.0.tar.gz") {
+		t.Errorf("failed = %d, errors = %v; want a failure naming the missing file", progress.Failed, progress.Errors)
+	}
+}
+
+func TestMirrorPyPIRefusesDigestMismatch(t *testing.T) {
+	upstream := newPyPIUpstream(t, map[string]string{mirrorWheelPath: strings.Repeat("0", 64)})
+
+	progress, _ := mirrorWithConfig(t, pypiMirrorConfig(upstream), "pkg:pypi/example-pkg@1.0?file_name=example_pkg-1.0-py3-none-any.whl")
+
+	if progress.Failed != 1 || progress.Completed != 0 {
+		t.Errorf("completed = %d, failed = %d, errors = %v; want the mismatching wheel refused", progress.Completed, progress.Failed, progress.Errors)
+	}
+}