Merge pull request #6846 from Checkmarx/joaom/kics-1254

feat(engine): add a timeout to decode results
Checkmarx · Jan 18, 2024 · 3e5dcdb · 3e5dcdb
2 parents 70d1949 + 288614c
commit 3e5dcdb
Show file tree

Hide file tree

Showing 3 changed files with 152 additions and 42 deletions.
diff --git a/pkg/engine/inspector.go b/pkg/engine/inspector.go
@@ -212,6 +212,7 @@ func (c *Inspector) Inspect(
 	}
 
 	queries := c.getQueriesByPlat(platforms)
+
 	for i, queryMeta := range queries {
 		currentQuery <- 1
 
@@ -240,6 +241,7 @@ func (c *Inspector) Inspect(
 		vuls, err := c.doRun(queryContext)
 
 		if err != nil {
+			fmt.Println()
 			sentryReport.ReportSentry(&sentryReport.Report{
 				Message:  fmt.Sprintf("Inspector. query executed with error, query=%s", query.Metadata.Query),
 				Err:      err,
@@ -260,7 +262,6 @@ func (c *Inspector) Inspect(
 
 		c.tracker.TrackQueryExecution(query.Metadata.Aggregation)
 	}
-
 	return vulnerabilities, nil
 }
 
@@ -343,11 +344,16 @@ func (c *Inspector) doRun(ctx *QueryContext) (vulns []model.Vulnerability, err e
 		Str("scanID", ctx.scanID).
 		Msgf("Inspector executed with result %+v, query=%s", results, ctx.Query.Metadata.Query)
 
-	return c.DecodeQueryResults(ctx, results)
+	timeoutCtxToDecode, cancelDecode := context.WithTimeout(ctx.Ctx, c.queryExecTimeout)
+	defer cancelDecode()
+	return c.DecodeQueryResults(ctx, timeoutCtxToDecode, results)
 }
 
 // DecodeQueryResults decodes the results into []model.Vulnerability
-func (c *Inspector) DecodeQueryResults(ctx *QueryContext, results rego.ResultSet) ([]model.Vulnerability, error) {
+func (c *Inspector) DecodeQueryResults(
+	ctx *QueryContext,
+	ctxTimeout context.Context,
+	results rego.ResultSet) ([]model.Vulnerability, error) {
 	if len(results) == 0 {
 		return nil, ErrNoResult
 	}
@@ -366,56 +372,81 @@ func (c *Inspector) DecodeQueryResults(ctx *QueryContext, results rego.ResultSet
 
 	vulnerabilities := make([]model.Vulnerability, 0, len(queryResultItems))
 	failedDetectLine := false
+	timeOut := false
 	for _, queryResultItem := range queryResultItems {
-		vulnerability, err := c.vb(ctx, c.tracker, queryResultItem, c.detector)
-		if err != nil && err.Error() == ErrNoResult.Error() {
-			// Ignoring bad results
-			continue
+		select {
+		case <-ctxTimeout.Done():
+			timeOut = true
+			break
+		default:
+			vulnerability, aux := getVulnerabilitiesFromQuery(ctx, c, queryResultItem)
+			if aux {
+				failedDetectLine = aux
+			}
+			if vulnerability != nil && !aux {
+				vulnerabilities = append(vulnerabilities, *vulnerability)
+			}
 		}
-		if err != nil {
-			sentryReport.ReportSentry(&sentryReport.Report{
-				Message:  fmt.Sprintf("Inspector can't save vulnerability, query=%s", ctx.Query.Metadata.Query),
-				Err:      err,
-				Location: "func decodeQueryResults()",
-				Platform: ctx.Query.Metadata.Platform,
-				Metadata: ctx.Query.Metadata.Metadata,
-				Query:    ctx.Query.Metadata.Query,
-			}, true)
+	}
 
-			if _, ok := c.failedQueries[ctx.Query.Metadata.Query]; !ok {
-				c.failedQueries[ctx.Query.Metadata.Query] = err
-			}
+	if timeOut {
+		fmt.Println()
+		log.Err(ctxTimeout.Err()).Msgf(
+			"Timeout processing the results of the query: %s %s",
+			ctx.Query.Metadata.Platform,
+			ctx.Query.Metadata.Query)
+	}
 
-			continue
-		}
-		file := ctx.Files[vulnerability.FileID]
-		if ShouldSkipVulnerability(file.Commands, vulnerability.QueryID) {
-			log.Debug().Msgf("Skipping vulnerability in file %s for query '%s':%s", file.FilePath, vulnerability.QueryName, vulnerability.QueryID)
-			continue
-		}
+	if failedDetectLine {
+		c.tracker.FailedDetectLine()
+	}
 
-		if vulnerability.Line == UndetectedVulnerabilityLine {
-			failedDetectLine = true
-		}
+	return vulnerabilities, nil
+}
 
-		if _, ok := c.excludeResults[vulnerability.SimilarityID]; ok {
-			log.Debug().
-				Msgf("Excluding result SimilarityID: %s", vulnerability.SimilarityID)
-			continue
-		} else if checkComment(vulnerability.Line, file.LinesIgnore) {
-			log.Debug().
-				Msgf("Excluding result Comment: %s", vulnerability.SimilarityID)
-			continue
+func getVulnerabilitiesFromQuery(ctx *QueryContext, c *Inspector, queryResultItem interface{}) (*model.Vulnerability, bool) {
+	vulnerability, err := c.vb(ctx, c.tracker, queryResultItem, c.detector)
+	if err != nil && err.Error() == ErrNoResult.Error() {
+		// Ignoring bad results
+		return nil, false
+	}
+	if err != nil {
+		sentryReport.ReportSentry(&sentryReport.Report{
+			Message:  fmt.Sprintf("Inspector can't save vulnerability, query=%s", ctx.Query.Metadata.Query),
+			Err:      err,
+			Location: "func decodeQueryResults()",
+			Platform: ctx.Query.Metadata.Platform,
+			Metadata: ctx.Query.Metadata.Metadata,
+			Query:    ctx.Query.Metadata.Query,
+		}, true)
+
+		if _, ok := c.failedQueries[ctx.Query.Metadata.Query]; !ok {
+			c.failedQueries[ctx.Query.Metadata.Query] = err
 		}
 
-		vulnerabilities = append(vulnerabilities, *vulnerability)
+		return nil, false
+	}
+	file := ctx.Files[vulnerability.FileID]
+	if ShouldSkipVulnerability(file.Commands, vulnerability.QueryID) {
+		log.Debug().Msgf("Skipping vulnerability in file %s for query '%s':%s", file.FilePath, vulnerability.QueryName, vulnerability.QueryID)
+		return nil, false
 	}
 
-	if failedDetectLine {
-		c.tracker.FailedDetectLine()
+	if vulnerability.Line == UndetectedVulnerabilityLine {
+		return nil, true
 	}
 
-	return vulnerabilities, nil
+	if _, ok := c.excludeResults[vulnerability.SimilarityID]; ok {
+		log.Debug().
+			Msgf("Excluding result SimilarityID: %s", vulnerability.SimilarityID)
+		return nil, false
+	} else if checkComment(vulnerability.Line, file.LinesIgnore) {
+		log.Debug().
+			Msgf("Excluding result Comment: %s", vulnerability.SimilarityID)
+		return nil, false
+	}
+
+	return vulnerability, false
 }
 
 // checkComment checks if the vulnerability should be skipped from comment

diff --git a/pkg/engine/inspector_test.go b/pkg/engine/inspector_test.go
@@ -3,6 +3,8 @@ package engine
 import (
 	"context"
 	"fmt"
+	"github.com/open-policy-agent/opa/rego"
+	"github.com/stretchr/testify/assert"
 	"io"
 	"os"
 	"path/filepath"
@@ -677,6 +679,81 @@ func TestShouldSkipFile(t *testing.T) {
 	}
 }
 
+func TestInspector_DecodeQueryResults(t *testing.T) {
+
+	//context
+	contextToUSe := context.Background()
+
+	//build inspector
+	c := newInspectorInstance(t, []string{})
+
+	type args struct {
+		queryContext QueryContext
+		regoResult   rego.ResultSet
+		timeDuration string
+	}
+	tests := []struct {
+		name     string
+		args     args
+		expected int
+	}{
+		{
+			name: "should_not_fail_when_timeout",
+			args: args{
+				queryContext: newQueryContext(contextToUSe),
+				regoResult:   newResultset(),
+				timeDuration: "0s",
+			},
+			expected: 0,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			//create a context with 0 second to timeout
+			timeoutDuration, _ := time.ParseDuration(tt.args.timeDuration)
+			myCtxTimeOut, _ := context.WithTimeout(contextToUSe, timeoutDuration)
+			result, err := c.DecodeQueryResults(&tt.args.queryContext, myCtxTimeOut, tt.args.regoResult)
+			assert.Nil(t, err, "Error not as expected")
+			assert.Equal(t, 0, len(result), "Array size is not as expected")
+		})
+	}
+}
+
+func newResultset() rego.ResultSet {
+	myValue := make(map[string]interface{})
+	myValue["documentId"] = "3a3be8f7-896e-4ef8-9db3-d6c19e60510b"
+	myValue["issueType"] = "IncorrectValue"
+	myValue["keyActualValue"] = "COPY --from referencesthe current FROM alias"
+	myValue["keyExpectedValue"] = "COPY --from should not references the current FROM alias"
+	myValue["searchKey"] = "{{ADD ${JAR_FILE} app.jar}}"
+
+	myBinding := make([]interface{}, 1)
+	myBinding[0] = myValue
+
+	myresult := rego.Result{
+		Bindings: map[string]interface{}{
+			"result": myBinding,
+		},
+	}
+	myResultSet := rego.ResultSet{myresult}
+	return myResultSet
+}
+
+func newQueryContext(ctx context.Context) QueryContext {
+	queryMetadata := model.QueryMetadata{
+		Platform: "myPlatform",
+		Query:    "myQuery"}
+	myQuery := PreparedQuery{
+		Metadata: queryMetadata,
+	}
+	queryContext := QueryContext{
+		Ctx:   ctx,
+		Query: &myQuery,
+	}
+	return queryContext
+}
+
 func newInspectorInstance(t *testing.T, queryPath []string) *Inspector {
 	querySource := source.NewFilesystemSource(queryPath, []string{""}, []string{""}, filepath.FromSlash("./assets/libraries"), true)
 	var vb = func(ctx *QueryContext, tracker Tracker, v interface{},

diff --git a/pkg/remediation/scan.go b/pkg/remediation/scan.go
@@ -177,7 +177,9 @@ func runQuery(r *runQueryInfo) []model.Vulnerability {
 		Files:         r.files.ToMap(),
 	}
 
-	decoded, err := r.inspector.DecodeQueryResults(queryCtx, results)
+	timeoutCtxToDecode, cancelDecode := context.WithTimeout(context.Background(), queryExecTimeout)
+	defer cancelDecode()
+	decoded, err := r.inspector.DecodeQueryResults(queryCtx, timeoutCtxToDecode, results)
 
 	if err != nil {
 		log.Err(err)