feat: add configurable tool call blocking and monitoring

internal/toolguard/toolguard.go

@@ -0,0 +1,291 @@
+// Package toolguard applies configurable blocking rules before tool execution.
+// It inspects tool names and arguments, including JSON string values and common
+// percent escapes. It does not resolve hosts or inspect redirects, files, or
+// arbitrary encoded payloads, and is not an exhaustive target authorization check.
+package toolguard
+
+import (
+	"encoding/json"
+	"fmt"
+	"regexp"
+	"sort"
+	"strings"
+	"sync/atomic"
+)
+
+const (
+	MaxRules         = 100
+	MaxIDLength      = 128
+	MaxNameLength    = 200
+	MaxPatternLength = 4096
+	MaxMessageLength = 4096
+
+	// The named group keeps surrounding boundary punctuation out of the reminder.
+	governmentDomainPattern = `(?i)(?:^|[^\p{L}\p{M}\p{N}_.-])(?P<match>(?:(?:[\p{L}\p{M}\p{N}_*-]+\.)+gov(?:\.[\p{L}\p{M}\p{N}_*-]+)*|gov(?:\.[\p{L}\p{M}\p{N}_*-]+)+|\.gov(?:\.[\p{L}\p{M}\p{N}_*-]+)*)\.?)(?:$|[^\p{L}\p{M}\p{N}_.-])`
+	defaultMessage          = "识别到 {match}，工具调用已被安全规则「{rule}」拦截，请检查目标与授权范围后再试。"
+)
+
+type Config struct {
+	Enabled bool   `json:"enabled" yaml:"enabled"`
+	Rules   []Rule `json:"rules" yaml:"rules"`
+}
+
+type Rule struct {
+	ID      string `json:"id" yaml:"id"`
+	Name    string `json:"name" yaml:"name"`
+	Enabled bool   `json:"enabled" yaml:"enabled"`
+	Pattern string `json:"pattern" yaml:"pattern"`
+	Message string `json:"message" yaml:"message"`
+}
+
+type Match struct {
+	RuleID      string `json:"ruleId"`
+	RuleName    string `json:"ruleName"`
+	MatchedText string `json:"matchedText"`
+	Message     string `json:"message"`
+}
+
+type compiledRule struct {
+	rule       Rule
+	pattern    *regexp.Regexp
+	matchGroup int
+}
+
+// Policy is an immutable snapshot safe for concurrent checks.
+type Policy struct {
+	config Config
+	rules  []compiledRule
+}
+
+// DefaultConfig enables a conservative government-domain rule, including .gov,
+// .gov.cn and wildcard forms such as *.gov.*. Additional rules can be configured.
+func DefaultConfig() Config {
+	return Config{
+		Enabled: true,
+		Rules: []Rule{{
+			ID:      "government-domains",
+			Name:    "政府网站保护",
+			Enabled: true,
+			Pattern: governmentDomainPattern,
+			Message: "识别到 {match}，禁止攻击政府网站。请检查目标与授权范围，并更换为已获授权的非政府目标。",
+		}},
+	}
+}
+
+// Compile validates every rule, including disabled ones. Limits are byte counts.
+// Rule order determines precedence. An optional named (?P<match>...) group selects
+// the text inserted into {match}; otherwise the entire regex match is used.
+func Compile(config Config) (*Policy, error) {
+	if len(config.Rules) > MaxRules {
+		return nil, fmt.Errorf("tool guard: at most %d rules are allowed", MaxRules)
+	}
+	policy := &Policy{config: cloneConfig(config)}
+	ids := make(map[string]struct{}, len(config.Rules))
+	for i, rule := range policy.config.Rules {
+		prefix := fmt.Sprintf("tool guard rule %d", i+1)
+		if strings.TrimSpace(rule.ID) == "" || len(rule.ID) > MaxIDLength {
+			return nil, fmt.Errorf("%s: id must be nonempty and at most %d bytes", prefix, MaxIDLength)
+		}
+		if rule.ID != strings.TrimSpace(rule.ID) {
+			return nil, fmt.Errorf("%s: id must not have surrounding whitespace", prefix)
+		}
+		if _, exists := ids[rule.ID]; exists {
+			return nil, fmt.Errorf("%s: duplicate id %q", prefix, rule.ID)
+		}
+		ids[rule.ID] = struct{}{}
+		if strings.TrimSpace(rule.Name) == "" || len(rule.Name) > MaxNameLength {
+			return nil, fmt.Errorf("%s: name must be nonempty and at most %d bytes", prefix, MaxNameLength)
+		}
+		if strings.TrimSpace(rule.Pattern) == "" || len(rule.Pattern) > MaxPatternLength {
+			return nil, fmt.Errorf("%s: pattern must be nonempty and at most %d bytes", prefix, MaxPatternLength)
+		}
+		if len(rule.Message) > MaxMessageLength {
+			return nil, fmt.Errorf("%s: message must be at most %d bytes", prefix, MaxMessageLength)
+		}
+		pattern, err := regexp.Compile(rule.Pattern)
+		if err != nil {
+			return nil, fmt.Errorf("%s (%s): invalid regular expression: %w", prefix, rule.ID, err)
+		}
+		if rule.Enabled {
+			policy.rules = append(policy.rules, compiledRule{rule: rule, pattern: pattern, matchGroup: pattern.SubexpIndex("match")})
+		}
+	}
+	return policy, nil
+}
+
+// Check returns the first blocking rule, or nil when the call is allowed. Args
+// should be JSON-compatible and must not be mutated while Check is running.
+func (p *Policy) Check(toolName string, args map[string]interface{}) *Match {
+	if p == nil || !p.config.Enabled || len(p.rules) == 0 {
+		return nil
+	}
+	candidates := candidateTexts(toolName, args)
+	for _, rule := range p.rules {
+		for _, candidate := range candidates {
+			indices := rule.pattern.FindStringSubmatchIndex(candidate)
+			if indices == nil {
+				continue
+			}
+			start, end := indices[0], indices[1]
+			if group := rule.matchGroup; group > 0 && indices[group*2] >= 0 {
+				start, end = indices[group*2], indices[group*2+1]
+			}
+			matched := candidate[start:end]
+			message := rule.rule.Message
+			if strings.TrimSpace(message) == "" {
+				message = defaultMessage
+			}
+			// A single replacement pass prevents matched text from introducing
+			// additional template substitutions.
+			message = strings.NewReplacer("{match}", matched, "{tool}", toolName, "{rule}", rule.rule.Name).Replace(message)
+			return &Match{RuleID: rule.rule.ID, RuleName: rule.rule.Name, MatchedText: matched, Message: message}
+		}
+	}
+	return nil
+}
+
+// Manager atomically replaces validated policy snapshots for live settings.
+type Manager struct {
+	policy atomic.Pointer[Policy]
+}
+
+func NewManager(config Config) (*Manager, error) {
+	m := &Manager{}
+	if err := m.Update(config); err != nil {
+		return nil, err
+	}
+	return m, nil
+}
+
+// Update retains the active policy if validation fails.
+func (m *Manager) Update(config Config) error {
+	policy, err := Compile(config)
+	if err != nil {
+		return err
+	}
+	m.policy.Store(policy)
+	return nil
+}
+
+func (m *Manager) Config() Config {
+	if m == nil {
+		return Config{}
+	}
+	if policy := m.policy.Load(); policy != nil {
+		return cloneConfig(policy.config)
+	}
+	return Config{}
+}
+
+func (m *Manager) Check(toolName string, args map[string]interface{}) *Match {
+	if m == nil {
+		return nil
+	}
+	return m.policy.Load().Check(toolName, args)
+}
+
+func cloneConfig(config Config) Config {
+	if config.Rules != nil {
+		rules := make([]Rule, len(config.Rules))
+		copy(rules, config.Rules)
+		config.Rules = rules
+	}
+	return config
+}
+
+func candidateTexts(toolName string, args map[string]interface{}) []string {
+	var candidates []string
+	seen := make(map[string]struct{})
+	add := func(value string) {
+		// Decode at most three rounds to cover common nested URL escaping
+		// without claiming support for arbitrarily encoded tool inputs.
+		for round := 0; round <= 3; round++ {
+			if _, exists := seen[value]; !exists {
+				seen[value] = struct{}{}
+				candidates = append(candidates, value)
+			}
+			decoded := decodePercentEscapes(value)
+			if decoded == value {
+				break
+			}
+			value = decoded
+		}
+	}
+	add(toolName)
+	var walk func(interface{}, int)
+	walk = func(value interface{}, remainingDepth int) {
+		if remainingDepth == 0 {
+			return
+		}
+		switch value := value.(type) {
+		case string:
+			add(value)
+		case map[string]interface{}:
+			keys := make([]string, 0, len(value))
+			for key := range value {
+				keys = append(keys, key)
+			}
+			sort.Strings(keys)
+			for _, key := range keys {
+				add(key)
+				walk(value[key], remainingDepth-1)
+			}
+		case []interface{}:
+			for _, item := range value {
+				walk(item, remainingDepth-1)
+			}
+		}
+	}
+	// Unmarshaling also normalizes typed slices/maps, json.RawMessage and
+	// escaped JSON keys/values into the recursive representation above.
+	if raw, err := json.Marshal(args); err == nil {
+		var value interface{}
+		if json.Unmarshal(raw, &value) == nil {
+			// encoding/json accepts at most 10,000 nesting levels. The decoded
+			// tree is acyclic, so inspect every accepted string/key at that depth.
+			walk(value, 10001)
+		}
+		add(string(raw))
+	} else {
+		// MCP rejects invalid JSON arguments independently; still inspect the
+		// ordinary values if a caller supplies a non-JSON value alongside them.
+		walk(args, 128)
+	}
+	return candidates
+}
+
+// Decode valid percent triplets even if another part of the string has a stray
+// percent sign. Whole-string URL unescaping otherwise misses such mixed inputs.
+func decodePercentEscapes(value string) string {
+	if !strings.Contains(value, "%") {
+		return value
+	}
+	var out strings.Builder
+	out.Grow(len(value))
+	for i := 0; i < len(value); i++ {
+		if value[i] == '%' && i+2 < len(value) {
+			hi, okHi := hexValue(value[i+1])
+			lo, okLo := hexValue(value[i+2])
+			if okHi && okLo {
+				out.WriteByte(hi<<4 | lo)
+				i += 2
+				continue
+			}
+		}
+		out.WriteByte(value[i])
+	}
+	return out.String()
+}
+
+func hexValue(value byte) (byte, bool) {
+	switch {
+	case value >= '0' && value <= '9':
+		return value - '0', true
+	case value >= 'a' && value <= 'f':
+		return value - 'a' + 10, true
+	case value >= 'A' && value <= 'F':
+		return value - 'A' + 10, true
+	default:
+		return 0, false
+	}
+}

