mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-09-15 14:15:27 +02:00
feat: add configurable tool call blocking and monitoring
This commit is contained in:
+34
-5
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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"
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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 = ©
|
||||
}
|
||||
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
|
||||
|
||||
@@ -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()
|
||||
}
|
||||
@@ -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()
|
||||
|
||||
// 读取现有配置文件并创建备份
|
||||
|
||||
@@ -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
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
@@ -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
|
||||
}
|
||||
|
||||
@@ -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}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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()
|
||||
}
|
||||
Reference in New Issue
Block a user