check-safety/proto_rbac.go

package main

import (
	"bufio"
	"fmt"
	"go/ast"
	"go/parser"
	"go/token"
	"os"
	"path/filepath"
	"regexp"
	"sort"
	"strings"

	"git.dev.alexdunmow.com/block/check-safety/internal/helpers"
	"golang.org/x/mod/modfile"
)

type envViolation struct {
	file   string
	line   int
	envVar string
}

type protoOwnershipViolation struct {
	root      string
	namespace string
	coreProto string
}

var (
	bufModulePathPattern = regexp.MustCompile(`^(?:-\s*)?path:\s*"?([^"\s]+)"?\s*$`)
	protoPackagePattern  = regexp.MustCompile(`^\s*package\s+([A-Za-z0-9_.]+)\s*;\s*$`)
	protoServicePattern  = regexp.MustCompile(`^\s*service\s+([A-Za-z0-9_]+)\s*\{?\s*$`)
	protoRPCPattern      = regexp.MustCompile(`^\s*rpc\s+([A-Za-z0-9_]+)\s*\(`)
)

func checkEnvReads(root string) []envViolation {
	var violations []envViolation

	if err := filepath.Walk(root, func(path string, info os.FileInfo, err error) error {
		if err != nil || info.IsDir() || !strings.HasSuffix(path, ".go") {
			return nil
		}
		if isAllowedEnvReader(path) {
			return nil
		}
		// Skip vendor
		if strings.Contains(path, "/vendor/") {
			return nil
		}

		fset := token.NewFileSet()
		f, err := parser.ParseFile(fset, path, nil, 0)
		if err != nil {
			return nil // skip unparseable files
		}

		ast.Inspect(f, func(n ast.Node) bool {
			call, ok := n.(*ast.CallExpr)
			if !ok {
				return true
			}

			// Match os.Getenv("SECRET_VAR") and os.LookupEnv("SECRET_VAR")
			sel, ok := call.Fun.(*ast.SelectorExpr)
			if !ok {
				return true
			}
			ident, ok := sel.X.(*ast.Ident)
			if !ok || ident.Name != "os" || (sel.Sel.Name != "Getenv" && sel.Sel.Name != "LookupEnv") {
				return true
			}
			if len(call.Args) != 1 {
				return true
			}
			lit, ok := call.Args[0].(*ast.BasicLit)
			if !ok || lit.Kind != token.STRING {
				return true
			}
			envVar := strings.Trim(lit.Value, `"`)
			if secretEnvVars[envVar] && !exemptSecretEnvVars[envVar] {
				pos := fset.Position(call.Pos())
				relPath, _ := filepath.Rel(root, pos.Filename)
				if relPath == "" {
					relPath = pos.Filename
				}
				violations = append(violations, envViolation{
					file:   relPath,
					line:   pos.Line,
					envVar: envVar,
				})
			}
			return true
		})

		return nil
	}); err != nil {
		violations = append(violations, envViolation{file: root, line: 0, envVar: err.Error()})
	}

	return violations
}

func checkPluginProtoOwnership(root string, display string) ([]protoOwnershipViolation, error) {
	if !isPluginModuleRoot(root) {
		return nil, nil
	}
	if hasLocalProtoSources(root) {
		return nil, nil
	}

	coreRepoRoot, err := resolveBlockNinjaRepoRootFromModule(root)
	if err != nil || coreRepoRoot == "" {
		return nil, nil
	}

	packages, err := extractProcedurePackages(root)
	if err != nil {
		return nil, err
	}
	if len(packages) == 0 {
		return nil, nil
	}

	corePackages, err := extractProcedurePackages(coreRepoRoot)
	if err != nil {
		return nil, err
	}
	corePackageSet := make(map[string]bool, len(corePackages))
	for _, pkg := range corePackages {
		corePackageSet[pkg] = true
	}

	var violations []protoOwnershipViolation
	for _, pkg := range packages {
		if strings.HasPrefix(pkg, "blockninja.") {
			continue
		}
		if !corePackageSet[pkg] {
			continue
		}

		nsRoot := strings.Split(pkg, ".")[0]
		coreProto := filepath.Join(coreRepoRoot, "proto", nsRoot)
		if !dirExists(coreProto) {
			coreProto = filepath.Join(coreRepoRoot, "proto")
		}

		violations = append(violations, protoOwnershipViolation{
			root:      display,
			namespace: pkg,
			coreProto: coreProto,
		})
	}

	return violations, nil
}

