tiny-proxy/internal/goproxy/handler.go

// Package goproxy implements a Go module proxy (the GOPROXY protocol) at
// /go/<vrepo>/<module>/@v/... and /go/<vrepo>/<module>/@latest.
//
// A go vrepo's upstreams are ordered routing rules. Each one serves the
// module paths its modules patterns match (GOPRIVATE syntax; none = every
// module) from one of:
//
//	goproxy   another module proxy, relayed (proxy.golang.org, Athens, ...)
//	gitea     a Gitea repository, read through its REST API
//	gitlab    a GitLab project, read through its REST API
//	git       any git server, read over the smart-HTTP protocol
//
// For the git-backed kinds the proxy builds the module files itself — the
// version list from tags, .info from commits, .mod from go.mod, .zip from
// the commit's tree — without a git binary. The repository is the rule's
// repo template under its url, with {n} replaced by the n-th element the
// pattern's globs captured; the rest of the module path is the module's
// directory in the repo.
//
// Rules are tried in order: a rule falls through to the next on "not
// found" (and, for goproxy rules, on 401/403/410/429 and gateway errors).
// Versioned .info/.mod/.zip files are cached forever, @v/list, @latest and
// branch/commit queries for metadata_ttl.
//
// With vanity_hosts set, ?go-get=1 lookups of a matching module path on
// one of those hosts are answered with a go-import meta tag pointing the go
// command at this proxy (go_get: mod) or at the repository (go_get: git).
package goproxy

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"html"
	"io"
	"log/slog"
	"net"
	"net/http"
	"os"
	"strconv"
	"strings"
	"sync"

	"golang.org/x/mod/module"

	"tiny-proxy/internal/auth"
	"tiny-proxy/internal/config"
	"tiny-proxy/internal/metrics"
	"tiny-proxy/internal/proxycache"
	"tiny-proxy/internal/upstream"
)

// rule is one upstream of a go vrepo.
type rule struct {
	conf   config.UpstreamConf
	client *http.Client

	mu      sync.Mutex
	sources map[string]source // by expanded repo path; git sources keep state
}

type repo struct {
	cfg   config.VRepo
	rules []*rule
}

type Handler struct {
	repos map[string]*repo
	cache *proxycache.Cache
}

// NewHandler collects go vrepos. clients[vrepo] must have one client per
// entry of vrepo.AllUpstreams(), in the same order.
func NewHandler(repos []config.VRepo, clients map[string][]*http.Client, cache *proxycache.Cache) *Handler {
	h := &Handler{repos: map[string]*repo{}, cache: cache}
	for _, rc := range repos {
		if rc.Type != config.VRepoGo {
			continue
		}
		rp := &repo{cfg: rc}
		for i, u := range rc.AllUpstreams() {
			rp.rules = append(rp.rules, &rule{conf: u, client: clients[rc.Name][i], sources: map[string]source{}})
		}
		h.repos[rc.Name] = rp
	}
	return h
}

// matched is a rule that matched a module path.
type matched struct {
	rule *rule
	idx  int
	root string // the module path prefix the pattern matched
	rest string
	repo string // expanded repo path, git-backed kinds only
}

// match returns the first pattern of r that matches mod.
func (r *rule) match(mod string) (matched, bool) {
	if len(r.conf.Modules) == 0 {
		return matched{rule: r}, true
	}
	for _, p := range r.conf.Modules {
		if caps, root, rest, ok := match(p, mod); ok {
			return matched{rule: r, root: root, rest: rest, repo: expand(r.conf.Repo, caps)}, true
		}
	}
	return matched{}, false
}

// source returns the (cached) source for an expanded repo path.
func (r *rule) source(repo string) (source, error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	if s, ok := r.sources[repo]; ok {
		return s, nil
	}
	var s source
	var err error
	switch r.conf.Kind {
	case config.GoKindGitea:
		s, err = newGiteaSource(r.client, r.conf.URL, repo, r.conf.Username, r.conf.Password)
	case config.GoKindGitlab:
		s, err = newGitlabSource(r.client, r.conf.URL, repo, r.conf.Password)
	case config.GoKindGit:
		s, err = newGitSource(r.client, r.conf.URL, repo, r.conf.Username, r.conf.Password)
	default:
		err = fmt.Errorf("kind %q has no source", r.conf.Kind)
	}
	if err != nil {
		return nil, err
	}
	r.sources[repo] = s
	return s, nil
}

// cloneURL is the repository URL a go_get: git vanity answer points at.
func (m matched) cloneURL() string {
	u := strings.TrimSuffix(m.rule.conf.URL, "/") + "/" + strings.Trim(m.repo, "/")
	if !strings.HasSuffix(u, ".git") {
		u += ".git"
	}
	return u
}

