trufflesecurity/trufflehog
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Improve the json-enumerator source (#5265)

Browse Source 
* 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.
This commit is contained in:
Brad Larsen
2026-09-03 13:44:35 -04:00
committed by GitHub
parent 8a9a4c6be8
commit f5370819b5
4 changed files with 51 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
go.mod
+1 -1
View File
@@ -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
pkg/config/config.go
+3
View File
@@ -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():
pkg/sources/json_enumerator/json_enumerator.go
+45 -6
View File
@@ -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)
}
}
pkg/sources/json_enumerator/json_enumerator_test.go
+2 -2
View File
@@ -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.