From eaf9fa5e3f9eed7ac38c4821be441743b6493b64 Mon Sep 17 00:00:00 2001 From: Teppei Fukuda Date: Wed, 25 Dec 2019 13:57:07 +0200 Subject: [PATCH] feat(cache): wrap kv cache (fanal#62) --- cache/cache.go | 49 ++++++++++++++++++--- cache/cache_test.go | 12 +++-- cmd/fanal/main.go | 8 ++-- extractor/docker/docker.go | 48 ++++---------------- extractor/docker/docker_test.go | 78 +++++++++++---------------------- utils/utils.go | 4 +- 6 files changed, 88 insertions(+), 111 deletions(-) diff --git a/cache/cache.go b/cache/cache.go index cf7ba6c27b..59978757b5 100644 --- a/cache/cache.go +++ b/cache/cache.go @@ -2,25 +2,64 @@ package cache import ( "os" + "path/filepath" + bolt "github.com/simar7/gokv/bbolt" + "github.com/simar7/gokv/encoding" + kvtypes "github.com/simar7/gokv/types" "golang.org/x/xerrors" ) type Cache interface { + Get(bucket, key string, value *[]byte) (found bool, err error) + Set(bucket, key string, value []byte) (err error) Clear() error } type RealCache struct { - Directory string + directory string + cache *bolt.Store } -func Initialize(cacheDir string) Cache { - return &RealCache{Directory: cacheDir} +func New(cacheDir string) (Cache, error) { + dir := filepath.Join(cacheDir, "fanal") + if err := os.MkdirAll(dir, 0700); err != nil { + return nil, xerrors.Errorf("unable to create cache dir: %w", err) + } + + cacheOptions := bolt.Options{ + RootBucketName: "fanal", + Path: filepath.Join(dir, "cache.db"), + Codec: encoding.Raw, + } + + kv, err := bolt.NewStore(cacheOptions) + if err != nil { + return nil, xerrors.Errorf("error initializing cache: %w", err) + } + + return &RealCache{directory: dir, cache: kv}, nil +} + +func (rc RealCache) Get(bucket, key string, value *[]byte) (bool, error) { + return rc.cache.Get(kvtypes.GetItemInput{ + BucketName: bucket, + Key: key, + Value: value, + }) +} + +func (rc RealCache) Set(bucket, key string, value []byte) error { + return rc.cache.BatchSet(kvtypes.BatchSetItemInput{ + BucketName: bucket, + Keys: []string{key}, + Values: value, + }) } func (rc RealCache) Clear() error { - if err := os.RemoveAll(rc.Directory); err != nil { - return xerrors.New("failed to remove cache") + if err := os.RemoveAll(rc.directory); err != nil { + return xerrors.Errorf("failed to remove cache: %w", err) } return nil } diff --git a/cache/cache_test.go b/cache/cache_test.go index d9eb4ddf90..3e8ce7ebe8 100644 --- a/cache/cache_test.go +++ b/cache/cache_test.go @@ -3,6 +3,7 @@ package cache import ( "io/ioutil" "os" + "path/filepath" "testing" "github.com/stretchr/testify/assert" @@ -10,13 +11,10 @@ import ( func TestRealCache_Clear(t *testing.T) { d, _ := ioutil.TempDir("", "TestRealCache_Clear") - c := Initialize(d) + defer os.RemoveAll(d) + c, err := New(d) + assert.NoError(t, err) assert.NoError(t, c.Clear()) - _, err := os.Stat(d) + _, err = os.Stat(filepath.Join(d, "fanal")) assert.True(t, os.IsNotExist(err)) - - t.Run("sad path, cache dir doesn't exist", func(t *testing.T) { - c := Initialize(".") - assert.Equal(t, "failed to remove cache", c.Clear().Error()) - }) } diff --git a/cmd/fanal/main.go b/cmd/fanal/main.go index a98244fc8a..c0423e97f9 100644 --- a/cmd/fanal/main.go +++ b/cmd/fanal/main.go @@ -48,8 +48,10 @@ func run() (err error) { clearCache := flag.Bool("clear", false, "clear cache") flag.Parse() - c := cache.Initialize(utils.CacheDir() + "/cache.db") - + c, err := cache.New(utils.CacheDir()) + if err != nil { + return err + } if *clearCache { if err = c.Clear(); err != nil { return xerrors.Errorf("%w", err) @@ -63,7 +65,7 @@ func run() (err error) { SkipPing: true, } - ext, err := docker.NewDockerExtractor(opt) + ext, err := docker.NewDockerExtractor(opt, c) if err != nil { return err } diff --git a/extractor/docker/docker.go b/extractor/docker/docker.go index 08ba2b8327..35d2a06d62 100644 --- a/extractor/docker/docker.go +++ b/extractor/docker/docker.go @@ -14,9 +14,8 @@ import ( "strings" "time" - "github.com/simar7/gokv/encoding" - "github.com/aquasecurity/fanal/analyzer/library" + "github.com/aquasecurity/fanal/cache" "github.com/aquasecurity/fanal/extractor" "github.com/aquasecurity/fanal/extractor/docker/token/ecr" "github.com/aquasecurity/fanal/extractor/docker/token/gcr" @@ -28,8 +27,6 @@ import ( "github.com/klauspost/compress/zstd" "github.com/knqyf263/nested" "github.com/opencontainers/go-digest" - bolt "github.com/simar7/gokv/bbolt" - kvtypes "github.com/simar7/gokv/types" "golang.org/x/xerrors" ) @@ -72,11 +69,11 @@ type layer struct { type Extractor struct { Client *client.Client - Cache *bolt.Store Option types.DockerOption + cache cache.Cache } -func NewDockerExtractorWithCache(option types.DockerOption, cacheOptions bolt.Options) (Extractor, error) { +func NewDockerExtractor(option types.DockerOption, cache cache.Cache) (Extractor, error) { RegisterRegistry(&gcr.GCR{}) RegisterRegistry(&ecr.ECR{}) @@ -85,26 +82,13 @@ func NewDockerExtractorWithCache(option types.DockerOption, cacheOptions bolt.Op return Extractor{}, xerrors.Errorf("error initializing docker extractor: %w", err) } - var kv *bolt.Store - if kv, err = bolt.NewStore(cacheOptions); err != nil { - return Extractor{}, xerrors.Errorf("error initializing cache: %w", err) - } - return Extractor{ Option: option, Client: cli, - Cache: kv, + cache: cache, }, nil } -func NewDockerExtractor(option types.DockerOption) (Extractor, error) { - return NewDockerExtractorWithCache(option, bolt.Options{ - RootBucketName: "fanal", - Path: utils.CacheDir() + "/cache.db", // TODO: Make this configurable via a public method - Codec: encoding.Raw, - }) -} - func applyLayers(layerPaths []string, filesInLayers map[string]extractor.FileMap, opqInLayers map[string]extractor.OPQDirs) (extractor.FileMap, error) { sep := "/" nestedMap := nested.Nested{} @@ -171,11 +155,7 @@ func (d Extractor) SaveLocalImage(ctx context.Context, imageName string) (io.Rea var storedReader io.Reader var storedImageBytes []byte - found, err := d.Cache.Get(kvtypes.GetItemInput{ - BucketName: KVImageBucket, - Key: imageName, - Value: &storedImageBytes, - }) + found, err := d.cache.Get(KVImageBucket, imageName, &storedImageBytes) if found { dec, _ := zstd.NewReader(nil) @@ -206,11 +186,7 @@ func (d Extractor) SaveLocalImage(ctx context.Context, imageName string) (io.Rea } dst := e.EncodeAll(savedImage, nil) - if err := d.Cache.BatchSet(kvtypes.BatchSetItemInput{ - BucketName: "imagebucket", - Keys: []string{imageName}, - Values: dst, - }); err != nil { + if err := d.cache.Set(KVImageBucket, imageName, dst); err != nil { log.Println(err) } } @@ -318,11 +294,7 @@ func (d Extractor) extractLayerWorker(dig digest.Digest, r *registry.Registry, c var cacheContent []byte var cacheBuf bytes.Buffer - found, _ := d.Cache.Get(kvtypes.GetItemInput{ - BucketName: LayerTarsBucket, - Key: string(dig), - Value: &cacheContent, - }) + found, _ := d.cache.Get(LayerTarsBucket, string(dig), &cacheContent) if found { b, errTar := extractTarFromTarZstd(cacheContent) @@ -424,11 +396,7 @@ func (d Extractor) storeLayerInCache(cacheBuf bytes.Buffer, dig digest.Digest) { _, _ = io.Copy(w, &cacheBuf) _ = w.Close() - if err := d.Cache.BatchSet(kvtypes.BatchSetItemInput{ - BucketName: LayerTarsBucket, - Keys: []string{string(dig)}, - Values: dst.Bytes(), - }); err != nil { + if err := d.cache.Set(LayerTarsBucket, string(dig), dst.Bytes()); err != nil { log.Printf("an error occurred while caching: %s", err) } } diff --git a/extractor/docker/docker_test.go b/extractor/docker/docker_test.go index 857d42335c..c103124b23 100644 --- a/extractor/docker/docker_test.go +++ b/extractor/docker/docker_test.go @@ -15,28 +15,29 @@ import ( "github.com/klauspost/compress/zstd" + "github.com/aquasecurity/fanal/cache" "github.com/aquasecurity/fanal/extractor" "github.com/aquasecurity/fanal/types" "github.com/docker/docker/client" "github.com/genuinetools/reg/registry" "github.com/opencontainers/go-digest" - bolt "github.com/simar7/gokv/bbolt" - kvtypes "github.com/simar7/gokv/types" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) // TODO: Use a memory based FS rather than actual fs // context: https://github.com/aquasecurity/fanal/pull/51#discussion_r352337762 -func setupCache() (*bolt.Store, *os.File, error) { - f, err := ioutil.TempFile(".", "Bolt_TestStore-*") +func setupCache() (cache.Cache, string, error) { + dir, err := ioutil.TempDir("", "Cache_TestStore-*") if err != nil { - return nil, nil, err + return nil, "", err } - s, err := bolt.NewStore(bolt.Options{ - Path: f.Name(), - }) - return s, f, err + c, err := cache.New(dir) + if err != nil { + return nil, "", err + } + return c, dir, nil } func TestExtractFromFile(t *testing.T) { @@ -238,27 +239,20 @@ func TestDockerExtractor_SaveLocalImage(t *testing.T) { assert.NoError(t, err) // setup cache - s, f, err := setupCache() - defer func() { - _ = f.Close() - _ = os.RemoveAll(f.Name()) - }() + cache, tmpDir, err := setupCache() + defer os.RemoveAll(tmpDir) assert.NoError(t, err) if tc.cacheHit { e, _ := zstd.NewWriter(nil, zstd.WithEncoderLevel(zstd.SpeedDefault)) dst := e.EncodeAll([]byte("foofromcache"), nil) - _ = s.Set(kvtypes.SetItemInput{ - BucketName: "imagebucket", - Key: "fooimage", - Value: dst, - }) + _ = cache.Set(KVImageBucket, "fooimage", dst) } de := Extractor{ Option: types.DockerOption{}, Client: c, - Cache: s, + cache: cache, } r, err := de.SaveLocalImage(context.TODO(), "fooimage") @@ -268,11 +262,7 @@ func TestDockerExtractor_SaveLocalImage(t *testing.T) { // check the cache for what was stored var actualValue []byte - found, err := de.Cache.Get(kvtypes.GetItemInput{ - BucketName: "imagebucket", - Key: "fooimage", - Value: &actualValue, - }) + found, err := de.cache.Get(KVImageBucket, "fooimage", &actualValue) assert.NoError(t, err, tc.name) assert.True(t, found, tc.name) @@ -379,12 +369,9 @@ func TestDockerExtractor_Extract(t *testing.T) { assert.NoError(t, err) // setup cache - s, f, err := setupCache() - defer func() { - _ = f.Close() - _ = os.RemoveAll(f.Name()) - }() + s, tmpDir, err := setupCache() assert.NoError(t, err) + defer os.RemoveAll(tmpDir) de := Extractor{ Option: types.DockerOption{ @@ -394,7 +381,7 @@ func TestDockerExtractor_Extract(t *testing.T) { Timeout: time.Second * 1000, }, Client: c, - Cache: s, + cache: s, } tsURL := strings.TrimPrefix(ts.URL, "http://") @@ -478,28 +465,17 @@ func TestDocker_ExtractLayerWorker(t *testing.T) { assert.NoError(t, err) // setup cache - s, f, err := setupCache() - defer func() { - _ = f.Close() - _ = os.RemoveAll(f.Name()) - }() - assert.NoError(t, err, tc.name) + s, tmpDir, err := setupCache() + require.NoError(t, err, tc.name) + defer os.RemoveAll(tmpDir) if tc.cacheHit { switch tc.garbageCache { case true: garbage, _ := ioutil.ReadFile("testdata/invalidgzvalidtar.tar.gz") - assert.NoError(t, s.Set(kvtypes.SetItemInput{ - BucketName: LayerTarsBucket, - Key: string(inputDigest), - Value: garbage, - }), tc.name) + assert.NoError(t, s.Set(LayerTarsBucket, string(inputDigest), garbage)) default: - assert.NoError(t, s.Set(kvtypes.SetItemInput{ - BucketName: LayerTarsBucket, - Key: string(inputDigest), - Value: goodtarzstdgolden, - }), tc.name) + assert.NoError(t, s.Set(LayerTarsBucket, string(inputDigest), goodtarzstdgolden)) } } @@ -511,7 +487,7 @@ func TestDocker_ExtractLayerWorker(t *testing.T) { Timeout: time.Second * 1000, }, Client: c, - Cache: s, + cache: s, } tsUrl := strings.TrimPrefix(ts.URL, "http://") @@ -542,11 +518,7 @@ func TestDocker_ExtractLayerWorker(t *testing.T) { // check cache contents var actualCacheContents []byte - found, err := s.Get(kvtypes.GetItemInput{ - BucketName: LayerTarsBucket, - Key: string(inputDigest), - Value: &actualCacheContents, - }) + found, err := s.Get(LayerTarsBucket, string(inputDigest), &actualCacheContents) assert.True(t, found, tc.name) assert.NoError(t, err, tc.name) diff --git a/utils/utils.go b/utils/utils.go index 665206e9af..d406ccce2b 100644 --- a/utils/utils.go +++ b/utils/utils.go @@ -5,7 +5,6 @@ import ( "fmt" "os" "os/exec" - "path/filepath" ) var ( @@ -17,8 +16,7 @@ func CacheDir() string { if err != nil { cacheDir = os.TempDir() } - dir := filepath.Join(cacheDir, "fanal") - return dir + return cacheDir } func StringInSlice(a string, list []string) bool {