tiny-proxy/internal/docker/handler_test.go

package docker

import (
	"net/http"
	"net/http/httptest"
	"strconv"
	"strings"
	"sync/atomic"
	"testing"

	"tiny-proxy/internal/config"
	"tiny-proxy/internal/proxycache"
	"tiny-proxy/internal/storage"
)

func TestParsePath(t *testing.T) {
	vrepo, name, kind, ref, ok := parsePath("hub/library/nginx/manifests/1.27")
	if !ok {
		t.Fatal("parse failed")
	}
	if vrepo != "hub" || name != "library/nginx" || kind != "manifests" || ref != "1.27" {
		t.Errorf("got vrepo=%q name=%q kind=%q ref=%q", vrepo, name, kind, ref)
	}

	vrepo, name, kind, ref, ok = parsePath("hub/library/nginx/blobs/sha256:abc")
	if !ok || kind != "blobs" || ref != "sha256:abc" {
		t.Errorf("blobs: ok=%v name=%q kind=%q ref=%q", ok, name, kind, ref)
	}

	for _, bad := range []string{"hub/nginx", "hub/a/b/c", "hub/library/nginx/tags/1.27", ""} {
		if _, _, _, _, ok := parsePath(bad); ok {
			t.Errorf("parsePath(%q) should fail", bad)
		}
	}
}

func TestNormalizeName(t *testing.T) {
	if got := normalizeName("https://registry-1.docker.io", "nginx"); got != "library/nginx" {
		t.Errorf("docker.io single name = %q", got)
	}
	if got := normalizeName("https://registry-1.docker.io", "org/nginx"); got != "org/nginx" {
		t.Errorf("docker.io two-part name = %q", got)
	}
	if got := normalizeName("https://ghcr.io", "nginx"); got != "nginx" {
		t.Errorf("ghcr single name = %q", got)
	}
	// Only the host counts: a chained tiny-proxy normalizes on its own side,
	// even when its vrepo path happens to mention docker.io.
	if got := normalizeName("https://repo.example/v2/docker.io", "nginx"); got != "nginx" {
		t.Errorf("chained single name = %q", got)
	}
	if got := normalizeName("https://mirror.docker.com", "nginx"); got != "library/nginx" {
		t.Errorf("docker.com single name = %q", got)
	}
}

func TestObjectURL(t *testing.T) {
	cases := []struct{ base, want string }{
		{"https://registry-1.docker.io", "https://registry-1.docker.io/v2/library/nginx/manifests/1.27"},
		{"https://registry-1.docker.io/", "https://registry-1.docker.io/v2/library/nginx/manifests/1.27"},
		// A /v2/ path already names the repository root.
		{"https://repo.example/v2/dockerhub", "https://repo.example/v2/dockerhub/library/nginx/manifests/1.27"},
		{"https://repo.example/v2/dockerhub/", "https://repo.example/v2/dockerhub/library/nginx/manifests/1.27"},
		{"https://repo.example/v2", "https://repo.example/v2/library/nginx/manifests/1.27"},
		// Any other path is a registry base that still needs /v2/.
		{"https://repo.example/registry", "https://repo.example/registry/v2/library/nginx/manifests/1.27"},
		{"https://repo.example/v2x", "https://repo.example/v2x/v2/library/nginx/manifests/1.27"},
	}
	for _, c := range cases {
		if got := objectURL(c.base, "library/nginx", "manifests", "1.27"); got != c.want {
			t.Errorf("objectURL(%q) = %q, want %q", c.base, got, c.want)
		}
	}
}

// newTestHandler wires one docker vrepo to a fake registry and returns the
// handler, the registry's request counter and the backing store.
func newTestHandler(t *testing.T, up config.UpstreamConf) (*Handler, *atomic.Int32, storage.Store) {
	t.Helper()
	var hits atomic.Int32
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		hits.Add(1)
		w.Header().Set("Content-Type", "application/octet-stream")
		w.Write([]byte("blob-bytes"))
	}))
	t.Cleanup(srv.Close)
	up.URL = srv.URL

	store, err := storage.NewLocalFS(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	rc := config.VRepo{
		Name:     "hub",
		Type:     config.VRepoDocker,
		Upstream: up,
		Auth:     config.AuthConfig{Mode: config.AuthPublic},
	}
	clients := map[string][]*http.Client{"hub": {srv.Client()}}
	return NewHandler([]config.VRepo{rc}, clients, proxycache.New(store)), &hits, store
}

