trufflehog/pkg/engine/ahocorasickcore_test.go
Cody Rose 7a156330b5
Support multiple detectors per match (#2065)
#1711 inadvertently removed the ability to match multiple custom detectors, or multiple detectors of the same type but different version, to a given keyword. (#2060 re-added support for multiple versions of detectors globally, and #2064 re-added support for multiple custom detectors globally, but neither fixed trufflehog's inability to support multiple such detectors for a given keyword match.) This PR re-adds the removed functionality (and narrows the AhoCorasickCore interface in the process.)
2023-11-03 12:26:18 -04:00

116 lines
3.3 KiB
Go

package engine
import (
"context"
"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"
)
const TestDetectorType = -1
type testDetectorV1 struct {
}
func (d testDetectorV1) FromData(ctx context.Context, verify bool, data []byte) ([]detectors.Result, error) {
return make([]detectors.Result, 0), nil
}
func (d testDetectorV1) Keywords() []string {
return []string{"a", "b"}
}
func (d testDetectorV1) Type() detectorspb.DetectorType {
return TestDetectorType
}
func (d testDetectorV1) Version() int {
return 1
}
type testDetectorV2 struct {
}
func (d testDetectorV2) FromData(ctx context.Context, verify bool, data []byte) ([]detectors.Result, error) {
return make([]detectors.Result, 0), nil
}
func (d testDetectorV2) Keywords() []string {
return []string{"a"}
}
func (d testDetectorV2) Type() detectorspb.DetectorType {
return TestDetectorType
}
func (d testDetectorV2) Version() int {
return 2
}
var _ detectors.Detector = (*testDetectorV1)(nil)
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{"a"},
Regex: map[string]string{"": ""},
})
assert.Nil(t, err)
allDetectors := []detectors.Detector{customDetector1, customDetector2}
ac := NewAhoCorasickCore(allDetectors)
detectorsMap := make(map[DetectorKey]detectors.Detector, 2)
ac.PopulateMatchingDetectors("a", detectorsMap)
matchingDetectors := make([]detectors.Detector, 0, 2)
for _, d := range detectorsMap {
matchingDetectors = append(matchingDetectors, d)
}
assert.ElementsMatch(t, allDetectors, matchingDetectors)
}
func TestAhoCorasickCore_MultipleDetectorVersionsMatchable(t *testing.T) {
v1 := testDetectorV1{}
v2 := testDetectorV2{}
allDetectors := []detectors.Detector{v1, v2}
ac := NewAhoCorasickCore(allDetectors)
detectorsMap := make(map[DetectorKey]detectors.Detector, 2)
ac.PopulateMatchingDetectors("a", detectorsMap)
matchingDetectors := make([]detectors.Detector, 0, 2)
for _, d := range detectorsMap {
matchingDetectors = append(matchingDetectors, d)
}
assert.ElementsMatch(t, allDetectors, matchingDetectors)
}
func TestAhoCorasickCore_NoDuplicateDetectorsMatched(t *testing.T) {
d := testDetectorV1{}
allDetectors := []detectors.Detector{d}
ac := NewAhoCorasickCore(allDetectors)
detectorsMap := make(map[DetectorKey]detectors.Detector, 2)
ac.PopulateMatchingDetectors("a a b b", detectorsMap)
matchingDetectors := make([]detectors.Detector, 0, 2)
for _, d := range detectorsMap {
matchingDetectors = append(matchingDetectors, d)
}
assert.ElementsMatch(t, allDetectors, matchingDetectors)
}