pmg/test/proxye2e/registry.go

package proxye2e

import (
	"encoding/json"
	"fmt"
	"net"
	"net/http"
	"net/http/httptest"
	"path"
	"strings"
	"sync"
	"time"

	"golang.org/x/mod/module"
)

type NpmVersion struct {
	Version     string
	PublishedAt time.Time
	Tarball     []byte
}

type NpmPackage struct {
	Name          string
	DistTagLatest string
	Versions      []NpmVersion
}

type PypiVersion struct {
	Version     string
	PublishedAt time.Time
	Bytes       []byte
}

type PypiPackage struct {
	Name     string
	Versions []PypiVersion
}

type GoVersion struct {
	Version     string
	PublishedAt time.Time
}

type GoModule struct {
	Path     string
	Versions []GoVersion
}

type RecordedRequest struct {
	Host   string
	Method string
	Path   string
}

// Registry is an in-process stand-in for the npm and PyPI registries. The proxy
// upstream is redirected here, so it answers for every registry hostname and
// records each request for routing assertions.
type Registry struct {
	mu       sync.Mutex
	npm      map[string]NpmPackage
	pypi     map[string]PypiPackage
	gomod    map[string]GoModule
	requests []RecordedRequest
	server   *httptest.Server
	goServer *httptest.Server
}

func newRegistry() *Registry {
	r := &Registry{
		npm:   map[string]NpmPackage{},
		pypi:  map[string]PypiPackage{},
		gomod: map[string]GoModule{},
	}
	r.server = httptest.NewTLSServer(http.HandlerFunc(r.serve))
	// Plain-HTTP GOPROXY endpoint for the interceptor's out-of-band .info
	// fetches, which go straight to the upstream base URL rather than through
	// the proxy under test. It also serves the /goproxy base path used to
	// exercise GOPROXY path-prefix handling.
	r.goServer = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		r.record(req)
		r.serveGoWithOptionalPrefix(w, req)
	}))
	return r
}

func (r *Registry) addr() string { return r.server.Listener.Addr().String() }

func (r *Registry) goBaseURL() string { return r.goServer.URL }

func (r *Registry) close() {
	r.server.Close()
	r.goServer.Close()
}

func (r *Registry) AddNpm(pkg NpmPackage) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.npm[pkg.Name] = pkg
}

func (r *Registry) AddPypi(pkg PypiPackage) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.pypi[normalizePypiName(pkg.Name)] = pkg
}

func (r *Registry) AddGoModule(mod GoModule) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.gomod[mod.Path] = mod
}

// Requests returns every request the proxy forwarded upstream, in order.
func (r *Registry) Requests() []RecordedRequest {
	r.mu.Lock()
	defer r.mu.Unlock()
	out := make([]RecordedRequest, len(r.requests))
	copy(out, r.requests)
	return out
}

// DownloadedTarball reports whether a tarball for the given npm package version
// was fetched from the registry.
func (r *Registry) DownloadedTarball(name, version string) bool {
	want := fmt.Sprintf("/%s/-/%s-%s.tgz", name, name, version)
	for _, req := range r.Requests() {
		if req.Path == want {
			return true
		}
	}
	return false
}

func (r *Registry) record(req *http.Request) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.requests = append(r.requests, RecordedRequest{Host: hostOnly(req.Host), Method: req.Method, Path: req.URL.Path})
}

func (r *Registry) serve(w http.ResponseWriter, req *http.Request) {
	r.record(req)

	switch hostOnly(req.Host) {
	case "registry.npmjs.org", "registry.yarnpkg.com":
		r.serveNpm(w, req)
	case "pypi.org":
		r.servePypiSimple(w, req)
	case "files.pythonhosted.org":
		r.servePypiFile(w, req)
	case "proxy.golang.org", "corp.example.com":
		r.serveGoWithOptionalPrefix(w, req)
	default:
		http.NotFound(w, req)
	}
}

// serveGoWithOptionalPrefix serves the GOPROXY protocol either at the root
// (proxy.golang.org) or under the /goproxy base path (corp.example.com and
// the corp base URL of the plain-HTTP go server).
func (r *Registry) serveGoWithOptionalPrefix(w http.ResponseWriter, req *http.Request) {
	if strings.HasPrefix(req.URL.Path, "/goproxy/") {
		http.StripPrefix("/goproxy", http.HandlerFunc(r.serveGo)).ServeHTTP(w, req)
		return
	}
	r.serveGo(w, req)
}

// DownloadedGoZip reports whether the module zip for the given path and
// version was fetched from the registry.
func (r *Registry) DownloadedGoZip(modulePath, version string) bool {
	want := "/" + goEscapePath(modulePath) + "/@v/" + goEscapeVersion(version) + ".zip"
	for _, req := range r.Requests() {
		if req.Path == want {
			return true
		}
	}
	return false
}