func getBlob(t *testing.T, h *Handler) (int, string) {
	t.Helper()
	rec := httptest.NewRecorder()
	req := httptest.NewRequest(http.MethodGet, "/v2/hub/library/nginx/blobs/sha256:abc", nil)
	h.ServeHTTP(rec, req)
	return rec.Code, rec.Body.String()
}

func TestBlobCachedByDefault(t *testing.T) {
	h, hits, store := newTestHandler(t, config.UpstreamConf{})
	for i := 0; i < 2; i++ {
		if code, body := getBlob(t, h); code != http.StatusOK || body != "blob-bytes" {
			t.Fatalf("request %d: code = %d body = %q", i, code, body)
		}
	}
	if got := hits.Load(); got != 1 {
		t.Errorf("upstream hits = %d, want 1 (blob is immutable)", got)
	}
	if _, err := store.Stat(t.Context(), "docker/hub/blobs/sha256:abc"); err != nil {
		t.Errorf("blob not cached: %v", err)
	}
}

func TestUpstreamCacheDisabled(t *testing.T) {
	no := false
	h, hits, store := newTestHandler(t, config.UpstreamConf{Cache: &no})
	for i := 0; i < 2; i++ {
		if code, body := getBlob(t, h); code != http.StatusOK || body != "blob-bytes" {
			t.Fatalf("request %d: code = %d body = %q", i, code, body)
		}
	}
	if got := hits.Load(); got != 2 {
		t.Errorf("upstream hits = %d, want 2 (cache: false)", got)
	}
	if _, err := store.Stat(t.Context(), "docker/hub/blobs/sha256:abc"); err != storage.ErrNotFound {
		t.Error("blob was stored despite cache: false")
	}
}

// stubRegistry counts requests, records the last path it was asked for and
// serves whatever its handler says.
type stubRegistry struct {
	srv  *httptest.Server
	hits atomic.Int32
	path atomic.Value // string
}

func newStubRegistry(t *testing.T, h http.HandlerFunc) *stubRegistry {
	t.Helper()
	s := &stubRegistry{}
	s.srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		s.hits.Add(1)
		s.path.Store(r.URL.Path)
		h(w, r)
	}))
	t.Cleanup(s.srv.Close)
	return s
}

func (s *stubRegistry) lastPath() string {
	p, _ := s.path.Load().(string)
	return p
}

// newMultiHandler wires one docker vrepo to an ordered upstream list.
func newMultiHandler(t *testing.T, ups ...config.UpstreamConf) *Handler {
	t.Helper()
	store, err := storage.NewLocalFS(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	rc := config.VRepo{
		Name:      "hub",
		Type:      config.VRepoDocker,
		Upstreams: ups,
		Auth:      config.AuthConfig{Mode: config.AuthPublic},
	}
	clients := map[string][]*http.Client{"hub": {}}
	for range ups {
		clients["hub"] = append(clients["hub"], http.DefaultClient)
	}
	return NewHandler([]config.VRepo{rc}, clients, proxycache.New(store))
}

func getManifest(t *testing.T, h *Handler, name string) (int, string) {
	t.Helper()
	rec := httptest.NewRecorder()
	h.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "/v2/hub/"+name+"/manifests/1.27", nil))
	return rec.Code, rec.Body.String()
}

func serving(body string) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) { w.Write([]byte(body)) }
}

func failing(code int) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) { http.Error(w, http.StatusText(code), code) }
}

