mesh-controller/internal/builder/registry_test.go

package builder

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
)

// An OCI registry as a content-addressed blob store, which is what it is.
//
// Against a stand-in rather than a real registry because what is under test is this side of the
// protocol — that the two steps happen in order, that the digest travels, that a blob already
// there is not sent again, and that a tag is never accepted as a pin.

type fakeRegistry struct {
	blobs    map[string][]byte
	uploads  int
	location string
}

func (f *fakeRegistry) serve(t *testing.T) *httptest.Server {
	t.Helper()
	if f.blobs == nil {
		f.blobs = map[string][]byte{}
	}
	mux := http.NewServeMux()
	server := httptest.NewServer(mux)
	mux.HandleFunc("/v2/", func(w http.ResponseWriter, r *http.Request) {
		switch {
		case r.Method == http.MethodHead && strings.Contains(r.URL.Path, "/blobs/sha256:"):
			digest := r.URL.Path[strings.LastIndex(r.URL.Path, "/")+1:]
			if _, ok := f.blobs[digest]; ok {
				w.WriteHeader(http.StatusOK)
				return
			}
			w.WriteHeader(http.StatusNotFound)
		case r.Method == http.MethodPost && strings.HasSuffix(r.URL.Path, "/blobs/uploads/"):
			f.uploads++
			where := f.location
			if where == "" {
				where = "/v2/upload/" + hex.EncodeToString([]byte("session"))
			}
			w.Header().Set("Location", where)
			w.WriteHeader(http.StatusAccepted)
		case r.Method == http.MethodPut:
			body, _ := io.ReadAll(r.Body)
			digest := r.URL.Query().Get("digest")
			sum := sha256.Sum256(body)
			if digest != "sha256:"+hex.EncodeToString(sum[:]) {
				// A registry verifies what it is given, which is why a corrupt blob is refused
				// here rather than by a machine unpacking it a week later.
				w.WriteHeader(http.StatusBadRequest)
				return
			}
			f.blobs[digest] = body
			w.WriteHeader(http.StatusCreated)
		default:
			w.WriteHeader(http.StatusNotFound)
		}
	})
	t.Cleanup(server.Close)
	return server
}

func registryFor(t *testing.T, f *fakeRegistry) Registry {
	t.Helper()
	server := f.serve(t)
	return Registry{Address: strings.TrimPrefix(server.URL, "http://")}
}

func TestAnArchiveIsStoredAndFetchableByItsDigest(t *testing.T) {
	f := &fakeRegistry{}
	r := registryFor(t, f)
	body := []byte("a theme")
	sum := sha256.Sum256(body)
	digest := "sha256:" + hex.EncodeToString(sum[:])

	where, err := r.PublishArchive(context.Background(), "shell/config", body, digest)
	if err != nil {
		t.Fatal(err)
	}
	if !strings.HasSuffix(where, "/blobs/"+digest) {
		t.Fatalf("what a machine is told to fetch is %q, which does not name the digest", where)
	}
	if string(f.blobs[digest]) != "a theme" {
		t.Fatalf("the registry holds %q", f.blobs[digest])
	}
}

func TestABlobAlreadyThereIsNotSentAgain(t *testing.T) {
	// Not an optimisation: blobs are named by their content, so re-uploading is asking the
	// registry to store what it already has under the name it already has.
	f := &fakeRegistry{}
	r := registryFor(t, f)
	body := []byte("a theme")
	sum := sha256.Sum256(body)
	digest := "sha256:" + hex.EncodeToString(sum[:])

	if _, err := r.PublishArchive(context.Background(), "shell/config", body, digest); err != nil {
		t.Fatal(err)
	}
	if _, err := r.PublishArchive(context.Background(), "shell/config", body, digest); err != nil {
		t.Fatal(err)
	}
	if f.uploads != 1 {
		t.Fatalf("the blob was uploaded %d times", f.uploads)
	}
}

func TestAnAbsoluteUploadLocationIsFollowed(t *testing.T) {
	// A registry may answer with a full URL or with a path. Both happen in the wild, and a client
	// that handles one produces a confusing failure against the other.
	f := &fakeRegistry{}
	server := f.serve(t)
	f.location = server.URL + "/v2/upload/session?state=abc"
	r := Registry{Address: strings.TrimPrefix(server.URL, "http://")}

	body := []byte("x")
	sum := sha256.Sum256(body)
	digest := "sha256:" + hex.EncodeToString(sum[:])
	if _, err := r.PublishArchive(context.Background(), "a/b", body, digest); err != nil {
		t.Fatal(err)
	}
	if _, ok := f.blobs[digest]; !ok {
		t.Fatal("the blob did not arrive when the location carried a query")
	}
}

func TestATagIsNeverAcceptedAsAPin(t *testing.T) {
	// A tag can be made to point at something else, and a bundle is applied on machines with no
	// mesh to ask about anything.
	r := Registry{Address: "registry.invalid", Run: func(
		_ context.Context, _, name string, args ...string) (string, error) {
		if name == "docker" && len(args) > 0 && args[0] == "inspect" {
			return "registry.invalid/shell/server:latest\n", nil
		}
		return "", nil
	}}
	_, err := r.PublishImage(context.Background(), "local", "shell/server")
	if err == nil {
		t.Fatal("an image referred to by tag was accepted")
	}
	if !strings.Contains(err.Error(), "pinned by digest") {
		t.Fatalf("refused for the wrong reason: %v", err)
	}
}

func TestAPushedImageComesBackPinned(t *testing.T) {
	pinned := "registry.invalid/shell/server@sha256:" + strings.Repeat("a", 64)
	var ran []string
	r := Registry{Address: "registry.invalid", Run: func(
		_ context.Context, _, name string, args ...string) (string, error) {
		ran = append(ran, name+" "+strings.Join(args, " "))
		if name == "docker" && len(args) > 0 && args[0] == "inspect" {
			return pinned + "\n", nil
		}
		return "", nil
	}}
	got, err := r.PublishImage(context.Background(), "local", "shell/server")
	if err != nil {
		t.Fatal(err)
	}
	if got != pinned {
		t.Fatalf("got %q", got)
	}
	if len(ran) < 2 || !strings.HasPrefix(ran[0], "docker tag") || !strings.HasPrefix(ran[1], "docker push") {
		t.Fatalf("it did not tag and then push: %v", ran)
	}
}