Dedupe results (#1479)

* init 4 dedupin

* use raw rather than rawv2

* rm comment

* comments

* nits

* clean up and use rawv2 too

* add decoder order test
This commit is contained in:
Zachary Rice
2023-07-11 15:48:00 -05:00
committed by GitHub
parent a123d5c5e1
commit b48ac24c46
3 changed files with 31 additions and 10 deletions
+1
View File
@@ -6,6 +6,7 @@ import (
func DefaultDecoders() []Decoder {
return []Decoder{
// UTF8 must be first for duplicate detection
&UTF8{},
&Base64{},
&UTF16{},
+20 -10
View File
@@ -2,6 +2,7 @@ package engine
import (
"bytes"
"fmt"
"reflect"
"runtime"
"strings"
@@ -233,6 +234,22 @@ func (e *Engine) BytesScanned() uint64 {
return e.bytesScanned
}
func (e *Engine) dedupeAndSend(chunkResults []detectors.ResultWithMetadata) {
dedupeMap := make(map[string]struct{})
for _, result := range chunkResults {
// dedupe by comparing the detector type, raw result, and source metadata
// NOTE: in order for the PLAIN decoder to maintain precedence, make sure UTF8 is the first decoder in the
// default decoders list
key := fmt.Sprintf("%s%s%s%+v", result.DetectorType.String(), result.Raw, result.RawV2, result.SourceMetadata)
if _, ok := dedupeMap[key]; ok {
continue
}
dedupeMap[key] = struct{}{}
e.results <- result
}
}
func (e *Engine) DetectorAvgTime() map[string][]time.Duration {
logger := context.Background().Logger()
avgTime := map[string][]time.Duration{}
@@ -257,6 +274,7 @@ func (e *Engine) DetectorAvgTime() map[string][]time.Duration {
func (e *Engine) detectorWorker(ctx context.Context) {
for originalChunk := range e.chunks {
for chunk := range sources.Chunker(originalChunk) {
var chunkResults []detectors.ResultWithMetadata
matchedKeywords := make(map[string]struct{})
atomic.AddUint64(&e.bytesScanned, uint64(len(chunk.Data)))
for _, decoder := range e.decoders {
@@ -273,21 +291,12 @@ func (e *Engine) detectorWorker(ctx context.Context) {
decoderType = detectorspb.DecoderType_UNKNOWN
}
original := chunk.Data
decoded := decoder.FromChunk(chunk)
if decoded == nil {
continue
}
if decoded == nil ||
// check if the decoded data is similar "enough" to the original data. If it is, then we can skip scanning the decoded data as
// it's likely already picked up by the PLAIN decoder. See related issue: https://github.com/trufflesecurity/trufflehog/issues/1450
(decoded != nil &&
decoderType == detectorspb.DecoderType_BASE64 && common.BytesEqual(original, decoded.Data, 40)) {
continue
}
// build a map of all keywords that were matched in the chunk
for _, m := range e.prefilter.FindAll(string(decoded.Data)) {
matchedKeywords[strings.ToLower(string(decoded.Data[m.Start():m.End()]))] = struct{}{}
@@ -343,7 +352,7 @@ func (e *Engine) detectorWorker(ctx context.Context) {
continue
}
result.DecoderType = decoderType
e.results <- detectors.CopyMetadata(resultChunk, result)
chunkResults = append(chunkResults, detectors.CopyMetadata(resultChunk, result))
}
if len(results) > 0 {
@@ -363,6 +372,7 @@ func (e *Engine) detectorWorker(ctx context.Context) {
}
}
}
e.dedupeAndSend(chunkResults)
}
atomic.AddUint64(&e.chunksScanned, 1)
}
+10
View File
@@ -3,6 +3,7 @@ package engine
import (
"testing"
"github.com/trufflesecurity/trufflehog/v3/pkg/decoders"
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors"
"github.com/trufflesecurity/trufflehog/v3/pkg/sources"
)
@@ -62,3 +63,12 @@ func TestFragmentLineOffset(t *testing.T) {
})
}
}
// Test to make sure that DefaultDecoders always returns the UTF8 decoder first.
// Technically a decoder test but we want this to run and fail in CI
func TestDefaultDecoders(t *testing.T) {
ds := decoders.DefaultDecoders()
if _, ok := ds[0].(*decoders.UTF8); !ok {
t.Errorf("DefaultDecoders() = %v, expected UTF8 decoder to be first", ds)
}
}