From f5370819b555a68f3eb60626ee3fd2fb8923b06f Mon Sep 17 00:00:00 2001 From: Brad Larsen Date: Thu, 3 Sep 2026 13:44:35 -0400 Subject: [PATCH] Improve the `json-enumerator` source (#5265) * implement SourceUnitEnumerator and SourceUnitChunker * natively support zstd-compressed inputs * allow it to be used in multi-source config Previously, multiple files would be processed sequentially, and zstd-compressed files would have to be decompressed externally. These changes allow for slightly more efficient scanning of zstd-compressed NDJSON inputs (less CPU required), and enable greater CPU utilization on multicore systems when scanning multiple NDJSON files. With these changes, I see a 1.18x speedup in wall clock time in a simple experiment of scanning 10 zstd-compressed NDJSON files (~9GiB on disk) with verification disabled. --- go.mod | 2 +- pkg/config/config.go | 3 ++ .../json_enumerator/json_enumerator.go | 51 ++++++++++++++++--- .../json_enumerator/json_enumerator_test.go | 4 +- 4 files changed, 51 insertions(+), 9 deletions(-) diff --git a/go.mod b/go.mod index fee1845ea..5b6ba4755 100644 --- a/go.mod +++ b/go.mod @@ -64,6 +64,7 @@ require ( github.com/joho/godotenv v1.5.1 github.com/jpillora/overseer v1.1.6 github.com/k0kubun/go-ansi v0.0.0-20180517002512-3bf9e2903213 + github.com/klauspost/compress v1.18.6 github.com/klauspost/pgzip v1.2.6 github.com/kylelemons/godebug v1.1.0 github.com/lestrrat-go/jwx/v3 v3.0.12 @@ -232,7 +233,6 @@ require ( github.com/jpillora/s3 v1.1.4 // indirect github.com/kevinburke/ssh_config v1.2.0 // indirect github.com/kjk/lzma v0.0.0-20161016003348-3fd93898850d // indirect - github.com/klauspost/compress v1.18.6 // indirect github.com/klauspost/cpuid/v2 v2.3.0 // indirect github.com/lestrrat-go/blackmagic v1.0.4 // indirect github.com/lestrrat-go/httpcc v1.0.1 // indirect diff --git a/pkg/config/config.go b/pkg/config/config.go index 47fe3868c..ef6b2d5a8 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -17,6 +17,7 @@ import ( "github.com/trufflesecurity/trufflehog/v3/pkg/sources/github" "github.com/trufflesecurity/trufflehog/v3/pkg/sources/gitlab" "github.com/trufflesecurity/trufflehog/v3/pkg/sources/jenkins" + "github.com/trufflesecurity/trufflehog/v3/pkg/sources/json_enumerator" "github.com/trufflesecurity/trufflehog/v3/pkg/sources/postman" "github.com/trufflesecurity/trufflehog/v3/pkg/sources/s3" ) @@ -97,6 +98,8 @@ func instantiateSourceFromType(sourceType string) (sources.Source, error) { source = new(filesystem.Source) case sourcespb.SourceType_SOURCE_TYPE_JENKINS.String(): source = new(jenkins.Source) + case sourcespb.SourceType_SOURCE_TYPE_JSON_ENUMERATOR.String(): + source = new(json_enumerator.Source) case sourcespb.SourceType_SOURCE_TYPE_GCS.String(): source = new(gcs.Source) case sourcespb.SourceType_SOURCE_TYPE_GCS_UNAUTHED.String(): diff --git a/pkg/sources/json_enumerator/json_enumerator.go b/pkg/sources/json_enumerator/json_enumerator.go index 991d9fb52..54e98dd51 100644 --- a/pkg/sources/json_enumerator/json_enumerator.go +++ b/pkg/sources/json_enumerator/json_enumerator.go @@ -7,8 +7,10 @@ import ( "fmt" "io" "os" + "strings" "unicode/utf8" + "github.com/klauspost/compress/zstd" "google.golang.org/protobuf/proto" "google.golang.org/protobuf/types/known/anypb" @@ -35,6 +37,7 @@ type Source struct { // Ensure the Source satisfies the interfaces at compile time var _ sources.Source = (*Source)(nil) var _ sources.SourceUnitUnmarshaller = (*Source)(nil) +var _ sources.SourceUnitEnumChunker = (*Source)(nil) func (s *Source) Type() sourcespb.SourceType { return SourceType } func (s *Source) SourceID() sources.SourceID { return s.sourceId } @@ -73,7 +76,7 @@ func (s *Source) Chunks( return nil } s.SetProgressComplete(i, len(s.paths), fmt.Sprintf("Path: %s", path), "") - if err := s.chunkJSONEnumerator(ctx, path, chunksChan); err != nil { + if err := s.chunkJSONEnumerator(ctx, path, sources.ChanReporter{Ch: chunksChan}); err != nil { ctx.Logger().Error(err, "error scanning JSON enumerator", "path", path) } } @@ -81,6 +84,34 @@ func (s *Source) Chunks( return nil } +// Enumerate implements the SourceUnitEnumerator interface, reporting each +// configured path as its own unit so paths can be chunked concurrently. +func (s *Source) Enumerate(ctx context.Context, reporter sources.UnitReporter) error { + for _, path := range s.paths { + f, err := os.Open(path) + if err != nil { + if err := reporter.UnitErr(ctx, err); err != nil { + return err + } + } else { + _ = f.Close() + if err := reporter.UnitOk(ctx, sources.CommonSourceUnit{ID: path}); err != nil { + return err + } + } + } + return nil +} + +// ChunkUnit implements the SourceUnitChunker interface. +func (s *Source) ChunkUnit(ctx context.Context, unit sources.SourceUnit, reporter sources.ChunkReporter) error { + path, _ := unit.SourceUnitID() + if err := s.chunkJSONEnumerator(ctx, path, reporter); err != nil { + return reporter.ChunkErr(ctx, err) + } + return nil +} + type jsonEntry struct { Metadata json.RawMessage Data []byte @@ -142,13 +173,11 @@ func (e *jsonEntry) UnmarshalJSON(data []byte) error { func (s *Source) chunkJSONEnumeratorReader( ctx context.Context, input io.Reader, - chunksChan chan *sources.Chunk, + reporter sources.ChunkReporter, ) error { decoder := json.NewDecoder(input) var entry jsonEntry - reporter := sources.ChanReporter{Ch: chunksChan} - for { if err := decoder.Decode(&entry); err != nil { if errors.Is(err, io.EOF) { @@ -190,7 +219,7 @@ func (s *Source) chunkJSONEnumeratorReader( func (s *Source) chunkJSONEnumerator( ctx context.Context, path string, - chunksChan chan *sources.Chunk, + reporter sources.ChunkReporter, ) error { ctx.Logger().V(3).Info("chunking JSON enumerator", "path", path) @@ -200,5 +229,15 @@ func (s *Source) chunkJSONEnumerator( } defer func() { _ = enumeratorFile.Close() }() - return s.chunkJSONEnumeratorReader(ctx, enumeratorFile, chunksChan) + if strings.HasSuffix(path, ".zst") || strings.HasSuffix(path, ".zstd") { + r, err := zstd.NewReader(enumeratorFile) + if err != nil { + return fmt.Errorf("unable to open zstd reader: %w", err) + } + defer r.Close() + return s.chunkJSONEnumeratorReader(ctx, r, reporter) + } else { + return s.chunkJSONEnumeratorReader(ctx, enumeratorFile, reporter) + } + } diff --git a/pkg/sources/json_enumerator/json_enumerator_test.go b/pkg/sources/json_enumerator/json_enumerator_test.go index fd215adf2..11d8f6ad9 100644 --- a/pkg/sources/json_enumerator/json_enumerator_test.go +++ b/pkg/sources/json_enumerator/json_enumerator_test.go @@ -105,7 +105,7 @@ func TestScanEnumerator(t *testing.T) { defer close(chunksChan) ctx := context.WithLogger(t.Context(), logr.Discard()) source := &Source{} - workerError = source.chunkJSONEnumeratorReader(ctx, readJSON, chunksChan) + workerError = source.chunkJSONEnumeratorReader(ctx, readJSON, sources.ChanReporter{Ch: chunksChan}) }() enc := json.NewEncoder(writeJSON) @@ -190,7 +190,7 @@ func TestScanEnumeratorAPKContentRouting(t *testing.T) { defer close(chunksChan) ctx := context.WithLogger(t.Context(), logr.Discard()) source := &Source{} - workerError = source.chunkJSONEnumeratorReader(ctx, readJSON, chunksChan) + workerError = source.chunkJSONEnumeratorReader(ctx, readJSON, sources.ChanReporter{Ch: chunksChan}) }() // No filename is set on the record; routing is content-based only.