dummy GHCR/OCI server

dummy_server.go

package ghcr

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"slices"
	"strings"
	"testing"

	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
	"github.com/pkg/errors"
	"oras.land/oras-go/v2"
	"oras.land/oras-go/v2/content/memory"

	"github.com/MaineK00n/vuls-data-update/pkg/dotgit/ls/remote"
	"github.com/google/uuid"
	"github.com/opencontainers/go-digest"
)

type User struct {
	Login string `json:"login"`
	Type  string `json:"type"`
}

type manifest struct {
	id     uint32
	digest string
	bin    []byte
	tags   []string
}

type DummyServer struct {
	owner     string
	pack      string
	users     map[string]User
	manifests map[string]manifest
	counter   uint32
}

func NewDummyServer(owner, pack string) *DummyServer {
	return &DummyServer{
		owner:     owner,
		pack:      pack,
		users:     make(map[string]User),
		manifests: make(map[string]manifest),
		counter:   123,
	}
}

func (d *DummyServer) PopulateUsers(us []User) {
	for _, u := range us {
		d.users[u.Login] = u
	}
}

func (d *DummyServer) PopulateManifest(tags []string) error {
	ctx := context.TODO()
	store := memory.New() // Just to calculate digest, not used as content store
	d.counter = d.counter + 1

	desc, err := oras.PackManifest(ctx, store, oras.PackManifestVersion1_1, "application/vnd.example.artifact.v1", oras.PackManifestOptions{
		ManifestAnnotations: map[string]string{"dummy-server-counter": fmt.Sprintf("%d", d.counter)}, // To make digests different
	})
	if err != nil {
		return errors.Wrap(err, "pack manifest")
	}

	rc, err := store.Fetch(ctx, desc)
	if err != nil {
		return errors.Wrap(err, "fetch manifest")
	}
	defer rc.Close()

	bin, err := io.ReadAll(rc)
	if err != nil {
		return errors.Wrap(err, "read manifest")
	}

	d.manifests[desc.Digest.String()] = manifest{
		id:     d.counter,
		digest: desc.Digest.String(),
		bin:    bin,
		tags:   tags,
	}

	return nil
}

