Compare commits
1 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| ddf25d359d |
@@ -164,7 +164,7 @@ trufflehog docker --image trufflesecurity/secrets --only-verified
|
||||
|
||||
+ All I see is `🐷🔑🐷 TruffleHog. Unearth your secrets. 🐷🔑🐷` and the program exits, what gives?
|
||||
+ That means no secrets were detected
|
||||
+ Why is the scan taking a long time when I scan a GitHub org
|
||||
+ Why is the scan is taking a long time when I scan a GitHub org
|
||||
+ Unauthenticated GitHub scans have rate limits. To improve your rate limits, include the `--token` flag with a personal access token
|
||||
+ It says a private key was verified, what does that mean?
|
||||
+ Check out our Driftwood blog post to learn how to do this, in short we've confirmed the key can be used live for SSH or SSL [Blog post](https://trufflesecurity.com/blog/driftwood-know-if-private-keys-are-sensitive/)
|
||||
@@ -219,7 +219,7 @@ Flags:
|
||||
--concurrency=10 Number of concurrent workers.
|
||||
--no-verification Don't verify the results.
|
||||
--only-verified Only output verified results.
|
||||
--filter-unverified Only output first unverified result per chunk per detector if there is more than one result.
|
||||
--filter-unverified Only output first unverified result per chunk per detector if there are more than one results.
|
||||
--config=CONFIG Path to configuration file.
|
||||
--print-avg-detector-time Print the average time spent on each detector.
|
||||
--no-update Don't check for updates.
|
||||
|
||||
@@ -47,7 +47,6 @@ require (
|
||||
github.com/h2non/filetype v1.1.3
|
||||
github.com/hashicorp/go-retryablehttp v0.7.4
|
||||
github.com/hashicorp/golang-lru v0.5.1
|
||||
github.com/hashicorp/golang-lru/v2 v2.0.7
|
||||
github.com/jlaffaye/ftp v0.2.0
|
||||
github.com/joho/godotenv v1.5.1
|
||||
github.com/jpillora/overseer v1.1.6
|
||||
|
||||
@@ -418,8 +418,6 @@ github.com/hashicorp/go-retryablehttp v0.7.4/go.mod h1:Jy/gPYAdjqffZ/yFGCFV2doI5
|
||||
github.com/hashicorp/golang-lru v0.5.0/go.mod h1:/m3WP610KZHVQ1SGc6re/UDhFvYD7pJ4Ao+sR/qLZy8=
|
||||
github.com/hashicorp/golang-lru v0.5.1 h1:0hERBMJE1eitiLkihrMvRVBYAkpHzc/J3QdDN+dAcgU=
|
||||
github.com/hashicorp/golang-lru v0.5.1/go.mod h1:/m3WP610KZHVQ1SGc6re/UDhFvYD7pJ4Ao+sR/qLZy8=
|
||||
github.com/hashicorp/golang-lru/v2 v2.0.7 h1:a+bsQ5rvGLjzHuww6tVxozPZFVghXaHOwFs4luLUK2k=
|
||||
github.com/hashicorp/golang-lru/v2 v2.0.7/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM=
|
||||
github.com/hpcloud/tail v1.0.0/go.mod h1:ab1qPbhIpdTxEkNHXyeSf5vhxWSCs/tWer42PpOxQnU=
|
||||
github.com/ianlancetaylor/demangle v0.0.0-20181102032728-5e5cf60278f6/go.mod h1:aSSvb/t6k1mPoxDqO4vJh6VOCGPwU4O0C2/Eqndh1Sc=
|
||||
github.com/ianlancetaylor/demangle v0.0.0-20210905161508-09a460cdf81d/go.mod h1:aYm2/VgdVmcIU8iMfdMvDMsRAQjcfZSKFby6HOFvi/w=
|
||||
|
||||
@@ -4,7 +4,6 @@ import (
|
||||
"bytes"
|
||||
"encoding/base64"
|
||||
|
||||
"github.com/trufflesecurity/trufflehog/v3/pkg/pb/detectorspb"
|
||||
"github.com/trufflesecurity/trufflehog/v3/pkg/sources"
|
||||
)
|
||||
|
||||
@@ -26,8 +25,7 @@ func init() {
|
||||
}
|
||||
}
|
||||
|
||||
func (d *Base64) FromChunk(chunk *sources.Chunk) *DecodableChunk {
|
||||
decodableChunk := &DecodableChunk{Chunk: chunk, DecoderType: detectorspb.DecoderType_BASE64}
|
||||
func (d *Base64) FromChunk(chunk *sources.Chunk) *sources.Chunk {
|
||||
encodedSubstrings := getSubstringsOfCharacterSet(chunk.Data, 20, b64CharsetMapping, b64EndChars)
|
||||
decodedSubstrings := make(map[string][]byte)
|
||||
|
||||
@@ -63,7 +61,7 @@ func (d *Base64) FromChunk(chunk *sources.Chunk) *DecodableChunk {
|
||||
}
|
||||
result.Write(chunk.Data[start:])
|
||||
chunk.Data = result.Bytes()
|
||||
return decodableChunk
|
||||
return chunk
|
||||
}
|
||||
|
||||
return nil
|
||||
|
||||
@@ -1,7 +1,6 @@
|
||||
package decoders
|
||||
|
||||
import (
|
||||
"github.com/trufflesecurity/trufflehog/v3/pkg/pb/detectorspb"
|
||||
"github.com/trufflesecurity/trufflehog/v3/pkg/sources"
|
||||
)
|
||||
|
||||
@@ -14,15 +13,8 @@ func DefaultDecoders() []Decoder {
|
||||
}
|
||||
}
|
||||
|
||||
// DecodableChunk is a chunk that includes the type of decoder used.
|
||||
// This allows us to avoid a type assertion on each decoder.
|
||||
type DecodableChunk struct {
|
||||
*sources.Chunk
|
||||
DecoderType detectorspb.DecoderType
|
||||
}
|
||||
|
||||
type Decoder interface {
|
||||
FromChunk(chunk *sources.Chunk) *DecodableChunk
|
||||
FromChunk(chunk *sources.Chunk) *sources.Chunk
|
||||
}
|
||||
|
||||
// Fuzz is an entrypoint for go-fuzz, which is an AFL-style fuzzing tool.
|
||||
|
||||
@@ -5,24 +5,22 @@ import (
|
||||
"encoding/binary"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/trufflesecurity/trufflehog/v3/pkg/pb/detectorspb"
|
||||
"github.com/trufflesecurity/trufflehog/v3/pkg/sources"
|
||||
)
|
||||
|
||||
type UTF16 struct{}
|
||||
|
||||
func (d *UTF16) FromChunk(chunk *sources.Chunk) *DecodableChunk {
|
||||
func (d *UTF16) FromChunk(chunk *sources.Chunk) *sources.Chunk {
|
||||
if chunk == nil || len(chunk.Data) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
decodableChunk := &DecodableChunk{Chunk: chunk, DecoderType: detectorspb.DecoderType_UTF16}
|
||||
if utf16Data, err := utf16ToUTF8(chunk.Data); err == nil {
|
||||
if len(utf16Data) == 0 {
|
||||
return nil
|
||||
}
|
||||
chunk.Data = utf16Data
|
||||
return decodableChunk
|
||||
return chunk
|
||||
}
|
||||
|
||||
return nil
|
||||
|
||||
@@ -4,25 +4,22 @@ import (
|
||||
"bytes"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/trufflesecurity/trufflehog/v3/pkg/pb/detectorspb"
|
||||
"github.com/trufflesecurity/trufflehog/v3/pkg/sources"
|
||||
)
|
||||
|
||||
type UTF8 struct{}
|
||||
|
||||
func (d *UTF8) FromChunk(chunk *sources.Chunk) *DecodableChunk {
|
||||
func (d *UTF8) FromChunk(chunk *sources.Chunk) *sources.Chunk {
|
||||
if chunk == nil || len(chunk.Data) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
decodableChunk := &DecodableChunk{Chunk: chunk, DecoderType: detectorspb.DecoderType_PLAIN}
|
||||
|
||||
if !utf8.Valid(chunk.Data) {
|
||||
chunk.Data = extractSubstrings(chunk.Data)
|
||||
return decodableChunk
|
||||
return chunk
|
||||
}
|
||||
|
||||
return decodableChunk
|
||||
return chunk
|
||||
}
|
||||
|
||||
// extractSubstrings performs similarly to the strings binutil,
|
||||
|
||||
@@ -2,7 +2,6 @@ package codeclimate
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"regexp"
|
||||
@@ -31,12 +30,6 @@ func (s Scanner) Keywords() []string {
|
||||
return []string{"codeclimate"}
|
||||
}
|
||||
|
||||
type response struct {
|
||||
Data struct {
|
||||
Id string `json:"id"`
|
||||
} `json:"data"`
|
||||
}
|
||||
|
||||
// FromData will find and optionally verify Codeclimate secrets in a given set of bytes.
|
||||
func (s Scanner) FromData(ctx context.Context, verify bool, data []byte) (results []detectors.Result, err error) {
|
||||
dataStr := string(data)
|
||||
@@ -65,14 +58,7 @@ func (s Scanner) FromData(ctx context.Context, verify bool, data []byte) (result
|
||||
if err == nil {
|
||||
defer res.Body.Close()
|
||||
if res.StatusCode >= 200 && res.StatusCode < 300 {
|
||||
var r response
|
||||
if err := json.NewDecoder(res.Body).Decode(&r); err != nil {
|
||||
s1.VerificationError = err
|
||||
continue
|
||||
}
|
||||
if r.Data.Id != "" {
|
||||
s1.Verified = true
|
||||
}
|
||||
s1.Verified = true
|
||||
} else {
|
||||
// This function will check false positives for common test words, but also it will make sure the key appears 'random' enough to be a real key
|
||||
if detectors.IsKnownFalsePositive(resMatch, detectors.DefaultFalsePositives, true) {
|
||||
|
||||
@@ -1,105 +0,0 @@
|
||||
package instamojo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"regexp"
|
||||
"strings"
|
||||
|
||||
"github.com/trufflesecurity/trufflehog/v3/pkg/common"
|
||||
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors"
|
||||
"github.com/trufflesecurity/trufflehog/v3/pkg/pb/detectorspb"
|
||||
)
|
||||
|
||||
type Scanner struct {
|
||||
client *http.Client
|
||||
}
|
||||
|
||||
// Ensure the Scanner satisfies the interface at compile time.
|
||||
var _ detectors.Detector = (*Scanner)(nil)
|
||||
|
||||
var (
|
||||
defaultClient = common.SaneHttpClient()
|
||||
// Make sure that your group is surrounded in boundary characters such as below to reduce false positives.
|
||||
//KeyPat is client_id
|
||||
keyPat = regexp.MustCompile(detectors.PrefixRegex([]string{"instamojo"}) + `\b([0-9a-zA-Z]{40})\b`)
|
||||
//Secretpat is Client_secret
|
||||
secretPat = regexp.MustCompile(detectors.PrefixRegex([]string{"instamojo"}) + `\b([0-9a-zA-Z]{128})\b`)
|
||||
)
|
||||
|
||||
// Keywords are used for efficiently pre-filtering chunks.
|
||||
// Use identifiers in the secret preferably, or the provider name.
|
||||
func (s Scanner) Keywords() []string {
|
||||
return []string{"instamojo"}
|
||||
}
|
||||
|
||||
// FromData will find and optionally verify Instamojo secrets in a given set of bytes.
|
||||
func (s Scanner) FromData(ctx context.Context, verify bool, data []byte) (results []detectors.Result, err error) {
|
||||
dataStr := string(data)
|
||||
|
||||
secretMatches := secretPat.FindAllStringSubmatch(dataStr, -1)
|
||||
clientIdmatches := keyPat.FindAllStringSubmatch(dataStr, -1)
|
||||
|
||||
for _, match := range secretMatches {
|
||||
if len(match) != 2 {
|
||||
continue
|
||||
}
|
||||
resSecret := strings.TrimSpace(match[1])
|
||||
|
||||
for _, clientIdMatch := range clientIdmatches {
|
||||
if len(clientIdMatch) != 2 {
|
||||
continue
|
||||
}
|
||||
resClientId := strings.TrimSpace(clientIdMatch[1])
|
||||
|
||||
s1 := detectors.Result{
|
||||
DetectorType: detectorspb.DetectorType_Instamojo,
|
||||
Raw: []byte(resClientId),
|
||||
}
|
||||
|
||||
if verify {
|
||||
client := s.client
|
||||
if client == nil {
|
||||
client = defaultClient
|
||||
}
|
||||
payload := strings.NewReader("grant_type=client_credentials&client_id=" + resClientId + "&client_secret=" + resSecret)
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, "POST", "https://api.instamojo.com/oauth2/token/", payload)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
req.Header.Add("Content-Type", "application/x-www-form-urlencoded")
|
||||
res, err := client.Do(req)
|
||||
if err == nil {
|
||||
defer res.Body.Close()
|
||||
bodyBytes, err := io.ReadAll(res.Body)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
body := string(bodyBytes)
|
||||
if (res.StatusCode >= 200 && res.StatusCode < 300) && strings.Contains(body, "access_token") {
|
||||
s1.Verified = true
|
||||
} else {
|
||||
s1.VerificationError = fmt.Errorf("unexpected HTTP response status %d", res.StatusCode)
|
||||
}
|
||||
} else {
|
||||
s1.VerificationError = err
|
||||
}
|
||||
}
|
||||
|
||||
if !s1.Verified && detectors.IsKnownFalsePositive(string(s1.Raw), detectors.DefaultFalsePositives, true) {
|
||||
continue
|
||||
}
|
||||
|
||||
results = append(results, s1)
|
||||
}
|
||||
}
|
||||
|
||||
return results, nil
|
||||
}
|
||||
|
||||
func (s Scanner) Type() detectorspb.DetectorType {
|
||||
return detectorspb.DetectorType_Instamojo
|
||||
}
|
||||
@@ -1,124 +0,0 @@
|
||||
//go:build detectors
|
||||
// +build detectors
|
||||
|
||||
package instamojo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/kylelemons/godebug/pretty"
|
||||
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors"
|
||||
|
||||
"github.com/trufflesecurity/trufflehog/v3/pkg/common"
|
||||
"github.com/trufflesecurity/trufflehog/v3/pkg/pb/detectorspb"
|
||||
)
|
||||
|
||||
func TestInstamojo_FromChunk(t *testing.T) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
|
||||
defer cancel()
|
||||
testSecrets, err := common.GetSecret(ctx, "trufflehog-testing", "detectors5")
|
||||
if err != nil {
|
||||
t.Fatalf("could not get test secrets from GCP: %s", err)
|
||||
}
|
||||
id := testSecrets.MustGetField("INSTAMOJO_CLIENT_ID")
|
||||
secret := testSecrets.MustGetField("INSTAMOJO_SECRET")
|
||||
inactiveSecret := testSecrets.MustGetField("INSTAMOJO_INACTIVE")
|
||||
|
||||
type args struct {
|
||||
ctx context.Context
|
||||
data []byte
|
||||
verify bool
|
||||
}
|
||||
tests := []struct {
|
||||
name string
|
||||
s Scanner
|
||||
args args
|
||||
want []detectors.Result
|
||||
wantErr bool
|
||||
wantVerificationErr bool
|
||||
}{
|
||||
{
|
||||
name: "found, verified",
|
||||
s: Scanner{},
|
||||
args: args{
|
||||
ctx: context.Background(),
|
||||
data: []byte(fmt.Sprintf("You can find a instamojo secret %s within id %s", secret, id)),
|
||||
verify: true,
|
||||
},
|
||||
want: []detectors.Result{
|
||||
{
|
||||
DetectorType: detectorspb.DetectorType_Instamojo,
|
||||
Verified: true,
|
||||
},
|
||||
},
|
||||
wantErr: false,
|
||||
wantVerificationErr: false,
|
||||
},
|
||||
{
|
||||
name: "found, unverified",
|
||||
s: Scanner{},
|
||||
args: args{
|
||||
ctx: context.Background(),
|
||||
data: []byte(fmt.Sprintf("You can find a instamojo secret %s within but not valid, within id %s", inactiveSecret, id)), // the secret would satisfy the regex but not pass validation
|
||||
verify: true,
|
||||
},
|
||||
want: []detectors.Result{
|
||||
{
|
||||
DetectorType: detectorspb.DetectorType_Instamojo,
|
||||
VerificationError: fmt.Errorf("unexpected HTTP response status 401"),
|
||||
},
|
||||
},
|
||||
wantErr: false,
|
||||
wantVerificationErr: true,
|
||||
},
|
||||
{
|
||||
name: "not found",
|
||||
s: Scanner{},
|
||||
args: args{
|
||||
ctx: context.Background(),
|
||||
data: []byte("You cannot find the secret within"),
|
||||
verify: true,
|
||||
},
|
||||
want: nil,
|
||||
wantErr: false,
|
||||
wantVerificationErr: true,
|
||||
},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got, err := tt.s.FromData(tt.args.ctx, tt.args.verify, tt.args.data)
|
||||
if (err != nil) != tt.wantErr {
|
||||
t.Errorf("Instamojo.FromData() error = %v, wantErr %v", err, tt.wantErr)
|
||||
return
|
||||
}
|
||||
for i := range got {
|
||||
if len(got[i].Raw) == 0 {
|
||||
t.Fatalf("no raw secret present: \n %+v", got[i])
|
||||
}
|
||||
got[i].Raw = nil
|
||||
}
|
||||
if diff := pretty.Compare(got, tt.want); diff != "" {
|
||||
t.Errorf("Instamojo.FromData() %s diff: (-got +want)\n%s", tt.name, diff)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkFromData(benchmark *testing.B) {
|
||||
ctx := context.Background()
|
||||
s := Scanner{}
|
||||
for name, data := range detectors.MustGetBenchmarkData() {
|
||||
benchmark.Run(name, func(b *testing.B) {
|
||||
b.ResetTimer()
|
||||
for n := 0; n < b.N; n++ {
|
||||
_, err := s.FromData(ctx, false, data)
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -1,84 +0,0 @@
|
||||
package ip2location
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"regexp"
|
||||
"strings"
|
||||
|
||||
"github.com/trufflesecurity/trufflehog/v3/pkg/common"
|
||||
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors"
|
||||
"github.com/trufflesecurity/trufflehog/v3/pkg/pb/detectorspb"
|
||||
)
|
||||
|
||||
type Scanner struct {}
|
||||
|
||||
// Ensure the Scanner satisfies the interface at compile time.
|
||||
var _ detectors.Detector = (*Scanner)(nil)
|
||||
|
||||
var (
|
||||
client = common.SaneHttpClient()
|
||||
// Make sure that your group is surrounded in boundary characters such as below to reduce false positives.
|
||||
keyPat = regexp.MustCompile(detectors.PrefixRegex([]string{"ip2location"}) + `\b([0-9A-Z]{32})\b`)
|
||||
)
|
||||
|
||||
// Keywords are used for efficiently pre-filtering chunks.
|
||||
// Use identifiers in the secret preferably, or the provider name.
|
||||
func (s Scanner) Keywords() []string {
|
||||
return []string{"ip2location"}
|
||||
}
|
||||
|
||||
// FromData will find and optionally verify Ip2location secrets in a given set of bytes.
|
||||
func (s Scanner) FromData(ctx context.Context, verify bool, data []byte) (results []detectors.Result, err error) {
|
||||
dataStr := string(data)
|
||||
|
||||
matches := keyPat.FindAllStringSubmatch(dataStr, -1)
|
||||
|
||||
for _, match := range matches {
|
||||
if len(match) != 2 {
|
||||
continue
|
||||
}
|
||||
resMatch := strings.TrimSpace(match[1])
|
||||
|
||||
s1 := detectors.Result{
|
||||
DetectorType: detectorspb.DetectorType_Ip2location,
|
||||
Raw: []byte(resMatch),
|
||||
}
|
||||
|
||||
if verify {
|
||||
req, err := http.NewRequestWithContext(ctx, "GET", "https://api.ip2location.io/?key=" + resMatch, nil)
|
||||
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
res, err := client.Do(req)
|
||||
if err == nil {
|
||||
defer res.Body.Close()
|
||||
if res.StatusCode >= 200 && res.StatusCode < 300 {
|
||||
s1.Verified = true
|
||||
} else if res.StatusCode == 401 {
|
||||
// The secret is determinately not verified (nothing to do)
|
||||
} else {
|
||||
s1.VerificationError = fmt.Errorf("unexpected HTTP response status %d", res.StatusCode)
|
||||
}
|
||||
} else {
|
||||
s1.VerificationError = err
|
||||
}
|
||||
}
|
||||
if !s1.Verified && detectors.IsKnownFalsePositive(resMatch, detectors.DefaultFalsePositives, true) {
|
||||
continue
|
||||
}
|
||||
|
||||
results = append(results, s1)
|
||||
}
|
||||
|
||||
return results, nil
|
||||
}
|
||||
|
||||
func (s Scanner) Type() detectorspb.DetectorType {
|
||||
return detectorspb.DetectorType_Ip2location
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -1,128 +0,0 @@
|
||||
//go:build detectors
|
||||
// +build detectors
|
||||
|
||||
package ip2location
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/google/go-cmp/cmp"
|
||||
"github.com/google/go-cmp/cmp/cmpopts"
|
||||
|
||||
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors"
|
||||
|
||||
"github.com/trufflesecurity/trufflehog/v3/pkg/common"
|
||||
"github.com/trufflesecurity/trufflehog/v3/pkg/pb/detectorspb"
|
||||
)
|
||||
|
||||
func TestIp2location_FromChunk(t *testing.T) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
|
||||
defer cancel()
|
||||
testSecrets, err := common.GetSecret(ctx, "trufflehog-testing", "detectors5")
|
||||
if err != nil {
|
||||
t.Fatalf("could not get test secrets from GCP: %s", err)
|
||||
}
|
||||
secret := testSecrets.MustGetField("IP2LOCATION")
|
||||
inactiveSecret := testSecrets.MustGetField("IP2LOCATION_INACTIVE")
|
||||
|
||||
type args struct {
|
||||
ctx context.Context
|
||||
data []byte
|
||||
verify bool
|
||||
}
|
||||
tests := []struct {
|
||||
name string
|
||||
s Scanner
|
||||
args args
|
||||
want []detectors.Result
|
||||
wantErr bool
|
||||
wantVerificationErr bool
|
||||
}{
|
||||
{
|
||||
name: "found, verified",
|
||||
s: Scanner{},
|
||||
args: args{
|
||||
ctx: context.Background(),
|
||||
data: []byte(fmt.Sprintf("You can find a ip2location secret %s within", secret)),
|
||||
verify: true,
|
||||
},
|
||||
want: []detectors.Result{
|
||||
{
|
||||
DetectorType: detectorspb.DetectorType_Ip2location,
|
||||
Verified: true,
|
||||
},
|
||||
},
|
||||
wantErr: false,
|
||||
wantVerificationErr: false,
|
||||
},
|
||||
{
|
||||
name: "found, unverified",
|
||||
s: Scanner{},
|
||||
args: args{
|
||||
ctx: context.Background(),
|
||||
data: []byte(fmt.Sprintf("You can find a ip2location secret %s within but not valid", inactiveSecret)), // the secret would satisfy the regex but not pass validation
|
||||
verify: true,
|
||||
},
|
||||
want: []detectors.Result{
|
||||
{
|
||||
DetectorType: detectorspb.DetectorType_Ip2location,
|
||||
Verified: false,
|
||||
},
|
||||
},
|
||||
wantErr: false,
|
||||
wantVerificationErr: false,
|
||||
},
|
||||
{
|
||||
name: "not found",
|
||||
s: Scanner{},
|
||||
args: args{
|
||||
ctx: context.Background(),
|
||||
data: []byte("You cannot find the secret within"),
|
||||
verify: true,
|
||||
},
|
||||
want: nil,
|
||||
wantErr: false,
|
||||
wantVerificationErr: false,
|
||||
},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got, err := tt.s.FromData(tt.args.ctx, tt.args.verify, tt.args.data)
|
||||
if (err != nil) != tt.wantErr {
|
||||
t.Errorf("Ip2location.FromData() error = %v, wantErr %v", err, tt.wantErr)
|
||||
return
|
||||
}
|
||||
for i := range got {
|
||||
if len(got[i].Raw) == 0 {
|
||||
t.Fatalf("no raw secret present: \n %+v", got[i])
|
||||
}
|
||||
if (got[i].VerificationError != nil) != tt.wantVerificationErr {
|
||||
t.Fatalf("wantVerificationError = %v, verification error = %v", tt.wantVerificationErr, got[i].VerificationError)
|
||||
}
|
||||
}
|
||||
ignoreOpts := cmpopts.IgnoreFields(detectors.Result{}, "Raw", "VerificationError")
|
||||
if diff := cmp.Diff(got, tt.want, ignoreOpts); diff != "" {
|
||||
t.Errorf("Ip2location.FromData() %s diff: (-got +want)\n%s", tt.name, diff)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkFromData(benchmark *testing.B) {
|
||||
ctx := context.Background()
|
||||
s := Scanner{}
|
||||
for name, data := range detectors.MustGetBenchmarkData() {
|
||||
benchmark.Run(name, func(b *testing.B) {
|
||||
b.ResetTimer()
|
||||
for n := 0; n < b.N; n++ {
|
||||
_, err := s.FromData(ctx, false, data)
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -1,89 +0,0 @@
|
||||
package ipinfo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"regexp"
|
||||
"strings"
|
||||
|
||||
"github.com/trufflesecurity/trufflehog/v3/pkg/common"
|
||||
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors"
|
||||
"github.com/trufflesecurity/trufflehog/v3/pkg/pb/detectorspb"
|
||||
)
|
||||
|
||||
type Scanner struct {
|
||||
client *http.Client
|
||||
}
|
||||
|
||||
// Ensure the Scanner satisfies the interface at compile time.
|
||||
var _ detectors.Detector = (*Scanner)(nil)
|
||||
|
||||
var (
|
||||
defaultClient = common.SaneHttpClient()
|
||||
// Make sure that your group is surrounded in boundary characters such as below to reduce false positives.
|
||||
keyPat = regexp.MustCompile(detectors.PrefixRegex([]string{"ipinfo"}) + `\b([a-f0-9]{14})\b`)
|
||||
)
|
||||
|
||||
// Keywords are used for efficiently pre-filtering chunks.
|
||||
// Use identifiers in the secret preferably, or the provider name.
|
||||
func (s Scanner) Keywords() []string {
|
||||
return []string{"ipinfo"}
|
||||
}
|
||||
|
||||
// FromData will find and optionally verify Ipinfo secrets in a given set of bytes.
|
||||
func (s Scanner) FromData(ctx context.Context, verify bool, data []byte) (results []detectors.Result, err error) {
|
||||
dataStr := string(data)
|
||||
|
||||
matches := keyPat.FindAllStringSubmatch(dataStr, -1)
|
||||
|
||||
for _, match := range matches {
|
||||
if len(match) != 2 {
|
||||
continue
|
||||
}
|
||||
resMatch := strings.TrimSpace(match[1])
|
||||
|
||||
s1 := detectors.Result{
|
||||
DetectorType: detectorspb.DetectorType_IPInfo,
|
||||
Raw: []byte(resMatch),
|
||||
}
|
||||
|
||||
if verify {
|
||||
client := s.client
|
||||
if client == nil {
|
||||
client = defaultClient
|
||||
}
|
||||
req, err := http.NewRequestWithContext(ctx, "GET", "https://ipinfo.io/json?token="+resMatch, nil)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
res, err := client.Do(req)
|
||||
if err == nil {
|
||||
fmt.Println(res.Status, resMatch)
|
||||
defer res.Body.Close()
|
||||
if res.StatusCode >= 200 && res.StatusCode < 300 {
|
||||
s1.Verified = true
|
||||
} else if res.StatusCode == 403 {
|
||||
// The secret is determinately not verified (nothing to do)
|
||||
} else {
|
||||
s1.VerificationError = fmt.Errorf("unexpected HTTP response status %d", res.StatusCode)
|
||||
}
|
||||
} else {
|
||||
s1.VerificationError = err
|
||||
}
|
||||
}
|
||||
|
||||
// This function will check false positives for common test words, but also it will make sure the key appears 'random' enough to be a real key.
|
||||
if !s1.Verified && detectors.IsKnownFalsePositive(resMatch, detectors.DefaultFalsePositives, true) {
|
||||
continue
|
||||
}
|
||||
|
||||
results = append(results, s1)
|
||||
}
|
||||
|
||||
return results, nil
|
||||
}
|
||||
|
||||
func (s Scanner) Type() detectorspb.DetectorType {
|
||||
return detectorspb.DetectorType_IPInfo
|
||||
}
|
||||
@@ -1,161 +0,0 @@
|
||||
//go:build detectors
|
||||
// +build detectors
|
||||
|
||||
package ipinfo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"github.com/google/go-cmp/cmp"
|
||||
"github.com/google/go-cmp/cmp/cmpopts"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors"
|
||||
|
||||
"github.com/trufflesecurity/trufflehog/v3/pkg/common"
|
||||
"github.com/trufflesecurity/trufflehog/v3/pkg/pb/detectorspb"
|
||||
)
|
||||
|
||||
func TestIpinfo_FromChunk(t *testing.T) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
|
||||
defer cancel()
|
||||
testSecrets, err := common.GetSecret(ctx, "trufflehog-testing", "detectors5")
|
||||
if err != nil {
|
||||
t.Fatalf("could not get test secrets from GCP: %s", err)
|
||||
}
|
||||
secret := testSecrets.MustGetField("IPINFO")
|
||||
inactiveSecret := testSecrets.MustGetField("IPINFO_INACTIVE")
|
||||
|
||||
type args struct {
|
||||
ctx context.Context
|
||||
data []byte
|
||||
verify bool
|
||||
}
|
||||
tests := []struct {
|
||||
name string
|
||||
s Scanner
|
||||
args args
|
||||
want []detectors.Result
|
||||
wantErr bool
|
||||
wantVerificationErr bool
|
||||
}{
|
||||
{
|
||||
name: "found, verified",
|
||||
s: Scanner{},
|
||||
args: args{
|
||||
ctx: context.Background(),
|
||||
data: []byte(fmt.Sprintf("You can find a ipinfo secret %s within", secret)),
|
||||
verify: true,
|
||||
},
|
||||
want: []detectors.Result{
|
||||
{
|
||||
DetectorType: detectorspb.DetectorType_IPInfo,
|
||||
Verified: true,
|
||||
},
|
||||
},
|
||||
wantErr: false,
|
||||
wantVerificationErr: false,
|
||||
},
|
||||
{
|
||||
name: "found, unverified",
|
||||
s: Scanner{},
|
||||
args: args{
|
||||
ctx: context.Background(),
|
||||
data: []byte(fmt.Sprintf("You can find a ipinfo secret %s within but not valid", inactiveSecret)), // the secret would satisfy the regex but not pass validation
|
||||
verify: true,
|
||||
},
|
||||
want: []detectors.Result{
|
||||
{
|
||||
DetectorType: detectorspb.DetectorType_IPInfo,
|
||||
Verified: false,
|
||||
},
|
||||
},
|
||||
wantErr: false,
|
||||
wantVerificationErr: false,
|
||||
},
|
||||
{
|
||||
name: "not found",
|
||||
s: Scanner{},
|
||||
args: args{
|
||||
ctx: context.Background(),
|
||||
data: []byte("You cannot find the secret within"),
|
||||
verify: true,
|
||||
},
|
||||
want: nil,
|
||||
wantErr: false,
|
||||
wantVerificationErr: false,
|
||||
},
|
||||
{
|
||||
name: "found, would be verified if not for timeout",
|
||||
s: Scanner{client: common.SaneHttpClientTimeOut(1 * time.Microsecond)},
|
||||
args: args{
|
||||
ctx: context.Background(),
|
||||
data: []byte(fmt.Sprintf("You can find a ipinfo secret %s within", secret)),
|
||||
verify: true,
|
||||
},
|
||||
want: []detectors.Result{
|
||||
{
|
||||
DetectorType: detectorspb.DetectorType_IPInfo,
|
||||
Verified: false,
|
||||
},
|
||||
},
|
||||
wantErr: false,
|
||||
wantVerificationErr: true,
|
||||
},
|
||||
{
|
||||
name: "found, verified but unexpected api surface",
|
||||
s: Scanner{client: common.ConstantResponseHttpClient(404, "")},
|
||||
args: args{
|
||||
ctx: context.Background(),
|
||||
data: []byte(fmt.Sprintf("You can find a ipinfo secret %s within", secret)),
|
||||
verify: true,
|
||||
},
|
||||
want: []detectors.Result{
|
||||
{
|
||||
DetectorType: detectorspb.DetectorType_IPInfo,
|
||||
Verified: false,
|
||||
},
|
||||
},
|
||||
wantErr: false,
|
||||
wantVerificationErr: true,
|
||||
},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got, err := tt.s.FromData(tt.args.ctx, tt.args.verify, tt.args.data)
|
||||
if (err != nil) != tt.wantErr {
|
||||
t.Errorf("Ipinfo.FromData() error = %v, wantErr %v", err, tt.wantErr)
|
||||
return
|
||||
}
|
||||
for i := range got {
|
||||
if len(got[i].Raw) == 0 {
|
||||
t.Fatalf("no raw secret present: \n %+v", got[i])
|
||||
}
|
||||
if (got[i].VerificationError != nil) != tt.wantVerificationErr {
|
||||
t.Fatalf("wantVerificationError = %v, verification error = %v", tt.wantVerificationErr, got[i].VerificationError)
|
||||
}
|
||||
}
|
||||
ignoreOpts := cmpopts.IgnoreFields(detectors.Result{}, "Raw", "VerificationError")
|
||||
if diff := cmp.Diff(got, tt.want, ignoreOpts); diff != "" {
|
||||
t.Errorf("Ipinfo.FromData() %s diff: (-got +want)\n%s", tt.name, diff)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkFromData(benchmark *testing.B) {
|
||||
ctx := context.Background()
|
||||
s := Scanner{}
|
||||
for name, data := range detectors.MustGetBenchmarkData() {
|
||||
benchmark.Run(name, func(b *testing.B) {
|
||||
b.ResetTimer()
|
||||
for n := 0; n < b.N; n++ {
|
||||
_, err := s.FromData(ctx, false, data)
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -47,7 +47,7 @@ func (s Scanner) FromData(ctx context.Context, verify bool, data []byte) (result
|
||||
}
|
||||
|
||||
if verify {
|
||||
req, err := http.NewRequestWithContext(ctx, "GET", "https://api.supernotes.app/v1/user", nil)
|
||||
req, err := http.NewRequestWithContext(ctx, "GET", "https://api.supernotes.app/v1/user/", nil)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
|
||||
@@ -1,111 +0,0 @@
|
||||
package engine
|
||||
|
||||
import (
|
||||
"strings"
|
||||
|
||||
ahocorasick "github.com/BobuSumisu/aho-corasick"
|
||||
|
||||
"github.com/trufflesecurity/trufflehog/v3/pkg/context"
|
||||
"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.
|
||||
type detectorKey struct {
|
||||
detectorType detectorspb.DetectorType
|
||||
version int
|
||||
}
|
||||
|
||||
// DetectorInfo is used to store a detector and whether it should be verified.
|
||||
type DetectorInfo struct {
|
||||
detectors.Detector
|
||||
ShouldVerify bool
|
||||
}
|
||||
|
||||
// AhoCorasickCore encapsulates the operations and data structures used for keyword matching via the
|
||||
// Aho-Corasick algorithm. It is responsible for constructing and managing the trie for efficient
|
||||
// substring searches, as well as mapping keywords to their associated detectors for rapid lookups.
|
||||
type AhoCorasickCore struct {
|
||||
// prefilter is a ahocorasick struct used for doing efficient string
|
||||
// matching given a set of words. (keywords from the rules in the config)
|
||||
prefilter ahocorasick.Trie
|
||||
// Maps for efficient lookups during detection.
|
||||
detectorTypeToDetectorInfo map[detectorKey]DetectorInfo
|
||||
detectors map[bool][]detectors.Detector
|
||||
keywordsToDetectors map[string][]detectorKey
|
||||
}
|
||||
|
||||
// NewAhoCorasickCore allocates and initializes a new instance of AhoCorasickCore.
|
||||
// It creates an empty keyword-to-detectors map for future string matching operations.
|
||||
// The map detectorTypeToDetectorInfo is pre-allocated based on the size of detectors
|
||||
// provided, for efficient storage and lookup of detector information.
|
||||
func NewAhoCorasickCore(detectors map[bool][]detectors.Detector) *AhoCorasickCore {
|
||||
return &AhoCorasickCore{
|
||||
keywordsToDetectors: make(map[string][]detectorKey),
|
||||
detectors: detectors,
|
||||
detectorTypeToDetectorInfo: make(map[detectorKey]DetectorInfo, len(detectors[true])+len(detectors[false])),
|
||||
}
|
||||
}
|
||||
|
||||
// Setup initializes the internal state of AhoCorasickCore to prepare it for keyword matching.
|
||||
// This involves pre-filtering setup and lookup optimization, critical for the engine's performance.
|
||||
func (ac *AhoCorasickCore) Setup(ctx context.Context) {
|
||||
// Prepare maps for fast detector lookups, instead of scanning through an array of detectors for every chunk.
|
||||
var keywords []string
|
||||
for verify, detectorsSet := range ac.detectors {
|
||||
for _, d := range detectorsSet {
|
||||
key := createDetectorKey(d)
|
||||
ac.detectorTypeToDetectorInfo[key] = DetectorInfo{Detector: d, ShouldVerify: verify}
|
||||
keywords = ac.extractAndMapKeywords(d, key, keywords)
|
||||
}
|
||||
}
|
||||
|
||||
// Use the Ahocorasick algorithm to create a trie structure for efficient keyword matching.
|
||||
// This ensures that we can rapidly match against a vast set of keywords without individually comparing each one.
|
||||
ac.prefilter = *ahocorasick.NewTrieBuilder().AddStrings(keywords).Build()
|
||||
ctx.Logger().V(4).Info("AhoCorasickCore Setup complete")
|
||||
}
|
||||
|
||||
// createDetectorKey creates a unique key for each detector. This key based on type and version,
|
||||
// it ensures faster lookups and reduces redundancy in our main detector store.
|
||||
func createDetectorKey(d detectors.Detector) detectorKey {
|
||||
detectorType := d.Type()
|
||||
var version int
|
||||
if v, ok := d.(detectors.Versioner); ok {
|
||||
version = v.Version()
|
||||
}
|
||||
return detectorKey{detectorType: detectorType, version: version}
|
||||
}
|
||||
|
||||
// extractAndMapKeywords captures keywords associated with each detector and maps them.
|
||||
// This allows us to quickly determine which detectors are relevant based on the presence of certain keywords.
|
||||
func (ac *AhoCorasickCore) extractAndMapKeywords(d detectors.Detector, key detectorKey, keywords []string) []string {
|
||||
for _, kw := range d.Keywords() {
|
||||
kwLower := strings.ToLower(kw)
|
||||
keywords = append(keywords, kwLower)
|
||||
ac.keywordsToDetectors[kwLower] = append(ac.keywordsToDetectors[kwLower], key)
|
||||
}
|
||||
return keywords
|
||||
}
|
||||
|
||||
// MatchString performs a string match using the Aho-Corasick algorithm, returning an array of matches.
|
||||
// Designed for internal use within the AhoCorasickCore component.
|
||||
func (ac *AhoCorasickCore) MatchString(input string) []*ahocorasick.Match {
|
||||
return ac.prefilter.MatchString(strings.ToLower(input))
|
||||
}
|
||||
|
||||
// PopulateDetectorsByMatch populates the given detectorMap based on the Aho-Corasick match results.
|
||||
// This method is designed to reuse the same map for performance optimization,
|
||||
// reducing the need for repeated allocations within each detector worker in the engine.
|
||||
func (ac *AhoCorasickCore) PopulateDetectorsByMatch(match *ahocorasick.Match, detectors map[detectorspb.DetectorType]DetectorInfo) bool {
|
||||
matchedKeys, ok := ac.keywordsToDetectors[match.MatchString()]
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
for _, key := range matchedKeys {
|
||||
detectors[key.detectorType] = ac.detectorTypeToDetectorInfo[key]
|
||||
}
|
||||
return true
|
||||
}
|
||||
@@ -4,7 +4,6 @@ import (
|
||||
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/anthropic"
|
||||
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/envoyapikey"
|
||||
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/huggingface"
|
||||
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/ip2location"
|
||||
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/ramp"
|
||||
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/salesforce"
|
||||
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/snowflake"
|
||||
@@ -1557,7 +1556,6 @@ func DefaultDetectors() []detectors.Detector {
|
||||
&anthropic.Scanner{},
|
||||
&sourcegraphcody.Scanner{},
|
||||
voiceflow.Scanner{},
|
||||
ip2location.Scanner{},
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
+110
-98
@@ -3,11 +3,14 @@ package engine
|
||||
import (
|
||||
"bytes"
|
||||
"fmt"
|
||||
"reflect"
|
||||
"runtime"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
ahocorasick "github.com/BobuSumisu/aho-corasick"
|
||||
lru "github.com/hashicorp/golang-lru"
|
||||
"google.golang.org/protobuf/proto"
|
||||
|
||||
@@ -63,8 +66,9 @@ type Engine struct {
|
||||
onlyVerified bool
|
||||
printAvgDetectorTime bool
|
||||
|
||||
// ahoCorasickHandler manages the Aho-Corasick trie and related keyword lookups.
|
||||
ahoCorasickCore *AhoCorasickCore
|
||||
// prefilter is a ahocorasick struct used for doing efficient string
|
||||
// matching given a set of words (keywords from the rules in the config)
|
||||
prefilter ahocorasick.Trie
|
||||
|
||||
// Engine synchronization primitives.
|
||||
sourceManager *sources.SourceManager
|
||||
@@ -89,10 +93,9 @@ type Engine struct {
|
||||
dedupeCache *lru.Cache
|
||||
}
|
||||
|
||||
// Option is used to configure the engine during initialization using functional options.
|
||||
type Option func(*Engine)
|
||||
type EngineOption func(*Engine)
|
||||
|
||||
func WithConcurrency(concurrency uint8) Option {
|
||||
func WithConcurrency(concurrency uint8) EngineOption {
|
||||
return func(e *Engine) {
|
||||
e.concurrency = concurrency
|
||||
}
|
||||
@@ -100,7 +103,7 @@ func WithConcurrency(concurrency uint8) Option {
|
||||
|
||||
const ignoreTag = "trufflehog:ignore"
|
||||
|
||||
func WithDetectors(verify bool, d ...detectors.Detector) Option {
|
||||
func WithDetectors(verify bool, d ...detectors.Detector) EngineOption {
|
||||
return func(e *Engine) {
|
||||
if e.detectors == nil {
|
||||
e.detectors = make(map[bool][]detectors.Detector)
|
||||
@@ -113,7 +116,7 @@ func WithDetectors(verify bool, d ...detectors.Detector) Option {
|
||||
}
|
||||
}
|
||||
|
||||
func WithDecoders(decoders ...decoders.Decoder) Option {
|
||||
func WithDecoders(decoders ...decoders.Decoder) EngineOption {
|
||||
return func(e *Engine) {
|
||||
e.decoders = decoders
|
||||
}
|
||||
@@ -121,14 +124,14 @@ func WithDecoders(decoders ...decoders.Decoder) Option {
|
||||
|
||||
// WithFilterUnverified sets the filterUnverified flag on the engine. If set to
|
||||
// true, the engine will only return the first unverified result for a chunk for a detector.
|
||||
func WithFilterUnverified(filter bool) Option {
|
||||
func WithFilterUnverified(filter bool) EngineOption {
|
||||
return func(e *Engine) {
|
||||
e.filterUnverified = filter
|
||||
}
|
||||
}
|
||||
|
||||
// WithFilterEntropy filters out unverified results using Shannon entropy.
|
||||
func WithFilterEntropy(entropy float64) Option {
|
||||
func WithFilterEntropy(entropy float64) EngineOption {
|
||||
return func(e *Engine) {
|
||||
if entropy > 0 {
|
||||
e.filterEntropy = &entropy
|
||||
@@ -138,7 +141,7 @@ func WithFilterEntropy(entropy float64) Option {
|
||||
|
||||
// WithOnlyVerified sets the onlyVerified flag on the engine. If set to true,
|
||||
// the engine will only print verified results.
|
||||
func WithOnlyVerified(onlyVerified bool) Option {
|
||||
func WithOnlyVerified(onlyVerified bool) EngineOption {
|
||||
return func(e *Engine) {
|
||||
e.onlyVerified = onlyVerified
|
||||
}
|
||||
@@ -150,7 +153,7 @@ func WithOnlyVerified(onlyVerified bool) Option {
|
||||
// the engine is configured to print the results.
|
||||
// Calculating the average time taken by each detector is an expensive operation
|
||||
// and should be avoided unless specified by the user.
|
||||
func WithPrintAvgDetectorTime(printAvgDetectorTime bool) Option {
|
||||
func WithPrintAvgDetectorTime(printAvgDetectorTime bool) EngineOption {
|
||||
return func(e *Engine) {
|
||||
e.printAvgDetectorTime = printAvgDetectorTime
|
||||
}
|
||||
@@ -160,7 +163,7 @@ func WithPrintAvgDetectorTime(printAvgDetectorTime bool) Option {
|
||||
// the filterFunc returns true, the detector will be included for scanning.
|
||||
// This option applies to the existing list of detectors configured, so the
|
||||
// order this option appears matters. All filtering happens before scanning.
|
||||
func WithFilterDetectors(filterFunc func(detectors.Detector) bool) Option {
|
||||
func WithFilterDetectors(filterFunc func(detectors.Detector) bool) EngineOption {
|
||||
return func(e *Engine) {
|
||||
// If no detectors are configured, do nothing.
|
||||
if e.detectors == nil {
|
||||
@@ -172,7 +175,7 @@ func WithFilterDetectors(filterFunc func(detectors.Detector) bool) Option {
|
||||
}
|
||||
|
||||
// WithPrinter sets the Printer on the engine.
|
||||
func WithPrinter(printer Printer) Option {
|
||||
func WithPrinter(printer Printer) EngineOption {
|
||||
return func(e *Engine) {
|
||||
e.printer = printer
|
||||
}
|
||||
@@ -266,60 +269,32 @@ func (e *Engine) DetectorAvgTime() map[string][]time.Duration {
|
||||
return avgTime
|
||||
}
|
||||
|
||||
// Start initializes and activates the engine's processing pipeline.
|
||||
// It sets up various default configurations, prepares lookup structures for
|
||||
// detectors, conducts basic sanity checks, and kickstarts all necessary workers.
|
||||
// Once started, the engine begins processing input data to identify secrets.
|
||||
func Start(ctx context.Context, options ...Option) (*Engine, error) {
|
||||
e := &Engine{}
|
||||
|
||||
if err := e.initialize(ctx, options...); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
e.setDefaults(ctx)
|
||||
ctx.Logger().V(4).Info("setting up aho-corasick core")
|
||||
e.ahoCorasickCore.Setup(ctx)
|
||||
e.sanityChecks(ctx)
|
||||
e.startWorkers(ctx)
|
||||
|
||||
return e, nil
|
||||
}
|
||||
|
||||
const defaultChannelBuffer = 1
|
||||
|
||||
// initialize prepares the engine's internal structures. The LRU cache optimizes
|
||||
// deduplication efforts, allowing the engine to quickly check if a chunk has
|
||||
// been processed before, thereby saving computational overhead.
|
||||
func (e *Engine) initialize(ctx context.Context, options ...Option) error {
|
||||
// TODO (ahrav): Determine the optimal cache size.
|
||||
const cacheSize = 512 // number of entries in the LRU cache
|
||||
// Start the engine with options.
|
||||
func Start(ctx context.Context, options ...EngineOption) (*Engine, error) {
|
||||
const (
|
||||
defaultChannelBuffer = 1
|
||||
// TODO (ahrav): Determine the optimal cache size.
|
||||
cacheSize = 512 // number of entries in the LRU cache
|
||||
)
|
||||
|
||||
cache, err := lru.New(cacheSize)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to initialize LRU cache: %w", err)
|
||||
return nil, fmt.Errorf("failed to initialize LRU cache: %w", err)
|
||||
}
|
||||
|
||||
// Channels are used for communication between different parts of the engine,
|
||||
// ensuring that data flows smoothly without race conditions.
|
||||
e.detectableChunksChan = make(chan detectableChunk, defaultChannelBuffer)
|
||||
e.results = make(chan detectors.ResultWithMetadata, defaultChannelBuffer)
|
||||
e.dedupeCache = cache
|
||||
e.printer = new(output.PlainPrinter)
|
||||
e.metrics = runtimeMetrics{Metrics: Metrics{scanStartTime: time.Now()}}
|
||||
e := &Engine{
|
||||
detectableChunksChan: make(chan detectableChunk, defaultChannelBuffer),
|
||||
results: make(chan detectors.ResultWithMetadata, defaultChannelBuffer),
|
||||
dedupeCache: cache,
|
||||
printer: new(output.PlainPrinter), // default printer
|
||||
metrics: runtimeMetrics{Metrics: Metrics{scanStartTime: time.Now()}},
|
||||
}
|
||||
|
||||
for _, option := range options {
|
||||
option(e)
|
||||
}
|
||||
ctx.Logger().V(4).Info("engine initialized")
|
||||
e.ahoCorasickCore = NewAhoCorasickCore(e.detectors)
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// setDefaults ensures that if specific engine properties aren't provided,
|
||||
// they're set to reasonable default values. It makes the engine robust to
|
||||
// incomplete configuration.
|
||||
func (e *Engine) setDefaults(ctx context.Context) {
|
||||
// Set defaults.
|
||||
if e.concurrency == 0 {
|
||||
numCPU := runtime.NumCPU()
|
||||
ctx.Logger().Info("No concurrency specified, defaulting to max", "cpu", numCPU)
|
||||
@@ -327,13 +302,12 @@ func (e *Engine) setDefaults(ctx context.Context) {
|
||||
}
|
||||
ctx.Logger().V(3).Info("engine started", "workers", e.concurrency)
|
||||
|
||||
// Create SourceManager.
|
||||
e.sourceManager = sources.NewManager(
|
||||
sources.WithConcurrentSources(int(e.concurrency)),
|
||||
sources.WithConcurrentUnits(int(e.concurrency)),
|
||||
sources.WithBufferedOutput(defaultChannelBuffer),
|
||||
)
|
||||
|
||||
// Default decoders handle common encoding formats.
|
||||
if len(e.decoders) == 0 {
|
||||
e.decoders = decoders.DefaultDecoders()
|
||||
}
|
||||
@@ -343,30 +317,46 @@ func (e *Engine) setDefaults(ctx context.Context) {
|
||||
e.detectors[true] = DefaultDetectors()
|
||||
e.detectors[false] = []detectors.Detector{}
|
||||
}
|
||||
ctx.Logger().V(4).Info("default engine options set")
|
||||
}
|
||||
|
||||
// Sanity check detectors for duplicate configuration. Only log in case
|
||||
// a detector has been configured in a way that isn't represented by
|
||||
// the DetectorID (type and version).
|
||||
func (e *Engine) sanityChecks(ctx context.Context) {
|
||||
dets := append(e.detectors[true], e.detectors[false]...)
|
||||
seenDetectors := make(map[config.DetectorID]struct{}, len(dets))
|
||||
for _, det := range dets {
|
||||
id := config.GetDetectorID(det)
|
||||
if _, ok := seenDetectors[id]; ok && id.ID != detectorspb.DetectorType_CustomRegex {
|
||||
ctx.Logger().Info("possible duplicate detector configured", "detector", id)
|
||||
// build ahocorasick prefilter for efficient string matching
|
||||
// on keywords
|
||||
keywords := []string{}
|
||||
for _, d := range e.detectors[false] {
|
||||
for _, kw := range d.Keywords() {
|
||||
keywords = append(keywords, strings.ToLower(kw))
|
||||
}
|
||||
}
|
||||
for _, d := range e.detectors[true] {
|
||||
for _, kw := range d.Keywords() {
|
||||
keywords = append(keywords, strings.ToLower(kw))
|
||||
}
|
||||
}
|
||||
e.prefilter = *ahocorasick.NewTrieBuilder().AddStrings(keywords).Build()
|
||||
|
||||
ctx.Logger().V(3).Info("loaded decoders", "count", len(e.decoders))
|
||||
ctx.Logger().V(3).Info("loaded detectors",
|
||||
"total", len(e.detectors[true])+len(e.detectors[false]),
|
||||
"verification_enabled", len(e.detectors[true]),
|
||||
"verification_disabled", len(e.detectors[false]),
|
||||
)
|
||||
|
||||
// Sanity check detectors for duplicate configuration. Only log in case
|
||||
// a detector has been configured in a way that isn't represented by
|
||||
// the DetectorID (type and version).
|
||||
{
|
||||
dets := append(e.detectors[true], e.detectors[false]...)
|
||||
seenDetectors := make(map[config.DetectorID]struct{}, len(dets))
|
||||
for _, det := range dets {
|
||||
id := config.GetDetectorID(det)
|
||||
if _, ok := seenDetectors[id]; ok && id.ID != detectorspb.DetectorType_CustomRegex {
|
||||
ctx.Logger().Info("possible duplicate detector configured", "detector", id)
|
||||
}
|
||||
seenDetectors[id] = struct{}{}
|
||||
}
|
||||
seenDetectors[id] = struct{}{}
|
||||
}
|
||||
}
|
||||
|
||||
// startWorkers initiates all necessary workers. Workers handle processing of
|
||||
// chunks concurrently. Separating the initialization of different types of
|
||||
// workers helps in scalability and makes it easier to diagnose issues.
|
||||
func (e *Engine) startWorkers(ctx context.Context) {
|
||||
// Scanner workers process input data and extract chunks for detectors.
|
||||
ctx.Logger().V(2).Info("starting scanner workers", "count", e.concurrency)
|
||||
// Run the Secret scanner workers and Notifier pipelines.
|
||||
for worker := uint64(0); worker < uint64(e.concurrency); worker++ {
|
||||
e.workersWg.Add(1)
|
||||
go func() {
|
||||
@@ -377,7 +367,6 @@ func (e *Engine) startWorkers(ctx context.Context) {
|
||||
}()
|
||||
}
|
||||
|
||||
// Detector workers apply keyword matching, regexes and API calls to detect secrets in chunks.
|
||||
const detectorWorkerMultiplier = 50
|
||||
ctx.Logger().V(2).Info("starting detector workers", "count", e.concurrency*detectorWorkerMultiplier)
|
||||
for worker := uint64(0); worker < uint64(e.concurrency*detectorWorkerMultiplier); worker++ {
|
||||
@@ -390,7 +379,6 @@ func (e *Engine) startWorkers(ctx context.Context) {
|
||||
}()
|
||||
}
|
||||
|
||||
// Notifier workers communicate detected issues to the user or any downstream systems.
|
||||
// We want 1/4th of the notifier workers as the number of scanner workers.
|
||||
const notifierWorkerRatio = 4
|
||||
maxNotifierWorkers := 1
|
||||
@@ -407,6 +395,8 @@ func (e *Engine) startWorkers(ctx context.Context) {
|
||||
e.notifyResults(ctx)
|
||||
}()
|
||||
}
|
||||
|
||||
return e, nil
|
||||
}
|
||||
|
||||
// Finish waits for running sources to complete and workers to finish scanning
|
||||
@@ -455,42 +445,64 @@ type detectableChunk struct {
|
||||
func (e *Engine) detectorWorker(ctx context.Context) {
|
||||
var wgDetect sync.WaitGroup
|
||||
|
||||
// Reuse the same map to avoid allocations.
|
||||
const avgDetectorsPerChunk = 2
|
||||
chunkSpecificDetectors := make(map[detectorspb.DetectorType]DetectorInfo, avgDetectorsPerChunk)
|
||||
for originalChunk := range e.ChunksChan() {
|
||||
for chunk := range sources.Chunker(originalChunk) {
|
||||
matchedKeywords := make(map[string]struct{})
|
||||
atomic.AddUint64(&e.metrics.BytesScanned, uint64(len(chunk.Data)))
|
||||
for _, decoder := range e.decoders {
|
||||
var decoderType detectorspb.DecoderType
|
||||
switch decoder.(type) {
|
||||
case *decoders.UTF8:
|
||||
decoderType = detectorspb.DecoderType_PLAIN
|
||||
case *decoders.Base64:
|
||||
decoderType = detectorspb.DecoderType_BASE64
|
||||
case *decoders.UTF16:
|
||||
decoderType = detectorspb.DecoderType_UTF16
|
||||
default:
|
||||
ctx.Logger().Info("unknown decoder type", "type", reflect.TypeOf(decoder).String())
|
||||
decoderType = detectorspb.DecoderType_UNKNOWN
|
||||
}
|
||||
|
||||
decoded := decoder.FromChunk(chunk)
|
||||
|
||||
if decoded == nil {
|
||||
ctx.Logger().V(4).Info("no decoder found for chunk", "chunk", chunk)
|
||||
continue
|
||||
}
|
||||
|
||||
for _, match := range e.ahoCorasickCore.MatchString(string(decoded.Chunk.Data)) {
|
||||
if !e.ahoCorasickCore.PopulateDetectorsByMatch(match, chunkSpecificDetectors) {
|
||||
continue
|
||||
}
|
||||
// build a map of all keywords that were matched in the chunk
|
||||
for _, m := range e.prefilter.MatchString(strings.ToLower(string(decoded.Data))) {
|
||||
matchedKeywords[strings.ToLower(m.MatchString())] = struct{}{}
|
||||
}
|
||||
|
||||
for k, detector := range chunkSpecificDetectors {
|
||||
decoded.Chunk.Verify = detector.ShouldVerify
|
||||
wgDetect.Add(1)
|
||||
e.detectableChunksChan <- detectableChunk{
|
||||
chunk: *decoded.Chunk,
|
||||
detector: detector,
|
||||
decoder: decoded.DecoderType,
|
||||
wgDoneFn: wgDetect.Done,
|
||||
for verify, detectorsSet := range e.detectors {
|
||||
for _, detector := range detectorsSet {
|
||||
chunkContainsKeyword := false
|
||||
for _, kw := range detector.Keywords() {
|
||||
if _, ok := matchedKeywords[strings.ToLower(kw)]; ok {
|
||||
chunkContainsKeyword = true
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if !chunkContainsKeyword {
|
||||
continue
|
||||
}
|
||||
|
||||
decoded.Verify = verify
|
||||
wgDetect.Add(1)
|
||||
e.detectableChunksChan <- detectableChunk{
|
||||
chunk: *decoded,
|
||||
detector: detector,
|
||||
decoder: decoderType,
|
||||
wgDoneFn: wgDetect.Done,
|
||||
}
|
||||
}
|
||||
delete(chunkSpecificDetectors, k)
|
||||
}
|
||||
}
|
||||
}
|
||||
atomic.AddUint64(&e.metrics.ChunksScanned, 1)
|
||||
}
|
||||
wgDetect.Wait()
|
||||
ctx.Logger().V(4).Info("finished scanning chunks")
|
||||
}
|
||||
|
||||
func (e *Engine) detectChunks(ctx context.Context) {
|
||||
|
||||
@@ -2,7 +2,6 @@ package engine
|
||||
|
||||
import (
|
||||
"os"
|
||||
"runtime"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
@@ -121,10 +120,9 @@ func BenchmarkGitEngine(b *testing.B) {
|
||||
defer cancel()
|
||||
|
||||
e, err := Start(ctx,
|
||||
WithConcurrency(uint8(runtime.NumCPU())),
|
||||
WithConcurrency(1),
|
||||
WithDecoders(decoders.DefaultDecoders()...),
|
||||
WithDetectors(false, DefaultDetectors()...),
|
||||
WithPrinter(new(discardPrinter)),
|
||||
)
|
||||
assert.Nil(b, err)
|
||||
|
||||
@@ -135,7 +133,6 @@ func BenchmarkGitEngine(b *testing.B) {
|
||||
}
|
||||
}()
|
||||
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
// TODO: this is measuring the time it takes to initialize the source
|
||||
// and not to do the full scan
|
||||
|
||||
@@ -1010,9 +1010,6 @@ const (
|
||||
DetectorType_SourcegraphCody DetectorType = 936
|
||||
DetectorType_Voiceflow DetectorType = 937
|
||||
DetectorType_Privacy DetectorType = 938
|
||||
DetectorType_IPInfo DetectorType = 939
|
||||
DetectorType_Ip2location DetectorType = 940
|
||||
DetectorType_Instamojo DetectorType = 941
|
||||
)
|
||||
|
||||
// Enum value maps for DetectorType.
|
||||
@@ -1953,9 +1950,6 @@ var (
|
||||
936: "SourcegraphCody",
|
||||
937: "Voiceflow",
|
||||
938: "Privacy",
|
||||
939: "IPInfo",
|
||||
940: "Ip2location",
|
||||
941: "Instamojo",
|
||||
}
|
||||
DetectorType_value = map[string]int32{
|
||||
"Alibaba": 0,
|
||||
@@ -2893,9 +2887,6 @@ var (
|
||||
"SourcegraphCody": 936,
|
||||
"Voiceflow": 937,
|
||||
"Privacy": 938,
|
||||
"IPInfo": 939,
|
||||
"Ip2location": 940,
|
||||
"Instamojo": 941,
|
||||
}
|
||||
)
|
||||
|
||||
@@ -3274,7 +3265,7 @@ var file_detectors_proto_rawDesc = []byte{
|
||||
0x44, 0x65, 0x63, 0x6f, 0x64, 0x65, 0x72, 0x54, 0x79, 0x70, 0x65, 0x12, 0x0b, 0x0a, 0x07, 0x55,
|
||||
0x4e, 0x4b, 0x4e, 0x4f, 0x57, 0x4e, 0x10, 0x00, 0x12, 0x09, 0x0a, 0x05, 0x50, 0x4c, 0x41, 0x49,
|
||||
0x4e, 0x10, 0x01, 0x12, 0x0a, 0x0a, 0x06, 0x42, 0x41, 0x53, 0x45, 0x36, 0x34, 0x10, 0x02, 0x12,
|
||||
0x09, 0x0a, 0x05, 0x55, 0x54, 0x46, 0x31, 0x36, 0x10, 0x03, 0x2a, 0xf3, 0x75, 0x0a, 0x0c, 0x44,
|
||||
0x09, 0x0a, 0x05, 0x55, 0x54, 0x46, 0x31, 0x36, 0x10, 0x03, 0x2a, 0xc4, 0x75, 0x0a, 0x0c, 0x44,
|
||||
0x65, 0x74, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x54, 0x79, 0x70, 0x65, 0x12, 0x0b, 0x0a, 0x07, 0x41,
|
||||
0x6c, 0x69, 0x62, 0x61, 0x62, 0x61, 0x10, 0x00, 0x12, 0x08, 0x0a, 0x04, 0x41, 0x4d, 0x51, 0x50,
|
||||
0x10, 0x01, 0x12, 0x07, 0x0a, 0x03, 0x41, 0x57, 0x53, 0x10, 0x02, 0x12, 0x09, 0x0a, 0x05, 0x41,
|
||||
@@ -4215,14 +4206,11 @@ var file_detectors_proto_rawDesc = []byte{
|
||||
0x0f, 0x53, 0x6f, 0x75, 0x72, 0x63, 0x65, 0x67, 0x72, 0x61, 0x70, 0x68, 0x43, 0x6f, 0x64, 0x79,
|
||||
0x10, 0xa8, 0x07, 0x12, 0x0e, 0x0a, 0x09, 0x56, 0x6f, 0x69, 0x63, 0x65, 0x66, 0x6c, 0x6f, 0x77,
|
||||
0x10, 0xa9, 0x07, 0x12, 0x0c, 0x0a, 0x07, 0x50, 0x72, 0x69, 0x76, 0x61, 0x63, 0x79, 0x10, 0xaa,
|
||||
0x07, 0x12, 0x0b, 0x0a, 0x06, 0x49, 0x50, 0x49, 0x6e, 0x66, 0x6f, 0x10, 0xab, 0x07, 0x12, 0x10,
|
||||
0x0a, 0x0b, 0x49, 0x70, 0x32, 0x6c, 0x6f, 0x63, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x10, 0xac, 0x07,
|
||||
0x12, 0x0e, 0x0a, 0x09, 0x49, 0x6e, 0x73, 0x74, 0x61, 0x6d, 0x6f, 0x6a, 0x6f, 0x10, 0xad, 0x07,
|
||||
0x42, 0x3d, 0x5a, 0x3b, 0x67, 0x69, 0x74, 0x68, 0x75, 0x62, 0x2e, 0x63, 0x6f, 0x6d, 0x2f, 0x74,
|
||||
0x72, 0x75, 0x66, 0x66, 0x6c, 0x65, 0x73, 0x65, 0x63, 0x75, 0x72, 0x69, 0x74, 0x79, 0x2f, 0x74,
|
||||
0x72, 0x75, 0x66, 0x66, 0x6c, 0x65, 0x68, 0x6f, 0x67, 0x2f, 0x76, 0x33, 0x2f, 0x70, 0x6b, 0x67,
|
||||
0x2f, 0x70, 0x62, 0x2f, 0x64, 0x65, 0x74, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x73, 0x70, 0x62, 0x62,
|
||||
0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33,
|
||||
0x07, 0x42, 0x3d, 0x5a, 0x3b, 0x67, 0x69, 0x74, 0x68, 0x75, 0x62, 0x2e, 0x63, 0x6f, 0x6d, 0x2f,
|
||||
0x74, 0x72, 0x75, 0x66, 0x66, 0x6c, 0x65, 0x73, 0x65, 0x63, 0x75, 0x72, 0x69, 0x74, 0x79, 0x2f,
|
||||
0x74, 0x72, 0x75, 0x66, 0x66, 0x6c, 0x65, 0x68, 0x6f, 0x67, 0x2f, 0x76, 0x33, 0x2f, 0x70, 0x6b,
|
||||
0x67, 0x2f, 0x70, 0x62, 0x2f, 0x64, 0x65, 0x74, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x73, 0x70, 0x62,
|
||||
0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33,
|
||||
}
|
||||
|
||||
var (
|
||||
|
||||
@@ -1,218 +0,0 @@
|
||||
package sources
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
lru "github.com/hashicorp/golang-lru/v2"
|
||||
"github.com/trufflesecurity/trufflehog/v3/pkg/context"
|
||||
)
|
||||
|
||||
// UnitHook implements JobProgressHook for tracking the progress of each
|
||||
// individual unit.
|
||||
type UnitHook struct {
|
||||
metrics *lru.Cache[string, *UnitMetrics]
|
||||
mu sync.Mutex
|
||||
NoopHook
|
||||
}
|
||||
|
||||
type UnitHookOpt func(*UnitHook)
|
||||
|
||||
func WithUnitHookCache(cache *lru.Cache[string, *UnitMetrics]) UnitHookOpt {
|
||||
return func(hook *UnitHook) { hook.metrics = cache }
|
||||
}
|
||||
|
||||
func NewUnitHook(ctx context.Context, opts ...UnitHookOpt) *UnitHook {
|
||||
// lru.NewWithEvict can only fail if the size is < 0.
|
||||
cache, _ := lru.NewWithEvict(1024, func(key string, value *UnitMetrics) {
|
||||
if value.handled {
|
||||
return
|
||||
}
|
||||
ctx.Logger().Error(fmt.Errorf("eviction"), "dropping unit metric",
|
||||
"id", key,
|
||||
"metric", value,
|
||||
)
|
||||
})
|
||||
hook := UnitHook{metrics: cache}
|
||||
for _, opt := range opts {
|
||||
opt(&hook)
|
||||
}
|
||||
return &hook
|
||||
}
|
||||
|
||||
// id is a helper method to generate an ID for the given job and unit.
|
||||
func (u *UnitHook) id(ref JobProgressRef, unit SourceUnit) string {
|
||||
unitID := ""
|
||||
if unit != nil {
|
||||
unitID = unit.SourceUnitID()
|
||||
}
|
||||
return fmt.Sprintf("%d/%d/%s", ref.SourceID, ref.JobID, unitID)
|
||||
}
|
||||
|
||||
func (u *UnitHook) StartUnitChunking(ref JobProgressRef, unit SourceUnit, start time.Time) {
|
||||
id := u.id(ref, unit)
|
||||
u.mu.Lock()
|
||||
defer u.mu.Unlock()
|
||||
|
||||
u.metrics.Add(id, &UnitMetrics{
|
||||
Unit: unit,
|
||||
Parent: ref,
|
||||
StartTime: start,
|
||||
})
|
||||
}
|
||||
|
||||
func (u *UnitHook) EndUnitChunking(ref JobProgressRef, unit SourceUnit, end time.Time) {
|
||||
id := u.id(ref, unit)
|
||||
u.mu.Lock()
|
||||
defer u.mu.Unlock()
|
||||
|
||||
metrics, ok := u.metrics.Get(id)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
metrics.EndTime = end
|
||||
}
|
||||
|
||||
func (u *UnitHook) ReportChunk(ref JobProgressRef, unit SourceUnit, chunk *Chunk) {
|
||||
id := u.id(ref, unit)
|
||||
u.mu.Lock()
|
||||
defer u.mu.Unlock()
|
||||
|
||||
metrics, ok := u.metrics.Get(id)
|
||||
if !ok && unit != nil {
|
||||
// The unit has been evicted.
|
||||
return
|
||||
} else if !ok && unit == nil {
|
||||
// This is a chunk from a non-unit source.
|
||||
metrics = &UnitMetrics{
|
||||
Unit: nil,
|
||||
Parent: ref,
|
||||
StartTime: ref.Snapshot().StartTime,
|
||||
}
|
||||
u.metrics.Add(id, metrics)
|
||||
}
|
||||
metrics.TotalChunks++
|
||||
metrics.TotalBytes += uint64(len(chunk.Data))
|
||||
}
|
||||
|
||||
func (u *UnitHook) ReportError(ref JobProgressRef, err error) {
|
||||
u.mu.Lock()
|
||||
defer u.mu.Unlock()
|
||||
|
||||
// Always add the error to the nil unit if it exists.
|
||||
if metrics, ok := u.metrics.Get(u.id(ref, nil)); ok {
|
||||
metrics.Errors = append(metrics.Errors, err)
|
||||
}
|
||||
|
||||
// Check if it's a ChunkError for a specific unit.
|
||||
var chunkErr ChunkError
|
||||
if !errors.As(err, &chunkErr) {
|
||||
return
|
||||
}
|
||||
id := u.id(ref, chunkErr.Unit)
|
||||
|
||||
metrics, ok := u.metrics.Get(id)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
metrics.Errors = append(metrics.Errors, err)
|
||||
}
|
||||
|
||||
func (u *UnitHook) Finish(ref JobProgressRef) {
|
||||
u.mu.Lock()
|
||||
defer u.mu.Unlock()
|
||||
// Clear out any metrics on this job. This covers the case for the
|
||||
// source running without unit support.
|
||||
prefix := u.id(ref, nil)
|
||||
for _, id := range u.metrics.Keys() {
|
||||
if !strings.HasPrefix(id, prefix) {
|
||||
continue
|
||||
}
|
||||
metric, ok := u.metrics.Get(id)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
// If the unit is nil, the source does not support units.
|
||||
// Use the overall job metrics instead.
|
||||
if metric.Unit == nil {
|
||||
snap := ref.Snapshot()
|
||||
metric.StartTime = snap.StartTime
|
||||
metric.EndTime = snap.EndTime
|
||||
metric.Errors = snap.Errors
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// UnitMetrics gets all the currently active or newly finished metrics for this
|
||||
// job. If a unit returned from this method has finished, it will be removed
|
||||
// from the cache and no longer returned in successive calls to UnitMetrics().
|
||||
func (u *UnitHook) UnitMetrics() []UnitMetrics {
|
||||
u.mu.Lock()
|
||||
defer u.mu.Unlock()
|
||||
output := make([]UnitMetrics, 0, u.metrics.Len())
|
||||
for _, id := range u.metrics.Keys() {
|
||||
metric, ok := u.metrics.Get(id)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
output = append(output, *metric)
|
||||
if metric.IsFinished() {
|
||||
metric.handled = true
|
||||
u.metrics.Remove(id)
|
||||
}
|
||||
}
|
||||
return output
|
||||
}
|
||||
|
||||
type UnitMetrics struct {
|
||||
Unit SourceUnit
|
||||
Parent JobProgressRef
|
||||
// Start and end time for chunking this unit.
|
||||
StartTime time.Time
|
||||
EndTime time.Time
|
||||
// Total number of chunks produced from this unit.
|
||||
TotalChunks uint64
|
||||
// Total number of bytes produced from this unit.
|
||||
TotalBytes uint64
|
||||
// All errors encountered by this unit.
|
||||
Errors []error
|
||||
// Flag to mark that these metrics were intentionally evicted from
|
||||
// the cache.
|
||||
handled bool
|
||||
}
|
||||
|
||||
func (u UnitMetrics) IsFinished() bool {
|
||||
return !u.EndTime.IsZero()
|
||||
}
|
||||
|
||||
// ElapsedTime is a convenience method that provides the elapsed time the job
|
||||
// has been running. If it hasn't started yet, 0 is returned. If it has
|
||||
// finished, the total time is returned.
|
||||
func (u UnitMetrics) ElapsedTime() time.Duration {
|
||||
if u.StartTime.IsZero() {
|
||||
return 0
|
||||
}
|
||||
if u.EndTime.IsZero() {
|
||||
return time.Since(u.StartTime)
|
||||
}
|
||||
return u.EndTime.Sub(u.StartTime)
|
||||
}
|
||||
|
||||
// NoopHook implements JobProgressHook by doing nothing. This is useful for
|
||||
// embedding in other structs to overwrite only the methods of the interface
|
||||
// that you care about.
|
||||
type NoopHook struct{}
|
||||
|
||||
func (NoopHook) Start(JobProgressRef, time.Time) {}
|
||||
func (NoopHook) End(JobProgressRef, time.Time) {}
|
||||
func (NoopHook) StartEnumerating(JobProgressRef, time.Time) {}
|
||||
func (NoopHook) EndEnumerating(JobProgressRef, time.Time) {}
|
||||
func (NoopHook) StartUnitChunking(JobProgressRef, SourceUnit, time.Time) {}
|
||||
func (NoopHook) EndUnitChunking(JobProgressRef, SourceUnit, time.Time) {}
|
||||
func (NoopHook) ReportError(JobProgressRef, error) {}
|
||||
func (NoopHook) ReportUnit(JobProgressRef, SourceUnit) {}
|
||||
func (NoopHook) ReportChunk(JobProgressRef, SourceUnit, *Chunk) {}
|
||||
func (NoopHook) Finish(JobProgressRef) {}
|
||||
@@ -3,10 +3,8 @@ package sources
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"sort"
|
||||
"testing"
|
||||
|
||||
lru "github.com/hashicorp/golang-lru/v2"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"google.golang.org/protobuf/types/known/anypb"
|
||||
|
||||
@@ -202,8 +200,6 @@ func (c *unitChunker) ChunkUnit(ctx context.Context, unit SourceUnit, rep ChunkR
|
||||
if err := rep.ChunkErr(ctx, fmt.Errorf(step.err)); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if step.output == "" {
|
||||
continue
|
||||
}
|
||||
if err := rep.ChunkOk(ctx, Chunk{Data: []byte(step.output)}); err != nil {
|
||||
@@ -314,114 +310,3 @@ func TestSourceManagerAvailableCapacity(t *testing.T) {
|
||||
<-ref.Done() // Wait for the job to finish.
|
||||
assert.Equal(t, 1337, mgr.AvailableCapacity())
|
||||
}
|
||||
|
||||
func TestSourceManagerUnitHook(t *testing.T) {
|
||||
hook := NewUnitHook(context.TODO())
|
||||
|
||||
input := []unitChunk{
|
||||
{unit: "one", output: "bar"},
|
||||
{unit: "two", err: "oh no"},
|
||||
{unit: "three", err: "not again"},
|
||||
}
|
||||
mgr := NewManager(
|
||||
WithBufferedOutput(8),
|
||||
WithSourceUnits(), WithConcurrentUnits(1),
|
||||
WithReportHook(hook),
|
||||
)
|
||||
source, err := buildDummy(&unitChunker{input})
|
||||
assert.NoError(t, err)
|
||||
ref, err := mgr.Run(context.Background(), "dummy", source)
|
||||
assert.NoError(t, err)
|
||||
<-ref.Done()
|
||||
|
||||
metrics := hook.UnitMetrics()
|
||||
assert.Equal(t, 3, len(metrics))
|
||||
sort.Slice(metrics, func(i, j int) bool {
|
||||
return metrics[i].EndTime.Before(metrics[j].EndTime)
|
||||
})
|
||||
m0, m1, m2 := metrics[0], metrics[1], metrics[2]
|
||||
|
||||
assert.Equal(t, "one", m0.Unit.SourceUnitID())
|
||||
assert.Equal(t, uint64(1), m0.TotalChunks)
|
||||
assert.Equal(t, uint64(3), m0.TotalBytes)
|
||||
assert.NotZero(t, m0.StartTime)
|
||||
assert.NotZero(t, m0.EndTime)
|
||||
assert.NotZero(t, m0.ElapsedTime())
|
||||
assert.Equal(t, 0, len(m0.Errors))
|
||||
|
||||
assert.Equal(t, "two", m1.Unit.SourceUnitID())
|
||||
assert.Equal(t, uint64(0), m1.TotalChunks)
|
||||
assert.Equal(t, uint64(0), m1.TotalBytes)
|
||||
assert.NotZero(t, m1.StartTime)
|
||||
assert.NotZero(t, m1.EndTime)
|
||||
assert.NotZero(t, m1.ElapsedTime())
|
||||
assert.Equal(t, 1, len(m1.Errors))
|
||||
|
||||
assert.Equal(t, "three", m2.Unit.SourceUnitID())
|
||||
assert.Equal(t, uint64(0), m2.TotalChunks)
|
||||
assert.Equal(t, uint64(0), m2.TotalBytes)
|
||||
assert.NotZero(t, m2.StartTime)
|
||||
assert.NotZero(t, m2.EndTime)
|
||||
assert.NotZero(t, m2.ElapsedTime())
|
||||
assert.Equal(t, 1, len(m2.Errors))
|
||||
}
|
||||
|
||||
// TestSourceManagerUnitHookNoBlock tests that the UnitHook drops metrics if
|
||||
// they aren't handled fast enough.
|
||||
func TestSourceManagerUnitHookNoBlock(t *testing.T) {
|
||||
var evictedKeys []string
|
||||
cache, _ := lru.NewWithEvict(1, func(key string, _ *UnitMetrics) {
|
||||
evictedKeys = append(evictedKeys, key)
|
||||
})
|
||||
hook := NewUnitHook(context.TODO(), WithUnitHookCache(cache))
|
||||
|
||||
input := []unitChunk{
|
||||
{unit: "one", output: "bar"},
|
||||
{unit: "two", err: "oh no"},
|
||||
{unit: "three", err: "not again"},
|
||||
}
|
||||
mgr := NewManager(
|
||||
WithBufferedOutput(8),
|
||||
WithSourceUnits(), WithConcurrentUnits(1),
|
||||
WithReportHook(hook),
|
||||
)
|
||||
source, err := buildDummy(&unitChunker{input})
|
||||
assert.NoError(t, err)
|
||||
ref, err := mgr.Run(context.Background(), "dummy", source)
|
||||
assert.NoError(t, err)
|
||||
<-ref.Done()
|
||||
|
||||
assert.Equal(t, 2, len(evictedKeys))
|
||||
metrics := hook.UnitMetrics()
|
||||
assert.Equal(t, 1, len(metrics))
|
||||
assert.Equal(t, "three", metrics[0].Unit.SourceUnitID())
|
||||
}
|
||||
|
||||
// TestSourceManagerUnitHookNoUnits tests whether the UnitHook works for
|
||||
// sources that don't support units.
|
||||
func TestSourceManagerUnitHookNoUnits(t *testing.T) {
|
||||
hook := NewUnitHook(context.TODO())
|
||||
|
||||
mgr := NewManager(
|
||||
WithBufferedOutput(8),
|
||||
WithReportHook(hook),
|
||||
)
|
||||
source, err := buildDummy(&counterChunker{count: 5})
|
||||
assert.NoError(t, err)
|
||||
|
||||
ref, err := mgr.Run(context.Background(), "dummy", source)
|
||||
assert.NoError(t, err)
|
||||
<-ref.Done()
|
||||
|
||||
metrics := hook.UnitMetrics()
|
||||
assert.Equal(t, 1, len(metrics))
|
||||
|
||||
m := metrics[0]
|
||||
assert.Equal(t, nil, m.Unit)
|
||||
assert.Equal(t, uint64(5), m.TotalChunks)
|
||||
assert.Equal(t, uint64(5), m.TotalBytes)
|
||||
assert.NotZero(t, m.StartTime)
|
||||
assert.NotZero(t, m.EndTime)
|
||||
assert.NotZero(t, m.ElapsedTime())
|
||||
assert.Equal(t, 0, len(m.Errors))
|
||||
}
|
||||
|
||||
@@ -947,9 +947,6 @@ enum DetectorType {
|
||||
SourcegraphCody = 936;
|
||||
Voiceflow = 937;
|
||||
Privacy = 938;
|
||||
IPInfo = 939;
|
||||
Ip2location = 940;
|
||||
Instamojo = 941;
|
||||
}
|
||||
|
||||
message Result {
|
||||
|
||||
Reference in New Issue
Block a user