From e6f8cd8b6a8a18810ff3cf0ae31e33df9576d487 Mon Sep 17 00:00:00 2001
From: Samarth Kishor <samarkis@amazon.com>
Date: Tue, 3 Dec 2024 17:03:55 -0500
Subject: [PATCH] Fix annotation loading for comments in the AST

---
 analysis/annotations/annotations.go |  7 ++-
 analysis/loadprogram/state.go       | 84 +++++++++++++++++++++++------
 2 files changed, 72 insertions(+), 19 deletions(-)

diff --git a/analysis/annotations/annotations.go b/analysis/annotations/annotations.go
index 1ceb4ca7..e47497f0 100644
--- a/analysis/annotations/annotations.go
+++ b/analysis/annotations/annotations.go
@@ -25,7 +25,6 @@ import (
 	"github.com/awslabs/ar-go-tools/analysis/config"
 	"github.com/awslabs/ar-go-tools/internal/funcutil"
 	"golang.org/x/exp/slices"
-	"golang.org/x/tools/go/packages"
 	"golang.org/x/tools/go/ssa"
 )
 
@@ -218,13 +217,13 @@ func (pa ProgramAnnotations) Iter(fx func(a Annotation)) {
 
 // CompleteFromSyntax takes a set of program annotations and adds additional non-ssa linked annotations
 // to the annotations
-func (pa ProgramAnnotations) CompleteFromSyntax(logger *config.LogGroup, pkg *packages.Package) {
-	for _, astFile := range pkg.Syntax {
+func (pa ProgramAnnotations) CompleteFromSyntax(logger *config.LogGroup, fset *token.FileSet, files []*ast.File) {
+	for _, astFile := range files {
 		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()))
+					pa.loadFileAnnotations(logger, annotationContents, fset.Position(comment.Pos()))
 				}
 			}
 		}
diff --git a/analysis/loadprogram/state.go b/analysis/loadprogram/state.go
index 86d2c0d5..b78869f3 100644
--- a/analysis/loadprogram/state.go
+++ b/analysis/loadprogram/state.go
@@ -16,6 +16,11 @@ package loadprogram
 
 import (
 	"fmt"
+	"go/ast"
+	"go/parser"
+	"go/token"
+	"io/fs"
+	"path/filepath"
 	"sync/atomic"
 
 	"github.com/awslabs/ar-go-tools/analysis/annotations"
@@ -67,25 +72,22 @@ func NewState(c *config.State) result.Result[State] {
 	}
 	// Load annotations by scanning all packages' syntax
 	pa, err := annotations.LoadAnnotations(c.Logger, program.AllPackages())
+	if err != nil {
+		return result.Err[State](fmt.Errorf("failed to load annotations from SSA"))
+	}
 
-	if pkgs != nil {
-		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
-				}
-				// TODO: find a way to not scan dependencies if there is demand. Currently, it is unlikely that some
-				// dependencies will contain argot annotations.
-				c.Logger.Debugf("Scan %s for annotations.\n", p.PkgPath)
-				pa.CompleteFromSyntax(c.Logger, p)
-				return true
-			})
-		}
+	dirs := c.Patterns
+	// If the project root is set, parse all Go files inside the root directory (recursively)
+	if root := c.Config.Root(); root != "" {
+		dirs = []string{root}
 	}
+
+	files, err := allAstFiles(dirs, program.Fset, pkgs)
 	if err != nil {
-		return result.Err[State](err)
+		return result.Err[State](fmt.Errorf("failed to parse AST files: %v", err))
 	}
+
+	pa.CompleteFromSyntax(c.Logger, program.Fset, files)
 	c.Logger.Infof("Loaded %d annotations from program\n", pa.Count())
 
 	report := config.NewReport()
@@ -173,3 +175,55 @@ func (wps *State) IncrementAndTestAlarms() bool {
 func (wps *State) TestAlarmCount() bool {
 	return wps.Config.MaxAlarms <= 0 || wps.numAlarms.Load() < int32(wps.Config.MaxAlarms)
 }
+
+func allAstFiles(dirs []string, fset *token.FileSet, pkgs []*packages.Package) ([]*ast.File, error) {
+	var files []*ast.File
+
+	// HACK Ideally, this loop shouldn't be necessary but sometimes
+	// analysisutil.VisitPackages will miss some AST files.
+	// Hopefully there's a better way to get all the AST files in the program
+	// without needing to parse it twice.
+	for _, dir := range dirs {
+		if err := filepath.Walk(dir, func(path string, info fs.FileInfo, err error) error {
+			if err != nil {
+				return err
+			}
+
+			if info.IsDir() {
+				// Parse all files in the directory, making sure to include comments
+				parsedDir, err := parser.ParseDir(fset, path, nil, parser.ParseComments)
+				if err != nil {
+					return fmt.Errorf("failed to parse dir %s: %v", path, err)
+				}
+
+				for _, pkg := range parsedDir {
+					for _, file := range pkg.Files {
+						files = append(files, file)
+					}
+				}
+			}
+
+			return nil
+		}); err != nil {
+			return nil, fmt.Errorf("failed to parse AST files: %v", err)
+		}
+	}
+
+	if pkgs != nil {
+		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
+				}
+
+				// TODO: Remove if there is no need for scanning dependencies
+				files = append(files, p.Syntax...)
+
+				return true
+			})
+		}
+	}
+
+	return files, nil
+}