tiny-proxy/internal/proxycache/proxycache.go

// Package proxycache implements the shared read-through cache used by all
// proxy handlers (docker, maven, github, gitea, debian, pacman): serve from
// storage when fresh, otherwise fetch from the upstream, tee the body into
// storage, and stream to the client.
package proxycache

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"strconv"
	"sync"
	"time"

	"tiny-proxy/internal/metrics"
	"tiny-proxy/internal/storage"
)

// replayHeaders are response headers captured on store and replayed on hits.
var replayHeaders = []string{"Docker-Content-Digest", "Etag", "Last-Modified"}

type Cache struct {
	Store storage.Store
	// Metrics records the process-wide counters (resume outcomes). The
	// per-vrepo cache counters come from the request context instead, since
	// only the handler knows which repo a key belongs to. Nil is fine; the
	// Inc methods are nil-receiver safe.
	Metrics *metrics.Registry

	mu    sync.Mutex
	locks map[string]*sync.Mutex

	// misses remembers keys the upstreams definitively do not have, until
	// the time stored for each. Memory only: it is short-lived by design,
	// and losing it on restart costs one more upstream round each.
	missMu    sync.Mutex
	misses    map[string]time.Time
	lastSweep time.Time

	// headSlots bounds the background downloads HEAD misses start, so a
	// client HEAD-ing hundreds of uncached artifacts cannot fan out that
	// many full downloads; past the bound a HEAD is answered by a HEAD.
	headSlots chan struct{}
}

// maxHeadStores is how many HEAD misses may be downloading into the store
// at once.
const maxHeadStores = 16

// MissTTLHeader carries, on a 404 answered from a remembered miss, the
// seconds that memory has left. A tiny-proxy chained in front remembers the
// 404 no longer than that, so misses do not stack up along a chain.
const MissTTLHeader = "X-Tinyproxy-Miss-Ttl"

// maxMisses caps the remembered misses, so a client probing random paths
// cannot grow the map without bound.
const maxMisses = 100_000

func New(store storage.Store) *Cache {
	return &Cache{
		Store:     store,
		locks:     map[string]*sync.Mutex{},
		misses:    map[string]time.Time{},
		headSlots: make(chan struct{}, maxHeadStores),
	}
}

// missed reports whether key is a remembered miss that has not expired, and
// how long it has left.
func (c *Cache) missed(key string) (time.Duration, bool) {
	c.missMu.Lock()
	defer c.missMu.Unlock()
	until, ok := c.misses[key]
	if !ok {
		return 0, false
	}
	left := time.Until(until)
	if left <= 0 {
		delete(c.misses, key)
		return 0, false
	}
	return left, true
}

// rememberMiss records that key is not found upstream, for ttl.
func (c *Cache) rememberMiss(key string, ttl time.Duration) {
	c.missMu.Lock()
	defer c.missMu.Unlock()
	now := time.Now()
	if len(c.misses) >= maxMisses {
		// Sweep expired entries at most once a minute: a full map under a
		// flood of new misses must not cost a 100k scan per insert.
		if now.Sub(c.lastSweep) > time.Minute {
			c.lastSweep = now
			for k, until := range c.misses {
				if now.After(until) {
					delete(c.misses, k)
				}
			}
		}
		// Still full: make room by dropping an arbitrary entry, which at
		// worst costs that path one more upstream walk.
		for k := range c.misses {
			if len(c.misses) < maxMisses {
				break
			}
			delete(c.misses, k)
		}
	}
	c.misses[key] = now.Add(ttl)
}

// keyLock serializes upstream fetches per cache key (single-flight).
func (c *Cache) keyLock(key string) *sync.Mutex {
	c.mu.Lock()
	defer c.mu.Unlock()
	l, ok := c.locks[key]
	if !ok {
		l = &sync.Mutex{}
		c.locks[key] = l
	}
	return l
}

