feat: add configurable tool call blocking and monitoring

2026-09-18 15:42:26 +02:00 · 2026-09-08 09:44:32 +08:00
parent c70da22de7
commit 6ad9ea2d13
54 changed files with 4635 additions and 152 deletions
@@ -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")
+	}
+}