mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-09-14 21:55:31 +02:00
feat: add configurable tool call blocking and monitoring
This commit is contained in:
@@ -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")
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user