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 cache hit/miss/error counters. Nil is fine; the
	// registry's Inc methods are nil-receiver safe.
	Metrics *metrics.Registry

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

func New(store storage.Store) *Cache {
	return &Cache{Store: store, locks: map[string]*sync.Mutex{}}
}

// 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
}

// 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)) {
	if c.serveHit(w, r, key, ttl) {
		c.Metrics.IncHit()
		return
	}

	lock := c.keyLock(key)
	lock.Lock()
	defer lock.Unlock()
	// Double-check: another request may have populated the key while we waited.
	if c.serveHit(w, r, key, ttl) {
		c.Metrics.IncHit()
		return
	}

	c.Metrics.IncMiss()
	res, err := fetch(r.Context())
	if err == nil && res.Resp == nil {
		err = errors.New("fetch returned no response")
	}
	if err != nil {
		c.Metrics.IncError()
		slog.ErrorContext(r.Context(), "upstream fetch failed", "key", key, "err", err)
		http.Error(w, "upstream error", http.StatusBadGateway)
		return
	}
	resp := res.Resp
	defer resp.Body.Close()

	if resp.StatusCode == http.StatusOK && r.Method == http.MethodGet && !res.NoStore {
		c.serveAndStore(w, r, key, res)
		return
	}
	proxyResponse(w, r, resp)
}

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) {
	ctx := r.Context()
	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
	}()

	writeMetaHeaders(w, &meta)
	// 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 {
		n, readErr := body.Read(buf)
		if n > 0 {
			chunk := buf[:n]
			if cacheLive {
				if err := writeAll(pw, chunk); err != nil {
					abort(err)
				}
			}
			if err := writeAll(w, 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 && 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
	switch {
	case failErr != nil:
		c.Metrics.IncError()
		c.Metrics.IncTransferFailure()
		slog.WarnContext(ctx, "upstream transfer failed, not caching",
			"key", key, "written", written, "declared", resp.ContentLength, "err", failErr)
	case putErr != nil:
		c.Metrics.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) {
	if body, ok := c.fetchHit(ctx, key, ttl); ok {
		c.Metrics.IncHit()
		return body, nil
	}

	lock := c.keyLock(key)
	lock.Lock()
	defer lock.Unlock()
	// Double-check: another goroutine may have populated the key.
	if body, ok := c.fetchHit(ctx, key, ttl); ok {
		c.Metrics.IncHit()
		return body, nil
	}

	c.Metrics.IncMiss()
	res, err := fetch(ctx)
	if err == nil && res.Resp == nil {
		err = errors.New("fetch returned no response")
	}
	if err != nil {
		c.Metrics.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 {
		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 {
		c.Metrics.IncError()
		c.Metrics.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 {
		c.Metrics.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)
}