tracee-rules: add flag for partial evaluation (#979)

Signed-off-by: Simar <simar@linux.com>

tracee-rules/benchmark/signature/rego/signatures.go

@@ -19,9 +19,9 @@ var (
 )
 
 func NewCodeInjectionSignature() (types.Signature, error) {
-	return regosig.NewRegoSignature(codeInjectionRego, helpersRego)
+	return regosig.NewRegoSignature(false, codeInjectionRego, helpersRego)
 }
 
 func NewAntiDebuggingSignature() (types.Signature, error) {
-	return regosig.NewRegoSignature(antiDebuggingPtracemeRego, helpersRego)
+	return regosig.NewRegoSignature(false, antiDebuggingPtracemeRego, helpersRego)
 }

tracee-rules/main.go

@@ -59,7 +59,7 @@ func main() {
 				}()
 			}
 
-			sigs, err := getSignatures(c.String("rules-dir"), c.StringSlice("rules"))
+			sigs, err := getSignatures(c.Bool("rego-partial-eval"), c.String("rules-dir"), c.StringSlice("rules"))
 			if err != nil {
 				return err
 			}
@@ -113,6 +113,10 @@ func main() {
 				Name:  "rules-dir",
 				Usage: "directory where to search for rules in OPA (.rego) or Go plugin (.so) formats",
 			},
+			&cli.BoolFlag{
+				Name:  "rego-partial-eval",
+				Usage: "enable partial evaluation of rego rules",
+			},
 			&cli.BoolFlag{
 				Name:  "list",
 				Usage: "print all available rules",

tracee-rules/signature.go

@@ -19,7 +19,7 @@ import (
 //go:embed signatures/rego/helpers.rego
 var regoHelpersCode string
 
-func getSignatures(rulesDir string, rules []string) ([]types.Signature, error) {
+func getSignatures(partialEval bool, rulesDir string, rules []string) ([]types.Signature, error) {
 	if rulesDir == "" {
 		exePath, err := os.Executable()
 		if err != nil {
@@ -31,7 +31,7 @@ func getSignatures(rulesDir string, rules []string) ([]types.Signature, error) {
 	if err != nil {
 		return nil, err
 	}
-	opasigs, err := findRegoSigs(rulesDir)
+	opasigs, err := findRegoSigs(partialEval, rulesDir)
 	if err != nil {
 		return nil, err
 	}
@@ -79,7 +79,7 @@ func findGoSigs(dir string) ([]types.Signature, error) {
 	return res, nil
 }
 
-func findRegoSigs(dir string) ([]types.Signature, error) {
+func findRegoSigs(partialEval bool, dir string) ([]types.Signature, error) {
 	regoHelpers := []string{regoHelpersCode}
 	filepath.WalkDir(dir, func(path string, d fs.DirEntry, err error) error {
 		if err != nil {
@@ -119,7 +119,7 @@ func findRegoSigs(dir string) ([]types.Signature, error) {
 			log.Printf("error reading file %s: %v", path, err)
 			return nil
 		}
-		sig, err := regosig.NewRegoSignature(append(regoHelpers, string(regoCode))...)
+		sig, err := regosig.NewRegoSignature(partialEval, append(regoHelpers, string(regoCode))...)
 		if err != nil {
 			newlineOffset := bytes.Index(regoCode, []byte("\n"))
 			if newlineOffset == -1 {

tracee-rules/signature_test.go

@@ -15,7 +15,7 @@ import (
 )
 
 func Test_getSignatures(t *testing.T) {
-	sigs, err := getSignatures("signatures/rego", []string{"TRC-2"})
+	sigs, err := getSignatures(false, "signatures/rego", []string{"TRC-2"})
 	require.NoError(t, err)
 	require.Equal(t, 1, len(sigs))
 
@@ -77,7 +77,7 @@ func Test_findRegoSigs(t *testing.T) {
 	require.NoError(t, err)
 
 	// find rego signatures
-	sigs, err := findRegoSigs(testRoot)
+	sigs, err := findRegoSigs(false, testRoot)
 	require.NoError(t, err)
 
 	assert.Equal(t, len(sigs), 2)

tracee-rules/signatures/rego/regosig/traceerego.go

@@ -35,7 +35,7 @@ const queryMetadata string = "data.%s.__rego_metadoc__"
 const packageNameRegex string = `package\s.*`
 
 // NewRegoSignature creates a new RegoSignature with the provided rego code string
-func NewRegoSignature(regoCodes ...string) (types.Signature, error) {
+func NewRegoSignature(partialEval bool, regoCodes ...string) (types.Signature, error) {
 	var err error
 	res := RegoSignature{}
 	regoMap := make(map[string]string)
@@ -62,13 +62,30 @@ func NewRegoSignature(regoCodes ...string) (types.Signature, error) {
 		return nil, err
 	}
 
-	res.matchPQ, err = rego.New(
-		rego.Compiler(res.compiledRego),
-		rego.Query(fmt.Sprintf(queryMatch, pkgName)),
-	).PrepareForEval(context.TODO())
-	if err != nil {
-		return nil, err
+	ctx := context.Background()
+	if partialEval {
+		pr, err := rego.New(
+			rego.Compiler(res.compiledRego),
+			rego.Query(fmt.Sprintf(queryMatch, pkgName)),
+		).PartialResult(ctx)
+		if err != nil {
+			return nil, err
+		}
+
+		res.matchPQ, err = pr.Rego().PrepareForEval(ctx)
+		if err != nil {
+			return nil, err
+		}
+	} else {
+		res.matchPQ, err = rego.New(
+			rego.Compiler(res.compiledRego),
+			rego.Query(fmt.Sprintf(queryMatch, pkgName)),
+		).PrepareForEval(ctx)
+		if err != nil {
+			return nil, err
+		}
 	}
+
 	res.metadata, err = res.getMetadata(pkgName)
 	if err != nil {
 		return nil, err

tracee-rules/signatures/rego/regosig/traceerego_test.go

@@ -35,7 +35,7 @@ __rego_metadoc__ := {
 }
 `
 
-	sig, err := NewRegoSignature(testRegoMeta)
+	sig, err := NewRegoSignature(false, testRegoMeta)
 	if err != nil {
 		t.Error(err)
 	}
@@ -66,7 +66,7 @@ tracee_selected_events[eventSelector] {
 	}
 }
 `
-	sig, err := NewRegoSignature(testRegoSelectedEvents)
+	sig, err := NewRegoSignature(false, testRegoSelectedEvents)
 	if err != nil {
 		t.Error(err)
 	}
@@ -119,7 +119,7 @@ tracee_match {
 		},
 	}
 	for _, st := range sts {
-		sig, err := NewRegoSignature(testRegoBool)
+		sig, err := NewRegoSignature(false, testRegoBool)
 		if err != nil {
 			t.Error(err)
 		}
@@ -229,7 +229,7 @@ tracee_match = res {
 		},
 	}
 	for _, st := range sts {
-		sig, err := NewRegoSignature(testRegoObj)
+		sig, err := NewRegoSignature(false, testRegoObj)
 		if err != nil {
 			t.Error(err)
 		}
@@ -268,7 +268,7 @@ func TestNewRegoSignature(t *testing.T) {
 
 	// assert basic attributes
 	for i, rc := range testRegoCodes {
-		gotSig, err := NewRegoSignature(rc)
+		gotSig, err := NewRegoSignature(false, rc)
 		require.NoError(t, err)
 
 		gotMetadata, err := gotSig.GetMetadata()