// Result is what a fetch callback returns: the upstream response plus the
// caching policy for it. The zero NoStore caches as usual.
type Result struct {
	Resp *http.Response
	// NoStore streams this response to the client without writing it to
	// storage. Handlers set it when the upstream that answered has
	// cache: false — normally because that upstream is itself a caching
	// layer, so storing the bytes here would just duplicate them.
	NoStore bool
	// Resume re-requests this object from the upstream that answered, with
	// the Range/If-Range headers given, so a transfer that dies mid-body is
	// continued instead of failing the client. Handlers build it with
	// upstream.Resumer so the resume can only ever hit the same upstream
	// and URL that served the first byte. Nil disables resuming.
	Resume func(ctx context.Context, hdr http.Header) (*http.Response, error)
	// ResumeAttempts caps how many times one transfer may be resumed, and
	// ResumeMaxWait caps a single wait between those attempts. Handlers
	// pass the answering upstream's retry budget.
	ResumeAttempts int
	ResumeMaxWait  time.Duration
	// NotFoundTTL, when positive, makes a 404 in Resp a remembered miss:
	// for that long the key is answered 404 without calling fetch. Handlers
	// set it only when every upstream really said "not found" — never when
	// one was rate limited or down, or an outage would be cached as absence.
	NotFoundTTL time.Duration
	// BodyForHead says Resp answers a GET the handler sent although the
	// client asked HEAD, so that a HEAD miss still lands in the cache: the
	// body is stored and only the headers go to the client.
	BodyForHead bool
}

// notFound answers a remembered miss that has left to go.
func notFound(w http.ResponseWriter, left time.Duration) {
	w.Header().Set(MissTTLHeader, strconv.Itoa(int((left+time.Second-1)/time.Second)))
	http.Error(w, "not found", http.StatusNotFound)
}

// headStoreKey marks the context of a HEAD fetch that may answer with a GET
// to be stored (see HeadStorable).
type headStoreKey struct{}

// HeadStorable reports whether the fetch running under ctx serves a HEAD
// that may be answered by a GET whose body is stored in the background
// (Result.BodyForHead). When false — a plain Serve, or too many such
// downloads already running — the fetch should send the HEAD as is.
func HeadStorable(ctx context.Context) bool {
	ok, _ := ctx.Value(headStoreKey{}).(bool)
	return ok
}

// ServeStoringHead is Serve for handlers whose fetch may answer a HEAD miss
// with a GET to be stored (see HeadStorable and Result.BodyForHead), so a
// HEAD does not reach the upstreams again and again. Such a fetch runs
// detached from the request, since its download outlives it.
func (c *Cache) ServeStoringHead(w http.ResponseWriter, r *http.Request, key string, ttl time.Duration, fetch func(ctx context.Context) (Result, error)) {
	c.serve(w, r, key, ttl, fetch, true)
}

// Serve answers r from the cache, fetching from upstream on a miss.
// fetch performs the upstream request with the same method as r and returns
// its response; ttl controls how long the cached object counts as fresh
// (0 = immutable). Non-2xx upstream responses, and responses the fetch
// marked NoStore, are proxied without caching.
func (c *Cache) Serve(w http.ResponseWriter, r *http.Request, key string, ttl time.Duration, fetch func(ctx context.Context) (Result, error)) {
	c.serve(w, r, key, ttl, fetch, false)
}

