Skip to content

Commit

Permalink
Remove files after creating SBOM (#107)
Browse files Browse the repository at this point in the history 
* wip: remove files after return

* wip: removing files after creating sboms

* do not continue on invalid sboms

* remove gh actions
  • Loading branch information
David Wertenteil authored Jul 24, 2023
1 parent cdc49c3 commit ad37eb1
Show file tree
Hide file tree
Showing 4 changed files with 64 additions and 29 deletions.
11 changes: 4 additions & 7 deletions pkg/filehandler/interface.go
View file
Original file line number Diff line number Diff line change
@@ -1,12 +1,9 @@
package filehandler

import (
"context"
)

type FileHandler interface {
AddFile(ctx context.Context, bucket, file string) error
AddFile(bucket, file string) error
AddFiles(bucket string, files map[string]bool) error
Close()
GetFiles(ctx context.Context, container string) (map[string]bool, error)
RemoveBucket(ctx context.Context, bucket string) error
GetFiles(container string) (map[string]bool, error)
RemoveBucket(bucket string) error
}
11 changes: 7 additions & 4 deletions pkg/filehandler/v1/boltdb.go
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
package filehandler

import (
"context"
"fmt"
"node-agent/pkg/filehandler"

Expand All @@ -24,7 +23,7 @@ func CreateBoltFileHandler() (*BoltFileHandler, error) {
return &BoltFileHandler{fileDB: db}, nil
}

func (b BoltFileHandler) AddFile(ctx context.Context, bucket, file string) error {
func (b BoltFileHandler) AddFile(bucket, file string) error {
return b.fileDB.Batch(func(tx *bolt.Tx) error {
b, err := tx.CreateBucketIfNotExists([]byte(bucket))
if err != nil {
Expand All @@ -38,7 +37,7 @@ func (b BoltFileHandler) Close() {
_ = b.fileDB.Close()
}

func (b BoltFileHandler) GetFiles(ctx context.Context, container string) (map[string]bool, error) {
func (b BoltFileHandler) GetFiles(container string) (map[string]bool, error) {
fileList := make(map[string]bool)
err := b.fileDB.View(func(tx *bolt.Tx) error {
b := tx.Bucket([]byte(container))
Expand All @@ -54,7 +53,7 @@ func (b BoltFileHandler) GetFiles(ctx context.Context, container string) (map[st
return fileList, err
}

func (b BoltFileHandler) RemoveBucket(ctx context.Context, bucket string) error {
func (b BoltFileHandler) RemoveBucket(bucket string) error {
return b.fileDB.Update(func(tx *bolt.Tx) error {
err := tx.DeleteBucket([]byte(bucket))
if err != nil {
Expand All @@ -64,3 +63,7 @@ func (b BoltFileHandler) RemoveBucket(ctx context.Context, bucket string) error
return nil
})
}
func (b BoltFileHandler) AddFiles(bucket string, files map[string]bool) error {
// do nothing
return nil
}
40 changes: 36 additions & 4 deletions pkg/filehandler/v1/inmemory.go
View file
Original file line number Diff line number Diff line change
@@ -1,13 +1,13 @@
package filehandler

import (
"context"
"fmt"
"node-agent/pkg/filehandler"
"sync"
)

const initFileListLength = 5000
const updateFileListLength = 200

type filesBucket struct {
lock *sync.RWMutex
Expand All @@ -27,7 +27,7 @@ func CreateInMemoryFileHandler() (*InMemoryFileHandler, error) {
}, nil
}

func (s *InMemoryFileHandler) AddFile(ctx context.Context, bucket, file string) error {
func (s *InMemoryFileHandler) AddFile(bucket, file string) error {
// Acquire a read lock first
s.mutex.RLock()
bucketFiles, ok := s.buckets[bucket]
Expand Down Expand Up @@ -71,7 +71,7 @@ func shallowCopyMapStringBool(m map[string]bool) map[string]bool {
return mCopy
}

func (s *InMemoryFileHandler) GetFiles(ctx context.Context, bucket string) (map[string]bool, error) {
func (s *InMemoryFileHandler) GetFiles(bucket string) (map[string]bool, error) {
s.mutex.RLock()
bucketFiles, ok := s.buckets[bucket]
s.mutex.RUnlock()
Expand All @@ -82,14 +82,46 @@ func (s *InMemoryFileHandler) GetFiles(ctx context.Context, bucket string) (map[

bucketFiles.lock.RLock()
copy := shallowCopyMapStringBool(bucketFiles.files)
bucketFiles.files = make(map[string]bool, updateFileListLength)
bucketFiles.lock.RUnlock()

return copy, nil
}
func (s *InMemoryFileHandler) RemoveBucket(ctx context.Context, bucket string) error {
func (s *InMemoryFileHandler) RemoveBucket(bucket string) error {
s.mutex.Lock()
delete(s.buckets, bucket)
s.mutex.Unlock()

return nil
}

func (s *InMemoryFileHandler) AddFiles(bucket string, files map[string]bool) error {
// Acquire a read lock first
s.mutex.RLock()
bucketFiles, ok := s.buckets[bucket]
s.mutex.RUnlock()

// If the bucket doesn't exist, acquire a write lock to create the new bucket
if !ok {
s.mutex.Lock()
// Double-check the bucket's existence to ensure another goroutine didn't already create it
bucketFiles, ok = s.buckets[bucket]
if !ok {
bucketFiles = &filesBucket{
lock: &sync.RWMutex{},
files: make(map[string]bool, initFileListLength),
}
s.buckets[bucket] = bucketFiles
}
s.mutex.Unlock()
}

// Acquire a write lock if the bucket already exists
bucketFiles.lock.Lock()
for file := range files {
bucketFiles.files[file] = true
}
bucketFiles.lock.Unlock()

return nil
}
31 changes: 17 additions & 14 deletions pkg/relevancymanager/v1/relevancy_manager.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -78,38 +78,40 @@ func (rm *RelevancyManager) afterTimerActions(ctx context.Context) error {
}
containerData := containerDataInterface.(watchedContainerData)

logger.L().Debug("getting files", helpers.String("container ID", afterTimerActionsData.containerID), helpers.String("k8s workload", containerData.k8sContainerID))

if rm.cfg.EnableRelevancy && afterTimerActionsData.service == RelevantCVEsService {
fileList, err := rm.fileHandler.GetFiles(ctx, containerData.k8sContainerID)
if err != nil {
logger.L().Debug("failed to get file list", helpers.String("container ID", afterTimerActionsData.containerID), helpers.String("k8s workload", containerData.k8sContainerID), helpers.Error(err))
continue
}
logger.L().Debug("fileList generated", helpers.String("container ID", afterTimerActionsData.containerID), helpers.String("k8s workload", containerData.k8sContainerID), helpers.String("file list", fmt.Sprintf("%v", fileList)))

ctxPostSBOM, spanPostSBOM := otel.Tracer("").Start(ctx, "PostFilterSBOM")
if err = <-containerData.syncChannel[StepGetSBOM]; err != nil {
if err := <-containerData.syncChannel[StepGetSBOM]; err != nil {
logger.L().Debug("failed to get SBOM", helpers.String("container ID", afterTimerActionsData.containerID), helpers.String("k8s workload", containerData.k8sContainerID), helpers.Error(err))
continue
}
if containerData.sbomClient == nil {
// it is possible that the sbom client was not created yet
logger.L().Debug("sbom client not yet created", helpers.String("container ID", afterTimerActionsData.containerID), helpers.String("k8s workload", containerData.k8sContainerID), helpers.Error(err))
logger.L().Debug("sbom client not yet created", helpers.String("container ID", afterTimerActionsData.containerID), helpers.String("k8s workload", containerData.k8sContainerID))
continue
}
if err = containerData.sbomClient.ValidateSBOM(ctx); err != nil {
if err := containerData.sbomClient.ValidateSBOM(ctx); err != nil {
logger.L().Warning("SBOM is incomplete", helpers.String("container ID", afterTimerActionsData.containerID), helpers.String("k8s workload", containerData.k8sContainerID), helpers.Error(err))
containerData.syncChannel[StepValidateSBOM] <- err
}

fileList, err := rm.fileHandler.GetFiles(containerData.k8sContainerID)
if err != nil {
logger.L().Debug("failed to get file list", helpers.String("container ID", afterTimerActionsData.containerID), helpers.String("k8s workload", containerData.k8sContainerID), helpers.Error(err))
continue
}
logger.L().Debug("fileList generated", helpers.String("container ID", afterTimerActionsData.containerID), helpers.String("k8s workload", containerData.k8sContainerID), helpers.String("file list", fmt.Sprintf("%v", fileList)))

if err = containerData.sbomClient.FilterSBOM(ctx, fileList); err != nil {
rm.fileHandler.AddFiles(containerData.k8sContainerID, fileList)
ctx, span := otel.Tracer("").Start(ctxPostSBOM, "FilterSBOM")
logger.L().Ctx(ctx).Warning("failed to filter SBOM", helpers.String("container ID", afterTimerActionsData.containerID), helpers.String("k8s workload", containerData.k8sContainerID), helpers.Error(err))
span.End()
continue
}
filterSBOMKey, err := containerData.instanceID.GetSlug()
if err != nil {
rm.fileHandler.AddFiles(containerData.k8sContainerID, fileList)
ctx, span := otel.Tracer("").Start(ctxPostSBOM, "filterSBOMKey")
logger.L().Ctx(ctx).Warning("failed to get filterSBOMKey for store filter SBOM", helpers.String("container ID", afterTimerActionsData.containerID), helpers.String("k8s workload", containerData.k8sContainerID), helpers.Error(err))
span.End()
Expand All @@ -118,6 +120,7 @@ func (rm *RelevancyManager) afterTimerActions(ctx context.Context) error {
// it is safe to use containerData.imageID directly since we needed it to retrieve the SBOM
if err = containerData.sbomClient.StoreFilterSBOM(ctx, containerData.imageID, filterSBOMKey); err != nil {
if !errors.Is(err, sbom.IsAlreadyExist()) {
rm.fileHandler.AddFiles(containerData.k8sContainerID, fileList)
ctx, span := otel.Tracer("").Start(ctxPostSBOM, "StoreFilterSBOM")
logger.L().Ctx(ctx).Error("failed to store filtered SBOM", helpers.String("container ID", afterTimerActionsData.containerID), helpers.String("k8s workload", containerData.k8sContainerID), helpers.Error(err))
span.End()
Expand All @@ -139,7 +142,7 @@ func (rm *RelevancyManager) deleteResources(watchedContainer watchedContainerDat
rm.watchedContainers.Delete(containerID)

// Remove container from the file DB
rm.fileHandler.RemoveBucket(context.Background(), containerID)
rm.fileHandler.RemoveBucket(containerID)
}

func (rm *RelevancyManager) getSBOM(ctx context.Context, container *containercollection.Container) {
Expand Down Expand Up @@ -345,7 +348,7 @@ func (rm *RelevancyManager) ReportContainerTerminated(ctx context.Context, conta
logger.L().Debug("container not found in memory", helpers.String("container ID", container.ID), helpers.String("k8s workload", k8sContainerID))
return
}
err := rm.fileHandler.RemoveBucket(ctx, k8sContainerID)
err := rm.fileHandler.RemoveBucket(k8sContainerID)
if err != nil {
logger.L().Error("failed to remove container bucket", helpers.Error(err), helpers.String("container ID", container.ID), helpers.String("k8s workload", k8sContainerID))
return
Expand All @@ -361,7 +364,7 @@ func (rm *RelevancyManager) ReportFileAccess(ctx context.Context, namespace, pod
return
}
k8sContainerID := utils.CreateK8sContainerID(namespace, pod, container)
err := rm.fileHandler.AddFile(ctx, k8sContainerID, file)
err := rm.fileHandler.AddFile(k8sContainerID, file)
if err != nil {
logger.L().Ctx(ctx).Error("failed to add file to container file list", helpers.Error(err), helpers.Interface("k8sContainerID", k8sContainerID), helpers.String("file", file))
}
Expand Down

0 comments on commit ad37eb1

Please sign in to comment.