ocitest.go

Documentation: cuelabs.dev/go/oci/ociregistry/ocitest

     1  // Copyright 2023 CUE Labs AG
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  // Package ocitest provides some helper types for writing ociregistry-related
    16  // tests. It's designed to be used alongside the [qt package].
    17  //
    18  // [qt package]: https://pkg.go.dev/github.com/go-quicktest/qt
    19  package ocitest
    20  
    21  import (
    22  	"bytes"
    23  	"context"
    24  	"encoding/json"
    25  	"fmt"
    26  	"io"
    27  	"sort"
    28  	"strings"
    29  	"testing"
    30  
    31  	"github.com/go-quicktest/qt"
    32  	"github.com/opencontainers/go-digest"
    33  
    34  	"cuelabs.dev/go/oci/ociregistry"
    35  )
    36  
    37  type Registry struct {
    38  	T *testing.T
    39  	R ociregistry.Interface
    40  }
    41  
    42  // NewRegistry returns a Registry instance that wraps r, providing
    43  // convenience methods for pushing and checking content
    44  // inside the given test instance.
    45  //
    46  // When a Must* method fails, it will fail using t.
    47  func NewRegistry(t *testing.T, r ociregistry.Interface) Registry {
    48  	return Registry{t, r}
    49  }
    50  
    51  // RegistryContent specifies the contents of a registry: a map from
    52  // repository name to the contents of that repository.
    53  type RegistryContent map[string]RepoContent
    54  
    55  // RepoContent specifies the content of a repository.
    56  // manifests and blobs are keyed by symbolic identifiers,
    57  // not used inside the registry itself, but instead
    58  // placeholders for the digest of the associated content.
    59  //
    60  // Digest strings inside manifests that are not valid digests
    61  // will be replaced by the calculated digest of the manifest or
    62  // blob with that identifier; the size and media type fields will also be
    63  // filled in.
    64  type RepoContent struct {
    65  	// Manifests maps from manifest identifier to the contents of the manifest.
    66  	// TODO support manifest indexes too.
    67  	Manifests map[string]ociregistry.Manifest
    68  
    69  	// Blobs maps from blob identifer to the contents of the blob.
    70  	Blobs map[string]string
    71  
    72  	// Tags maps from tag name to manifest identifier.
    73  	Tags map[string]string
    74  }
    75  
    76  // PushedRepoContent mirrors RepoContent but, instead
    77  // of describing content that is to be pushed, describes the
    78  // content that has been pushed.
    79  type PushedRepoContent struct {
    80  	// Manifests holds an entry for each manifest identifier
    81  	// with the descriptor for that manifest.
    82  	Manifests map[string]ociregistry.Descriptor
    83  
    84  	// ManifestData holds the actually pushed data for each manifest.
    85  	ManifestData map[string][]byte
    86  
    87  	// Blobs holds an entry for each blob identifier
    88  	// with the descriptor for that manifest.
    89  	Blobs map[string]ociregistry.Descriptor
    90  }
    91  
    92  // PushContent pushes all the content in rc to r.
    93  //
    94  // It returns a map mapping repository name to the descriptors
    95  // describing the content that has actually been pushed.
    96  func PushContent(r ociregistry.Interface, rc RegistryContent) (map[string]PushedRepoContent, error) {
    97  	regContent := make(map[string]PushedRepoContent)
    98  	for repo, repoc := range rc {
    99  		prc, err := PushRepoContent(r, repo, repoc)
   100  		if err != nil {
   101  			return nil, fmt.Errorf("cannot push content for repository %q: %v", repo, err)
   102  		}
   103  		regContent[repo] = prc
   104  	}
   105  	return regContent, nil
   106  }
   107  
   108  // PushRepoContent pushes the content for a single repository.
   109  func PushRepoContent(r ociregistry.Interface, repo string, repoc RepoContent) (PushedRepoContent, error) {
   110  	ctx := context.Background()
   111  	prc := PushedRepoContent{
   112  		Manifests:    make(map[string]ociregistry.Descriptor),
   113  		ManifestData: make(map[string][]byte),
   114  		Blobs:        make(map[string]ociregistry.Descriptor),
   115  	}
   116  
   117  	for id, blob := range repoc.Blobs {
   118  		prc.Blobs[id] = ociregistry.Descriptor{
   119  			Digest:    digest.FromString(blob),
   120  			Size:      int64(len(blob)),
   121  			MediaType: "application/binary",
   122  		}
   123  	}
   124  	manifests, manifestSeq, err := completedManifests(repoc, prc.Blobs)
   125  	if err != nil {
   126  		return PushedRepoContent{}, err
   127  	}
   128  	for id, content := range manifests {
   129  		prc.Manifests[id] = content.desc
   130  		prc.ManifestData[id] = content.data
   131  	}
   132  	// First push all the blobs:
   133  	for id, content := range repoc.Blobs {
   134  		_, err := r.PushBlob(ctx, repo, prc.Blobs[id], strings.NewReader(content))
   135  		if err != nil {
   136  			return PushedRepoContent{}, fmt.Errorf("cannot push blob %q in repo %q: %v", id, repo, err)
   137  		}
   138  	}
   139  	// Then push the manifests that refer to the blobs.
   140  	for _, mc := range manifestSeq {
   141  		_, err := r.PushManifest(ctx, repo, "", mc.data, mc.desc.MediaType)
   142  		if err != nil {
   143  			return PushedRepoContent{}, fmt.Errorf("cannot push manifest %q in repo %q: %v", mc.id, repo, err)
   144  		}
   145  	}
   146  	// Then push any tags.
   147  	for tag, id := range repoc.Tags {
   148  		mc, ok := manifests[id]
   149  		if !ok {
   150  			return PushedRepoContent{}, fmt.Errorf("tag %q refers to unknown manifest id %q", tag, id)
   151  		}
   152  		_, err := r.PushManifest(ctx, repo, tag, mc.data, mc.desc.MediaType)
   153  		if err != nil {
   154  			return PushedRepoContent{}, fmt.Errorf("cannot push tag %q in repo %q: %v", id, repo, err)
   155  		}
   156  	}
   157  	return prc, nil
   158  }
   159  
   160  // PushContent pushes all the content in rc to r.
   161  //
   162  // It returns a map mapping repository name to the descriptors
   163  // describing the content that has actually been pushed.
   164  func (r Registry) MustPushContent(rc RegistryContent) map[string]PushedRepoContent {
   165  	prc, err := PushContent(r.R, rc)
   166  	qt.Assert(r.T, qt.IsNil(err))
   167  	return prc
   168  }
   169  
   170  type manifestContent struct {
   171  	id   string
   172  	data []byte
   173  	desc ociregistry.Descriptor
   174  }
   175  
   176  // completedManifests calculates the content of all the manifests and returns
   177  // them all, keyed by id, and a partially ordered sequence suitable
   178  // for pushing to a registry in bottom-up order.
   179  func completedManifests(repoc RepoContent, blobs map[string]ociregistry.Descriptor) (map[string]manifestContent, []manifestContent, error) {
   180  	manifests := make(map[string]manifestContent)
   181  	manifestSeq := make([]manifestContent, 0, len(repoc.Manifests))
   182  	// subject relationships can be arbitrarily deep, so continue iterating until
   183  	// all the levels are completed. If at any point we can't make progress, we
   184  	// know there's a problem and panic.
   185  	required := make(map[string]bool)
   186  	for {
   187  		madeProgress := false
   188  		needMore := false
   189  		need := func(digest ociregistry.Digest) {
   190  			needMore = true
   191  			if !required[string(digest)] {
   192  				required[string(digest)] = true
   193  				madeProgress = true
   194  			}
   195  		}
   196  		for id, m := range repoc.Manifests {
   197  			if _, ok := manifests[id]; ok {
   198  				continue
   199  			}
   200  			m1 := m
   201  			if m1.Subject != nil {
   202  				mc, ok := manifests[string(m1.Subject.Digest)]
   203  				if !ok {
   204  					need(m1.Subject.Digest)
   205  					continue
   206  				}
   207  				m1.Subject = ref(*m1.Subject)
   208  				*m1.Subject = mc.desc
   209  				madeProgress = true
   210  			}
   211  			m1.Config = fillBlobDescriptor(m.Config, blobs)
   212  			m1.Layers = make([]ociregistry.Descriptor, len(m.Layers))
   213  			for i, desc := range m.Layers {
   214  				m1.Layers[i] = fillBlobDescriptor(desc, blobs)
   215  			}
   216  			data, err := json.Marshal(m1)
   217  			if err != nil {
   218  				panic(err)
   219  			}
   220  			mc := manifestContent{
   221  				id:   id,
   222  				data: data,
   223  				desc: ociregistry.Descriptor{
   224  					Digest:    digest.FromBytes(data),
   225  					Size:      int64(len(data)),
   226  					MediaType: m.MediaType,
   227  				},
   228  			}
   229  			manifests[id] = mc
   230  			madeProgress = true
   231  			manifestSeq = append(manifestSeq, mc)
   232  		}
   233  		if !needMore {
   234  			return manifests, manifestSeq, nil
   235  		}
   236  		if !madeProgress {
   237  			for m := range required {
   238  				if _, ok := manifests[m]; ok {
   239  					delete(required, m)
   240  				}
   241  			}
   242  			return nil, nil, fmt.Errorf("no manifest found for ids %s", strings.Join(mapKeys(required), ", "))
   243  		}
   244  	}
   245  }
   246  
   247  func fillManifestDescriptors(m ociregistry.Manifest, blobs map[string]ociregistry.Descriptor) ociregistry.Manifest {
   248  	m.Config = fillBlobDescriptor(m.Config, blobs)
   249  	m.Layers = append([]ociregistry.Descriptor(nil), m.Layers...)
   250  	for i, desc := range m.Layers {
   251  		m.Layers[i] = fillBlobDescriptor(desc, blobs)
   252  	}
   253  	return m
   254  }
   255  
   256  func fillBlobDescriptor(d ociregistry.Descriptor, blobs map[string]ociregistry.Descriptor) ociregistry.Descriptor {
   257  	blobDesc, ok := blobs[string(d.Digest)]
   258  	if !ok {
   259  		panic(fmt.Errorf("no blob found with id %q", d.Digest))
   260  	}
   261  	d.Digest = blobDesc.Digest
   262  	d.Size = blobDesc.Size
   263  	if d.MediaType == "" {
   264  		d.MediaType = blobDesc.MediaType
   265  	}
   266  	return d
   267  }
   268  
   269  func (r Registry) MustPushBlob(repo string, data []byte) ociregistry.Descriptor {
   270  	desc := ociregistry.Descriptor{
   271  		Digest:    digest.FromBytes(data),
   272  		Size:      int64(len(data)),
   273  		MediaType: "application/octet-stream",
   274  	}
   275  	desc1, err := r.R.PushBlob(context.Background(), repo, desc, bytes.NewReader(data))
   276  	qt.Assert(r.T, qt.IsNil(err))
   277  	return desc1
   278  }
   279  
   280  func (r Registry) MustPushManifest(repo string, jsonObject any, tag string) ([]byte, ociregistry.Descriptor) {
   281  	data, err := json.Marshal(jsonObject)
   282  	qt.Assert(r.T, qt.IsNil(err))
   283  	var mt struct {
   284  		MediaType string `json:"mediaType,omitempty"`
   285  	}
   286  	err = json.Unmarshal(data, &mt)
   287  	qt.Assert(r.T, qt.IsNil(err))
   288  	qt.Assert(r.T, qt.Not(qt.Equals(mt.MediaType, "")))
   289  	desc := ociregistry.Descriptor{
   290  		Digest:    digest.FromBytes(data),
   291  		Size:      int64(len(data)),
   292  		MediaType: mt.MediaType,
   293  	}
   294  	desc1, err := r.R.PushManifest(context.Background(), repo, tag, data, mt.MediaType)
   295  	qt.Assert(r.T, qt.IsNil(err))
   296  	qt.Check(r.T, qt.Equals(desc1.Digest, desc.Digest))
   297  	qt.Check(r.T, qt.Equals(desc1.Size, desc.Size))
   298  	qt.Check(r.T, qt.Equals(desc1.MediaType, desc.MediaType))
   299  	return data, desc1
   300  }
   301  
   302  type Repo struct {
   303  	T    *testing.T
   304  	Name string
   305  	R    ociregistry.Interface
   306  }
   307  
   308  // HasContent returns a checker that checks r matches the expected
   309  // data and has the expected content type. If wantMediaType is
   310  // empty, application/octet-stream will be expected.
   311  func HasContent(r ociregistry.BlobReader, wantData []byte, wantMediaType string) qt.Checker {
   312  	if wantMediaType == "" {
   313  		wantMediaType = "application/octet-stream"
   314  	}
   315  	return contentChecker{
   316  		r:             r,
   317  		wantData:      wantData,
   318  		wantMediaType: wantMediaType,
   319  	}
   320  }
   321  
   322  type contentChecker struct {
   323  	r             ociregistry.BlobReader
   324  	wantData      []byte
   325  	wantMediaType string
   326  }
   327  
   328  func (c contentChecker) Args() []qt.Arg {
   329  	return []qt.Arg{{
   330  		Name:  "reader",
   331  		Value: c.r,
   332  	}, {
   333  		Name:  "data",
   334  		Value: c.wantData,
   335  	}, {
   336  		Name:  "mediaType",
   337  		Value: c.wantMediaType,
   338  	}}
   339  }
   340  
   341  func (c contentChecker) Check(note func(key string, value any)) error {
   342  	desc := c.r.Descriptor()
   343  	gotData, err := io.ReadAll(c.r)
   344  	if err != nil {
   345  		return qt.BadCheckf("error reading data: %v", err)
   346  	}
   347  	if got, want := desc.Size, int64(len(c.wantData)); got != want {
   348  		note("actual data", gotData)
   349  		return fmt.Errorf("mismatched content length (got %d want %d)", got, want)
   350  	}
   351  	if got, want := desc.Digest, digest.FromBytes(c.wantData); got != want {
   352  		note("actual data", gotData)
   353  		return fmt.Errorf("mismatched digest (got %v want %v)", got, want)
   354  	}
   355  	if !bytes.Equal(gotData, c.wantData) {
   356  		note("actual data", gotData)
   357  		return fmt.Errorf("mismatched content")
   358  	}
   359  	if got, want := desc.MediaType, c.wantMediaType; got != want {
   360  		note("actual media type", desc.MediaType)
   361  		return fmt.Errorf("media type mismatch")
   362  	}
   363  	return nil
   364  }
   365  
   366  func ref[T any](x T) *T {
   367  	return &x
   368  }
   369  
   370  func mapKeys[V any](m map[string]V) []string {
   371  	keys := make([]string, 0, len(m))
   372  	for k := range m {
   373  		keys = append(keys, k)
   374  	}
   375  	sort.Strings(keys)
   376  	return keys
   377  }
   378
View as plain text