func (d *DummyServer) Start(t *testing.T) (string, func()) {
	tlss := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		t.Logf("%s: %s\n", r.Method, r.URL.Path)

		switch {
		case r.URL.Path == fmt.Sprintf("/users/%s", d.owner):
			switch r.Method {
			case http.MethodGet:
				w.WriteHeader(http.StatusOK)
				u, found := d.users[d.owner]
				if !found {
					t.Errorf("user not found: %s", d.owner)
					http.Error(w, "User not found", http.StatusNotFound)
					return
				}
				if err := json.NewEncoder(w).Encode(u); err != nil {
					t.Errorf("write response: %v", err)
				}
			default:
				t.Errorf("unexpected request received. method: %s, URL: %s", r.Method, r.URL.String())
				http.Error(w, "Bad Request", http.StatusBadRequest)
			}
		case strings.HasPrefix(r.URL.Path, fmt.Sprintf("/orgs/%s/packages/container/%s/versions", d.owner, d.pack)):
			switch r.Method {
			case http.MethodGet:
				vs := make([]remote.Version, 0, len(d.manifests))

				for k, m := range d.manifests {
					vs = append(vs, remote.Version{
						ID:   int(m.id),
						Name: k,
						URL:  fmt.Sprintf("https://api.github.com/orgs/%s/packages/container/%s/versions/%d", d.owner, d.pack, m.id),
						Metadata: &remote.Metadata{
							PackageType: "container",
							Container: &remote.Container{
								Tags: m.tags,
							},
						},
					})
				}
				w.WriteHeader(http.StatusOK)
				err := json.NewEncoder(w).Encode(vs)
				if err != nil {
					t.Errorf("write response: %v", err)
				}
			case http.MethodDelete:
				id := strings.TrimPrefix(r.URL.Path, fmt.Sprintf("/orgs/%s/packages/container/%s/versions/", d.owner, d.pack))
				for digest, m := range d.manifests {
					if fmt.Sprintf("%d", m.id) != id {
						continue
					}
					delete(d.manifests, digest)
					break
				}

				w.WriteHeader(http.StatusNoContent)
			default:
				t.Errorf("unexpected request received. method: %s, URL: %s", r.Method, r.URL.String())
				http.Error(w, "Bad Request", http.StatusBadRequest)
			}

		case strings.HasPrefix(r.URL.Path, fmt.Sprintf("/v2/%s/%s/blobs/uploads/", d.owner, d.pack)):
			switch r.Method {
			case http.MethodPost:
				sessionID := uuid.NewString()
				w.Header().Set("Location", fmt.Sprintf("/v2/%s/%s/blobs/uploads/%s", d.owner, d.pack, sessionID))
				w.Header().Set("Docker-Upload-UUID", sessionID)
				w.WriteHeader(http.StatusAccepted)
			case http.MethodPut:
				queryDigest := r.URL.Query().Get("digest")
				if queryDigest == "" {
					t.Errorf("digest query parameter is required. actual: %q", r.URL.RawQuery)
					http.Error(w, "digest query parameter is required", http.StatusBadRequest)
					return
				}

				bodyBytes, err := io.ReadAll(r.Body)
				if err != nil {
					t.Errorf("read request body: %v", err)
					http.Error(w, "failed to read body", http.StatusInternalServerError)
					return
				}
				calculatedDigest := digest.FromBytes(bodyBytes)
				if calculatedDigest.String() != queryDigest {
					t.Errorf("digest mismatch: query=%s, calculated=%s", queryDigest, calculatedDigest)
					http.Error(w, "digest mismatch", http.StatusBadRequest)
					return
				}

				w.Header().Set("Location", fmt.Sprintf("/v2/%s/%s/blobs/%s", d.owner, d.pack, calculatedDigest))
				w.Header().Set("Docker-Content-Digest", calculatedDigest.String())
				w.WriteHeader(http.StatusCreated)

			default:
				t.Errorf("unexpected request received. method: %s, URL: %s", r.Method, r.URL.String())
				http.Error(w, "Bad Request", http.StatusBadRequest)
			}

		case strings.HasPrefix(r.URL.Path, fmt.Sprintf("/v2/%s/%s/blobs/", d.owner, d.pack)):
			switch r.Method {
			case http.MethodHead:
				w.WriteHeader(http.StatusNotFound)
			default:
				t.Errorf("unexpected request received. method: %s, URL: %s", r.Method, r.URL.String())
				http.Error(w, "Bad Request", http.StatusBadRequest)
			}

		case strings.HasPrefix(r.URL.Path, fmt.Sprintf("/v2/%s/%s/manifests/", d.owner, d.pack)):
			key := strings.TrimPrefix(r.URL.Path, fmt.Sprintf("/v2/%s/%s/manifests/", d.owner, d.pack))
			switch r.Method {
			case http.MethodPut:
				bin, err := io.ReadAll(r.Body)
				if err != nil {
					t.Errorf("read request body: %v", err)
				}

				if strings.HasPrefix(key, "sha256:") {
					d.counter = d.counter + 1
					d.manifests[key] = manifest{
						id:     d.counter,
						digest: key,
						bin:    bin,
					}
				} else {
					calculatedDigest := digest.FromBytes(bin)
					m, ok := d.manifests[calculatedDigest.String()]
					if !ok {
						t.Errorf("manifest not found: %s", calculatedDigest.String())
						http.Error(w, "manifest not found", http.StatusBadRequest)
						return
					}

					m.tags = append(m.tags, key)
					d.manifests[calculatedDigest.String()] = m
				}

				w.Header().Set("Location", r.URL.Path)
				w.WriteHeader(http.StatusCreated)
			case http.MethodGet:
				var m *manifest
				for k, v := range d.manifests {
					if k == key {
						m = &v
						break
					}
					if slices.Contains(v.tags, key) {
						m = &v
						break
					}
				}
				if m == nil {
					w.WriteHeader(http.StatusNotFound)
					return
				}

				w.Header().Set("Content-Type", ocispec.MediaTypeImageManifest)
				w.Header().Set("Docker-Content-Digest", m.digest)
				w.WriteHeader(http.StatusOK)
				_, err := w.Write(m.bin)
				if err != nil {
					t.Errorf("write response: %v", err)
				}
			default:
				t.Errorf("unexpected request received. method: %s, URL: %s", r.Method, r.URL.String())
				http.Error(w, "Bad Request", http.StatusBadRequest)
			}

		}
	}))

	originalTransport := http.DefaultTransport
	http.DefaultTransport = tlss.Client().Transport

	return tlss.URL, func() {
		tlss.Close()
		http.DefaultTransport = originalTransport
	}
}