pull: fix session updating on resolver

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

source/containerimage/pull.go

@@ -161,9 +161,10 @@ 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)
+	} else {
+		p.ResolverOpt.Auth.AddSession(g)
 	}
 	_, desc, err := p.Puller.Resolve(ctx)
 	if err != nil {
@@ -201,9 +202,10 @@ 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)
+	} else {
+		p.ResolverOpt.Auth.AddSession(g)
 	}
 
 	layerNeedsTypeWindows := false

util/pull/resolver.go

@@ -36,7 +36,7 @@ func NewResolver(g session.Group, opt ResolverOpt) remotes.Resolver {
 	}
 
 	r := resolver.New(opt.Hosts, opt.Auth)
-	r = cache.Add(opt.Ref, r, g)
+	r = cache.Add(opt.Ref, r, opt.Auth, g)
 
 	return withLocal(r, opt.ImageStore, opt.Mode)
 }
@@ -117,6 +117,7 @@ type cachedResolver struct {
 	counter int64
 	timeout time.Time
 	remotes.Resolver
+	auth *resolver.SessionAuthenticator
 }
 
 func (cr *cachedResolver) Resolve(ctx context.Context, ref string) (name string, desc ocispec.Descriptor, err error) {
@@ -124,26 +125,23 @@ func (cr *cachedResolver) Resolve(ctx context.Context, ref string) (name string,
 	return cr.Resolver.Resolve(ctx, ref)
 }
 
-func (r *resolverCache) Add(ref string, resolver remotes.Resolver, g session.Group) remotes.Resolver {
+func (r *resolverCache) Add(ref string, resolver remotes.Resolver, auth *resolver.SessionAuthenticator, g session.Group) *cachedResolver {
 	r.mu.Lock()
 	defer r.mu.Unlock()
 
-	res := &cachedResolver{Resolver: resolver, timeout: time.Now().Add(time.Minute)}
+	ref = r.repo(ref)
 
-	for _, sid := range session.AllSessionIDs(g) {
-		ref = r.repo(ref) + "-" + sid
-
-		cr, ok := r.m[ref]
-		res = &cr
-		cr.timeout = time.Now().Add(time.Minute)
-		if ok {
-			continue
-		}
-
-		cr.Resolver = resolver
-		r.m[ref] = cr
+	cr, ok := r.m[ref]
+	cr.timeout = time.Now().Add(time.Minute)
+	if ok {
+		cr.auth.AddSession(g)
+		return &cr
 	}
-	return res
+
+	cr.Resolver = resolver
+	cr.auth = auth
+	r.m[ref] = cr
+	return &cr
 }
 
 func (r *resolverCache) repo(refStr string) string {
@@ -154,17 +152,16 @@ func (r *resolverCache) repo(refStr string) string {
 	return ref.Name()
 }
 
-func (r *resolverCache) Get(ref string, g session.Group) remotes.Resolver {
+func (r *resolverCache) Get(ref string, g session.Group) *cachedResolver {
 	r.mu.Lock()
 	defer r.mu.Unlock()
 
-	for _, sid := range session.AllSessionIDs(g) {
-		ref = r.repo(ref) + "-" + sid
+	ref = r.repo(ref)
 
-		cr, ok := r.m[ref]
-		if ok {
-			return &cr
-		}
+	cr, ok := r.m[ref]
+	if ok {
+		cr.auth.AddSession(g)
+		return &cr
 	}
 	return nil
 }

util/resolver/resolver.go

@@ -150,25 +150,34 @@ func NewRegistryConfig(m map[string]config.RegistryConfig) docker.RegistryHosts
 }
 
 type SessionAuthenticator struct {
-	sm *session.Manager
-	g  session.Group
-	mu sync.Mutex
+	sm     *session.Manager
+	groups []session.Group
+	mu     sync.RWMutex
 }
 
 func NewSessionAuthenticator(sm *session.Manager, g session.Group) *SessionAuthenticator {
-	return &SessionAuthenticator{sm: sm, g: g}
+	return &SessionAuthenticator{sm: sm, groups: []session.Group{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)
+	a.mu.RLock()
+	defer a.mu.RUnlock()
+
+	var err error
+	for i := len(a.groups) - 1; i >= 0; i-- {
+		var user, secret string
+		user, secret, err = auth.CredentialsFunc(a.sm, a.groups[i])(h)
+		if err != nil {
+			continue
+		}
+		return user, secret, nil
+	}
+	return "", "", err
 }
 
-func (a *SessionAuthenticator) SetSession(g session.Group) {
+func (a *SessionAuthenticator) AddSession(g session.Group) {
 	a.mu.Lock()
-	a.g = g
+	a.groups = append(a.groups, g)
 	a.mu.Unlock()
 }