func TestFallbackToSecondUpstream(t *testing.T) {
	for _, code := range []int{
		http.StatusNotFound, http.StatusUnauthorized, http.StatusForbidden,
		http.StatusTooManyRequests, http.StatusBadGateway,
		http.StatusServiceUnavailable, http.StatusGatewayTimeout,
	} {
		t.Run(strconv.Itoa(code), func(t *testing.T) {
			first := newStubRegistry(t, failing(code))
			second := newStubRegistry(t, serving("manifest"))
			h := newMultiHandler(t,
				config.UpstreamConf{URL: first.srv.URL},
				config.UpstreamConf{URL: second.srv.URL})

			if got, body := getManifest(t, h, "org/app"); got != http.StatusOK || body != "manifest" {
				t.Fatalf("got %d %q, want 200 manifest", got, body)
			}
			if first.hits.Load() != 1 || second.hits.Load() != 1 {
				t.Errorf("hits: first=%d second=%d, want 1 and 1", first.hits.Load(), second.hits.Load())
			}
		})
	}
}

func TestTransportErrorFallsThrough(t *testing.T) {
	dead := newStubRegistry(t, serving("never"))
	dead.srv.Close() // connections now fail
	second := newStubRegistry(t, serving("manifest"))
	h := newMultiHandler(t,
		config.UpstreamConf{URL: dead.srv.URL},
		config.UpstreamConf{URL: second.srv.URL})

	if code, body := getManifest(t, h, "org/app"); code != http.StatusOK || body != "manifest" {
		t.Fatalf("got %d %q, want 200 manifest", code, body)
	}
}

// With nothing answering at all, the failure is a gateway error, not a
// missing image.
func TestAllUnreachableReturns502(t *testing.T) {
	dead := newStubRegistry(t, serving("never"))
	dead.srv.Close()
	h := newMultiHandler(t, config.UpstreamConf{URL: dead.srv.URL})

	if code, _ := getManifest(t, h, "org/app"); code != http.StatusBadGateway {
		t.Fatalf("got %d, want 502", code)
	}
}

func TestServerErrorDoesNotFallThrough(t *testing.T) {
	broken := newStubRegistry(t, failing(http.StatusInternalServerError))
	second := newStubRegistry(t, serving("manifest"))
	h := newMultiHandler(t,
		config.UpstreamConf{URL: broken.srv.URL},
		config.UpstreamConf{URL: second.srv.URL})

	if code, _ := getManifest(t, h, "org/app"); code != http.StatusInternalServerError {
		t.Fatalf("got %d, want 500 (not a fall-through status)", code)
	}
	if second.hits.Load() != 0 {
		t.Errorf("second upstream consulted %d times after 500, want 0", second.hits.Load())
	}
}

func TestLastUpstreamGatewayErrorReturnedAsIs(t *testing.T) {
	missing := newStubRegistry(t, http.NotFound)
	down := newStubRegistry(t, failing(http.StatusServiceUnavailable))
	h := newMultiHandler(t,
		config.UpstreamConf{URL: missing.srv.URL},
		config.UpstreamConf{URL: down.srv.URL})

	if code, _ := getManifest(t, h, "org/app"); code != http.StatusServiceUnavailable {
		t.Fatalf("got %d, want 503 from last upstream", code)
	}
}

func TestAllMissReturns404(t *testing.T) {
	first := newStubRegistry(t, http.NotFound)
	second := newStubRegistry(t, http.NotFound)
	h := newMultiHandler(t,
		config.UpstreamConf{URL: first.srv.URL},
		config.UpstreamConf{URL: second.srv.URL})

	if code, _ := getManifest(t, h, "org/app"); code != http.StatusNotFound {
		t.Fatalf("got %d, want 404", code)
	}
}

// A 404 after a rate-limited upstream must not be reported as "missing".
func TestRateLimitedThen404Returns429(t *testing.T) {
	limited := newStubRegistry(t, func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Retry-After", "30")
		http.Error(w, "slow down", http.StatusTooManyRequests)
	})
	missing := newStubRegistry(t, http.NotFound)
	h := newMultiHandler(t,
		config.UpstreamConf{URL: limited.srv.URL},
		config.UpstreamConf{URL: missing.srv.URL})

	rec := httptest.NewRecorder()
	h.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "/v2/hub/org/app/manifests/1.27", nil))
	if rec.Code != http.StatusTooManyRequests {
		t.Fatalf("got %d, want 429", rec.Code)
	}
	if got := rec.Header().Get("Retry-After"); got != "30" {
		t.Errorf("Retry-After = %q, want 30 carried over from upstream", got)
	}
}