func extractBufProcedures(root string, allowedPackagePrefixes ...string) ([]string, error) {
	moduleDirs, err := resolveBufModuleDirs(root)
	if err != nil {
		return nil, err
	}

	procedureSet := make(map[string]bool)
	for _, moduleDir := range moduleDirs {
		err := filepath.Walk(moduleDir, func(path string, info os.FileInfo, err error) error {
			if err != nil {
				return err
			}
			if info.IsDir() {
				if info.Name() == "vendor" || info.Name() == ".git" {
					return filepath.SkipDir
				}
				return nil
			}
			if !strings.HasSuffix(path, ".proto") {
				return nil
			}

			procedures, err := extractProtoProcedures(path)
			if err != nil {
				return err
			}
			for _, procedure := range procedures {
				if !procedureMatchesPackagePrefixes(procedure, allowedPackagePrefixes) {
					continue
				}
				procedureSet[procedure] = true
			}
			return nil
		})
		if err != nil {
			return nil, err
		}
	}

	procedures := make([]string, 0, len(procedureSet))
	for procedure := range procedureSet {
		procedures = append(procedures, procedure)
	}
	if len(procedures) == 0 {
		return extractConnectProceduresRecursive(root, allowedPackagePrefixes...)
	}
	sort.Strings(procedures)
	return procedures, nil
}

func extractProcedurePackages(root string) ([]string, error) {
	procedures, err := extractBufProcedures(root)
	if err != nil {
		return nil, err
	}

	pkgs := make(map[string]bool)
	for _, procedure := range procedures {
		trimmed := strings.TrimPrefix(procedure, "/")
		parts := strings.SplitN(trimmed, "/", 2)
		if len(parts) != 2 || parts[0] == "" {
			continue
		}
		serviceName := parts[0]
		lastDot := strings.LastIndex(serviceName, ".")
		if lastDot <= 0 {
			continue
		}
		pkgs[serviceName[:lastDot]] = true
	}

	out := make([]string, 0, len(pkgs))
	for pkg := range pkgs {
		out = append(out, pkg)
	}
	sort.Strings(out)
	return out, nil
}

func procedureMatchesPackagePrefixes(procedure string, allowedPackagePrefixes []string) bool {
	if len(allowedPackagePrefixes) == 0 {
		return true
	}

	trimmed := strings.TrimPrefix(procedure, "/")
	parts := strings.SplitN(trimmed, "/", 2)
	if len(parts) != 2 {
		return false
	}
	pkg := parts[0]

	for _, prefix := range allowedPackagePrefixes {
		if strings.HasPrefix(pkg, prefix) {
			return true
		}
	}

	return false
}

func resolveBufModuleDirs(root string) ([]string, error) {
	if !dirExists(root) {
		return nil, fmt.Errorf("root does not exist: %s", root)
	}

	bufPath := filepath.Join(root, "buf.yaml")
	if fileExists(bufPath) {
		paths, err := parseBufModulePaths(bufPath)
		if err != nil {
			return nil, err
		}
		if len(paths) > 0 {
			var dirs []string
			for _, relPath := range paths {
				moduleDir := filepath.Join(root, filepath.Clean(relPath))
				if dirExists(moduleDir) {
					dirs = append(dirs, moduleDir)
				}
			}
			if len(dirs) > 0 {
				return dirs, nil
			}
		}
	}

	fallback := filepath.Join(root, "proto")
	if dirExists(fallback) {
		return []string{fallback}, nil
	}

	return []string{root}, nil
}

