mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-09-16 06:35:29 +02:00
240 lines
8.8 KiB
Go
240 lines
8.8 KiB
Go
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)
|
|
}
|
|
}
|