trufflehog/pkg/engine/ahocorasickcore_test.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)
}