func hasLocalProtoSources(root string) bool {
	moduleDirs, err := resolveBufModuleDirs(root)
	if err != nil {
		return false
	}

	for _, moduleDir := range moduleDirs {
		found := false
		if err := filepath.Walk(moduleDir, func(path string, info os.FileInfo, err error) error {
			if err != nil {
				return nil
			}
			if info.IsDir() {
				switch info.Name() {
				case ".git", "vendor", "node_modules":
					return filepath.SkipDir
				}
				return nil
			}
			if strings.HasSuffix(path, ".proto") {
				found = true
				return filepath.SkipAll
			}
			return nil
		}); err != nil {
			continue
		}
		if found {
			return true
		}
	}

	return false
}

func readModulePath(root string) (string, error) {
	modPath, ok := findNearestGoModuleRoot(root)
	if !ok {
		return "", fmt.Errorf("no go.mod found for %s", root)
	}
	data, err := os.ReadFile(filepath.Join(modPath, "go.mod"))
	if err != nil {
		return "", err
	}
	file, err := modfile.Parse("go.mod", data, nil)
	if err != nil {
		return "", err
	}
	if file.Module == nil {
		return "", fmt.Errorf("go.mod missing module declaration")
	}
	return file.Module.Mod.Path, nil
}

func resolveBlockNinjaRepoRootFromModule(root string) (string, error) {
	moduleRoot, ok := findNearestGoModuleRoot(root)
	if !ok {
		return "", fmt.Errorf("no go.mod found for %s", root)
	}
	data, err := os.ReadFile(filepath.Join(moduleRoot, "go.mod"))
	if err != nil {
		return "", err
	}
	file, err := modfile.Parse("go.mod", data, nil)
	if err != nil {
		return "", err
	}

	for _, replace := range file.Replace {
		if replace.Old.Path != "git.dev.alexdunmow.com/block/cms" {
			continue
		}
		if replace.New.Version != "" {
			continue
		}

		corePath := replace.New.Path
		if !filepath.IsAbs(corePath) {
			corePath = filepath.Join(moduleRoot, corePath)
		}
		corePath = filepath.Clean(corePath)
		if dirExists(filepath.Join(corePath, "cmd", "check-safety")) {
			return filepath.Clean(filepath.Join(corePath, "..")), nil
		}
		if dirExists(filepath.Join(corePath, "proto")) {
			return corePath, nil
		}
	}

	return "", nil
}

func parseBufModulePaths(path string) ([]string, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return nil, err
	}

	var paths []string
	inModules := false
	modulesIndent := 0

	for rawLine := range strings.SplitSeq(string(data), "\n") {
		line := strings.TrimRight(rawLine, "\r")
		trimmed := strings.TrimSpace(line)
		if trimmed == "" || strings.HasPrefix(trimmed, "#") {
			continue
		}

		indent := len(line) - len(strings.TrimLeft(line, " "))
		if !inModules {
			if trimmed == "modules:" {
				inModules = true
				modulesIndent = indent
			}
			continue
		}

		if indent <= modulesIndent && !strings.HasPrefix(trimmed, "-") {
			break
		}
		if match := bufModulePathPattern.FindStringSubmatch(trimmed); len(match) == 2 {
			paths = append(paths, match[1])
		}
	}

	return paths, nil
}

func extractProtoProcedures(path string) ([]string, error) {
	file, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer helpers.LogDeferredError(nil, "close proto source file", file.Close, "path", path)

	var (
		pkg          string
		service      string
		serviceDepth int
		procedures   []string
	)

	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		line := stripProtoLineComment(scanner.Text())
		trimmed := strings.TrimSpace(line)
		if trimmed == "" {
			continue
		}

		if pkg == "" {
			if match := protoPackagePattern.FindStringSubmatch(trimmed); len(match) == 2 {
				pkg = match[1]
			}
		}

		if service == "" {
			if match := protoServicePattern.FindStringSubmatch(trimmed); len(match) == 2 {
				service = match[1]
				serviceDepth = braceDelta(line)
				if serviceDepth <= 0 {
					serviceDepth = 1
				}
			}
			continue
		}

		if pkg != "" {
			if match := protoRPCPattern.FindStringSubmatch(trimmed); len(match) == 2 {
				procedures = append(procedures, fmt.Sprintf("/%s.%s/%s", pkg, service, match[1]))
			}
		}

		serviceDepth += braceDelta(line)
		if serviceDepth <= 0 {
			service = ""
			serviceDepth = 0
		}
	}

	if err := scanner.Err(); err != nil {
		return nil, err
	}

	return procedures, nil
}

