feat: add configurable tool call blocking and monitoring

This commit is contained in:
Ed1s0nZ
2026-09-08 09:44:32 +08:00
parent c70da22de7
commit 6ad9ea2d13
54 changed files with 4635 additions and 152 deletions
+11 -7
View File
@@ -23,6 +23,7 @@ import (
"cyberstrike-ai/internal/mcp/builtin"
"cyberstrike-ai/internal/openai"
"cyberstrike-ai/internal/security"
"cyberstrike-ai/internal/toolguard"
"github.com/cloudwego/eino/schema"
"github.com/gin-gonic/gin"
@@ -95,6 +96,7 @@ type ConfigHandler struct {
db *database.DB
logger *zap.Logger
mu sync.RWMutex
toolGuard *toolguard.Manager
lastEmbeddingConfig *config.EmbeddingConfig // 上一次的嵌入模型配置(用于检测变更)
}
@@ -347,13 +349,13 @@ func (h *ConfigHandler) GetConfig(c *gin.Context) {
subAgentCount = len(agents.MergeYAMLAndMarkdown(h.config.MultiAgent.SubAgents, load.SubAgents))
}
multiPub := config.MultiAgentPublic{
Enabled: h.config.MultiAgent.Enabled,
RobotDefaultAgentMode: config.NormalizeRobotAgentMode(h.config.MultiAgent),
BatchUseMultiAgent: h.config.MultiAgent.BatchUseMultiAgent,
SubAgentCount: subAgentCount,
Orchestration: config.NormalizeMultiAgentOrchestration(h.config.MultiAgent.Orchestration),
PlanExecuteLoopMaxIterations: h.config.MultiAgent.PlanExecuteLoopMaxIterations,
SummarizationUserIntentLedgerMaxRunes: h.config.MultiAgent.EinoMiddleware.SummarizationUserIntentLedgerMaxRunesEffective(),
Enabled: h.config.MultiAgent.Enabled,
RobotDefaultAgentMode: config.NormalizeRobotAgentMode(h.config.MultiAgent),
BatchUseMultiAgent: h.config.MultiAgent.BatchUseMultiAgent,
SubAgentCount: subAgentCount,
Orchestration: config.NormalizeMultiAgentOrchestration(h.config.MultiAgent.Orchestration),
PlanExecuteLoopMaxIterations: h.config.MultiAgent.PlanExecuteLoopMaxIterations,
SummarizationUserIntentLedgerMaxRunes: h.config.MultiAgent.EinoMiddleware.SummarizationUserIntentLedgerMaxRunesEffective(),
SummarizationUserIntentLedgerEntryMaxRunes: h.config.MultiAgent.EinoMiddleware.SummarizationUserIntentLedgerEntryMaxRunesEffective(),
LatestUserMessageMaxRunes: h.config.MultiAgent.EinoMiddleware.LatestUserMessageMaxRunesEffective(),
LatestUserMessageHeadRunes: h.config.MultiAgent.EinoMiddleware.LatestUserMessageHeadRunesEffective(),
@@ -1746,6 +1748,8 @@ func (h *ConfigHandler) ApplyConfig(c *gin.Context) {
// saveConfig 保存配置到文件
func (h *ConfigHandler) saveConfig() error {
configFileMu.Lock()
defer configFileMu.Unlock()
h.config.NormalizeAIProviderProfiles()
// 读取现有配置文件并创建备份
+8
View File
@@ -0,0 +1,8 @@
package handler
import "sync"
// configFileMu serializes complete read-modify-write transactions across
// handlers that share config.yaml. Per-handler locks cannot prevent lost
// updates when another settings page saves a different YAML section.
var configFileMu sync.Mutex
+2
View File
@@ -379,6 +379,8 @@ func (h *ExternalMCPHandler) isEnabled(cfg config.ExternalMCPServerConfig) bool
// saveConfig 保存配置到文件
func (h *ExternalMCPHandler) saveConfig() error {
configFileMu.Lock()
defer configFileMu.Unlock()
data, err := os.ReadFile(h.configPath)
if err != nil {
return fmt.Errorf("读取配置文件失败: %w", err)
+19 -8
View File
@@ -76,6 +76,7 @@ type MonitorStatsSummary struct {
TotalCalls int `json:"totalCalls"`
SuccessCalls int `json:"successCalls"`
FailedCalls int `json:"failedCalls"`
BlockedCalls int `json:"blockedCalls"`
LastCallTime *time.Time `json:"lastCallTime,omitempty"`
ToolCount int `json:"toolCount"`
}
@@ -171,6 +172,8 @@ func summarizeAccessibleExecutionPage(executions []*mcp.ToolExecution, topN int)
stat.FailedCalls++
} else if exec.Status == "completed" {
stat.SuccessCalls++
} else if exec.Status == mcp.ToolExecutionStatusBlocked {
stat.BlockedCalls++
}
started := exec.StartTime
if stat.LastCallTime == nil || started.After(*stat.LastCallTime) {
@@ -448,6 +451,7 @@ func dbStatsSummaryToMonitor(result *database.ToolStatsSummaryResult) *MonitorSt
TotalCalls: result.Summary.TotalCalls,
SuccessCalls: result.Summary.SuccessCalls,
FailedCalls: result.Summary.FailedCalls,
BlockedCalls: result.Summary.BlockedCalls,
ToolCount: result.Summary.ToolCount,
}
if result.Summary.LastCallTime != nil {
@@ -472,6 +476,7 @@ func summarizeToolStats(stats map[string]*mcp.ToolStats, topN int) (*MonitorStat
summary.TotalCalls += stat.TotalCalls
summary.SuccessCalls += stat.SuccessCalls
summary.FailedCalls += stat.FailedCalls
summary.BlockedCalls += stat.BlockedCalls
if stat.LastCallTime != nil && (summary.LastCallTime == nil || stat.LastCallTime.After(*summary.LastCallTime)) {
t := *stat.LastCallTime
summary.LastCallTime = &t
@@ -528,6 +533,7 @@ func (h *MonitorHandler) loadStatsMap() map[string]*mcp.ToolStats {
existing.TotalCalls += v.TotalCalls
existing.SuccessCalls += v.SuccessCalls
existing.FailedCalls += v.FailedCalls
existing.BlockedCalls += v.BlockedCalls
// 使用最新的调用时间
if v.LastCallTime != nil && (existing.LastCallTime == nil || v.LastCallTime.After(*existing.LastCallTime)) {
existing.LastCallTime = v.LastCallTime
@@ -734,9 +740,10 @@ func (h *MonitorHandler) GetStats(c *gin.Context) {
// CallsTimelinePoint 调用趋势数据点
type CallsTimelinePoint struct {
T time.Time `json:"t"`
Total int `json:"total"`
Failed int `json:"failed"`
T time.Time `json:"t"`
Total int `json:"total"`
Failed int `json:"failed"`
Blocked int `json:"blocked"`
}
// CallsTimelineSummary 调用趋势汇总
@@ -778,7 +785,7 @@ func truncateToBucket(t time.Time, bucketSize time.Duration, dailyBuckets bool)
return t.Truncate(bucketSize)
}
func buildCallsTimelinePoints(cfg callsTimelineConfig, buckets map[time.Time]struct{ total, failed int }) []CallsTimelinePoint {
func buildCallsTimelinePoints(cfg callsTimelineConfig, buckets map[time.Time]struct{ total, failed, blocked int }) []CallsTimelinePoint {
now := time.Now()
start := truncateToBucket(now.Add(-cfg.duration), cfg.bucketSize, cfg.dailyBuckets)
end := truncateToBucket(now, cfg.bucketSize, cfg.dailyBuckets)
@@ -787,9 +794,10 @@ func buildCallsTimelinePoints(cfg callsTimelineConfig, buckets map[time.Time]str
for current := start; !current.After(end); current = current.Add(cfg.bucketSize) {
val := buckets[current]
points = append(points, CallsTimelinePoint{
T: current,
Total: val.total,
Failed: val.failed,
T: current,
Total: val.total,
Failed: val.failed,
Blocked: val.blocked,
})
}
return points
@@ -797,7 +805,7 @@ func buildCallsTimelinePoints(cfg callsTimelineConfig, buckets map[time.Time]str
func (h *MonitorHandler) loadCallsTimeline(cfg callsTimelineConfig) []CallsTimelinePoint {
since := time.Now().Add(-cfg.duration)
bucketMap := make(map[time.Time]struct{ total, failed int })
bucketMap := make(map[time.Time]struct{ total, failed, blocked int })
if h.db != nil {
dbBuckets, err := h.db.LoadCallsTimeline(since, cfg.dailyBuckets)
@@ -809,6 +817,7 @@ func (h *MonitorHandler) loadCallsTimeline(cfg callsTimelineConfig) []CallsTimel
entry := bucketMap[key]
entry.total += b.Total
entry.failed += b.Failed
entry.blocked += b.Blocked
bucketMap[key] = entry
}
return buildCallsTimelinePoints(cfg, bucketMap)
@@ -824,6 +833,8 @@ func (h *MonitorHandler) loadCallsTimeline(cfg callsTimelineConfig) []CallsTimel
entry.total++
if monitorStatusCountsAsFailed(exec.Status) {
entry.failed++
} else if exec.Status == mcp.ToolExecutionStatusBlocked {
entry.blocked++
}
bucketMap[key] = entry
}
+175
View File
@@ -0,0 +1,175 @@
package handler
import (
"bytes"
"encoding/json"
"fmt"
"io"
"net/http"
"os"
"path/filepath"
"cyberstrike-ai/internal/toolguard"
"github.com/gin-gonic/gin"
"gopkg.in/yaml.v3"
)
func (h *ConfigHandler) SetToolGuard(manager *toolguard.Manager) {
h.mu.Lock()
defer h.mu.Unlock()
h.toolGuard = manager
}
func (h *ConfigHandler) GetToolGuard(c *gin.Context) {
h.mu.RLock()
defer h.mu.RUnlock()
if h.toolGuard == nil {
c.JSON(http.StatusServiceUnavailable, gin.H{"error": "调用拦截服务未初始化"})
return
}
c.JSON(http.StatusOK, h.toolGuard.Config())
}
// decodeToolGuardRequest bounds both config and dry-run inputs, rejects unknown
// fields and trailing JSON, and never invokes an actual tool.
func decodeToolGuardRequest(c *gin.Context, dst interface{}) error {
c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, 1<<20)
decoder := json.NewDecoder(c.Request.Body)
decoder.DisallowUnknownFields()
decoder.UseNumber()
if err := decoder.Decode(dst); err != nil {
return err
}
if err := decoder.Decode(new(interface{})); err != io.EOF {
return fmt.Errorf("请求必须只包含一个 JSON 对象")
}
return nil
}
func (h *ConfigHandler) UpdateToolGuard(c *gin.Context) {
var req struct {
Enabled *bool `json:"enabled"`
Rules *[]toolguard.Rule `json:"rules"`
}
if err := decodeToolGuardRequest(c, &req); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": "无效的调用拦截配置: " + err.Error()})
return
}
if req.Enabled == nil || req.Rules == nil {
c.JSON(http.StatusBadRequest, gin.H{"error": "必须明确提供 enabled 和 rules;清空规则请提供空数组"})
return
}
cfg := toolguard.Config{Enabled: *req.Enabled, Rules: *req.Rules}
if _, err := toolguard.Compile(cfg); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
h.mu.Lock()
defer h.mu.Unlock()
if h.toolGuard == nil {
c.JSON(http.StatusServiceUnavailable, gin.H{"error": "调用拦截服务未初始化"})
return
}
// Commit the file first; a validation/write failure must leave the current
// effective policy and in-memory config intact.
if err := h.saveToolGuardConfig(cfg); err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": "保存调用拦截配置失败: " + err.Error()})
return
}
if err := h.toolGuard.Update(cfg); err != nil {
// The same immutable input was compiled above, so this cannot fail
// unless validation gains an additional runtime dependency.
c.JSON(http.StatusInternalServerError, gin.H{"error": "应用调用拦截配置失败: " + err.Error()})
return
}
h.config.ToolGuard = &cfg
if h.audit != nil {
h.audit.RecordOK(c, "config", "tool_guard_update", "更新调用拦截规则", "config", "tool_guard", map[string]interface{}{
"enabled": cfg.Enabled, "rule_count": len(cfg.Rules),
})
}
c.JSON(http.StatusOK, h.toolGuard.Config())
}
func (h *ConfigHandler) TestToolGuard(c *gin.Context) {
var req struct {
Config *toolguard.Config `json:"config"`
ToolName string `json:"toolName"`
Arguments map[string]interface{} `json:"arguments"`
}
if err := decodeToolGuardRequest(c, &req); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": "无效的试匹配参数: " + err.Error()})
return
}
if req.Config == nil {
c.JSON(http.StatusBadRequest, gin.H{"error": "请提供待测试的 config"})
return
}
policy, err := toolguard.Compile(*req.Config)
if err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
if match := policy.Check(req.ToolName, req.Arguments); match != nil {
c.JSON(http.StatusOK, gin.H{"blocked": true, "match": match})
return
}
c.JSON(http.StatusOK, gin.H{"blocked": false})
}
// saveToolGuardConfig changes only this YAML section, preserving unrelated
// settings/comments and file permissions. Rename makes the write atomic.
// h.mu protects the runtime configuration; configFileMu also covers independent
// writers such as ExternalMCPHandler.
func (h *ConfigHandler) saveToolGuardConfig(cfg toolguard.Config) error {
configFileMu.Lock()
defer configFileMu.Unlock()
path, err := filepath.EvalSymlinks(h.configPath)
if err != nil {
return err
}
doc, err := loadYAMLDocument(path)
if err != nil {
return err
}
var node yaml.Node
if err := node.Encode(cfg); err != nil {
return err
}
_, value := ensureKeyValue(doc.Content[0], "tool_guard")
*value = node
var buf bytes.Buffer
encoder := yaml.NewEncoder(&buf)
encoder.SetIndent(2)
if err := encoder.Encode(doc); err != nil {
return err
}
if err := encoder.Close(); err != nil {
return err
}
info, err := os.Stat(path)
if err != nil {
return err
}
tmp, err := os.CreateTemp(filepath.Dir(path), ".tool-guard-*.yaml")
if err != nil {
return err
}
defer os.Remove(tmp.Name())
defer tmp.Close()
if err := tmp.Chmod(info.Mode().Perm()); err != nil {
return err
}
if _, err := tmp.Write(buf.Bytes()); err != nil {
return err
}
if err := tmp.Sync(); err != nil {
return err
}
if err := tmp.Close(); err != nil {
return err
}
return os.Rename(tmp.Name(), path)
}
+193
View File
@@ -0,0 +1,193 @@
package handler
import (
"bytes"
"encoding/json"
"net/http"
"net/http/httptest"
"os"
"path/filepath"
"reflect"
"strings"
"sync"
"testing"
"cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/database"
"cyberstrike-ai/internal/security"
"cyberstrike-ai/internal/toolguard"
"github.com/gin-gonic/gin"
"go.uber.org/zap"
)
func newToolGuardTestHandler(t *testing.T) *ConfigHandler {
t.Helper()
path := filepath.Join(t.TempDir(), "config.yaml")
if err := os.WriteFile(path, []byte("# keep this comment\nserver:\n port: 8123\nhitl:\n tool_whitelist: [read_file]\n"), 0600); err != nil {
t.Fatal(err)
}
manager, err := toolguard.NewManager(toolguard.DefaultConfig())
if err != nil {
t.Fatal(err)
}
return &ConfigHandler{configPath: path, config: &config.Config{}, toolGuard: manager}
}
func toolGuardRequest(t *testing.T, handler gin.HandlerFunc, body interface{}) *httptest.ResponseRecorder {
t.Helper()
data, err := json.Marshal(body)
if err != nil {
t.Fatal(err)
}
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest(http.MethodPut, "/api/tool-guard", bytes.NewReader(data))
c.Request.Header.Set("Content-Type", "application/json")
handler(c)
return w
}
func TestToolGuardSavePersistsAndAppliesWithoutChangingHITL(t *testing.T) {
h := newToolGuardTestHandler(t)
cfg := toolguard.DefaultConfig()
cfg.Rules[0].Message = "识别到 {match},禁止攻击政府网站,请检查目标。"
w := toolGuardRequest(t, h.UpdateToolGuard, cfg)
if w.Code != http.StatusOK {
t.Fatalf("save: %d %s", w.Code, w.Body.String())
}
loaded, err := config.Load(h.configPath)
if err != nil {
t.Fatal(err)
}
if !reflect.DeepEqual(loaded.EffectiveToolGuard(), cfg) || !reflect.DeepEqual(h.toolGuard.Config(), cfg) {
t.Fatal("saved and effective policies differ")
}
if loaded.Server.Port != 8123 || !reflect.DeepEqual(loaded.Hitl.ToolWhitelist, []string{"read_file"}) {
t.Fatal("unrelated configuration was changed")
}
info, _ := os.Stat(h.configPath)
data, _ := os.ReadFile(h.configPath)
if info.Mode().Perm() != 0600 || !strings.Contains(string(data), "# keep this comment") {
t.Fatal("file permissions or comments were lost")
}
match := h.toolGuard.Check("scan", map[string]interface{}{"target": "agency.gov.cn"})
if match == nil || !strings.Contains(match.Message, "agency.gov.cn") {
t.Fatalf("updated message not applied: %+v", match)
}
cfg.Enabled = false
w = toolGuardRequest(t, h.UpdateToolGuard, cfg)
if w.Code != http.StatusOK || h.toolGuard.Check("scan", map[string]interface{}{"target": "agency.gov"}) != nil {
t.Fatal("explicitly disabling protection did not apply")
}
}
func TestToolGuardInvalidAndFailedSaveKeepEffectivePolicy(t *testing.T) {
h := newToolGuardTestHandler(t)
before, _ := os.ReadFile(h.configPath)
cfg := toolguard.DefaultConfig()
cfg.Enabled = false
cfg.Rules[0].Pattern = "["
for _, body := range []interface{}{cfg, map[string]interface{}{}, nil, map[string]interface{}{"enabled": false, "rules": nil}} {
w := toolGuardRequest(t, h.UpdateToolGuard, body)
if w.Code != http.StatusBadRequest {
t.Fatalf("invalid update accepted: %d %s", w.Code, w.Body.String())
}
}
after, _ := os.ReadFile(h.configPath)
if !bytes.Equal(before, after) || !h.toolGuard.Config().Enabled {
t.Fatal("invalid input changed protection")
}
h.configPath = filepath.Join(t.TempDir(), "missing", "config.yaml")
cfg = toolguard.DefaultConfig()
cfg.Enabled = false
w := toolGuardRequest(t, h.UpdateToolGuard, cfg)
if w.Code != http.StatusInternalServerError || !h.toolGuard.Config().Enabled || h.config.ToolGuard != nil {
t.Fatal("failed persistence changed live configuration")
}
}
func TestToolGuardDryRunUsesUnsavedPolicyWithoutMutation(t *testing.T) {
h := newToolGuardTestHandler(t)
cfg := toolguard.DefaultConfig()
cfg.Rules[0].Pattern = "example\\.org"
w := toolGuardRequest(t, h.TestToolGuard, map[string]interface{}{
"config": cfg, "toolName": "scan", "arguments": map[string]interface{}{"target": "example.org"},
})
var got struct {
Blocked bool `json:"blocked"`
Match *toolguard.Match `json:"match"`
}
if w.Code != http.StatusOK || json.Unmarshal(w.Body.Bytes(), &got) != nil || !got.Blocked || got.Match == nil || got.Match.MatchedText != "example.org" {
t.Fatalf("dry run failed: %s", w.Body.String())
}
if !reflect.DeepEqual(h.toolGuard.Config(), toolguard.DefaultConfig()) || h.config.ToolGuard != nil {
t.Fatal("dry run changed live configuration")
}
w = toolGuardRequest(t, h.TestToolGuard, map[string]interface{}{"config": cfg, "arguments": []string{"example.org"}})
if w.Code != http.StatusBadRequest {
t.Fatal("non-object tool arguments accepted")
}
}
func TestToolGuardRoutesEnforceConfigurationPermissions(t *testing.T) {
gin.SetMode(gin.TestMode)
for _, tc := range []struct {
method, path, permission, scope string
want int
}{
{"GET", "/api/tool-guard", "hitl:read", database.RBACScopeAll, 403},
{"PUT", "/api/tool-guard", "hitl:write", database.RBACScopeAll, 403},
{"GET", "/api/tool-guard", "config:read", database.RBACScopeAll, 200},
{"POST", "/api/tool-guard/test", "config:read", database.RBACScopeAll, 200},
{"PUT", "/api/tool-guard", "config:write", database.RBACScopeAll, 200},
{"PUT", "/api/tool-guard", "config:write", database.RBACScopeOwn, 403},
} {
t.Run(tc.method+tc.permission+tc.scope, func(t *testing.T) {
r := gin.New()
r.Use(func(c *gin.Context) {
c.Set(security.ContextSessionKey, security.Session{UserID: "test", Permissions: map[string]bool{tc.permission: true}, Scope: tc.scope})
})
r.Use(security.RBACMiddleware(&database.DB{}))
r.Handle(tc.method, tc.path, func(c *gin.Context) { c.Status(200) })
w := httptest.NewRecorder()
r.ServeHTTP(w, httptest.NewRequest(tc.method, tc.path, nil))
if w.Code != tc.want {
t.Fatalf("got %d, want %d: %s", w.Code, tc.want, w.Body.String())
}
})
}
}
func TestToolGuardConcurrentOtherSettingsSavePreservesPolicy(t *testing.T) {
h := newToolGuardTestHandler(t)
external := &ExternalMCPHandler{configPath: h.configPath, config: h.config, logger: zap.NewNop()}
cfg := toolguard.DefaultConfig()
cfg.Rules[0].Message = "持久化策略 {match}"
var wg sync.WaitGroup
errors := make(chan error, 2)
for _, save := range []func() error{func() error { return h.saveToolGuardConfig(cfg) }, external.saveConfig} {
wg.Add(1)
go func(save func() error) {
defer wg.Done()
for i := 0; i < 20; i++ {
if err := save(); err != nil {
errors <- err
return
}
}
}(save)
}
wg.Wait()
close(errors)
for err := range errors {
t.Fatal(err)
}
loaded, err := config.Load(h.configPath)
if err != nil {
t.Fatal(err)
}
if !reflect.DeepEqual(loaded.EffectiveToolGuard(), cfg) {
t.Fatal("another settings save overwrote the tool guard policy")
}
}