Go source parsing to detect web handlers

detect_handler.go

package main

import (
	"flag"
	"fmt"
	"go/ast"
	"go/parser"
	"go/token"
)

var (
	dirFlag string
)

func main() {
	flag.StringVar(&dirFlag, "dir", "./", "Dir where to parse go files")
	flag.Parse()

	fset := token.NewFileSet()
	packages, err := parser.ParseDir(fset, dirFlag, nil, 0)
	if err != nil {
		fmt.Println(err)
		return
	}

	for _, pkg := range packages {
		for _, f := range pkg.Files {
			visitor := &callExprVisitor{f}
			ast.Walk(visitor, f)
		}
	}
}

type callExprVisitor struct {
	f *ast.File
}

func (v *callExprVisitor) Visit(n ast.Node) ast.Visitor {
	switch node := n.(type) {
	case *ast.File:
		return v
	case *ast.CallExpr:
		if args := node.Args; len(args) == 2 {
			if isStringLit(args[0]) && isHandleFunc(args[1]) {
				fmt.Printf("detected handle func: %v\n", args[0])
			}
		}
	}
	return v
}

func isStringLit(node ast.Node) bool {
	switch n := node.(type) {
	case *ast.BasicLit:
		if n.Kind == token.STRING {
			return true
		}
	}
	return false
}
func isHandleFunc(node ast.Node) bool {
	switch fn := node.(type) {
	case *ast.FuncLit:
		params := fn.Type.Params
		if params.NumFields() != 2 {
			return false
		}
		fields := params.List
		if isSelector(fields[0].Type, "http", "ResponseWriter") &&
			isPtrSelector(fields[1].Type, "http", "Request") {
			return true
		}
	}
	return false
}

func isPtrSelector(expr ast.Node, x, sel string) bool {
	switch n := expr.(type) {
	case *ast.StarExpr:
		return isSelector(n.X, x, sel)
	}
	return false
}

func isSelector(expr ast.Node, x, sel string) bool {
	switch n := expr.(type) {
	case *ast.SelectorExpr:
		if fmt.Sprint(n.X) == x && sel == n.Sel.Name {
			return true
		}
	}
	return false
}