allow specifying media type (oras-project#14)

* Support media type * update CLI * optionally allow all media types * no media types means all * log name change * default blob media type
dcomas · Dec 28, 2018 · 8088466 · 8088466
1 parent 73bedce
commit 8088466
Show file tree

Hide file tree

Showing 7 changed files with 119 additions and 63 deletions.
diff --git a/cmd/oras/pull.go b/cmd/oras/pull.go
@@ -13,9 +13,11 @@ import (
 )
 
 type pullOptions struct {
-	targetRef string
-	output    string
-	verbose   bool
+	targetRef          string
+	allowedMediaTypes  []string
+	allowAllMediaTypes bool
+	output             string
+	verbose            bool
 
 	debug    bool
 	username string
@@ -34,6 +36,8 @@ func pullCmd() *cobra.Command {
 		},
 	}
 
+	cmd.Flags().StringArrayVarP(&opts.allowedMediaTypes, "media-type", "t", nil, "allowed media types to be pulled")
+	cmd.Flags().BoolVarP(&opts.allowAllMediaTypes, "allow-all", "a", false, "allow all media types to be pulled")
 	cmd.Flags().StringVarP(&opts.output, "output", "o", "", "output directory")
 	cmd.Flags().BoolVarP(&opts.verbose, "verbose", "v", false, "verbose output")
 
@@ -47,18 +51,23 @@ func runPull(opts pullOptions) error {
 	if opts.debug {
 		logrus.SetLevel(logrus.DebugLevel)
 	}
+	if opts.allowAllMediaTypes {
+		opts.allowedMediaTypes = nil
+	} else if len(opts.allowedMediaTypes) == 0 {
+		opts.allowedMediaTypes = []string{oras.DefaultBlobMediaType}
+	}
 
 	resolver := newResolver(opts.username, opts.password)
-	contents, err := oras.Pull(context.Background(), resolver, opts.targetRef)
+	blobs, err := oras.Pull(context.Background(), resolver, opts.targetRef, opts.allowedMediaTypes...)
 	if err != nil {
 		return err
 	}
 
-	for name, content := range contents {
+	for name, blob := range blobs {
 		if opts.output != "" {
 			name = path.Join(opts.output, name)
 		}
-		if err := ioutil.WriteFile(name, content, 0644); err != nil {
+		if err := ioutil.WriteFile(name, blob.Content, 0644); err != nil {
 			return err
 		}
 		if opts.verbose {

diff --git a/cmd/oras/push.go b/cmd/oras/push.go
@@ -3,6 +3,7 @@ package main
 import (
 	"context"
 	"io/ioutil"
+	"strings"
 
 	"github.com/shizhMSFT/oras/pkg/oras"
 
@@ -12,7 +13,7 @@ import (
 
 type pushOptions struct {
 	targetRef string
-	filenames []string
+	fileRefs  []string
 
 	debug    bool
 	username string
@@ -27,7 +28,7 @@ func pushCmd() *cobra.Command {
 		Args:  cobra.MinimumNArgs(2),
 		RunE: func(cmd *cobra.Command, args []string) error {
 			opts.targetRef = args[0]
-			opts.filenames = args[1:]
+			opts.fileRefs = args[1:]
 			return runPush(opts)
 		},
 	}
@@ -45,14 +46,23 @@ func runPush(opts pushOptions) error {
 
 	resolver := newResolver(opts.username, opts.password)
 
-	contents := make(map[string][]byte)
-	for _, filename := range opts.filenames {
+	blobs := make(map[string]oras.Blob)
+	for _, fileRef := range opts.fileRefs {
+		ref := strings.SplitN(fileRef, ":", 2)
+		filename := ref[0]
+		var mediaType string
+		if len(ref) == 2 {
+			mediaType = ref[1]
+		}
 		content, err := ioutil.ReadFile(filename)
 		if err != nil {
 			return err
 		}
-		contents[filename] = content
+		blobs[filename] = oras.Blob{
+			MediaType: mediaType,
+			Content:   content,
+		}
 	}
 
-	return oras.Push(context.Background(), resolver, opts.targetRef, contents)
+	return oras.Push(context.Background(), resolver, opts.targetRef, blobs)
 }
diff --git a/pkg/oras/blob.go b/pkg/oras/blob.go
@@ -0,0 +1,12 @@
+package oras
+
+import ocispec "github.com/opencontainers/image-spec/specs-go/v1"
+
+// DefaultBlobMediaType specifies the default blob media type
+const DefaultBlobMediaType = ocispec.MediaTypeImageLayer
+
+// Blob refers a blob with a media type
+type Blob struct {
+	MediaType string
+	Content   []byte
+}
diff --git a/pkg/oras/errors.go b/pkg/oras/errors.go
@@ -6,5 +6,5 @@ import "errors"
 var (
 	ErrNotFound          = errors.New("not_found")
 	ErrResolverUndefined = errors.New("resolver_undefined")
-	ErrEmptyContents     = errors.New("empty_contents")
+	ErrEmptyBlobs        = errors.New("empty_blobs")
 )
diff --git a/pkg/oras/oras_test.go b/pkg/oras/oras_test.go
@@ -66,39 +66,39 @@ func (suite *ORASTestSuite) SetupSuite() {
 func (suite *ORASTestSuite) Test_0_Push() {
 	var err error
 	var ref string
-	var contents map[string][]byte
+	var blobs map[string]Blob
 
-	err = Push(newContext(), nil, ref, contents)
+	err = Push(newContext(), nil, ref, blobs)
 	suite.NotNil(err, "error pushing with empty resolver")
 
-	err = Push(newContext(), newResolver(), ref, contents)
+	err = Push(newContext(), newResolver(), ref, blobs)
 	suite.NotNil(err, "error pushing when context missing hostname")
 
 	ref = fmt.Sprintf("%s/empty:test", suite.DockerRegistryHost)
-	err = Push(newContext(), newResolver(), ref, contents)
-	suite.NotNil(ErrEmptyContents, err, "error pushing with empty contents")
+	err = Push(newContext(), newResolver(), ref, blobs)
+	suite.NotNil(ErrEmptyBlobs, err, "error pushing with empty blobs")
 
-	// Load contents with test chart tgz (as single layer)
-	contents = make(map[string][]byte)
+	// Load blobs with test chart tgz (as single layer)
+	blobs = make(map[string]Blob)
 	content, err := ioutil.ReadFile(testTarball)
 	suite.Nil(err, "no error loading test chart")
 	basename := filepath.Base(testTarball)
-	contents[basename] = content
+	blobs[basename] = Blob{Content: content}
 
 	ref = fmt.Sprintf("%s/chart-tgz:test", suite.DockerRegistryHost)
-	err = Push(newContext(), newResolver(), ref, contents)
+	err = Push(newContext(), newResolver(), ref, blobs)
 	suite.Nil(err, "no error pushing test chart tgz (as single layer)")
 
-	// Load contents with test chart dir (each file as layer)
-	contents = make(map[string][]byte)
+	// Load blobs with test chart dir (each file as layer)
+	blobs = make(map[string]Blob)
 	var ff = func(pathX string, infoX os.FileInfo, errX error) error {
 		if !infoX.IsDir() {
 			filename := filepath.Join(filepath.Dir(pathX), infoX.Name())
 			content, err := ioutil.ReadFile(filename)
 			if err != nil {
 				return err
 			}
-			contents[filepath.ToSlash(filename)] = content
+			blobs[filepath.ToSlash(filename)] = Blob{Content: content}
 		}
 		return nil
 	}
@@ -109,48 +109,48 @@ func (suite *ORASTestSuite) Test_0_Push() {
 	os.Chdir(cwd)
 
 	ref = fmt.Sprintf("%s/chart-dir:test", suite.DockerRegistryHost)
-	err = Push(newContext(), newResolver(), ref, contents)
+	err = Push(newContext(), newResolver(), ref, blobs)
 	suite.Nil(err, "no error pushing test chart dir (each file as layer)")
 }
 
-// Pull files and verify contents
+// Pull files and verify blobs
 func (suite *ORASTestSuite) Test_1_Pull() {
 	var err error
 	var ref string
-	var contents map[string][]byte
+	var blobs map[string]Blob
 
-	contents, err = Pull(newContext(), nil, ref)
+	blobs, err = Pull(newContext(), nil, ref)
 	suite.NotNil(err, "error pulling with empty resolver")
-	suite.Nil(contents, "contents nil pulling with empty resolver")
+	suite.Nil(blobs, "blobs nil pulling with empty resolver")
 
 	// Pull non-existant
 	ref = fmt.Sprintf("%s/nonexistant:test", suite.DockerRegistryHost)
-	contents, err = Pull(newContext(), newResolver(), ref)
+	blobs, err = Pull(newContext(), newResolver(), ref)
 	suite.NotNil(err, "error pulling non-existant ref")
-	suite.Nil(contents, "contents empty with error")
+	suite.Nil(blobs, "blobs empty with error")
 
 	// Pull chart-tgz
 	ref = fmt.Sprintf("%s/chart-tgz:test", suite.DockerRegistryHost)
-	contents, err = Pull(newContext(), newResolver(), ref)
+	blobs, err = Pull(newContext(), newResolver(), ref)
 	suite.Nil(err, "no error pulling chart-tgz ref")
 
-	// Verify the contents, single layer/file
+	// Verify the blobs, single layer/file
 	content, err := ioutil.ReadFile(testTarball)
 	suite.Nil(err, "no error loading test chart")
-	suite.Equal(content, contents[filepath.Base(testTarball)], ".tgz content matches on pull")
+	suite.Equal(content, blobs[filepath.Base(testTarball)].Content, ".tgz content matches on pull")
 
 	// Pull chart-dir
 	ref = fmt.Sprintf("%s/chart-dir:test", suite.DockerRegistryHost)
-	contents, err = Pull(newContext(), newResolver(), ref)
+	blobs, err = Pull(newContext(), newResolver(), ref)
 	suite.Nil(err, "no error pulling chart-dir ref")
 
-	// Verify the contents, multiple layers/files
+	// Verify the blobs, multiple layers/files
 	cwd, _ := os.Getwd()
 	os.Chdir(testDir)
 	for _, filename := range testDirFiles {
 		content, err = ioutil.ReadFile(filename)
 		suite.Nil(err, fmt.Sprintf("no error loading %s", filename))
-		suite.Equal(content, contents[filename], fmt.Sprintf("%s content matches on pull", filename))
+		suite.Equal(content, blobs[filename].Content, fmt.Sprintf("%s content matches on pull", filename))
 	}
 	os.Chdir(cwd)
 }

diff --git a/pkg/oras/pull.go b/pkg/oras/pull.go
@@ -10,7 +10,7 @@ import (
 )
 
 // Pull pull files from the remote
-func Pull(ctx context.Context, resolver remotes.Resolver, ref string) (map[string][]byte, error) {
+func Pull(ctx context.Context, resolver remotes.Resolver, ref string, allowedMediaTypes ...string) (map[string]Blob, error) {
 	if resolver == nil {
 		return nil, ErrResolverUndefined
 	}
@@ -24,44 +24,65 @@ func Pull(ctx context.Context, resolver remotes.Resolver, ref string) (map[strin
 		return nil, err
 	}
 
-	var layers []ocispec.Descriptor
+	var blobs []ocispec.Descriptor
 	picker := images.HandlerFunc(func(ctx context.Context, desc ocispec.Descriptor) ([]ocispec.Descriptor, error) {
-		if desc.MediaType == ocispec.MediaTypeImageLayer {
-			layers = append(layers, desc)
+		if isAllowedMediaType(desc.MediaType, allowedMediaTypes...) {
+			blobs = append(blobs, desc)
+			return nil, nil
 		}
 		return nil, nil
 	})
 	store := NewMemoryStore()
-	handlers := images.Handlers(filterHandler(), store.FetchHandler(fetcher), picker, images.ChildrenHandler(store))
+	handlers := images.Handlers(
+		filterHandler(allowedMediaTypes...),
+		store.FetchHandler(fetcher),
+		picker,
+		images.ChildrenHandler(store),
+	)
 	if err := images.Dispatch(ctx, handlers, desc); err != nil {
 		return nil, err
 	}
 
-	res := make(map[string][]byte)
-	for _, layer := range layers {
-		if content, ok := store.Get(layer); ok {
-			if name, ok := layer.Annotations[ocispec.AnnotationTitle]; ok && len(name) > 0 {
-				res[name] = content
+	res := make(map[string]Blob)
+	for _, blob := range blobs {
+		if content, ok := store.Get(blob); ok {
+			if name, ok := blob.Annotations[ocispec.AnnotationTitle]; ok && len(name) > 0 {
+				res[name] = Blob{
+					MediaType: blob.MediaType,
+					Content:   content,
+				}
 			}
 		}
 	}
 
 	return res, nil
 }
 
-func filterHandler() images.HandlerFunc {
+func filterHandler(allowedMediaTypes ...string) images.HandlerFunc {
 	return func(ctx context.Context, desc ocispec.Descriptor) ([]ocispec.Descriptor, error) {
-		switch desc.MediaType {
-		case ocispec.MediaTypeImageManifest, ocispec.MediaTypeImageIndex:
+		switch {
+		case isAllowedMediaType(desc.MediaType, ocispec.MediaTypeImageManifest, ocispec.MediaTypeImageIndex):
 			return nil, nil
-		case ocispec.MediaTypeImageLayer:
+		case isAllowedMediaType(desc.MediaType, allowedMediaTypes...):
 			if name, ok := desc.Annotations[ocispec.AnnotationTitle]; ok && len(name) > 0 {
 				return nil, nil
 			}
-			log.G(ctx).Warnf("layer_no_name: %v", desc.Digest)
+			log.G(ctx).Warnf("blob_no_name: %v", desc.Digest)
 		default:
 			log.G(ctx).Warnf("unknown_type: %v", desc.MediaType)
 		}
 		return nil, images.ErrStopHandler
 	}
 }
+
+func isAllowedMediaType(mediaType string, allowedMediaTypes ...string) bool {
+	if len(allowedMediaTypes) == 0 {
+		return true
+	}
+	for _, allowedMediaType := range allowedMediaTypes {
+		if mediaType == allowedMediaType {
+			return true
+		}
+	}
+	return false
+}
diff --git a/pkg/oras/push.go b/pkg/oras/push.go
@@ -12,29 +12,29 @@ import (
 )
 
 // Push pushes files to the remote
-func Push(ctx context.Context, resolver remotes.Resolver, ref string, contents map[string][]byte) error {
+func Push(ctx context.Context, resolver remotes.Resolver, ref string, blobs map[string]Blob) error {
 	if resolver == nil {
 		return ErrResolverUndefined
 	}
 
-	if contents == nil {
-		return ErrEmptyContents
+	if blobs == nil {
+		return ErrEmptyBlobs
 	}
 
 	pusher, err := resolver.Pusher(ctx, ref)
 	if err != nil {
 		return err
 	}
 
-	desc, provider, err := pack(contents)
+	desc, provider, err := pack(blobs)
 	if err != nil {
 		return err
 	}
 
 	return remotes.PushContent(ctx, pusher, desc, provider, nil)
 }
 
-func pack(contents map[string][]byte) (ocispec.Descriptor, content.Provider, error) {
+func pack(blobs map[string]Blob) (ocispec.Descriptor, content.Provider, error) {
 	store := NewMemoryStore()
 
 	// Config
@@ -48,16 +48,20 @@ func pack(contents map[string][]byte) (ocispec.Descriptor, content.Provider, err
 
 	// Layer
 	var layers []ocispec.Descriptor
-	for name, content := range contents {
+	for name, blob := range blobs {
+		mediaType := blob.MediaType
+		if mediaType == "" {
+			mediaType = DefaultBlobMediaType
+		}
 		layer := ocispec.Descriptor{
-			MediaType: ocispec.MediaTypeImageLayer,
-			Digest:    digest.FromBytes(content),
-			Size:      int64(len(content)),
+			MediaType: mediaType,
+			Digest:    digest.FromBytes(blob.Content),
+			Size:      int64(len(blob.Content)),
 			Annotations: map[string]string{
 				ocispec.AnnotationTitle: name,
 			},
 		}
-		store.Set(layer, content)
+		store.Set(layer, blob.Content)
 		layers = append(layers, layer)
 	}