func stripProtoLineComment(line string) string {
	if before, _, ok := strings.Cut(line, "//"); ok {
		return before
	}
	return line
}

func braceDelta(line string) int {
	return strings.Count(line, "{") - strings.Count(line, "}")
}

func extractPluginRoleMethods(root string) (map[string]string, error) {
	methods := make(map[string]string)
	fset := token.NewFileSet()

	err := filepath.Walk(root, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if info.IsDir() {
			switch info.Name() {
			case ".git", "vendor", "node_modules":
				return filepath.SkipDir
			}
			return nil
		}
		if !strings.HasSuffix(path, ".go") || strings.HasSuffix(path, "_test.go") {
			return nil
		}

		file, err := parser.ParseFile(fset, path, nil, 0)
		if err != nil {
			return nil
		}

		ast.Inspect(file, func(n ast.Node) bool {
			comp, ok := n.(*ast.CompositeLit)
			if !ok || !isRoleMapType(comp.Type) {
				return true
			}

			for _, elt := range comp.Elts {
				kv, ok := elt.(*ast.KeyValueExpr)
				if !ok {
					continue
				}
				lit, ok := kv.Key.(*ast.BasicLit)
				if !ok || lit.Kind != token.STRING {
					continue
				}
				method := strings.Trim(lit.Value, `"`)
				if strings.HasPrefix(method, "/") && strings.Contains(method[1:], "/") {
					methods[method] = normalizeRoleName(roleExprString(kv.Value))
				}
			}
			return true
		})

		return nil
	})
	if err != nil {
		return nil, err
	}

	return methods, nil
}

func isRoleMapType(expr ast.Expr) bool {
	mapType, ok := expr.(*ast.MapType)
	if !ok {
		return false
	}

	keyIdent, ok := mapType.Key.(*ast.Ident)
	if !ok || keyIdent.Name != "string" {
		return false
	}

	switch valueType := mapType.Value.(type) {
	case *ast.Ident:
		return valueType.Name == "Role"
	case *ast.SelectorExpr:
		return valueType.Sel != nil && valueType.Sel.Name == "Role"
	default:
		return false
	}
}

// extractConnectProceduresRecursive walks a tree, parses *.connect.go files, and extracts
// procedure string constants. This is a fallback for plugin repos that ship generated Connect
// stubs but not local proto sources.
func extractConnectProceduresRecursive(root string, allowedPackagePrefixes ...string) ([]string, error) {
	procedureSet := make(map[string]bool)
	fset := token.NewFileSet()
	err := filepath.Walk(root, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if info.IsDir() {
			switch info.Name() {
			case ".git", "vendor", "node_modules":
				return filepath.SkipDir
			}
			return nil
		}
		if !strings.HasSuffix(path, ".connect.go") {
			return nil
		}

		f, err := parser.ParseFile(fset, path, nil, 0)
		if err != nil {
			return nil
		}

		for _, decl := range f.Decls {
			genDecl, ok := decl.(*ast.GenDecl)
			if !ok || genDecl.Tok != token.CONST {
				continue
			}
			for _, spec := range genDecl.Specs {
				vs, ok := spec.(*ast.ValueSpec)
				if !ok || len(vs.Names) == 0 || len(vs.Values) == 0 {
					continue
				}
				name := vs.Names[0].Name
				if !strings.HasSuffix(name, "Procedure") {
					continue
				}
				lit, ok := vs.Values[0].(*ast.BasicLit)
				if !ok || lit.Kind != token.STRING {
					continue
				}
				val := strings.Trim(lit.Value, `"`)
				if procedureMatchesPackagePrefixes(val, allowedPackagePrefixes) {
					procedureSet[val] = true
				}
			}
		}
		return nil
	})
	if err != nil {
		return nil, err
	}

	procedures := make([]string, 0, len(procedureSet))
	for procedure := range procedureSet {
		procedures = append(procedures, procedure)
	}
	sort.Strings(procedures)
	return procedures, nil
}