func (c *Cache) serve(w http.ResponseWriter, r *http.Request, key string, ttl time.Duration, fetch func(ctx context.Context) (Result, error), storeHead bool) {
	// The vrepo serving this request was recorded by its handler, so the
	// cache outcome lands on that repo's series rather than one global
	// hit rate that averages a warm mirror with a cold one.
	stats := metrics.VRepoFromContext(r.Context())
	if c.serveHit(w, r, key, ttl) {
		stats.IncHit()
		return
	}
	if left, ok := c.missed(key); ok {
		c.Metrics.IncRememberedMiss()
		notFound(w, left)
		return
	}

	lock := c.keyLock(key)
	lock.Lock()
	// A HEAD miss may be answered by storing a GET in the background, which
	// then owns the lock and the context until the body has landed.
	background := false
	defer func() {
		if !background {
			lock.Unlock()
		}
	}()
	// Double-check: another request may have populated the key — or found
	// it missing everywhere — while we waited.
	if c.serveHit(w, r, key, ttl) {
		stats.IncHit()
		return
	}
	if left, ok := c.missed(key); ok {
		c.Metrics.IncRememberedMiss()
		notFound(w, left)
		return
	}

	stats.IncMiss()
	// A HEAD whose fetch may turn into a download to store (BodyForHead)
	// takes a slot, and runs detached from the request: the client is done
	// the moment it has the headers, and the download outlives it.
	ctx, cancel := r.Context(), context.CancelFunc(func() {})
	slot := false
	if storeHead && r.Method == http.MethodHead {
		select {
		case c.headSlots <- struct{}{}:
			slot = true
			ctx, cancel = context.WithTimeout(context.WithoutCancel(ctx), headStoreTimeout)
			ctx = context.WithValue(ctx, headStoreKey{}, true)
		default:
		}
	}
	release := func() {
		cancel()
		if slot {
			<-c.headSlots
		}
	}
	defer func() {
		if !background {
			release()
		}
	}()
	res, err := fetch(ctx)
	if err == nil && res.Resp == nil {
		err = errors.New("fetch returned no response")
	}
	if err != nil {
		stats.IncError()
		slog.ErrorContext(ctx, "upstream fetch failed", "key", key, "err", err)
		http.Error(w, "upstream error", http.StatusBadGateway)
		return
	}
	resp := res.Resp

	if resp.StatusCode == http.StatusNotFound && res.NotFoundTTL > 0 {
		c.rememberMiss(key, res.NotFoundTTL)
	}
	if resp.StatusCode == http.StatusOK && !res.NoStore {
		switch {
		case r.Method == http.MethodGet:
			defer resp.Body.Close()
			c.serveAndStore(w, r, key, res)
			return
		case r.Method == http.MethodHead && res.BodyForHead && slot:
			// The client gets its headers and its connection back now;
			// the body goes to storage behind it, still under the key
			// lock so a request for the same key waits for it.
			meta := metaFromResponse(resp)
			writeMetaHeaders(w, &meta)
			background = true
			go func() {
				defer lock.Unlock()
				defer release()
				defer resp.Body.Close()
				// Out here no handler recovers a panic: it would take
				// the whole process down, not one request.
				defer func() {
					if p := recover(); p != nil {
						slog.ErrorContext(ctx, "background store panicked", "key", key, "panic", p)
					}
				}()
				c.store(ctx, key, res, nil)
			}()
			return
		}
	}
	defer resp.Body.Close()
	proxyResponse(w, r, resp)
}

// headStoreTimeout bounds the background download a HEAD miss starts.
const headStoreTimeout = 10 * time.Minute

func (c *Cache) serveHit(w http.ResponseWriter, r *http.Request, key string, ttl time.Duration) bool {
	meta, err := c.Store.Stat(r.Context(), key)
	if err != nil {
		return false
	}
	if meta.Expired(ttl) {
		return false
	}
	rc, meta, err := c.Store.Get(r.Context(), key)
	if err != nil {
		return false
	}
	defer rc.Close()
	writeMetaHeaders(w, meta)
	if r.Method == http.MethodHead {
		return true
	}
	if _, err := io.Copy(w, rc); err != nil {
		slog.WarnContext(r.Context(), "stream from cache failed", "key", key, "err", err)
	}
	return true
}

// serveAndStore streams a 200 OK upstream body to the client while tee-ing
// it into storage.
//
// The cache write is committed only when the whole object arrived. A
// transfer that dies mid-body — a reset connection, a truncated TLS stream, a
// server that closes early — aborts the pipe with an error so the store
// discards its temp file. Committing it instead is what poisons a cache
// permanently: immutable artifacts never expire, so a truncated .deb would be
// served as the real thing forever.
func (c *Cache) serveAndStore(w http.ResponseWriter, r *http.Request, key string, res Result) {
	meta := metaFromResponse(res.Resp)
	writeMetaHeaders(w, &meta)
	c.store(r.Context(), key, res, w)
}

