github.com/devseccon/trivy@v0.47.1-0.20231123133102-bd902a0bd996/pkg/fanal/cache/redis.go (about)

     1  package cache
     2  
     3  import (
     4  	"context"
     5  	"encoding/json"
     6  	"fmt"
     7  	"time"
     8  
     9  	"github.com/go-redis/redis/v8"
    10  	"github.com/hashicorp/go-multierror"
    11  	"golang.org/x/xerrors"
    12  
    13  	"github.com/devseccon/trivy/pkg/fanal/types"
    14  )
    15  
    16  var _ Cache = &RedisCache{}
    17  
    18  const (
    19  	redisPrefix = "fanal"
    20  )
    21  
    22  type RedisCache struct {
    23  	client     *redis.Client
    24  	expiration time.Duration
    25  }
    26  
    27  func NewRedisCache(options *redis.Options, expiration time.Duration) RedisCache {
    28  	return RedisCache{
    29  		client:     redis.NewClient(options),
    30  		expiration: expiration,
    31  	}
    32  }
    33  
    34  func (c RedisCache) PutArtifact(artifactID string, artifactConfig types.ArtifactInfo) error {
    35  	key := fmt.Sprintf("%s::%s::%s", redisPrefix, artifactBucket, artifactID)
    36  	b, err := json.Marshal(artifactConfig)
    37  	if err != nil {
    38  		return xerrors.Errorf("failed to marshal artifact JSON: %w", err)
    39  	}
    40  	if err := c.client.Set(context.TODO(), key, string(b), c.expiration).Err(); err != nil {
    41  		return xerrors.Errorf("unable to store artifact information in Redis cache (%s): %w", artifactID, err)
    42  	}
    43  	return nil
    44  }
    45  
    46  func (c RedisCache) PutBlob(blobID string, blobInfo types.BlobInfo) error {
    47  	b, err := json.Marshal(blobInfo)
    48  	if err != nil {
    49  		return xerrors.Errorf("failed to marshal blob JSON: %w", err)
    50  	}
    51  	key := fmt.Sprintf("%s::%s::%s", redisPrefix, blobBucket, blobID)
    52  	if err := c.client.Set(context.TODO(), key, string(b), c.expiration).Err(); err != nil {
    53  		return xerrors.Errorf("unable to store blob information in Redis cache (%s): %w", blobID, err)
    54  	}
    55  	return nil
    56  }
    57  func (c RedisCache) DeleteBlobs(blobIDs []string) error {
    58  	var errs error
    59  	for _, blobID := range blobIDs {
    60  		key := fmt.Sprintf("%s::%s::%s", redisPrefix, artifactBucket, blobID)
    61  		if err := c.client.Del(context.TODO(), key).Err(); err != nil {
    62  			errs = multierror.Append(errs, xerrors.Errorf("unable to delete blob %s: %w", blobID, err))
    63  		}
    64  	}
    65  	return errs
    66  }
    67  
    68  func (c RedisCache) GetArtifact(artifactID string) (types.ArtifactInfo, error) {
    69  	key := fmt.Sprintf("%s::%s::%s", redisPrefix, artifactBucket, artifactID)
    70  	val, err := c.client.Get(context.TODO(), key).Bytes()
    71  	if err == redis.Nil {
    72  		return types.ArtifactInfo{}, xerrors.Errorf("artifact (%s) is missing in Redis cache", artifactID)
    73  	} else if err != nil {
    74  		return types.ArtifactInfo{}, xerrors.Errorf("failed to get artifact from the Redis cache: %w", err)
    75  	}
    76  
    77  	var info types.ArtifactInfo
    78  	err = json.Unmarshal(val, &info)
    79  	if err != nil {
    80  		return types.ArtifactInfo{}, xerrors.Errorf("failed to unmarshal artifact (%s) from Redis value: %w", artifactID, err)
    81  	}
    82  	return info, nil
    83  }
    84  
    85  func (c RedisCache) GetBlob(blobID string) (types.BlobInfo, error) {
    86  	key := fmt.Sprintf("%s::%s::%s", redisPrefix, blobBucket, blobID)
    87  	val, err := c.client.Get(context.TODO(), key).Bytes()
    88  	if err == redis.Nil {
    89  		return types.BlobInfo{}, xerrors.Errorf("blob (%s) is missing in Redis cache", blobID)
    90  	} else if err != nil {
    91  		return types.BlobInfo{}, xerrors.Errorf("failed to get blob from the Redis cache: %w", err)
    92  	}
    93  
    94  	var blobInfo types.BlobInfo
    95  	if err = json.Unmarshal(val, &blobInfo); err != nil {
    96  		return types.BlobInfo{}, xerrors.Errorf("failed to unmarshal blob (%s) from Redis value: %w", blobID, err)
    97  	}
    98  	return blobInfo, nil
    99  }
   100  
   101  func (c RedisCache) MissingBlobs(artifactID string, blobIDs []string) (bool, []string, error) {
   102  	var missingArtifact bool
   103  	var missingBlobIDs []string
   104  	for _, blobID := range blobIDs {
   105  		blobInfo, err := c.GetBlob(blobID)
   106  		if err != nil {
   107  			// error means cache missed blob info
   108  			missingBlobIDs = append(missingBlobIDs, blobID)
   109  			continue
   110  		}
   111  		if blobInfo.SchemaVersion != types.BlobJSONSchemaVersion {
   112  			missingBlobIDs = append(missingBlobIDs, blobID)
   113  		}
   114  	}
   115  	// get artifact info
   116  	artifactInfo, err := c.GetArtifact(artifactID)
   117  	// error means cache missed artifact info
   118  	if err != nil {
   119  		return true, missingBlobIDs, nil
   120  	}
   121  	if artifactInfo.SchemaVersion != types.ArtifactJSONSchemaVersion {
   122  		missingArtifact = true
   123  	}
   124  	return missingArtifact, missingBlobIDs, nil
   125  }
   126  
   127  func (c RedisCache) Close() error {
   128  	return c.client.Close()
   129  }
   130  
   131  func (c RedisCache) Clear() error {
   132  	ctx := context.Background()
   133  
   134  	var cursor uint64
   135  	for {
   136  		var keys []string
   137  		var err error
   138  		keys, cursor, err = c.client.Scan(ctx, cursor, redisPrefix+"::*", 100).Result()
   139  		if err != nil {
   140  			return xerrors.Errorf("failed to perform prefix scanning: %w", err)
   141  		}
   142  		if err = c.client.Unlink(ctx, keys...).Err(); err != nil {
   143  			return xerrors.Errorf("failed to unlink redis keys: %w", err)
   144  		}
   145  		if cursor == 0 {
   146  			break
   147  		}
   148  	}
   149  	return nil
   150  }