Add in-memory caching pkg (#1189)

* Add in-memory caching lib, used by the GCS source.

* Fix static check.

* Add test for NewWithData.

* update comment.

* update comments.

* remove unused dep.

* address comments.

* Add exists method.

* fix test.

go.mod

@@ -45,6 +45,7 @@ require (
 	github.com/lib/pq v1.10.7
 	github.com/mattn/go-sqlite3 v1.14.16
 	github.com/mholt/archiver/v4 v4.0.0-alpha.7
+	github.com/patrickmn/go-cache v2.1.0+incompatible
 	github.com/paulbellamy/ratecounter v0.2.0
 	github.com/petar-dambovaliev/aho-corasick v0.0.0-20211021192214-5ab2d9280aa9
 	github.com/pkg/errors v0.9.1

go.sum

@@ -283,6 +283,8 @@ github.com/onsi/gomega v1.10.1/go.mod h1:iN09h71vgCQne3DLsj+A5owkum+a2tYe+TOCB1y
 github.com/onsi/gomega v1.23.0 h1:/oxKu9c2HVap+F3PfKort2Hw5DEU+HGlW8n+tguWsys=
 github.com/onsi/gomega v1.23.0/go.mod h1:Z/NWtiqwBrwUt4/2loMmHL63EDLnYHmVbuBpDr2vQAg=
 github.com/op/go-logging v0.0.0-20160315200505-970db520ece7/go.mod h1:HzydrMdWErDVzsI23lYNej1Htcns9BCg93Dk0bBINWk=
+github.com/patrickmn/go-cache v2.1.0+incompatible h1:HRMgzkcYKYpi3C8ajMPV8OFXaaRUnok+kx1WdO15EQc=
+github.com/patrickmn/go-cache v2.1.0+incompatible/go.mod h1:3Qf8kWWT7OJRJbdiICTKqZju1ZixQ/KpMGzzAfe6+WQ=
 github.com/paulbellamy/ratecounter v0.2.0 h1:2L/RhJq+HA8gBQImDXtLPrDXK5qAj6ozWVK/zFXVJGs=
 github.com/paulbellamy/ratecounter v0.2.0/go.mod h1:Hfx1hDpSGoqxkVVpBi/IlYD7kChlfo5C6hzIHwPqfFE=
 github.com/petar-dambovaliev/aho-corasick v0.0.0-20211021192214-5ab2d9280aa9 h1:lL+y4Xv20pVlCGyLzNHRC0I0rIHhIL1lTvHizoS/dU8=

pkg/cache/cache.go

@@ -0,0 +1,20 @@
+// Package cache provides an interface which can be implemented by different cache types.
+package cache
+
+// Cache is used to store key/value pairs.
+type Cache interface {
+	// Set stores the given key/value pair.
+	Set(string, string)
+	// Get returns the value for the given key and a boolean indicating if the key was found.
+	Get(string) (string, bool)
+	// Exists returns true if the given key exists in the cache.
+	Exists(string) bool
+	// Delete the given key from the cache.
+	Delete(string)
+	// Clear all key/value pairs from the cache.
+	Clear()
+	// Count the number of key/value pairs in the cache.
+	Count() int
+	// Contents returns all keys in the cache encoded as a string.
+	Contents() string
+}

pkg/cache/memory/memory.go

@@ -0,0 +1,85 @@
+package memory
+
+import (
+	"strings"
+	"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
+)
+
+// Cache is a wrapper around the go-cache library.
+type Cache struct {
+	c *cache.Cache
+}
+
+// New constructs a new in-memory cache.
+func New() *Cache {
+	c := cache.New(expirationInterval, purgeInterval)
+	return &Cache{c: c}
+}
+
+// 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))
+
+	items := make(map[string]cache.Item, len(data))
+	for _, d := range data {
+		items[d] = cache.Item{Object: d, Expiration: int64(defaultExpiration)}
+	}
+
+	c := cache.NewFrom(expirationInterval, purgeInterval, items)
+	return &Cache{c: c}
+}
+
+// Set adds a key-value pair to the cache.
+func (c *Cache) Set(key, value string) {
+	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
+}
+
+// Exists returns true if the given key exists in the cache.
+func (c *Cache) Exists(key string) bool {
+	_, ok := c.c.Get(key)
+	return ok
+}
+
+// Delete removes the key-value pair from the cache.
+func (c *Cache) Delete(key string) {
+	c.c.Delete(key)
+}
+
+// Clear removes all key-value pairs from the cache.
+func (c *Cache) Clear() {
+	c.c.Flush()
+}
+
+// Count returns the number of key-value pairs in the cache.
+func (c *Cache) Count() int {
+	return c.c.ItemCount()
+}
+
+// Contents returns all key-value pairs in the cache encodes as a string.
+func (c *Cache) Contents() string {
+	items := c.c.Items()
+	res := make([]string, 0, len(items))
+	for k := range items {
+		res = append(res, k)
+	}
+	return strings.Join(res, ",")
+}

pkg/cache/memory/memory_test.go

@@ -0,0 +1,153 @@
+package memory
+
+import (
+	"fmt"
+	"sort"
+	"strings"
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+
+	logContext "github.com/trufflesecurity/trufflehog/v3/pkg/context"
+)
+
+func TestCache(t *testing.T) {
+	c := New()
+
+	// Test set and get.
+	c.Set("key1", "key1")
+	v, ok := c.Get("key1")
+	if !ok || v != "key1" {
+		t.Fatalf("Unexpected value for key1: %v, %v", v, ok)
+	}
+
+	// Test exists.
+	if !c.Exists("key1") {
+		t.Fatalf("Expected key1 to exist")
+	}
+
+	// Test the count.
+	if c.Count() != 1 {
+		t.Fatalf("Unexpected count: %d", c.Count())
+	}
+
+	// Test delete.
+	c.Delete("key1")
+	v, ok = c.Get("key1")
+	if ok || v != "" {
+		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 != "" {
+		t.Fatalf("Unexpected value for key10 after clear: %v, %v", v, ok)
+	}
+
+	// Test contents.
+	keys := []string{"key1", "key2", "key3"}
+	for _, k := range keys {
+		c.Set(k, k)
+	}
+
+	items := c.Contents()
+	sort.Strings(keys)
+	res := strings.Split(items, ",")
+	sort.Strings(res)
+
+	if len(keys) != len(res) {
+		t.Fatalf("Unexpected length of items: %d", len(res))
+	}
+	if !cmp.Equal(keys, res) {
+		t.Fatalf("Unexpected items: %v", res)
+	}
+}
+
+func TestCache_NewWithData(t *testing.T) {
+	c := NewWithData(logContext.Background(), []string{"key1", "key2", "key3"})
+
+	// Test the count.
+	if c.Count() != 3 {
+		t.Fatalf("Unexpected count: %d", c.Count())
+	}
+
+	// Test contents.
+	keys := []string{"key1", "key2", "key3"}
+	items := c.Contents()
+	sort.Strings(keys)
+	res := strings.Split(items, ",")
+	sort.Strings(res)
+
+	if len(keys) != len(res) {
+		t.Fatalf("Unexpected length of items: %d", len(res))
+	}
+	if !cmp.Equal(keys, res) {
+		t.Fatalf("Unexpected items: %v", res)
+	}
+}
+
+func setupBenchmarks(b *testing.B) *Cache {
+	b.Helper()
+
+	c := New()
+
+	for i := 0; i < 500_000; i++ {
+		key := fmt.Sprintf("key%d", i)
+		c.Set(key, key)
+	}
+
+	return c
+}
+
+func BenchmarkSet(b *testing.B) {
+	c := New()
+
+	for i := 0; i < b.N; i++ {
+		key := fmt.Sprintf("key%d", i)
+		c.Set(key, key)
+	}
+}
+
+func BenchmarkGet(b *testing.B) {
+	c := setupBenchmarks(b)
+	b.ResetTimer()
+
+	for i := 0; i < b.N; i++ {
+		key := fmt.Sprintf("key%d", i)
+		c.Get(key)
+	}
+}
+
+func BenchmarkDelete(b *testing.B) {
+	c := setupBenchmarks(b)
+	b.ResetTimer()
+
+	for i := 0; i < b.N; i++ {
+		key := fmt.Sprintf("key%d", i)
+		c.Delete(key)
+	}
+}
+
+func BenchmarkCount(b *testing.B) {
+	c := setupBenchmarks(b)
+	b.ResetTimer()
+
+	for i := 0; i < b.N; i++ {
+		c.Count()
+	}
+}
+
+func BenchmarkContents(b *testing.B) {
+	c := setupBenchmarks(b)
+	b.ResetTimer()
+
+	var s string
+
+	for i := 0; i < b.N; i++ {
+		s = c.Contents()
+	}
+
+	_ = s
+}