git.stella-ops.org/tools/stella-callgraph-go/framework.go

// Framework detection for Go projects
package main

import (
	"golang.org/x/tools/go/ssa"
	"strings"
)

// FrameworkPattern defines detection patterns for a framework
type FrameworkPattern struct {
	Name          string
	Packages      []string
	EntrypointFns []string
	HandlerType   string
}

// Known Go web frameworks
var frameworkPatterns = []FrameworkPattern{
	{
		Name:          "net/http",
		Packages:      []string{"net/http"},
		EntrypointFns: []string{"HandleFunc", "Handle", "ListenAndServe"},
		HandlerType:   "http_handler",
	},
	{
		Name:          "gin",
		Packages:      []string{"github.com/gin-gonic/gin"},
		EntrypointFns: []string{"GET", "POST", "PUT", "DELETE", "PATCH", "Run"},
		HandlerType:   "http_handler",
	},
	{
		Name:          "echo",
		Packages:      []string{"github.com/labstack/echo"},
		EntrypointFns: []string{"GET", "POST", "PUT", "DELETE", "PATCH", "Start"},
		HandlerType:   "http_handler",
	},
	{
		Name:          "fiber",
		Packages:      []string{"github.com/gofiber/fiber"},
		EntrypointFns: []string{"Get", "Post", "Put", "Delete", "Patch", "Listen"},
		HandlerType:   "http_handler",
	},
	{
		Name:          "chi",
		Packages:      []string{"github.com/go-chi/chi"},
		EntrypointFns: []string{"Get", "Post", "Put", "Delete", "Patch", "Route"},
		HandlerType:   "http_handler",
	},
	{
		Name:          "mux",
		Packages:      []string{"github.com/gorilla/mux"},
		EntrypointFns: []string{"HandleFunc", "Handle", "NewRouter"},
		HandlerType:   "http_handler",
	},
	{
		Name:          "grpc",
		Packages:      []string{"google.golang.org/grpc"},
		EntrypointFns: []string{"RegisterServer", "NewServer"},
		HandlerType:   "grpc_method",
	},
	{
		Name:          "cobra",
		Packages:      []string{"github.com/spf13/cobra"},
		EntrypointFns: []string{"Execute", "AddCommand", "Run"},
		HandlerType:   "cli_command",
	},
}

// DetectFramework checks if a function is related to a known framework
func DetectFramework(fn *ssa.Function) *FrameworkPattern {
	if fn.Pkg == nil {
		return nil
	}

	pkgPath := fn.Pkg.Pkg.Path()

	for _, pattern := range frameworkPatterns {
		for _, pkg := range pattern.Packages {
			if strings.Contains(pkgPath, pkg) {
				return &pattern
			}
		}
	}

	return nil
}

// DetectFrameworkEntrypoint checks if a call is a framework route registration
func DetectFrameworkEntrypoint(call *ssa.Call) *Entrypoint {
	callee := call.Call.StaticCallee()
	if callee == nil || callee.Pkg == nil {
		return nil
	}

	pkgPath := callee.Pkg.Pkg.Path()
	fnName := callee.Name()

	for _, pattern := range frameworkPatterns {
		for _, pkg := range pattern.Packages {
			if strings.Contains(pkgPath, pkg) {
				for _, epFn := range pattern.EntrypointFns {
					if fnName == epFn {
						nodeID := makeSymbolID(callee)
						return &Entrypoint{
							ID:   nodeID,
							Type: pattern.HandlerType,
						}
					}
				}
			}
		}
	}

	return nil
}

// IsHTTPHandler checks if a function signature matches http.Handler
func IsHTTPHandler(fn *ssa.Function) bool {
	sig := fn.Signature

	// Check for (http.ResponseWriter, *http.Request) signature
	if sig.Params().Len() == 2 {
		p0 := sig.Params().At(0).Type().String()
		p1 := sig.Params().At(1).Type().String()

		if strings.Contains(p0, "ResponseWriter") && strings.Contains(p1, "Request") {
			return true
		}
	}

	// Check for gin.Context, echo.Context, fiber.Ctx, etc.
	if sig.Params().Len() >= 1 {
		p0 := sig.Params().At(0).Type().String()
		if strings.Contains(p0, "gin.Context") ||
			strings.Contains(p0, "echo.Context") ||
			strings.Contains(p0, "fiber.Ctx") ||
			strings.Contains(p0, "chi.") {
			return true
		}
	}

	return false
}