Extend memory cache

pkg/cache/cache.go

@@ -4,9 +4,9 @@ package cache
 // Cache is used to store key/value pairs.
 type Cache interface {
 	// Set stores the given key/value pair.
-	Set(string, string)
+	Set(string, any)
 	// Get returns the value for the given key and a boolean indicating if the key was found.
-	Get(string) (string, bool)
+	Get(string) (any, bool)
 	// Exists returns true if the given key exists in the cache.
 	Exists(string) bool
 	// Delete the given key from the cache.
@@ -18,7 +18,7 @@ type Cache interface {
 	// Keys returns all keys in the cache.
 	Keys() []string
 	// Values returns all values in the cache.
-	Values() []string
+	Values() []any
 	// Contents returns all keys in the cache encoded as a string.
 	Contents() string
 }

pkg/cache/memory/memory.go

@@ -5,52 +5,77 @@ import (
 	"time"
 
 	"github.com/patrickmn/go-cache"
-
-	"github.com/trufflesecurity/trufflehog/v3/pkg/context"
 )
 
 const (
-	expirationInterval = 12 * time.Hour
-	purgeInterval      = 13 * time.Hour
-	defaultExpiration  = cache.DefaultExpiration
+	defaultExpirationInterval = 12 * time.Hour
+	defaultPurgeInterval      = 13 * time.Hour
+	defaultExpiration         = cache.DefaultExpiration
 )
 
-// Cache is a wrapper around the go-cache library.
+// Cache wraps the go-cache library to provide an in-memory key-value store.
 type Cache struct {
-	c *cache.Cache
+	c             *cache.Cache
+	expiration    time.Duration
+	purgeInterval time.Duration
+}
+
+// CacheOption defines a function type used for configuring a Cache.
+type CacheOption func(*Cache)
+
+// WithExpirationInterval returns a CacheOption to set the expiration interval of cache items.
+// The interval determines the duration a cached item remains in the cache before it is expired.
+func WithExpirationInterval(interval time.Duration) CacheOption {
+	return func(c *Cache) { c.expiration = interval }
+}
+
+// WithPurgeInterval returns a CacheOption to set the interval at which the cache purges expired items.
+// Regular purging helps in freeing up memory by removing stale entries.
+func WithPurgeInterval(interval time.Duration) CacheOption {
+	return func(c *Cache) { c.purgeInterval = interval }
 }
 
-// New constructs a new in-memory cache.
-func New() *Cache {
-	c := cache.New(expirationInterval, purgeInterval)
-	return &Cache{c: c}
+// New constructs a new in-memory cache instance with optional configurations.
+// By default, it sets the expiration and purge intervals to 12 and 13 hours, respectively.
+// These defaults can be overridden using the functional options: WithExpirationInterval and WithPurgeInterval.
+func New(opts ...CacheOption) *Cache {
+	return NewWithData(nil, opts...)
+}
+
+// CacheEntry represents a single entry in the cache, consisting of a key and its corresponding value.
+type CacheEntry struct {
+	// Key is the unique identifier for the entry.
+	Key string
+	// Value is the data stored in the entry.
+	Value any
 }
 
 // NewWithData constructs a new in-memory cache with existing data.
-func NewWithData(ctx context.Context, data []string) *Cache {
-	ctx.Logger().V(3).Info("Loading cache", "num-items", len(data))
+// It also accepts CacheOption parameters to override default configuration values.
+func NewWithData(data []CacheEntry, opts ...CacheOption) *Cache {
+	instance := &Cache{expiration: defaultExpirationInterval, purgeInterval: defaultPurgeInterval}
+	for _, opt := range opts {
+		opt(instance)
+	}
 
+	// Convert data slice to map required by go-cache.
 	items := make(map[string]cache.Item, len(data))
 	for _, d := range data {
-		items[d] = cache.Item{Object: d, Expiration: int64(defaultExpiration)}
+		items[d.Key] = cache.Item{Object: d.Value, Expiration: int64(defaultExpiration)}
 	}
 
-	c := cache.NewFrom(expirationInterval, purgeInterval, items)
-	return &Cache{c: c}
+	instance.c = cache.NewFrom(instance.expiration, instance.purgeInterval, items)
+	return instance
 }
 
 // Set adds a key-value pair to the cache.
-func (c *Cache) Set(key, value string) {
+func (c *Cache) Set(key string, value any) {
 	c.c.Set(key, value, defaultExpiration)
 }
 
 // Get returns the value for the given key.
-func (c *Cache) Get(key string) (string, bool) {
-	res, ok := c.c.Get(key)
-	if !ok {
-		return "", ok
-	}
-	return res.(string), ok
+func (c *Cache) Get(key string) (any, bool) {
+	return c.c.Get(key)
 }
 
 // Exists returns true if the given key exists in the cache.
@@ -85,11 +110,11 @@ func (c *Cache) Keys() []string {
 }
 
 // Values returns all values in the cache.
-func (c *Cache) Values() []string {
+func (c *Cache) Values() []any {
 	items := c.c.Items()
-	res := make([]string, 0, len(items))
+	res := make([]any, 0, len(items))
 	for _, v := range items {
-		res = append(res, v.Object.(string))
+		res = append(res, v.Object)
 	}
 	return res
 }

pkg/cache/memory/memory_test.go

@@ -7,8 +7,6 @@ import (
 	"testing"
 
 	"github.com/google/go-cmp/cmp"
-
-	logContext "github.com/trufflesecurity/trufflehog/v3/pkg/context"
 )
 
 func TestCache(t *testing.T) {
@@ -34,15 +32,15 @@ func TestCache(t *testing.T) {
 	// Test delete.
 	c.Delete("key1")
 	v, ok = c.Get("key1")
-	if ok || v != "" {
+	if ok || v != nil {
 		t.Fatalf("Unexpected value for key1 after delete: %v, %v", v, ok)
 	}
 
 	// Test clear.
 	c.Set("key10", "key10")
 	c.Clear()
 	v, ok = c.Get("key10")
-	if ok || v != "" {
+	if ok || v != nil {
 		t.Fatalf("Unexpected value for key10 after clear: %v, %v", v, ok)
 	}
 
@@ -60,7 +58,10 @@ func TestCache(t *testing.T) {
 	}
 
 	// Test getting only the values.
-	vals := c.Values()
+	vals := make([]string, 0, c.Count())
+	for _, v := range c.Values() {
+		vals = append(vals, v.(string))
+	}
 	sort.Strings(vals)
 	sort.Strings(values)
 	if !cmp.Equal(values, vals) {
@@ -82,7 +83,8 @@ func TestCache(t *testing.T) {
 }
 
 func TestCache_NewWithData(t *testing.T) {
-	c := NewWithData(logContext.Background(), []string{"key1", "key2", "key3"})
+	data := []CacheEntry{{"key1", "value1"}, {"key2", "value2"}, {"key3", "value3"}}
+	c := NewWithData(data)
 
 	// Test the count.
 	if c.Count() != 3 {

pkg/sources/gcs/gcs.go

@@ -97,7 +97,7 @@ func newPersistableCache(increment int, cache cache.Cache, p *sources.Progress)
 
 // Set overrides the cache Set method of the cache to enable the persistence
 // of the cache contents the Progress of the source at given increments.
-func (c *persistableCache) Set(key, val string) {
+func (c *persistableCache) Set(key string, val any) {
 	c.Cache.Set(key, val)
 	if ok, contents := c.shouldPersist(); ok {
 		c.Progress.EncodedResumeInfo = contents
@@ -297,7 +297,14 @@ func (s *Source) Chunks(ctx context.Context, chunksChan chan *sources.Chunk, _ .
 func (s *Source) setupCache(ctx context.Context) *persistableCache {
 	var c cache.Cache
 	if s.Progress.EncodedResumeInfo != "" {
-		c = memory.NewWithData(ctx, strings.Split(s.Progress.EncodedResumeInfo, ","))
+		keys := strings.Split(s.Progress.EncodedResumeInfo, ",")
+		entries := make([]memory.CacheEntry, len(keys))
+		for i, val := range keys {
+			entries[i] = memory.CacheEntry{Key: val, Value: val}
+		}
+
+		c = memory.NewWithData(entries)
+		ctx.Logger().V(3).Info("Loaded cache", "num_entries", len(entries))
 	} else {
 		c = memory.New()
 	}

pkg/sources/github/github.go

@@ -472,7 +472,15 @@ func (s *Source) enumerate(ctx context.Context, apiEndpoint string) (*github.Cli
 		return nil, errors.Errorf("Invalid configuration given for source. Name: %s, Type: %s", s.name, s.Type())
 	}
 
-	s.repos = s.filteredRepoCache.Values()
+	s.repos = make([]string, 0, s.filteredRepoCache.Count())
+	for _, repo := range s.filteredRepoCache.Values() {
+		r, ok := repo.(string)
+		if !ok {
+			ctx.Logger().Error(fmt.Errorf("type assertion failed"), "unexpected value in cache", "repo", repo)
+			continue
+		}
+		s.repos = append(s.repos, r)
+	}
 	githubReposEnumerated.WithLabelValues(s.name).Set(float64(len(s.repos)))
 	s.log.Info("Completed enumeration", "num_repos", len(s.repos), "num_orgs", s.orgsCache.Count(), "num_members", len(s.memberCache))