Shahzad Haider
30d5bb91af
Test / test (push) Waiting to run
Test / test-community (push) Waiting to run
Lint / golangci-lint (push) Waiting to run
Lint / man-page-staleness (push) Waiting to run
Lint / semgrep (push) Waiting to run
Lint / checksecretparts (push) Waiting to run
Scan for secrets / test (push) Waiting to run
Release / mark-latest (push) Canceled after 0s
Release / Release (push) Canceled after 0s
* s3: surface bucket listing failures for explicitly configured buckets Listing failures were logged at V(3) whenever a role was assumed, hiding access and role-assumption errors for buckets the user explicitly asked to scan. Suppression now applies only in list-all-buckets mode, and a bucket_list_errors_total metric records every listing failure. * s3: accumulate object count across role passes scanBuckets runs once per configured role and reset its object counter each pass, so the final progress message only reflected the last role's count. Multi-role scans could report 0 objects scanned even when earlier roles scanned objects.
795 lines
26 KiB
Go
795 lines
26 KiB
Go
package s3
|
||
|
||
import (
|
||
"encoding/json"
|
||
"fmt"
|
||
"slices"
|
||
"strings"
|
||
"sync"
|
||
"sync/atomic"
|
||
"time"
|
||
|
||
"github.com/aws/aws-sdk-go-v2/aws"
|
||
"github.com/aws/aws-sdk-go-v2/config"
|
||
"github.com/aws/aws-sdk-go-v2/credentials"
|
||
"github.com/aws/aws-sdk-go-v2/credentials/stscreds"
|
||
s3manager "github.com/aws/aws-sdk-go-v2/feature/s3/manager"
|
||
"github.com/aws/aws-sdk-go-v2/service/s3"
|
||
s3types "github.com/aws/aws-sdk-go-v2/service/s3/types"
|
||
"github.com/aws/aws-sdk-go-v2/service/sts"
|
||
"github.com/go-errors/errors"
|
||
"golang.org/x/sync/errgroup"
|
||
"google.golang.org/protobuf/proto"
|
||
"google.golang.org/protobuf/types/known/anypb"
|
||
|
||
"github.com/trufflesecurity/trufflehog/v3/pkg/common"
|
||
"github.com/trufflesecurity/trufflehog/v3/pkg/context"
|
||
"github.com/trufflesecurity/trufflehog/v3/pkg/handlers"
|
||
"github.com/trufflesecurity/trufflehog/v3/pkg/log"
|
||
"github.com/trufflesecurity/trufflehog/v3/pkg/pb/source_metadatapb"
|
||
"github.com/trufflesecurity/trufflehog/v3/pkg/pb/sourcespb"
|
||
"github.com/trufflesecurity/trufflehog/v3/pkg/sanitizer"
|
||
"github.com/trufflesecurity/trufflehog/v3/pkg/sources"
|
||
)
|
||
|
||
const (
|
||
SourceType = sourcespb.SourceType_SOURCE_TYPE_S3
|
||
|
||
defaultAWSRegion = "us-east-1"
|
||
defaultMaxObjectSize = 250 * 1024 * 1024 // 250 MiB
|
||
maxObjectSizeLimit = 250 * 1024 * 1024 // 250 MiB
|
||
)
|
||
|
||
type Source struct {
|
||
name string
|
||
sourceID sources.SourceID
|
||
jobID sources.JobID
|
||
verify bool
|
||
concurrency int
|
||
conn *sourcespb.S3
|
||
|
||
sources.Progress
|
||
metricsCollector metricsCollector
|
||
|
||
errorCount *sync.Map
|
||
jobPool *errgroup.Group
|
||
maxObjectSize int64
|
||
}
|
||
|
||
// Ensure the Source satisfies the interfaces at compile time
|
||
var _ sources.Source = (*Source)(nil)
|
||
var _ sources.SourceUnitUnmarshaller = (*Source)(nil)
|
||
var _ sources.Validator = (*Source)(nil)
|
||
var _ sources.SourceUnitEnumChunker = (*Source)(nil)
|
||
|
||
// Type returns the type of source
|
||
func (s *Source) Type() sourcespb.SourceType { return SourceType }
|
||
|
||
func (s *Source) SourceID() sources.SourceID { return s.sourceID }
|
||
|
||
func (s *Source) JobID() sources.JobID { return s.jobID }
|
||
|
||
// Init returns an initialized AWS source
|
||
func (s *Source) Init(
|
||
ctx context.Context,
|
||
name string,
|
||
jobID sources.JobID,
|
||
sourceID sources.SourceID,
|
||
verify bool,
|
||
connection *anypb.Any,
|
||
concurrency int,
|
||
) error {
|
||
s.name = name
|
||
s.sourceID = sourceID
|
||
s.jobID = jobID
|
||
s.verify = verify
|
||
s.concurrency = concurrency
|
||
s.errorCount = &sync.Map{}
|
||
s.jobPool = &errgroup.Group{}
|
||
s.jobPool.SetLimit(concurrency)
|
||
|
||
var conn sourcespb.S3
|
||
if err := anypb.UnmarshalTo(connection, &conn, proto.UnmarshalOptions{}); err != nil {
|
||
return fmt.Errorf("error unmarshalling connection: %w", err)
|
||
}
|
||
s.conn = &conn
|
||
|
||
s.metricsCollector = metricsInstance
|
||
|
||
s.setMaxObjectSize(conn.GetMaxObjectSize())
|
||
|
||
if len(conn.GetBuckets()) > 0 && len(conn.GetIgnoreBuckets()) > 0 {
|
||
return errors.New("either a bucket include list or a bucket ignore list can be specified, but not both")
|
||
}
|
||
|
||
return nil
|
||
}
|
||
|
||
func (s *Source) Validate(ctx context.Context) []error {
|
||
var errs []error
|
||
visitor := func(c context.Context, defaultRegionClient *s3.Client, roleArn string, buckets []string) error {
|
||
roleErrs := s.validateBucketAccess(c, defaultRegionClient, roleArn, buckets)
|
||
if len(roleErrs) > 0 {
|
||
errs = append(errs, roleErrs...)
|
||
}
|
||
return nil
|
||
}
|
||
|
||
if err := s.visitRoles(ctx, visitor); err != nil {
|
||
errs = append(errs, err)
|
||
}
|
||
|
||
return errs
|
||
}
|
||
|
||
// setMaxObjectSize sets the maximum size of objects that will be scanned. If
|
||
// not set, set to a negative number, or set larger than the
|
||
// maxObjectSizeLimit, the defaultMaxObjectSizeLimit will be used.
|
||
func (s *Source) setMaxObjectSize(maxObjectSize int64) {
|
||
if maxObjectSize <= 0 || maxObjectSize > maxObjectSizeLimit {
|
||
s.maxObjectSize = defaultMaxObjectSize
|
||
} else {
|
||
s.maxObjectSize = maxObjectSize
|
||
}
|
||
}
|
||
|
||
func (s *Source) newClient(ctx context.Context, region, roleArn string) (*s3.Client, error) {
|
||
var credsProvider aws.CredentialsProvider
|
||
switch cred := s.conn.GetCredential().(type) {
|
||
case *sourcespb.S3_SessionToken:
|
||
credsProvider = credentials.NewStaticCredentialsProvider(
|
||
cred.SessionToken.GetKey(),
|
||
cred.SessionToken.GetSecret(),
|
||
cred.SessionToken.GetSessionToken(),
|
||
)
|
||
log.RedactGlobally(cred.SessionToken.GetSecret())
|
||
log.RedactGlobally(cred.SessionToken.GetSessionToken())
|
||
case *sourcespb.S3_AccessKey:
|
||
credsProvider = credentials.NewStaticCredentialsProvider(cred.AccessKey.GetKey(), cred.AccessKey.GetSecret(), "")
|
||
log.RedactGlobally(cred.AccessKey.GetSecret())
|
||
case *sourcespb.S3_Unauthenticated:
|
||
credsProvider = aws.AnonymousCredentials{}
|
||
default:
|
||
// In all other cases, the AWS SDK will follow its normal waterfall logic to pick up credentials (i.e. they can
|
||
// come from the environment or the credentials file or whatever else AWS gets up to).
|
||
}
|
||
|
||
if roleArn != "" {
|
||
// The config loaded here will be used to retrieve and refresh temporary credentials from AssumeRole
|
||
cfg, err := config.LoadDefaultConfig(ctx, config.WithRegion(region), config.WithCredentialsProvider(credsProvider))
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
stsClient := sts.NewFromConfig(cfg)
|
||
provider := stscreds.NewAssumeRoleProvider(stsClient, roleArn, func(options *stscreds.AssumeRoleOptions) {
|
||
options.RoleSessionName = "trufflehog"
|
||
})
|
||
// From https://docs.aws.amazon.com/sdk-for-go/v2/developer-guide/configure-gosdk.html#specify-credentials-programmatically:
|
||
// "If you explicitly configure a provider on aws.Config directly,
|
||
// you must also explicitly wrap the provider with this type using NewCredentialsCache"
|
||
credsProvider = aws.NewCredentialsCache(provider)
|
||
}
|
||
|
||
cfg, err := config.LoadDefaultConfig(
|
||
ctx,
|
||
config.WithRegion(region),
|
||
config.WithCredentialsProvider(credsProvider),
|
||
)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
return s3.NewFromConfig(cfg, func(options *s3.Options) {
|
||
options.DisableLogOutputChecksumValidationSkipped = true
|
||
}), nil
|
||
}
|
||
|
||
// getBucketsToScan returns a list of S3 buckets to scan.
|
||
// If the connection has a list of buckets specified, those are returned.
|
||
// Otherwise, it lists all buckets the client has access to and filters out the ignored ones.
|
||
// The list of buckets is sorted lexicographically to ensure consistent ordering,
|
||
// which allows resuming scanning from the same place if the scan is interrupted.
|
||
//
|
||
// Note: The IAM identity needs the s3:ListBuckets permission.
|
||
func (s *Source) getBucketsToScan(ctx context.Context, client *s3.Client) ([]string, error) {
|
||
if buckets := s.conn.GetBuckets(); len(buckets) > 0 {
|
||
slices.Sort(buckets)
|
||
return buckets, nil
|
||
}
|
||
|
||
ignore := make(map[string]struct{}, len(s.conn.GetIgnoreBuckets()))
|
||
for _, bucket := range s.conn.GetIgnoreBuckets() {
|
||
ignore[bucket] = struct{}{}
|
||
}
|
||
|
||
res, err := client.ListBuckets(ctx, &s3.ListBucketsInput{})
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
var bucketsToScan []string
|
||
for _, bucket := range res.Buckets {
|
||
name := *bucket.Name
|
||
if _, ignored := ignore[name]; !ignored {
|
||
bucketsToScan = append(bucketsToScan, name)
|
||
}
|
||
}
|
||
slices.Sort(bucketsToScan)
|
||
|
||
return bucketsToScan, nil
|
||
}
|
||
|
||
// pageMetadata contains metadata about a single page of S3 objects being scanned.
|
||
type pageMetadata struct {
|
||
bucket string // The name of the S3 bucket being scanned
|
||
role string // The AWS role ARN used for scanning
|
||
pageNumber int // Current page number in the pagination sequence
|
||
client *s3.Client // AWS S3 client configured for the appropriate region
|
||
page *s3.ListObjectsV2Output // Contains the list of S3 objects in this page
|
||
}
|
||
|
||
// processingState tracks the state of concurrent S3 object processing.
|
||
type processingState struct {
|
||
errorCount *sync.Map // Thread-safe map tracking errors per prefix
|
||
objectCount *uint64 // Total number of objects processed
|
||
}
|
||
|
||
// resumePosition tracks where to restart scanning S3 buckets and objects after an interruption.
|
||
// It encapsulates all the information needed to resume a scan from its last known position.
|
||
type resumePosition struct {
|
||
bucket string // The bucket name we were processing
|
||
index int // Index in the buckets slice where we should resume
|
||
startAfter string // The last processed object key within the bucket
|
||
isNewScan bool // True if we're starting a fresh scan
|
||
exactMatch bool // True if we found the exact bucket we were previously processing
|
||
role string // The role used during the previous scan
|
||
}
|
||
|
||
// determineResumePosition calculates where to resume scanning from based on the last saved checkpoint
|
||
// and the current list of available buckets to scan. It handles several scenarios:
|
||
//
|
||
// 1. If getting the resume point fails or there is no previous bucket saved (CurrentBucket is empty),
|
||
// we start a new scan from the beginning, this is the safest option.
|
||
//
|
||
// 2. If the previous bucket exists in our current scan list (exactMatch=true),
|
||
// we resume from that exact position and use the StartAfter value
|
||
// to continue from the last processed object within that bucket.
|
||
//
|
||
// 3. If the previous bucket is not found in our current scan list (exactMatch=false), this typically means:
|
||
// - The bucket was deleted since our last scan
|
||
// - The bucket was explicitly excluded from this scan's configuration
|
||
// - The IAM role no longer has access to the bucket
|
||
// - The bucket name changed due to a configuration update
|
||
// In this case, we use binary search to find the closest position where the bucket would have been,
|
||
// allowing us to resume from the nearest available point in our sorted bucket list rather than
|
||
// restarting the entire scan.
|
||
func determineResumePosition(ctx context.Context, tracker *Checkpointer, buckets []string) resumePosition {
|
||
resumePoint, err := tracker.ResumePoint(ctx)
|
||
if err != nil {
|
||
ctx.Logger().Error(err, "failed to get resume point; starting from the beginning")
|
||
return resumePosition{isNewScan: true}
|
||
}
|
||
|
||
if resumePoint.CurrentBucket == "" {
|
||
return resumePosition{isNewScan: true}
|
||
}
|
||
|
||
startIdx, found := slices.BinarySearch(buckets, resumePoint.CurrentBucket)
|
||
return resumePosition{
|
||
bucket: resumePoint.CurrentBucket,
|
||
startAfter: resumePoint.StartAfter,
|
||
index: startIdx,
|
||
exactMatch: found,
|
||
role: resumePoint.Role,
|
||
}
|
||
}
|
||
|
||
// scanBuckets scans the given buckets using the given role and adds the number
|
||
// of objects it scanned to totalObjectCount. The counter is owned by Chunks and
|
||
// shared across role passes so that the completion message reflects the whole
|
||
// scan, not just the last role's pass.
|
||
func (s *Source) scanBuckets(
|
||
ctx context.Context,
|
||
client *s3.Client,
|
||
role string,
|
||
bucketsToScan []string,
|
||
chunksChan chan *sources.Chunk,
|
||
totalObjectCount *uint64,
|
||
) {
|
||
if role != "" {
|
||
ctx = context.WithValue(ctx, "role", role)
|
||
}
|
||
|
||
checkpointer := NewCheckpointer(ctx, &s.Progress, false)
|
||
pos := determineResumePosition(ctx, checkpointer, bucketsToScan)
|
||
switch {
|
||
case pos.isNewScan:
|
||
ctx.Logger().Info("Starting new scan from beginning")
|
||
case !pos.exactMatch:
|
||
ctx.Logger().Info(
|
||
"Resume bucket no longer available, starting from closest position",
|
||
"original_bucket", pos.bucket,
|
||
"position", pos.index,
|
||
"role", pos.role,
|
||
)
|
||
default:
|
||
ctx.Logger().Info(
|
||
"Resuming scan from previous scan's bucket",
|
||
"bucket", pos.bucket,
|
||
"position", pos.index,
|
||
"role", pos.role,
|
||
)
|
||
}
|
||
|
||
bucketsToScanCount := len(bucketsToScan)
|
||
for bucketIdx := pos.index; bucketIdx < bucketsToScanCount; bucketIdx++ {
|
||
bucket := bucketsToScan[bucketIdx]
|
||
|
||
s.SetProgressComplete(
|
||
bucketIdx,
|
||
len(bucketsToScan),
|
||
fmt.Sprintf("Bucket: %s", bucket),
|
||
s.EncodedResumeInfo,
|
||
)
|
||
|
||
var startAfter *string
|
||
if bucket == pos.bucket && pos.startAfter != "" && role == pos.role {
|
||
startAfter = &pos.startAfter
|
||
ctx.Logger().V(3).Info(
|
||
"Resuming bucket scan",
|
||
"start_after", pos.startAfter,
|
||
"bucket", bucket,
|
||
)
|
||
}
|
||
|
||
objectCount := s.scanBucket(ctx, client, role, bucket, sources.ChanReporter{Ch: chunksChan}, startAfter, checkpointer)
|
||
*totalObjectCount += objectCount
|
||
}
|
||
|
||
s.SetProgressComplete(
|
||
len(bucketsToScan),
|
||
len(bucketsToScan),
|
||
fmt.Sprintf("Completed scanning source %s. %d objects scanned.", s.name, *totalObjectCount),
|
||
"",
|
||
)
|
||
}
|
||
|
||
func (s *Source) scanBucket(
|
||
ctx context.Context,
|
||
client *s3.Client,
|
||
role string,
|
||
bucket string,
|
||
reporter sources.ChunkReporter,
|
||
startAfter *string,
|
||
checkpointer *Checkpointer,
|
||
) uint64 {
|
||
s.metricsCollector.RecordBucketForRole(role)
|
||
|
||
ctx = context.WithValue(ctx, "bucket", bucket)
|
||
|
||
if common.IsDone(ctx) {
|
||
ctx.Logger().Error(ctx.Err(), "context done, while scanning bucket")
|
||
return 0
|
||
}
|
||
|
||
ctx.Logger().V(3).Info("Scanning bucket")
|
||
|
||
regionalClient, err := s.getRegionalClientForBucket(ctx, client, role, bucket)
|
||
if err != nil {
|
||
ctx.Logger().Error(err, "could not get regional client for bucket")
|
||
return 0
|
||
}
|
||
|
||
errorCount := sync.Map{}
|
||
|
||
input := &s3.ListObjectsV2Input{Bucket: &bucket}
|
||
if startAfter != nil {
|
||
input.StartAfter = startAfter
|
||
}
|
||
|
||
pageNumber := 1
|
||
paginator := s3.NewListObjectsV2Paginator(regionalClient, input)
|
||
var objectCount uint64
|
||
for paginator.HasMorePages() {
|
||
output, err := paginator.NextPage(ctx)
|
||
if err != nil {
|
||
if s.listErrorsAreExpected(role) {
|
||
// Our documentation blesses specifying a role to assume without specifying buckets to scan, which will
|
||
// often cause this to happen a lot (because in that case the scanner tries to scan every bucket in the
|
||
// account, but the role probably doesn't have access to all of them). This makes it expected behavior
|
||
// and therefore not an error.
|
||
ctx.Logger().V(3).Info("could not list objects in bucket", "err", err)
|
||
} else {
|
||
// This can also be a failure to assume the role itself: role credentials
|
||
// are resolved lazily, so the first request that needs them surfaces the
|
||
// STS error here rather than at client construction.
|
||
ctx.Logger().Error(err, "could not list objects in bucket")
|
||
}
|
||
s.metricsCollector.RecordBucketListError(bucket, role)
|
||
break
|
||
}
|
||
pageMetadata := pageMetadata{
|
||
bucket: bucket,
|
||
role: role,
|
||
pageNumber: pageNumber,
|
||
client: regionalClient,
|
||
page: output,
|
||
}
|
||
processingState := processingState{
|
||
errorCount: &errorCount,
|
||
objectCount: &objectCount,
|
||
}
|
||
s.pageChunker(ctx, pageMetadata, processingState, reporter, checkpointer)
|
||
|
||
pageNumber++
|
||
}
|
||
return objectCount
|
||
}
|
||
|
||
// listErrorsAreExpected reports whether a failure to list a bucket's objects
|
||
// should be suppressed rather than logged as an error. When a role is assumed
|
||
// without an explicit bucket list, the scanner attempts every bucket in the
|
||
// account and is expected to be denied on some of them. When buckets are
|
||
// explicitly configured, a listing failure means a configured target is being
|
||
// silently skipped, so it is always an error, even under an assumed role.
|
||
func (s *Source) listErrorsAreExpected(role string) bool {
|
||
return role != "" && len(s.conn.GetBuckets()) == 0
|
||
}
|
||
|
||
// Chunks emits chunks of bytes over a channel.
|
||
func (s *Source) Chunks(ctx context.Context, chunksChan chan *sources.Chunk, _ ...sources.ChunkingTarget) error {
|
||
var totalObjectCount uint64
|
||
visitor := func(c context.Context, defaultRegionClient *s3.Client, roleArn string, buckets []string) error {
|
||
s.scanBuckets(c, defaultRegionClient, roleArn, buckets, chunksChan, &totalObjectCount)
|
||
return nil
|
||
}
|
||
|
||
return s.visitRoles(ctx, visitor)
|
||
}
|
||
|
||
func (s *Source) getRegionalClientForBucket(
|
||
ctx context.Context,
|
||
defaultRegionClient *s3.Client,
|
||
role string,
|
||
bucket string,
|
||
) (*s3.Client, error) {
|
||
region, err := s3manager.GetBucketRegion(ctx, defaultRegionClient, bucket)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("could not get s3 region for bucket: %s: %w", bucket, err)
|
||
}
|
||
|
||
if region == defaultAWSRegion {
|
||
return defaultRegionClient, nil
|
||
}
|
||
|
||
regionalClient, err := s.newClient(ctx, region, role)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("could not create regional s3 client for bucket %s: %w", bucket, err)
|
||
}
|
||
|
||
return regionalClient, nil
|
||
}
|
||
|
||
// pageChunker emits chunks onto the given channel from a page.
|
||
func (s *Source) pageChunker(
|
||
ctx context.Context,
|
||
metadata pageMetadata,
|
||
state processingState,
|
||
reporter sources.ChunkReporter,
|
||
checkpointer *Checkpointer,
|
||
) {
|
||
checkpointer.Reset() // Reset the checkpointer for each PAGE
|
||
ctx = context.WithValues(ctx, "bucket", metadata.bucket, "page_number", metadata.pageNumber)
|
||
for objIdx, obj := range metadata.page.Contents {
|
||
ctx = context.WithValues(ctx, "key", *obj.Key, "size", *obj.Size)
|
||
if common.IsDone(ctx) {
|
||
return
|
||
}
|
||
|
||
// Skip GLACIER and GLACIER_IR objects.
|
||
if obj.StorageClass == s3types.ObjectStorageClassGlacier || obj.StorageClass == s3types.ObjectStorageClassGlacierIr {
|
||
ctx.Logger().V(5).Info("Skipping object in storage class", "storage_class", obj.StorageClass)
|
||
s.metricsCollector.RecordObjectSkipped(metadata.bucket, "storage_class", float64(*obj.Size))
|
||
if err := checkpointer.UpdateObjectCompletion(ctx, objIdx, metadata.bucket, metadata.role, metadata.page.Contents); err != nil {
|
||
ctx.Logger().Error(err, "could not update progress for glacier object")
|
||
}
|
||
continue
|
||
}
|
||
|
||
// Ignore large files.
|
||
if *obj.Size > s.maxObjectSize {
|
||
ctx.Logger().V(5).Info("Skipping large file", "max_object_size", s.maxObjectSize)
|
||
s.metricsCollector.RecordObjectSkipped(metadata.bucket, "size_limit", float64(*obj.Size))
|
||
if err := checkpointer.UpdateObjectCompletion(ctx, objIdx, metadata.bucket, metadata.role, metadata.page.Contents); err != nil {
|
||
ctx.Logger().Error(err, "could not update progress for large file")
|
||
}
|
||
continue
|
||
}
|
||
|
||
// File empty file.
|
||
if *obj.Size == 0 {
|
||
ctx.Logger().V(5).Info("Skipping empty file")
|
||
s.metricsCollector.RecordObjectSkipped(metadata.bucket, "empty_file", 0)
|
||
if err := checkpointer.UpdateObjectCompletion(ctx, objIdx, metadata.bucket, metadata.role, metadata.page.Contents); err != nil {
|
||
ctx.Logger().Error(err, "could not update progress for empty file")
|
||
}
|
||
continue
|
||
}
|
||
|
||
// Skip incompatible extensions.
|
||
if common.SkipFile(*obj.Key) {
|
||
ctx.Logger().V(5).Info("Skipping file with incompatible extension")
|
||
s.metricsCollector.RecordObjectSkipped(metadata.bucket, "incompatible_extension", float64(*obj.Size))
|
||
if err := checkpointer.UpdateObjectCompletion(ctx, objIdx, metadata.bucket, metadata.role, metadata.page.Contents); err != nil {
|
||
ctx.Logger().Error(err, "could not update progress for incompatible file")
|
||
}
|
||
continue
|
||
}
|
||
|
||
s.jobPool.Go(func() error {
|
||
defer common.RecoverWithExit(ctx)
|
||
if common.IsDone(ctx) {
|
||
return ctx.Err()
|
||
}
|
||
|
||
if strings.HasSuffix(*obj.Key, "/") {
|
||
ctx.Logger().V(5).Info("Skipping directory")
|
||
s.metricsCollector.RecordObjectSkipped(metadata.bucket, "directory", float64(*obj.Size))
|
||
return nil
|
||
}
|
||
|
||
path := strings.Split(*obj.Key, "/")
|
||
prefix := strings.Join(path[:len(path)-1], "/")
|
||
|
||
nErr, ok := state.errorCount.Load(prefix)
|
||
if !ok {
|
||
nErr = 0
|
||
}
|
||
if nErr.(int) > 3 {
|
||
ctx.Logger().V(2).Info("Skipped due to excessive errors")
|
||
return nil
|
||
}
|
||
// Make sure we use a separate context for the GetObjectWithContext call.
|
||
// This ensures that the timeout is isolated and does not affect any downstream operations. (e.g. HandleFile)
|
||
const getObjectTimeout = 30 * time.Second
|
||
objCtx, cancel := context.WithTimeout(ctx, getObjectTimeout)
|
||
defer cancel()
|
||
|
||
res, err := metadata.client.GetObject(objCtx, &s3.GetObjectInput{
|
||
Bucket: &metadata.bucket,
|
||
Key: obj.Key,
|
||
})
|
||
if err != nil {
|
||
if strings.Contains(err.Error(), "AccessDenied") {
|
||
ctx.Logger().Error(err, "could not get S3 object; access denied")
|
||
s.metricsCollector.RecordObjectSkipped(metadata.bucket, "access_denied", float64(*obj.Size))
|
||
} else {
|
||
ctx.Logger().Error(err, "could not get S3 object")
|
||
s.metricsCollector.RecordObjectError(metadata.bucket)
|
||
}
|
||
// According to the documentation for GetObjectWithContext,
|
||
// the response can be non-nil even if there was an error.
|
||
// It's uncertain if the body will be nil in such cases,
|
||
// but we'll close it if it's not.
|
||
if res != nil && res.Body != nil {
|
||
_ = res.Body.Close()
|
||
}
|
||
|
||
nErr, ok := state.errorCount.Load(prefix)
|
||
if !ok {
|
||
nErr = 0
|
||
}
|
||
if nErr.(int) > 3 {
|
||
ctx.Logger().V(3).Info("Skipped due to excessive errors")
|
||
return nil
|
||
}
|
||
nErr = nErr.(int) + 1
|
||
state.errorCount.Store(prefix, nErr)
|
||
// too many consecutive errors on this page
|
||
if nErr.(int) > 3 {
|
||
ctx.Logger().V(2).Info("Too many consecutive errors, excluding prefix", "prefix", prefix)
|
||
}
|
||
return nil
|
||
}
|
||
defer func() { _ = res.Body.Close() }()
|
||
|
||
email := "Unknown"
|
||
if obj.Owner != nil {
|
||
email = *obj.Owner.DisplayName
|
||
}
|
||
modified := obj.LastModified.String()
|
||
chunkSkel := &sources.Chunk{
|
||
SourceType: s.Type(),
|
||
SourceName: s.name,
|
||
SourceID: s.SourceID(),
|
||
JobID: s.JobID(),
|
||
SourceMetadata: &source_metadatapb.MetaData{
|
||
Data: &source_metadatapb.MetaData_S3{
|
||
S3: &source_metadatapb.S3{
|
||
Bucket: metadata.bucket,
|
||
File: sanitizer.UTF8(*obj.Key),
|
||
Link: sanitizer.UTF8(makeS3Link(metadata.bucket, metadata.client.Options().Region, *obj.Key)),
|
||
Email: sanitizer.UTF8(email),
|
||
Timestamp: sanitizer.UTF8(modified),
|
||
},
|
||
},
|
||
},
|
||
SourceVerify: s.verify,
|
||
}
|
||
|
||
if err := handlers.HandleFile(ctx, res.Body, chunkSkel, reporter); err != nil {
|
||
ctx.Logger().Error(err, "error handling file")
|
||
s.metricsCollector.RecordObjectError(metadata.bucket)
|
||
return nil
|
||
}
|
||
atomic.AddUint64(state.objectCount, 1)
|
||
ctx.Logger().V(5).Info("S3 object scanned.", "object_count", state.objectCount)
|
||
nErr, ok = state.errorCount.Load(prefix)
|
||
if !ok {
|
||
nErr = 0
|
||
}
|
||
if nErr.(int) > 0 {
|
||
state.errorCount.Store(prefix, 0)
|
||
}
|
||
// Update progress after successful processing.
|
||
if err := checkpointer.UpdateObjectCompletion(ctx, objIdx, metadata.bucket, metadata.role, metadata.page.Contents); err != nil {
|
||
ctx.Logger().Error(err, "could not update progress for scanned object")
|
||
}
|
||
s.metricsCollector.RecordObjectScanned(metadata.bucket, float64(*obj.Size))
|
||
return nil
|
||
})
|
||
}
|
||
_ = s.jobPool.Wait()
|
||
}
|
||
|
||
func (s *Source) validateBucketAccess(ctx context.Context, client *s3.Client, roleArn string, buckets []string) []error {
|
||
shouldHaveAccessToAllBuckets := roleArn == ""
|
||
wasAbleToListAnyBucket := false
|
||
var errs []error
|
||
|
||
for _, bucket := range buckets {
|
||
if common.IsDone(ctx) {
|
||
return append(errs, ctx.Err())
|
||
}
|
||
|
||
regionalClient, err := s.getRegionalClientForBucket(ctx, client, roleArn, bucket)
|
||
if err != nil {
|
||
errs = append(errs, fmt.Errorf("could not get regional client for bucket %q: %w", bucket, err))
|
||
continue
|
||
}
|
||
|
||
_, err = regionalClient.ListObjectsV2(ctx, &s3.ListObjectsV2Input{Bucket: &bucket})
|
||
if err == nil {
|
||
wasAbleToListAnyBucket = true
|
||
} else if shouldHaveAccessToAllBuckets {
|
||
errs = append(errs, fmt.Errorf("could not list objects in bucket %q: %w", bucket, err))
|
||
}
|
||
}
|
||
|
||
if !wasAbleToListAnyBucket {
|
||
if roleArn == "" {
|
||
errs = append(errs, errors.New("could not list objects in any bucket"))
|
||
} else {
|
||
errs = append(errs, fmt.Errorf("role %q could not list objects in any bucket", roleArn))
|
||
}
|
||
}
|
||
|
||
return errs
|
||
}
|
||
|
||
// visitRoles iterates over the configured AWS roles and calls the provided function
|
||
// for each role, passing in the default S3 client, the role ARN, and the list of
|
||
// buckets to scan.
|
||
//
|
||
// The provided function parameter typically implements the core scanning logic
|
||
// and must handle context cancellation appropriately.
|
||
//
|
||
// If no roles are configured, it will call the function with an empty role ARN.
|
||
func (s *Source) visitRoles(
|
||
ctx context.Context,
|
||
f func(c context.Context, defaultRegionClient *s3.Client, roleArn string, buckets []string) error,
|
||
) error {
|
||
roles := s.conn.GetRoles()
|
||
if len(roles) == 0 {
|
||
roles = []string{""}
|
||
}
|
||
|
||
for _, role := range roles {
|
||
s.metricsCollector.RecordRoleScanned(role)
|
||
|
||
client, err := s.newClient(ctx, defaultAWSRegion, role)
|
||
if err != nil {
|
||
return fmt.Errorf("could not create s3 client: %w", err)
|
||
}
|
||
|
||
bucketsToScan, err := s.getBucketsToScan(ctx, client)
|
||
if err != nil {
|
||
return fmt.Errorf("role %q could not list any s3 buckets for scanning: %w", role, err)
|
||
}
|
||
|
||
if err := f(ctx, client, role, bucketsToScan); err != nil {
|
||
return err
|
||
}
|
||
}
|
||
|
||
return nil
|
||
}
|
||
|
||
// makeS3Link creates a S3 virtual-hosted–style URIs. They have the format of:
|
||
// https://[bucket-name].s3.[region-code].amazonaws.com/[key-name]
|
||
//
|
||
// See https://docs.aws.amazon.com/AmazonS3/latest/userguide/VirtualHosting.html#virtual-hosted-style-access
|
||
func makeS3Link(bucket, region, key string) string {
|
||
return fmt.Sprintf("https://%s.s3.%s.amazonaws.com/%s", bucket, region, key)
|
||
}
|
||
|
||
// Enumerate implements SourceUnitEnumerator interface. This implementation visits
|
||
// each configured role and passes each s3 bucket as a source unit
|
||
func (s *Source) Enumerate(ctx context.Context, reporter sources.UnitReporter) error {
|
||
visitor := func(c context.Context, defaultRegionClient *s3.Client, roleArn string, buckets []string) error {
|
||
for _, bucket := range buckets {
|
||
if common.IsDone(ctx) {
|
||
return ctx.Err()
|
||
}
|
||
|
||
unit := S3SourceUnit{
|
||
Bucket: bucket,
|
||
Role: roleArn,
|
||
}
|
||
|
||
if err := reporter.UnitOk(ctx, unit); err != nil {
|
||
return err
|
||
}
|
||
}
|
||
return nil
|
||
}
|
||
|
||
return s.visitRoles(ctx, visitor)
|
||
}
|
||
|
||
// ChunkUnit implements SourceUnitChunker interface. This implementation scans
|
||
// the given S3 bucket source unit and emits chunks for each object found.
|
||
// It supports sub-unit resumption by utilizing the checkpointer to track progress.
|
||
func (s *Source) ChunkUnit(ctx context.Context, unit sources.SourceUnit, reporter sources.ChunkReporter) error {
|
||
|
||
s3unit, ok := unit.(S3SourceUnit)
|
||
if !ok {
|
||
return fmt.Errorf("expected *S3SourceUnit, got %T", unit)
|
||
}
|
||
// unitID is a combination of bucket name and role ARN
|
||
unitID, _ := s3unit.SourceUnitID()
|
||
defaultClient, err := s.newClient(ctx, defaultAWSRegion, s3unit.Role)
|
||
if err != nil {
|
||
return fmt.Errorf("could not create s3 client for bucket %s and role %s: %w", s3unit.Bucket, s3unit.Role, err)
|
||
}
|
||
|
||
checkpointer := NewCheckpointer(ctx, &s.Progress, true)
|
||
|
||
var startAfterPtr *string
|
||
startAfter := s.GetEncodedResumeInfoFor(unitID)
|
||
if startAfter != "" {
|
||
ctx.Logger().V(3).Info(
|
||
"Resuming unit scan",
|
||
"start_after", startAfter,
|
||
"unitID", unitID,
|
||
)
|
||
startAfterPtr = &startAfter
|
||
}
|
||
defer s.ClearEncodedResumeInfoFor(unitID)
|
||
s.scanBucket(ctx, defaultClient, s3unit.Role, s3unit.Bucket, reporter, startAfterPtr, checkpointer)
|
||
return nil
|
||
}
|
||
|
||
func (s *Source) UnmarshalSourceUnit(data []byte) (sources.SourceUnit, error) {
|
||
var unit S3SourceUnit
|
||
if err := json.Unmarshal(data, &unit); err != nil {
|
||
return nil, err
|
||
}
|
||
unitID, kind := unit.SourceUnitID()
|
||
if unitID == "" || kind != SourceUnitKindBucket {
|
||
return nil, fmt.Errorf("not an S3SourceUnit")
|
||
}
|
||
return unit, nil
|
||
}
|