map.go

Documentation: edge-infra.dev/pkg/f8n/warehouse/oci/mutate

     1  package mutate
     2  
     3  import (
     4  	"errors"
     5  	"fmt"
     6  
     7  	v1 "github.com/google/go-containerregistry/pkg/v1"
     8  	"github.com/google/go-containerregistry/pkg/v1/empty"
     9  	"github.com/google/go-containerregistry/pkg/v1/mutate"
    10  
    11  	"edge-infra.dev/pkg/f8n/warehouse"
    12  	"edge-infra.dev/pkg/f8n/warehouse/oci"
    13  	"edge-infra.dev/pkg/f8n/warehouse/oci/layer"
    14  )
    15  
    16  // SkipAll is used as a return value from map Fns to indicate that traversal
    17  // should end immediately, returning the input artifact unmodified.
    18  //
    19  //nolint:revive // Sentinel error values for controlling functionality by user can omit Err prefix. See the stdlib fs packages for examples.
    20  var SkipAll = errors.New("end the traversal immediately")
    21  
    22  // Map performs a depth-first traversal of the input artifact's graph, visiting
    23  // each node in reverse topological order from the input artifact's root node:
    24  //
    25  // - Dependencies are always visited before the nodes that depend on them:
    26  //   - When the input is a [v1.ImageIndex], [v1.Image] attached directly to the
    27  //     root index with the same name as the index will be visited last, because
    28  //     they represent cluster variants of the package artifact being mapped,
    29  //     and must be visited last. In a standard DFS, those descriptors
    30  //     could be visited on the first pass.
    31  //
    32  // - Functions are called after transformation, in a bottom up fashion.
    33  func Map(a oci.Artifact, fns *Fns) (oci.Artifact, error) {
    34  	switch t := a.(type) {
    35  	case oci.Unwrapper:
    36  		return Map(t.Unwrap(), fns)
    37  	case v1.ImageIndex:
    38  		return fns.mapIdx(t, nil)
    39  	case v1.Image:
    40  		return fns.mapImg(t, nil)
    41  	default:
    42  		return nil, fmt.Errorf("%w: unexpected Artifact type %T", oci.ErrInvalidArtifact, t)
    43  	}
    44  }
    45  
    46  // Fns defines functions for each OCI entity to be called during graph traversal
    47  // in [Map] operations. Functions can mutate the input OCI entity by returning
    48  // a new entity. Returning nil causes the entity to be dropped.
    49  //
    50  // The input OCI entity (idx, img, l) is passed in _after_ transforming all of
    51  // it's child nodes. The parent OCI entity passed in is _before_ transformation,
    52  // as it cannot be visited until after all of its child nodes have.
    53  type Fns struct {
    54  	Index func(idx v1.ImageIndex, parent v1.ImageIndex) (v1.ImageIndex, error)
    55  	Image func(img v1.Image, parent v1.ImageIndex) (v1.Image, error)
    56  	Layer func(l layer.Layer, parent v1.Image) (layer.Layer, error)
    57  }
    58  
    59  func (f *Fns) mapIdx(idx v1.ImageIndex, parent v1.ImageIndex) (v1.ImageIndex, error) {
    60  	m, err := idx.IndexManifest()
    61  	if err != nil {
    62  		return nil, err
    63  	}
    64  
    65  	var (
    66  		rootImgs   []v1.Image
    67  		xArtifacts []oci.Artifact // transformed idx constituents
    68  		changed    = false        // track whether or not children are mutated
    69  	)
    70  
    71  	// First pass, visit all dependencies before dependents:
    72  	// - Visit any descriptors referencing images which don't have the same
    73  	//   name as the index
    74  	for _, desc := range m.Manifests {
    75  		switch {
    76  		case desc.MediaType.IsIndex():
    77  			ii, err := idx.ImageIndex(desc.Digest)
    78  			if err != nil {
    79  				return nil, err
    80  			}
    81  
    82  			x, err := f.mapIdx(ii, idx)
    83  
    84  			switch {
    85  			case errors.Is(err, SkipAll):
    86  				return idx, nil
    87  			case err != nil:
    88  				return nil, err
    89  			case x == nil:
    90  				changed = true
    91  				continue
    92  			}
    93  
    94  			changed = changed || (x != ii)
    95  			xArtifacts = append(xArtifacts, x)
    96  		case desc.MediaType.IsImage():
    97  			img, err := idx.Image(desc.Digest)
    98  			if err != nil {
    99  				return nil, err
   100  			}
   101  
   102  			// Stash desc for some fun later, this is the descriptor referencing
   103  			// the root nodes image, which needs to be visited last
   104  			if descRefsIdx(m, desc) {
   105  				rootImgs = append(rootImgs, img)
   106  				continue
   107  			}
   108  
   109  			// Otherwise, we can go ahead and visit it, since it is a leaf node
   110  			x, err := f.mapImg(img, idx)
   111  
   112  			switch {
   113  			case errors.Is(err, SkipAll):
   114  				return idx, nil
   115  			case err != nil:
   116  				return nil, err
   117  			case x == nil:
   118  				changed = true
   119  				continue
   120  			}
   121  
   122  			changed = changed || (x != img)
   123  			xArtifacts = append(xArtifacts, x)
   124  		default:
   125  			return nil, fmt.Errorf("%w: unexpected mediaType %s",
   126  				oci.ErrInvalidArtifact, desc.MediaType)
   127  		}
   128  	}
   129  
   130  	// Now that we have traversed everything except for the root package image
   131  	// descriptors, we can visit them.
   132  	for _, img := range rootImgs {
   133  		x, err := f.mapImg(img, idx)
   134  
   135  		switch {
   136  		case errors.Is(err, SkipAll):
   137  			return idx, nil
   138  		case err != nil:
   139  			return nil, err
   140  		case x == nil:
   141  			changed = true
   142  			continue
   143  		}
   144  
   145  		changed = changed || (x != img)
   146  		xArtifacts = append(xArtifacts, x)
   147  	}
   148  
   149  	if !changed {
   150  		return f.idxFn(idx, parent)
   151  	}
   152  
   153  	// TODO(aw185176): should probably explicitly preserve mediaType here once
   154  	// google/go-containerregistry cuts a release containing fix for empty.Index
   155  	// mediaType
   156  	n, err := AppendManifests(empty.Index, xArtifacts...)
   157  	if err != nil {
   158  		return nil, err
   159  	}
   160  
   161  	// Visit the index now that all of its children have been visited
   162  	return f.idxFn(mutate.Annotations(n, m.Annotations).(v1.ImageIndex), parent)
   163  }
   164  
   165  func (f *Fns) mapImg(img v1.Image, parent v1.ImageIndex) (v1.Image, error) {
   166  	// short circuit if there is no chance we change the contents of img
   167  	switch {
   168  	case f.Image == nil && f.Layer == nil:
   169  		return img, nil
   170  	case f.Layer == nil:
   171  		return f.imgFn(img, parent)
   172  	}
   173  
   174  	var (
   175  		changed = false        // track if image contents were mutated
   176  		xLayers = []v1.Layer{} // transformed layers
   177  	)
   178  
   179  	layers, err := layer.FromImage(img)
   180  	if err != nil {
   181  		return nil, err
   182  	}
   183  
   184  	for _, layer := range layers {
   185  		x, err := f.layerFn(layer, img)
   186  
   187  		switch {
   188  		case errors.Is(err, SkipAll):
   189  			return img, nil
   190  		case err != nil:
   191  			return nil, err
   192  		case x == nil:
   193  			changed = true
   194  			continue
   195  		}
   196  
   197  		changed = changed || (x != layer)
   198  		xLayers = append(xLayers, x)
   199  	}
   200  
   201  	if !changed {
   202  		return f.imgFn(img, parent)
   203  	}
   204  
   205  	// TODO: generalize image mutation, mutate.Image ?
   206  	n, err := ReplaceLayers(img, xLayers...)
   207  	if err != nil {
   208  		return nil, err
   209  	}
   210  
   211  	return f.imgFn(n, parent)
   212  }
   213  
   214  // idxFn wraps Fns.Index with a nil check, to simplify mapping logic
   215  func (f *Fns) idxFn(idx v1.ImageIndex, parent v1.ImageIndex) (v1.ImageIndex, error) {
   216  	if f.Index != nil {
   217  		return f.Index(idx, parent)
   218  	}
   219  	return idx, nil
   220  }
   221  
   222  // imgFn wraps Fns.Image with a nil check, to simplify mapping logic
   223  func (f *Fns) imgFn(img v1.Image, parent v1.ImageIndex) (v1.Image, error) {
   224  	if f.Image != nil {
   225  		return f.Image(img, parent)
   226  	}
   227  	return img, nil
   228  }
   229  
   230  // layerFn wraps Fns.Layer with a nil check, to simplify mapping logic
   231  func (f *Fns) layerFn(l layer.Layer, parent v1.Image) (layer.Layer, error) {
   232  	if f.Layer != nil {
   233  		return f.Layer(l, parent)
   234  	}
   235  	return l, nil
   236  }
   237  
   238  // descRefsIdx returns true if the input descriptor is referencing the input
   239  // image manifest, e.g., provider variants of a package
   240  func descRefsIdx(idx *v1.IndexManifest, d v1.Descriptor) bool {
   241  	return idx.Annotations[warehouse.AnnotationName] ==
   242  		d.Annotations[warehouse.AnnotationRefName]
   243  }
   244
View as plain text