From 3c5f85abd10c949a238e89efef667c3a40b17303 Mon Sep 17 00:00:00 2001
From: Clayton Coleman <ccoleman@redhat.com>
Date: Tue, 19 Sep 2017 21:25:37 -0400
Subject: [PATCH] Allow clients to request specific manifest media types

The current registry/client sends the registered manifest types in
random order. Allow clients to request a single specific manifest type
or a preferred order as per the HTTP spec.

Signed-off-by: Clayton Coleman <ccoleman@redhat.com>
---
 registry.go                        | 15 ++++++++
 registry/client/repository.go      | 16 ++++++--
 registry/client/repository_test.go | 61 ++++++++++++++++++++++++++++++
 3 files changed, 88 insertions(+), 4 deletions(-)

diff --git a/registry.go b/registry.go
index c34207d0..a3a80ab8 100644
--- a/registry.go
+++ b/registry.go
@@ -73,6 +73,21 @@ func (o WithTagOption) Apply(m ManifestService) error {
 	return nil
 }
 
+// WithManifestMediaTypes lists the media types the client wishes
+// the server to provide.
+func WithManifestMediaTypes(mediaTypes []string) ManifestServiceOption {
+	return WithManifestMediaTypesOption{mediaTypes}
+}
+
+// WithManifestMediaTypesOption holds a list of accepted media types
+type WithManifestMediaTypesOption struct{ MediaTypes []string }
+
+// Apply conforms to the ManifestServiceOption interface
+func (o WithManifestMediaTypesOption) Apply(m ManifestService) error {
+	// no implementation
+	return nil
+}
+
 // Repository is a named collection of manifests and layers.
 type Repository interface {
 	// Named returns the name of the repository.
diff --git a/registry/client/repository.go b/registry/client/repository.go
index f9499552..ed29cf3f 100644
--- a/registry/client/repository.go
+++ b/registry/client/repository.go
@@ -418,18 +418,22 @@ func (ms *manifests) Get(ctx context.Context, dgst digest.Digest, options ...dis
 		ref         reference.Named
 		err         error
 		contentDgst *digest.Digest
+		mediaTypes  []string
 	)
 
 	for _, option := range options {
-		if opt, ok := option.(distribution.WithTagOption); ok {
+		switch opt := option.(type) {
+		case distribution.WithTagOption:
 			digestOrTag = opt.Tag
 			ref, err = reference.WithTag(ms.name, opt.Tag)
 			if err != nil {
 				return nil, err
 			}
-		} else if opt, ok := option.(contentDigestOption); ok {
+		case contentDigestOption:
 			contentDgst = opt.digest
-		} else {
+		case distribution.WithManifestMediaTypesOption:
+			mediaTypes = opt.MediaTypes
+		default:
 			err := option.Apply(ms)
 			if err != nil {
 				return nil, err
@@ -445,6 +449,10 @@ func (ms *manifests) Get(ctx context.Context, dgst digest.Digest, options ...dis
 		}
 	}
 
+	if len(mediaTypes) == 0 {
+		mediaTypes = distribution.ManifestMediaTypes()
+	}
+
 	u, err := ms.ub.BuildManifestURL(ref)
 	if err != nil {
 		return nil, err
@@ -455,7 +463,7 @@ func (ms *manifests) Get(ctx context.Context, dgst digest.Digest, options ...dis
 		return nil, err
 	}
 
-	for _, t := range distribution.ManifestMediaTypes() {
+	for _, t := range mediaTypes {
 		req.Header.Add("Accept", t)
 	}
 
diff --git a/registry/client/repository_test.go b/registry/client/repository_test.go
index e5a6dad0..3bcf63e1 100644
--- a/registry/client/repository_test.go
+++ b/registry/client/repository_test.go
@@ -9,6 +9,8 @@ import (
 	"log"
 	"net/http"
 	"net/http/httptest"
+	"reflect"
+	"sort"
 	"strconv"
 	"strings"
 	"testing"
@@ -815,6 +817,65 @@ func TestManifestFetchWithEtag(t *testing.T) {
 	}
 }
 
+func TestManifestFetchWithAccept(t *testing.T) {
+	ctx := context.Background()
+	repo, _ := reference.WithName("test.example.com/repo")
+	_, dgst, _ := newRandomSchemaV1Manifest(repo, "latest", 6)
+	headers := make(chan []string, 1)
+	s := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
+		headers <- req.Header["Accept"]
+	}))
+	defer close(headers)
+	defer s.Close()
+
+	r, err := NewRepository(repo, s.URL, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	ms, err := r.Manifests(ctx)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	testCases := []struct {
+		// the media types we send
+		mediaTypes []string
+		// the expected Accept headers the server should receive
+		expect []string
+		// whether to sort the request and response values for comparison
+		sort bool
+	}{
+		{
+			mediaTypes: []string{},
+			expect:     distribution.ManifestMediaTypes(),
+			sort:       true,
+		},
+		{
+			mediaTypes: []string{"test1", "test2"},
+			expect:     []string{"test1", "test2"},
+		},
+		{
+			mediaTypes: []string{"test1"},
+			expect:     []string{"test1"},
+		},
+		{
+			mediaTypes: []string{""},
+			expect:     []string{""},
+		},
+	}
+	for _, testCase := range testCases {
+		ms.Get(ctx, dgst, distribution.WithManifestMediaTypes(testCase.mediaTypes))
+		actual := <-headers
+		if testCase.sort {
+			sort.Strings(actual)
+			sort.Strings(testCase.expect)
+		}
+		if !reflect.DeepEqual(actual, testCase.expect) {
+			t.Fatalf("unexpected Accept header values: %v", actual)
+		}
+	}
+}
+
 func TestManifestDelete(t *testing.T) {
 	repo, _ := reference.WithName("test.example.com/repo/delete")
 	_, dgst1, _ := newRandomSchemaV1Manifest(repo, "latest", 6)