// request is a parsed GOPROXY request path.
type request struct {
	escaped string // module path as sent (case-encoded), part of the cache key
	mod     string // module path
	file    string // "list", "@latest", "<version>.info|.mod|.zip"
	version string // decoded version for .info/.mod/.zip, else ""
	ext     string // ".info", ".mod", ".zip" or ""
}

// parseRequest splits "<module>/@v/<file>" or "<module>/@latest".
func parseRequest(p string) (request, bool) {
	var rq request
	if esc, ok := strings.CutSuffix(p, "/@latest"); ok {
		rq.escaped, rq.file = esc, "@latest"
	} else {
		i := strings.LastIndex(p, "/@v/")
		if i < 0 {
			return rq, false
		}
		rq.escaped, rq.file = p[:i], p[i+len("/@v/"):]
		if rq.file != "list" {
			for _, ext := range []string{".info", ".mod", ".zip"} {
				if v, ok := strings.CutSuffix(rq.file, ext); ok && v != "" && !strings.Contains(v, "/") {
					var err error
					if rq.version, err = module.UnescapeVersion(v); err != nil {
						return rq, false
					}
					rq.ext = ext
				}
			}
			if rq.ext == "" {
				return rq, false
			}
		}
	}
	mod, err := module.UnescapePath(rq.escaped)
	if err != nil {
		return rq, false
	}
	rq.mod = mod
	return rq, true
}

// immutable reports whether the answer can never change: files of a
// canonical version. Lists, @latest and branch/commit queries can.
func (rq request) immutable() bool {
	return rq.ext != "" && isVersion(rq.version)
}

func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet && r.Method != http.MethodHead {
		http.Error(w, "read-only repository", http.StatusMethodNotAllowed)
		return
	}
	rest := strings.TrimPrefix(r.URL.Path, "/go/")
	vrepo, p, ok := strings.Cut(rest, "/")
	if !ok || !validPath(p) {
		http.Error(w, "bad path", http.StatusBadRequest)
		return
	}
	rp, exists := h.repos[vrepo]
	if !exists {
		http.Error(w, "unknown repository", http.StatusNotFound)
		return
	}
	if !auth.Check(rp.cfg.Auth, r) {
		auth.Unauthorized(w)
		return
	}
	metrics.MarkVRepo(r.Context(), vrepo)

	rq, ok := parseRequest(p)
	if !ok {
		http.Error(w, "not found: "+p, http.StatusNotFound)
		return
	}
	// .mod and .zip only exist for canonical versions; asking for a branch
	// there is a client error, not something to look up upstream.
	if (rq.ext == ".mod" || rq.ext == ".zip") && !isVersion(rq.version) {
		http.Error(w, "not found: "+p, http.StatusNotFound)
		return
	}
	var routes []matched
	for i, ru := range rp.rules {
		if m, ok := ru.match(rq.mod); ok {
			m.idx = i
			routes = append(routes, m)
		}
	}
	if len(routes) == 0 {
		http.Error(w, "no rule for module "+rq.mod, http.StatusNotFound)
		return
	}

	key := "go/" + vrepo + "/" + p
	ttl := rp.cfg.MutableTTL()
	if rq.immutable() {
		ttl = 0
	}
	h.cache.Serve(w, r, key, ttl, func(ctx context.Context) (proxycache.Result, error) {
		return rp.fetch(ctx, r, vrepo, p, rq, routes)
	})
}

