From c04850088f84fdf638cf5884f914957919465e3b Mon Sep 17 00:00:00 2001
From: ChristopherHX <christopher.homberger@web.de>
Date: Fri, 29 Mar 2024 03:07:20 +0100
Subject: [PATCH] fix: cache adjust restore order of exact key matches (#2267)

* wip: adjust restore order

* fixup

* add tests

* cleanup

* fix typo

---------

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
(cherry picked from commit f825e42ce200fc4973c3f28797ffba942d322d38)
---
 pkg/artifactcache/handler.go      |  11 ++++
 pkg/artifactcache/handler_test.go | 104 ++++++++++++++++++++++++++++++
 2 files changed, 115 insertions(+)

diff --git a/pkg/artifactcache/handler.go b/pkg/artifactcache/handler.go
index 395a39e..065c7dd 100644
--- a/pkg/artifactcache/handler.go
+++ b/pkg/artifactcache/handler.go
@@ -352,6 +352,17 @@ func (h *Handler) middleware(handler httprouter.Handle) httprouter.Handle {
 func findCache(db *bolthold.Store, keys []string, version string) (*Cache, error) {
 	cache := &Cache{}
 	for _, prefix := range keys {
+		// if a key in the list matches exactly, don't return partial matches
+		if err := db.FindOne(cache,
+			bolthold.Where("Key").Eq(prefix).
+				And("Version").Eq(version).
+				And("Complete").Eq(true).
+				SortBy("CreatedAt").Reverse()); err == nil || !errors.Is(err, bolthold.ErrNotFound) {
+			if err != nil {
+				return nil, fmt.Errorf("find cache: %w", err)
+			}
+			return cache, nil
+		}
 		prefixPattern := fmt.Sprintf("^%s", regexp.QuoteMeta(prefix))
 		re, err := regexp.Compile(prefixPattern)
 		if err != nil {
diff --git a/pkg/artifactcache/handler_test.go b/pkg/artifactcache/handler_test.go
index b418f00..252c420 100644
--- a/pkg/artifactcache/handler_test.go
+++ b/pkg/artifactcache/handler_test.go
@@ -422,6 +422,110 @@ func TestHandler(t *testing.T) {
 			assert.Equal(t, key+"_abc", got.CacheKey)
 		}
 	})
+
+	t.Run("exact keys are preferred (key 0)", func(t *testing.T) {
+		version := "c19da02a2bd7e77277f1ac29ab45c09b7d46a4ee758284e26bb3045ad11d9d20"
+		key := strings.ToLower(t.Name())
+		keys := [3]string{
+			key + "_a",
+			key + "_a_b_c",
+			key + "_a_b",
+		}
+		contents := [3][]byte{
+			make([]byte, 100),
+			make([]byte, 200),
+			make([]byte, 300),
+		}
+		for i := range contents {
+			_, err := rand.Read(contents[i])
+			require.NoError(t, err)
+			uploadCacheNormally(t, base, keys[i], version, contents[i])
+			time.Sleep(time.Second) // ensure CreatedAt of caches are different
+		}
+
+		reqKeys := strings.Join([]string{
+			key + "_a",
+			key + "_a_b",
+		}, ",")
+
+		resp, err := http.Get(fmt.Sprintf("%s/cache?keys=%s&version=%s", base, reqKeys, version))
+		require.NoError(t, err)
+		require.Equal(t, 200, resp.StatusCode)
+
+		/*
+			Expect `key_a` because:
+			- `key_a` matches `key_a`, `key_a_b` and `key_a_b_c`, but `key_a` is an exact match.
+			- `key_a_b` matches `key_a_b` and `key_a_b_c`, but previous key had a match
+		*/
+		expect := 0
+
+		got := struct {
+			ArchiveLocation string `json:"archiveLocation"`
+			CacheKey        string `json:"cacheKey"`
+		}{}
+		require.NoError(t, json.NewDecoder(resp.Body).Decode(&got))
+		assert.Equal(t, keys[expect], got.CacheKey)
+
+		contentResp, err := http.Get(got.ArchiveLocation)
+		require.NoError(t, err)
+		require.Equal(t, 200, contentResp.StatusCode)
+		content, err := io.ReadAll(contentResp.Body)
+		require.NoError(t, err)
+		assert.Equal(t, contents[expect], content)
+	})
+
+	t.Run("exact keys are preferred (key 1)", func(t *testing.T) {
+		version := "c19da02a2bd7e77277f1ac29ab45c09b7d46a4ee758284e26bb3045ad11d9d20"
+		key := strings.ToLower(t.Name())
+		keys := [3]string{
+			key + "_a",
+			key + "_a_b_c",
+			key + "_a_b",
+		}
+		contents := [3][]byte{
+			make([]byte, 100),
+			make([]byte, 200),
+			make([]byte, 300),
+		}
+		for i := range contents {
+			_, err := rand.Read(contents[i])
+			require.NoError(t, err)
+			uploadCacheNormally(t, base, keys[i], version, contents[i])
+			time.Sleep(time.Second) // ensure CreatedAt of caches are different
+		}
+
+		reqKeys := strings.Join([]string{
+			"------------------------------------------------------",
+			key + "_a",
+			key + "_a_b",
+		}, ",")
+
+		resp, err := http.Get(fmt.Sprintf("%s/cache?keys=%s&version=%s", base, reqKeys, version))
+		require.NoError(t, err)
+		require.Equal(t, 200, resp.StatusCode)
+
+		/*
+			Expect `key_a` because:
+			- `------------------------------------------------------` doesn't match any caches.
+			- `key_a` matches `key_a`, `key_a_b` and `key_a_b_c`, but `key_a` is an exact match.
+			- `key_a_b` matches `key_a_b` and `key_a_b_c`, but previous key had a match
+		*/
+		expect := 0
+
+		got := struct {
+			ArchiveLocation string `json:"archiveLocation"`
+			CacheKey        string `json:"cacheKey"`
+		}{}
+		require.NoError(t, json.NewDecoder(resp.Body).Decode(&got))
+		assert.Equal(t, keys[expect], got.CacheKey)
+
+		contentResp, err := http.Get(got.ArchiveLocation)
+		require.NoError(t, err)
+		require.Equal(t, 200, contentResp.StatusCode)
+		content, err := io.ReadAll(contentResp.Body)
+		require.NoError(t, err)
+		assert.Equal(t, contents[expect], content)
+	})
 }
 
 func uploadCacheNormally(t *testing.T, base, key, version string, content []byte) {