pull: allow separate sessions for different parts of pull

Signed-off-by: Tonis Tiigi <tonistiigi@gmail.com>

cache/remotecache/registry/registry.go

@@ -37,7 +37,7 @@ func ResolveCacheExporterFunc(sm *session.Manager, hosts docker.RegistryHosts) r
 		if err != nil {
 			return nil, err
 		}
-		remote := resolver.New(hosts, sm, g)
+		remote := resolver.New(hosts, resolver.NewSessionAuthenticator(sm, g))
 		pusher, err := remote.Pusher(ctx, ref)
 		if err != nil {
 			return nil, err
@@ -52,7 +52,7 @@ func ResolveCacheImporterFunc(sm *session.Manager, cs content.Store, hosts docke
 		if err != nil {
 			return nil, specs.Descriptor{}, err
 		}
-		remote := resolver.New(hosts, sm, g)
+		remote := resolver.New(hosts, resolver.NewSessionAuthenticator(sm, g))
 		xref, desc, err := remote.Resolve(ctx, ref)
 		if err != nil {
 			return nil, specs.Descriptor{}, err

source/containerimage/pull.go

@@ -23,6 +23,7 @@ import (
 	"github.com/moby/buildkit/util/leaseutil"
 	"github.com/moby/buildkit/util/progress"
 	"github.com/moby/buildkit/util/pull"
+	"github.com/moby/buildkit/util/resolver"
 	"github.com/moby/buildkit/util/winlayers"
 	digest "github.com/opencontainers/go-digest"
 	"github.com/opencontainers/image-spec/identity"
@@ -79,11 +80,11 @@ func (is *Source) ResolveImageConfig(ctx context.Context, ref string, opt llb.Re
 
 	res, err := is.g.Do(ctx, key, func(ctx context.Context) (interface{}, error) {
 		dgst, dt, err := imageutil.Config(ctx, ref, pull.NewResolver(g, pull.ResolverOpt{
-			Hosts:          is.RegistryHosts,
-			SessionManager: sm,
-			ImageStore:     is.ImageStore,
-			Mode:           rm,
-			Ref:            ref,
+			Hosts:      is.RegistryHosts,
+			Auth:       resolver.NewSessionAuthenticator(sm, g),
+			ImageStore: is.ImageStore,
+			Mode:       rm,
+			Ref:        ref,
 		}), is.ContentStore, is.LeaseManager, opt.Platform)
 		if err != nil {
 			return nil, err
@@ -122,11 +123,11 @@ func (is *Source) Resolve(ctx context.Context, id source.Identifier, sm *session
 		id:            imageIdentifier,
 		LeaseManager:  is.LeaseManager,
 		ResolverOpt: pull.ResolverOpt{
-			Hosts:          is.RegistryHosts,
-			SessionManager: sm,
-			ImageStore:     is.ImageStore,
-			Mode:           imageIdentifier.ResolveMode,
-			Ref:            imageIdentifier.Reference.String(),
+			Hosts:      is.RegistryHosts,
+			Auth:       resolver.NewSessionAuthenticator(sm, nil),
+			ImageStore: is.ImageStore,
+			Mode:       imageIdentifier.ResolveMode,
+			Ref:        imageIdentifier.Reference.String(),
 		},
 	}
 	return p, nil
@@ -160,6 +161,7 @@ func mainManifestKey(ctx context.Context, desc specs.Descriptor, platform specs.
 }
 
 func (p *puller) CacheKey(ctx context.Context, g session.Group, index int) (string, bool, error) {
+	p.ResolverOpt.Auth.SetSession(g)
 	if p.Puller.Resolver == nil {
 		p.Puller.Resolver = pull.NewResolver(g, p.ResolverOpt)
 	}
@@ -199,6 +201,7 @@ func (p *puller) CacheKey(ctx context.Context, g session.Group, index int) (stri
 }
 
 func (p *puller) Snapshot(ctx context.Context, g session.Group) (ir cache.ImmutableRef, err error) {
+	p.ResolverOpt.Auth.SetSession(g)
 	if p.Puller.Resolver == nil {
 		p.Puller.Resolver = pull.NewResolver(g, p.ResolverOpt)
 	}

util/pull/resolver.go

@@ -23,19 +23,19 @@ func init() {
 }
 
 type ResolverOpt struct {
-	Hosts          docker.RegistryHosts
-	SessionManager *session.Manager
-	ImageStore     images.Store
-	Mode           source.ResolveMode
-	Ref            string
+	Hosts      docker.RegistryHosts
+	Auth       *resolver.SessionAuthenticator
+	ImageStore images.Store
+	Mode       source.ResolveMode
+	Ref        string
 }
 
 func NewResolver(g session.Group, opt ResolverOpt) remotes.Resolver {
 	if res := cache.Get(opt.Ref, g); res != nil {
 		return withLocal(res, opt.ImageStore, opt.Mode)
 	}
 
-	r := resolver.New(opt.Hosts, opt.SessionManager, g)
+	r := resolver.New(opt.Hosts, opt.Auth)
 	r = cache.Add(opt.Ref, r, g)
 
 	return withLocal(r, opt.ImageStore, opt.Mode)

util/push/push.go

@@ -42,7 +42,7 @@ func Push(ctx context.Context, sm *session.Manager, sid string, cs content.Store
 		ref = reference.TagNameOnly(parsed).String()
 	}
 
-	resolver := resolver.New(hosts, sm, session.NewGroup(sid))
+	resolver := resolver.New(hosts, resolver.NewSessionAuthenticator(sm, session.NewGroup(sid)))
 
 	pusher, err := resolver.Pusher(ctx, ref)
 	if err != nil {

util/resolver/resolver.go

@@ -10,6 +10,7 @@ import (
 	"path/filepath"
 	"runtime"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/containerd/containerd/remotes"
@@ -148,13 +149,36 @@ func NewRegistryConfig(m map[string]config.RegistryConfig) docker.RegistryHosts
 	)
 }
 
-func New(hosts docker.RegistryHosts, sm *session.Manager, g session.Group) remotes.Resolver {
+type SessionAuthenticator struct {
+	sm *session.Manager
+	g  session.Group
+	mu sync.Mutex
+}
+
+func NewSessionAuthenticator(sm *session.Manager, g session.Group) *SessionAuthenticator {
+	return &SessionAuthenticator{sm: sm, g: g}
+}
+
+func (a *SessionAuthenticator) credentials(h string) (string, string, error) {
+	a.mu.Lock()
+	g := a.g
+	a.mu.Unlock()
+	return auth.CredentialsFunc(a.sm, g)(h)
+}
+
+func (a *SessionAuthenticator) SetSession(g session.Group) {
+	a.mu.Lock()
+	a.g = g
+	a.mu.Unlock()
+}
+
+func New(hosts docker.RegistryHosts, auth *SessionAuthenticator) remotes.Resolver {
 	return docker.NewResolver(docker.ResolverOptions{
-		Hosts: hostsWithCredentials(hosts, sm, g),
+		Hosts: hostsWithCredentials(hosts, auth),
 	})
 }
 
-func hostsWithCredentials(hosts docker.RegistryHosts, sm *session.Manager, g session.Group) docker.RegistryHosts {
+func hostsWithCredentials(hosts docker.RegistryHosts, auth *SessionAuthenticator) docker.RegistryHosts {
 	return func(domain string) ([]docker.RegistryHost, error) {
 		res, err := hosts(domain)
 		if err != nil {
@@ -166,7 +190,7 @@ func hostsWithCredentials(hosts docker.RegistryHosts, sm *session.Manager, g ses
 
 		a := docker.NewDockerAuthorizer(
 			docker.WithAuthClient(res[0].Client),
-			docker.WithAuthCreds(auth.CredentialsFunc(sm, g)),
+			docker.WithAuthCreds(auth.credentials),
 		)
 		for i := range res {
 			res[i].Authorizer = a