Update resolver code

Signed-off-by: Derek McGowan <derek@mcgstyle.net>

remotes/docker/resolver.go

@@ -18,18 +18,13 @@ package docker
 
 import (
 	"context"
-	"encoding/json"
-	"fmt"
-	"io"
-	"io/ioutil"
 	"net/http"
 	"net/url"
 	"path"
 	"strconv"
 	"strings"
-	"sync"
-	"time"
 
+	"github.com/containerd/containerd/errdefs"
 	"github.com/containerd/containerd/images"
 	"github.com/containerd/containerd/log"
 	"github.com/containerd/containerd/reference"
@@ -51,19 +46,37 @@ var (
 	ErrInvalidAuthorization = errors.New("authorization failed")
 )
 
-type dockerResolver struct {
-	credentials func(string) (string, string, error)
-	host        func(string) (string, error)
-	plainHTTP   bool
-	client      *http.Client
-	tracker     StatusTracker
+// Authorizer is used to authorize HTTP requests based on 401 HTTP responses.
+// An Authorizer is responsible for caching tokens or credentials used by
+// requests.
+type Authorizer interface {
+	// Authorize sets the appropriate `Authorization` header on the given
+	// request.
+	//
+	// If no authorization is found for the request, the request remains
+	// unmodified. It may also add an `Authorization` header as
+	//  "bearer <some bearer token>"
+	//  "basic <base64 encoded credentials>"
+	Authorize(context.Context, *http.Request) error
+
+	// AddResponses adds a 401 response for the authorizer to consider when
+	// authorizing requests. The last response should be unauthorized and
+	// the previous requests are used to consider redirects and retries
+	// that may have led to the 401.
+	//
+	// If response is not handled, returns `ErrNotImplemented`
+	AddResponses(context.Context, []*http.Response) error
 }
 
 // ResolverOptions are used to configured a new Docker register resolver
 type ResolverOptions struct {
+	// Authorizer is used to authorize registry requests
+	Authorizer Authorizer
+
 	// Credentials provides username and secret given a host.
 	// If username is empty but a secret is given, that secret
 	// is interpretted as a long lived token.
+	// Deprecated: use Authorizer
 	Credentials func(string) (string, string, error)
 
 	// Host provides the hostname given a namespace.
@@ -89,22 +102,31 @@ func DefaultHost(ns string) (string, error) {
 	return ns, nil
 }
 
+type dockerResolver struct {
+	auth      Authorizer
+	host      func(string) (string, error)
+	plainHTTP bool
+	client    *http.Client
+	tracker   StatusTracker
+}
+
 // NewResolver returns a new resolver to a Docker registry
 func NewResolver(options ResolverOptions) remotes.Resolver {
-	tracker := options.Tracker
-	if tracker == nil {
-		tracker = NewInMemoryTracker()
+	if options.Tracker == nil {
+		options.Tracker = NewInMemoryTracker()
 	}
-	host := options.Host
-	if host == nil {
-		host = DefaultHost
+	if options.Host == nil {
+		options.Host = DefaultHost
+	}
+	if options.Authorizer == nil {
+		options.Authorizer = NewAuthorizer(options.Client, options.Credentials)
 	}
 	return &dockerResolver{
-		credentials: options.Credentials,
-		host:        host,
-		plainHTTP:   options.PlainHTTP,
-		client:      options.Client,
-		tracker:     tracker,
+		auth:      options.Authorizer,
+		host:      options.Host,
+		plainHTTP: options.PlainHTTP,
+		client:    options.Client,
+		tracker:   options.Tracker,
 	}
 }
 
@@ -272,18 +294,14 @@ type dockerBase struct {
 	refspec reference.Spec
 	base    url.URL
 
-	client           *http.Client
-	useBasic         bool
-	username, secret string
-	token            string
-	mu               sync.Mutex
+	client *http.Client
+	auth   Authorizer
 }
 
 func (r *dockerResolver) base(refspec reference.Spec) (*dockerBase, error) {
 	var (
-		err              error
-		base             url.URL
-		username, secret string
+		err  error
+		base url.URL
 	)
 
 	host := refspec.Hostname()
@@ -300,61 +318,40 @@ func (r *dockerResolver) base(refspec reference.Spec) (*dockerBase, error) {
 		base.Scheme = "http"
 	}
 
-	if r.credentials != nil {
-		username, secret, err = r.credentials(base.Host)
-		if err != nil {
-			return nil, err
-		}
-	}
-
 	prefix := strings.TrimPrefix(refspec.Locator, host+"/")
 	base.Path = path.Join("/v2", prefix)
 
 	return &dockerBase{
-		refspec:  refspec,
-		base:     base,
-		client:   r.client,
-		username: username,
-		secret:   secret,
+		refspec: refspec,
+		base:    base,
+		client:  r.client,
+		auth:    r.auth,
 	}, nil
 }
 
-func (r *dockerBase) getToken() string {
-	r.mu.Lock()
-	defer r.mu.Unlock()
-
-	return r.token
-}
-
-func (r *dockerBase) setToken(token string) bool {
-	r.mu.Lock()
-	defer r.mu.Unlock()
-
-	changed := r.token != token
-	r.token = token
-
-	return changed
-}
-
 func (r *dockerBase) url(ps ...string) string {
 	url := r.base
 	url.Path = path.Join(url.Path, path.Join(ps...))
 	return url.String()
 }
 
-func (r *dockerBase) authorize(req *http.Request) {
-	token := r.getToken()
-	if r.useBasic {
-		req.SetBasicAuth(r.username, r.secret)
-	} else if token != "" {
-		req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", token))
+func (r *dockerBase) authorize(ctx context.Context, req *http.Request) error {
+	// Check if has header for host
+	if r.auth != nil {
+		if err := r.auth.Authorize(ctx, req); err != nil {
+			return err
+		}
 	}
+
+	return nil
 }
 
 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")
-	r.authorize(req)
+	if err := r.authorize(ctx, req); err != nil {
+		return nil, errors.Wrap(err, "failed to authorize")
+	}
 	resp, err := ctxhttp.Do(ctx, r.client, req)
 	if err != nil {
 		return nil, errors.Wrap(err, "failed to do request")
@@ -392,23 +389,14 @@ func (r *dockerBase) retryRequest(ctx context.Context, req *http.Request, respon
 	last := responses[len(responses)-1]
 	if last.StatusCode == http.StatusUnauthorized {
 		log.G(ctx).WithField("header", last.Header.Get("WWW-Authenticate")).Debug("Unauthorized")
-		for _, c := range parseAuthHeader(last.Header) {
-			if c.scheme == bearerAuth {
-				if err := invalidAuthorization(c, responses); err != nil {
-					r.setToken("")
-					return nil, err
-				}
-				if err := r.setTokenAuth(ctx, c.parameters); err != nil {
-					return nil, err
-				}
-				return copyRequest(req)
-			} else if c.scheme == basicAuth {
-				if r.username != "" && r.secret != "" {
-					r.useBasic = true
-				}
+		if r.auth != nil {
+			if err := r.auth.AddResponses(ctx, responses); err == nil {
 				return copyRequest(req)
+			} else if !errdefs.IsNotImplemented(err) {
+				return nil, err
 			}
 		}
+
 		return nil, nil
 	} else if last.StatusCode == http.StatusMethodNotAllowed && req.Method == http.MethodHead {
 		// Support registries which have not properly implemented the HEAD method for
@@ -424,30 +412,6 @@ func (r *dockerBase) retryRequest(ctx context.Context, req *http.Request, respon
 	return nil, nil
 }
 
-func invalidAuthorization(c challenge, responses []*http.Response) error {
-	errStr := c.parameters["error"]
-	if errStr == "" {
-		return nil
-	}
-
-	n := len(responses)
-	if n == 1 || (n > 1 && !sameRequest(responses[n-2].Request, responses[n-1].Request)) {
-		return nil
-	}
-
-	return errors.Wrapf(ErrInvalidAuthorization, "server message: %s", errStr)
-}
-
-func sameRequest(r1, r2 *http.Request) bool {
-	if r1.Method != r2.Method {
-		return false
-	}
-	if *r1.URL != *r2.URL {
-		return false
-	}
-	return true
-}
-
 func copyRequest(req *http.Request) (*http.Request, error) {
 	ireq := *req
 	if ireq.GetBody != nil {
@@ -459,167 +423,3 @@ func copyRequest(req *http.Request) (*http.Request, error) {
 	}
 	return &ireq, nil
 }
-
-func (r *dockerBase) setTokenAuth(ctx context.Context, params map[string]string) error {
-	realm, ok := params["realm"]
-	if !ok {
-		return errors.New("no realm specified for token auth challenge")
-	}
-
-	realmURL, err := url.Parse(realm)
-	if err != nil {
-		return fmt.Errorf("invalid token auth challenge realm: %s", err)
-	}
-
-	to := tokenOptions{
-		realm:   realmURL.String(),
-		service: params["service"],
-	}
-
-	to.scopes = getTokenScopes(ctx, params)
-	if len(to.scopes) == 0 {
-		return errors.Errorf("no scope specified for token auth challenge")
-	}
-
-	var token string
-	if r.secret != "" {
-		// Credential information is provided, use oauth POST endpoint
-		token, err = r.fetchTokenWithOAuth(ctx, to)
-		if err != nil {
-			return errors.Wrap(err, "failed to fetch oauth token")
-		}
-	} else {
-		// Do request anonymously
-		token, err = r.fetchToken(ctx, to)
-		if err != nil {
-			return errors.Wrap(err, "failed to fetch anonymous token")
-		}
-	}
-	r.setToken(token)
-
-	return nil
-}
-
-type tokenOptions struct {
-	realm   string
-	service string
-	scopes  []string
-}
-
-type postTokenResponse struct {
-	AccessToken  string    `json:"access_token"`
-	RefreshToken string    `json:"refresh_token"`
-	ExpiresIn    int       `json:"expires_in"`
-	IssuedAt     time.Time `json:"issued_at"`
-	Scope        string    `json:"scope"`
-}
-
-func (r *dockerBase) fetchTokenWithOAuth(ctx context.Context, to tokenOptions) (string, error) {
-	form := url.Values{}
-	form.Set("scope", strings.Join(to.scopes, " "))
-	form.Set("service", to.service)
-	// TODO: Allow setting client_id
-	form.Set("client_id", "containerd-dist-tool")
-
-	if r.username == "" {
-		form.Set("grant_type", "refresh_token")
-		form.Set("refresh_token", r.secret)
-	} else {
-		form.Set("grant_type", "password")
-		form.Set("username", r.username)
-		form.Set("password", r.secret)
-	}
-
-	resp, err := ctxhttp.PostForm(ctx, r.client, to.realm, form)
-	if err != nil {
-		return "", err
-	}
-	defer resp.Body.Close()
-
-	// Registries without support for POST may return 404 for POST /v2/token.
-	// As of September 2017, GCR is known to return 404.
-	// As of February 2018, JFrog Artifactory is known to return 401.
-	if (resp.StatusCode == 405 && r.username != "") || resp.StatusCode == 404 || resp.StatusCode == 401 {
-		return r.fetchToken(ctx, to)
-	} else if resp.StatusCode < 200 || resp.StatusCode >= 400 {
-		b, _ := ioutil.ReadAll(io.LimitReader(resp.Body, 64000)) // 64KB
-		log.G(ctx).WithFields(logrus.Fields{
-			"status": resp.Status,
-			"body":   string(b),
-		}).Debugf("token request failed")
-		// TODO: handle error body and write debug output
-		return "", errors.Errorf("unexpected status: %s", resp.Status)
-	}
-
-	decoder := json.NewDecoder(resp.Body)
-
-	var tr postTokenResponse
-	if err = decoder.Decode(&tr); err != nil {
-		return "", fmt.Errorf("unable to decode token response: %s", err)
-	}
-
-	return tr.AccessToken, nil
-}
-
-type getTokenResponse struct {
-	Token        string    `json:"token"`
-	AccessToken  string    `json:"access_token"`
-	ExpiresIn    int       `json:"expires_in"`
-	IssuedAt     time.Time `json:"issued_at"`
-	RefreshToken string    `json:"refresh_token"`
-}
-
-// getToken fetches a token using a GET request
-func (r *dockerBase) fetchToken(ctx context.Context, to tokenOptions) (string, error) {
-	req, err := http.NewRequest("GET", to.realm, nil)
-	if err != nil {
-		return "", err
-	}
-
-	reqParams := req.URL.Query()
-
-	if to.service != "" {
-		reqParams.Add("service", to.service)
-	}
-
-	for _, scope := range to.scopes {
-		reqParams.Add("scope", scope)
-	}
-
-	if r.secret != "" {
-		req.SetBasicAuth(r.username, r.secret)
-	}
-
-	req.URL.RawQuery = reqParams.Encode()
-
-	resp, err := ctxhttp.Do(ctx, r.client, req)
-	if err != nil {
-		return "", err
-	}
-	defer resp.Body.Close()
-
-	if resp.StatusCode < 200 || resp.StatusCode >= 400 {
-		// TODO: handle error body and write debug output
-		return "", errors.Errorf("unexpected status: %s", resp.Status)
-	}
-
-	decoder := json.NewDecoder(resp.Body)
-
-	var tr getTokenResponse
-	if err = decoder.Decode(&tr); err != nil {
-		return "", fmt.Errorf("unable to decode token response: %s", err)
-	}
-
-	// `access_token` is equivalent to `token` and if both are specified
-	// the choice is undefined.  Canonicalize `access_token` by sticking
-	// things in `token`.
-	if tr.AccessToken != "" {
-		tr.Token = tr.AccessToken
-	}
-
-	if tr.Token == "" {
-		return "", ErrNoToken
-	}
-
-	return tr.Token, nil
-}