[bug] - Create a new context with timeout per request (#3163)

* Create a new context with timeout per request * match timeout * use context timeout * reduce timeout
trufflesecurity · Aug 2, 2024 · 0a3451a · 0a3451a
1 parent f939572
commit 0a3451a
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 10 deletions.
diff --git a/pkg/detectors/privatekey/ssh.go b/pkg/detectors/privatekey/ssh.go
@@ -7,7 +7,6 @@ import (
 	"fmt"
 	"net"
 	"strings"
-	"time"
 
 	"golang.org/x/crypto/ssh"
 )
@@ -35,8 +34,7 @@ func firstResponseFromSSH(ctx context.Context, parsedKey any, username, hostport
 
 	// Verify the server fingerprint to ensure that there is no MITM replay attack
 	config := &ssh.ClientConfig{
-		Timeout: 5 * time.Second,
-		User:    username,
+		User: username,
 		Auth: []ssh.AuthMethod{
 			ssh.PublicKeys(signer),
 		},
@@ -87,7 +85,7 @@ func firstResponseFromSSH(ctx context.Context, parsedKey any, username, hostport
 }
 
 func sshDialWithContext(ctx context.Context, network, addr string, config *ssh.ClientConfig) (*ssh.Client, error) {
-	d := net.Dialer{Timeout: config.Timeout}
+	d := net.Dialer{}
 	conn, err := d.DialContext(ctx, network, addr)
 	if err != nil {
 		return nil, fmt.Errorf("error dialing %s: %w", addr, err)

diff --git a/pkg/engine/engine.go b/pkg/engine/engine.go
@@ -882,11 +882,15 @@ func (e *Engine) verificationOverlapWorker(ctx context.Context) {
 			// DO NOT VERIFY at this stage of the pipeline.
 			matchedBytes := detector.Matches()
 			for _, match := range matchedBytes {
+				ctx, cancel := context.WithTimeout(ctx, time.Second*2)
 				results, err := detector.FromData(ctx, false, match)
-				ctx.Logger().Error(
-					err, "error finding results in chunk during verification overlap",
-					"detector", detector.Key.Type().String(),
-				)
+				cancel()
+				if err != nil {
+					ctx.Logger().Error(
+						err, "error finding results in chunk during verification overlap",
+						"detector", detector.Key.Type().String(),
+					)
+				}
 
 				if len(results) == 0 {
 					continue
@@ -980,9 +984,7 @@ func (e *Engine) detectChunk(ctx context.Context, data detectableChunk) {
 	if e.printAvgDetectorTime {
 		start = time.Now()
 	}
-	ctx, cancel := context.WithTimeout(ctx, time.Second*10)
 	defer common.Recover(ctx)
-	defer cancel()
 
 	isFalsePositive := detectors.GetFalsePositiveCheck(data.detector)
 
@@ -996,7 +998,10 @@ func (e *Engine) detectChunk(ctx context.Context, data detectableChunk) {
 	for _, matchBytes := range matches {
 		matchCount++
 		detectBytesPerMatch.Observe(float64(len(matchBytes)))
+
+		ctx, cancel := context.WithTimeout(ctx, time.Second*10)
 		results, err := data.detector.Detector.FromData(ctx, data.chunk.Verify, matchBytes)
+		cancel()
 		if err != nil {
 			ctx.Logger().Error(
 				err, "error finding results in chunk",