// extractRBACMethodRoles parses interceptor.go and extracts all keys and role values from the
// MethodRoles map (orchestrator/helpdesk) or the methodRolesSeed map (CMS — the static seed
// table behind the atomic-pointer copy-on-write live table; see interceptor.go).
func extractRBACMethodRoles(file string) (map[string]string, error) {
	fset := token.NewFileSet()
	f, err := parser.ParseFile(fset, file, nil, 0)
	if err != nil {
		return nil, err
	}

	methods := make(map[string]string)

	ast.Inspect(f, func(n ast.Node) bool {
		// Find: var MethodRoles = map[string]Role{ ... }       (orchestrator/helpdesk)
		// or:   var methodRolesSeed = map[string]Role{ ... }   (CMS copy-on-write seed)
		vs, ok := n.(*ast.ValueSpec)
		if !ok || len(vs.Names) == 0 {
			return true
		}
		if name := vs.Names[0].Name; name != "MethodRoles" && name != "methodRolesSeed" {
			return true
		}
		if len(vs.Values) == 0 {
			return true
		}
		comp, ok := vs.Values[0].(*ast.CompositeLit)
		if !ok {
			return true
		}
		for _, elt := range comp.Elts {
			kv, ok := elt.(*ast.KeyValueExpr)
			if !ok {
				continue
			}
			lit, ok := kv.Key.(*ast.BasicLit)
			if !ok || lit.Kind != token.STRING {
				continue
			}
			method := strings.Trim(lit.Value, `"`)
			methods[method] = normalizeRoleName(roleExprString(kv.Value))
		}
		return false
	})

	return methods, nil
}

func roleExprString(expr ast.Expr) string {
	switch value := expr.(type) {
	case *ast.BasicLit:
		if value.Kind == token.STRING {
			return strings.Trim(value.Value, `"`)
		}
	case *ast.Ident:
		return value.Name
	case *ast.SelectorExpr:
		if value.Sel != nil {
			return value.Sel.Name
		}
	}
	return ""
}

func normalizeRoleName(raw string) string {
	switch raw {
	case "", "RolePublic":
		return "public"
	}
	if after, ok := strings.CutPrefix(raw, "Role"); ok {
		raw = after
	}
	return strings.ToLower(raw)
}

func isExcludedServiceMethod(method string) bool {
	parts := strings.Split(method, "/")
	if len(parts) < 3 {
		return false
	}
	svcParts := strings.Split(parts[1], ".")
	svcName := svcParts[len(svcParts)-1]
	return excludedServices[svcName]
}

func countRBACRoles(protoMethods []string, roleByMethod map[string]string) map[string]int {
	counts := make(map[string]int)
	for _, method := range protoMethods {
		if isExcludedServiceMethod(method) {
			continue
		}
		role, ok := roleByMethod[method]
		if !ok {
			continue
		}
		counts[role]++
	}
	return counts
}

func printRoleBreakdown(roleCounts map[string]int) {
	orderedRoles := []string{"public", "viewer", "user", "admin", "superadmin", "cms"}
	printed := make(map[string]bool)
	for _, role := range orderedRoles {
		if count, ok := roleCounts[role]; ok {
			fmt.Printf("    role %-10s %d\n", role, count)
			printed[role] = true
		}
	}

	var extraRoles []string
	for role := range roleCounts {
		if printed[role] {
			continue
		}
		extraRoles = append(extraRoles, role)
	}
	sort.Strings(extraRoles)
	for _, role := range extraRoles {
		fmt.Printf("    role %-10s %d\n", role, roleCounts[role])
	}
}