From 0d7d97d131095ae9436a3cdab94194a57344bc6a Mon Sep 17 00:00:00 2001
From: Nikita Pivkin <nikita.pivkin@smartforce.io>
Date: Sat, 20 Apr 2024 10:00:18 +0700
Subject: [PATCH] refactor(misconf): simplify the retrieval of module
 annotations (#6528)

---
 pkg/iac/rego/metadata.go | 17 ++++++-----------
 1 file changed, 6 insertions(+), 11 deletions(-)

diff --git a/pkg/iac/rego/metadata.go b/pkg/iac/rego/metadata.go
index 141bb4070562..6d6996dd6828 100644
--- a/pkg/iac/rego/metadata.go
+++ b/pkg/iac/rego/metadata.go
@@ -8,6 +8,7 @@ import (
 	"github.com/mitchellh/mapstructure"
 	"github.com/open-policy-agent/opa/ast"
 	"github.com/open-policy-agent/opa/rego"
+	"github.com/samber/lo"
 
 	"github.com/aquasecurity/trivy/pkg/iac/framework"
 	"github.com/aquasecurity/trivy/pkg/iac/providers"
@@ -16,6 +17,8 @@ import (
 	iacTypes "github.com/aquasecurity/trivy/pkg/iac/types"
 )
 
+const annotationScopePackage = "package"
+
 type StaticMetadata struct {
 	ID                 string
 	AVDID              string
@@ -234,17 +237,9 @@ func NewMetadataRetriever(compiler *ast.Compiler) *MetadataRetriever {
 }
 
 func (m *MetadataRetriever) findPackageAnnotations(module *ast.Module) *ast.Annotations {
-	annotationSet := m.compiler.GetAnnotationSet()
-	if annotationSet == nil {
-		return nil
-	}
-	for _, annotation := range annotationSet.Flatten() {
-		if annotation.GetPackage().Path.String() != module.Package.Path.String() || annotation.Annotations.Scope != "package" {
-			continue
-		}
-		return annotation.Annotations
-	}
-	return nil
+	return lo.FindOrElse(module.Annotations, nil, func(a *ast.Annotations) bool {
+		return a.Scope == annotationScopePackage
+	})
 }
 
 func (m *MetadataRetriever) RetrieveMetadata(ctx context.Context, module *ast.Module, contents ...any) (*StaticMetadata, error) {