diff --git a/go.mod b/go.mod index 851236b52..671c3778a 100644 --- a/go.mod +++ b/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 diff --git a/go.sum b/go.sum index ab2fff415..6a75a8b70 100644 --- a/go.sum +++ b/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= diff --git a/pkg/cache/cache.go b/pkg/cache/cache.go new file mode 100644 index 000000000..b9ebf6d9e --- /dev/null +++ b/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 +} diff --git a/pkg/cache/memory/memory.go b/pkg/cache/memory/memory.go new file mode 100644 index 000000000..49375067e --- /dev/null +++ b/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, ",") +} diff --git a/pkg/cache/memory/memory_test.go b/pkg/cache/memory/memory_test.go new file mode 100644 index 000000000..70f38a77c --- /dev/null +++ b/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 +}