pkg-proxy/internal/handler/conda.go

package handler

import (
	"compress/gzip"
	"encoding/json"
	"io"
	"net/http"
	"slices"
	"strings"
	"time"

	"github.com/git-pkgs/proxy/internal/retention"
	"github.com/git-pkgs/purl"
)

const (
	condaUpstream = "https://conda.anaconda.org"
	minCondaParts = 3 // name-version-build requires at least 3 hyphen-separated parts
)

// Packages are stored under the name {channel}/{name}, which purl writes as
// pkg:conda/{channel}%2F{name}. Package rules name the channel as the
// namespace, pkg:conda/conda-forge/numpy, and are mapped to that form.
func init() {
	retention.Default.Register(retention.Spec{Key: "conda", CanonicalPackage: condaRetentionKey})
}

// condaRetentionKey maps a package rule to the stored package PURL. It
// takes the channel from the namespace, or from the name in the stored
// form. Downloads are routed as /{channel}/{arch}/{filename}, so the channel
// is always one path segment; a key without a channel, with more segments
// or with an empty one is refused, since no stored package looks like that.
// The package name is lowercased the way conda treats it; the channel stays
// as the URL had it.
func condaRetentionKey(p *purl.PURL) (string, bool) {
	if p == nil {
		return "", false
	}
	full := p.Name
	if p.Namespace != "" {
		full = p.Namespace + "/" + p.Name
	}
	segments := strings.Split(full, "/")
	if len(segments) != 2 || slices.Contains(segments, "") {
		return "", false
	}
	channel, name := segments[0], strings.ToLower(segments[1])
	return retention.DefaultCanonical("conda")(purl.New(p.Type, "", channel+"/"+name, "", nil))
}

// CondaHandler handles Conda/Anaconda registry protocol requests.
type CondaHandler struct {
	proxy       *Proxy
	upstreamURL string
	proxyURL    string
}

// NewCondaHandler creates a new Conda protocol handler.
func NewCondaHandler(proxy *Proxy, proxyURL string) *CondaHandler {
	return &CondaHandler{
		proxy:       proxy,
		upstreamURL: condaUpstream,
		proxyURL:    strings.TrimSuffix(proxyURL, "/"),
	}
}

// NewCondaHandlerWithUpstream creates a Conda handler with a custom upstream.
func NewCondaHandlerWithUpstream(proxy *Proxy, proxyURL, upstreamURL string) *CondaHandler {
	h := NewCondaHandler(proxy, proxyURL)
	h.upstreamURL = configuredUpstreamURL(upstreamURL, condaUpstream)
	return h
}

// Routes returns the HTTP handler for Conda requests.
func (h *CondaHandler) Routes() http.Handler {
	mux := http.NewServeMux()

	// Channel index (repodata)
	mux.HandleFunc("GET /{channel}/{arch}/repodata.json", h.handleRepodata)
	mux.HandleFunc("GET /{channel}/{arch}/repodata.json.bz2", h.proxyCached)
	mux.HandleFunc("GET /{channel}/{arch}/current_repodata.json", h.handleRepodata)

	// Package downloads (cache these)
	mux.HandleFunc("GET /{channel}/{arch}/{filename}", h.handleDownload)

	return mux
}

// handleDownload serves a package file, fetching and caching from upstream if needed.
func (h *CondaHandler) handleDownload(w http.ResponseWriter, r *http.Request) {
	channel := r.PathValue("channel")
	arch := r.PathValue("arch")
	filename := r.PathValue("filename")

	// Only cache actual package files
	if !h.isPackageFile(filename) {
		h.proxyUpstream(w, r)
		return
	}

	name, version := h.parseFilename(filename)
	if name == "" {
		h.proxyUpstream(w, r)
		return
	}

	// Include channel in package name
	packageName := channel + "/" + name

	h.proxy.Logger.Info("conda download request",
		"channel", channel, "arch", arch, "name", name, "version", version)

	upstreamURL := h.upstreamURL + r.URL.Path

	result, err := h.proxy.GetOrFetchArtifactFromURL(r.Context(), "conda", packageName, version, filename, upstreamURL)
	if err != nil {
		h.proxy.serveArtifactError(w, err, "failed to fetch package")
		return
	}

	ServeArtifactRequest(w, r, result)
}

// isPackageFile returns true if the filename is a Conda package.
func (h *CondaHandler) isPackageFile(filename string) bool {
	return strings.HasSuffix(filename, ".tar.bz2") || strings.HasSuffix(filename, ".conda")
}

// parseFilename extracts name and version from a Conda package filename.
// Conda filenames are: {name}-{version}-{build}.{ext}
// e.g., "numpy-1.24.0-py311h64a7726_0.conda"
func (h *CondaHandler) parseFilename(filename string) (name, version string) {
	// Remove extension
	base := filename
	for _, ext := range []string{".conda", ".tar.bz2"} {
		if strings.HasSuffix(base, ext) {
			base = strings.TrimSuffix(base, ext)
			break
		}
	}

	// Split by hyphens, the format is name-version-build
	// The name can contain hyphens, so we need to find version-build at the end
	parts := strings.Split(base, "-")
	if len(parts) < minCondaParts {
		return "", ""
	}

	// Build is the last part, version is second to last
	// Everything before is the name
	build := parts[len(parts)-1]
	version = parts[len(parts)-2]
	name = strings.Join(parts[:len(parts)-2], "-")

	// Version should start with a digit
	if len(version) == 0 || version[0] < '0' || version[0] > '9' {
		return "", ""
	}

	// Include build in version for uniqueness
	_ = build

	return name, version
}

