From f83447b516a566f69e395f941056e514aacc37c8 Mon Sep 17 00:00:00 2001
From: Tonis Tiigi <tonistiigi@gmail.com>
Date: Tue, 17 Oct 2023 18:22:59 -0700
Subject: [PATCH] solver: fix possible concurrent map access on cache export

Signed-off-by: Tonis Tiigi <tonistiigi@gmail.com>
(cherry picked from commit e99bfa9839f265c5c2f00f25f2551d824426a4f7)
---
 solver/cachekey.go      |  4 +++-
 solver/combinedcache.go | 43 ++++++++++++++++++++++++-----------------
 solver/exporter.go      | 13 +++++++------
 3 files changed, 35 insertions(+), 25 deletions(-)

diff --git a/solver/cachekey.go b/solver/cachekey.go
index 398368716ad2..9617789de052 100644
--- a/solver/cachekey.go
+++ b/solver/cachekey.go
@@ -55,15 +55,17 @@ func (ck *CacheKey) Output() Index {
 }
 
 func (ck *CacheKey) clone() *CacheKey {
+	ck.mu.RLock()
 	nk := &CacheKey{
 		ID:     ck.ID,
 		digest: ck.digest,
 		vtx:    ck.vtx,
 		output: ck.output,
-		ids:    map[*cacheManager]string{},
+		ids:    make(map[*cacheManager]string, len(ck.ids)),
 	}
 	for cm, id := range ck.ids {
 		nk.ids[cm] = id
 	}
+	ck.mu.RUnlock()
 	return nk
 }
diff --git a/solver/combinedcache.go b/solver/combinedcache.go
index ffffbb39536a..aed8aa90d584 100644
--- a/solver/combinedcache.go
+++ b/solver/combinedcache.go
@@ -101,34 +101,41 @@ func (cm *combinedCacheManager) Save(key *CacheKey, s Result, createdAt time.Tim
 }
 
 func (cm *combinedCacheManager) Records(ctx context.Context, ck *CacheKey) ([]*CacheRecord, error) {
+	ck.mu.RLock()
 	if len(ck.ids) == 0 {
+		ck.mu.RUnlock()
 		return nil, errors.Errorf("no results")
 	}
 
+	cms := make([]*cacheManager, 0, len(ck.ids))
+	for cm := range ck.ids {
+		cms = append(cms, cm)
+	}
+	ck.mu.RUnlock()
+
 	records := map[string]*CacheRecord{}
 	var mu sync.Mutex
 
 	eg, _ := errgroup.WithContext(context.TODO())
-	for c := range ck.ids {
-		func(c *cacheManager) {
-			eg.Go(func() error {
-				recs, err := c.Records(ctx, ck)
-				if err != nil {
-					return err
-				}
-				mu.Lock()
-				for _, rec := range recs {
-					if _, ok := records[rec.ID]; !ok || c == cm.main {
-						if c == cm.main {
-							rec.Priority = 1
-						}
-						records[rec.ID] = rec
+	for _, c := range cms {
+		c := c
+		eg.Go(func() error {
+			recs, err := c.Records(ctx, ck)
+			if err != nil {
+				return err
+			}
+			mu.Lock()
+			for _, rec := range recs {
+				if _, ok := records[rec.ID]; !ok || c == cm.main {
+					if c == cm.main {
+						rec.Priority = 1
 					}
+					records[rec.ID] = rec
 				}
-				mu.Unlock()
-				return nil
-			})
-		}(c)
+			}
+			mu.Unlock()
+			return nil
+		})
 	}
 
 	if err := eg.Wait(); err != nil {
diff --git a/solver/exporter.go b/solver/exporter.go
index 78ce77c2d2f5..44e788ee202b 100644
--- a/solver/exporter.go
+++ b/solver/exporter.go
@@ -85,8 +85,9 @@ func (e *exporter) ExportTo(ctx context.Context, t CacheExporterTarget, opt Cach
 		r        CacheExporterRecord
 		selector digest.Digest
 	}
+	k := e.k.clone() // protect against *CacheKey internal ids mutation from other exports
 
-	recKey := rootKey(e.k.Digest(), e.k.Output())
+	recKey := rootKey(k.Digest(), k.Output())
 	rec := t.Add(recKey)
 	allRec := []CacheExporterRecord{rec}
 
@@ -97,7 +98,7 @@ func (e *exporter) ExportTo(ctx context.Context, t CacheExporterTarget, opt Cach
 	}
 
 	exportRecord := opt.ExportRoots
-	if len(e.k.Deps()) > 0 {
+	if len(deps) > 0 {
 		exportRecord = true
 	}
 
@@ -126,7 +127,7 @@ func (e *exporter) ExportTo(ctx context.Context, t CacheExporterTarget, opt Cach
 		if opt.CompressionOpt != nil {
 			for _, r := range remotes { // record all remaining remotes as well
 				rec := t.Add(recKey)
-				rec.AddResult(e.k.vtx, int(e.k.output), v.CreatedAt, r)
+				rec.AddResult(k.vtx, int(k.output), v.CreatedAt, r)
 				variants = append(variants, rec)
 			}
 		}
@@ -147,7 +148,7 @@ func (e *exporter) ExportTo(ctx context.Context, t CacheExporterTarget, opt Cach
 			if opt.CompressionOpt != nil {
 				for _, r := range remotes { // record all remaining remotes as well
 					rec := t.Add(recKey)
-					rec.AddResult(e.k.vtx, int(e.k.output), v.CreatedAt, r)
+					rec.AddResult(k.vtx, int(k.output), v.CreatedAt, r)
 					variants = append(variants, rec)
 				}
 			}
@@ -155,7 +156,7 @@ func (e *exporter) ExportTo(ctx context.Context, t CacheExporterTarget, opt Cach
 
 		if remote != nil {
 			for _, rec := range allRec {
-				rec.AddResult(e.k.vtx, int(e.k.output), v.CreatedAt, remote)
+				rec.AddResult(k.vtx, int(k.output), v.CreatedAt, remote)
 			}
 		}
 		allRec = append(allRec, variants...)
@@ -198,7 +199,7 @@ func (e *exporter) ExportTo(ctx context.Context, t CacheExporterTarget, opt Cach
 			}
 		}
 
-		for cm, id := range e.k.ids {
+		for cm, id := range k.ids {
 			if _, err := addBacklinks(t, rec, cm, id, bkm); err != nil {
 				return nil, err
 			}