Add custom headers option to dockerResolver

0xch4z · 0xch4z · commit 068880ff1670 · 2018-10-31T19:42:42.000-04:00
Signed-off-by: Charles Kenney &lt;charlesc.kenney@gmail.com&gt;
diff --git a/remotes/docker/resolver.go b/remotes/docker/resolver.go
@@ -29,6 +29,7 @@ import (
 	"github.com/containerd/containerd/log"
 	"github.com/containerd/containerd/reference"
 	"github.com/containerd/containerd/remotes"
+	"github.com/containerd/containerd/version"
 	digest "github.com/opencontainers/go-digest"
 	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
 	"github.com/pkg/errors"
@@ -82,6 +83,9 @@ type ResolverOptions struct {
 	// Host provides the hostname given a namespace.
 	Host func(string) (string, error)
 
+	// Headers are the HTTP request header fields sent by the resolver
+	Headers http.Header
+
 	// PlainHTTP specifies to use plain http and not https
 	PlainHTTP bool
 
@@ -105,6 +109,7 @@ func DefaultHost(ns string) (string, error) {
 type dockerResolver struct {
 	auth      Authorizer
 	host      func(string) (string, error)
+	headers   http.Header
 	plainHTTP bool
 	client    *http.Client
 	tracker   StatusTracker
@@ -118,12 +123,27 @@ func NewResolver(options ResolverOptions) remotes.Resolver {
 	if options.Host == nil {
 		options.Host = DefaultHost
 	}
+	if options.Headers == nil {
+		options.Headers = make(http.Header)
+	}
+	if _, ok := options.Headers["Accept"]; !ok {
+		// set headers for all the types we support for resolution.
+		options.Headers.Set("Accept", strings.Join([]string{
+			images.MediaTypeDockerSchema2Manifest,
+			images.MediaTypeDockerSchema2ManifestList,
+			ocispec.MediaTypeImageManifest,
+			ocispec.MediaTypeImageIndex, "*"}, ", "))
+	}
+	if _, ok := options.Headers["User-Agent"]; !ok {
+		options.Headers.Set("User-Agent", "containerd/"+version.Version)
+	}
 	if options.Authorizer == nil {
 		options.Authorizer = NewAuthorizer(options.Client, options.Credentials)
 	}
 	return &dockerResolver{
 		auth:      options.Authorizer,
 		host:      options.Host,
+		headers:   options.Headers,
 		plainHTTP: options.PlainHTTP,
 		client:    options.Client,
 		tracker:   options.Tracker,
@@ -182,12 +202,7 @@ func (r *dockerResolver) Resolve(ctx context.Context, ref string) (string, ocisp
 			return "", ocispec.Descriptor{}, err
 		}
 
-		// set headers for all the types we support for resolution.
-		req.Header.Set("Accept", strings.Join([]string{
-			images.MediaTypeDockerSchema2Manifest,
-			images.MediaTypeDockerSchema2ManifestList,
-			ocispec.MediaTypeImageManifest,
-			ocispec.MediaTypeImageIndex, "*"}, ", "))
+		req.Header = r.headers
 
 		log.G(ctx).Debug("resolving")
 		resp, err := fetcher.doRequestWithRetries(ctx, req, nil)