// handleRepodata proxies repodata.json, applying cooldown filtering when enabled.
func (h *CondaHandler) handleRepodata(w http.ResponseWriter, r *http.Request) {
	if h.proxy.Cooldown == nil || !h.proxy.Cooldown.Enabled() {
		h.proxyCached(w, r)
		return
	}

	upstreamURL := h.upstreamURL + r.URL.Path

	h.proxy.Logger.Debug("fetching repodata for cooldown filtering", "url", upstreamURL)

	req, err := http.NewRequestWithContext(r.Context(), http.MethodGet, upstreamURL, nil)
	if err != nil {
		http.Error(w, "failed to create request", http.StatusInternalServerError)
		return
	}
	req.Header.Set(headerAcceptEncoding, "gzip")

	resp, err := h.proxy.HTTPClient.Do(req)
	if err != nil {
		if writeUpstreamHeld(w, err) {
			return
		}
		h.proxy.Logger.Error("upstream request failed", "error", err)
		http.Error(w, "upstream request failed", http.StatusBadGateway)
		return
	}
	defer func() { _ = resp.Body.Close() }()

	if resp.StatusCode != http.StatusOK {
		h.proxy.relayResponse(w, r, resp, nil)
		return
	}

	// Explicitly requesting gzip disables the transport's automatic decoding.
	// Cooldown needs JSON, so decode here and bound the decompressed size too.
	var reader io.Reader = resp.Body
	if strings.EqualFold(strings.TrimSpace(resp.Header.Get(headerContentEncoding)), "gzip") {
		decoded, err := gzip.NewReader(resp.Body)
		if err != nil {
			http.Error(w, "failed to decode response", http.StatusBadGateway)
			return
		}
		defer func() { _ = decoded.Close() }()
		reader = decoded
	}
	body, err := h.proxy.ReadMetadata(reader)
	if err != nil {
		http.Error(w, "failed to read response", http.StatusBadGateway)
		return
	}

	filtered, err := h.applyCooldownFiltering(body)
	if err != nil {
		h.proxy.Logger.Warn("failed to filter repodata, proxying original", "error", err)
		w.Header().Set(headerContentType, "application/json")
		_, _ = w.Write(body)
		return
	}

	w.Header().Set(headerContentType, "application/json")
	_, _ = w.Write(filtered)
}

// condaTimestampDivisor converts Conda's millisecond timestamps to seconds.
const condaTimestampDivisor = 1000

// applyCooldownFiltering removes entries from repodata.json that were
// published too recently based on their timestamp field.
func (h *CondaHandler) applyCooldownFiltering(body []byte) ([]byte, error) {
	if h.proxy.Cooldown == nil || !h.proxy.Cooldown.Enabled() {
		return body, nil
	}

	var repodata map[string]any
	if err := json.Unmarshal(body, &repodata); err != nil {
		return nil, err
	}

	for _, key := range []string{"packages", "packages.conda"} {
		packages, ok := repodata[key].(map[string]any)
		if !ok {
			continue
		}

		for filename, entry := range packages {
			entryMap, ok := entry.(map[string]any)
			if !ok {
				continue
			}

			ts, ok := entryMap["timestamp"].(float64)
			if !ok || ts == 0 {
				continue
			}

			publishedAt := time.Unix(int64(ts)/condaTimestampDivisor, 0)

			name, _ := entryMap["name"].(string)
			if name == "" {
				continue
			}

			packagePURL := canonicalPackagePURL("conda", name)

			if !h.proxy.Cooldown.IsAllowed("conda", packagePURL, publishedAt) {
				version, _ := entryMap["version"].(string)
				h.proxy.Logger.Info("cooldown: filtering conda package",
					"name", name, "version", version, "filename", filename)
				delete(packages, filename)
			}
		}
	}

	return json.Marshal(repodata)
}

// proxyCached forwards a metadata request with caching.
func (h *CondaHandler) proxyCached(w http.ResponseWriter, r *http.Request) {
	cacheKey := strings.TrimPrefix(r.URL.Path, "/")
	cacheKey = strings.ReplaceAll(cacheKey, "/", "_")
	// Large JSON indexes can exceed the metadata limit uncompressed. Conda
	// clients decode Content-Encoding, so cache and relay gzip verbatim. The
	// already-compressed repodata.json.bz2 route must stay on identity.
	acceptEncoding := "identity"
	if strings.HasSuffix(r.URL.Path, ".json") {
		acceptEncoding = "gzip"
	}
	h.proxy.proxyCachedWithEncoding(w, r, h.upstreamURL+r.URL.Path, "conda", cacheKey, acceptEncoding, "*/*")
}

// proxyUpstream forwards a request to Anaconda without caching.
func (h *CondaHandler) proxyUpstream(w http.ResponseWriter, r *http.Request) {
	h.proxy.ProxyUpstream(w, r, h.upstreamURL+r.URL.Path, []string{headerAcceptEncoding})
}