// serveGo implements a minimal GOPROXY protocol endpoint: .info (with publish
// time), .mod and .zip per registered module version, plus /sumdb/* which the
// real proxy serves for checksum-database lookups.
func (r *Registry) serveGo(w http.ResponseWriter, req *http.Request) {
	p := strings.TrimPrefix(req.URL.Path, "/")

	if strings.HasPrefix(p, "sumdb/") {
		_, _ = w.Write([]byte("e2e-sumdb"))
		return
	}

	escapedPath, versionPart, found := strings.Cut(p, "/@v/")
	if !found {
		http.NotFound(w, req)
		return
	}

	modulePath, err := module.UnescapePath(escapedPath)
	if err != nil {
		http.NotFound(w, req)
		return
	}

	r.mu.Lock()
	mod, ok := r.gomod[modulePath]
	r.mu.Unlock()
	if !ok {
		http.NotFound(w, req)
		return
	}

	ext := path.Ext(versionPart)
	version, err := module.UnescapeVersion(strings.TrimSuffix(versionPart, ext))
	if err != nil {
		http.NotFound(w, req)
		return
	}

	var published time.Time
	versionFound := false
	for _, v := range mod.Versions {
		if v.Version == version {
			published = v.PublishedAt
			versionFound = true
			break
		}
	}
	if !versionFound {
		http.NotFound(w, req)
		return
	}

	switch ext {
	case ".info":
		w.Header().Set("Content-Type", "application/json")
		body, _ := json.Marshal(map[string]string{
			"Version": version,
			"Time":    published.UTC().Format(time.RFC3339),
		})
		_, _ = w.Write(body)
	case ".mod":
		_, _ = fmt.Fprintf(w, "module %s\n", modulePath)
	case ".zip":
		w.Header().Set("Content-Type", "application/zip")
		_, _ = w.Write([]byte("e2e-module-zip"))
	default:
		http.NotFound(w, req)
	}
}

func goEscapePath(p string) string {
	escaped, err := module.EscapePath(p)
	if err != nil {
		return p
	}
	return escaped
}

func goEscapeVersion(v string) string {
	escaped, err := module.EscapeVersion(v)
	if err != nil {
		return v
	}
	return escaped
}

func (r *Registry) serveNpm(w http.ResponseWriter, req *http.Request) {
	path := strings.Trim(req.URL.Path, "/")

	if strings.Contains(path, "/-/") {
		w.Header().Set("Content-Type", "application/octet-stream")
		_, _ = w.Write([]byte("e2e-tarball"))
		return
	}

	r.mu.Lock()
	pkg, ok := r.npm[path]
	r.mu.Unlock()
	if !ok {
		http.NotFound(w, req)
		return
	}

	w.Header().Set("Content-Type", "application/json")
	_, _ = w.Write(buildPackument(pkg))
}

func (r *Registry) servePypiSimple(w http.ResponseWriter, req *http.Request) {
	name := strings.Trim(strings.TrimPrefix(req.URL.Path, "/simple/"), "/")

	r.mu.Lock()
	pkg, ok := r.pypi[normalizePypiName(name)]
	r.mu.Unlock()
	if !ok {
		http.NotFound(w, req)
		return
	}

	w.Header().Set("Content-Type", pypiSimpleContentType)
	_, _ = w.Write(buildPypiSimple(pkg))
}

func (r *Registry) servePypiFile(w http.ResponseWriter, _ *http.Request) {
	w.Header().Set("Content-Type", "application/octet-stream")
	_, _ = w.Write([]byte("e2e-wheel"))
}

const pypiSimpleContentType = "application/vnd.pypi.simple.v1+json"

func buildPackument(pkg NpmPackage) []byte {
	versions := map[string]any{}
	times := map[string]string{}
	for _, v := range pkg.Versions {
		versions[v.Version] = map[string]any{
			"name":    pkg.Name,
			"version": v.Version,
			"dist": map[string]any{
				"tarball": fmt.Sprintf("https://registry.npmjs.org/%s/-/%s-%s.tgz", pkg.Name, pkg.Name, v.Version),
			},
		}
		times[v.Version] = v.PublishedAt.UTC().Format(time.RFC3339)
	}

	latest := pkg.DistTagLatest
	if latest == "" && len(pkg.Versions) > 0 {
		latest = pkg.Versions[len(pkg.Versions)-1].Version
	}

	doc := map[string]any{
		"name":      pkg.Name,
		"dist-tags": map[string]string{"latest": latest},
		"versions":  versions,
		"time":      times,
	}

	body, _ := json.Marshal(doc)
	return body
}

func buildPypiSimple(pkg PypiPackage) []byte {
	norm := normalizePypiName(pkg.Name)
	files := []map[string]any{}
	for _, v := range pkg.Versions {
		filename := fmt.Sprintf("%s-%s.tar.gz", norm, v.Version)
		files = append(files, map[string]any{
			"filename":    filename,
			"url":         fmt.Sprintf("https://files.pythonhosted.org/packages/source/%c/%s/%s", norm[0], norm, filename),
			"hashes":      map[string]string{},
			"upload-time": v.PublishedAt.UTC().Format(time.RFC3339Nano),
		})
	}

	doc := map[string]any{
		"meta":  map[string]any{"api-version": "1.0"},
		"name":  norm,
		"files": files,
	}

	body, _ := json.Marshal(doc)
	return body
}

func hostOnly(host string) string {
	if h, _, err := net.SplitHostPort(host); err == nil {
		return h
	}
	return host
}