remotes: support cross-repo-push

With distribution source label in content store, select the longest
common prefix components as condidate mount blob source and try to push
with mount blob.

Fix #2964

Signed-off-by: Wei Fu <fuweid89@gmail.com>

remotes/docker/authorizer.go

@@ -44,7 +44,8 @@ type dockerAuthorizer struct {
 	ua     string
 	mu     sync.Mutex
 
-	auth map[string]string
+	// indexed by host name
+	handlers map[string]*authHandler
 }
 
 // NewAuthorizer creates a Docker authorizer using the provided function to
@@ -53,116 +54,226 @@ func NewAuthorizer(client *http.Client, f func(string) (string, string, error))
 	if client == nil {
 		client = http.DefaultClient
 	}
+
 	return &dockerAuthorizer{
 		credentials: f,
 		client:      client,
 		ua:          "containerd/" + version.Version,
-		auth:        map[string]string{},
+		handlers:    make(map[string]*authHandler),
 	}
 }
 
+// Authorize handles auth request.
 func (a *dockerAuthorizer) Authorize(ctx context.Context, req *http.Request) error {
-	// TODO: Lookup matching challenge and scope rather than just host
+	// skip if there is no auth handler
-	if auth := a.getAuth(req.URL.Host); auth != "" {
+	ah := a.getAuthHandler(req.URL.Host)
-		req.Header.Set("Authorization", auth)
+	if ah == nil {
+		return nil
 	}
 
+	auth, err := ah.authorize(ctx)
+	if err != nil {
+		return err
+	}
+
+	req.Header.Set("Authorization", auth)
 	return nil
 }
 
+func (a *dockerAuthorizer) getAuthHandler(host string) *authHandler {
+	a.mu.Lock()
+	defer a.mu.Unlock()
+
+	return a.handlers[host]
+}
+
 func (a *dockerAuthorizer) AddResponses(ctx context.Context, responses []*http.Response) error {
 	last := responses[len(responses)-1]
 	host := last.Request.URL.Host
+
+	a.mu.Lock()
+	defer a.mu.Unlock()
 	for _, c := range parseAuthHeader(last.Header) {
 		if c.scheme == bearerAuth {
 			if err := invalidAuthorization(c, responses); err != nil {
-				// TODO: Clear token
+				delete(a.handlers, host)
-				a.setAuth(host, "")
 				return err
 			}
 
-			// TODO(dmcg): Store challenge, not token
+			// reuse existing handler
-			// Move token fetching to authorize
+			//
-			return a.setTokenAuth(ctx, host, c.parameters)
+			// assume that one registry will return the common
+			// challenge information, including realm and service.
+			// and the resource scope is only different part
+			// which can be provided by each request.
+			if _, ok := a.handlers[host]; ok {
+				return nil
+			}
+
+			common, err := a.generateTokenOptions(ctx, host, c)
+			if err != nil {
+				return err
+			}
+
+			a.handlers[host] = newAuthHandler(a.client, a.ua, c.scheme, common)
+			return nil
 		} else if c.scheme == basicAuth && a.credentials != nil {
-			// TODO: Resolve credentials on authorize
 			username, secret, err := a.credentials(host)
 			if err != nil {
 				return err
 			}
+
 			if username != "" && secret != "" {
-				auth := username + ":" + secret
+				common := tokenOptions{
-				a.setAuth(host, fmt.Sprintf("Basic %s", base64.StdEncoding.EncodeToString([]byte(auth))))
+					username: username,
+					secret:   secret,
+				}
+
+				a.handlers[host] = newAuthHandler(a.client, a.ua, c.scheme, common)
 				return nil
 			}
 		}
 	}
-
 	return errors.Wrap(errdefs.ErrNotImplemented, "failed to find supported auth scheme")
 }
 
-func (a *dockerAuthorizer) getAuth(host string) string {
+func (a *dockerAuthorizer) generateTokenOptions(ctx context.Context, host string, c challenge) (tokenOptions, error) {
-	a.mu.Lock()
+	realm, ok := c.parameters["realm"]
-	defer a.mu.Unlock()
-
-	return a.auth[host]
-}
-
-func (a *dockerAuthorizer) setAuth(host string, auth string) bool {
-	a.mu.Lock()
-	defer a.mu.Unlock()
-
-	changed := a.auth[host] != auth
-	a.auth[host] = auth
-
-	return changed
-}
-
-func (a *dockerAuthorizer) setTokenAuth(ctx context.Context, host string, params map[string]string) error {
-	realm, ok := params["realm"]
 	if !ok {
-		return errors.New("no realm specified for token auth challenge")
+		return tokenOptions{}, errors.New("no realm specified for token auth challenge")
 	}
 
 	realmURL, err := url.Parse(realm)
 	if err != nil {
-		return errors.Wrap(err, "invalid token auth challenge realm")
+		return tokenOptions{}, errors.Wrap(err, "invalid token auth challenge realm")
 	}
 
 	to := tokenOptions{
 		realm:   realmURL.String(),
-		service: params["service"],
+		service: c.parameters["service"],
 	}
 
-	to.scopes = getTokenScopes(ctx, params)
+	scope, ok := c.parameters["scope"]
-	if len(to.scopes) == 0 {
+	if !ok {
-		return errors.Errorf("no scope specified for token auth challenge")
+		return tokenOptions{}, errors.Errorf("no scope specified for token auth challenge")
 	}
+	to.scopes = append(to.scopes, scope)
 
 	if a.credentials != nil {
 		to.username, to.secret, err = a.credentials(host)
 		if err != nil {
-			return err
+			return tokenOptions{}, err
 		}
 	}
+	return to, nil
+}
 
-	var token string
+// authResult is used to control limit rate.
+type authResult struct {
+	sync.WaitGroup
+	token string
+	err   error
+}
+
+// authHandler is used to handle auth request per registry server.
+type authHandler struct {
+	sync.Mutex
+
+	ua string
+
+	client *http.Client
+
+	// only support basic and bearer schemes
+	scheme authenticationScheme
+
+	// common contains common challenge answer
+	common tokenOptions
+
+	// scopedTokens caches token indexed by scopes, which used in
+	// bearer auth case
+	scopedTokens map[string]*authResult
+}
+
+func newAuthHandler(client *http.Client, ua string, scheme authenticationScheme, opts tokenOptions) *authHandler {
+	if client == nil {
+		client = http.DefaultClient
+	}
+
+	return &authHandler{
+		ua:           ua,
+		client:       client,
+		scheme:       scheme,
+		common:       opts,
+		scopedTokens: map[string]*authResult{},
+	}
+}
+
+func (ah *authHandler) authorize(ctx context.Context) (string, error) {
+	switch ah.scheme {
+	case basicAuth:
+		return ah.doBasicAuth(ctx)
+	case bearerAuth:
+		return ah.doBearerAuth(ctx)
+	default:
+		return "", errors.Wrap(errdefs.ErrNotImplemented, "failed to find supported auth scheme")
+	}
+}
+
+func (ah *authHandler) doBasicAuth(ctx context.Context) (string, error) {
+	username, secret := ah.common.username, ah.common.secret
+
+	if username == "" || secret == "" {
+		return "", fmt.Errorf("failed to handle basic auth because missing username or secret")
+	}
+
+	auth := base64.StdEncoding.EncodeToString([]byte(username + ":" + secret))
+	return fmt.Sprintf("%s %s", "Basic", auth), nil
+}
+
+func (ah *authHandler) doBearerAuth(ctx context.Context) (string, error) {
+	// copy common tokenOptions
+	to := ah.common
+
+	to.scopes = getTokenScopes(ctx, to.scopes)
+	if len(to.scopes) == 0 {
+		return "", errors.Errorf("no scope specified for token auth challenge")
+	}
+
+	// Docs: https://docs.docker.com/registry/spec/auth/scope
+	scoped := strings.Join(to.scopes, " ")
+
+	ah.Lock()
+	if r, exist := ah.scopedTokens[scoped]; exist {
+		ah.Unlock()
+		r.Wait()
+		return r.token, r.err
+	}
+
+	// only one fetch token job
+	r := new(authResult)
+	r.Add(1)
+	ah.scopedTokens[scoped] = r
+	ah.Unlock()
+
+	// fetch token for the resource scope
+	var (
+		token string
+		err   error
+	)
 	if to.secret != "" {
-		// Credential information is provided, use oauth POST endpoint
+		// credential information is provided, use oauth POST endpoint
-		token, err = a.fetchTokenWithOAuth(ctx, to)
+		token, err = ah.fetchTokenWithOAuth(ctx, to)
-		if err != nil {
+		err = errors.Wrap(err, "failed to fetch oauth token")
-			return errors.Wrap(err, "failed to fetch oauth token")
-		}
 	} else {
-		// Do request anonymously
+		// do request anonymously
-		token, err = a.fetchToken(ctx, to)
+		token, err = ah.fetchToken(ctx, to)
-		if err != nil {
+		err = errors.Wrap(err, "failed to fetch anonymous token")
-			return errors.Wrap(err, "failed to fetch anonymous token")
-		}
 	}
-	a.setAuth(host, fmt.Sprintf("Bearer %s", token))
+	token = fmt.Sprintf("%s %s", "Bearer", token)
 
-	return nil
+	r.token, r.err = token, err
+	r.Done()
+	return r.token, r.err
 }
 
 type tokenOptions struct {
@@ -181,7 +292,7 @@ type postTokenResponse struct {
 	Scope        string    `json:"scope"`
 }
 
-func (a *dockerAuthorizer) fetchTokenWithOAuth(ctx context.Context, to tokenOptions) (string, error) {
+func (ah *authHandler) fetchTokenWithOAuth(ctx context.Context, to tokenOptions) (string, error) {
 	form := url.Values{}
 	form.Set("scope", strings.Join(to.scopes, " "))
 	form.Set("service", to.service)
@@ -202,11 +313,11 @@ func (a *dockerAuthorizer) fetchTokenWithOAuth(ctx context.Context, to tokenOpti
 		return "", err
 	}
 	req.Header.Set("Content-Type", "application/x-www-form-urlencoded; charset=utf-8")
-	if a.ua != "" {
+	if ah.ua != "" {
-		req.Header.Set("User-Agent", a.ua)
+		req.Header.Set("User-Agent", ah.ua)
 	}
 
-	resp, err := ctxhttp.Do(ctx, a.client, req)
+	resp, err := ctxhttp.Do(ctx, ah.client, req)
 	if err != nil {
 		return "", err
 	}
@@ -216,7 +327,7 @@ func (a *dockerAuthorizer) fetchTokenWithOAuth(ctx context.Context, to tokenOpti
 	// 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 && to.username != "") || resp.StatusCode == 404 || resp.StatusCode == 401 {
-		return a.fetchToken(ctx, to)
+		return ah.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{
@@ -245,15 +356,15 @@ type getTokenResponse struct {
 	RefreshToken string    `json:"refresh_token"`
 }
 
-// getToken fetches a token using a GET request
+// fetchToken fetches a token using a GET request
-func (a *dockerAuthorizer) fetchToken(ctx context.Context, to tokenOptions) (string, error) {
+func (ah *authHandler) fetchToken(ctx context.Context, to tokenOptions) (string, error) {
 	req, err := http.NewRequest("GET", to.realm, nil)
 	if err != nil {
 		return "", err
 	}
 
-	if a.ua != "" {
+	if ah.ua != "" {
-		req.Header.Set("User-Agent", a.ua)
+		req.Header.Set("User-Agent", ah.ua)
 	}
 
 	reqParams := req.URL.Query()
@@ -272,7 +383,7 @@ func (a *dockerAuthorizer) fetchToken(ctx context.Context, to tokenOptions) (str
 
 	req.URL.RawQuery = reqParams.Encode()
 
-	resp, err := ctxhttp.Do(ctx, a.client, req)
+	resp, err := ctxhttp.Do(ctx, ah.client, req)
 	if err != nil {
 		return "", err
 	}

remotes/docker/handler.go

@@ -110,3 +110,45 @@ func appendDistributionSourceLabel(originLabel, repo string) string {
 func distributionSourceLabelKey(source string) string {
 	return fmt.Sprintf("%s.%s", labelDistributionSource, source)
 }
+
+// selectRepositoryMountCandidate will select the repo which has longest
+// common prefix components as the candidate.
+func selectRepositoryMountCandidate(refspec reference.Spec, sources map[string]string) string {
+	u, err := url.Parse("dummy://" + refspec.Locator)
+	if err != nil {
+		// NOTE: basically, it won't be error here
+		return ""
+	}
+
+	source, target := u.Hostname(), strings.TrimPrefix(u.Path, "/")
+	repoLabel, ok := sources[distributionSourceLabelKey(source)]
+	if !ok || repoLabel == "" {
+		return ""
+	}
+
+	n, match := 0, ""
+	components := strings.Split(target, "/")
+	for _, repo := range strings.Split(repoLabel, ",") {
+		// the target repo is not a candidate
+		if repo == target {
+			continue
+		}
+
+		if l := commonPrefixComponents(components, repo); l >= n {
+			n, match = l, repo
+		}
+	}
+	return match
+}
+
+func commonPrefixComponents(components []string, target string) int {
+	targetComponents := strings.Split(target, "/")
+
+	i := 0
+	for ; i < len(components) && i < len(targetComponents); i++ {
+		if components[i] != targetComponents[i] {
+			break
+		}
+	}
+	return i
+}

remotes/docker/pusher.go

@@ -116,8 +116,6 @@ func (p dockerPusher) Push(ctx context.Context, desc ocispec.Descriptor) (conten
 		}
 	}
 
-	// TODO: Lookup related objects for cross repository push
-
 	if isManifest {
 		var putPath string
 		if p.tag != "" {
@@ -132,21 +130,57 @@ func (p dockerPusher) Push(ctx context.Context, desc ocispec.Descriptor) (conten
 		}
 		req.Header.Add("Content-Type", desc.MediaType)
 	} else {
-		// TODO: Do monolithic upload if size is small
-
 		// Start upload request
 		req, err = http.NewRequest(http.MethodPost, p.url("blobs", "uploads")+"/", nil)
 		if err != nil {
 			return nil, err
 		}
 
-		resp, err := p.doRequestWithRetries(ctx, req, nil)
+		var resp *http.Response
-		if err != nil {
+		if fromRepo := selectRepositoryMountCandidate(p.refspec, desc.Annotations); fromRepo != "" {
-			return nil, err
+			req = requestWithMountFrom(req, desc.Digest.String(), fromRepo)
+			pctx := contextWithAppendPullRepositoryScope(ctx, fromRepo)
+
+			// NOTE: the fromRepo might be private repo and
+			// auth service still can grant token without error.
+			// but the post request will fail because of 401.
+			//
+			// for the private repo, we should remove mount-from
+			// query and send the request again.
+			resp, err = p.doRequest(pctx, req)
+			if err != nil {
+				return nil, err
+			}
+
+			if resp.StatusCode == http.StatusUnauthorized {
+				log.G(ctx).Debugf("failed to mount from repository %s", fromRepo)
+
+				resp.Body.Close()
+				resp = nil
+
+				req, err = removeMountFromQuery(req)
+				if err != nil {
+					return nil, err
+				}
+			}
+		}
+
+		if resp == nil {
+			resp, err = p.doRequestWithRetries(ctx, req, nil)
+			if err != nil {
+				return nil, err
+			}
 		}
 
 		switch resp.StatusCode {
 		case http.StatusOK, http.StatusAccepted, http.StatusNoContent:
+		case http.StatusCreated:
+			p.tracker.SetStatus(ref, Status{
+				Status: content.Status{
+					Ref: ref,
+				},
+			})
+			return nil, errors.Wrapf(errdefs.ErrAlreadyExists, "content %v on remote", desc.Digest)
 		default:
 			// TODO: log error
 			return nil, errors.Errorf("unexpected response: %s", resp.Status)
@@ -320,3 +354,25 @@ func (pw *pushWriter) Truncate(size int64) error {
 	// TODO: always error on manifest
 	return errors.New("cannot truncate remote upload")
 }
+
+func requestWithMountFrom(req *http.Request, mount, from string) *http.Request {
+	q := req.URL.Query()
+
+	q.Set("mount", mount)
+	q.Set("from", from)
+	req.URL.RawQuery = q.Encode()
+	return req
+}
+
+func removeMountFromQuery(req *http.Request) (*http.Request, error) {
+	req, err := copyRequest(req)
+	if err != nil {
+		return nil, err
+	}
+
+	q := req.URL.Query()
+	q.Del("mount")
+	q.Del("from")
+	req.URL.RawQuery = q.Encode()
+	return req, nil
+}

remotes/docker/scope.go

@@ -18,6 +18,7 @@ package docker
 
 import (
 	"context"
+	"fmt"
 	"net/url"
 	"sort"
 	"strings"
@@ -53,24 +54,38 @@ func contextWithRepositoryScope(ctx context.Context, refspec reference.Spec, pus
 	return context.WithValue(ctx, tokenScopesKey{}, []string{s}), nil
 }
 
-// getTokenScopes returns deduplicated and sorted scopes from ctx.Value(tokenScopesKey{}) and params["scope"].
+// contextWithAppendPullRepositoryScope is used to append repository pull
-func getTokenScopes(ctx context.Context, params map[string]string) []string {
+// scope into existing scopes indexed by the tokenScopesKey{}.
+func contextWithAppendPullRepositoryScope(ctx context.Context, repo string) context.Context {
+	var scopes []string
+
+	if v := ctx.Value(tokenScopesKey{}); v != nil {
+		scopes = append(scopes, v.([]string)...)
+	}
+	scopes = append(scopes, fmt.Sprintf("repository:%s:pull", repo))
+	return context.WithValue(ctx, tokenScopesKey{}, scopes)
+}
+
+// getTokenScopes returns deduplicated and sorted scopes from ctx.Value(tokenScopesKey{}) and common scopes.
+func getTokenScopes(ctx context.Context, common []string) []string {
 	var scopes []string
 	if x := ctx.Value(tokenScopesKey{}); x != nil {
 		scopes = append(scopes, x.([]string)...)
 	}
-	if scope, ok := params["scope"]; ok {
+
-		for _, s := range scopes {
+	scopes = append(scopes, common...)
-			// Note: this comparison is unaware of the scope grammar (https://docs.docker.com/registry/spec/auth/scope/)
-			// So, "repository:foo/bar:pull,push" != "repository:foo/bar:push,pull", although semantically they are equal.
-			if s == scope {
-				// already appended
-				goto Sort
-			}
-		}
-		scopes = append(scopes, scope)
-	}
-Sort:
 	sort.Strings(scopes)
-	return scopes
+
+	l := 0
+	for idx := 1; idx < len(scopes); idx++ {
+		// Note: this comparison is unaware of the scope grammar (https://docs.docker.com/registry/spec/auth/scope/)
+		// So, "repository:foo/bar:pull,push" != "repository:foo/bar:push,pull", although semantically they are equal.
+		if scopes[l] == scopes[idx] {
+			continue
+		}
+
+		l++
+		scopes[l] = scopes[idx]
+	}
+	return scopes[:l+1]
 }

remotes/docker/scope_test.go

@@ -17,6 +17,7 @@
 package docker
 
 import (
+	"context"
 	"testing"
 
 	"github.com/containerd/containerd/reference"
@@ -54,3 +55,42 @@ func TestRepositoryScope(t *testing.T) {
 		})
 	}
 }
+
+func TestGetTokenScopes(t *testing.T) {
+	testCases := []struct {
+		scopesInCtx  []string
+		commonScopes []string
+		expected     []string
+	}{
+		{
+			scopesInCtx:  []string{},
+			commonScopes: []string{"repository:foo/bar:pull"},
+			expected:     []string{"repository:foo/bar:pull"},
+		},
+		{
+			scopesInCtx:  []string{"repository:foo/bar:pull,push"},
+			commonScopes: []string{},
+			expected:     []string{"repository:foo/bar:pull,push"},
+		},
+		{
+			scopesInCtx:  []string{"repository:foo/bar:pull"},
+			commonScopes: []string{"repository:foo/bar:pull"},
+			expected:     []string{"repository:foo/bar:pull"},
+		},
+		{
+			scopesInCtx:  []string{"repository:foo/bar:pull"},
+			commonScopes: []string{"repository:foo/bar:pull,push"},
+			expected:     []string{"repository:foo/bar:pull", "repository:foo/bar:pull,push"},
+		},
+		{
+			scopesInCtx:  []string{"repository:foo/bar:pull"},
+			commonScopes: []string{"repository:foo/bar:pull,push", "repository:foo/bar:pull"},
+			expected:     []string{"repository:foo/bar:pull", "repository:foo/bar:pull,push"},
+		},
+	}
+	for _, tc := range testCases {
+		ctx := context.WithValue(context.TODO(), tokenScopesKey{}, tc.scopesInCtx)
+		actual := getTokenScopes(ctx, tc.commonScopes)
+		assert.DeepEqual(t, tc.expected, actual)
+	}
+}

remotes/handlers.go

@@ -156,7 +156,7 @@ func push(ctx context.Context, provider content.Provider, pusher Pusher, desc oc
 //
 // Base handlers can be provided which will be called before any push specific
 // handlers.
-func PushContent(ctx context.Context, pusher Pusher, desc ocispec.Descriptor, provider content.Provider, platform platforms.MatchComparer, wrapper func(h images.Handler) images.Handler) error {
+func PushContent(ctx context.Context, pusher Pusher, desc ocispec.Descriptor, store content.Store, platform platforms.MatchComparer, wrapper func(h images.Handler) images.Handler) error {
 	var m sync.Mutex
 	manifestStack := []ocispec.Descriptor{}
 
@@ -173,10 +173,14 @@ func PushContent(ctx context.Context, pusher Pusher, desc ocispec.Descriptor, pr
 		}
 	})
 
-	pushHandler := PushHandler(pusher, provider)
+	pushHandler := PushHandler(pusher, store)
+
+	platformFilterhandler := images.FilterPlatforms(images.ChildrenHandler(store), platform)
+
+	annotateHandler := annotateDistributionSourceHandler(platformFilterhandler, store)
 
 	var handler images.Handler = images.Handlers(
-		images.FilterPlatforms(images.ChildrenHandler(provider), platform),
+		annotateHandler,
 		filterHandler,
 		pushHandler,
 	)
@@ -241,3 +245,45 @@ func FilterManifestByPlatformHandler(f images.HandlerFunc, m platforms.Matcher)
 		return descs, nil
 	}
 }
+
+// annotateDistributionSourceHandler add distribution source label into
+// annotation of config or blob descriptor.
+func annotateDistributionSourceHandler(f images.HandlerFunc, manager content.Manager) images.HandlerFunc {
+	return func(ctx context.Context, desc ocispec.Descriptor) ([]ocispec.Descriptor, error) {
+		children, err := f(ctx, desc)
+		if err != nil {
+			return nil, err
+		}
+
+		// only add distribution source for the config or blob data descriptor
+		switch desc.MediaType {
+		case images.MediaTypeDockerSchema2Manifest, ocispec.MediaTypeImageManifest,
+			images.MediaTypeDockerSchema2ManifestList, ocispec.MediaTypeImageIndex:
+		default:
+			return children, nil
+		}
+
+		for i := range children {
+			child := children[i]
+
+			info, err := manager.Info(ctx, child.Digest)
+			if err != nil {
+				return nil, err
+			}
+
+			for k, v := range info.Labels {
+				if !strings.HasPrefix(k, "containerd.io/distribution.source.") {
+					continue
+				}
+
+				if child.Annotations == nil {
+					child.Annotations = map[string]string{}
+				}
+				child.Annotations[k] = v
+			}
+
+			children[i] = child
+		}
+		return children, nil
+	}
+}