From 827ec8eac229527f0058df12c0bf9fdbc53c71d6 Mon Sep 17 00:00:00 2001
From: Nikita Pivkin <nikita.pivkin@smartforce.io>
Date: Fri, 23 Aug 2024 13:43:10 +0600
Subject: [PATCH] feat(misconf): ignore duplicate checks (#7317)

Signed-off-by: nikpivkin <nikita.pivkin@smartforce.io>
---
 pkg/iac/rego/embed.go      | 28 ++++++++++++++++++----
 pkg/iac/rego/embed_test.go | 48 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 71 insertions(+), 5 deletions(-)

diff --git a/pkg/iac/rego/embed.go b/pkg/iac/rego/embed.go
index 6f542d9a0b2b..e582508a396b 100644
--- a/pkg/iac/rego/embed.go
+++ b/pkg/iac/rego/embed.go
@@ -2,6 +2,7 @@ package rego
 
 import (
 	"context"
+	"fmt"
 	"io/fs"
 	"path/filepath"
 	"strings"
@@ -10,6 +11,7 @@ import (
 
 	checks "github.com/aquasecurity/trivy-checks"
 	"github.com/aquasecurity/trivy/pkg/iac/rules"
+	"github.com/aquasecurity/trivy/pkg/log"
 )
 
 func init() {
@@ -47,17 +49,34 @@ func RegisterRegoRules(modules map[string]*ast.Module) {
 	}
 
 	retriever := NewMetadataRetriever(compiler)
+	regoCheckIDs := make(map[string]struct{})
+
 	for _, module := range modules {
 		metadata, err := retriever.RetrieveMetadata(ctx, module)
 		if err != nil {
+			log.Warn("Failed to retrieve metadata", log.String("avdid", metadata.AVDID), log.Err(err))
 			continue
 		}
+
 		if metadata.AVDID == "" {
+			log.Warn("Check ID is empty", log.FilePath(module.Package.Location.File))
 			continue
 		}
-		rules.Register(
-			metadata.ToRule(),
-		)
+
+		if !metadata.Deprecated {
+			regoCheckIDs[metadata.AVDID] = struct{}{}
+		}
+
+		rules.Register(metadata.ToRule())
+	}
+
+	for _, check := range rules.GetRegistered() {
+		if !check.Deprecated && check.CanCheck() {
+			if _, exists := regoCheckIDs[check.AVDID]; exists {
+				log.Warn("Ignore duplicate Go check", log.String("avdid", check.AVDID))
+				rules.Deregister(check)
+			}
+		}
 	}
 }
 
@@ -95,8 +114,7 @@ func LoadPoliciesFromDirs(target fs.FS, paths ...string) (map[string]*ast.Module
 				ProcessAnnotation: true,
 			})
 			if err != nil {
-				// s.debug.Log("Failed to load module: %s, err: %s", filepath.ToSlash(path), err.Error())
-				return err
+				return fmt.Errorf("failed to parse Rego module: %w", err)
 			}
 			modules[path] = module
 			return nil
diff --git a/pkg/iac/rego/embed_test.go b/pkg/iac/rego/embed_test.go
index 15001119a934..5b6368dec2eb 100644
--- a/pkg/iac/rego/embed_test.go
+++ b/pkg/iac/rego/embed_test.go
@@ -2,6 +2,7 @@ package rego
 
 import (
 	"testing"
+	"testing/fstest"
 
 	"github.com/open-policy-agent/opa/ast"
 	"github.com/stretchr/testify/assert"
@@ -10,6 +11,7 @@ import (
 	checks "github.com/aquasecurity/trivy-checks"
 	"github.com/aquasecurity/trivy/pkg/iac/rules"
 	"github.com/aquasecurity/trivy/pkg/iac/scan"
+	"github.com/aquasecurity/trivy/pkg/iac/state"
 )
 
 func Test_EmbeddedLoading(t *testing.T) {
@@ -204,3 +206,49 @@ deny[res]{
 		})
 	}
 }
+
+func Test_IgnoreDuplicateChecks(t *testing.T) {
+	rules.Reset()
+
+	r := scan.Rule{
+		AVDID: "TEST001",
+		Check: func(s *state.State) (results scan.Results) {
+			for _, bucket := range s.AWS.S3.Buckets {
+				if bucket.Name.Value() == "evil" {
+					results.Add("Bucket name should not be evil", bucket.Name)
+				}
+			}
+			return
+		},
+	}
+	reg := rules.Register(r)
+	defer rules.Deregister(reg)
+
+	fsys := fstest.MapFS{
+		"test.rego": &fstest.MapFile{
+			Data: []byte(`
+# METADATA
+# title: "Test rego"
+# scope: package
+# schemas:
+# - input: schema["cloud"]
+# custom:
+#   avd_id: TEST001
+#   severity: LOW
+package user.test001
+
+deny[res] {
+	res := result.new("test", {})
+}
+`),
+		},
+	}
+
+	modules, err := LoadPoliciesFromDirs(fsys, ".")
+	require.NoError(t, err)
+
+	RegisterRegoRules(modules)
+	registered := rules.GetRegistered()
+	assert.Len(t, registered, 1)
+	assert.Equal(t, "TEST001", registered[0].AVDID)
+}