CalvinBackup/CyberStrikeAI
1
0
Fork 0
mirror of https://github.com/Ed1s0nZ/CyberStrikeAI.git synced 2026-08-05 02:18:46 +02:00
Code Issues Packages Projects Releases Wiki Activity

Add files via upload

Browse Source
This commit is contained in:
公明
2026-07-24 15:52:16 +08:00
committed by GitHub
parent 8cb317cbd6
commit 5cbd828cad
6 changed files with 103 additions and 70 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/mcp/external_manager.go
+48 -43
View File
@@ -674,48 +674,6 @@ func (m *ExternalMCPManager) updateToolCache(name string, tools []Tool) {
// CallTool 调用外部MCP工具（返回执行ID）
func (m *ExternalMCPManager) CallTool(ctx context.Context, toolName string, args map[string]interface{}) (*ToolResult, string, error) {
_, authenticated := authctx.PrincipalFromContext(ctx)
m.mu.RLock()
authorizer := m.toolAuthorizer
m.mu.RUnlock()
if authorizer != nil {
if err := authorizer(ctx, toolName, args); err != nil {
return nil, "", fmt.Errorf("external tool authorization denied: %w", err)
}
} else if authenticated {
return nil, "", fmt.Errorf("external tool authorization policy is not configured")
}
// 解析工具名称：name::toolName
var mcpName, actualToolName string
if idx := findSubstring(toolName, "::"); idx > 0 {
mcpName = toolName[:idx]
actualToolName = toolName[idx+2:]
} else {
return nil, "", fmt.Errorf("无效的工具名称格式: %s", toolName)
}
client, exists := m.GetClient(mcpName)
if !exists {
return nil, "", fmt.Errorf("外部MCP客户端不存在: %s", mcpName)
}
if err := m.checkExternalMCPCircuit(mcpName); err != nil {
return nil, "", err
}
// 检查连接状态，如果未连接或状态为error，不允许调用
if !client.IsConnected() {
status := client.GetStatus()
if status == "error" {
// 获取错误信息（如果有）
errorMsg := m.GetError(mcpName)
if errorMsg != "" {
return nil, "", fmt.Errorf("外部MCP连接失败: %s (错误: %s)", mcpName, errorMsg)
}
return nil, "", fmt.Errorf("外部MCP连接失败: %s", mcpName)
}
return nil, "", fmt.Errorf("外部MCP客户端未连接: %s (状态: %s)", mcpName, status)
}
if m.executionService == nil {
m.executionService = NewExecutionService(m.storage, m.logger)
m.executionService.ConfigureToolResultMaxBytes(m.toolResultMaxBytes)
@@ -725,12 +683,57 @@ func (m *ExternalMCPManager) CallTool(ctx context.Context, toolName string, args
if principal, ok := authctx.PrincipalFromContext(ctx); ok {
ownerUserID = principal.UserID
}
var mcpName, actualToolName string
var client ExternalMCPClient
handle, err := m.executionService.Submit(ctx, ExecutionRequest{
ToolName: toolName,
Arguments: args,
ConversationID: MCPConversationIDFromContext(ctx),
OwnerUserID: ownerUserID,
PreRun: func(runCtx context.Context, exec *ToolExecution) (func(), error) {
_, authenticated := authctx.PrincipalFromContext(runCtx)
m.mu.RLock()
authorizer := m.toolAuthorizer
m.mu.RUnlock()
if authorizer != nil {
if err := authorizer(runCtx, toolName, args); err != nil {
return nil, fmt.Errorf("external tool authorization denied: %w", err)
}
} else if authenticated {
return nil, fmt.Errorf("external tool authorization policy is not configured")
}
// 解析工具名称：name::toolName
if idx := findSubstring(toolName, "::"); idx > 0 {
mcpName = toolName[:idx]
actualToolName = toolName[idx+2:]
} else {
return nil, fmt.Errorf("无效的工具名称格式: %s", toolName)
}
var exists bool
client, exists = m.GetClient(mcpName)
if !exists {
return nil, fmt.Errorf("外部MCP客户端不存在: %s", mcpName)
}
if err := m.checkExternalMCPCircuit(mcpName); err != nil {
return nil, err
}
// 检查连接状态，如果未连接或状态为error，不允许调用
if !client.IsConnected() {
status := client.GetStatus()
if status == "error" {
// 获取错误信息（如果有）
errorMsg := m.GetError(mcpName)
if errorMsg != "" {
return nil, fmt.Errorf("外部MCP连接失败: %s (错误: %s)", mcpName, errorMsg)
}
return nil, fmt.Errorf("外部MCP连接失败: %s", mcpName)
}
return nil, fmt.Errorf("外部MCP客户端未连接: %s (状态: %s)", mcpName, status)
}
release, acquireErr := m.acquireExternalMCPCallSlot(runCtx, mcpName)
if acquireErr != nil {
return nil, acquireErr
@@ -746,7 +749,9 @@ func (m *ExternalMCPManager) CallTool(ctx context.Context, toolName string, args
},
OnDone: func(exec *ToolExecution) {
failed := exec != nil && exec.Status != ToolExecutionStatusCompleted && exec.Status != ToolExecutionStatusCancelled
m.recordExternalMCPResult(mcpName, failed)
if mcpName != "" {
m.recordExternalMCPResult(mcpName, failed)
}
m.updateStats(toolName, failed)
},
})
internal/mcp/external_manager_test.go
+11 -1
View File
@@ -20,10 +20,20 @@ func TestExternalManagerEnforcesConfiguredAuthorizer(t *testing.T) {
return errors.New("denied by policy")
})
ctx := authctx.WithPrincipal(context.Background(), authctx.NewPrincipal("u1", "user", "assigned", map[string]bool{"agent:execute": true}))
_, _, err := manager.CallTool(ctx, "server::tool", map[string]interface{}{})
_, executionID, err := manager.CallTool(ctx, "server::tool", map[string]interface{}{})
if err == nil || !strings.Contains(err.Error(), "authorization denied") {
t.Fatalf("external call bypassed authorizer: %v", err)
}
if executionID == "" {
t.Fatal("denied external call should still return an execution id")
}
execution, ok := manager.GetExecution(executionID)
if !ok || execution == nil {
t.Fatalf("missing denied external execution %q", executionID)
}
if execution.Status != ToolExecutionStatusFailed || !strings.Contains(execution.Error, "denied by policy") {
t.Fatalf("denied external execution = %#v, want failed with policy error", execution)
}
}
func TestExternalMCPManager_AddOrUpdateConfig(t *testing.T) {
internal/mcp/server.go
+15 -19
View File
@@ -895,25 +895,6 @@ func (s *Server) GetAllTools() []Tool {
// CallTool 直接调用工具（用于内部调用）
func (s *Server) CallTool(ctx context.Context, toolName string, args map[string]interface{}) (*ToolResult, string, error) {
_, authenticated := authctx.PrincipalFromContext(ctx)
s.mu.RLock()
authorizer := s.toolAuthorizer
s.mu.RUnlock()
if authorizer != nil {
if err := authorizer(ctx, toolName, args); err != nil {
return nil, "", fmt.Errorf("tool authorization denied: %w", err)
}
} else if authenticated {
return nil, "", errors.New("tool authorization policy is not configured")
}
s.mu.RLock()
handler, exists := s.tools[toolName]
s.mu.RUnlock()
if !exists {
return nil, "", fmt.Errorf("工具 %s 未找到", toolName)
}
if s.executionService == nil {
s.executionService = NewExecutionService(s.storage, s.logger)
s.executionService.ConfigureToolResultMaxBytes(s.toolResultMaxBytes)
@@ -929,6 +910,21 @@ func (s *Server) CallTool(ctx context.Context, toolName string, args map[string]
ConversationID: MCPConversationIDFromContext(ctx),
OwnerUserID: ownerUserID,
Run: func(runCtx context.Context) (*ToolResult, error) {
_, authenticated := authctx.PrincipalFromContext(runCtx)
s.mu.RLock()
authorizer := s.toolAuthorizer
handler, exists := s.tools[toolName]
s.mu.RUnlock()
if authorizer != nil {
if err := authorizer(runCtx, toolName, args); err != nil {
return nil, fmt.Errorf("tool authorization denied: %w", err)
}
} else if authenticated {
return nil, errors.New("tool authorization policy is not configured")
}
if !exists {
return nil, fmt.Errorf("工具 %s 未找到", toolName)
}
return handler(runCtx, args)
},
OnDone: func(exec *ToolExecution) {
internal/mcp/server_authorization_test.go
+12 -1
View File
@@ -23,9 +23,20 @@ func TestToolAuthorizerIsUniversalAndExecutionKeepsOwner(t *testing.T) {
}
return nil
})
if _, _, err := server.CallTool(context.Background(), "echo", nil); err == nil {
_, deniedExecutionID, err := server.CallTool(context.Background(), "echo", nil)
if err == nil {
t.Fatal("tool call without principal was allowed")
}
if deniedExecutionID == "" {
t.Fatal("denied tool call should still return an execution id")
}
deniedExecution, ok := server.GetExecution(deniedExecutionID)
if !ok || deniedExecution == nil {
t.Fatalf("missing denied execution %q", deniedExecutionID)
}
if deniedExecution.Status != ToolExecutionStatusFailed || !strings.Contains(deniedExecution.Error, "principal required") {
t.Fatalf("denied execution = %#v, want failed with authorization error", deniedExecution)
}
ctx := authctx.WithPrincipal(context.Background(), authctx.NewPrincipal("u1", "user", "assigned", map[string]bool{"mcp:execute": true}))
_, executionID, err := server.CallTool(ctx, "echo", nil)
if err != nil {