internal/toolguard/toolguard_test.go

@@ -0,0 +1,231 @@
+package toolguard
+
+import (
+	"encoding/json"
+	"strings"
+	"sync"
+	"testing"
+)
+
+func TestDefaultGovernmentProtection(t *testing.T) {
+	policy, err := Compile(DefaultConfig())
+	if err != nil {
+		t.Fatal(err)
+	}
+	for _, test := range []struct {
+		input string
+		match string
+	}{
+		{"https://agency.gov/login", "agency.gov"},
+		{"https://www.agency.gov.cn:443/login", "www.agency.gov.cn"},
+		{"curl https://EXAMPLE.GOV.UK/a", "EXAMPLE.GOV.UK"},
+		{"*.gov", "*.gov"},
+		{"*.gov.*", "*.gov.*"},
+		{".gov", ".gov"},
+		{".gov.*", ".gov.*"},
+		{"https://政务.gov.cn/", "政务.gov.cn"},
+		{"https://gov.cn/", "gov.cn"},
+		{"https://agency.gov./", "agency.gov."},
+		{"https://agency%2egov/a", "agency.gov"},
+		{"https://agency%252Egov/a", "agency.gov"},
+		{"echo 100% && curl https://agency%2egov/a", "agency.gov"},
+	} {
+		t.Run(test.input, func(t *testing.T) {
+			match := policy.Check("http_request", map[string]interface{}{"target": test.input})
+			if match == nil || match.MatchedText != test.match {
+				t.Fatalf("Check = %+v, want match %q", match, test.match)
+			}
+			if !strings.Contains(match.Message, test.match) || !strings.Contains(match.Message, "禁止攻击政府网站") {
+				t.Fatalf("unexpected reminder: %q", match.Message)
+			}
+		})
+	}
+	for _, input := range []string{
+		"https://example.com/", "https://government.example/", "https://agency.govt/",
+		"https://agency.gov-example.com/", "https://agency.gov_cn/", "governance", "gov",
+		".government", ".govx",
+	} {
+		t.Run("allowed "+input, func(t *testing.T) {
+			if match := policy.Check("http_request", map[string]interface{}{"target": input}); match != nil {
+				t.Fatalf("unexpected match: %+v", match)
+			}
+		})
+	}
+}
+
+func TestNestedArgumentsAndJSONEscapes(t *testing.T) {
+	policy, _ := Compile(DefaultConfig())
+	for _, args := range []map[string]interface{}{
+		{"targets": []interface{}{map[string]interface{}{"target": "https://agency.gov"}}},
+		{"targets": []string{"https://agency.gov"}},
+		{"targets": map[string]string{"target": "https://agency.gov"}},
+		{"https://agency.gov": true},
+		{"payload": json.RawMessage(`{"target":"https://agency\u002egov"}`)},
+		{"payload": json.RawMessage(`{"https://agency\u002egov":true}`)},
+		{"bad_value": make(chan string), "target": "https://agency.gov"},
+	} {
+		if match := policy.Check("request", args); match == nil || match.MatchedText != "agency.gov" {
+			t.Fatalf("Check(%v) = %+v", args, match)
+		}
+	}
+}
+
+func TestDeepNestedJSONEscapes(t *testing.T) {
+	policy, _ := Compile(DefaultConfig())
+	// Deep nesting must not hide a domain represented with JSON Unicode escapes.
+	payload := strings.Repeat("[", 200) + `"https://agency\u002egov"` + strings.Repeat("]", 200)
+	match := policy.Check("request", map[string]interface{}{"payload": json.RawMessage(payload)})
+	if match == nil || match.MatchedText != "agency.gov" {
+		t.Fatalf("deep JSON value was not checked: %+v", match)
+	}
+}
+
+func TestRuleOrderingAndInputCoverage(t *testing.T) {
+	config := Config{Enabled: true, Rules: []Rule{
+		{ID: "first", Name: "First", Enabled: true, Pattern: "payload-risk", Message: "{rule}/{tool}/{match}"},
+		{ID: "second", Name: "Second", Enabled: true, Pattern: "tool-risk"},
+	}}
+	policy, _ := Compile(config)
+	match := policy.Check("tool-risk", map[string]interface{}{"value": "payload-risk"})
+	if match == nil || match.RuleID != "first" || match.Message != "First/tool-risk/payload-risk" {
+		t.Fatalf("rule order or reminder incorrect: %+v", match)
+	}
+	if match = policy.Check("tool-risk", nil); match == nil || match.RuleID != "second" || match.Message == "" {
+		t.Fatalf("tool name not checked: %+v", match)
+	}
+	config.Rules[0].Pattern = `"port":443`
+	policy, _ = Compile(config)
+	if match = policy.Check("request", map[string]interface{}{"port": 443}); match == nil || match.MatchedText != `"port":443` {
+		t.Fatalf("serialized arguments not checked: %+v", match)
+	}
+	config.Rules[0].Pattern = `^risk.+$`
+	policy, _ = Compile(config)
+	if match = policy.Check("request", map[string]interface{}{"z": "risk-z", "a": "risk-a"}); match == nil || match.MatchedText != "risk-a" {
+		t.Fatalf("field traversal is not deterministic: %+v", match)
+	}
+}
+
+func TestTemplateReplacementDoesNotExpandMatchedText(t *testing.T) {
+	config := Config{Enabled: true, Rules: []Rule{{
+		ID: "template", Name: "Rule", Enabled: true, Pattern: `\{tool\}`, Message: "{match}; {tool}; {rule}",
+	}}}
+	policy, _ := Compile(config)
+	match := policy.Check("request", map[string]interface{}{"value": "{tool}"})
+	if match == nil || match.Message != "{tool}; request; Rule" {
+		t.Fatalf("template expansion was recursive: %+v", match)
+	}
+}
+
+func TestPercentDecodingBudgetAppliesToEachInput(t *testing.T) {
+	config := Config{Enabled: true, Rules: []Rule{{
+		ID: "domain", Name: "Domain", Enabled: true, Pattern: `^agency\.gov$`,
+	}}}
+	policy, _ := Compile(config)
+	// A deeply encoded value may finish its decoding budget at an intermediate
+	// string, but must not prevent an independent field from decoding further.
+	match := policy.Check("request", map[string]interface{}{
+		"a": "agency%2525252egov", "b": "agency%2egov",
+	})
+	if match == nil || match.MatchedText != "agency.gov" {
+		t.Fatalf("an earlier value suppressed decoding of another field: %+v", match)
+	}
+}
+
+func TestDisabledSettings(t *testing.T) {
+	config := DefaultConfig()
+	config.Enabled = false
+	policy, _ := Compile(config)
+	args := map[string]interface{}{"target": "agency.gov"}
+	if policy.Check("request", args) != nil {
+		t.Fatal("disabled policy blocked the call")
+	}
+	config.Enabled = true
+	config.Rules[0].Enabled = false
+	policy, _ = Compile(config)
+	if policy.Check("request", args) != nil {
+		t.Fatal("disabled rule blocked the call")
+	}
+	config.Rules = []Rule{}
+	policy, _ = Compile(config)
+	if policy.Check("request", args) != nil {
+		t.Fatal("empty policy blocked the call")
+	}
+}
+
+func TestCompileValidation(t *testing.T) {
+	for _, test := range []struct {
+		name   string
+		change func(*Config)
+	}{
+		{"invalid regex", func(c *Config) { c.Rules[0].Pattern = "[" }},
+		{"disabled invalid regex", func(c *Config) { c.Enabled = false; c.Rules[0].Enabled = false; c.Rules[0].Pattern = "[" }},
+		{"unsupported lookahead", func(c *Config) { c.Rules[0].Pattern = "x(?=y)" }},
+		{"empty regex", func(c *Config) { c.Rules[0].Pattern = " " }},
+		{"oversized regex", func(c *Config) { c.Rules[0].Pattern = strings.Repeat("a", MaxPatternLength+1) }},
+		{"empty id", func(c *Config) { c.Rules[0].ID = " " }},
+		{"padded id", func(c *Config) { c.Rules[0].ID = " id" }},
+		{"oversized id", func(c *Config) { c.Rules[0].ID = strings.Repeat("a", MaxIDLength+1) }},
+		{"duplicate id", func(c *Config) { c.Rules = append(c.Rules, c.Rules[0]) }},
+		{"empty name", func(c *Config) { c.Rules[0].Name = " " }},
+		{"oversized name", func(c *Config) { c.Rules[0].Name = strings.Repeat("a", MaxNameLength+1) }},
+		{"oversized message", func(c *Config) { c.Rules[0].Message = strings.Repeat("a", MaxMessageLength+1) }},
+		{"too many rules", func(c *Config) { c.Rules = make([]Rule, MaxRules+1) }},
+	} {
+		t.Run(test.name, func(t *testing.T) {
+			config := DefaultConfig()
+			test.change(&config)
+			if _, err := Compile(config); err == nil {
+				t.Fatal("expected validation error")
+			}
+		})
+	}
+}
+
+func TestManagerUsesImmutableValidatedSnapshots(t *testing.T) {
+	config := DefaultConfig()
+	manager, err := NewManager(config)
+	if err != nil {
+		t.Fatal(err)
+	}
+	args := map[string]interface{}{"target": "agency.gov"}
+	config.Rules[0].Pattern = "safe"
+	snapshot := manager.Config()
+	snapshot.Rules[0].Enabled = false
+	if manager.Check("request", args) == nil {
+		t.Fatal("external config mutation changed the active policy")
+	}
+	invalid := DefaultConfig()
+	invalid.Rules[0].Pattern = "["
+	if err := manager.Update(invalid); err == nil || manager.Check("request", args) == nil {
+		t.Fatal("invalid update did not preserve protection")
+	}
+	disabled := DefaultConfig()
+	disabled.Enabled = false
+	if err := manager.Update(disabled); err != nil || manager.Check("request", args) != nil {
+		t.Fatal("valid update did not take effect")
+	}
+}
+
+func TestManagerConcurrentUpdatesAndChecks(t *testing.T) {
+	manager, _ := NewManager(DefaultConfig())
+	var workers sync.WaitGroup
+	for worker := 0; worker < 4; worker++ {
+		workers.Add(1)
+		go func() {
+			defer workers.Done()
+			for i := 0; i < 100; i++ {
+				if match := manager.Check("request", map[string]interface{}{"target": "agency.gov"}); match == nil {
+					t.Error("an update created an unprotected interval")
+					return
+				}
+				config := manager.Config()
+				config.Rules[0].Message = "Block {match}"
+				if err := manager.Update(config); err != nil {
+					t.Error(err)
+					return
+				}
+			}
+		}()
+	}
+	workers.Wait()
+}