Support multiple custom detectors (#2064)

#1711 accidentally removed the ability to support multiple custom detectors. This PR partially adds back this capability: Multiple custom detectors are now supported overall, but only one custom detector can be returned for a given keyword match.

pkg/custom_detectors/custom_detectors.go

@@ -20,20 +20,20 @@ import (
 // for poorly defined regexps.
 const maxTotalMatches = 100
 
-// customRegexWebhook is a CustomRegex with webhook validation that is
+// CustomRegexWebhook is a CustomRegex with webhook validation that is
 // guaranteed to be valid (assuming the data is not changed after
 // initialization).
-type customRegexWebhook struct {
+type CustomRegexWebhook struct {
 	*custom_detectorspb.CustomRegex
 }
 
 // Ensure the Scanner satisfies the interface at compile time.
-var _ detectors.Detector = (*customRegexWebhook)(nil)
+var _ detectors.Detector = (*CustomRegexWebhook)(nil)
 
-// NewWebhookCustomRegex initializes and validates a customRegexWebhook. An
+// NewWebhookCustomRegex initializes and validates a CustomRegexWebhook. An
 // unexported type is intentionally returned here to ensure the values have
 // been validated.
-func NewWebhookCustomRegex(pb *custom_detectorspb.CustomRegex) (*customRegexWebhook, error) {
+func NewWebhookCustomRegex(pb *custom_detectorspb.CustomRegex) (*CustomRegexWebhook, error) {
 	// TODO: Return all validation errors.
 	if err := ValidateKeywords(pb.Keywords); err != nil {
 		return nil, err
@@ -52,12 +52,12 @@ func NewWebhookCustomRegex(pb *custom_detectorspb.CustomRegex) (*customRegexWebh
 	}
 
 	// TODO: Copy only necessary data out of pb.
-	return &customRegexWebhook{pb}, nil
+	return &CustomRegexWebhook{pb}, nil
 }
 
 var httpClient = common.SaneHttpClient()
 
-func (c *customRegexWebhook) FromData(ctx context.Context, verify bool, data []byte) (results []detectors.Result, err error) {
+func (c *CustomRegexWebhook) FromData(ctx context.Context, verify bool, data []byte) (results []detectors.Result, err error) {
 	dataStr := string(data)
 	regexMatches := make(map[string][][]string, len(c.GetRegex()))
 
@@ -109,7 +109,7 @@ func (c *customRegexWebhook) FromData(ctx context.Context, verify bool, data []b
 	return results, nil
 }
 
-func (c *customRegexWebhook) createResults(ctx context.Context, match map[string][]string, verify bool, results chan<- detectors.Result) error {
+func (c *CustomRegexWebhook) createResults(ctx context.Context, match map[string][]string, verify bool, results chan<- detectors.Result) error {
 	if common.IsDone(ctx) {
 		// TODO: Log we're possibly leaving out results.
 		return ctx.Err()
@@ -180,7 +180,7 @@ func (c *customRegexWebhook) createResults(ctx context.Context, match map[string
 	}
 }
 
-func (c *customRegexWebhook) Keywords() []string {
+func (c *CustomRegexWebhook) Keywords() []string {
 	return c.GetKeywords()
 }
 
@@ -245,6 +245,6 @@ func permutateMatches(regexMatches map[string][][]string) []map[string][]string
 	return matches
 }
 
-func (c *customRegexWebhook) Type() detectorspb.DetectorType {
+func (c *CustomRegexWebhook) Type() detectorspb.DetectorType {
 	return detectorspb.DetectorType_CustomRegex
 }

pkg/custom_detectors/custom_detectors_test.go

@@ -105,7 +105,7 @@ func TestCustomDetectorsParsing(t *testing.T) {
 }
 
 func TestFromData_InvalidRegEx(t *testing.T) {
-	c := &customRegexWebhook{
+	c := &CustomRegexWebhook{
 		&custom_detectorspb.CustomRegex{
 			Name:     "Internal bi tool",
 			Keywords: []string{"secret_v1_", "pat_v2_"},

pkg/engine/ahocorasickcore.go

@@ -4,17 +4,20 @@ import (
 	"strings"
 
 	ahocorasick "github.com/BobuSumisu/aho-corasick"
-
+	"github.com/trufflesecurity/trufflehog/v3/pkg/custom_detectors"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/detectors"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/pb/detectorspb"
 )
 
 // detectorKey is used to identify a detector in the keywordsToDetectors map.
 // Multiple detectors can have the same detector type but different versions.
-// This allows us to identify a detector by its type and version.
+// This allows us to identify a detector by its type and version. An
+// additional (optional) field is provided to disambiguate multiple custom
+// detectors.
 type detectorKey struct {
-	detectorType detectorspb.DetectorType
-	version      int
+	detectorType       detectorspb.DetectorType
+	version            int
+	customDetectorName string
 }
 
 // AhoCorasickCore encapsulates the operations and data structures used for keyword matching via the
@@ -85,5 +88,9 @@ func createDetectorKey(d detectors.Detector) detectorKey {
 	if v, ok := d.(detectors.Versioner); ok {
 		version = v.Version()
 	}
-	return detectorKey{detectorType: detectorType, version: version}
+	var customDetectorName string
+	if r, ok := d.(*custom_detectors.CustomRegexWebhook); ok {
+		customDetectorName = r.GetName()
+	}
+	return detectorKey{detectorType: detectorType, version: version, customDetectorName: customDetectorName}
 }

pkg/engine/ahocorasickcore_test.go

@@ -5,7 +5,9 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/trufflesecurity/trufflehog/v3/pkg/custom_detectors"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/detectors"
+	"github.com/trufflesecurity/trufflehog/v3/pkg/pb/custom_detectorspb"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/pb/detectorspb"
 )
 
@@ -54,6 +56,53 @@ var _ detectors.Detector = (*testDetectorV2)(nil)
 var _ detectors.Versioner = (*testDetectorV1)(nil)
 var _ detectors.Versioner = (*testDetectorV2)(nil)
 
+func TestAhoCorasickCore_MultipleCustomDetectorsMatchable(t *testing.T) {
+	customDetector1, err := custom_detectors.NewWebhookCustomRegex(&custom_detectorspb.CustomRegex{
+		Name:     "custom detector 1",
+		Keywords: []string{"a"},
+		Regex:    map[string]string{"": ""},
+	})
+	assert.Nil(t, err)
+
+	customDetector2, err := custom_detectors.NewWebhookCustomRegex(&custom_detectorspb.CustomRegex{
+		Name:     "custom detector 2",
+		Keywords: []string{"b"},
+		Regex:    map[string]string{"": ""},
+	})
+	assert.Nil(t, err)
+
+	testCases := []struct {
+		matchString string
+		detector    detectors.Detector
+	}{
+		{
+			matchString: "a",
+			detector:    customDetector1,
+		},
+		{
+			matchString: "b",
+			detector:    customDetector2,
+		},
+	}
+
+	var allDetectors []detectors.Detector
+	for _, tt := range testCases {
+		allDetectors = append(allDetectors, tt.detector)
+	}
+
+	ac := NewAhoCorasickCore(allDetectors)
+
+	for _, tt := range testCases {
+		matches := ac.MatchString(tt.matchString)
+		assert.Equal(t, 1, len(matches))
+
+		matchingDetectors := make(map[detectorspb.DetectorType]detectors.Detector)
+		ac.PopulateDetectorsByMatch(matches[0], matchingDetectors)
+		assert.Equal(t, 1, len(matchingDetectors))
+		assert.Equal(t, tt.detector, matchingDetectors[detectorspb.DetectorType_CustomRegex])
+	}
+}
+
 func TestAhoCorasickCore_MultipleDetectorVersionsMatchable(t *testing.T) {
 	testCases := []struct {
 		matchString string