main.go

package main

import (
	"go/ast"
	"go/token"

	"golang.org/x/tools/go/analysis"
	"golang.org/x/tools/go/analysis/singlechecker"
)

var Analyzer = &analysis.Analyzer{
	Name: "forloopcheck",
	Doc:  "checks for potentially infinite loops without proper guardrails",
	Run:  run,
}

func run(pass *analysis.Pass) (interface{}, error) {
	// Create a map to store parent nodes
	parents := make(map[ast.Node]ast.Node)

	// First pass: build parent relationships
	for _, file := range pass.Files {
		ast.Inspect(file, func(n ast.Node) bool {
			if n == nil {
				return true
			}
			children := []ast.Node{}
			ast.Inspect(n, func(child ast.Node) bool {
				if child != nil && child != n {
					children = append(children, child)
				}
				return true
			})
			for _, child := range children {
				parents[child] = n
			}
			return true
		})
	}

	// Add helper function to check for timeout guards
	hasTimeoutGuard := func(node ast.Node) bool {
		var hasTimeout bool
		ast.Inspect(node, func(n ast.Node) bool {
			switch v := n.(type) {
			case *ast.SelectStmt:
				// Check for time.After or context deadline in select cases
				for _, clause := range v.Body.List {
					if commClause, ok := clause.(*ast.CommClause); ok {
						if commClause.Comm != nil {
							if recv, ok := commClause.Comm.(*ast.ExprStmt); ok {
								if call, ok := recv.X.(*ast.UnaryExpr); ok {
									if callExpr, ok := call.X.(*ast.CallExpr); ok {
										if sel, ok := callExpr.Fun.(*ast.SelectorExpr); ok {
											// Check for time.After
											if sel.Sel.Name == "After" {
												hasTimeout = true
												return false
											}
											// Check for context deadline
											if sel.Sel.Name == "Done" {
												hasTimeout = true
												return false
											}
										}
									}
								}
							}
						}
					}
				}
			case *ast.CallExpr:
				// Check for context deadline checks
				if sel, ok := v.Fun.(*ast.SelectorExpr); ok {
					if sel.Sel.Name == "Deadline" || sel.Sel.Name == "Done" {
						hasTimeout = true
						return false
					}
				}
			}
			return true
		})
		return hasTimeout
	}

	// Second pass: analyze loops
	for _, file := range pass.Files {
		ast.Inspect(file, func(n ast.Node) bool {
			forStmt, ok := n.(*ast.ForStmt)
			if !ok {
				return true
			}

			// Allow traditional for i := 0; i < N; i++ style loops
			if forStmt.Init != nil && forStmt.Cond != nil && forStmt.Post != nil {
				return true
			}

			// Allow range-based loops
			if _, ok := n.(*ast.RangeStmt); ok {
				return true
			}

			// Variables for tracking loop safety
			var (
				hasGuardIncrement     bool
				hasGuardCheck         bool
				hasFunctionCall       bool
				modifiesConditionVars bool
				guardVars             = make(map[string]bool)
				condVars              = make(map[string]bool)
			)

			// Find guard variables declared before the loop
			if parent, ok := parents[forStmt]; ok {
				if block, ok := parent.(*ast.BlockStmt); ok {
					for _, stmt := range block.List {
						if assign, ok := stmt.(*ast.AssignStmt); ok {
							for _, lhs := range assign.Lhs {
								if ident, ok := lhs.(*ast.Ident); ok {
									guardVars[ident.Name] = true
								}
							}
						}
					}
				}
			}

			// Also check for guard variables declared inside the loop
			ast.Inspect(forStmt.Body, func(n ast.Node) bool {
				if assign, ok := n.(*ast.AssignStmt); ok {
					for _, lhs := range assign.Lhs {
						if ident, ok := lhs.(*ast.Ident); ok {
							guardVars[ident.Name] = true
						}
					}
				}
				return true
			})

			// Second pass: Check for guard increments and break conditions
			ast.Inspect(forStmt.Body, func(n ast.Node) bool {
				switch v := n.(type) {
				case *ast.IncDecStmt:
					if ident, ok := v.X.(*ast.Ident); ok {
						if guardVars[ident.Name] {
							hasGuardIncrement = true
						}
					}
				case *ast.AssignStmt:
					if len(v.Lhs) == 1 {
						if ident, ok := v.Lhs[0].(*ast.Ident); ok {
							if guardVars[ident.Name] {
								if v.Tok == token.ADD_ASSIGN || v.Tok == token.SUB_ASSIGN {
									hasGuardIncrement = true
								}
							}
						}
					}
				case *ast.IfStmt:
					// Check if the condition uses a guard variable
					ast.Inspect(v.Cond, func(n ast.Node) bool {
						if binExpr, ok := n.(*ast.BinaryExpr); ok {
							if ident, ok := binExpr.X.(*ast.Ident); ok && guardVars[ident.Name] {
								// Look for break statement in the if body
								ast.Inspect(v.Body, func(n ast.Node) bool {
									if br, ok := n.(*ast.BranchStmt); ok && br.Tok == token.BREAK {
										hasGuardCheck = true
										return false
									}
									return true
								})
							} else if ident, ok := binExpr.Y.(*ast.Ident); ok && guardVars[ident.Name] {
								// Check the right side of the comparison too
								ast.Inspect(v.Body, func(n ast.Node) bool {
									if br, ok := n.(*ast.BranchStmt); ok && br.Tok == token.BREAK {
										hasGuardCheck = true
										return false
									}
									return true
								})
							}
						}
						return true
					})
				}
				return true
			})

			// Check if this is a condition-based loop (for cond {})
			if forStmt.Cond != nil {
				// Check if the condition contains a function call or channel operation that could be non-deterministic
				ast.Inspect(forStmt.Cond, func(n ast.Node) bool {
					switch v := n.(type) {
					case *ast.CallExpr:
						hasFunctionCall = true
						return false
					case *ast.UnaryExpr:
						// Check for channel receive operations (<-ch)
						if v.Op == token.ARROW {
							hasFunctionCall = true // We'll reuse this flag for both function calls and channel ops
							return false
						}
					}
					return true
				})

				// If there's no function call or channel operation, it's likely safe
				if !hasFunctionCall {
					return true
				}

				// Collect variables used in condition
				ast.Inspect(forStmt.Cond, func(n ast.Node) bool {
					if ident, ok := n.(*ast.Ident); ok {
						condVars[ident.Name] = true
					}
					return true
				})

				// Check if any of these variables are modified in the loop body
				ast.Inspect(forStmt.Body, func(n ast.Node) bool {
					switch v := n.(type) {
					case *ast.AssignStmt:
						for _, lhs := range v.Lhs {
							if ident, ok := lhs.(*ast.Ident); ok {
								if condVars[ident.Name] {
									modifiesConditionVars = true
									return false
								}
							}
						}
					case *ast.IncDecStmt:
						if ident, ok := v.X.(*ast.Ident); ok {
							if condVars[ident.Name] {
								modifiesConditionVars = true
								return false
							}
						}
					}
					return true
				})

				// Check for guard variables and break conditions
				var hasGuardCheck bool

				ast.Inspect(forStmt.Body, func(n ast.Node) bool {
					if ifStmt, ok := n.(*ast.IfStmt); ok {
						ast.Inspect(ifStmt.Body, func(n ast.Node) bool {
							if br, ok := n.(*ast.BranchStmt); ok && br.Tok == token.BREAK {
								// Check if the if condition uses a guard variable
								ast.Inspect(ifStmt.Cond, func(n ast.Node) bool {
									if binExpr, ok := n.(*ast.BinaryExpr); ok {
										if ident, ok := binExpr.X.(*ast.Ident); ok && guardVars[ident.Name] {
											hasGuardCheck = true
										} else if ident, ok := binExpr.Y.(*ast.Ident); ok && guardVars[ident.Name] {
											hasGuardCheck = true
										}
									}
									return true
								})
							}
							return true
						})
					}
					return true
				})

				// If we have a guard check with break, it's safe
				if hasGuardCheck {
					return true
				}

				// Only report if there's no guard check and no condition variable modification
				if !hasGuardCheck && !modifiesConditionVars {
					pass.Reportf(forStmt.Pos(), "potentially infinite loop: condition uses function call but loop body doesn't modify condition variables or have guard rails")
				}
				return true
			}

			// For infinite loops (for {}) we require either traditional guards or timeout guards
			if forStmt.Cond == nil {
				if !hasGuardIncrement || !hasGuardCheck {
					// Check for timeout-based guards before reporting
					if !hasTimeoutGuard(forStmt.Body) {
						pass.Reportf(forStmt.Pos(), "potentially infinite loop without proper guardrails or timeout mechanism")
					}
				}
			}

			// For condition-based loops, check timeout guards as well
			if forStmt.Cond != nil && hasFunctionCall {
				if !hasGuardCheck && !modifiesConditionVars && !hasTimeoutGuard(forStmt.Body) {
					pass.Reportf(forStmt.Pos(), "potentially infinite loop: condition uses function call but loop body doesn't modify condition variables, have guard rails, or timeout mechanism")
				}
			}

			return true
		})
	}

	return nil, nil
}

func main() {
	singlechecker.Main(Analyzer)
}