agentic-coding-workflow/internal/acp/handlers/server_grep.go

package handlers

import (
	"bufio"
	"context"
	"encoding/json"
	"io/fs"
	"os"
	"path/filepath"
	"regexp"
	"strings"
)

// Grep caps. These bound the cost of a single grep request to avoid a DoS via a
// crafted pattern / huge tree. All are overridable on the handler.
const (
	DefaultGrepMaxMatches  = 200
	DefaultGrepMaxFiles    = 5000
	DefaultGrepMaxFileSize = 1 << 20 // 1 MiB: skip larger files
	DefaultGrepMaxLineLen  = 1000    // truncate long matched lines
)

// GrepHandler handles fs/grep: a bounded regexp search under the worktree using a
// pure-Go walk (no shell, no git grep interpolation) confined to the pathsafe
// root. Binary files and oversize files are skipped.
type GrepHandler struct {
	MaxMatches  int
	MaxFiles    int
	MaxFileSize int64
}

// NewGrepHandler constructs a GrepHandler with default caps.
func NewGrepHandler() *GrepHandler {
	return &GrepHandler{
		MaxMatches:  DefaultGrepMaxMatches,
		MaxFiles:    DefaultGrepMaxFiles,
		MaxFileSize: DefaultGrepMaxFileSize,
	}
}

func (h *GrepHandler) maxMatches() int {
	if h.MaxMatches > 0 {
		return h.MaxMatches
	}
	return DefaultGrepMaxMatches
}

func (h *GrepHandler) maxFiles() int {
	if h.MaxFiles > 0 {
		return h.MaxFiles
	}
	return DefaultGrepMaxFiles
}

func (h *GrepHandler) maxFileSize() int64 {
	if h.MaxFileSize > 0 {
		return h.MaxFileSize
	}
	return DefaultGrepMaxFileSize
}

type fsGrepParams struct {
	SessionID     string `json:"sessionId"`
	Pattern       string `json:"pattern"`
	Path          string `json:"path,omitempty"` // sub-path to scope the search; defaults to cwd root
	CaseSensitive bool   `json:"case_sensitive,omitempty"`
}

type grepMatch struct {
	File string `json:"file"` // worktree-relative
	Line int    `json:"line"` // 1-based
	Text string `json:"text"`
}

type fsGrepResult struct {
	Matches   []grepMatch `json:"matches"`
	Truncated bool        `json:"truncated"`
}

// Handle runs a regexp search. Returns FsError for path escape (-32001), invalid
// pattern (-32602), or read failure. Match/file caps are honored.
func (h *GrepHandler) Handle(ctx context.Context, sess SessionContext, paramsJSON json.RawMessage) FsResult {
	var p fsGrepParams
	if err := json.Unmarshal(paramsJSON, &p); err != nil {
		return FsResult{Err: &FsError{Code: -32602, Message: "invalid params: " + err.Error()}}
	}
	if strings.TrimSpace(p.Pattern) == "" {
		return FsResult{Err: &FsError{Code: -32602, Message: "pattern required"}}
	}
	expr := p.Pattern
	if !p.CaseSensitive {
		expr = "(?i)" + expr
	}
	re, err := regexp.Compile(expr)
	if err != nil {
		return FsResult{Err: &FsError{Code: -32602, Message: "invalid pattern: " + err.Error()}}
	}

	searchPath := p.Path
	if searchPath == "" {
		searchPath = "."
	}
	root, err := resolveSafePath(sess.CwdPath, searchPath)
	if err != nil {
		return FsResult{Err: &FsError{Code: -32001, Message: err.Error()}}
	}

	cwd := filepath.Clean(sess.CwdPath)
	out := fsGrepResult{Matches: make([]grepMatch, 0, 16)}
	maxMatches := h.maxMatches()
	maxFiles := h.maxFiles()
	maxSize := h.maxFileSize()
	filesScanned := 0

	walkErr := filepath.WalkDir(root, func(path string, d fs.DirEntry, werr error) error {
		if werr != nil {
			return nil // skip unreadable entries
		}
		select {
		case <-ctx.Done():
			return filepath.SkipAll
		default:
		}
		if d.IsDir() {
			if isSkippedDir(d.Name()) {
				return filepath.SkipDir
			}
			return nil
		}
		if filesScanned >= maxFiles {
			out.Truncated = true
			return filepath.SkipAll
		}
		info, ierr := d.Info()
		if ierr != nil || info.Size() > maxSize {
			return nil
		}
		filesScanned++
		rel, rerr := filepath.Rel(cwd, path)
		if rerr != nil {
			rel = path
		}
		rel = filepath.ToSlash(rel)
		if grepFile(path, rel, re, &out, maxMatches) {
			out.Truncated = true
			return filepath.SkipAll
		}
		return nil
	})
	if walkErr != nil && walkErr != filepath.SkipAll {
		return FsResult{Err: &FsError{Code: -32603, Message: "grep walk: " + walkErr.Error()}}
	}

	res, _ := json.Marshal(out)
	return FsResult{OK: res}
}

// grepFile scans one file line-by-line. Returns true when the global match cap
// was reached (caller should stop the walk).
func grepFile(abs, rel string, re *regexp.Regexp, out *fsGrepResult, maxMatches int) bool {
	f, err := os.Open(abs)
	if err != nil {
		return false
	}
	defer f.Close()

	sc := bufio.NewScanner(f)
	sc.Buffer(make([]byte, 0, 64*1024), 1<<20)
	lineNo := 0
	for sc.Scan() {
		lineNo++
		line := sc.Bytes()
		// Skip binary content (NUL in line).
		if hasNUL(line) {
			return false
		}
		if re.Match(line) {
			text := string(line)
			if len(text) > DefaultGrepMaxLineLen {
				text = text[:DefaultGrepMaxLineLen] + "...[truncated]"
			}
			out.Matches = append(out.Matches, grepMatch{File: rel, Line: lineNo, Text: text})
			if len(out.Matches) >= maxMatches {
				return true
			}
		}
	}
	return false
}

func hasNUL(b []byte) bool {
	for _, c := range b {
		if c == 0 {
			return true
		}
	}
	return false
}

// skippedDirs are never descended into during grep.
var skippedDirs = map[string]struct{}{
	".git": {}, "node_modules": {}, "vendor": {}, "dist": {}, "build": {},
	".venv": {}, "__pycache__": {}, "third_party": {},
}

func isSkippedDir(name string) bool {
	_, ok := skippedDirs[name]
	return ok
}