// fetch walks the matching rules in order. A goproxy rule falls through on
// the statuses shouldFallthrough lists; a git-backed rule on "not found"
// and on errors talking to its forge. The last rule's answer is returned
// as-is. A 404 is remembered (not_found_ttl) only when every rule said
// "not found" — never when one was throttled or down.
func (rp *repo) fetch(ctx context.Context, r *http.Request, vrepo, p string, rq request, routes []matched) (proxycache.Result, error) {
	limited, retryAfter := false, ""
	definite := true // every rule so far answered "not found"
	for i, m := range routes {
		last := i+1 == len(routes)
		up := m.rule
		if up.conf.Kind == config.GoKindGoproxy {
			upstreamURL := strings.TrimSuffix(up.conf.URL, "/") + "/" + p
			req, err := http.NewRequestWithContext(ctx, r.Method, upstreamURL, nil)
			if err != nil {
				return proxycache.Result{}, err
			}
			resp, err := up.client.Do(req)
			if err != nil {
				slog.WarnContext(ctx, "go upstream error, trying next", "vrepo", vrepo, "path", p, "url", up.conf.URL, "err", err)
				definite = false
				continue
			}
			if resp.StatusCode == http.StatusTooManyRequests && !limited {
				limited, retryAfter = true, resp.Header.Get("Retry-After")
			}
			if shouldFallthrough(resp.StatusCode) && !last {
				resp.Body.Close()
				if resp.StatusCode != http.StatusNotFound && resp.StatusCode != http.StatusGone {
					definite = false
				}
				slog.DebugContext(ctx, "go upstream miss, trying next", "vrepo", vrepo, "path", p, "url", up.conf.URL, "status", resp.StatusCode)
				continue
			}
			if limited && (resp.StatusCode == http.StatusNotFound || resp.StatusCode == http.StatusGone) {
				resp.Body.Close()
				return proxycache.Result{Resp: rateLimitedResponse(r, p, retryAfter)}, nil
			}
			return proxycache.Result{
				Resp:           resp,
				NoStore:        !up.conf.CacheEnabled(),
				Resume:         upstream.Resumer(up.client, r.Method, upstreamURL),
				ResumeAttempts: up.conf.RetryAttempts(),
				ResumeMaxWait:  up.conf.RetryMaxWait(),
			}, nil
		}

		resp, err := build(ctx, r, m, rq)
		if errors.Is(err, errNotFound) {
			slog.DebugContext(ctx, "go rule miss, trying next", "vrepo", vrepo, "module", rq.mod, "rule", m.idx, "repo", m.repo, "err", err)
			continue
		}
		if err != nil {
			slog.WarnContext(ctx, "go rule error, trying next", "vrepo", vrepo, "module", rq.mod, "rule", m.idx, "repo", m.repo, "err", err)
			definite = false
			continue
		}
		return proxycache.Result{Resp: resp, NoStore: !up.conf.CacheEnabled()}, nil
	}
	switch {
	case limited:
		return proxycache.Result{Resp: rateLimitedResponse(r, p, retryAfter)}, nil
	case !definite:
		return proxycache.Result{Resp: statusResponse(r, http.StatusBadGateway, "upstream error: "+p+"\n", nil)}, nil
	}
	res := proxycache.Result{Resp: statusResponse(r, http.StatusNotFound, "not found: "+p+"\n", nil)}
	res.NotFoundTTL = rp.cfg.NotFoundCacheTTL()
	return res, nil
}

// build answers rq from a git-backed rule.
func build(ctx context.Context, r *http.Request, m matched, rq request) (*http.Response, error) {
	src, err := m.rule.source(m.repo)
	if err != nil {
		return nil, err
	}
	b := &builder{src: src, lay: layoutFor(rq.mod, m.root, m.rest)}
	switch {
	case rq.file == "list":
		vs, err := b.list(ctx)
		if err != nil {
			return nil, err
		}
		var body strings.Builder
		for _, v := range vs {
			body.WriteString(v + "\n")
		}
		return bodyResponse(r, "text/plain; charset=utf-8", body.String()), nil
	case rq.file == "@latest":
		in, err := b.latest(ctx)
		if err != nil {
			return nil, err
		}
		return jsonResponse(r, in)
	case rq.ext == ".info":
		in, err := b.stat(ctx, rq.version)
		if err != nil {
			return nil, err
		}
		return jsonResponse(r, in)
	case rq.ext == ".mod":
		data, err := b.mod(ctx, rq.version)
		if err != nil {
			return nil, err
		}
		return bodyResponse(r, "text/plain; charset=utf-8", string(data)), nil
	case rq.ext == ".zip":
		return zipResponse(ctx, r, b, rq.version)
	}
	return nil, errNotFound
}

// zipResponse builds the module zip into a temp file, removed once the
// response body is closed.
func zipResponse(ctx context.Context, r *http.Request, b *builder, version string) (*http.Response, error) {
	f, err := os.CreateTemp("", "tiny-proxy-go-mod-*.zip")
	if err != nil {
		return nil, err
	}
	fail := func(err error) (*http.Response, error) {
		f.Close()
		os.Remove(f.Name())
		return nil, err
	}
	if err := b.zip(ctx, version, f); err != nil {
		return fail(err)
	}
	size, err := f.Seek(0, io.SeekCurrent)
	if err != nil {
		return fail(err)
	}
	if _, err := f.Seek(0, io.SeekStart); err != nil {
		return fail(err)
	}
	resp := statusResponse(r, http.StatusOK, "", http.Header{"Content-Type": {"application/zip"}})
	resp.Body = tempBody{f}
	resp.ContentLength = size
	return resp, nil
}

type tempBody struct{ *os.File }

func (t tempBody) Close() error {
	t.File.Close()
	return os.Remove(t.Name())
}

