feat: add configurable tool call blocking and monitoring

This commit is contained in:
Ed1s0nZ
2026-09-08 09:44:32 +08:00
parent c70da22de7
commit 6ad9ea2d13
54 changed files with 4635 additions and 152 deletions
+34 -5
View File
@@ -512,6 +512,7 @@ type ToolExecutionResult struct {
Result string
ExecutionID string
IsError bool
Blocked bool
}
func buildToolFailureMessage(toolName, detail string, err error) string {
@@ -612,6 +613,7 @@ func (a *Agent) executeToolViaMCP(ctx context.Context, toolName string, args map
Result: resultStr,
ExecutionID: executionID,
IsError: result != nil && result.IsError,
Blocked: result != nil && result.Blocked,
}, nil
}
@@ -815,6 +817,10 @@ func (a *Agent) UpdateMCPExecutionDisplayResult(executionID, resultText string)
tr := &mcp.ToolResult{
Content: []mcp.Content{{Type: "text", Text: text}},
}
if exec := a.mcpExecution(executionID); exec != nil && exec.Result != nil {
tr.IsError = exec.Result.IsError
tr.Blocked = exec.Result.Blocked
}
if a.mcpServer != nil {
_ = a.mcpServer.UpdateToolExecutionResult(executionID, tr)
}
@@ -823,16 +829,39 @@ func (a *Agent) UpdateMCPExecutionDisplayResult(executionID, resultText string)
// MCPExecutionResultText returns the monitor-facing result text after storage
// guards such as large-output spilling have been applied.
func (a *Agent) MCPExecutionResultText(executionID string) string {
if a == nil || a.mcpServer == nil || strings.TrimSpace(executionID) == "" {
return ""
}
exec, ok := a.mcpServer.GetExecution(executionID)
if !ok || exec == nil || exec.Result == nil {
exec := a.mcpExecution(executionID)
if exec == nil || exec.Result == nil {
return ""
}
return mcp.ToolResultPlainText(exec.Result)
}
// MCPExecutionStatus returns the recorded outcome independently of model-facing
// text reduction, which can remove the original refusal wording.
func (a *Agent) MCPExecutionStatus(executionID string) string {
if exec := a.mcpExecution(executionID); exec != nil {
return exec.Status
}
return ""
}
func (a *Agent) mcpExecution(executionID string) *mcp.ToolExecution {
if a == nil || strings.TrimSpace(executionID) == "" {
return nil
}
if a.mcpServer != nil {
if exec, ok := a.mcpServer.GetExecution(executionID); ok && exec != nil {
return exec
}
}
if a.externalMCPMgr != nil {
if exec, ok := a.externalMCPMgr.GetExecution(executionID); ok {
return exec
}
}
return nil
}
// CancelMCPToolExecutionWithNote 取消一次进行中的 MCP 工具(先内部后外部),与监控页「终止工具」一致;note 非空时合并进返回给模型的文本。
func (a *Agent) CancelMCPToolExecutionWithNote(executionID, note string) bool {
executionID = strings.TrimSpace(executionID)
+11
View File
@@ -33,6 +33,7 @@ import (
"cyberstrike-ai/internal/robot"
"cyberstrike-ai/internal/security"
"cyberstrike-ai/internal/skillpackage"
"cyberstrike-ai/internal/toolguard"
"github.com/gin-gonic/gin"
"github.com/google/uuid"
@@ -76,6 +77,10 @@ type App struct {
// New 创建新应用
func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error) {
toolGuard, err := toolguard.NewManager(cfg.EffectiveToolGuard())
if err != nil {
return nil, fmt.Errorf("初始化调用拦截规则: %w", err)
}
if err := multiagent.InitADK(); err != nil {
return nil, fmt.Errorf("初始化 Eino ADK: %w", err)
}
@@ -147,6 +152,7 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
// 创建MCP服务器(带数据库持久化)
mcpServer := mcp.NewServerWithStorage(log.Logger, db)
mcpServer.SetToolAuthorizer(mcpToolAuthorizer(db))
mcpServer.SetToolGuard(toolGuard)
mcpServer.ConfigureHTTPToolCallTimeoutFromAgentMinutes(cfg.Agent.ToolTimeoutMinutes)
mcpServer.ConfigureToolWaitTimeoutSeconds(cfg.Agent.ToolWaitTimeoutSeconds)
mcpServer.ConfigureToolResultMaxBytes(cfg.MultiAgent.EinoMiddleware.ReductionMaxLengthForTruncEffective())
@@ -170,6 +176,7 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
// 创建外部MCP管理器(使用与内部MCP服务器相同的存储)
externalMCPMgr := mcp.NewExternalMCPManagerWithStorage(log.Logger, db)
externalMCPMgr.SetToolAuthorizer(externalMCPToolAuthorizer())
externalMCPMgr.SetToolGuard(toolGuard)
externalMCPMgr.ConfigureToolWaitTimeoutSeconds(cfg.Agent.ToolWaitTimeoutSeconds)
externalMCPMgr.ConfigureToolResultMaxBytes(cfg.MultiAgent.EinoMiddleware.ReductionMaxLengthForTruncEffective())
externalMCPMgr.ConfigureToolResultSpillRoot(cfg.MultiAgent.EinoMiddleware.ReductionRootDir)
@@ -412,6 +419,7 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
registerWebshellManagementTools(mcpServer, db, webshellHandler, log.Logger)
configHandler := handler.NewConfigHandler(configPath, cfg, mcpServer, executor, agent, attackChainHandler, externalMCPMgr, log.Logger)
configHandler.SetDB(db)
configHandler.SetToolGuard(toolGuard)
configHandler.SetAudit(auditSvc)
agentHandler.SetHitlToolWhitelistSaver(configHandler)
agentHandler.SetHitlAuditStrategySaver(configHandler)
@@ -1054,6 +1062,9 @@ func setupRoutes(
// 配置管理
protected.GET("/config", configHandler.GetConfig)
protected.GET("/tool-guard", configHandler.GetToolGuard)
protected.PUT("/tool-guard", configHandler.UpdateToolGuard)
protected.POST("/tool-guard/test", configHandler.TestToolGuard)
protected.GET("/config/tools", configHandler.GetTools)
protected.GET("/config/tools/:name/schema", configHandler.GetToolSchema)
protected.PUT("/config", configHandler.UpdateConfig)
+10
View File
@@ -13,6 +13,7 @@ import (
"strings"
"cyberstrike-ai/internal/termout"
"cyberstrike-ai/internal/toolguard"
"gopkg.in/yaml.v3"
)
@@ -30,6 +31,7 @@ type Config struct {
Shodan SpaceSearchConfig `yaml:"shodan,omitempty" json:"shodan,omitempty"`
Agent AgentConfig `yaml:"agent"`
Hitl HitlConfig `yaml:"hitl,omitempty" json:"hitl,omitempty"`
ToolGuard *toolguard.Config `yaml:"tool_guard,omitempty" json:"tool_guard,omitempty"`
Security SecurityConfig `yaml:"security"`
Database DatabaseConfig `yaml:"database"`
Auth AuthConfig `yaml:"auth"`
@@ -1439,6 +1441,14 @@ func Load(path string) (*Config, error) {
if err := yaml.Unmarshal(data, &cfg); err != nil {
return nil, fmt.Errorf("解析配置文件失败: %w", err)
}
if cfg.ToolGuard != nil {
if err := validateToolGuardYAML(data); err != nil {
return nil, fmt.Errorf("调用拦截配置无效: %w", err)
}
}
if _, err := toolguard.Compile(cfg.EffectiveToolGuard()); err != nil {
return nil, fmt.Errorf("调用拦截配置无效: %w", err)
}
if cfg.Auth.SessionDurationHours <= 0 {
cfg.Auth.SessionDurationHours = 12
+39
View File
@@ -0,0 +1,39 @@
package config
import (
"fmt"
"cyberstrike-ai/internal/toolguard"
"gopkg.in/yaml.v3"
)
// EffectiveToolGuard enables the default government-domain protection for old
// configurations as well as new installs. An explicit config may disable it.
func (c *Config) EffectiveToolGuard() toolguard.Config {
if c.ToolGuard == nil {
return toolguard.DefaultConfig()
}
return *c.ToolGuard
}
// validateToolGuardYAML requires an explicit decision for both protection and
// its rules whenever a non-null section is supplied. Otherwise a typo or partial
// section could silently turn the enabled-by-default protection off. Pointer
// fields distinguish false/[] from omitted or null values, and the YAML decoder
// continues to support aliases and merged configuration mappings.
func validateToolGuardYAML(data []byte) error {
var document struct {
ToolGuard *struct {
Enabled *bool `yaml:"enabled"`
Rules *[]toolguard.Rule `yaml:"rules"`
} `yaml:"tool_guard"`
}
if err := yaml.Unmarshal(data, &document); err != nil {
return err
}
if section := document.ToolGuard; section != nil && (section.Enabled == nil || section.Rules == nil) {
return fmt.Errorf("tool_guard 必须明确提供 enabled 和 rules;清空规则请提供空数组")
}
return nil
}
+45
View File
@@ -0,0 +1,45 @@
package config
import (
"os"
"path/filepath"
"testing"
)
func TestLoadToolGuardDefaultsAndValidation(t *testing.T) {
for _, tc := range []struct {
name, yaml string
enabled, wantErr bool
}{
{"legacy config", "server: {port: 8080}\n", true, false},
{"null section", "tool_guard: null\n", true, false},
{"implicit null section", "tool_guard:\n", true, false},
{"explicit off", "tool_guard: {enabled: false, rules: []}\n", false, false},
{"explicit empty", "tool_guard: {enabled: true, rules: []}\n", true, false},
{"merged explicit config", "guard_defaults: &guard_defaults {enabled: false, rules: []}\ntool_guard: {<<: *guard_defaults}\n", false, false},
{"empty section", "tool_guard: {}\n", false, true},
{"missing enabled", "tool_guard: {rules: []}\n", false, true},
{"null enabled", "tool_guard: {enabled: null, rules: []}\n", false, true},
{"missing rules while off", "tool_guard: {enabled: false}\n", false, true},
{"missing rules while on", "tool_guard: {enabled: true}\n", false, true},
{"null rules", "tool_guard: {enabled: false, rules: null}\n", false, true},
{"mistyped enabled field", "tool_guard: {enable: false, rules: []}\n", false, true},
{"malformed rules while off", "tool_guard: {enabled: false, rules: disabled}\n", false, true},
{"malformed rule while off", "tool_guard: {enabled: false, rules: [invalid]}\n", false, true},
{"invalid pattern", "tool_guard:\n enabled: false\n rules:\n - {id: invalid, name: invalid, enabled: false, pattern: '['}\n", false, true},
} {
t.Run(tc.name, func(t *testing.T) {
path := filepath.Join(t.TempDir(), "config.yaml")
if err := os.WriteFile(path, []byte(tc.yaml), 0600); err != nil {
t.Fatal(err)
}
cfg, err := Load(path)
if (err != nil) != tc.wantErr {
t.Fatalf("load error: %v", err)
}
if err == nil && cfg.EffectiveToolGuard().Enabled != tc.enabled {
t.Fatal("wrong effective enabled state")
}
})
}
}
+120
View File
@@ -0,0 +1,120 @@
package database
import (
"fmt"
"path/filepath"
"testing"
"time"
"cyberstrike-ai/internal/mcp"
"go.uber.org/zap"
)
func TestBlockedExecutionPersistenceStatsAndReconciliation(t *testing.T) {
db, conversationID, _ := setupProcessDetailsSummaryTest(t)
now := time.Now()
for _, status := range []string{"completed", "failed", "blocked", "cancelled"} {
result := &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: "policy message"}}, IsError: status != "completed", Blocked: status == "blocked"}
if err := db.SaveToolExecution(&mcp.ToolExecution{ID: status, ToolName: "test", Status: status, Result: result, StartTime: now.Add(-time.Minute), EndTime: &now, ConversationID: conversationID}); err != nil {
t.Fatal(err)
}
}
if err := db.UpdateToolStats("test", 4, 1, 1, &now); err != nil {
t.Fatal(err)
}
if err := db.UpdateToolExecutionResult("blocked", &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: "reduced output"}}}); err != nil {
t.Fatal(err)
}
reloaded, err := db.GetToolExecution("blocked")
if err != nil || reloaded.Status != "blocked" || !reloaded.Result.Blocked || !reloaded.Result.IsError || reloaded.Result.Content[0].Text != "reduced output" {
t.Fatalf("reduction/storage lost blocked classification: %#v err=%v", reloaded, err)
}
count, err := db.CancelOrphanedRunningToolExecutions(now, "restart")
if err != nil || count != 0 {
t.Fatalf("terminal blocks reclassified as orphaned: count=%d err=%v", count, err)
}
page, err := db.LoadToolExecutionListPage(0, 10, "blocked", "")
if err != nil || len(page) != 1 || page[0].ID != "blocked" {
t.Fatalf("blocked status filter failed: %#v err=%v", page, err)
}
summary, err := db.LoadToolStatsSummary(1)
if err != nil {
t.Fatal(err)
}
if summary.Summary.TotalCalls != 4 || summary.Summary.SuccessCalls != 1 || summary.Summary.FailedCalls != 1 || summary.Summary.BlockedCalls != 1 || summary.TopTools[0].BlockedCalls != 1 {
t.Fatalf("incorrect summary: %#v top=%#v", summary.Summary, summary.TopTools)
}
stats, err := db.LoadToolStats()
if err != nil || stats["test"].BlockedCalls != 1 || stats["test"].FailedCalls != 1 {
t.Fatalf("incorrect legacy stats: %#v err=%v", stats, err)
}
for _, daily := range []bool{false, true} {
buckets, err := db.LoadCallsTimeline(now.Add(-time.Hour), daily)
if err != nil || len(buckets) != 1 || buckets[0].Total != 4 || buckets[0].Failed != 1 || buckets[0].Blocked != 1 {
t.Fatalf("incorrect timeline daily=%v: %#v err=%v", daily, buckets, err)
}
}
}
func TestLegacyToolGuardBlockMigrationIsStrictAndIdempotent(t *testing.T) {
path := filepath.Join(t.TempDir(), "legacy-guard.db")
db, err := NewDB(path, zap.NewNop())
if err != nil {
t.Fatal(err)
}
now := time.Now()
refusal := "工具调用已被安全规则拦截:识别到 example.gov,禁止操作。\n规则: 政府网站保护 (government-domains)\n匹配内容: \"example.gov\""
for i, reason := range []string{
refusal,
"upstream returned: " + refusal,
"工具调用已被安全规则拦截:regular error without the envelope",
"工具调用已被安全规则拦截:malformed match\n规则: Rule (id)\n匹配内容: unquoted",
} {
if err := db.SaveToolExecution(&mcp.ToolExecution{ID: fmt.Sprint(i), ToolName: "test", Status: "failed", Error: reason, StartTime: now, EndTime: &now}); err != nil {
t.Fatal(err)
}
}
if err := db.UpdateToolStats("test", 4, 0, 4, &now); err != nil {
t.Fatal(err)
}
if err := db.Close(); err != nil {
t.Fatal(err)
}
for run := 0; run < 2; run++ {
db, err = NewDB(path, zap.NewNop())
if err != nil {
t.Fatal(err)
}
exec, err := db.GetToolExecution("0")
if err != nil || exec.Status != "blocked" || !exec.Result.Blocked || !exec.Result.IsError || exec.Result.Content[0].Text != refusal {
t.Fatalf("migration did not retain refusal: %#v err=%v", exec, err)
}
stats, err := db.LoadToolStats()
if err != nil || stats["test"].TotalCalls != 4 || stats["test"].FailedCalls != 3 || stats["test"].BlockedCalls != 1 {
t.Fatalf("migration run=%d stats=%#v err=%v", run, stats, err)
}
count, err := db.CountToolExecutions("failed", "")
if err != nil || count != 3 {
t.Fatalf("migration changed unrelated failures: count=%d err=%v", count, err)
}
if err := db.Close(); err != nil {
t.Fatal(err)
}
}
}
func TestToolResultStatusFromPayloadDistinguishesBlocked(t *testing.T) {
for _, tc := range []struct {
payload map[string]interface{}
want string
}{
{map[string]interface{}{"blocked": true, "success": false, "isError": true}, "blocked"},
{map[string]interface{}{"status": "blocked", "success": false}, "blocked"},
{map[string]interface{}{"success": false, "isError": true, "result": "工具调用已被安全规则拦截"}, "failed"},
{map[string]interface{}{"success": true}, "completed"},
} {
if got := toolResultStatusFromPayload(tc.payload, "tool_result"); got != tc.want {
t.Fatalf("payload=%#v status=%s want=%s", tc.payload, got, tc.want)
}
}
}
+3
View File
@@ -1637,6 +1637,9 @@ func toolResultStatusFromPayload(payload map[string]interface{}, eventType strin
if eventType != "tool_result" {
return ""
}
if blocked, _ := payload["blocked"].(bool); blocked || strings.EqualFold(processDetailString(payload, "status"), "blocked") {
return "blocked"
}
if status := processDetailString(payload, "status"); strings.EqualFold(status, "background_running") {
return "background_running"
}
+4
View File
@@ -155,6 +155,10 @@ func NewDB(dbPath string, logger *zap.Logger) (*DB, error) {
_ = db.Close()
return nil, fmt.Errorf("初始化表失败: %w", err)
}
if err := database.migrateLegacyToolGuardBlocks(); err != nil {
_ = db.Close()
return nil, fmt.Errorf("迁移历史安全拦截记录失败: %w", err)
}
database.startPassiveCheckpointLoop("conversations")
return database, nil
+33 -9
View File
@@ -91,6 +91,15 @@ func (db *DB) UpdateToolExecutionResult(id string, result *mcp.ToolResult) error
if id == "" || result == nil {
return nil
}
var status string
if err := db.QueryRow(`SELECT status FROM tool_executions WHERE id = ?`, id).Scan(&status); err != nil && err != sql.ErrNoRows {
return err
}
if status == mcp.ToolExecutionStatusBlocked {
copy := *result
copy.Blocked, copy.IsError = true, true
result = &copy
}
resultBytes, err := json.Marshal(result)
if err != nil {
return err
@@ -276,6 +285,7 @@ type ToolStatsSummary struct {
TotalCalls int
SuccessCalls int
FailedCalls int
BlockedCalls int
LastCallTime *time.Time
ToolCount int
}
@@ -304,6 +314,7 @@ func (db *DB) LoadToolStatsSummary(topN int) (*ToolStatsSummaryResult, error) {
SELECT COUNT(*),
COALESCE(SUM(CASE WHEN status = 'completed' THEN 1 ELSE 0 END), 0),
COALESCE(SUM(CASE WHEN status IN ('failed', 'hard_timeout', 'orphaned') THEN 1 ELSE 0 END), 0),
COALESCE(SUM(CASE WHEN status = 'blocked' THEN 1 ELSE 0 END), 0),
MAX(start_time),
COUNT(DISTINCT tool_name)
FROM tool_executions
@@ -313,6 +324,7 @@ func (db *DB) LoadToolStatsSummary(topN int) (*ToolStatsSummaryResult, error) {
&result.Summary.TotalCalls,
&result.Summary.SuccessCalls,
&result.Summary.FailedCalls,
&result.Summary.BlockedCalls,
&lastCallRaw,
&result.Summary.ToolCount,
)
@@ -334,6 +346,7 @@ func (db *DB) LoadToolStatsSummary(topN int) (*ToolStatsSummaryResult, error) {
COUNT(*) AS total_calls,
SUM(CASE WHEN status = 'completed' THEN 1 ELSE 0 END) AS success_calls,
SUM(CASE WHEN status IN ('failed', 'hard_timeout', 'orphaned') THEN 1 ELSE 0 END) AS failed_calls,
SUM(CASE WHEN status = 'blocked' THEN 1 ELSE 0 END) AS blocked_calls,
MAX(start_time) AS last_call_time
FROM tool_executions
GROUP BY tool_name
@@ -354,6 +367,7 @@ func (db *DB) LoadToolStatsSummary(topN int) (*ToolStatsSummaryResult, error) {
&stat.TotalCalls,
&stat.SuccessCalls,
&stat.FailedCalls,
&stat.BlockedCalls,
&lastCallTime,
); err != nil {
db.logger.Warn("加载 Top 工具统计失败", zap.Error(err))
@@ -385,8 +399,9 @@ func (db *DB) LoadToolStatsSummaryForAccess(topN int, access RBACListAccess) (*T
err := db.QueryRow(`SELECT COUNT(*),
COALESCE(SUM(CASE WHEN status = 'completed' THEN 1 ELSE 0 END), 0),
COALESCE(SUM(CASE WHEN status IN ('failed', 'hard_timeout', 'orphaned') THEN 1 ELSE 0 END), 0),
COALESCE(SUM(CASE WHEN status = 'blocked' THEN 1 ELSE 0 END), 0),
MAX(start_time), COUNT(DISTINCT tool_name)`+fromSQL, args...).Scan(
&result.Summary.TotalCalls, &result.Summary.SuccessCalls, &result.Summary.FailedCalls,
&result.Summary.TotalCalls, &result.Summary.SuccessCalls, &result.Summary.FailedCalls, &result.Summary.BlockedCalls,
&lastCall, &result.Summary.ToolCount,
)
if err != nil {
@@ -398,7 +413,8 @@ func (db *DB) LoadToolStatsSummaryForAccess(topN int, access RBACListAccess) (*T
}
rows, err := db.Query(`SELECT tool_name, COUNT(*),
SUM(CASE WHEN status = 'completed' THEN 1 ELSE 0 END),
SUM(CASE WHEN status IN ('failed', 'hard_timeout', 'orphaned') THEN 1 ELSE 0 END), MAX(start_time)`+
SUM(CASE WHEN status IN ('failed', 'hard_timeout', 'orphaned') THEN 1 ELSE 0 END),
SUM(CASE WHEN status = 'blocked' THEN 1 ELSE 0 END), MAX(start_time)`+
fromSQL+` GROUP BY tool_name ORDER BY COUNT(*) DESC, tool_name ASC LIMIT ?`, append(args, topN)...)
if err != nil {
return nil, err
@@ -407,7 +423,7 @@ func (db *DB) LoadToolStatsSummaryForAccess(topN int, access RBACListAccess) (*T
for rows.Next() {
var stat mcp.ToolStats
var last sql.NullString
if err := rows.Scan(&stat.ToolName, &stat.TotalCalls, &stat.SuccessCalls, &stat.FailedCalls, &last); err != nil {
if err := rows.Scan(&stat.ToolName, &stat.TotalCalls, &stat.SuccessCalls, &stat.FailedCalls, &stat.BlockedCalls, &last); err != nil {
return nil, err
}
if last.Valid {
@@ -916,8 +932,11 @@ func (db *DB) SaveToolStats(toolName string, stats *mcp.ToolStats) error {
// LoadToolStats 加载所有工具统计信息
func (db *DB) LoadToolStats() (map[string]*mcp.ToolStats, error) {
query := `
SELECT tool_name, total_calls, success_calls, failed_calls, last_call_time
FROM tool_stats
SELECT stats.tool_name, total_calls, success_calls, failed_calls, last_call_time,
COALESCE(blocked.calls, 0)
FROM tool_stats stats
LEFT JOIN (SELECT tool_name, COUNT(*) AS calls FROM tool_executions WHERE status = 'blocked' GROUP BY tool_name) blocked
ON blocked.tool_name = stats.tool_name
`
rows, err := db.Query(query)
@@ -937,6 +956,7 @@ func (db *DB) LoadToolStats() (map[string]*mcp.ToolStats, error) {
&stat.SuccessCalls,
&stat.FailedCalls,
&lastCallTime,
&stat.BlockedCalls,
)
if err != nil {
db.logger.Warn("加载统计信息失败", zap.Error(err))
@@ -989,6 +1009,7 @@ type CallsTimelineBucket struct {
BucketTime time.Time
Total int
Failed int
Blocked int
}
// truncateCallsTimelineBucket 将时间截断到趋势图桶边界(本地时区,与 handler 侧 truncateToBucket 一致)
@@ -1008,7 +1029,8 @@ func (db *DB) LoadCallsTimeline(since time.Time, dailyBuckets bool) ([]CallsTime
query = `
SELECT date(start_time, 'localtime') AS bucket,
COUNT(*) AS total,
SUM(CASE WHEN status IN ('failed', 'hard_timeout', 'orphaned') THEN 1 ELSE 0 END) AS failed
SUM(CASE WHEN status IN ('failed', 'hard_timeout', 'orphaned') THEN 1 ELSE 0 END) AS failed,
SUM(CASE WHEN status = 'blocked' THEN 1 ELSE 0 END) AS blocked
FROM tool_executions
WHERE start_time >= ?
GROUP BY bucket
@@ -1018,7 +1040,8 @@ func (db *DB) LoadCallsTimeline(since time.Time, dailyBuckets bool) ([]CallsTime
query = `
SELECT strftime('%Y-%m-%d %H:00:00', start_time, 'localtime') AS bucket,
COUNT(*) AS total,
SUM(CASE WHEN status IN ('failed', 'hard_timeout', 'orphaned') THEN 1 ELSE 0 END) AS failed
SUM(CASE WHEN status IN ('failed', 'hard_timeout', 'orphaned') THEN 1 ELSE 0 END) AS failed,
SUM(CASE WHEN status = 'blocked' THEN 1 ELSE 0 END) AS blocked
FROM tool_executions
WHERE start_time >= ?
GROUP BY bucket
@@ -1035,8 +1058,8 @@ func (db *DB) LoadCallsTimeline(since time.Time, dailyBuckets bool) ([]CallsTime
buckets := make([]CallsTimelineBucket, 0)
for rows.Next() {
var bucketStr string
var total, failed int
if err := rows.Scan(&bucketStr, &total, &failed); err != nil {
var total, failed, blocked int
if err := rows.Scan(&bucketStr, &total, &failed, &blocked); err != nil {
db.logger.Warn("加载调用趋势失败", zap.Error(err))
continue
}
@@ -1049,6 +1072,7 @@ func (db *DB) LoadCallsTimeline(since time.Time, dailyBuckets bool) ([]CallsTime
BucketTime: bucketTime,
Total: total,
Failed: failed,
Blocked: blocked,
})
}
return buckets, nil
+84
View File
@@ -0,0 +1,84 @@
package database
import (
"encoding/json"
"strconv"
"strings"
"cyberstrike-ai/internal/mcp"
)
const legacyToolGuardPrefix = "工具调用已被安全规则拦截"
// Only the exact envelope emitted by the old local guard is recognized here.
// New executions use the structured marker and never infer policy from text.
func isLegacyToolGuardRefusal(text string) bool {
if !strings.HasPrefix(text, legacyToolGuardPrefix+"") && !strings.HasPrefix(text, legacyToolGuardPrefix+"\n规则: ") {
return false
}
matchIndex := strings.LastIndex(text, "\n匹配内容: ")
if matchIndex < 0 {
return false
}
if _, err := strconv.Unquote(text[matchIndex+len("\n匹配内容: "):]); err != nil {
return false
}
ruleIndex := strings.LastIndex(text[:matchIndex], "\n规则: ")
if ruleIndex < 0 {
return false
}
rule := text[ruleIndex+len("\n规则: ") : matchIndex]
idIndex := strings.LastIndex(rule, " (")
return idIndex > 0 && strings.HasSuffix(rule, ")") && len(rule[idIndex+2:len(rule)-1]) > 0 && !strings.Contains(rule, "\n")
}
// migrateLegacyToolGuardBlocks is idempotent because only failed records qualify.
// Keeping status and accumulated failure counts in one transaction makes monitor
// filters, badges and statistics agree immediately after upgrading.
func (db *DB) migrateLegacyToolGuardBlocks() error {
tx, err := db.Begin()
if err != nil {
return err
}
defer tx.Rollback()
rows, err := tx.Query(`SELECT id, tool_name, error, COALESCE(result, '') FROM tool_executions WHERE status = 'failed' AND error LIKE ?`, legacyToolGuardPrefix+"%")
if err != nil {
return err
}
type record struct{ id, tool, reason, result string }
var records []record
for rows.Next() {
var r record
if err := rows.Scan(&r.id, &r.tool, &r.reason, &r.result); err != nil {
rows.Close()
return err
}
if isLegacyToolGuardRefusal(r.reason) {
records = append(records, r)
}
}
err = rows.Err()
rows.Close()
if err != nil {
return err
}
for _, r := range records {
var result mcp.ToolResult
_ = json.Unmarshal([]byte(r.result), &result)
if len(result.Content) == 0 {
result.Content = []mcp.Content{{Type: "text", Text: r.reason}}
}
result.Blocked, result.IsError = true, true
encoded, err := json.Marshal(result)
if err != nil {
return err
}
if _, err := tx.Exec(`UPDATE tool_executions SET status = 'blocked', result = ? WHERE id = ?`, string(encoded), r.id); err != nil {
return err
}
if _, err := tx.Exec(`UPDATE tool_stats SET failed_calls = MAX(0, failed_calls - 1) WHERE tool_name = ?`, r.tool); err != nil {
return err
}
}
return tx.Commit()
}
+11 -7
View File
@@ -23,6 +23,7 @@ import (
"cyberstrike-ai/internal/mcp/builtin"
"cyberstrike-ai/internal/openai"
"cyberstrike-ai/internal/security"
"cyberstrike-ai/internal/toolguard"
"github.com/cloudwego/eino/schema"
"github.com/gin-gonic/gin"
@@ -95,6 +96,7 @@ type ConfigHandler struct {
db *database.DB
logger *zap.Logger
mu sync.RWMutex
toolGuard *toolguard.Manager
lastEmbeddingConfig *config.EmbeddingConfig // 上一次的嵌入模型配置(用于检测变更)
}
@@ -347,13 +349,13 @@ func (h *ConfigHandler) GetConfig(c *gin.Context) {
subAgentCount = len(agents.MergeYAMLAndMarkdown(h.config.MultiAgent.SubAgents, load.SubAgents))
}
multiPub := config.MultiAgentPublic{
Enabled: h.config.MultiAgent.Enabled,
RobotDefaultAgentMode: config.NormalizeRobotAgentMode(h.config.MultiAgent),
BatchUseMultiAgent: h.config.MultiAgent.BatchUseMultiAgent,
SubAgentCount: subAgentCount,
Orchestration: config.NormalizeMultiAgentOrchestration(h.config.MultiAgent.Orchestration),
PlanExecuteLoopMaxIterations: h.config.MultiAgent.PlanExecuteLoopMaxIterations,
SummarizationUserIntentLedgerMaxRunes: h.config.MultiAgent.EinoMiddleware.SummarizationUserIntentLedgerMaxRunesEffective(),
Enabled: h.config.MultiAgent.Enabled,
RobotDefaultAgentMode: config.NormalizeRobotAgentMode(h.config.MultiAgent),
BatchUseMultiAgent: h.config.MultiAgent.BatchUseMultiAgent,
SubAgentCount: subAgentCount,
Orchestration: config.NormalizeMultiAgentOrchestration(h.config.MultiAgent.Orchestration),
PlanExecuteLoopMaxIterations: h.config.MultiAgent.PlanExecuteLoopMaxIterations,
SummarizationUserIntentLedgerMaxRunes: h.config.MultiAgent.EinoMiddleware.SummarizationUserIntentLedgerMaxRunesEffective(),
SummarizationUserIntentLedgerEntryMaxRunes: h.config.MultiAgent.EinoMiddleware.SummarizationUserIntentLedgerEntryMaxRunesEffective(),
LatestUserMessageMaxRunes: h.config.MultiAgent.EinoMiddleware.LatestUserMessageMaxRunesEffective(),
LatestUserMessageHeadRunes: h.config.MultiAgent.EinoMiddleware.LatestUserMessageHeadRunesEffective(),
@@ -1746,6 +1748,8 @@ func (h *ConfigHandler) ApplyConfig(c *gin.Context) {
// saveConfig 保存配置到文件
func (h *ConfigHandler) saveConfig() error {
configFileMu.Lock()
defer configFileMu.Unlock()
h.config.NormalizeAIProviderProfiles()
// 读取现有配置文件并创建备份
+8
View File
@@ -0,0 +1,8 @@
package handler
import "sync"
// configFileMu serializes complete read-modify-write transactions across
// handlers that share config.yaml. Per-handler locks cannot prevent lost
// updates when another settings page saves a different YAML section.
var configFileMu sync.Mutex
+2
View File
@@ -379,6 +379,8 @@ func (h *ExternalMCPHandler) isEnabled(cfg config.ExternalMCPServerConfig) bool
// saveConfig 保存配置到文件
func (h *ExternalMCPHandler) saveConfig() error {
configFileMu.Lock()
defer configFileMu.Unlock()
data, err := os.ReadFile(h.configPath)
if err != nil {
return fmt.Errorf("读取配置文件失败: %w", err)
+19 -8
View File
@@ -76,6 +76,7 @@ type MonitorStatsSummary struct {
TotalCalls int `json:"totalCalls"`
SuccessCalls int `json:"successCalls"`
FailedCalls int `json:"failedCalls"`
BlockedCalls int `json:"blockedCalls"`
LastCallTime *time.Time `json:"lastCallTime,omitempty"`
ToolCount int `json:"toolCount"`
}
@@ -171,6 +172,8 @@ func summarizeAccessibleExecutionPage(executions []*mcp.ToolExecution, topN int)
stat.FailedCalls++
} else if exec.Status == "completed" {
stat.SuccessCalls++
} else if exec.Status == mcp.ToolExecutionStatusBlocked {
stat.BlockedCalls++
}
started := exec.StartTime
if stat.LastCallTime == nil || started.After(*stat.LastCallTime) {
@@ -448,6 +451,7 @@ func dbStatsSummaryToMonitor(result *database.ToolStatsSummaryResult) *MonitorSt
TotalCalls: result.Summary.TotalCalls,
SuccessCalls: result.Summary.SuccessCalls,
FailedCalls: result.Summary.FailedCalls,
BlockedCalls: result.Summary.BlockedCalls,
ToolCount: result.Summary.ToolCount,
}
if result.Summary.LastCallTime != nil {
@@ -472,6 +476,7 @@ func summarizeToolStats(stats map[string]*mcp.ToolStats, topN int) (*MonitorStat
summary.TotalCalls += stat.TotalCalls
summary.SuccessCalls += stat.SuccessCalls
summary.FailedCalls += stat.FailedCalls
summary.BlockedCalls += stat.BlockedCalls
if stat.LastCallTime != nil && (summary.LastCallTime == nil || stat.LastCallTime.After(*summary.LastCallTime)) {
t := *stat.LastCallTime
summary.LastCallTime = &t
@@ -528,6 +533,7 @@ func (h *MonitorHandler) loadStatsMap() map[string]*mcp.ToolStats {
existing.TotalCalls += v.TotalCalls
existing.SuccessCalls += v.SuccessCalls
existing.FailedCalls += v.FailedCalls
existing.BlockedCalls += v.BlockedCalls
// 使用最新的调用时间
if v.LastCallTime != nil && (existing.LastCallTime == nil || v.LastCallTime.After(*existing.LastCallTime)) {
existing.LastCallTime = v.LastCallTime
@@ -734,9 +740,10 @@ func (h *MonitorHandler) GetStats(c *gin.Context) {
// CallsTimelinePoint 调用趋势数据点
type CallsTimelinePoint struct {
T time.Time `json:"t"`
Total int `json:"total"`
Failed int `json:"failed"`
T time.Time `json:"t"`
Total int `json:"total"`
Failed int `json:"failed"`
Blocked int `json:"blocked"`
}
// CallsTimelineSummary 调用趋势汇总
@@ -778,7 +785,7 @@ func truncateToBucket(t time.Time, bucketSize time.Duration, dailyBuckets bool)
return t.Truncate(bucketSize)
}
func buildCallsTimelinePoints(cfg callsTimelineConfig, buckets map[time.Time]struct{ total, failed int }) []CallsTimelinePoint {
func buildCallsTimelinePoints(cfg callsTimelineConfig, buckets map[time.Time]struct{ total, failed, blocked int }) []CallsTimelinePoint {
now := time.Now()
start := truncateToBucket(now.Add(-cfg.duration), cfg.bucketSize, cfg.dailyBuckets)
end := truncateToBucket(now, cfg.bucketSize, cfg.dailyBuckets)
@@ -787,9 +794,10 @@ func buildCallsTimelinePoints(cfg callsTimelineConfig, buckets map[time.Time]str
for current := start; !current.After(end); current = current.Add(cfg.bucketSize) {
val := buckets[current]
points = append(points, CallsTimelinePoint{
T: current,
Total: val.total,
Failed: val.failed,
T: current,
Total: val.total,
Failed: val.failed,
Blocked: val.blocked,
})
}
return points
@@ -797,7 +805,7 @@ func buildCallsTimelinePoints(cfg callsTimelineConfig, buckets map[time.Time]str
func (h *MonitorHandler) loadCallsTimeline(cfg callsTimelineConfig) []CallsTimelinePoint {
since := time.Now().Add(-cfg.duration)
bucketMap := make(map[time.Time]struct{ total, failed int })
bucketMap := make(map[time.Time]struct{ total, failed, blocked int })
if h.db != nil {
dbBuckets, err := h.db.LoadCallsTimeline(since, cfg.dailyBuckets)
@@ -809,6 +817,7 @@ func (h *MonitorHandler) loadCallsTimeline(cfg callsTimelineConfig) []CallsTimel
entry := bucketMap[key]
entry.total += b.Total
entry.failed += b.Failed
entry.blocked += b.Blocked
bucketMap[key] = entry
}
return buildCallsTimelinePoints(cfg, bucketMap)
@@ -824,6 +833,8 @@ func (h *MonitorHandler) loadCallsTimeline(cfg callsTimelineConfig) []CallsTimel
entry.total++
if monitorStatusCountsAsFailed(exec.Status) {
entry.failed++
} else if exec.Status == mcp.ToolExecutionStatusBlocked {
entry.blocked++
}
bucketMap[key] = entry
}
+175
View File
@@ -0,0 +1,175 @@
package handler
import (
"bytes"
"encoding/json"
"fmt"
"io"
"net/http"
"os"
"path/filepath"
"cyberstrike-ai/internal/toolguard"
"github.com/gin-gonic/gin"
"gopkg.in/yaml.v3"
)
func (h *ConfigHandler) SetToolGuard(manager *toolguard.Manager) {
h.mu.Lock()
defer h.mu.Unlock()
h.toolGuard = manager
}
func (h *ConfigHandler) GetToolGuard(c *gin.Context) {
h.mu.RLock()
defer h.mu.RUnlock()
if h.toolGuard == nil {
c.JSON(http.StatusServiceUnavailable, gin.H{"error": "调用拦截服务未初始化"})
return
}
c.JSON(http.StatusOK, h.toolGuard.Config())
}
// decodeToolGuardRequest bounds both config and dry-run inputs, rejects unknown
// fields and trailing JSON, and never invokes an actual tool.
func decodeToolGuardRequest(c *gin.Context, dst interface{}) error {
c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, 1<<20)
decoder := json.NewDecoder(c.Request.Body)
decoder.DisallowUnknownFields()
decoder.UseNumber()
if err := decoder.Decode(dst); err != nil {
return err
}
if err := decoder.Decode(new(interface{})); err != io.EOF {
return fmt.Errorf("请求必须只包含一个 JSON 对象")
}
return nil
}
func (h *ConfigHandler) UpdateToolGuard(c *gin.Context) {
var req struct {
Enabled *bool `json:"enabled"`
Rules *[]toolguard.Rule `json:"rules"`
}
if err := decodeToolGuardRequest(c, &req); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": "无效的调用拦截配置: " + err.Error()})
return
}
if req.Enabled == nil || req.Rules == nil {
c.JSON(http.StatusBadRequest, gin.H{"error": "必须明确提供 enabled 和 rules;清空规则请提供空数组"})
return
}
cfg := toolguard.Config{Enabled: *req.Enabled, Rules: *req.Rules}
if _, err := toolguard.Compile(cfg); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
h.mu.Lock()
defer h.mu.Unlock()
if h.toolGuard == nil {
c.JSON(http.StatusServiceUnavailable, gin.H{"error": "调用拦截服务未初始化"})
return
}
// Commit the file first; a validation/write failure must leave the current
// effective policy and in-memory config intact.
if err := h.saveToolGuardConfig(cfg); err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": "保存调用拦截配置失败: " + err.Error()})
return
}
if err := h.toolGuard.Update(cfg); err != nil {
// The same immutable input was compiled above, so this cannot fail
// unless validation gains an additional runtime dependency.
c.JSON(http.StatusInternalServerError, gin.H{"error": "应用调用拦截配置失败: " + err.Error()})
return
}
h.config.ToolGuard = &cfg
if h.audit != nil {
h.audit.RecordOK(c, "config", "tool_guard_update", "更新调用拦截规则", "config", "tool_guard", map[string]interface{}{
"enabled": cfg.Enabled, "rule_count": len(cfg.Rules),
})
}
c.JSON(http.StatusOK, h.toolGuard.Config())
}
func (h *ConfigHandler) TestToolGuard(c *gin.Context) {
var req struct {
Config *toolguard.Config `json:"config"`
ToolName string `json:"toolName"`
Arguments map[string]interface{} `json:"arguments"`
}
if err := decodeToolGuardRequest(c, &req); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": "无效的试匹配参数: " + err.Error()})
return
}
if req.Config == nil {
c.JSON(http.StatusBadRequest, gin.H{"error": "请提供待测试的 config"})
return
}
policy, err := toolguard.Compile(*req.Config)
if err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
if match := policy.Check(req.ToolName, req.Arguments); match != nil {
c.JSON(http.StatusOK, gin.H{"blocked": true, "match": match})
return
}
c.JSON(http.StatusOK, gin.H{"blocked": false})
}
// saveToolGuardConfig changes only this YAML section, preserving unrelated
// settings/comments and file permissions. Rename makes the write atomic.
// h.mu protects the runtime configuration; configFileMu also covers independent
// writers such as ExternalMCPHandler.
func (h *ConfigHandler) saveToolGuardConfig(cfg toolguard.Config) error {
configFileMu.Lock()
defer configFileMu.Unlock()
path, err := filepath.EvalSymlinks(h.configPath)
if err != nil {
return err
}
doc, err := loadYAMLDocument(path)
if err != nil {
return err
}
var node yaml.Node
if err := node.Encode(cfg); err != nil {
return err
}
_, value := ensureKeyValue(doc.Content[0], "tool_guard")
*value = node
var buf bytes.Buffer
encoder := yaml.NewEncoder(&buf)
encoder.SetIndent(2)
if err := encoder.Encode(doc); err != nil {
return err
}
if err := encoder.Close(); err != nil {
return err
}
info, err := os.Stat(path)
if err != nil {
return err
}
tmp, err := os.CreateTemp(filepath.Dir(path), ".tool-guard-*.yaml")
if err != nil {
return err
}
defer os.Remove(tmp.Name())
defer tmp.Close()
if err := tmp.Chmod(info.Mode().Perm()); err != nil {
return err
}
if _, err := tmp.Write(buf.Bytes()); err != nil {
return err
}
if err := tmp.Sync(); err != nil {
return err
}
if err := tmp.Close(); err != nil {
return err
}
return os.Rename(tmp.Name(), path)
}
+193
View File
@@ -0,0 +1,193 @@
package handler
import (
"bytes"
"encoding/json"
"net/http"
"net/http/httptest"
"os"
"path/filepath"
"reflect"
"strings"
"sync"
"testing"
"cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/database"
"cyberstrike-ai/internal/security"
"cyberstrike-ai/internal/toolguard"
"github.com/gin-gonic/gin"
"go.uber.org/zap"
)
func newToolGuardTestHandler(t *testing.T) *ConfigHandler {
t.Helper()
path := filepath.Join(t.TempDir(), "config.yaml")
if err := os.WriteFile(path, []byte("# keep this comment\nserver:\n port: 8123\nhitl:\n tool_whitelist: [read_file]\n"), 0600); err != nil {
t.Fatal(err)
}
manager, err := toolguard.NewManager(toolguard.DefaultConfig())
if err != nil {
t.Fatal(err)
}
return &ConfigHandler{configPath: path, config: &config.Config{}, toolGuard: manager}
}
func toolGuardRequest(t *testing.T, handler gin.HandlerFunc, body interface{}) *httptest.ResponseRecorder {
t.Helper()
data, err := json.Marshal(body)
if err != nil {
t.Fatal(err)
}
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodPut, "/api/tool-guard", bytes.NewReader(data))
c.Request.Header.Set("Content-Type", "application/json")
handler(c)
return w
}
func TestToolGuardSavePersistsAndAppliesWithoutChangingHITL(t *testing.T) {
h := newToolGuardTestHandler(t)
cfg := toolguard.DefaultConfig()
cfg.Rules[0].Message = "识别到 {match},禁止攻击政府网站,请检查目标。"
w := toolGuardRequest(t, h.UpdateToolGuard, cfg)
if w.Code != http.StatusOK {
t.Fatalf("save: %d %s", w.Code, w.Body.String())
}
loaded, err := config.Load(h.configPath)
if err != nil {
t.Fatal(err)
}
if !reflect.DeepEqual(loaded.EffectiveToolGuard(), cfg) || !reflect.DeepEqual(h.toolGuard.Config(), cfg) {
t.Fatal("saved and effective policies differ")
}
if loaded.Server.Port != 8123 || !reflect.DeepEqual(loaded.Hitl.ToolWhitelist, []string{"read_file"}) {
t.Fatal("unrelated configuration was changed")
}
info, _ := os.Stat(h.configPath)
data, _ := os.ReadFile(h.configPath)
if info.Mode().Perm() != 0600 || !strings.Contains(string(data), "# keep this comment") {
t.Fatal("file permissions or comments were lost")
}
match := h.toolGuard.Check("scan", map[string]interface{}{"target": "agency.gov.cn"})
if match == nil || !strings.Contains(match.Message, "agency.gov.cn") {
t.Fatalf("updated message not applied: %+v", match)
}
cfg.Enabled = false
w = toolGuardRequest(t, h.UpdateToolGuard, cfg)
if w.Code != http.StatusOK || h.toolGuard.Check("scan", map[string]interface{}{"target": "agency.gov"}) != nil {
t.Fatal("explicitly disabling protection did not apply")
}
}
func TestToolGuardInvalidAndFailedSaveKeepEffectivePolicy(t *testing.T) {
h := newToolGuardTestHandler(t)
before, _ := os.ReadFile(h.configPath)
cfg := toolguard.DefaultConfig()
cfg.Enabled = false
cfg.Rules[0].Pattern = "["
for _, body := range []interface{}{cfg, map[string]interface{}{}, nil, map[string]interface{}{"enabled": false, "rules": nil}} {
w := toolGuardRequest(t, h.UpdateToolGuard, body)
if w.Code != http.StatusBadRequest {
t.Fatalf("invalid update accepted: %d %s", w.Code, w.Body.String())
}
}
after, _ := os.ReadFile(h.configPath)
if !bytes.Equal(before, after) || !h.toolGuard.Config().Enabled {
t.Fatal("invalid input changed protection")
}
h.configPath = filepath.Join(t.TempDir(), "missing", "config.yaml")
cfg = toolguard.DefaultConfig()
cfg.Enabled = false
w := toolGuardRequest(t, h.UpdateToolGuard, cfg)
if w.Code != http.StatusInternalServerError || !h.toolGuard.Config().Enabled || h.config.ToolGuard != nil {
t.Fatal("failed persistence changed live configuration")
}
}
func TestToolGuardDryRunUsesUnsavedPolicyWithoutMutation(t *testing.T) {
h := newToolGuardTestHandler(t)
cfg := toolguard.DefaultConfig()
cfg.Rules[0].Pattern = "example\\.org"
w := toolGuardRequest(t, h.TestToolGuard, map[string]interface{}{
"config": cfg, "toolName": "scan", "arguments": map[string]interface{}{"target": "example.org"},
})
var got struct {
Blocked bool `json:"blocked"`
Match *toolguard.Match `json:"match"`
}
if w.Code != http.StatusOK || json.Unmarshal(w.Body.Bytes(), &got) != nil || !got.Blocked || got.Match == nil || got.Match.MatchedText != "example.org" {
t.Fatalf("dry run failed: %s", w.Body.String())
}
if !reflect.DeepEqual(h.toolGuard.Config(), toolguard.DefaultConfig()) || h.config.ToolGuard != nil {
t.Fatal("dry run changed live configuration")
}
w = toolGuardRequest(t, h.TestToolGuard, map[string]interface{}{"config": cfg, "arguments": []string{"example.org"}})
if w.Code != http.StatusBadRequest {
t.Fatal("non-object tool arguments accepted")
}
}
func TestToolGuardRoutesEnforceConfigurationPermissions(t *testing.T) {
gin.SetMode(gin.TestMode)
for _, tc := range []struct {
method, path, permission, scope string
want int
}{
{"GET", "/api/tool-guard", "hitl:read", database.RBACScopeAll, 403},
{"PUT", "/api/tool-guard", "hitl:write", database.RBACScopeAll, 403},
{"GET", "/api/tool-guard", "config:read", database.RBACScopeAll, 200},
{"POST", "/api/tool-guard/test", "config:read", database.RBACScopeAll, 200},
{"PUT", "/api/tool-guard", "config:write", database.RBACScopeAll, 200},
{"PUT", "/api/tool-guard", "config:write", database.RBACScopeOwn, 403},
} {
t.Run(tc.method+tc.permission+tc.scope, func(t *testing.T) {
r := gin.New()
r.Use(func(c *gin.Context) {
c.Set(security.ContextSessionKey, security.Session{UserID: "test", Permissions: map[string]bool{tc.permission: true}, Scope: tc.scope})
})
r.Use(security.RBACMiddleware(&database.DB{}))
r.Handle(tc.method, tc.path, func(c *gin.Context) { c.Status(200) })
w := httptest.NewRecorder()
r.ServeHTTP(w, httptest.NewRequest(tc.method, tc.path, nil))
if w.Code != tc.want {
t.Fatalf("got %d, want %d: %s", w.Code, tc.want, w.Body.String())
}
})
}
}
func TestToolGuardConcurrentOtherSettingsSavePreservesPolicy(t *testing.T) {
h := newToolGuardTestHandler(t)
external := &ExternalMCPHandler{configPath: h.configPath, config: h.config, logger: zap.NewNop()}
cfg := toolguard.DefaultConfig()
cfg.Rules[0].Message = "持久化策略 {match}"
var wg sync.WaitGroup
errors := make(chan error, 2)
for _, save := range []func() error{func() error { return h.saveToolGuardConfig(cfg) }, external.saveConfig} {
wg.Add(1)
go func(save func() error) {
defer wg.Done()
for i := 0; i < 20; i++ {
if err := save(); err != nil {
errors <- err
return
}
}
}(save)
}
wg.Wait()
close(errors)
for err := range errors {
t.Fatal(err)
}
loaded, err := config.Load(h.configPath)
if err != nil {
t.Fatal(err)
}
if !reflect.DeepEqual(loaded.EffectiveToolGuard(), cfg) {
t.Fatal("another settings save overwrote the tool guard policy")
}
}
+86
View File
@@ -0,0 +1,86 @@
package mcp
import (
"context"
"encoding/json"
"strings"
"testing"
"time"
sdkmcp "github.com/modelcontextprotocol/go-sdk/mcp"
)
func TestBlockedExecutionIsTerminalAndNotFailed(t *testing.T) {
for _, blocked := range []bool{true, false} {
name := "error"
want := ToolExecutionStatusFailed
if blocked {
name, want = "blocked", ToolExecutionStatusBlocked
}
t.Run(name, func(t *testing.T) {
service := NewExecutionService(nil, nil)
handle, err := service.Submit(context.Background(), ExecutionRequest{
ToolName: "test",
Run: func(context.Context) (*ToolResult, error) {
// Identical text must not turn ordinary failures into policy blocks.
return &ToolResult{Content: []Content{{Type: "text", Text: toolGuardBlockedPrefix}}, IsError: true, Blocked: blocked}, nil
},
})
if err != nil {
t.Fatal(err)
}
snap, err := service.Wait(context.Background(), handle.ID, time.Second)
if err != nil || snap.Execution.Status != want || snap.Execution.Result.Blocked != blocked || snap.Execution.Error == "" {
t.Fatalf("incorrect classification: snapshot=%#v err=%v", snap, err)
}
if !isExecutionTerminal(want) || executionStatusCountsAsFailed(want) == blocked {
t.Fatalf("incorrect terminal/failure classification for %s", want)
}
if service.Cancel(handle.ID, "cancel after completion") {
t.Fatal("terminal execution must not be cancellable")
}
after, _ := service.Get(handle.ID)
if after.Execution.Status != want {
t.Fatalf("cancel reclassified terminal execution: %s", after.Execution.Status)
}
})
}
}
func TestBlockedMarkerSurvivesNormalizationAndMCPProtocol(t *testing.T) {
original := &ToolResult{Content: []Content{{Type: "text", Text: strings.Repeat("refused ", 2000)}}, IsError: true, Blocked: true}
bounded := NormalizeToolResultForStorageWithSpill(original, 1000, ToolResultSpillConfig{RootDir: t.TempDir(), ExecutionID: "blocked"})
if !bounded.Blocked || !bounded.IsError || ToolResultPlainText(bounded) == ToolResultPlainText(original) {
t.Fatal("normalization must retain classification while bounding long output")
}
wire, err := json.Marshal(CallToolResponse{Content: bounded.Content, IsError: bounded.IsError, Blocked: bounded.Blocked, Meta: toolResultProtocolMeta(bounded)})
if err != nil {
t.Fatal(err)
}
var decoded ToolResult
if err := json.Unmarshal(wire, &decoded); err != nil || !decoded.Blocked || !decoded.IsError {
t.Fatalf("application protocol lost block marker: %#v err=%v", decoded, err)
}
var sdkResult sdkmcp.CallToolResult
if err := json.Unmarshal(wire, &sdkResult); err != nil {
t.Fatal(err)
}
converted := sdkCallToolResultToOurs(&sdkResult)
if !converted.Blocked || !converted.IsError {
t.Fatalf("SDK round trip lost block marker: %#v", converted)
}
}
func TestToolStatsSeparateBlockedFromFailures(t *testing.T) {
server := NewServer(nil)
manager := NewExternalMCPManager(nil)
for _, status := range []string{ToolExecutionStatusCompleted, ToolExecutionStatusFailed, ToolExecutionStatusBlocked, ToolExecutionStatusCancelled} {
server.updateStats("test", status)
manager.updateStats("test", status)
}
for name, stat := range map[string]*ToolStats{"internal": server.stats["test"], "external": manager.stats["test"]} {
if stat.TotalCalls != 4 || stat.SuccessCalls != 1 || stat.FailedCalls != 1 || stat.BlockedCalls != 1 {
t.Fatalf("%s stats = %#v", name, stat)
}
}
}
+2
View File
@@ -308,9 +308,11 @@ func sdkCallToolResultToOurs(res *mcp.CallToolResult) *ToolResult {
return &ToolResult{Content: []Content{}}
}
content := sdkContentToOurs(res.Content)
blocked, _ := res.Meta[toolGuardBlockedMetaKey].(bool)
return &ToolResult{
Content: content,
IsError: res.IsError,
Blocked: blocked,
}
}
+3
View File
@@ -190,6 +190,9 @@ func formatExecutionForModel(exec *ToolExecution, opts executionFormatOptions) s
if exec.Result != nil {
payload["result"] = ToolResultPlainText(exec.Result)
payload["is_error"] = exec.Result.IsError
if exec.Result.Blocked {
payload["blocked"] = true
}
}
if opts.includePartialOutput && exec.PartialOutput != "" {
partial := tailStringBytes(exec.PartialOutput, opts.partialMaxBytes)
+38 -8
View File
@@ -18,6 +18,7 @@ const (
ToolExecutionStatusQueued = "queued"
ToolExecutionStatusRunning = "running"
ToolExecutionStatusCompleted = "completed"
ToolExecutionStatusBlocked = "blocked"
ToolExecutionStatusFailed = "failed"
ToolExecutionStatusCancelled = "cancelled"
ToolExecutionStatusHardTimeout = "hard_timeout"
@@ -224,6 +225,10 @@ func (s *ExecutionService) markEntryRunning(entry *executionEntry) {
func (s *ExecutionService) finishEntry(ctx context.Context, entry *executionEntry, result *ToolResult, err error, onDone ExecutionDoneFunc) {
id := entry.exec.ID
var blockedErr *toolGuardBlockError
if errors.As(err, &blockedErr) {
result, err = blockedErr.result, nil
}
cancelledWithUserNote := s.applyAbortUserNoteToCancelledToolResult(id, &result, &err)
now := time.Now()
@@ -258,6 +263,10 @@ func (s *ExecutionService) finishEntry(ctx context.Context, entry *executionEntr
entry.exec.Status = ToolExecutionStatusFailed
entry.exec.Error = err.Error()
}
} else if result != nil && result.Blocked {
entry.exec.Status = ToolExecutionStatusBlocked
entry.exec.Error = firstToolResultText(result, "工具调用已被安全规则拦截")
entry.exec.Result = result
} else if result != nil && result.IsError {
if cancelledWithUserNote {
entry.exec.Status = ToolExecutionStatusCancelled
@@ -318,10 +327,11 @@ func (s *ExecutionService) Wait(ctx context.Context, executionID string, timeout
if entry == nil {
return s.getPersistedSnapshot(executionID)
}
if isExecutionTerminal(entry.exec.Status) {
return &ExecutionSnapshot{Execution: cloneToolExecution(entry.exec)}, nil
select {
case <-entry.done:
return s.snapshotEntry(entry), nil
default:
}
var timeoutCh <-chan time.Time
var timer *time.Timer
if timeout > 0 {
@@ -332,18 +342,26 @@ func (s *ExecutionService) Wait(ctx context.Context, executionID string, timeout
select {
case <-entry.done:
return &ExecutionSnapshot{Execution: cloneToolExecution(entry.exec)}, nil
return s.snapshotEntry(entry), nil
case <-timeoutCh:
return &ExecutionSnapshot{Execution: cloneToolExecution(entry.exec)}, ErrExecutionWaitTimeout
return s.snapshotEntry(entry), ErrExecutionWaitTimeout
case <-ctxDone(ctx):
return &ExecutionSnapshot{Execution: cloneToolExecution(entry.exec)}, ctx.Err()
return s.snapshotEntry(entry), ctx.Err()
}
}
// snapshotEntry synchronizes snapshots with worker state and partial output
// updates. Wait uses done to also observe persistence and completion callbacks.
func (s *ExecutionService) snapshotEntry(entry *executionEntry) *ExecutionSnapshot {
s.mu.Lock()
defer s.mu.Unlock()
return &ExecutionSnapshot{Execution: cloneToolExecution(entry.exec)}
}
func (s *ExecutionService) Get(executionID string) (*ExecutionSnapshot, error) {
entry := s.getEntry(executionID)
if entry != nil {
return &ExecutionSnapshot{Execution: cloneToolExecution(entry.exec)}, nil
return s.snapshotEntry(entry), nil
}
return s.getPersistedSnapshot(executionID)
}
@@ -464,6 +482,9 @@ func (s *ExecutionService) applyAbortUserNoteToCancelledToolResult(executionID s
}
hasErr := err != nil && *err != nil
hasRes := result != nil && *result != nil
if hasRes && (*result).Blocked {
return false
}
if !hasErr && !hasRes {
return false
}
@@ -549,7 +570,16 @@ func isBackgroundWaitToolResult(result *ToolResult) bool {
func isExecutionTerminal(status string) bool {
switch strings.TrimSpace(strings.ToLower(status)) {
case ToolExecutionStatusCompleted, ToolExecutionStatusFailed, ToolExecutionStatusCancelled, ToolExecutionStatusHardTimeout, ToolExecutionStatusOrphaned:
case ToolExecutionStatusCompleted, ToolExecutionStatusBlocked, ToolExecutionStatusFailed, ToolExecutionStatusCancelled, ToolExecutionStatusHardTimeout, ToolExecutionStatusOrphaned:
return true
default:
return false
}
}
func executionStatusCountsAsFailed(status string) bool {
switch status {
case ToolExecutionStatusFailed, ToolExecutionStatusHardTimeout, ToolExecutionStatusOrphaned:
return true
default:
return false
+44 -8
View File
@@ -11,6 +11,7 @@ import (
"cyberstrike-ai/internal/authctx"
"cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/toolguard"
"go.uber.org/zap"
)
@@ -74,6 +75,7 @@ type ExternalMCPManager struct {
reconnectLastTry map[string]time.Time
reconnectAttempts map[string]int
toolAuthorizer func(context.Context, string, map[string]interface{}) error
toolGuard *toolguard.Manager
executionService *ExecutionService
toolWaitTimeout time.Duration
toolResultMaxBytes int
@@ -96,6 +98,23 @@ func (m *ExternalMCPManager) SetToolAuthorizer(authorizer func(context.Context,
m.mu.Unlock()
}
// SetToolGuard installs safety rules evaluated before dispatch to external MCPs.
func (m *ExternalMCPManager) SetToolGuard(guard *toolguard.Manager) {
if m == nil {
return
}
m.mu.Lock()
m.toolGuard = guard
m.mu.Unlock()
}
func (m *ExternalMCPManager) checkToolGuard(toolName string, args map[string]interface{}) *ToolResult {
m.mu.RLock()
guard := m.toolGuard
m.mu.RUnlock()
return toolGuardBlockedResult(guard, toolName, args)
}
// NewExternalMCPManagerWithStorage 创建外部MCP管理器(带持久化存储)
func NewExternalMCPManagerWithStorage(logger *zap.Logger, storage MonitorStorage) *ExternalMCPManager {
manager := &ExternalMCPManager{
@@ -685,6 +704,7 @@ func (m *ExternalMCPManager) CallTool(ctx context.Context, toolName string, args
}
var mcpName, actualToolName string
var client ExternalMCPClient
var blockedByGuard bool
handle, err := m.executionService.Submit(ctx, ExecutionRequest{
ToolName: toolName,
Arguments: args,
@@ -702,6 +722,10 @@ func (m *ExternalMCPManager) CallTool(ctx context.Context, toolName string, args
} else if authenticated {
return nil, fmt.Errorf("external tool authorization policy is not configured")
}
if blocked := m.checkToolGuard(toolName, args); blocked != nil {
blockedByGuard = true
return nil, &toolGuardBlockError{result: blocked}
}
// 解析工具名称:name::toolName
if idx := findSubstring(toolName, "::"); idx > 0 {
@@ -741,6 +765,11 @@ func (m *ExternalMCPManager) CallTool(ctx context.Context, toolName string, args
return release, nil
},
Run: func(runCtx context.Context) (*ToolResult, error) {
// Rules may have changed while this execution waited for a slot.
if blocked := m.checkToolGuard(toolName, args); blocked != nil {
blockedByGuard = true
return blocked, nil
}
result, callErr := client.CallTool(runCtx, actualToolName, args)
if callErr != nil {
m.handleConnectionDead(mcpName, client, callErr)
@@ -748,11 +777,13 @@ func (m *ExternalMCPManager) CallTool(ctx context.Context, toolName string, args
return result, callErr
},
OnDone: func(exec *ToolExecution) {
failed := exec != nil && exec.Status != ToolExecutionStatusCompleted && exec.Status != ToolExecutionStatusCancelled
if mcpName != "" {
failed := exec != nil && executionStatusCountsAsFailed(exec.Status)
if mcpName != "" && !blockedByGuard && (exec == nil || exec.Status != ToolExecutionStatusBlocked) {
m.recordExternalMCPResult(mcpName, failed)
}
m.updateStats(toolName, failed)
if exec != nil {
m.updateStats(toolName, exec.Status)
}
},
})
if err != nil {
@@ -941,6 +972,9 @@ func (m *ExternalMCPManager) applyAbortUserNoteToCancelledToolResult(executionID
}
hasErr := err != nil && *err != nil
hasRes := result != nil && *result != nil
if hasRes && (*result).Blocked {
return false
}
if !hasErr && !hasRes {
return false
}
@@ -1098,15 +1132,15 @@ func (m *ExternalMCPManager) ActiveRunningExecutionIDs() map[string]struct{} {
}
// updateStats 更新统计信息
func (m *ExternalMCPManager) updateStats(toolName string, failed bool) {
func (m *ExternalMCPManager) updateStats(toolName string, status string) {
now := time.Now()
if m.storage != nil {
totalCalls := 1
successCalls := 0
failedCalls := 0
if failed {
if executionStatusCountsAsFailed(status) {
failedCalls = 1
} else {
} else if status == ToolExecutionStatusCompleted {
successCalls = 1
}
if err := m.storage.UpdateToolStats(toolName, totalCalls, successCalls, failedCalls, &now); err != nil {
@@ -1128,10 +1162,12 @@ func (m *ExternalMCPManager) updateStats(toolName string, failed bool) {
stats.TotalCalls++
stats.LastCallTime = &now
if failed {
if executionStatusCountsAsFailed(status) {
stats.FailedCalls++
} else {
} else if status == ToolExecutionStatusCompleted {
stats.SuccessCalls++
} else if status == ToolExecutionStatusBlocked {
stats.BlockedCalls++
}
}
@@ -72,7 +72,9 @@ func TestExternalMCPManager_CallToolBoundedWaitThenContinue(t *testing.T) {
manager.ConfigureToolWaitTimeoutSeconds(1)
manager.toolWaitTimeout = 10 * time.Millisecond
client := newBlockingExternalMCPClient("slow result ready")
manager.mu.Lock()
manager.clients["lab"] = client
manager.mu.Unlock()
callCtx, callCancel := context.WithCancel(context.Background())
result, executionID, err := manager.CallTool(callCtx, "lab::slow_tool", map[string]interface{}{"target": "example"})
@@ -117,7 +119,9 @@ func TestExecutionControlWaitToolReturnsCompletedResult(t *testing.T) {
manager := NewExternalMCPManager(zap.NewNop())
manager.toolWaitTimeout = 10 * time.Millisecond
client := newBlockingExternalMCPClient("control wait result")
manager.mu.Lock()
manager.clients["lab"] = client
manager.mu.Unlock()
result, executionID, err := manager.CallTool(context.Background(), "lab::slow_tool", nil)
if err != nil {
@@ -157,7 +161,9 @@ func TestExternalMCPManager_PerServerConcurrencyLimitsWorkers(t *testing.T) {
CircuitCooldown: time.Second,
})
client := newBlockingExternalMCPClient("ok")
manager.mu.Lock()
manager.clients["lab"] = client
manager.mu.Unlock()
done1 := make(chan struct{})
go func() {
@@ -217,7 +223,9 @@ func TestExternalMCPManager_CircuitBreakerOpensAfterFailures(t *testing.T) {
CircuitFailureThreshold: 1,
CircuitCooldown: time.Minute,
})
manager.mu.Lock()
manager.clients["lab"] = &failingExternalMCPClient{}
manager.mu.Unlock()
_, _, err := manager.CallTool(context.Background(), "lab::fail_tool", nil)
if err == nil || !strings.Contains(err.Error(), "boom") {
+54 -16
View File
@@ -16,6 +16,7 @@ import (
"cyberstrike-ai/internal/authctx"
"cyberstrike-ai/internal/mcp/builtin"
"cyberstrike-ai/internal/toolguard"
"github.com/google/uuid"
"go.uber.org/zap"
@@ -53,6 +54,7 @@ type Server struct {
httpToolTimeoutMinutes *int
httpToolTimeoutMu sync.RWMutex
toolAuthorizer func(context.Context, string, map[string]interface{}) error
toolGuard *toolguard.Manager
executionService *ExecutionService
toolWaitTimeout time.Duration
toolResultMaxBytes int
@@ -72,6 +74,23 @@ func (s *Server) SetToolAuthorizer(authorizer func(context.Context, string, map[
s.mu.Unlock()
}
// SetToolGuard installs the runtime safety rules shared by HTTP and internal calls.
func (s *Server) SetToolGuard(guard *toolguard.Manager) {
if s == nil {
return
}
s.mu.Lock()
s.toolGuard = guard
s.mu.Unlock()
}
func (s *Server) checkToolGuard(toolName string, args map[string]interface{}) *ToolResult {
s.mu.RLock()
guard := s.toolGuard
s.mu.RUnlock()
return toolGuardBlockedResult(guard, toolName, args)
}
type sseClient struct {
id string
send chan []byte
@@ -566,7 +585,7 @@ func (s *Server) handleCallTool(requestCtx context.Context, msg *Message) *Messa
s.mu.Unlock()
}
s.updateStats(req.Name, true)
s.updateStats(req.Name, ToolExecutionStatusFailed)
return &Message{
ID: msg.ID,
@@ -590,10 +609,13 @@ func (s *Server) handleCallTool(requestCtx context.Context, msg *Message) *Messa
zap.Any("arguments", req.Arguments),
)
result, err := handler(execCtx, req.Arguments)
result := s.checkToolGuard(req.Name, req.Arguments)
var err error
if result == nil {
result, err = handler(execCtx, req.Arguments)
}
cancelledWithUserNote := s.applyAbortUserNoteToCancelledToolResult(executionID, &result, &err)
now := time.Now()
var failed bool
var finalResult *ToolResult
s.mu.Lock()
@@ -604,13 +626,15 @@ func (s *Server) handleCallTool(requestCtx context.Context, msg *Message) *Messa
st, msg := executionStatusAndMessage(err)
execution.Status = st
execution.Error = msg
failed = st != "cancelled"
} else if result != nil && result.Blocked {
execution.Status = ToolExecutionStatusBlocked
execution.Error = firstToolResultText(result, toolGuardBlockedPrefix)
execution.Result = result
} else if result != nil && result.IsError {
if cancelledWithUserNote {
execution.Status = "cancelled"
execution.Error = ""
execution.Result = result
failed = false
} else {
execution.Status = "failed"
if len(result.Content) > 0 {
@@ -619,7 +643,6 @@ func (s *Server) handleCallTool(requestCtx context.Context, msg *Message) *Messa
execution.Error = "工具执行返回错误结果"
}
execution.Result = result
failed = true
}
} else {
execution.Status = "completed"
@@ -631,7 +654,6 @@ func (s *Server) handleCallTool(requestCtx context.Context, msg *Message) *Messa
}
}
execution.Result = result
failed = false
}
finalResult = execution.Result
@@ -643,7 +665,7 @@ func (s *Server) handleCallTool(requestCtx context.Context, msg *Message) *Messa
}
}
s.updateStats(req.Name, failed)
s.updateStats(req.Name, execution.Status)
if s.storage != nil {
s.mu.Lock()
@@ -683,6 +705,8 @@ func (s *Server) handleCallTool(requestCtx context.Context, msg *Message) *Messa
errorResult, _ := json.Marshal(CallToolResponse{
Content: finalResult.Content,
IsError: true,
Blocked: finalResult.Blocked,
Meta: toolResultProtocolMeta(finalResult),
})
return &Message{
ID: msg.ID,
@@ -719,15 +743,15 @@ func (s *Server) handleCallTool(requestCtx context.Context, msg *Message) *Messa
}
// updateStats 更新统计信息
func (s *Server) updateStats(toolName string, failed bool) {
func (s *Server) updateStats(toolName string, status string) {
now := time.Now()
if s.storage != nil {
totalCalls := 1
successCalls := 0
failedCalls := 0
if failed {
if executionStatusCountsAsFailed(status) {
failedCalls = 1
} else {
} else if status == ToolExecutionStatusCompleted {
successCalls = 1
}
if err := s.storage.UpdateToolStats(toolName, totalCalls, successCalls, failedCalls, &now); err != nil {
@@ -749,10 +773,12 @@ func (s *Server) updateStats(toolName string, failed bool) {
stats.TotalCalls++
stats.LastCallTime = &now
if failed {
if executionStatusCountsAsFailed(status) {
stats.FailedCalls++
} else {
} else if status == ToolExecutionStatusCompleted {
stats.SuccessCalls++
} else if status == ToolExecutionStatusBlocked {
stats.BlockedCalls++
}
}
@@ -925,11 +951,15 @@ func (s *Server) CallTool(ctx context.Context, toolName string, args map[string]
if !exists {
return nil, fmt.Errorf("工具 %s 未找到", toolName)
}
if blocked := s.checkToolGuard(toolName, args); blocked != nil {
return blocked, nil
}
return handler(runCtx, args)
},
OnDone: func(exec *ToolExecution) {
failed := exec != nil && exec.Status != ToolExecutionStatusCompleted && exec.Status != ToolExecutionStatusCancelled
s.updateStats(toolName, failed)
if exec != nil {
s.updateStats(toolName, exec.Status)
}
},
})
if err != nil {
@@ -1111,7 +1141,7 @@ func (s *Server) FinishToolExecution(ctx context.Context, executionID, toolName
}
}
s.updateStats(exec.ToolName, failed)
s.updateStats(exec.ToolName, exec.Status)
if s.storage != nil {
s.mu.Lock()
@@ -1155,6 +1185,11 @@ func (s *Server) UpdateToolExecutionResult(executionID string, result *ToolResul
if executionID == "" || result == nil {
return nil
}
if previous, ok := s.GetExecution(executionID); ok && previous != nil &&
(previous.Status == ToolExecutionStatusBlocked || previous.Result != nil && previous.Result.Blocked) {
result = cloneToolResult(result)
result.Blocked, result.IsError = true, true
}
s.mu.Lock()
spill := ToolResultSpillConfig{
RootDir: s.spillRootDir,
@@ -1270,6 +1305,9 @@ func (s *Server) applyAbortUserNoteToCancelledToolResult(executionID string, res
}
hasErr := err != nil && *err != nil
hasRes := result != nil && *result != nil
if hasRes && (*result).Blocked {
return false
}
if !hasErr && !hasRes {
return false
}
+41
View File
@@ -0,0 +1,41 @@
package mcp
import (
"fmt"
"strings"
"cyberstrike-ai/internal/toolguard"
)
const toolGuardBlockedPrefix = "工具调用已被安全规则拦截"
const toolGuardBlockedMetaKey = "cyberstrike.ai/blocked"
// toolGuardBlockError carries structured policy results through pre-run hooks.
type toolGuardBlockError struct{ result *ToolResult }
func (e *toolGuardBlockError) Error() string { return ToolResultPlainText(e.result) }
func toolResultProtocolMeta(result *ToolResult) map[string]interface{} {
if result != nil && result.Blocked {
return map[string]interface{}{toolGuardBlockedMetaKey: true}
}
return nil
}
// toolGuardBlockedResult uses the standard MCP error result so the refusal is
// visible both to the model and in persisted execution monitoring records.
func toolGuardBlockedResult(guard *toolguard.Manager, toolName string, args map[string]interface{}) *ToolResult {
if guard == nil {
return nil
}
match := guard.Check(toolName, args)
if match == nil {
return nil
}
message := toolGuardBlockedPrefix
if custom := strings.TrimSpace(match.Message); custom != "" {
message += "" + custom
}
message += fmt.Sprintf("\n规则: %s (%s)\n匹配内容: %q", match.RuleName, match.RuleID, match.MatchedText)
return &ToolResult{Content: []Content{{Type: "text", Text: message}}, IsError: true, Blocked: true}
}
+239
View File
@@ -0,0 +1,239 @@
package mcp
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"strings"
"sync/atomic"
"testing"
"time"
"cyberstrike-ai/internal/toolguard"
"go.uber.org/zap"
)
func testToolGuard(t *testing.T, enabled bool) *toolguard.Manager {
t.Helper()
guard, err := toolguard.NewManager(toolguard.DefaultConfig())
if err != nil {
t.Fatal(err)
}
if err := guard.Update(toolguard.Config{Enabled: enabled, Rules: []toolguard.Rule{{
ID: "government", Name: "政府网站保护", Enabled: true,
Pattern: `(?i)[a-z0-9.-]+\.gov(?:\.[a-z0-9.-]+)?`,
Message: "识别到 {match},禁止攻击政府网站,请检查目标授权。",
}}}); err != nil {
t.Fatal(err)
}
return guard
}
func assertGuardRefusal(t *testing.T, result *ToolResult, err error) {
t.Helper()
message := ToolResultPlainText(result)
if err != nil {
t.Fatalf("expected structured refusal, got error: %v", err)
} else if result == nil || !result.IsError || !result.Blocked {
t.Fatalf("expected tool error result, got %#v", result)
}
for _, text := range []string{toolGuardBlockedPrefix, "禁止攻击政府网站", "agency.gov.cn", "government"} {
if !strings.Contains(message, text) {
t.Errorf("refusal %q missing %q", message, text)
}
}
}
func TestServerToolGuardBlocksBeforeHandlerAndUpdatesLive(t *testing.T) {
storage := newInMemoryMonitorStorage()
server := NewServerWithStorage(zap.NewNop(), storage)
guard := testToolGuard(t, true)
server.SetToolGuard(guard)
var calls, authorized atomic.Int32
server.SetToolAuthorizer(func(context.Context, string, map[string]interface{}) error {
authorized.Add(1)
return nil
})
server.RegisterTool(Tool{Name: "scan"}, func(context.Context, map[string]interface{}) (*ToolResult, error) {
calls.Add(1)
return &ToolResult{Content: []Content{{Type: "text", Text: "ok"}}}, nil
})
args := map[string]interface{}{"command": "scan https://agency.gov.cn"}
result, executionID, err := server.CallTool(context.Background(), "scan", args)
assertGuardRefusal(t, result, err)
if calls.Load() != 0 || authorized.Load() != 1 {
t.Fatalf("calls=%d authorized=%d, want 0 and 1", calls.Load(), authorized.Load())
}
execution, err := storage.GetToolExecution(executionID)
if err != nil || execution == nil || execution.Status != ToolExecutionStatusBlocked || !strings.Contains(execution.Error, toolGuardBlockedPrefix) {
t.Fatalf("expected persisted blocked execution, got %#v, err=%v", execution, err)
}
result, _, err = server.CallTool(context.Background(), "scan", map[string]interface{}{"target": "example.org"})
if err != nil || result.IsError || calls.Load() != 1 {
t.Fatalf("allowed target did not execute: result=%#v calls=%d err=%v", result, calls.Load(), err)
}
cfg := guard.Config()
cfg.Rules[0].Enabled = false
if err := guard.Update(cfg); err != nil {
t.Fatal(err)
}
result, _, err = server.CallTool(context.Background(), "scan", args)
if err != nil || result.IsError || calls.Load() != 2 {
t.Fatalf("disabled rule did not take effect: result=%#v calls=%d err=%v", result, calls.Load(), err)
}
}
func TestHTTPToolGuardReturnsMCPErrorAndPersistsRefusal(t *testing.T) {
storage := newInMemoryMonitorStorage()
server := NewServerWithStorage(zap.NewNop(), storage)
server.SetToolGuard(testToolGuard(t, true))
var calls int
server.RegisterTool(Tool{Name: "scan"}, func(context.Context, map[string]interface{}) (*ToolResult, error) {
calls++
return &ToolResult{Content: []Content{{Type: "text", Text: "ok"}}}, nil
})
for _, tc := range []struct {
target string
blocked bool
}{
{target: "https://agency.gov.cn", blocked: true},
{target: "https://example.org", blocked: false},
} {
body, err := json.Marshal(map[string]interface{}{
"jsonrpc": "2.0", "id": 1, "method": "tools/call",
"params": map[string]interface{}{"name": "scan", "arguments": map[string]interface{}{"target": tc.target}},
})
if err != nil {
t.Fatal(err)
}
recorder := httptest.NewRecorder()
server.HandleHTTP(recorder, httptest.NewRequest(http.MethodPost, "/api/mcp", strings.NewReader(string(body))))
var response Message
if err := json.Unmarshal(recorder.Body.Bytes(), &response); err != nil {
t.Fatal(err)
}
if recorder.Code != http.StatusOK || response.Error != nil {
t.Fatalf("expected MCP tool result, status=%d body=%s", recorder.Code, recorder.Body)
}
var result ToolResult
if err := json.Unmarshal(response.Result, &result); err != nil {
t.Fatal(err)
}
if tc.blocked {
assertGuardRefusal(t, &result, nil)
if calls != 0 {
t.Fatal("HTTP tool handler ran for a blocked target")
}
executions, err := storage.LoadToolExecutions()
if err != nil || len(executions) != 1 || executions[0].Status != ToolExecutionStatusBlocked || !strings.Contains(executions[0].Error, toolGuardBlockedPrefix) {
t.Fatalf("expected persisted HTTP refusal, got %#v err=%v", executions, err)
}
} else if result.IsError || calls != 1 {
t.Fatalf("allowed HTTP target did not execute: result=%#v calls=%d", result, calls)
}
}
}
func TestExternalToolGuardBlocksBeforeClientAndUpdatesLive(t *testing.T) {
manager := NewExternalMCPManager(zap.NewNop())
t.Cleanup(manager.StopAll)
guard := testToolGuard(t, true)
manager.SetToolGuard(guard)
client := newBlockingExternalMCPClient("ok")
close(client.release)
manager.mu.Lock()
manager.clients["lab"] = client
manager.mu.Unlock()
args := map[string]interface{}{"target": "https://agency.gov.cn"}
result, executionID, err := manager.CallTool(context.Background(), "lab::slow_tool", args)
assertGuardRefusal(t, result, err)
if client.count.Load() != 0 {
t.Fatal("external client ran for a blocked target")
}
execution, ok := manager.GetExecution(executionID)
if !ok || execution.Status != ToolExecutionStatusBlocked || !strings.Contains(execution.Error, toolGuardBlockedPrefix) {
t.Fatalf("expected blocked external execution, got %#v", execution)
}
cfg := guard.Config()
cfg.Enabled = false
if err := guard.Update(cfg); err != nil {
t.Fatal(err)
}
result, _, err = manager.CallTool(context.Background(), "lab::slow_tool", args)
if err != nil || result.IsError || client.count.Load() != 1 {
t.Fatalf("disabled guard did not take effect: result=%#v calls=%d err=%v", result, client.count.Load(), err)
}
}
func TestExternalToolGuardRechecksQueuedCallsWithoutTrippingCircuit(t *testing.T) {
manager := NewExternalMCPManager(zap.NewNop())
t.Cleanup(manager.StopAll)
manager.toolWaitTimeout = 10 * time.Millisecond
manager.ConfigureResilience(ExternalMCPResilienceConfig{
MaxConcurrentPerServer: 1, MaxConcurrentTotal: 4,
CircuitFailureThreshold: 1, CircuitCooldown: time.Minute,
})
guard := testToolGuard(t, false)
manager.SetToolGuard(guard)
client := newBlockingExternalMCPClient("ok")
close(client.release)
manager.mu.Lock()
manager.clients["lab"] = client
manager.mu.Unlock()
// Occupy the provider slot so the call passes its initial policy check and
// remains queued until a live rule update is applied.
release, err := manager.acquireExternalMCPCallSlot(context.Background(), "lab")
if err != nil {
t.Fatal(err)
}
released := false
t.Cleanup(func() {
if !released {
release()
}
})
_, executionID, err := manager.CallTool(context.Background(), "lab::slow_tool", map[string]interface{}{"target": "agency.gov.cn"})
if err != nil || executionID == "" {
t.Fatalf("failed to queue external call: id=%q err=%v", executionID, err)
}
deadline := time.After(time.Second)
ticker := time.NewTicker(time.Millisecond)
defer ticker.Stop()
for len(manager.globalSemaphore) != 2 {
select {
case <-deadline:
t.Fatal("execution did not reach the provider slot queue")
case <-ticker.C:
}
}
cfg := guard.Config()
cfg.Enabled = true
if err := guard.Update(cfg); err != nil {
t.Fatal(err)
}
release()
released = true
snapshot, err := manager.executionService.Wait(context.Background(), executionID, time.Second)
if err != nil || snapshot == nil || snapshot.Execution == nil || snapshot.Execution.Status != ToolExecutionStatusBlocked {
t.Fatalf("expected queued execution to be blocked on policy recheck, got %#v err=%v", snapshot, err)
}
assertGuardRefusal(t, snapshot.Execution.Result, nil)
if client.count.Load() != 0 {
t.Fatal("queued call bypassed the updated guard")
}
manager.mu.RLock()
runtime := manager.serverRuntimes["lab"]
failures, openUntil := runtime.consecutiveFailures, runtime.circuitOpenUntil
manager.mu.RUnlock()
if failures != 0 || !openUntil.IsZero() {
t.Fatalf("local policy refusal affected provider circuit: failures=%d openUntil=%v", failures, openUntil)
}
result, _, err := manager.CallTool(context.Background(), "lab::slow_tool", map[string]interface{}{"target": "example.org"})
if err != nil || result.IsError || client.count.Load() != 1 {
t.Fatalf("allowed call failed after policy refusal: result=%#v calls=%d err=%v", result, client.count.Load(), err)
}
}
+9 -3
View File
@@ -116,6 +116,9 @@ type ToolCall struct {
type ToolResult struct {
Content []Content `json:"content"`
IsError bool `json:"isError,omitempty"`
// Blocked means policy stopped the call before execution. IsError remains
// true for MCP/model handling, while monitoring uses a distinct status.
Blocked bool `json:"blocked,omitempty"`
}
// Content 表示内容
@@ -184,8 +187,10 @@ type CallToolRequest struct {
// CallToolResponse 调用工具响应
type CallToolResponse struct {
Content []Content `json:"content"`
IsError bool `json:"isError,omitempty"`
Content []Content `json:"content"`
IsError bool `json:"isError,omitempty"`
Blocked bool `json:"blocked,omitempty"`
Meta map[string]interface{} `json:"_meta,omitempty"`
}
// ToolExecution 工具执行记录
@@ -193,7 +198,7 @@ type ToolExecution struct {
ID string `json:"id"`
ToolName string `json:"toolName"`
Arguments map[string]interface{} `json:"arguments"`
Status string `json:"status"` // pending, running, completed, failed, cancelled
Status string `json:"status"` // queued, running, completed, blocked, failed, cancelled, hard_timeout, orphaned
Result *ToolResult `json:"result,omitempty"`
Error string `json:"error,omitempty"`
StartTime time.Time `json:"startTime"`
@@ -216,6 +221,7 @@ type ToolStats struct {
TotalCalls int `json:"totalCalls"`
SuccessCalls int `json:"successCalls"`
FailedCalls int `json:"failedCalls"`
BlockedCalls int `json:"blockedCalls"`
LastCallTime *time.Time `json:"lastCallTime,omitempty"`
}
@@ -0,0 +1,87 @@
package multiagent
import (
"context"
"testing"
"cyberstrike-ai/internal/agent"
"cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/mcp"
"cyberstrike-ai/internal/toolguard"
"go.uber.org/zap"
)
func TestEinoToolResultPreservesBlockedOutcomeAfterTextReduction(t *testing.T) {
for _, blocked := range []bool{true, false} {
name := "execution_error"
if blocked {
name = "safety_block"
}
t.Run(name, func(t *testing.T) {
ctx := context.Background()
logger := zap.NewNop()
server := mcp.NewServer(logger)
guard, err := toolguard.NewManager(toolguard.DefaultConfig())
if err != nil {
t.Fatal(err)
}
server.SetToolGuard(guard)
calls := 0
server.RegisterTool(mcp.Tool{Name: "inspect"}, func(context.Context, map[string]interface{}) (*mcp.ToolResult, error) {
calls++
return &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: "execution failed"}}, IsError: true}, nil
})
ag := agent.NewAgent(&config.OpenAIConfig{}, &config.AgentConfig{}, server, nil, logger, 1)
target := "example.org"
if blocked {
target = "example.gov"
}
result, err := ag.ExecuteMCPToolForConversation(ctx, "conv-block", "inspect", map[string]interface{}{"target": target})
if err != nil || result == nil || !result.IsError || result.Blocked != blocked {
t.Fatalf("agent result = %#v, error = %v", result, err)
}
if blocked && calls != 0 || !blocked && calls != 1 {
t.Fatalf("handler calls = %d, blocked = %v", calls, blocked)
}
binder := NewMCPExecutionBinder()
binder.Bind("call-block", result.ExecutionID)
var event map[string]interface{}
emitter := newEinoToolResultProgressEmitter(einoToolResultProgressEmitterConfig{
FilesystemMonitorAgent: ag,
MCPExecutionBinder: binder,
Progress: func(eventType, _ string, data interface{}) {
if eventType == "tool_result" {
event = data.(map[string]interface{})
}
},
})
// The reduced text deliberately contains no refusal wording. A blocked
// outcome must survive even if reduction also loses the error prefix.
const reduced = "The request did not run."
if !emitter.Emit(ctx, "inspect", reduced, "call-block", !blocked, "worker") {
t.Fatal("missing tool result event")
}
if event["success"] != false || event["isError"] != true || event["result"] != reduced {
t.Fatalf("event = %#v", event)
}
if blocked && (event["blocked"] != true || event["status"] != "blocked" || event["executionId"] != result.ExecutionID) {
t.Fatalf("blocked event lost its classification: %#v", event)
}
if !blocked && event["blocked"] != nil {
t.Fatalf("ordinary failure was classified as blocked: %#v", event)
}
exec, ok := server.GetExecution(result.ExecutionID)
if !ok || exec.Result == nil || !exec.Result.IsError || exec.Result.Blocked != blocked {
t.Fatalf("display update lost result flags: %#v", exec)
}
wantStatus := "failed"
if blocked {
wantStatus = "blocked"
}
if ag.MCPExecutionStatus(result.ExecutionID) != wantStatus {
t.Fatalf("execution status = %q, want %q", exec.Status, wantStatus)
}
})
}
}
@@ -8,6 +8,7 @@ import (
"cyberstrike-ai/internal/agent"
"cyberstrike-ai/internal/einomcp"
"cyberstrike-ai/internal/mcp"
"github.com/cloudwego/eino/adk"
)
@@ -133,6 +134,15 @@ func (e *einoToolResultProgressEmitter) Emit(ctx context.Context, toolName, cont
}
if e.filesystemMonitorAgent != nil && e.mcpExecutionBinder != nil {
if execID := e.mcpExecutionBinder.ExecutionID(toolCallID); execID != "" {
// Use the execution record rather than parsing the rendered result:
// reduction can rewrite text without changing the safety decision.
if e.filesystemMonitorAgent.MCPExecutionStatus(execID) == mcp.ToolExecutionStatusBlocked {
data["blocked"] = true
data["status"] = mcp.ToolExecutionStatusBlocked
data["success"] = false
data["isError"] = true
data["executionId"] = execID
}
e.filesystemMonitorAgent.UpdateMCPExecutionDisplayResult(execID, content)
}
}
+6
View File
@@ -129,6 +129,10 @@ func permissionForRequest(method, fullPath string) string {
return "notification:read"
}
return "notification:write"
case path == "/tool-guard/test" && method == http.MethodPost:
return "config:read"
case path == "/tool-guard":
return crudPermission(method, "config")
case strings.HasPrefix(path, "/config"):
return crudPermission(method, "config")
case strings.HasPrefix(path, "/terminal"):
@@ -208,6 +212,8 @@ func resourceAllowed(c *gin.Context, db *database.DB) bool {
}
path := strings.TrimPrefix(c.FullPath(), "/api")
switch {
case path == "/tool-guard" && isMutationMethod(c.Request.Method):
return session.Scope == database.RBACScopeAll
case path == "/monitor/stats", path == "/monitor/calls-timeline":
// These APIs currently operate on process-global state. Until every MCP
// invocation and persisted execution record carries an immutable owner,
+291
View File
@@ -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
}
}
+231
View File
@@ -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()
}