From c1c4af0e6e426f64160fc3285df862ed201f998d Mon Sep 17 00:00:00 2001
From: Jon Johnson <jon.johnson@chainguard.dev>
Date: Mon, 6 Nov 2023 12:39:19 -0800
Subject: [PATCH] Add more locking around on-disk image cache

Concurrent writes/reads to the cache might result in partially written
files, which we'd like to avoid.

On top of that, we generally don't want builds to fail if something goes
wrong with the cache, so now we just log any cache errors and do a
fallback.

Signed-off-by: Jon Johnson <jon.johnson@chainguard.dev>
---
 pkg/commands/cache.go  | 13 +++++++++++--
 pkg/commands/config.go | 29 +++++++++++++++++++++--------
 2 files changed, 32 insertions(+), 10 deletions(-)

diff --git a/pkg/commands/cache.go b/pkg/commands/cache.go
index 5ce5b74cda..e39420949e 100644
--- a/pkg/commands/cache.go
+++ b/pkg/commands/cache.go
@@ -38,8 +38,12 @@ import (
 type imageCache struct {
 	// In memory
 	cache sync.Map
-	p     *layout.Path
 
+	// On disk
+	mu sync.Mutex
+	p  *layout.Path
+
+	// Over the network
 	puller *remote.Puller
 }
 
@@ -101,10 +105,15 @@ func (i *imageCache) get(ctx context.Context, ref name.Reference, missFunc baseF
 			key = dig.String()
 		}
 
+		// Use a pretty broad lock on the on-disk cache to avoid races.
+		i.mu.Lock()
+		defer i.mu.Unlock()
+
 		ii, err := i.p.ImageIndex()
 		if err != nil {
-			return nil, err
+			return nil, fmt.Errorf("loading cache index: %w", err)
 		}
+
 		h, err := v1.NewHash(key)
 		if err != nil {
 			return nil, err
diff --git a/pkg/commands/config.go b/pkg/commands/config.go
index 432d98cd07..896dbc489e 100644
--- a/pkg/commands/config.go
+++ b/pkg/commands/config.go
@@ -75,11 +75,6 @@ func getBaseImage(bo *options.BuildOptions) build.GetBase {
 	}
 
 	fetch := func(ctx context.Context, ref name.Reference) (build.Result, error) {
-		// For ko.local, look in the daemon.
-		if ref.Context().RegistryStr() == publish.LocalDomain {
-			return daemon.Image(ref)
-		}
-
 		ropt = append(ropt, remote.WithContext(ctx))
 
 		desc, err := remote.Get(ref, ropt...)
@@ -91,6 +86,7 @@ func getBaseImage(bo *options.BuildOptions) build.GetBase {
 		}
 		return desc.Image()
 	}
+
 	return func(ctx context.Context, s string) (name.Reference, build.Result, error) {
 		s = strings.TrimPrefix(s, build.StrictScheme)
 		// Viper configuration file keys are case insensitive, and are
@@ -112,9 +108,26 @@ func getBaseImage(bo *options.BuildOptions) build.GetBase {
 			return nil, nil, fmt.Errorf("parsing base image (%q): %w", baseImage, err)
 		}
 
-		result, err := cache.get(ctx, ref, fetch)
-		if err != nil {
-			return ref, result, err
+		var result build.Result
+
+		// For ko.local, look in the daemon.
+		if ref.Context().RegistryStr() == publish.LocalDomain {
+			result, err = daemon.Image(ref)
+			if err != nil {
+				return nil, nil, fmt.Errorf("loading %s from daemon: %w", ref, err)
+			}
+		} else {
+			result, err = cache.get(ctx, ref, fetch)
+			if err != nil {
+				// We don't expect this to fail, usually, but the cache should also not be fatal.
+				// Log it so people can complain about it and we can fix the cache.
+				log.Printf("cache.get(%q) failed with %v", ref.String(), err)
+
+				result, err = fetch(ctx, ref)
+				if err != nil {
+					return nil, nil, fmt.Errorf("pulling %s: %w", ref, err)
+				}
+			}
 		}
 
 		if _, ok := ref.(name.Digest); ok {