// TestChainedUpstream covers the tiny-proxy-to-tiny-proxy form: a /v2/ URL
// is the repository root, so the image path is appended without another
// /v2/ and without library/ (the chained proxy adds that itself;
// TestNormalizeName covers the Docker Hub side).
func TestChainedUpstream(t *testing.T) {
	chained := newStubRegistry(t, serving("from-chain"))
	h := newMultiHandler(t, config.UpstreamConf{URL: chained.srv.URL + "/v2/dockerhub"})

	if code, body := getManifest(t, h, "nginx"); code != http.StatusOK || body != "from-chain" {
		t.Fatalf("got %d %q, want 200 from-chain", code, body)
	}
	if got, want := chained.lastPath(), "/v2/dockerhub/nginx/manifests/1.27"; got != want {
		t.Errorf("chained upstream path = %q, want %q", got, want)
	}
}

// TestAcceptForwarded: registries pick the manifest media type from Accept,
// so every upstream in the walk must see the client's header.
func TestAcceptForwarded(t *testing.T) {
	const accept = "application/vnd.oci.image.index.v1+json"
	var seen [2]atomic.Value
	first := newStubRegistry(t, func(w http.ResponseWriter, r *http.Request) {
		seen[0].Store(r.Header.Get("Accept"))
		http.NotFound(w, r)
	})
	second := newStubRegistry(t, func(w http.ResponseWriter, r *http.Request) {
		seen[1].Store(r.Header.Get("Accept"))
		w.Write([]byte("index"))
	})
	h := newMultiHandler(t,
		config.UpstreamConf{URL: first.srv.URL},
		config.UpstreamConf{URL: second.srv.URL})

	req := httptest.NewRequest(http.MethodGet, "/v2/hub/org/app/manifests/1.27", nil)
	req.Header.Set("Accept", accept)
	rec := httptest.NewRecorder()
	h.ServeHTTP(rec, req)
	if rec.Code != http.StatusOK {
		t.Fatalf("got %d, want 200", rec.Code)
	}
	for i := range seen {
		if got, _ := seen[i].Load().(string); got != accept {
			t.Errorf("upstream %d Accept = %q, want %q", i, got, accept)
		}
	}
}

// TestPerUpstreamCachePolicy: the answering upstream's cache flag decides,
// so a chained proxy marked cache: false is passed through while the
// direct fallback's answers are stored.
func TestPerUpstreamCachePolicy(t *testing.T) {
	chained := newStubRegistry(t, func(w http.ResponseWriter, r *http.Request) {
		if strings.HasSuffix(r.URL.Path, "/blobs/sha256:chain") {
			w.Write([]byte("chain-blob"))
			return
		}
		http.NotFound(w, r)
	})
	direct := newStubRegistry(t, serving("direct-blob"))
	no := false
	store, err := storage.NewLocalFS(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	ups := []config.UpstreamConf{
		{URL: chained.srv.URL + "/v2/dockerhub", Cache: &no},
		{URL: direct.srv.URL},
	}
	rc := config.VRepo{Name: "hub", Type: config.VRepoDocker, Upstreams: ups,
		Auth: config.AuthConfig{Mode: config.AuthPublic}}
	h := NewHandler([]config.VRepo{rc},
		map[string][]*http.Client{"hub": {http.DefaultClient, http.DefaultClient}}, proxycache.New(store))

	for _, digest := range []string{"sha256:chain", "sha256:direct"} {
		rec := httptest.NewRecorder()
		h.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "/v2/hub/org/app/blobs/"+digest, nil))
		if rec.Code != http.StatusOK {
			t.Fatalf("%s: got %d, want 200", digest, rec.Code)
		}
	}
	if _, err := store.Stat(t.Context(), "docker/hub/blobs/sha256:chain"); err != storage.ErrNotFound {
		t.Error("blob from cache: false upstream was stored")
	}
	if _, err := store.Stat(t.Context(), "docker/hub/blobs/sha256:direct"); err != nil {
		t.Errorf("blob from caching upstream missing from store: %v", err)
	}
}