From bb008728001199665d0a6a484d064372487e7150 Mon Sep 17 00:00:00 2001
From: Derek McGowan <derek@mcgstyle.net>
Date: Tue, 11 Jun 2019 22:57:36 +0800
Subject: [PATCH] Add user agent header to all requests

Currently the user agent is only being used on the initial
resolve request, then switching to the default user agent.
This ensures the correct user agent is always used. There is
a larger fix in progress which does this is a cleaner way, but
the scope of this change is fixing the user agent issue.

Signed-off-by: Derek McGowan <derek@mcgstyle.net>
---
 remotes/docker/authorizer.go | 22 +++++++++++++++++-----
 remotes/docker/resolver.go   | 14 ++++++++++++--
 2 files changed, 29 insertions(+), 7 deletions(-)

diff --git a/remotes/docker/authorizer.go b/remotes/docker/authorizer.go
index 73adb5a2f..3097cc3b5 100644
--- a/remotes/docker/authorizer.go
+++ b/remotes/docker/authorizer.go
@@ -31,6 +31,7 @@ import (
 
 	"github.com/containerd/containerd/errdefs"
 	"github.com/containerd/containerd/log"
+	"github.com/containerd/containerd/version"
 	"github.com/pkg/errors"
 	"github.com/sirupsen/logrus"
 	"golang.org/x/net/context/ctxhttp"
@@ -40,6 +41,7 @@ type dockerAuthorizer struct {
 	credentials func(string) (string, string, error)
 
 	client *http.Client
+	ua     string
 	mu     sync.Mutex
 
 	auth map[string]string
@@ -54,6 +56,7 @@ func NewAuthorizer(client *http.Client, f func(string) (string, string, error))
 	return &dockerAuthorizer{
 		credentials: f,
 		client:      client,
+		ua:          "containerd/" + version.Version,
 		auth:        map[string]string{},
 	}
 }
@@ -194,11 +197,16 @@ func (a *dockerAuthorizer) fetchTokenWithOAuth(ctx context.Context, to tokenOpti
 		form.Set("password", to.secret)
 	}
 
-	resp, err := ctxhttp.Post(
-		ctx, a.client, to.realm,
-		"application/x-www-form-urlencoded; charset=utf-8",
-		strings.NewReader(form.Encode()),
-	)
+	req, err := http.NewRequest("POST", to.realm, strings.NewReader(form.Encode()))
+	if err != nil {
+		return "", err
+	}
+	req.Header.Set("Content-Type", "application/x-www-form-urlencoded; charset=utf-8")
+	if a.ua != "" {
+		req.Header.Set("User-Agent", a.ua)
+	}
+
+	resp, err := ctxhttp.Do(ctx, a.client, req)
 	if err != nil {
 		return "", err
 	}
@@ -244,6 +252,10 @@ func (a *dockerAuthorizer) fetchToken(ctx context.Context, to tokenOptions) (str
 		return "", err
 	}
 
+	if a.ua != "" {
+		req.Header.Set("User-Agent", a.ua)
+	}
+
 	reqParams := req.URL.Query()
 
 	if to.service != "" {
diff --git a/remotes/docker/resolver.go b/remotes/docker/resolver.go
index 00e1c8556..45282fb51 100644
--- a/remotes/docker/resolver.go
+++ b/remotes/docker/resolver.go
@@ -111,6 +111,7 @@ type dockerResolver struct {
 	auth      Authorizer
 	host      func(string) (string, error)
 	headers   http.Header
+	uagent    string
 	plainHTTP bool
 	client    *http.Client
 	tracker   StatusTracker
@@ -135,16 +136,22 @@ func NewResolver(options ResolverOptions) remotes.Resolver {
 			ocispec.MediaTypeImageManifest,
 			ocispec.MediaTypeImageIndex, "*"}, ", "))
 	}
-	if _, ok := options.Headers["User-Agent"]; !ok {
-		options.Headers.Set("User-Agent", "containerd/"+version.Version)
+	ua := options.Headers.Get("User-Agent")
+	if ua != "" {
+		options.Headers.Del("User-Agent")
+	} else {
+		ua = "containerd/" + version.Version
 	}
+
 	if options.Authorizer == nil {
 		options.Authorizer = NewAuthorizer(options.Client, options.Credentials)
+		options.Authorizer.(*dockerAuthorizer).ua = ua
 	}
 	return &dockerResolver{
 		auth:      options.Authorizer,
 		host:      options.Host,
 		headers:   options.Headers,
+		uagent:    ua,
 		plainHTTP: options.PlainHTTP,
 		client:    options.Client,
 		tracker:   options.Tracker,
@@ -351,6 +358,7 @@ func (r *dockerResolver) Pusher(ctx context.Context, ref string) (remotes.Pusher
 type dockerBase struct {
 	refspec reference.Spec
 	base    url.URL
+	uagent  string
 
 	client *http.Client
 	auth   Authorizer
@@ -382,6 +390,7 @@ func (r *dockerResolver) base(refspec reference.Spec) (*dockerBase, error) {
 	return &dockerBase{
 		refspec: refspec,
 		base:    base,
+		uagent:  r.uagent,
 		client:  r.client,
 		auth:    r.auth,
 	}, nil
@@ -407,6 +416,7 @@ func (r *dockerBase) authorize(ctx context.Context, req *http.Request) error {
 func (r *dockerBase) doRequest(ctx context.Context, req *http.Request) (*http.Response, error) {
 	ctx = log.WithLogger(ctx, log.G(ctx).WithField("url", req.URL.String()))
 	log.G(ctx).WithField("request.headers", req.Header).WithField("request.method", req.Method).Debug("do request")
+	req.Header.Set("User-Agent", r.uagent)
 	if err := r.authorize(ctx, req); err != nil {
 		return nil, errors.Wrap(err, "failed to authorize")
 	}