// store tees a 200 OK upstream body into storage and, when client is not
// nil, to the client. With no client, it stops as soon as the store fails:
// there is nobody left to read the rest.
func (c *Cache) store(ctx context.Context, key string, res Result, client io.Writer) {
	resp := res.Resp
	meta := metaFromResponse(resp)

	body := c.upstreamBody(ctx, key, res)
	defer body.Close()

	pr, pw := io.Pipe()
	putDone := make(chan error, 1)
	go func() {
		// Size 0 so the store records the bytes that actually landed
		// rather than the length the upstream advertised: the two disagree
		// exactly when something went wrong, and a sidecar claiming more
		// than the file holds makes every later hit fail.
		putMeta := meta
		putMeta.Size = 0
		err := c.Store.Put(ctx, key, pr, putMeta)
		if err != nil {
			// Unblock the tee: if Put fails before draining the pipe
			// (e.g. unwritable storage), the pipe writer would block
			// forever and the client would hang.
			pr.CloseWithError(err)
		}
		putDone <- err
	}()

	// Hand-rolled tee: read upstream, write to the cache pipe (until it
	// fails — caching is best-effort) and always to the client.
	// io.TeeReader would drop a chunk on pipe error because it consumes
	// bytes from the upstream before the failing write reports n=0.
	buf := make([]byte, 32*1024)
	var written int64
	var copyErr, failErr error
	cacheLive := true
	abort := func(err error) {
		if cacheLive {
			cacheLive = false
			pw.CloseWithError(err)
		}
	}
	for client != nil || cacheLive {
		n, readErr := body.Read(buf)
		if n > 0 {
			chunk := buf[:n]
			if cacheLive {
				if err := writeAll(pw, chunk); err != nil {
					abort(err)
				}
			}
			if client != nil {
				if err := writeAll(client, chunk); err != nil {
					copyErr = err
					abort(err)
					break
				}
			}
			written += int64(n)
		}
		if readErr != nil {
			if !errors.Is(readErr, io.EOF) {
				failErr = readErr
			}
			break
		}
	}
	// A clean close that still delivered fewer bytes than advertised is a
	// failed transfer too, and the one the client is most likely to mistake
	// for a good file.
	if failErr == nil && copyErr == nil && (client != nil || cacheLive) && resp.ContentLength >= 0 && written != resp.ContentLength {
		failErr = fmt.Errorf("upstream sent %d bytes, declared %d", written, resp.ContentLength)
	}
	if failErr != nil {
		abort(failErr)
	}
	if cacheLive {
		pw.Close()
	}

	putErr := <-putDone
	stats := metrics.VRepoFromContext(ctx)
	switch {
	case failErr != nil:
		stats.IncError()
		stats.IncTransferFailure()
		slog.WarnContext(ctx, "upstream transfer failed, not caching",
			"key", key, "written", written, "declared", resp.ContentLength, "err", failErr)
	case putErr != nil:
		stats.IncError()
		slog.ErrorContext(ctx, "cache store failed", "key", key, "err", putErr)
	}
	if copyErr != nil {
		slog.WarnContext(ctx, "stream to client failed", "key", key, "err", copyErr)
	}
}

// upstreamBody returns the reader serveAndStore should tee from: the raw
// body, or one that transparently resumes when the transfer drops, if the
// handler supplied a way to re-request from the upstream that answered.
func (c *Cache) upstreamBody(ctx context.Context, key string, res Result) io.ReadCloser {
	if res.Resume == nil || res.ResumeAttempts <= 0 {
		return res.Resp.Body
	}
	return &resumeReader{
		body:     res.Resp.Body,
		resume:   res.Resume,
		ctx:      ctx,
		ifRange:  validator(res.Resp),
		attempts: res.ResumeAttempts,
		maxWait:  res.ResumeMaxWait,
		key:      key,
		metrics:  c.Metrics,
	}
}

// writeAll writes the whole chunk, treating a short write as an error.
// io.Writer permits n < len(p) with a nil error; silently dropping the
// remainder would put a hole in a cached object.
func writeAll(w io.Writer, p []byte) error {
	n, err := w.Write(p)
	if err != nil {
		return err
	}
	if n != len(p) {
		return io.ErrShortWrite
	}
	return nil
}

// proxyResponse copies an upstream response to the client unchanged
// (used for non-200 statuses, HEAD misses and partial responses).
func proxyResponse(w http.ResponseWriter, r *http.Request, resp *http.Response) {
	copyHeader(w.Header(), resp.Header)
	w.WriteHeader(resp.StatusCode)
	if r.Method != http.MethodHead {
		io.Copy(w, resp.Body)
	}
}

func metaFromResponse(resp *http.Response) storage.ObjectMeta {
	meta := storage.ObjectMeta{
		ContentType: resp.Header.Get("Content-Type"),
		Size:        resp.ContentLength,
		StoredAt:    time.Now(),
		Headers:     map[string]string{},
	}
	// Chunked responses report -1; leave 0 so the store records the
	// actual byte count after the tee.
	if meta.Size < 0 {
		meta.Size = 0
	}
	for _, h := range replayHeaders {
		if v := resp.Header.Get(h); v != "" {
			meta.Headers[h] = v
		}
	}
	return meta
}