func jsonResponse(r *http.Request, v any) (*http.Response, error) {
	data, err := json.Marshal(v)
	if err != nil {
		return nil, err
	}
	return bodyResponse(r, "application/json", string(data)+"\n"), nil
}

func bodyResponse(r *http.Request, contentType, body string) *http.Response {
	return statusResponse(r, http.StatusOK, body, http.Header{"Content-Type": {contentType}})
}

// Vanity wraps next, answering ?go-get=1 lookups for the vanity hosts of
// go vrepos; every other request goes to next.
func (h *Handler) Vanity(next http.Handler) http.Handler {
	hosts := map[string][]*repo{}
	for _, rp := range h.repos {
		for _, host := range rp.cfg.VanityHosts {
			hosts[strings.ToLower(host)] = append(hosts[strings.ToLower(host)], rp)
		}
	}
	if len(hosts) == 0 {
		return next
	}
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		host := strings.ToLower(r.Host)
		if h, _, err := net.SplitHostPort(host); err == nil {
			host = h
		}
		repos, ok := hosts[host]
		if !ok || r.URL.Query().Get("go-get") != "1" || r.Method != http.MethodGet {
			next.ServeHTTP(w, r)
			return
		}
		importPath := host + strings.TrimSuffix(r.URL.Path, "/")
		for _, rp := range repos {
			for _, ru := range rp.rules {
				m, ok := ru.match(importPath)
				if !ok || m.root == "" {
					continue // a catch-all rule names no root to announce
				}
				content := m.root + " mod " + baseURL(r) + "/go/" + rp.cfg.Name
				if ru.conf.GoGet == config.GoGetGit {
					content = m.root + " git " + m.cloneURL()
				}
				w.Header().Set("Content-Type", "text/html; charset=utf-8")
				fmt.Fprintf(w, "<!DOCTYPE html>\n<html><head><meta name=\"go-import\" content=\"%s\"></head><body>go get %s</body></html>\n",
					html.EscapeString(content), html.EscapeString(importPath))
				return
			}
		}
		http.Error(w, "no rule for "+importPath, http.StatusNotFound)
	})
}

// baseURL is this proxy's external base URL as the client reached it,
// honoring X-Forwarded-Proto from a TLS-terminating ingress.
func baseURL(r *http.Request) string {
	scheme := "http"
	if r.TLS != nil {
		scheme = "https"
	}
	if p := r.Header.Get("X-Forwarded-Proto"); p == "https" || p == "http" {
		scheme = p
	}
	return scheme + "://" + r.Host
}

// shouldFallthrough reports whether a goproxy upstream status means "try
// the next rule". It matches distro and maven, plus 410 Gone, which module
// proxies (proxy.golang.org among them) answer for modules they do not
// have. Kept per package on purpose, as there.
func shouldFallthrough(code int) bool {
	switch code {
	case http.StatusNotFound, http.StatusGone, http.StatusUnauthorized, http.StatusForbidden,
		http.StatusTooManyRequests, http.StatusBadGateway,
		http.StatusServiceUnavailable, http.StatusGatewayTimeout:
		return true
	}
	return false
}

// rateLimitedResponse is the 429 returned when throttling, not absence, is
// why the path could not be served.
func rateLimitedResponse(r *http.Request, p, retryAfter string) *http.Response {
	header := http.Header{}
	if retryAfter != "" {
		header.Set("Retry-After", retryAfter)
	}
	return statusResponse(r, http.StatusTooManyRequests, "rate limited by upstream: "+p+"\n", header)
}

// statusResponse builds a synthetic response for proxycache to relay.
func statusResponse(r *http.Request, code int, body string, header http.Header) *http.Response {
	if header == nil {
		header = http.Header{}
	}
	if header.Get("Content-Type") == "" {
		header.Set("Content-Type", "text/plain; charset=utf-8")
	}
	return &http.Response{
		Status:        strconv.Itoa(code) + " " + http.StatusText(code),
		StatusCode:    code,
		Proto:         "HTTP/1.1",
		ProtoMajor:    1,
		ProtoMinor:    1,
		Header:        header,
		Body:          io.NopCloser(strings.NewReader(body)),
		ContentLength: int64(len(body)),
		Request:       r,
	}
}

// validPath rejects traversal and malformed paths; cache keys become
// localfs paths.
func validPath(p string) bool {
	if p == "" || strings.HasPrefix(p, "/") || strings.Contains(p, "\\") {
		return false
	}
	for _, seg := range strings.Split(p, "/") {
		if seg == "" || seg == "." || seg == ".." {
			return false
		}
	}
	return true
}