diff --git a/internal/app/app.go b/internal/app/app.go index d3dfc32d..213edce1 100644 --- a/internal/app/app.go +++ b/internal/app/app.go @@ -8,6 +8,7 @@ import ( "fmt" "net" "net/http" + "net/url" "os" "path/filepath" "strings" @@ -16,6 +17,7 @@ import ( "cyberstrike-ai/internal/agent" "cyberstrike-ai/internal/audit" + "cyberstrike-ai/internal/authctx" "cyberstrike-ai/internal/c2" "cyberstrike-ai/internal/config" "cyberstrike-ai/internal/database" @@ -123,6 +125,7 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error // 创建MCP服务器(带数据库持久化) mcpServer := mcp.NewServerWithStorage(log.Logger, db) + mcpServer.SetToolAuthorizer(mcpToolAuthorizer(db)) mcpServer.ConfigureHTTPToolCallTimeoutFromAgentMinutes(cfg.Agent.ToolTimeoutMinutes) // 创建安全工具执行器 @@ -146,6 +149,7 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error // 创建外部MCP管理器(使用与内部MCP服务器相同的存储) externalMCPMgr := mcp.NewExternalMCPManagerWithStorage(log.Logger, db) + externalMCPMgr.SetToolAuthorizer(externalMCPToolAuthorizer()) if cfg.ExternalMCP.Servers != nil { externalMCPMgr.LoadConfigs(&cfg.ExternalMCP) // 启动所有启用的外部MCP客户端 @@ -372,7 +376,7 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error vulnerabilityHandler.SetAudit(auditSvc) webshellHandler := handler.NewWebShellHandler(log.Logger, db) webshellHandler.SetAudit(auditSvc) - chatUploadsHandler := handler.NewChatUploadsHandler(log.Logger) + chatUploadsHandler := handler.NewChatUploadsHandler(log.Logger, db) chatUploadsHandler.SetAudit(auditSvc) registerWebshellTools(mcpServer, db, webshellHandler, log.Logger) registerWebshellManagementTools(mcpServer, db, webshellHandler, log.Logger) @@ -541,6 +545,7 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error terminalHandler, app.c2Handler, auditHandler, + auditSvc, rbacHandler, mcpServer, authManager, @@ -554,17 +559,30 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error // mcpHandlerWithAuth 在鉴权通过后转发到 MCP 处理;若配置了 auth_header 则校验请求头,否则直接放行 func (a *App) mcpHandlerWithAuth(w http.ResponseWriter, r *http.Request) { cfg := a.config.MCP - if cfg.AuthHeader != "" { - actual := []byte(r.Header.Get(cfg.AuthHeader)) - expected := []byte(cfg.AuthHeaderValue) - if subtle.ConstantTimeCompare(actual, expected) != 1 { - a.logger.Logger.Debug("MCP 鉴权失败:header 缺失或值不匹配", zap.String("header", cfg.AuthHeader)) - w.Header().Set("Content-Type", "application/json") - w.WriteHeader(http.StatusUnauthorized) - w.Write([]byte(`{"error":"unauthorized"}`)) + if authHeader := strings.TrimSpace(r.Header.Get("Authorization")); len(authHeader) > 7 && strings.EqualFold(authHeader[:7], "Bearer ") { + if session, ok := a.auth.ValidateToken(strings.TrimSpace(authHeader[7:])); ok && session.Permissions["mcp:execute"] { + principal := authctx.NewPrincipalWithScopes(session.UserID, session.Username, session.Scope, session.Permissions, session.PermissionScopes) + a.mcpServer.HandleHTTP(w, r.WithContext(authctx.WithPrincipal(r.Context(), principal))) return } } + if !cfg.AllowGlobalAccess || strings.TrimSpace(cfg.AuthHeader) == "" || strings.TrimSpace(cfg.AuthHeaderValue) == "" { + http.Error(w, "use an authorized user bearer token; global MCP service access is disabled", http.StatusUnauthorized) + return + } + if subtle.ConstantTimeCompare([]byte(r.Header.Get(cfg.AuthHeader)), []byte(cfg.AuthHeaderValue)) != 1 { + a.logger.Logger.Debug("MCP 鉴权失败:header 缺失或值不匹配", zap.String("header", cfg.AuthHeader)) + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusUnauthorized) + w.Write([]byte(`{"error":"unauthorized"}`)) + return + } + permissions := make(map[string]bool, len(security.PermissionCatalog)) + for permission := range security.PermissionCatalog { + permissions[permission] = true + } + principal := authctx.NewPrincipal("service:mcp", "mcp-service", database.RBACScopeAll, permissions) + r = r.WithContext(authctx.WithPrincipal(r.Context(), principal)) a.mcpServer.HandleHTTP(w, r) } @@ -825,6 +843,7 @@ func setupRoutes( terminalHandler *handler.TerminalHandler, c2Handler *handler.C2Handler, auditHandler *handler.AuditHandler, + auditSvc *audit.Service, rbacHandler *handler.RBACHandler, mcpServer *mcp.Server, authManager *security.AuthManager, @@ -835,8 +854,9 @@ func setupRoutes( // 认证相关路由 authRoutes := api.Group("/auth") + loginRL := security.NewRateLimiter(10, 1*time.Minute) { - authRoutes.POST("/login", authHandler.Login) + authRoutes.POST("/login", security.RateLimitMiddleware(loginRL), authHandler.Login) authRoutes.POST("/logout", security.AuthMiddleware(authManager), authHandler.Logout) authRoutes.POST("/change-password", security.AuthMiddleware(authManager), security.RequirePermission("auth:self"), authHandler.ChangePassword) authRoutes.GET("/validate", security.AuthMiddleware(authManager), authHandler.Validate) @@ -856,7 +876,15 @@ func setupRoutes( protected := api.Group("") protected.Use(security.AuthMiddleware(authManager)) - protected.Use(security.RBACMiddleware(app.db)) + protected.Use(security.RBACMiddlewareWithDenyHook(app.db, func(c *gin.Context, reason, permission string) { + if auditSvc != nil { + auditSvc.Record(c, audit.Entry{ + Level: "warn", Category: "rbac", Action: "access_denied", Result: "failure", + Message: "RBAC 拒绝访问", ResourceType: "route", ResourceID: c.FullPath(), + Detail: map[string]interface{}{"reason": reason, "permission": permission, "method": c.Request.Method}, + }) + } + })) { protected.GET("/rbac/me", rbacHandler.Me) protected.GET("/rbac/metadata", rbacHandler.Metadata) @@ -1490,7 +1518,13 @@ func registerWebshellManagementTools(mcpServer *mcp.Server, db *database.DB, web }, } listHandler := func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) { - connections, err := db.ListWebshellConnections() + connections := []database.WebShellConnection{} + var err error + if principal, ok := authctx.PrincipalFromContext(ctx); ok { + connections, err = db.ListWebshellConnectionsForAccess(principal.UserID, principal.ScopeFor("webshell:read")) + } else { + return &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: "缺少认证身份"}}, IsError: true}, nil + } if err != nil { return &mcp.ToolResult{ Content: []mcp.Content{{Type: "text", Text: "获取连接列表失败: " + err.Error()}}, @@ -1605,6 +1639,10 @@ func registerWebshellManagementTools(mcpServer *mcp.Server, db *database.DB, web IsError: true, }, nil } + if principal, ok := authctx.PrincipalFromContext(ctx); ok { + _ = db.SetResourceOwner("webshell", conn.ID, principal.UserID) + _ = db.AssignResourceToUser(principal.UserID, "webshell", conn.ID) + } return &mcp.ToolResult{ Content: []mcp.Content{{ @@ -2000,8 +2038,17 @@ func initializeKnowledge( // corsMiddleware CORS中间件 func corsMiddleware() gin.HandlerFunc { return func(c *gin.Context) { - c.Writer.Header().Set("Access-Control-Allow-Origin", "*") - c.Writer.Header().Set("Access-Control-Allow-Credentials", "true") + origin := strings.TrimSpace(c.GetHeader("Origin")) + if origin != "" { + parsed, err := url.Parse(origin) + if err != nil || parsed.Host == "" || !strings.EqualFold(parsed.Host, c.Request.Host) { + c.AbortWithStatusJSON(http.StatusForbidden, gin.H{"error": "cross-origin request denied"}) + return + } + c.Writer.Header().Set("Access-Control-Allow-Origin", origin) + c.Writer.Header().Set("Access-Control-Allow-Credentials", "true") + c.Writer.Header().Add("Vary", "Origin") + } c.Writer.Header().Set("Access-Control-Allow-Headers", "Content-Type, Content-Length, Accept-Encoding, X-CSRF-Token, Authorization, accept, origin, Cache-Control, X-Requested-With") c.Writer.Header().Set("Access-Control-Allow-Methods", "POST, OPTIONS, GET, PUT, DELETE") diff --git a/internal/app/c2_tools.go b/internal/app/c2_tools.go index 92f53591..b9d95f30 100644 --- a/internal/app/c2_tools.go +++ b/internal/app/c2_tools.go @@ -4,11 +4,13 @@ import ( "context" "encoding/json" "fmt" + "path/filepath" "strconv" "strings" "time" "cyberstrike-ai/internal/agent" + "cyberstrike-ai/internal/authctx" "cyberstrike-ai/internal/c2" "cyberstrike-ai/internal/database" "cyberstrike-ai/internal/mcp" @@ -66,16 +68,16 @@ tcp_reverse 默认仅接受 CSB1 加密 Beacon(AES-GCM + ImplantToken)才登 InputSchema: map[string]interface{}{ "type": "object", "properties": map[string]interface{}{ - "action": map[string]interface{}{"type": "string", "description": "操作: list/get/create/update/start/stop/delete", "enum": []string{"list", "get", "create", "update", "start", "stop", "delete"}}, - "listener_id": map[string]interface{}{"type": "string", "description": "监听器 ID(get/update/start/stop/delete 需要)"}, - "name": map[string]interface{}{"type": "string", "description": "监听器名称(create/update)"}, - "type": map[string]interface{}{"type": "string", "description": "监听器类型(create)", "enum": []string{"tcp_reverse", "http_beacon", "https_beacon", "websocket"}}, + "action": map[string]interface{}{"type": "string", "description": "操作: list/get/create/update/start/stop/delete", "enum": []string{"list", "get", "create", "update", "start", "stop", "delete"}}, + "listener_id": map[string]interface{}{"type": "string", "description": "监听器 ID(get/update/start/stop/delete 需要)"}, + "name": map[string]interface{}{"type": "string", "description": "监听器名称(create/update)"}, + "type": map[string]interface{}{"type": "string", "description": "监听器类型(create)", "enum": []string{"tcp_reverse", "http_beacon", "https_beacon", "websocket"}}, "bind_host": map[string]interface{}{"type": "string", "description": "绑定地址,默认 127.0.0.1;外网监听常用 0.0.0.0"}, "callback_host": map[string]interface{}{"type": "string", "description": "可选:植入端/Payload 回连主机名(公网 IP 或域名)。写入 config_json;生成 oneliner/beacon 时优先于 bind_host。update 时传入空字符串可清除"}, - "bind_port": map[string]interface{}{"type": "integer", "description": fmt.Sprintf("绑定端口(create 必填)。须 ≠ %d(当前本服务 Web/API 端口,配置 server.port)", webListenPort), "minimum": 1, "maximum": 65535}, - "profile_id": map[string]interface{}{"type": "string", "description": "Malleable Profile ID"}, - "remark": map[string]interface{}{"type": "string", "description": "备注"}, - "config": map[string]interface{}{"type": "object", "description": "高级配置(beacon 路径/TLS/OPSEC 等),create/update 可用。tcp_reverse 可选 allow_legacy_shell:true 允许未加密经典 shell(默认 false)"}, + "bind_port": map[string]interface{}{"type": "integer", "description": fmt.Sprintf("绑定端口(create 必填)。须 ≠ %d(当前本服务 Web/API 端口,配置 server.port)", webListenPort), "minimum": 1, "maximum": 65535}, + "profile_id": map[string]interface{}{"type": "string", "description": "Malleable Profile ID"}, + "remark": map[string]interface{}{"type": "string", "description": "备注"}, + "config": map[string]interface{}{"type": "object", "description": "高级配置(beacon 路径/TLS/OPSEC 等),create/update 可用。tcp_reverse 可选 allow_legacy_shell:true 允许未加密经典 shell(默认 false)"}, }, "required": []string{"action"}, }, @@ -85,7 +87,7 @@ tcp_reverse 默认仅接受 CSB1 加密 Beacon(AES-GCM + ImplantToken)才登 switch action { case "list": - listeners, err := m.DB().ListC2Listeners() + listeners, err := m.DB().ListC2ListenersForAccess(c2ToolAccess(ctx)) if err != nil { return makeC2Result(nil, err) } @@ -128,6 +130,10 @@ tcp_reverse 默认仅接受 CSB1 加密 Beacon(AES-GCM + ImplantToken)才登 if err != nil { return makeC2Result(nil, err) } + if principal, ok := authctx.PrincipalFromContext(ctx); ok { + _ = m.DB().SetResourceOwner("c2_listener", listener.ID, principal.UserID) + _ = m.DB().AssignResourceToUser(principal.UserID, "c2_listener", listener.ID) + } implantToken := listener.ImplantToken listener.EncryptionKey = "" listener.ImplantToken = "" @@ -264,7 +270,7 @@ func registerC2SessionTool(s *mcp.Server, m *c2.Manager, l *zap.Logger) { if v, ok := params["suspicious"].(bool); ok && v { filter.Suspicious = true } - sessions, err := m.DB().ListC2Sessions(filter) + sessions, err := m.DB().ListC2SessionsForAccess(filter, c2ToolAccess(ctx)) return makeC2Result(map[string]interface{}{"sessions": sessions, "count": len(sessions)}, err) case "get": @@ -494,7 +500,7 @@ func registerC2TaskManageTool(s *mcp.Server, m *c2.Manager, l *zap.Logger) { if limit := int(getFloat64(params, "limit")); limit > 0 { filter.Limit = limit } - tasks, err := m.DB().ListC2Tasks(filter) + tasks, err := m.DB().ListC2TasksForAccess(filter, c2ToolAccess(ctx)) return makeC2Result(map[string]interface{}{"tasks": tasks, "count": len(tasks)}, err) case "cancel": @@ -602,6 +608,9 @@ func registerC2PayloadTool(s *mcp.Server, m *c2.Manager, l *zap.Logger, webListe if err != nil { return makeC2Result(nil, err) } + if principal, ok := authctx.PrincipalFromContext(ctx); ok { + _ = m.DB().RecordC2PayloadArtifact(filepath.Base(result.OutputPath), result.PayloadID, result.ListenerID, principal.UserID) + } return makeC2Result(map[string]interface{}{ "payload_id": result.PayloadID, "download_path": result.DownloadPath, "os": result.OS, "arch": result.Arch, "size_bytes": result.SizeBytes, @@ -648,11 +657,19 @@ func registerC2EventTool(s *mcp.Server, m *c2.Manager, l *zap.Logger) { filter.Since = &t } } - events, err := m.DB().ListC2Events(filter) + events, err := m.DB().ListC2EventsForAccess(filter, c2ToolAccess(ctx)) return makeC2Result(map[string]interface{}{"events": events, "count": len(events)}, err) }) } +func c2ToolAccess(ctx context.Context) database.RBACListAccess { + principal, ok := authctx.PrincipalFromContext(ctx) + if !ok { + return database.RBACListAccess{Scope: database.RBACScopeAssigned} + } + return database.RBACListAccess{UserID: principal.UserID, Scope: principal.ScopeFor("c2:read")} +} + // ============================================================================ // c2_profile — Malleable Profile 管理工具(新增) // ============================================================================ diff --git a/internal/app/cors_security_test.go b/internal/app/cors_security_test.go new file mode 100644 index 00000000..1860dea3 --- /dev/null +++ b/internal/app/cors_security_test.go @@ -0,0 +1,34 @@ +package app + +import ( + "net/http" + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" +) + +func TestCORSMiddlewareAllowsSameOriginAndRejectsForeignOrigin(t *testing.T) { + gin.SetMode(gin.TestMode) + router := gin.New() + router.Use(corsMiddleware()) + router.GET("/test", func(c *gin.Context) { c.Status(http.StatusNoContent) }) + + same := httptest.NewRequest(http.MethodGet, "http://app.example/test", nil) + same.Host = "app.example" + same.Header.Set("Origin", "http://app.example") + sameW := httptest.NewRecorder() + router.ServeHTTP(sameW, same) + if sameW.Code != http.StatusNoContent || sameW.Header().Get("Access-Control-Allow-Origin") != "http://app.example" { + t.Fatalf("same-origin response = %d, allow-origin=%q", sameW.Code, sameW.Header().Get("Access-Control-Allow-Origin")) + } + + foreign := httptest.NewRequest(http.MethodGet, "http://app.example/test", nil) + foreign.Host = "app.example" + foreign.Header.Set("Origin", "https://evil.example") + foreignW := httptest.NewRecorder() + router.ServeHTTP(foreignW, foreign) + if foreignW.Code != http.StatusForbidden { + t.Fatalf("foreign-origin response = %d, want %d", foreignW.Code, http.StatusForbidden) + } +} diff --git a/internal/app/mcp_authorization.go b/internal/app/mcp_authorization.go new file mode 100644 index 00000000..3d3c815a --- /dev/null +++ b/internal/app/mcp_authorization.go @@ -0,0 +1,249 @@ +package app + +import ( + "context" + "fmt" + "strings" + + "cyberstrike-ai/internal/agent" + "cyberstrike-ai/internal/authctx" + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/mcp" + "cyberstrike-ai/internal/mcp/builtin" +) + +func mcpToolAuthorizer(db *database.DB) func(context.Context, string, map[string]interface{}) error { + return func(ctx context.Context, toolName string, args map[string]interface{}) error { + principal, ok := authctx.PrincipalFromContext(ctx) + if !ok { + return fmt.Errorf("missing authenticated principal") + } + require := func(permission string) error { + if !principal.HasPermission(permission) { + return fmt.Errorf("missing permission %s", permission) + } + return nil + } + resource := func(permission, resourceType, argument string) error { + if err := require(permission); err != nil { + return err + } + id := mcpAuthorizationString(args, argument) + if id == "" || db == nil || !db.UserCanAccessResource(principal.UserID, principal.ScopeFor(permission), resourceType, id) { + return fmt.Errorf("no access to %s %s", resourceType, id) + } + return nil + } + + switch toolName { + case builtin.ToolWebshellExec, builtin.ToolWebshellFileWrite: + return resource("webshell:write", "webshell", "connection_id") + case builtin.ToolWebshellFileList, builtin.ToolWebshellFileRead: + return resource("webshell:read", "webshell", "connection_id") + case builtin.ToolManageWebshellList: + return require("webshell:read") + case builtin.ToolManageWebshellAdd: + return require("webshell:write") + case builtin.ToolManageWebshellUpdate, builtin.ToolManageWebshellTest: + return resource("webshell:write", "webshell", "connection_id") + case builtin.ToolManageWebshellDelete: + return resource("webshell:delete", "webshell", "connection_id") + case builtin.ToolRecordVulnerability: + if err := require("vulnerability:write"); err != nil { + return err + } + conversationID := mcpAuthorizationString(args, "conversation_id") + if conversationID == "" { + conversationID = mcpAuthorizationConversationID(ctx) + } + if conversationID == "" || db == nil || !db.UserCanAccessResource(principal.UserID, principal.ScopeFor("vulnerability:write"), "conversation", conversationID) { + return fmt.Errorf("no access to conversation %s", conversationID) + } + return nil + case builtin.ToolListVulnerabilities: + if err := require("vulnerability:read"); err != nil { + return err + } + conversationID := mcpAuthorizationConversationID(ctx) + if conversationID == "" || db == nil || !db.UserCanAccessResource(principal.UserID, principal.ScopeFor("vulnerability:read"), "conversation", conversationID) { + return fmt.Errorf("no access to conversation %s", conversationID) + } + return nil + case builtin.ToolGetVulnerability: + return resource("vulnerability:read", "vulnerability", "id") + case builtin.ToolUpsertProjectFact, builtin.ToolDeprecateProjectFact, builtin.ToolRestoreProjectFact: + return authorizeProjectTool(ctx, principal, db, "project:write") + case builtin.ToolGetProjectFact, builtin.ToolListProjectFacts, builtin.ToolSearchProjectFacts: + return authorizeProjectTool(ctx, principal, db, "project:read") + case builtin.ToolListKnowledgeRiskTypes, builtin.ToolSearchKnowledgeBase: + return require("knowledge:read") + case builtin.ToolAnalyzeImage: + return require("agent:execute") + case builtin.ToolBatchTaskList: + return require("tasks:read") + case builtin.ToolBatchTaskGet: + return resource("tasks:read", "batch_task", "queue_id") + case builtin.ToolBatchTaskCreate: + if err := require("tasks:write"); err != nil { + return err + } + if projectID := mcpAuthorizationString(args, "project_id"); projectID != "" && (db == nil || !db.UserCanAccessResource(principal.UserID, principal.ScopeFor("tasks:write"), "project", projectID)) { + return fmt.Errorf("no access to project %s", projectID) + } + return nil + case builtin.ToolBatchTaskDelete, builtin.ToolBatchTaskRemove: + return resource("tasks:delete", "batch_task", "queue_id") + case builtin.ToolBatchTaskStart, builtin.ToolBatchTaskRerun, builtin.ToolBatchTaskPause, + builtin.ToolBatchTaskUpdateMetadata, builtin.ToolBatchTaskUpdateSchedule, + builtin.ToolBatchTaskScheduleEnabled, builtin.ToolBatchTaskAdd, builtin.ToolBatchTaskUpdate: + return resource("tasks:write", "batch_task", "queue_id") + case builtin.ToolC2Listener: + return authorizeC2Action(principal, db, args, "c2_listener", "listener_id") + case builtin.ToolC2Session, builtin.ToolC2Task, builtin.ToolC2File: + if toolName == builtin.ToolC2File && mcpAuthorizationString(args, "action") == "get_result" { + return authorizeC2Action(principal, db, args, "c2_task", "task_id") + } + return authorizeC2Action(principal, db, args, "c2_session", "session_id") + case builtin.ToolC2TaskManage: + return authorizeC2Action(principal, db, args, "c2_task", "task_id") + case builtin.ToolC2Payload: + return resource("c2:write", "c2_listener", "listener_id") + case builtin.ToolC2Event: + if id := mcpAuthorizationString(args, "session_id"); id != "" { + return resource("c2:read", "c2_session", "session_id") + } + if principal.ScopeFor("c2:read") != database.RBACScopeAll { + return fmt.Errorf("unfiltered C2 event list requires global scope") + } + return require("c2:read") + case builtin.ToolC2Profile: + // Profiles are process-global and do not yet have an owner. Writes are + // therefore reserved for global scope; reads require c2:read. + if mcpAuthorizationString(args, "action") == "list" || mcpAuthorizationString(args, "action") == "get" { + return require("c2:read") + } + permission := "c2:write" + if mcpAuthorizationString(args, "action") == "delete" { + permission = "c2:delete" + } + if principal.ScopeFor(permission) != database.RBACScopeAll { + return fmt.Errorf("C2 profile mutation requires global scope") + } + if mcpAuthorizationString(args, "action") == "delete" { + return require("c2:delete") + } + return require("c2:write") + default: + if builtin.IsBuiltinTool(toolName) { + return fmt.Errorf("no authorization policy registered for builtin tool %s", toolName) + } + if principal.HasPermission("agent:local-execute") { + return nil + } + return fmt.Errorf("missing agent:local-execute") + } + } +} + +func externalMCPToolAuthorizer() func(context.Context, string, map[string]interface{}) error { + return func(ctx context.Context, toolName string, _ map[string]interface{}) error { + principal, ok := authctx.PrincipalFromContext(ctx) + if !ok { + return fmt.Errorf("missing authenticated principal") + } + if !principal.HasPermission("mcp:external:execute") { + return fmt.Errorf("missing permission mcp:external:execute") + } + if principal.ScopeFor("mcp:external:execute") != database.RBACScopeAll { + return fmt.Errorf("external MCP invocation requires global scope") + } + if strings.TrimSpace(toolName) == "" { + return fmt.Errorf("missing external tool name") + } + return nil + } +} + +func authorizeC2Action(principal authctx.Principal, db *database.DB, args map[string]interface{}, resourceType, argument string) error { + action := mcpAuthorizationString(args, "action") + permission := "c2:write" + if action == "list" || action == "get" || action == "get_result" || action == "wait" { + permission = "c2:read" + } else if action == "delete" || action == "delete_batch" { + permission = "c2:delete" + } + if !principal.HasPermission(permission) { + return fmt.Errorf("missing permission %s", permission) + } + id := mcpAuthorizationString(args, argument) + if action == "delete_batch" { + ids := mcpAuthorizationStrings(args, argument+"s") + if len(ids) == 0 { + return fmt.Errorf("missing resource identifiers %ss", argument) + } + for _, candidate := range ids { + if db == nil || !db.UserCanAccessResource(principal.UserID, principal.ScopeFor(permission), resourceType, candidate) { + return fmt.Errorf("no access to %s %s", resourceType, candidate) + } + } + return nil + } + if id == "" { + if action == "create" || action == "list" { + return nil + } + return fmt.Errorf("missing resource identifier %s", argument) + } + if db == nil || !db.UserCanAccessResource(principal.UserID, principal.ScopeFor(permission), resourceType, id) { + return fmt.Errorf("no access to %s %s", resourceType, id) + } + return nil +} + +func mcpAuthorizationStrings(args map[string]interface{}, key string) []string { + values := []string{} + switch raw := args[key].(type) { + case []string: + for _, value := range raw { + if value = strings.TrimSpace(value); value != "" { + values = append(values, value) + } + } + case []interface{}: + for _, item := range raw { + if value, ok := item.(string); ok { + if value = strings.TrimSpace(value); value != "" { + values = append(values, value) + } + } + } + } + return values +} + +func authorizeProjectTool(ctx context.Context, principal authctx.Principal, db *database.DB, permission string) error { + if !principal.HasPermission(permission) { + return fmt.Errorf("missing permission %s", permission) + } + conversationID := mcpAuthorizationConversationID(ctx) + if conversationID == "" || db == nil || !db.UserCanAccessResource(principal.UserID, principal.ScopeFor(permission), "conversation", conversationID) { + return fmt.Errorf("no access to conversation %s", conversationID) + } + projectID, err := db.GetConversationProjectID(conversationID) + if err != nil || strings.TrimSpace(projectID) == "" || !db.UserCanAccessResource(principal.UserID, principal.ScopeFor(permission), "project", projectID) { + return fmt.Errorf("no access to project %s", projectID) + } + return nil +} + +func mcpAuthorizationConversationID(ctx context.Context) string { + if id := strings.TrimSpace(agent.ConversationIDFromContext(ctx)); id != "" { + return id + } + return strings.TrimSpace(mcp.MCPConversationIDFromContext(ctx)) +} + +func mcpAuthorizationString(args map[string]interface{}, key string) string { + value, _ := args[key].(string) + return strings.TrimSpace(value) +} diff --git a/internal/app/mcp_authorization_test.go b/internal/app/mcp_authorization_test.go new file mode 100644 index 00000000..0324ca26 --- /dev/null +++ b/internal/app/mcp_authorization_test.go @@ -0,0 +1,97 @@ +package app + +import ( + "context" + "path/filepath" + "strings" + "testing" + "time" + + "cyberstrike-ai/internal/authctx" + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/mcp/builtin" + "cyberstrike-ai/internal/security" + + "go.uber.org/zap" +) + +func TestMCPToolAuthorizerEnforcesPermissionAndResource(t *testing.T) { + db, err := database.NewDB(filepath.Join(t.TempDir(), "mcp-authz.db"), zap.NewNop()) + if err != nil { + t.Fatal(err) + } + defer db.Close() + user, err := db.CreateRBACUser("mcp-user", "MCP User", "hash", true, nil) + if err != nil { + t.Fatal(err) + } + for _, id := range []string{"ws_allowed", "ws_hidden"} { + if err := db.CreateWebshellConnection(&database.WebShellConnection{ID: id, URL: "http://127.0.0.1/" + id, Type: "php", Method: "post", CmdParam: "cmd", CreatedAt: time.Now()}); err != nil { + t.Fatal(err) + } + } + if err := db.AssignResourceToUser(user.ID, "webshell", "ws_allowed"); err != nil { + t.Fatal(err) + } + + principal := authctx.NewPrincipal(user.ID, user.Username, database.RBACScopeAssigned, map[string]bool{"mcp:write": true, "webshell:write": true}) + ctx := authctx.WithPrincipal(context.Background(), principal) + authorize := mcpToolAuthorizer(db) + if err := authorize(ctx, builtin.ToolWebshellExec, map[string]interface{}{"connection_id": "ws_allowed"}); err != nil { + t.Fatalf("allowed resource denied: %v", err) + } + if err := authorize(ctx, builtin.ToolWebshellExec, map[string]interface{}{"connection_id": "ws_hidden"}); err == nil { + t.Fatal("foreign webshell resource was allowed") + } + if err := authorize(ctx, builtin.ToolManageWebshellDelete, map[string]interface{}{"connection_id": "ws_allowed"}); err == nil { + t.Fatal("delete without webshell:delete was allowed") + } +} + +func TestEveryBuiltinMCPToolHasExplicitAuthorizationPolicy(t *testing.T) { + db, err := database.NewDB(filepath.Join(t.TempDir(), "mcp-policy-inventory.db"), zap.NewNop()) + if err != nil { + t.Fatal(err) + } + defer db.Close() + permissions := map[string]bool{} + for permission := range security.PermissionCatalog { + permissions[permission] = true + } + ctx := authctx.WithPrincipal(context.Background(), authctx.NewPrincipal("admin", "admin", database.RBACScopeAll, permissions)) + authorize := mcpToolAuthorizer(db) + args := map[string]interface{}{ + "action": "get", "connection_id": "x", "queue_id": "x", "listener_id": "x", + "session_id": "x", "task_id": "x", "id": "x", "conversation_id": "x", + } + for _, toolName := range builtin.GetAllBuiltinTools() { + err := authorize(ctx, toolName, args) + if err != nil && strings.Contains(err.Error(), "no authorization policy registered") { + t.Errorf("builtin tool %s has no explicit policy", toolName) + } + } +} + +func TestExternalMCPRequiresDedicatedPermission(t *testing.T) { + authorize := externalMCPToolAuthorizer() + ctx := authctx.WithPrincipal(context.Background(), authctx.NewPrincipal("u1", "user", database.RBACScopeAssigned, map[string]bool{"agent:execute": true})) + if err := authorize(ctx, "server::tool", nil); err == nil { + t.Fatal("agent:execute alone authorized an external MCP tool") + } + ctx = authctx.WithPrincipal(context.Background(), authctx.NewPrincipal("u1", "user", database.RBACScopeAll, map[string]bool{"mcp:external:execute": true})) + if err := authorize(ctx, "server::tool", nil); err != nil { + t.Fatalf("dedicated external MCP permission rejected: %v", err) + } +} + +func TestConfiguredCommandToolRequiresLocalExecutePermission(t *testing.T) { + authorize := mcpToolAuthorizer(nil) + agentOnly := authctx.WithPrincipal(context.Background(), authctx.NewPrincipal("u1", "user", database.RBACScopeAssigned, map[string]bool{"agent:execute": true})) + if err := authorize(agentOnly, "nmap_scan", nil); err == nil { + t.Fatal("agent:execute alone authorized a configured command tool") + } + local := authctx.WithPrincipal(context.Background(), authctx.NewPrincipal("u1", "user", database.RBACScopeAssigned, map[string]bool{"agent:local-execute": true})) + if err := authorize(local, "nmap_scan", nil); err != nil { + t.Fatalf("agent:local-execute rejected: %v", err) + } +} diff --git a/internal/app/mcp_http_auth_test.go b/internal/app/mcp_http_auth_test.go new file mode 100644 index 00000000..d741b342 --- /dev/null +++ b/internal/app/mcp_http_auth_test.go @@ -0,0 +1,55 @@ +package app + +import ( + "bytes" + "net/http" + "net/http/httptest" + "path/filepath" + "testing" + + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/mcp" + "cyberstrike-ai/internal/security" + + "go.uber.org/zap" +) + +func TestStandaloneMCPPrefersUserRBACAndDisablesGlobalTokenByDefault(t *testing.T) { + db, err := database.NewDB(filepath.Join(t.TempDir(), "mcp-http-auth.db"), zap.NewNop()) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = db.Close() }) + auth, err := security.NewAuthManager("admin-secret", 12) + if err != nil { + t.Fatal(err) + } + if err := auth.AttachRBACStore(db); err != nil { + t.Fatal(err) + } + token, _, err := auth.Authenticate("admin", "admin-secret") + if err != nil { + t.Fatal(err) + } + server := mcp.NewServer(zap.NewNop()) + server.SetToolAuthorizer(mcpToolAuthorizer(db)) + a := &App{config: &config.Config{MCP: config.MCPConfig{AuthHeader: "X-MCP-Token", AuthHeaderValue: "static-secret"}}, auth: auth, mcpServer: server} + body := []byte(`{"jsonrpc":"2.0","id":1,"method":"tools/list"}`) + + userReq := httptest.NewRequest(http.MethodPost, "/mcp", bytes.NewReader(body)) + userReq.Header.Set("Authorization", "Bearer "+token) + userW := httptest.NewRecorder() + a.mcpHandlerWithAuth(userW, userReq) + if userW.Code != http.StatusOK { + t.Fatalf("user bearer status = %d: %s", userW.Code, userW.Body.String()) + } + + staticReq := httptest.NewRequest(http.MethodPost, "/mcp", bytes.NewReader(body)) + staticReq.Header.Set("X-MCP-Token", "static-secret") + staticW := httptest.NewRecorder() + a.mcpHandlerWithAuth(staticW, staticReq) + if staticW.Code != http.StatusUnauthorized { + t.Fatalf("global static token status = %d, want 401", staticW.Code) + } +} diff --git a/internal/app/vulnerability_tools.go b/internal/app/vulnerability_tools.go index 6fdb917f..dd9c097a 100644 --- a/internal/app/vulnerability_tools.go +++ b/internal/app/vulnerability_tools.go @@ -6,6 +6,7 @@ import ( "strings" "cyberstrike-ai/internal/agent" + "cyberstrike-ai/internal/authctx" "cyberstrike-ai/internal/database" "cyberstrike-ai/internal/mcp" "cyberstrike-ai/internal/mcp/builtin" @@ -313,6 +314,10 @@ func registerRecordVulnerabilityTool(mcpServer *mcp.Server, db *database.DB, log } return textResult(fmt.Sprintf("记录漏洞失败: %v", err), true), nil } + if principal, ok := authctx.PrincipalFromContext(ctx); ok { + _ = db.SetResourceOwner("vulnerability", created.ID, principal.UserID) + _ = db.AssignResourceToUser(principal.UserID, "vulnerability", created.ID) + } if logger != nil { logger.Info("漏洞记录成功", diff --git a/internal/authctx/principal.go b/internal/authctx/principal.go new file mode 100644 index 00000000..f74199ac --- /dev/null +++ b/internal/authctx/principal.go @@ -0,0 +1,72 @@ +package authctx + +import ( + "context" + "strings" +) + +// Principal is the immutable authorization identity propagated beyond the +// transport layer into Agent, MCP and background task contexts. +type Principal struct { + UserID string + Username string + Permissions map[string]bool + PermissionScopes map[string]string + Scope string +} + +type principalContextKey struct{} + +func NewPrincipal(userID, username, scope string, permissions map[string]bool) Principal { + return NewPrincipalWithScopes(userID, username, scope, permissions, nil) +} + +func NewPrincipalWithScopes(userID, username, scope string, permissions map[string]bool, permissionScopes map[string]string) Principal { + permissionCopy := make(map[string]bool, len(permissions)) + scopeCopy := make(map[string]string, len(permissionScopes)) + for permission, allowed := range permissions { + if allowed { + permissionCopy[permission] = true + if permissionScope := strings.TrimSpace(permissionScopes[permission]); permissionScope != "" { + scopeCopy[permission] = permissionScope + } + } + } + return Principal{ + UserID: strings.TrimSpace(userID), Username: strings.TrimSpace(username), + Scope: strings.TrimSpace(scope), Permissions: permissionCopy, PermissionScopes: scopeCopy, + } +} + +func WithPrincipal(ctx context.Context, principal Principal) context.Context { + if ctx == nil { + ctx = context.Background() + } + if strings.TrimSpace(principal.UserID) == "" { + return ctx + } + return context.WithValue(ctx, principalContextKey{}, principal) +} + +func PrincipalFromContext(ctx context.Context) (Principal, bool) { + if ctx == nil { + return Principal{}, false + } + principal, ok := ctx.Value(principalContextKey{}).(Principal) + return principal, ok && strings.TrimSpace(principal.UserID) != "" +} + +func (p Principal) HasPermission(permission string) bool { + return p.Permissions[strings.TrimSpace(permission)] +} + +// ScopeFor returns the scope attached to the permission that authorizes the +// current action. Falling back to Scope keeps explicit service principals and +// legacy callers compatible without reintroducing cross-role scope widening. +func (p Principal) ScopeFor(permission string) string { + permission = strings.TrimSpace(permission) + if scope := strings.TrimSpace(p.PermissionScopes[permission]); scope != "" { + return scope + } + return strings.TrimSpace(p.Scope) +} diff --git a/internal/mcp/external_manager.go b/internal/mcp/external_manager.go index cf2ff30b..9f69ade0 100644 --- a/internal/mcp/external_manager.go +++ b/internal/mcp/external_manager.go @@ -8,6 +8,7 @@ import ( "sync/atomic" "time" + "cyberstrike-ai/internal/authctx" "cyberstrike-ai/internal/config" "github.com/google/uuid" @@ -37,29 +38,30 @@ type listToolsInflight struct { // ExternalMCPManager 外部MCP管理器 type ExternalMCPManager struct { - clients map[string]ExternalMCPClient - configs map[string]config.ExternalMCPServerConfig - logger *zap.Logger - storage MonitorStorage // 可选的持久化存储 - executions map[string]*ToolExecution // 执行记录 - stats map[string]*ToolStats // 工具统计信息 - errors map[string]string // 错误信息 - toolCounts map[string]int // 工具数量缓存 - toolCountsMu sync.RWMutex // 工具数量缓存的锁 - toolCache map[string]toolListCacheEntry // 工具列表缓存:MCP名称 -> 工具列表 - toolCacheMu sync.RWMutex // 工具列表缓存的锁 - listToolsMu sync.Mutex + clients map[string]ExternalMCPClient + configs map[string]config.ExternalMCPServerConfig + logger *zap.Logger + storage MonitorStorage // 可选的持久化存储 + executions map[string]*ToolExecution // 执行记录 + stats map[string]*ToolStats // 工具统计信息 + errors map[string]string // 错误信息 + toolCounts map[string]int // 工具数量缓存 + toolCountsMu sync.RWMutex // 工具数量缓存的锁 + toolCache map[string]toolListCacheEntry // 工具列表缓存:MCP名称 -> 工具列表 + toolCacheMu sync.RWMutex // 工具列表缓存的锁 + listToolsMu sync.Mutex listToolsInflight map[string]*listToolsInflight - stopRefresh chan struct{} // 停止后台刷新的信号 - refreshWg sync.WaitGroup // 等待后台刷新goroutine完成 - refreshing atomic.Bool // 防止 refreshToolCounts 并发堆积 - mu sync.RWMutex + stopRefresh chan struct{} // 停止后台刷新的信号 + refreshWg sync.WaitGroup // 等待后台刷新goroutine完成 + refreshing atomic.Bool // 防止 refreshToolCounts 并发堆积 + mu sync.RWMutex runningCancels map[string]context.CancelFunc abortUserNotes map[string]string reconnectMu sync.Mutex reconnecting map[string]bool reconnectLastTry map[string]time.Time reconnectAttempts map[string]int + toolAuthorizer func(context.Context, string, map[string]interface{}) error } // NewExternalMCPManager 创建外部MCP管理器 @@ -67,16 +69,24 @@ func NewExternalMCPManager(logger *zap.Logger) *ExternalMCPManager { return NewExternalMCPManagerWithStorage(logger, nil) } +// SetToolAuthorizer installs the policy decision point for all external MCP +// invocations. App wiring configures this before any Agent can call a tool. +func (m *ExternalMCPManager) SetToolAuthorizer(authorizer func(context.Context, string, map[string]interface{}) error) { + m.mu.Lock() + m.toolAuthorizer = authorizer + m.mu.Unlock() +} + // NewExternalMCPManagerWithStorage 创建外部MCP管理器(带持久化存储) func NewExternalMCPManagerWithStorage(logger *zap.Logger, storage MonitorStorage) *ExternalMCPManager { manager := &ExternalMCPManager{ - clients: make(map[string]ExternalMCPClient), - configs: make(map[string]config.ExternalMCPServerConfig), - logger: logger, - storage: storage, - executions: make(map[string]*ToolExecution), - stats: make(map[string]*ToolStats), - errors: make(map[string]string), + clients: make(map[string]ExternalMCPClient), + configs: make(map[string]config.ExternalMCPServerConfig), + logger: logger, + storage: storage, + executions: make(map[string]*ToolExecution), + stats: make(map[string]*ToolStats), + errors: make(map[string]string), toolCounts: make(map[string]int), toolCache: make(map[string]toolListCacheEntry), listToolsInflight: make(map[string]*listToolsInflight), @@ -554,6 +564,17 @@ 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 { @@ -591,6 +612,10 @@ func (m *ExternalMCPManager) CallTool(ctx context.Context, toolName string, args Status: "running", StartTime: time.Now(), } + if principal, ok := authctx.PrincipalFromContext(ctx); ok { + execution.OwnerUserID = principal.UserID + } + execution.ConversationID = MCPConversationIDFromContext(ctx) m.mu.Lock() m.executions[executionID] = execution diff --git a/internal/mcp/external_manager_test.go b/internal/mcp/external_manager_test.go index c7260f1d..2ccf722b 100644 --- a/internal/mcp/external_manager_test.go +++ b/internal/mcp/external_manager_test.go @@ -2,14 +2,30 @@ package mcp import ( "context" + "errors" + "strings" "testing" "time" + "cyberstrike-ai/internal/authctx" "cyberstrike-ai/internal/config" "go.uber.org/zap" ) +func TestExternalManagerEnforcesConfiguredAuthorizer(t *testing.T) { + manager := NewExternalMCPManager(zap.NewNop()) + t.Cleanup(manager.StopAll) + manager.SetToolAuthorizer(func(context.Context, string, map[string]interface{}) error { + 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{}{}) + if err == nil || !strings.Contains(err.Error(), "authorization denied") { + t.Fatalf("external call bypassed authorizer: %v", err) + } +} + func TestExternalMCPManager_AddOrUpdateConfig(t *testing.T) { logger := zap.NewNop() manager := NewExternalMCPManager(logger) diff --git a/internal/mcp/server.go b/internal/mcp/server.go index 69d67bb8..6d97d999 100644 --- a/internal/mcp/server.go +++ b/internal/mcp/server.go @@ -14,6 +14,8 @@ import ( "sync" "time" + "cyberstrike-ai/internal/authctx" + "github.com/google/uuid" "go.uber.org/zap" ) @@ -49,6 +51,18 @@ type Server struct { // nil 表示未配置,沿用默认 30 分钟;指向 0 表示不限制;>0 为分钟数。 httpToolTimeoutMinutes *int httpToolTimeoutMu sync.RWMutex + toolAuthorizer func(context.Context, string, map[string]interface{}) error +} + +// SetToolAuthorizer installs the common policy decision point for every +// user-attributed tool call, whether it originates from HTTP or an Agent. +func (s *Server) SetToolAuthorizer(authorizer func(context.Context, string, map[string]interface{}) error) { + if s == nil { + return + } + s.mu.Lock() + s.toolAuthorizer = authorizer + s.mu.Unlock() } type sseClient struct { @@ -111,21 +125,24 @@ func (s *Server) ConfigureHTTPToolCallTimeoutFromAgentMinutes(minutes int) { s.httpToolTimeoutMinutes = &v } -func (s *Server) effectiveHTTPToolCallDeadline() (context.Context, context.CancelFunc) { +func (s *Server) effectiveHTTPToolCallDeadline(parent context.Context) (context.Context, context.CancelFunc) { const defaultDur = 30 * time.Minute + if parent == nil { + parent = context.Background() + } if s == nil { - return context.WithTimeout(context.Background(), defaultDur) + return context.WithTimeout(parent, defaultDur) } s.httpToolTimeoutMu.RLock() mPtr := s.httpToolTimeoutMinutes s.httpToolTimeoutMu.RUnlock() if mPtr == nil { - return context.WithTimeout(context.Background(), defaultDur) + return context.WithTimeout(parent, defaultDur) } if *mPtr <= 0 { - return context.WithCancel(context.Background()) + return context.WithCancel(parent) } - return context.WithTimeout(context.Background(), time.Duration(*mPtr)*time.Minute) + return context.WithTimeout(parent, time.Duration(*mPtr)*time.Minute) } // RegisterTool 注册工具 @@ -196,7 +213,7 @@ func (s *Server) HandleHTTP(w http.ResponseWriter, r *http.Request) { return } - response := s.handleMessage(&msg) + response := s.handleMessage(r.Context(), &msg) w.Header().Set("Content-Type", "application/json") json.NewEncoder(w).Encode(response) } @@ -223,7 +240,7 @@ func (s *Server) serveSSESessionMessage(w http.ResponseWriter, r *http.Request, return } - response := s.handleMessage(&msg) + response := s.handleMessage(r.Context(), &msg) if response == nil { w.WriteHeader(http.StatusAccepted) return @@ -317,7 +334,7 @@ func (s *Server) removeSSEClient(id string) { } // handleMessage 处理MCP消息 -func (s *Server) handleMessage(msg *Message) *Message { +func (s *Server) handleMessage(ctx context.Context, msg *Message) *Message { // 检查是否是通知(notification)- 通知没有id字段,不需要响应 isNotification := msg.ID.Value() == nil || msg.ID.String() == "" @@ -332,7 +349,7 @@ func (s *Server) handleMessage(msg *Message) *Message { case "tools/list": return s.handleListTools(msg) case "tools/call": - return s.handleCallTool(msg) + return s.handleCallTool(ctx, msg) case "prompts/list": return s.handleListPrompts(msg) case "prompts/get": @@ -433,7 +450,7 @@ func (s *Server) handleListTools(msg *Message) *Message { } // handleCallTool 处理工具调用请求 -func (s *Server) handleCallTool(msg *Message) *Message { +func (s *Server) handleCallTool(requestCtx context.Context, msg *Message) *Message { var req CallToolRequest if err := json.Unmarshal(msg.Params, &req); err != nil { return &Message{ @@ -443,6 +460,17 @@ func (s *Server) handleCallTool(msg *Message) *Message { Error: &Error{Code: -32602, Message: "Invalid params"}, } } + _, authenticated := authctx.PrincipalFromContext(requestCtx) + s.mu.RLock() + authorizer := s.toolAuthorizer + s.mu.RUnlock() + if authorizer != nil { + if err := authorizer(requestCtx, req.Name, req.Arguments); err != nil { + return &Message{ID: msg.ID, Type: MessageTypeError, Version: "2.0", Error: &Error{Code: -32003, Message: "Forbidden", Data: err.Error()}} + } + } else if authenticated { + return &Message{ID: msg.ID, Type: MessageTypeError, Version: "2.0", Error: &Error{Code: -32003, Message: "Tool authorization policy is not configured"}} + } executionID := uuid.New().String() execution := &ToolExecution{ @@ -452,6 +480,10 @@ func (s *Server) handleCallTool(msg *Message) *Message { Status: "running", StartTime: time.Now(), } + if principal, ok := authctx.PrincipalFromContext(requestCtx); ok { + execution.OwnerUserID = principal.UserID + } + execution.ConversationID = MCPConversationIDFromContext(requestCtx) s.mu.Lock() s.executions[executionID] = execution @@ -495,7 +527,7 @@ func (s *Server) handleCallTool(msg *Message) *Message { } } - baseCtx, timeoutCancel := s.effectiveHTTPToolCallDeadline() + baseCtx, timeoutCancel := s.effectiveHTTPToolCallDeadline(requestCtx) defer timeoutCancel() execCtx, runCancel := context.WithCancel(baseCtx) s.registerRunningCancel(executionID, runCancel) @@ -809,6 +841,17 @@ 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() @@ -826,6 +869,10 @@ func (s *Server) CallTool(ctx context.Context, toolName string, args map[string] Status: "running", StartTime: time.Now(), } + if principal, ok := authctx.PrincipalFromContext(ctx); ok { + execution.OwnerUserID = principal.UserID + } + execution.ConversationID = MCPConversationIDFromContext(ctx) s.mu.Lock() s.executions[executionID] = execution @@ -922,7 +969,7 @@ func (s *Server) CallTool(ctx context.Context, toolName string, args map[string] } // BeginToolExecution 创建 running 状态的执行记录,供 Eino 等非 CallTool 路径在工具开始时落库。 -func (s *Server) BeginToolExecution(toolName string, args map[string]interface{}) string { +func (s *Server) BeginToolExecution(ctx context.Context, toolName string, args map[string]interface{}) string { if s == nil { return "" } @@ -937,6 +984,10 @@ func (s *Server) BeginToolExecution(toolName string, args map[string]interface{} Status: "running", StartTime: time.Now(), } + if principal, ok := authctx.PrincipalFromContext(ctx); ok { + execution.OwnerUserID = principal.UserID + } + execution.ConversationID = MCPConversationIDFromContext(ctx) s.mu.Lock() s.executions[executionID] = execution @@ -952,7 +1003,7 @@ func (s *Server) BeginToolExecution(toolName string, args map[string]interface{} } // FinishToolExecution 完成先前 BeginToolExecution 创建的记录;executionID 为空时等同 RecordCompletedToolInvocation。 -func (s *Server) FinishToolExecution(executionID, toolName string, args map[string]interface{}, resultText string, invokeErr error) string { +func (s *Server) FinishToolExecution(ctx context.Context, executionID, toolName string, args map[string]interface{}, resultText string, invokeErr error) string { if s == nil { return "" } @@ -961,7 +1012,7 @@ func (s *Server) FinishToolExecution(executionID, toolName string, args map[stri } id := strings.TrimSpace(executionID) if id == "" { - return s.RecordCompletedToolInvocation(toolName, args, resultText, invokeErr) + id = uuid.New().String() } now := time.Now() @@ -984,6 +1035,12 @@ func (s *Server) FinishToolExecution(executionID, toolName string, args map[stri if len(args) > 0 { exec.Arguments = args } + if principal, ok := authctx.PrincipalFromContext(ctx); ok { + exec.OwnerUserID = principal.UserID + } + if conversationID := MCPConversationIDFromContext(ctx); conversationID != "" { + exec.ConversationID = conversationID + } exec.EndTime = &now if exec.StartTime.IsZero() { exec.StartTime = now @@ -1027,8 +1084,8 @@ func (s *Server) FinishToolExecution(executionID, toolName string, args map[stri // RecordCompletedToolInvocation 将已在其它路径完成的工具调用写入监控存储(格式与 CallTool 结束后一致), // 用于 Eino ADK filesystem execute 等未经过 CallTool 的场景;返回 executionId 供助手消息 mcpExecutionIds 关联。 -func (s *Server) RecordCompletedToolInvocation(toolName string, args map[string]interface{}, resultText string, invokeErr error) string { - return s.FinishToolExecution("", toolName, args, resultText, invokeErr) +func (s *Server) RecordCompletedToolInvocation(ctx context.Context, toolName string, args map[string]interface{}, resultText string, invokeErr error) string { + return s.FinishToolExecution(ctx, "", toolName, args, resultText, invokeErr) } // UpdateToolExecutionResult 将监控库中的工具结果更新为送入模型的展示正文(如 reduction 后的 persisted-output)。 @@ -1519,7 +1576,7 @@ func (s *Server) HandleStdio() error { } // 处理消息 - response := s.handleMessage(&msg) + response := s.handleMessage(context.Background(), &msg) // 如果是通知(response 为 nil),不需要发送响应 if response == nil { diff --git a/internal/mcp/server_authorization_test.go b/internal/mcp/server_authorization_test.go new file mode 100644 index 00000000..51c44283 --- /dev/null +++ b/internal/mcp/server_authorization_test.go @@ -0,0 +1,36 @@ +package mcp + +import ( + "context" + "errors" + "testing" + + "cyberstrike-ai/internal/authctx" + + "go.uber.org/zap" +) + +func TestToolAuthorizerIsUniversalAndExecutionKeepsOwner(t *testing.T) { + server := NewServer(zap.NewNop()) + server.RegisterTool(Tool{Name: "echo", InputSchema: map[string]interface{}{"type": "object"}}, func(ctx context.Context, args map[string]interface{}) (*ToolResult, error) { + return &ToolResult{Content: []Content{{Type: "text", Text: "ok"}}}, nil + }) + server.SetToolAuthorizer(func(ctx context.Context, toolName string, args map[string]interface{}) error { + if _, ok := authctx.PrincipalFromContext(ctx); !ok { + return errors.New("principal required") + } + return nil + }) + if _, _, err := server.CallTool(context.Background(), "echo", nil); err == nil { + t.Fatal("tool call without principal was allowed") + } + 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 { + t.Fatal(err) + } + execution, ok := server.GetExecution(executionID) + if !ok || execution.OwnerUserID != "u1" { + t.Fatalf("execution owner = %#v, want u1", execution) + } +} diff --git a/internal/mcp/types.go b/internal/mcp/types.go index 03567cf4..d8217286 100644 --- a/internal/mcp/types.go +++ b/internal/mcp/types.go @@ -201,6 +201,7 @@ type ToolExecution struct { Duration time.Duration `json:"duration,omitempty"` // ConversationID 仅 API 展示用(进行中的 Agent 任务),不写入 tool_executions 表。 ConversationID string `json:"conversationId,omitempty"` + OwnerUserID string `json:"-"` } // ToolStats 工具统计信息