func writeMetaHeaders(w http.ResponseWriter, meta *storage.ObjectMeta) {
	h := w.Header()
	if meta.ContentType != "" {
		h.Set("Content-Type", meta.ContentType)
	}
	if meta.Size > 0 {
		h.Set("Content-Length", strconv.FormatInt(meta.Size, 10))
	}
	for k, v := range meta.Headers {
		h.Set(k, v)
	}
	w.WriteHeader(http.StatusOK)
}

func copyHeader(dst, src http.Header) {
	for k, vs := range src {
		for _, v := range vs {
			dst.Add(k, v)
		}
	}
}

// IsNotFound reports whether err is a storage cache miss.
func IsNotFound(err error) bool {
	return errors.Is(err, storage.ErrNotFound)
}

// Fetch returns the cached body for key, fetching from upstream on a miss or
// when the cached object is older than ttl (0 = immutable). It is the
// in-process counterpart of Serve, for values a handler needs to inspect
// (e.g. alias resolution) rather than stream to a client. Only 200 OK
// bodies are cached, and only when the fetch did not mark them NoStore;
// other statuses return an *UpstreamError.
func (c *Cache) Fetch(ctx context.Context, key string, ttl time.Duration, fetch func(ctx context.Context) (Result, error)) ([]byte, error) {
	stats := metrics.VRepoFromContext(ctx)
	if body, ok := c.fetchHit(ctx, key, ttl); ok {
		stats.IncHit()
		return body, nil
	}
	if _, ok := c.missed(key); ok {
		c.Metrics.IncRememberedMiss()
		return nil, &UpstreamError{StatusCode: http.StatusNotFound}
	}

	lock := c.keyLock(key)
	lock.Lock()
	defer lock.Unlock()
	// Double-check: another goroutine may have populated the key, or found
	// it missing.
	if body, ok := c.fetchHit(ctx, key, ttl); ok {
		stats.IncHit()
		return body, nil
	}
	if _, ok := c.missed(key); ok {
		c.Metrics.IncRememberedMiss()
		return nil, &UpstreamError{StatusCode: http.StatusNotFound}
	}

	stats.IncMiss()
	res, err := fetch(ctx)
	if err == nil && res.Resp == nil {
		err = errors.New("fetch returned no response")
	}
	if err != nil {
		stats.IncError()
		return nil, err
	}
	resp := res.Resp
	defer resp.Body.Close()
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	if resp.StatusCode != http.StatusOK {
		if resp.StatusCode == http.StatusNotFound && res.NotFoundTTL > 0 {
			c.rememberMiss(key, res.NotFoundTTL)
		}
		return nil, &UpstreamError{StatusCode: resp.StatusCode, Body: body}
	}
	// io.ReadAll already reports a torn connection; this catches the
	// quieter case of a clean close after a short body.
	if resp.ContentLength >= 0 && int64(len(body)) != resp.ContentLength {
		stats.IncError()
		stats.IncTransferFailure()
		return nil, fmt.Errorf("upstream sent %d bytes, declared %d", len(body), resp.ContentLength)
	}
	if res.NoStore {
		return body, nil
	}
	meta := metaFromResponse(resp)
	meta.Size = int64(len(body))
	if err := c.Store.Put(ctx, key, bytes.NewReader(body), meta); err != nil {
		stats.IncError()
		slog.ErrorContext(ctx, "cache store failed", "key", key, "err", err)
	}
	return body, nil
}

func (c *Cache) fetchHit(ctx context.Context, key string, ttl time.Duration) ([]byte, bool) {
	meta, err := c.Store.Stat(ctx, key)
	if err != nil || meta.Expired(ttl) {
		return nil, false
	}
	rc, _, err := c.Store.Get(ctx, key)
	if err != nil {
		return nil, false
	}
	defer rc.Close()
	body, err := io.ReadAll(rc)
	if err != nil {
		return nil, false
	}
	return body, true
}

// UpstreamError reports a non-200 upstream response from Fetch.
type UpstreamError struct {
	StatusCode int
	Body       []byte
}

func (e *UpstreamError) Error() string {
	return fmt.Sprintf("upstream status %d", e.StatusCode)
}