Use lru cache to expire entrypoint items(#332)

Use a lru cache that will only keep the most recently-used entrypoints
cached and drop entrypoints haven't been used recently, in order to
present a risk of DoS.

1. replace internal map in original Cache with lru.Cache
2. fix incorrect entrypoint_test package name

Signed-off-by: mathspanda <[email protected]>

pkg/reconciler/v1alpha1/taskrun/entrypoint/entrypoint.go

@@ -20,15 +20,15 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
-	"sync"
 
 	"github.com/google/go-containerregistry/pkg/authn"
 	"github.com/google/go-containerregistry/pkg/name"
 	"github.com/google/go-containerregistry/pkg/v1/remote"
-	"github.com/knative/build/pkg/apis/build/v1alpha1"
+	"github.com/hashicorp/golang-lru"
 	corev1 "k8s.io/api/core/v1"
 
 	"github.com/knative/build-pipeline/pkg/reconciler/v1alpha1/taskrun/config"
+	"github.com/knative/build/pkg/apis/build/v1alpha1"
 )
 
 const (
@@ -41,6 +41,8 @@ const (
 	InitContainerName = "place-tools"
 	ProcessLogFile    = "/tools/process-log.txt"
 	MarkerFile        = "/tools/marker-file.txt"
+
+	cacheSize = 1024
 )
 
 var toolsMount = corev1.VolumeMount{
@@ -49,35 +51,28 @@ var toolsMount = corev1.VolumeMount{
 }
 
 // Cache is a simple caching mechanism allowing for caching the results of
-// getting the Entrypoint of a container image from a remote registry. It
-// is synchronized via a mutex so that we can share a single Cache across
-// each worker thread that the reconciler is running. The mutex is necessary
-// due to the possibility of a panic if two workers were to attempt to read and
-// write to the internal map at the same time.
+// getting the Entrypoint of a container image from a remote registry. The
+// internal lru cache is thread-safe.
 type Cache struct {
-	mtx   sync.RWMutex
-	cache map[string][]string
+	lru *lru.Cache
 }
 
 // NewCache is a simple helper function that returns a pointer to a Cache that
-// has had the internal cache map initialized.
-func NewCache() *Cache {
-	return &Cache{
-		cache: make(map[string][]string),
-	}
+// has had the internal fixed-sized lru cache initialized.
+func NewCache() (*Cache, error) {
+	lru, err := lru.New(cacheSize)
+	return &Cache{lru}, err
 }
 
 func (c *Cache) get(sha string) ([]string, bool) {
-	c.mtx.RLock()
-	ep, ok := c.cache[sha]
-	c.mtx.RUnlock()
-	return ep, ok
+	if ep, ok := c.lru.Get(sha); ok {
+		return ep.([]string), true
+	}
+	return nil, false
 }
 
 func (c *Cache) set(sha string, ep []string) {
-	c.mtx.Lock()
-	c.cache[sha] = ep
-	c.mtx.Unlock()
+	c.lru.Add(sha, ep)
 }
 
 // AddCopyStep will prepend a BuildStep (Container) that will

pkg/reconciler/v1alpha1/taskrun/entrypoint/entrypoint_test.go

@@ -1,4 +1,4 @@
-package entrypoint_test
+package entrypoint
 
 import (
 	"context"
@@ -10,14 +10,17 @@ import (
 	"strings"
 	"testing"
 
-	v1 "github.com/google/go-containerregistry/pkg/v1"
+	"github.com/google/go-containerregistry/pkg/v1"
 	"github.com/google/go-containerregistry/pkg/v1/partial"
 	"github.com/google/go-containerregistry/pkg/v1/types"
 	"github.com/knative/build/pkg/apis/build/v1alpha1"
 	corev1 "k8s.io/api/core/v1"
 
 	"github.com/knative/build-pipeline/pkg/reconciler/v1alpha1/taskrun/config"
-	"github.com/knative/build-pipeline/pkg/reconciler/v1alpha1/taskrun/entrypoint"
+)
+
+const (
+	exceedCacheSize = 10
 )
 
 func TestAddEntrypoint(t *testing.T) {
@@ -44,12 +47,12 @@ func TestAddEntrypoint(t *testing.T) {
 		`{"args":["abcd"],"process_log":"/tools/process-log.txt","marker_file":"/tools/marker-file.txt"}`,
 		`{"args":["abcd","efgh"],"process_log":"/tools/process-log.txt","marker_file":"/tools/marker-file.txt"}`,
 	}
-	err := entrypoint.RedirectSteps(inputs)
+	err := RedirectSteps(inputs)
 	if err != nil {
 		t.Errorf("failed to get resources: %v", err)
 	}
 	for i, input := range inputs {
-		if len(input.Command) == 0 || input.Command[0] != entrypoint.BinaryLocation {
+		if len(input.Command) == 0 || input.Command[0] != BinaryLocation {
 			t.Errorf("command incorrectly set: %q", input.Command)
 		}
 		if len(input.Args) > 0 {
@@ -59,13 +62,13 @@ func TestAddEntrypoint(t *testing.T) {
 			t.Error("there should be atleast one envvar")
 		}
 		for _, e := range input.Env {
-			if e.Name == entrypoint.JSONConfigEnvVar && e.Value != envVarStrings[i] {
+			if e.Name == JSONConfigEnvVar && e.Value != envVarStrings[i] {
 				t.Errorf("envvar \n%s\n does not match \n%s", e.Value, envVarStrings[i])
 			}
 		}
 		found := false
 		for _, vm := range input.VolumeMounts {
-			if vm.Name == entrypoint.MountName {
+			if vm.Name == MountName {
 				found = true
 				break
 			}
@@ -165,7 +168,11 @@ func TestGetRemoteEntrypoint(t *testing.T) {
 	}))
 	defer server.Close()
 	image := path.Join(strings.TrimPrefix(server.URL, "http://"), "image:latest")
-	ep, err := entrypoint.GetRemoteEntrypoint(entrypoint.NewCache(), image)
+	entrypointCache, err := NewCache()
+	if err != nil {
+		t.Fatalf("couldn't create new entrypoint cache: %v", err)
+	}
+	ep, err := GetRemoteEntrypoint(entrypointCache, image)
 	if err != nil {
 		t.Errorf("couldn't get entrypoint remote: %v", err)
 	}
@@ -174,6 +181,31 @@ func TestGetRemoteEntrypoint(t *testing.T) {
 	}
 }
 
+func TestEntrypointCacheLRU(t *testing.T) {
+	entrypoint := []string{"/bin/expected", "entrypoint"}
+	entrypointCache, err := NewCache()
+	if err != nil {
+		t.Fatalf("couldn't create new entrypoint cache: %v", err)
+	}
+
+	for i := 0; i < cacheSize+exceedCacheSize; i++ {
+		image := fmt.Sprintf("image%d:latest", i)
+		entrypointCache.set(image, entrypoint)
+	}
+	for i := 0; i < exceedCacheSize; i++ {
+		image := fmt.Sprintf("image%d:latest", i)
+		if _, ok := entrypointCache.get(image); ok {
+			t.Errorf("entrypoint of image %s should be expired", image)
+		}
+	}
+	for i := exceedCacheSize; i < cacheSize+exceedCacheSize; i++ {
+		image := fmt.Sprintf("image%d:latest", i)
+		if _, ok := entrypointCache.get(image); !ok {
+			t.Errorf("entrypoint of image %s shouldn't be expired", image)
+		}
+	}
+}
+
 func TestAddCopyStep(t *testing.T) {
 	cfg := &config.Config{
 		Entrypoint: &config.Entrypoint{
@@ -194,11 +226,11 @@ func TestAddCopyStep(t *testing.T) {
 	}
 
 	expectedSteps := len(bs.Steps) + 1
-	entrypoint.AddCopyStep(ctx, bs)
+	AddCopyStep(ctx, bs)
 	if len(bs.Steps) != 3 {
 		t.Errorf("BuildSpec has the wrong step count: %d should be %d", len(bs.Steps), expectedSteps)
 	}
-	if bs.Steps[0].Name != entrypoint.InitContainerName {
-		t.Errorf("entrypoint is incorrect: %s should be %s", bs.Steps[0].Name, entrypoint.InitContainerName)
+	if bs.Steps[0].Name != InitContainerName {
+		t.Errorf("entrypoint is incorrect: %s should be %s", bs.Steps[0].Name, InitContainerName)
 	}
 }

pkg/reconciler/v1alpha1/taskrun/taskrun.go

@@ -387,7 +387,10 @@ func (c *Reconciler) createBuildPod(ctx context.Context, tr *v1alpha1.TaskRun, t
 
 	// For each step with no entrypoint set, try to populate it with the info
 	// from the remote registry
-	cache := entrypoint.NewCache()
+	cache, err := entrypoint.NewCache()
+	if err != nil {
+		return nil, fmt.Errorf("couldn't create new entrypoint cache: %v", err)
+	}
 	bSpec := bs.DeepCopy()
 	for i := range bSpec.Steps {
 		step := &bSpec.Steps[i]