From 9450072ac24d581a7c68073be10b034d76b2b9fb Mon Sep 17 00:00:00 2001
From: Victor Nicolet <victornl@amazon.com>
Date: Thu, 24 Oct 2024 20:50:05 -0400
Subject: [PATCH] Scanning imported pacakges for annotations.

---
 analysis/annotations/annotations.go      | 19 ++++++------
 analysis/annotations/annotations_test.go |  2 +-
 analysis/dataflow/state.go               | 15 +++++++--
 internal/analysisutil/iterators.go       | 39 ++++++++++++++++++++++++
 4 files changed, 63 insertions(+), 12 deletions(-)
 create mode 100644 internal/analysisutil/iterators.go

diff --git a/analysis/annotations/annotations.go b/analysis/annotations/annotations.go
index e7180e80..1ceb4ca7 100644
--- a/analysis/annotations/annotations.go
+++ b/analysis/annotations/annotations.go
@@ -211,19 +211,20 @@ func (pa ProgramAnnotations) Iter(fx func(a Annotation)) {
 	for _, cst := range pa.Globals {
 		funcutil.Iter(cst, fx)
 	}
+	for _, cst := range pa.Positional {
+		fx(cst)
+	}
 }
 
 // CompleteFromSyntax takes a set of program annotations and adds additional non-ssa linked annotations
 // to the annotations
-func (pa ProgramAnnotations) CompleteFromSyntax(logger *config.LogGroup, pkgs []*packages.Package) {
-	for _, pkg := range pkgs {
-		for _, astFile := range pkg.Syntax {
-			pa.loadPackageDocAnnotations(astFile.Doc)
-			for _, comments := range astFile.Comments {
-				for _, comment := range comments.List {
-					if annotationContents := extractAnnotation(comment); annotationContents != nil {
-						pa.loadFileAnnotations(logger, annotationContents, pkg.Fset.Position(comment.Pos()))
-					}
+func (pa ProgramAnnotations) CompleteFromSyntax(logger *config.LogGroup, pkg *packages.Package) {
+	for _, astFile := range pkg.Syntax {
+		pa.loadPackageDocAnnotations(astFile.Doc)
+		for _, comments := range astFile.Comments {
+			for _, comment := range comments.List {
+				if annotationContents := extractAnnotation(comment); annotationContents != nil {
+					pa.loadFileAnnotations(logger, annotationContents, pkg.Fset.Position(comment.Pos()))
 				}
 			}
 		}
diff --git a/analysis/annotations/annotations_test.go b/analysis/annotations/annotations_test.go
index b45e1e1a..67651e0e 100644
--- a/analysis/annotations/annotations_test.go
+++ b/analysis/annotations/annotations_test.go
@@ -36,7 +36,7 @@ func TestLoadAnnotations(t *testing.T) {
 	testProgram.Prog.Build()
 	logger := config.NewLogGroup(testProgram.Config)
 	a, err := annotations.LoadAnnotations(logger, testProgram.Prog.AllPackages())
-	a.CompleteFromSyntax(logger, testProgram.Pkgs)
+	a.CompleteFromSyntax(logger, testProgram.Pkgs[0])
 	if err != nil {
 		t.Errorf("error loading annotations: %s", err)
 	}
diff --git a/analysis/dataflow/state.go b/analysis/dataflow/state.go
index dc788639..3e27a624 100644
--- a/analysis/dataflow/state.go
+++ b/analysis/dataflow/state.go
@@ -25,6 +25,7 @@ import (
 	"github.com/awslabs/ar-go-tools/analysis/config"
 	"github.com/awslabs/ar-go-tools/analysis/lang"
 	"github.com/awslabs/ar-go-tools/analysis/summaries"
+	"github.com/awslabs/ar-go-tools/internal/analysisutil"
 	"github.com/awslabs/ar-go-tools/internal/pointer"
 	"golang.org/x/tools/go/callgraph"
 	"golang.org/x/tools/go/callgraph/cha"
@@ -120,10 +121,20 @@ func NewAnalyzerState(p *ssa.Program, pkgs []*packages.Package, l *config.LogGro
 	steps []func(*AnalyzerState)) (*AnalyzerState, error) {
 	var allContracts []Contract
 
-	// Load annotations
+	// Load annotations byt scanning all package's syntax
 	pa, err := annotations.LoadAnnotations(l, p.AllPackages())
 	if pkgs != nil {
-		pa.CompleteFromSyntax(l, pkgs)
+		for _, pkg := range pkgs {
+			analysisutil.VisitPackages(pkg, func(p *packages.Package) bool {
+				// Don't scan stdlib for annotations!
+				if summaries.IsStdPackageName(p.Name) {
+					return false
+				}
+				l.Debugf("Scan %s for annotations.\n", p.PkgPath)
+				pa.CompleteFromSyntax(l, p)
+				return true
+			})
+		}
 	}
 
 	if err != nil {
diff --git a/internal/analysisutil/iterators.go b/internal/analysisutil/iterators.go
new file mode 100644
index 00000000..2fcddc16
--- /dev/null
+++ b/internal/analysisutil/iterators.go
@@ -0,0 +1,39 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package analysisutil
+
+import (
+	"golang.org/x/tools/go/packages"
+)
+
+// VisitPackages calls f recursively on the root package provided, and then all the imports provided f returns
+// true for that package. If f returns false, the imports will not be visited.
+func VisitPackages(root *packages.Package, f func(p *packages.Package) bool) {
+	seen := map[*packages.Package]bool{}
+	queue := []*packages.Package{root}
+	for len(queue) > 0 {
+		cur := queue[0]
+		queue = queue[1:]
+		if seen[cur] {
+			continue
+		}
+		if f(cur) {
+			for _, importedPkg := range cur.Imports {
+				queue = append(queue, importedPkg)
+			}
+		}
+		seen[cur] = true
+	}
+}