From 1b64f5d8a01f527f5fa63817cc998b2e9c19482a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=85=AC=E6=98=8E?= <83812544+Ed1s0nZ@users.noreply.github.com> Date: Fri, 10 Jul 2026 18:54:04 +0800 Subject: [PATCH] Add files via upload --- internal/audit/service.go | 7 +- internal/handler/agent.go | 53 ++++- internal/handler/audit.go | 28 ++- internal/handler/auth.go | 16 +- internal/handler/batch_queue_executor.go | 13 +- internal/handler/batch_task_mcp.go | 20 +- internal/handler/c2.go | 16 ++ internal/handler/chat_uploads.go | 93 +++++++- .../handler/eino_empty_response_continue.go | 4 +- internal/handler/eino_single_agent.go | 11 +- internal/handler/external_mcp.go | 25 ++- internal/handler/group.go | 45 +++- internal/handler/monitor.go | 121 +++++++++- internal/handler/multi_agent.go | 10 +- internal/handler/multi_agent_prepare.go | 36 ++- internal/handler/notification.go | 56 +++-- internal/handler/notification_rbac_test.go | 45 ++++ internal/handler/project.go | 6 + internal/handler/rbac.go | 36 ++- internal/handler/rbac_boundary_test.go | 206 ++++++++++++++++++ internal/handler/rbac_context.go | 12 + internal/handler/robot.go | 149 +++++++++---- internal/handler/robot_rbac_test.go | 64 ++++++ internal/handler/vulnerability.go | 10 +- internal/handler/webshell.go | 27 ++- internal/handler/webshell_rbac_test.go | 12 + internal/handler/workflow_integration.go | 6 +- internal/handler/workflow_run.go | 38 ++++ 28 files changed, 1040 insertions(+), 125 deletions(-) create mode 100644 internal/handler/notification_rbac_test.go create mode 100644 internal/handler/rbac_boundary_test.go create mode 100644 internal/handler/rbac_context.go create mode 100644 internal/handler/robot_rbac_test.go diff --git a/internal/audit/service.go b/internal/audit/service.go index a6cc1203..60ea11b3 100644 --- a/internal/audit/service.go +++ b/internal/audit/service.go @@ -63,7 +63,12 @@ func (s *Service) Record(c *gin.Context, e Entry) { } } if strings.TrimSpace(e.Actor) == "" { - e.Actor = "admin" + if c != nil { + e.Actor = strings.TrimSpace(c.GetString(security.ContextUsernameKey)) + } + if e.Actor == "" { + e.Actor = "admin" + } } maxDetail := s.cfg.Audit.MaxDetailBytesEffective() detail := SanitizeDetail(e.Detail, maxDetail) diff --git a/internal/handler/agent.go b/internal/handler/agent.go index ac2ea714..087144bd 100644 --- a/internal/handler/agent.go +++ b/internal/handler/agent.go @@ -18,6 +18,7 @@ import ( "cyberstrike-ai/internal/agent" "cyberstrike-ai/internal/audit" + "cyberstrike-ai/internal/authctx" "cyberstrike-ai/internal/config" "cyberstrike-ai/internal/database" "cyberstrike-ai/internal/mcp/builtin" @@ -728,7 +729,22 @@ func (h *AgentHandler) runRobotMultiAgentWithRetry( } // ProcessMessageForRobot 供机器人（企业微信/钉钉/飞书）调用：Eino 单/多代理执行路径（含 progressCallback、过程详情），仅不发送 SSE，最后返回完整回复 -func (h *AgentHandler) ProcessMessageForRobot(ctx context.Context, platform, conversationID, message, role string) (response string, convID string, err error) { +func (h *AgentHandler) ProcessMessageForRobot(ctx context.Context, platform, ownerUserID, conversationID, message, role string) (response string, convID string, err error) { + permissions := map[string]bool{ + "agent:execute": true, + "chat:read": true, "chat:write": true, "chat:delete": true, + "project:read": true, "project:write": true, + "knowledge:read": true, + "vulnerability:read": true, "vulnerability:write": true, + "attackchain:read": true, "attackchain:write": true, + "workflow:read": true, + "hitl:read": true, "hitl:write": true, + } + ownerUserID = strings.TrimSpace(ownerUserID) + if ownerUserID == "" { + return "", "", fmt.Errorf("robot owner identity is required") + } + ctx = authctx.WithPrincipal(ctx, authctx.NewPrincipal(ownerUserID, "robot-service", database.RBACScopeOwn, permissions)) if conversationID == "" { title := safeTruncateString(message, 50) src := "robot" @@ -737,13 +753,17 @@ func (h *AgentHandler) ProcessMessageForRobot(ctx context.Context, platform, con } meta := audit.ConversationCreateMeta(src) meta.ProjectID = effectiveProjectID(h.config, "") + if meta.ProjectID != "" && !h.db.UserCanAccessResource(ownerUserID, database.RBACScopeOwn, "project", meta.ProjectID) { + meta.ProjectID = "" + } conv, createErr := h.db.CreateConversation(title, meta) if createErr != nil { return "", "", fmt.Errorf("创建对话失败: %w", createErr) } conversationID = conv.ID + _ = h.db.SetResourceOwner("conversation", conversationID, ownerUserID) } else { - if _, getErr := h.db.GetConversation(conversationID); getErr != nil { + if _, getErr := h.db.GetConversation(conversationID); getErr != nil || !h.db.UserCanAccessResource(ownerUserID, database.RBACScopeOwn, "conversation", conversationID) { return "", "", fmt.Errorf("对话不存在") } } @@ -1489,6 +1509,10 @@ func (h *AgentHandler) CancelAgentLoop(c *gin.Context) { c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) return } + if !h.agentConversationAllowed(c, req.ConversationID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } if req.ContinueAfter { if h.tasks.GetTask(req.ConversationID) == nil { @@ -1570,6 +1594,10 @@ func (h *AgentHandler) SubscribeAgentTaskEvents(c *gin.Context) { c.JSON(http.StatusBadRequest, gin.H{"error": "conversationId is required"}) return } + if !h.agentConversationAllowed(c, conversationID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } if h.tasks.GetTask(conversationID) == nil { c.JSON(http.StatusNotFound, gin.H{"error": "no active task for this conversation"}) return @@ -1648,6 +1676,9 @@ func (h *AgentHandler) enrichCompletedTasksWithConversationTitles(tasks []*Compl // ListAgentTasks 列出所有运行中的任务 func (h *AgentHandler) ListAgentTasks(c *gin.Context) { tasks := h.tasks.GetActiveTasks() + tasks = filterSlice(tasks, func(task *AgentTask) bool { + return task != nil && h.agentConversationAllowed(c, task.ConversationID) + }) h.enrichAgentTasksWithConversationTitles(tasks) c.JSON(http.StatusOK, gin.H{ "tasks": tasks, @@ -1657,12 +1688,30 @@ func (h *AgentHandler) ListAgentTasks(c *gin.Context) { // ListCompletedTasks 列出最近完成的任务历史 func (h *AgentHandler) ListCompletedTasks(c *gin.Context) { tasks := h.tasks.GetCompletedTasks() + tasks = filterSlice(tasks, func(task *CompletedTask) bool { + return task != nil && h.agentConversationAllowed(c, task.ConversationID) + }) h.enrichCompletedTasksWithConversationTitles(tasks) c.JSON(http.StatusOK, gin.H{ "tasks": tasks, }) } +func (h *AgentHandler) agentConversationAllowed(c *gin.Context, conversationID string) bool { + session, ok := security.CurrentSession(c) + return ok && h.db != nil && h.db.UserCanAccessResource(session.UserID, session.Scope, "conversation", strings.TrimSpace(conversationID)) +} + +func filterSlice[T any](items []T, keep func(T) bool) []T { + out := make([]T, 0, len(items)) + for _, item := range items { + if keep(item) { + out = append(out, item) + } + } + return out +} + // BatchTaskRequest 批量任务请求 type BatchTaskRequest struct { Title string `json:"title"` // 任务标题（可选） diff --git a/internal/handler/audit.go b/internal/handler/audit.go index 7cb4dd47..6baba312 100644 --- a/internal/handler/audit.go +++ b/internal/handler/audit.go @@ -6,6 +6,7 @@ import ( "cyberstrike-ai/internal/audit" "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/security" "github.com/gin-gonic/gin" "go.uber.org/zap" @@ -32,8 +33,8 @@ func (h *AuditHandler) Meta(c *gin.Context) { retentionDays = h.audit.RetentionDays() } c.JSON(http.StatusOK, gin.H{ - "enabled": enabled, - "retention_days": retentionDays, + "enabled": enabled, + "retention_days": retentionDays, "default_page_size": 20, "max_page_size": 100, "max_export": 5000, @@ -46,7 +47,7 @@ func (h *AuditHandler) Summary(c *gin.Context) { c.JSON(http.StatusInternalServerError, gin.H{"error": "database unavailable"}) return } - base := auditFilterFromQuery(c) + base := auditFilterForAccess(c, auditFilterFromQuery(c)) total, err := h.db.CountAuditLogs(base) if err != nil { c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) @@ -68,9 +69,9 @@ func (h *AuditHandler) Summary(c *gin.Context) { return } c.JSON(http.StatusOK, gin.H{ - "total": total, - "failures": failures, - "recent_7d": recent7d, + "total": total, + "failures": failures, + "recent_7d": recent7d, "has_filters": c.Query("category") != "" || c.Query("action") != "" || c.Query("result") != "" || c.Query("q") != "" || c.Query("since") != "" || c.Query("until") != "", }) @@ -82,7 +83,7 @@ func (h *AuditHandler) ListLogs(c *gin.Context) { c.JSON(http.StatusInternalServerError, gin.H{"error": "database unavailable"}) return } - filter := auditFilterFromQuery(c) + filter := auditFilterForAccess(c, auditFilterFromQuery(c)) page, pageSize := auditPaginationFromQuery(c) filter.Limit = pageSize filter.Offset = (page - 1) * pageSize @@ -116,6 +117,10 @@ func (h *AuditHandler) GetLog(c *gin.Context) { c.JSON(http.StatusNotFound, gin.H{"error": "审计记录不存在"}) return } + if session, ok := security.CurrentSession(c); !ok || (session.Scope != database.RBACScopeAll && row.Actor != session.Username) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } audit.ApplyResourceAvailability(h.db, row) c.JSON(http.StatusOK, gin.H{"log": row}) } @@ -126,7 +131,7 @@ func (h *AuditHandler) ExportLogs(c *gin.Context) { c.JSON(http.StatusInternalServerError, gin.H{"error": "database unavailable"}) return } - filter := auditFilterFromQuery(c) + filter := auditFilterForAccess(c, auditFilterFromQuery(c)) filter.Limit = 5000 filter.Offset = 0 @@ -145,3 +150,10 @@ func (h *AuditHandler) ExportLogs(c *gin.Context) { "logs": logs, }) } + +func auditFilterForAccess(c *gin.Context, filter database.ListAuditLogsFilter) database.ListAuditLogsFilter { + if session, ok := security.CurrentSession(c); ok && session.Scope != database.RBACScopeAll { + filter.Actor = session.Username + } + return filter +} diff --git a/internal/handler/auth.go b/internal/handler/auth.go index c36e72e2..b2ce9ee9 100644 --- a/internal/handler/auth.go +++ b/internal/handler/auth.go @@ -64,6 +64,7 @@ func (h *AuthHandler) Login(c *gin.Context) { Action: "login", Result: "failure", Message: "登录失败：密码错误", + Actor: strings.TrimSpace(req.Username), }) } c.JSON(http.StatusUnauthorized, gin.H{"error": "密码错误"}) @@ -78,6 +79,7 @@ func (h *AuthHandler) Login(c *gin.Context) { Result: "success", SessionHint: audit.HintFromToken(token), Message: "登录成功", + Actor: session.Username, Detail: map[string]interface{}{ "expires_at": expiresAt.UTC().Format(time.RFC3339), }, @@ -93,9 +95,10 @@ func (h *AuthHandler) Login(c *gin.Context) { "username": session.Username, "display_name": session.DisplayName, }, - "roles": session.Roles, - "permissions": permissionKeys(session.Permissions), - "scope": session.Scope, + "roles": session.Roles, + "permissions": permissionKeys(session.Permissions), + "permission_scopes": session.PermissionScopes, + "scope": session.Scope, }) } @@ -234,9 +237,10 @@ func (h *AuthHandler) Validate(c *gin.Context) { "username": session.Username, "display_name": session.DisplayName, }, - "roles": session.Roles, - "permissions": permissionKeys(session.Permissions), - "scope": session.Scope, + "roles": session.Roles, + "permissions": permissionKeys(session.Permissions), + "permission_scopes": session.PermissionScopes, + "scope": session.Scope, }) } diff --git a/internal/handler/batch_queue_executor.go b/internal/handler/batch_queue_executor.go index 4709d172..7f770e31 100644 --- a/internal/handler/batch_queue_executor.go +++ b/internal/handler/batch_queue_executor.go @@ -11,6 +11,7 @@ import ( "cyberstrike-ai/internal/agent" "cyberstrike-ai/internal/audit" + "cyberstrike-ai/internal/authctx" "cyberstrike-ai/internal/config" "cyberstrike-ai/internal/mcp" "cyberstrike-ai/internal/multiagent" @@ -109,6 +110,13 @@ func (h *AgentHandler) tryFinalizeBatchQueue(queueID string) { // executeOneBatchSubTask 执行单条批量子任务（各自独立会话）。 func (h *AgentHandler) executeOneBatchSubTask(queueID string, queue *BatchTaskQueue, task *BatchTask) { + ownerUserID := h.db.GetResourceOwner("batch_task", queueID) + access, accessErr := h.db.ResolveRBACAccess(ownerUserID) + if accessErr != nil || access == nil || !access.User.Enabled { + h.batchTaskManager.UpdateTaskStatus(queueID, task.ID, BatchTaskStatusFailed, "", "队列所有者不存在或已禁用") + return + } + principal := authctx.NewPrincipalWithScopes(access.User.ID, access.User.Username, access.Scope, access.Permissions, access.PermissionScopes) title := safeTruncateString(task.Message, 50) batchMeta := audit.ConversationCreateMeta("batch_task") batchMeta.ProjectID = effectiveProjectID(h.config, queue.ProjectID) @@ -119,6 +127,8 @@ func (h *AgentHandler) executeOneBatchSubTask(queueID string, queue *BatchTaskQu return } conversationID := conv.ID + _ = h.db.SetResourceOwner("conversation", conversationID, access.User.ID) + _ = h.db.AssignResourceToUser(access.User.ID, "conversation", conversationID) h.batchTaskManager.UpdateTaskStatusWithConversationID(queueID, task.ID, BatchTaskStatusRunning, "", "", conversationID) @@ -156,7 +166,8 @@ func (h *AgentHandler) executeOneBatchSubTask(queueID string, queue *BatchTaskQu h.logger.Info("执行批量任务", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.String("message", task.Message), zap.String("role", queue.Role), zap.String("conversationId", conversationID)) - baseCtx, cancelWithCause := context.WithCancelCause(context.Background()) + principalCtx := authctx.WithPrincipal(context.Background(), principal) + baseCtx, cancelWithCause := context.WithCancelCause(principalCtx) taskCtx, timeoutCancel := context.WithTimeout(baseCtx, 6*time.Hour) registered := false diff --git a/internal/handler/batch_task_mcp.go b/internal/handler/batch_task_mcp.go index e024903b..3c714901 100644 --- a/internal/handler/batch_task_mcp.go +++ b/internal/handler/batch_task_mcp.go @@ -9,7 +9,9 @@ import ( "strings" "time" + "cyberstrike-ai/internal/authctx" "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/database" "cyberstrike-ai/internal/mcp" "cyberstrike-ai/internal/mcp/builtin" @@ -74,7 +76,14 @@ func RegisterBatchTaskMCPTools(mcpServer *mcp.Server, h *AgentHandler, logger *z if offset > 100000 { offset = 100000 } - queues, total, err := h.batchTaskManager.ListQueues(pageSize, offset, status, keyword) + queues := []*BatchTaskQueue{} + total := 0 + var err error + if principal, ok := authctx.PrincipalFromContext(ctx); ok { + queues, total, err = h.batchTaskManager.ListQueuesForAccess(pageSize, offset, status, keyword, principal.UserID, principal.ScopeFor("tasks:read")) + } else { + return batchMCPTextResult("缺少认证身份", true), nil + } if err != nil { return batchMCPTextResult(fmt.Sprintf("列出队列失败: %v", err), true), nil } @@ -215,11 +224,20 @@ func RegisterBatchTaskMCPTools(mcpServer *mcp.Server, h *AgentHandler, logger *z executeNow = false } projectID := strings.TrimSpace(mcpArgString(args, "project_id")) + if principal, ok := authctx.PrincipalFromContext(ctx); ok && projectID != "" && principal.ScopeFor("tasks:write") != database.RBACScopeAll { + if h.db == nil || !h.db.UserCanAccessResource(principal.UserID, principal.ScopeFor("tasks:write"), "project", projectID) { + return batchMCPTextResult("无权访问目标项目", true), nil + } + } concurrency := int(mcpArgFloat(args, "concurrency")) queue, createErr := h.batchTaskManager.CreateBatchQueue(title, role, agentMode, scheduleMode, cronExpr, projectID, nextRunAt, concurrency, tasks) if createErr != nil { return batchMCPTextResult("创建队列失败: "+createErr.Error(), true), nil } + if principal, ok := authctx.PrincipalFromContext(ctx); ok && h.db != nil { + _ = h.db.SetResourceOwner("batch_task", queue.ID, principal.UserID) + _ = h.db.AssignResourceToUser(principal.UserID, "batch_task", queue.ID) + } started := false if executeNow { ok, err := h.startBatchQueueExecution(queue.ID, false) diff --git a/internal/handler/c2.go b/internal/handler/c2.go index fe977e30..b2640a69 100644 --- a/internal/handler/c2.go +++ b/internal/handler/c2.go @@ -536,6 +536,14 @@ func (h *C2Handler) CreateTask(c *gin.Context) { c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) return } + if conversationID := strings.TrimSpace(req.ConversationID); conversationID != "" { + session, ok := security.CurrentSession(c) + if !ok || !h.mgr().DB().UserCanAccessResource(session.UserID, session.Scope, "conversation", conversationID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权关联目标对话"}) + return + } + req.ConversationID = conversationID + } input := c2.EnqueueTaskInput{ SessionID: req.SessionID, @@ -722,6 +730,9 @@ func (h *C2Handler) PayloadBuild(c *gin.Context) { c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) return } + if session, ok := security.CurrentSession(c); ok { + _ = h.mgr().DB().RecordC2PayloadArtifact(filepath.Base(result.OutputPath), result.PayloadID, result.ListenerID, session.UserID) + } c.JSON(http.StatusOK, gin.H{ "payload": result, @@ -740,6 +751,11 @@ func (h *C2Handler) PayloadDownload(c *gin.Context) { c.JSON(http.StatusBadRequest, gin.H{"error": "invalid payload id"}) return } + session, ok := security.CurrentSession(c) + if !ok || !h.mgr().DB().UserCanAccessC2Payload(session.UserID, session.Scope, filename) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } builder := c2.NewPayloadBuilder(h.mgr(), h.logger, "", "") storageDir := builder.GetPayloadStoragePath() diff --git a/internal/handler/chat_uploads.go b/internal/handler/chat_uploads.go index 7ca91ebc..d8e4bd20 100644 --- a/internal/handler/chat_uploads.go +++ b/internal/handler/chat_uploads.go @@ -13,6 +13,8 @@ import ( "unicode/utf8" "cyberstrike-ai/internal/audit" + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/security" "github.com/gin-gonic/gin" "go.uber.org/zap" @@ -27,6 +29,7 @@ const ( type ChatUploadsHandler struct { logger *zap.Logger audit *audit.Service + db *database.DB } // SetAudit wires platform audit logging. @@ -35,8 +38,32 @@ func (h *ChatUploadsHandler) SetAudit(s *audit.Service) { } // NewChatUploadsHandler 创建处理器 -func NewChatUploadsHandler(logger *zap.Logger) *ChatUploadsHandler { - return &ChatUploadsHandler{logger: logger} +func NewChatUploadsHandler(logger *zap.Logger, databases ...*database.DB) *ChatUploadsHandler { + h := &ChatUploadsHandler{logger: logger} + if len(databases) > 0 { + h.db = databases[0] + } + return h +} + +func (h *ChatUploadsHandler) pathAllowed(c *gin.Context, relativePath string) bool { + session, ok := security.CurrentSession(c) + if !ok || h.db == nil { + return false + } + if session.Scope == database.RBACScopeAll { + return true + } + rel := filepath.ToSlash(filepath.Clean(filepath.FromSlash(strings.TrimSpace(relativePath)))) + rel = strings.Trim(rel, "/") + if conversationID, ownerUserID, found := h.db.GetChatUploadArtifact(rel); found { + return strings.TrimSpace(ownerUserID) == session.UserID || h.db.UserCanAccessResource(session.UserID, session.Scope, "conversation", conversationID) + } + parts := strings.Split(strings.Trim(rel, "/"), "/") + if len(parts) < 2 || parts[1] == "" || parts[1] == "_manual" { + return false + } + return h.db.UserCanAccessResource(session.UserID, session.Scope, "conversation", parts[1]) } func (h *ChatUploadsHandler) absRoot() (string, error) { @@ -175,6 +202,21 @@ func (h *ChatUploadsHandler) List(c *gin.Context) { } folders = filteredFolders } + files = filterSlice(files, func(file ChatUploadFileItem) bool { + return h.pathAllowed(c, file.RelativePath) + }) + folders = filterSlice(folders, func(folder string) bool { + if h.pathAllowed(c, folder) { + return true + } + prefix := strings.TrimSuffix(folder, "/") + "/" + for _, file := range files { + if strings.HasPrefix(file.RelativePath, prefix) { + return true + } + } + return false + }) sort.Strings(folders) sort.Slice(files, func(i, j int) bool { return files[i].ModifiedUnix > files[j].ModifiedUnix @@ -185,6 +227,10 @@ func (h *ChatUploadsHandler) List(c *gin.Context) { // Download GET /api/chat-uploads/download?path=... func (h *ChatUploadsHandler) Download(c *gin.Context) { p := c.Query("path") + if !h.pathAllowed(c, p) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } abs, err := h.resolveUnderChatUploads(p) if err != nil { c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) @@ -209,6 +255,10 @@ func (h *ChatUploadsHandler) Delete(c *gin.Context) { c.JSON(http.StatusBadRequest, gin.H{"error": "invalid body"}) return } + if !h.pathAllowed(c, body.Path) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } abs, err := h.resolveUnderChatUploads(body.Path) if err != nil { c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) @@ -238,6 +288,7 @@ func (h *ChatUploadsHandler) Delete(c *gin.Context) { return } } + _ = h.db.DeleteChatUploadArtifactPath(filepath.ToSlash(filepath.Clean(filepath.FromSlash(body.Path)))) if h.audit != nil { h.audit.RecordOK(c, "file", "delete", "删除对话附件", "chat_upload", body.Path, nil) } @@ -272,6 +323,10 @@ func (h *ChatUploadsHandler) Mkdir(c *gin.Context) { if parent == "." { parent = "" } + if !h.pathAllowed(c, parent) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } root, err := h.absRoot() if err != nil { @@ -327,6 +382,10 @@ func (h *ChatUploadsHandler) Rename(c *gin.Context) { c.JSON(http.StatusBadRequest, gin.H{"error": "invalid body"}) return } + if !h.pathAllowed(c, body.Path) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } newName := strings.TrimSpace(body.NewName) if newName == "" || strings.ContainsAny(newName, `/\`) { c.JSON(http.StatusBadRequest, gin.H{"error": "invalid newName"}) @@ -354,6 +413,8 @@ func (h *ChatUploadsHandler) Rename(c *gin.Context) { return } newRel, _ := filepath.Rel(root, newAbs) + oldRel := filepath.ToSlash(filepath.Clean(filepath.FromSlash(body.Path))) + _ = h.db.RenameChatUploadArtifactPath(oldRel, filepath.ToSlash(newRel)) c.JSON(http.StatusOK, gin.H{"ok": true, "relativePath": filepath.ToSlash(newRel)}) } @@ -365,6 +426,10 @@ type chatUploadContentBody struct { // GetContent GET /api/chat-uploads/content?path=... func (h *ChatUploadsHandler) GetContent(c *gin.Context) { p := c.Query("path") + if !h.pathAllowed(c, p) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } abs, err := h.resolveUnderChatUploads(p) if err != nil { c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) @@ -398,6 +463,10 @@ func (h *ChatUploadsHandler) PutContent(c *gin.Context) { c.JSON(http.StatusBadRequest, gin.H{"error": "invalid body"}) return } + if !h.pathAllowed(c, body.Path) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } if !utf8.ValidString(body.Content) { c.JSON(http.StatusBadRequest, gin.H{"error": "content must be valid UTF-8"}) return @@ -444,6 +513,10 @@ func (h *ChatUploadsHandler) Upload(c *gin.Context) { var targetDir string targetRel := strings.TrimSpace(c.PostForm("relativeDir")) if targetRel != "" { + if !h.pathAllowed(c, targetRel) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } absDir, err := h.resolveUnderChatUploads(targetRel) if err != nil { c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) @@ -467,13 +540,17 @@ func (h *ChatUploadsHandler) Upload(c *gin.Context) { targetDir = absDir } else { convID := strings.TrimSpace(c.PostForm("conversationId")) + dateStr := time.Now().Format("2006-01-02") + if !h.pathAllowed(c, filepath.ToSlash(filepath.Join(dateStr, convID))) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } convDir := convID if convDir == "" { convDir = "_manual" } else { convDir = strings.ReplaceAll(convDir, string(filepath.Separator), "_") } - dateStr := time.Now().Format("2006-01-02") targetDir = filepath.Join(root, dateStr, convDir) if err := os.MkdirAll(targetDir, 0755); err != nil { c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) @@ -514,6 +591,16 @@ func (h *ChatUploadsHandler) Upload(c *gin.Context) { } rel, _ := filepath.Rel(root, fullPath) absSaved, _ := filepath.Abs(fullPath) + if session, ok := security.CurrentSession(c); ok { + conversationID := strings.TrimSpace(c.PostForm("conversationId")) + if conversationID == "" { + parts := strings.Split(filepath.ToSlash(rel), "/") + if len(parts) >= 2 { + conversationID = parts[1] + } + } + _ = h.db.UpsertChatUploadArtifact(filepath.ToSlash(rel), conversationID, session.UserID) + } if h.audit != nil { h.audit.RecordOK(c, "file", "upload", "上传对话附件", "chat_upload", filepath.ToSlash(rel), map[string]interface{}{ "name": unique, diff --git a/internal/handler/eino_empty_response_continue.go b/internal/handler/eino_empty_response_continue.go index f63e3b1c..98139854 100644 --- a/internal/handler/eino_empty_response_continue.go +++ b/internal/handler/eino_empty_response_continue.go @@ -13,11 +13,11 @@ import ( ) // rebindEinoRunningTask 中断并继续 / 空正文续跑：重建 cancel 链与超时 ctx，保持任务 running。 -func (h *AgentHandler) rebindEinoRunningTask(conversationID string, timeoutCancel context.CancelFunc) (context.Context, context.CancelCauseFunc, context.Context, context.CancelFunc) { +func (h *AgentHandler) rebindEinoRunningTask(parent context.Context, conversationID string, timeoutCancel context.CancelFunc) (context.Context, context.CancelCauseFunc, context.Context, context.CancelFunc) { if timeoutCancel != nil { timeoutCancel() } - baseCtx, cancelWithCause := context.WithCancelCause(context.Background()) + baseCtx, cancelWithCause := context.WithCancelCause(detachedAgentContext(parent)) h.tasks.BindTaskCancel(conversationID, cancelWithCause) taskCtx, newTimeoutCancel := context.WithTimeout(baseCtx, 600*time.Minute) h.tasks.UpdateTaskStatus(conversationID, "running") diff --git a/internal/handler/eino_single_agent.go b/internal/handler/eino_single_agent.go index f1d2bccb..d9faf2be 100644 --- a/internal/handler/eino_single_agent.go +++ b/internal/handler/eino_single_agent.go @@ -116,7 +116,7 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) { "userMessageId": prep.UserMessageID, }) } - if h.runRoleWorkflowStreamIfBound(&req, prep, sendEvent) { + if h.runRoleWorkflowStreamIfBound(c, &req, prep, sendEvent) { return } @@ -154,7 +154,7 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) { var result *multiagent.RunResult var runErr error - baseCtx, cancelWithCause = context.WithCancelCause(context.Background()) + baseCtx, cancelWithCause = context.WithCancelCause(detachedAgentContext(c.Request.Context())) taskCtx, timeoutCancel := context.WithTimeout(baseCtx, 600*time.Minute) if _, err := h.tasks.StartTask(conversationID, req.Message, cancelWithCause); err != nil { @@ -247,7 +247,7 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) { if h.tryContinueOnEinoEmptyResponse(taskCtx, mw, conversationID, result, &emptyResponseContinueAttempt, &curHistory, &curFinalMessage, progressCallback) { mainIterationOffset += segmentMainIterationMax timeoutCancel() - baseCtx, cancelWithCause, taskCtx, timeoutCancel = h.rebindEinoRunningTask(conversationID, timeoutCancel) + baseCtx, cancelWithCause, taskCtx, timeoutCancel = h.rebindEinoRunningTask(taskCtx, conversationID, timeoutCancel) continue } timeoutCancel() @@ -279,7 +279,7 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) { }) mainIterationOffset += segmentMainIterationMax timeoutCancel() - baseCtx, cancelWithCause = context.WithCancelCause(context.Background()) + baseCtx, cancelWithCause = context.WithCancelCause(detachedAgentContext(baseCtx)) h.tasks.BindTaskCancel(conversationID, cancelWithCause) taskCtx, timeoutCancel = context.WithTimeout(baseCtx, 600*time.Minute) h.tasks.UpdateTaskStatus(conversationID, "running") @@ -381,7 +381,8 @@ func (h *AgentHandler) EinoSingleAgentLoop(c *gin.Context) { prep, err := h.prepareMultiAgentSession(&req, c, "eino_agent") if err != nil { - c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + status, msg := multiAgentHTTPErrorStatus(err) + c.JSON(status, gin.H{"error": msg}) return } h.activateHITLForConversation(prep.ConversationID, req.Hitl) diff --git a/internal/handler/external_mcp.go b/internal/handler/external_mcp.go index 931c9e09..fff5784d 100644 --- a/internal/handler/external_mcp.go +++ b/internal/handler/external_mcp.go @@ -9,6 +9,7 @@ import ( "cyberstrike-ai/internal/audit" "cyberstrike-ai/internal/config" "cyberstrike-ai/internal/mcp" + "cyberstrike-ai/internal/security" "github.com/gin-gonic/gin" "go.uber.org/zap" @@ -67,7 +68,7 @@ func (h *ExternalMCPHandler) GetExternalMCPs(c *gin.Context) { errorMsg := externalMCPStatusError(h.manager, name, status) result[name] = ExternalMCPResponse{ - Config: cfg, + Config: externalMCPConfigForResponse(c, cfg), Status: status, ToolCount: toolCount, Error: errorMsg, @@ -113,13 +114,33 @@ func (h *ExternalMCPHandler) GetExternalMCP(c *gin.Context) { } c.JSON(http.StatusOK, ExternalMCPResponse{ - Config: cfg, + Config: externalMCPConfigForResponse(c, cfg), Status: status, ToolCount: toolCount, Error: externalMCPStatusError(h.manager, name, status), }) } +func externalMCPConfigForResponse(c *gin.Context, cfg config.ExternalMCPServerConfig) config.ExternalMCPServerConfig { + if security.SessionHasPermission(c, "mcp:write") { + return cfg + } + copyCfg := cfg + if len(cfg.Env) > 0 { + copyCfg.Env = make(map[string]string, len(cfg.Env)) + for key := range cfg.Env { + copyCfg.Env[key] = "***" + } + } + if len(cfg.Headers) > 0 { + copyCfg.Headers = make(map[string]string, len(cfg.Headers)) + for key := range cfg.Headers { + copyCfg.Headers[key] = "***" + } + } + return copyCfg +} + // externalMCPStatusError 在 error/disconnected 状态下返回最近错误（含断连原因）。 func externalMCPStatusError(manager *mcp.ExternalMCPManager, name, status string) string { if status != "error" && status != "disconnected" { diff --git a/internal/handler/group.go b/internal/handler/group.go index db44e228..23045e32 100644 --- a/internal/handler/group.go +++ b/internal/handler/group.go @@ -44,7 +44,8 @@ func (h *GroupHandler) CreateGroup(c *gin.Context) { return } - group, err := h.db.CreateGroup(req.Name, req.Icon) + session, _ := security.CurrentSession(c) + group, err := h.db.CreateGroup(req.Name, req.Icon, session.UserID) if err != nil { h.logger.Error("创建分组失败", zap.Error(err)) // 如果是名称重复错误，返回400状态码 @@ -61,7 +62,8 @@ func (h *GroupHandler) CreateGroup(c *gin.Context) { // ListGroups 列出所有分组 func (h *GroupHandler) ListGroups(c *gin.Context) { - groups, err := h.db.ListGroups() + session, _ := security.CurrentSession(c) + groups, err := h.db.ListGroupsForAccess(session.UserID, session.Scope) if err != nil { h.logger.Error("获取分组列表失败", zap.Error(err)) c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) @@ -74,6 +76,10 @@ func (h *GroupHandler) ListGroups(c *gin.Context) { // GetGroup 获取分组 func (h *GroupHandler) GetGroup(c *gin.Context) { id := c.Param("id") + if !h.groupAllowed(c, id) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } group, err := h.db.GetGroup(id) if err != nil { @@ -94,6 +100,10 @@ type UpdateGroupRequest struct { // UpdateGroup 更新分组 func (h *GroupHandler) UpdateGroup(c *gin.Context) { id := c.Param("id") + if !h.groupAllowed(c, id) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } var req UpdateGroupRequest if err := c.ShouldBindJSON(&req); err != nil { @@ -130,6 +140,10 @@ func (h *GroupHandler) UpdateGroup(c *gin.Context) { // DeleteGroup 删除分组 func (h *GroupHandler) DeleteGroup(c *gin.Context) { id := c.Param("id") + if !h.groupAllowed(c, id) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } if err := h.db.DeleteGroup(id); err != nil { h.logger.Error("删除分组失败", zap.Error(err)) @@ -157,6 +171,10 @@ func (h *GroupHandler) AddConversationToGroup(c *gin.Context) { c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) return } + if !h.groupAllowed(c, req.GroupID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该分组"}) + return + } if err := h.db.AddConversationToGroup(req.ConversationID, req.GroupID); err != nil { h.logger.Error("添加对话到分组失败", zap.Error(err)) @@ -171,6 +189,10 @@ func (h *GroupHandler) AddConversationToGroup(c *gin.Context) { func (h *GroupHandler) RemoveConversationFromGroup(c *gin.Context) { conversationID := c.Param("conversationId") groupID := c.Param("id") + if !h.groupAllowed(c, groupID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该分组"}) + return + } if !h.groupConversationAllowed(c, conversationID) { c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) return @@ -198,6 +220,10 @@ type GroupConversation struct { // GetGroupConversations 获取分组中的所有对话 func (h *GroupHandler) GetGroupConversations(c *gin.Context) { groupID := c.Param("id") + if !h.groupAllowed(c, groupID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该分组"}) + return + } searchQuery := c.Query("search") // 获取搜索参数 var conversations []*database.Conversation @@ -256,7 +282,7 @@ func (h *GroupHandler) GetAllMappings(c *gin.Context) { } filtered := mappings[:0] for _, mapping := range mappings { - if h.groupConversationAllowed(c, mapping.ConversationID) { + if h.groupConversationAllowed(c, mapping.ConversationID) && h.groupAllowed(c, mapping.GroupID) { filtered = append(filtered, mapping) } } @@ -300,6 +326,10 @@ type UpdateGroupPinnedRequest struct { // UpdateGroupPinned 更新分组置顶状态 func (h *GroupHandler) UpdateGroupPinned(c *gin.Context) { groupID := c.Param("id") + if !h.groupAllowed(c, groupID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该分组"}) + return + } var req UpdateGroupPinnedRequest if err := c.ShouldBindJSON(&req); err != nil { @@ -325,6 +355,10 @@ type UpdateConversationPinnedInGroupRequest struct { func (h *GroupHandler) UpdateConversationPinnedInGroup(c *gin.Context) { groupID := c.Param("id") conversationID := c.Param("conversationId") + if !h.groupAllowed(c, groupID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该分组"}) + return + } if !h.groupConversationAllowed(c, conversationID) { c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) return @@ -352,3 +386,8 @@ func (h *GroupHandler) groupConversationAllowed(c *gin.Context, conversationID s } return h.db.UserCanAccessResource(session.UserID, session.Scope, "conversation", conversationID) } + +func (h *GroupHandler) groupAllowed(c *gin.Context, groupID string) bool { + session, ok := security.CurrentSession(c) + return ok && h.db.UserCanAccessGroup(session.UserID, session.Scope, groupID) +} diff --git a/internal/handler/monitor.go b/internal/handler/monitor.go index 8f597bbf..fe5af636 100644 --- a/internal/handler/monitor.go +++ b/internal/handler/monitor.go @@ -120,9 +120,22 @@ func (h *MonitorHandler) Monitor(c *gin.Context) { // 解析工具筛选参数（兼容 mcp__tool 与内部 mcp::tool） toolName := normalizeToolNameFilter(c.Query("tool")) - executions, total := h.loadExecutionListWithPagination(page, pageSize, status, toolName) + access := notificationAccessFromContext(c) + executions, total := h.loadExecutionListWithPagination(page, pageSize, status, toolName, access) h.enrichExecutionsConversationID(executions) - summary, topTools := h.loadStatsSummary(monitorPageTopTools) + var summary *MonitorStatsSummary + var topTools []*mcp.ToolStats + if access.Scope == database.RBACScopeAll { + summary, topTools = h.loadStatsSummary(monitorPageTopTools) + } else if h.db != nil { + if scoped, err := h.db.LoadToolStatsSummaryForAccess(monitorPageTopTools, access); err == nil { + summary, topTools = dbStatsSummaryToMonitor(scoped), scoped.TopTools + } else { + summary, topTools = summarizeAccessibleExecutionPage(executions, monitorPageTopTools) + } + } else { + summary, topTools = summarizeAccessibleExecutionPage(executions, monitorPageTopTools) + } totalPages := (total + pageSize - 1) / pageSize if totalPages == 0 { @@ -142,6 +155,31 @@ func (h *MonitorHandler) Monitor(c *gin.Context) { }) } +func summarizeAccessibleExecutionPage(executions []*mcp.ToolExecution, topN int) (*MonitorStatsSummary, []*mcp.ToolStats) { + stats := map[string]*mcp.ToolStats{} + for _, exec := range executions { + if exec == nil { + continue + } + stat := stats[exec.ToolName] + if stat == nil { + stat = &mcp.ToolStats{ToolName: exec.ToolName} + stats[exec.ToolName] = stat + } + stat.TotalCalls++ + if exec.Status == "failed" || exec.Status == "cancelled" { + stat.FailedCalls++ + } else if exec.Status == "completed" { + stat.SuccessCalls++ + } + started := exec.StartTime + if stat.LastCallTime == nil || started.After(*stat.LastCallTime) { + stat.LastCallTime = &started + } + } + return summarizeToolStats(stats, topN) +} + func (h *MonitorHandler) monitorRetentionDays() int { if h.monitorRetention != nil { return h.monitorRetention.RetentionDays() @@ -154,9 +192,9 @@ func (h *MonitorHandler) loadExecutions() []*mcp.ToolExecution { return executions } -func (h *MonitorHandler) loadExecutionListWithPagination(page, pageSize int, status, toolName string) ([]*mcp.ToolExecution, int) { +func (h *MonitorHandler) loadExecutionListWithPagination(page, pageSize int, status, toolName string, access database.RBACListAccess) ([]*mcp.ToolExecution, int) { if h.db == nil { - allExecutions := h.mcpServer.GetAllExecutions() + allExecutions := filterToolExecutionsForAccess(h.mcpServer.GetAllExecutions(), access, h.db) if status != "" || toolName != "" { filtered := make([]*mcp.ToolExecution, 0) for _, exec := range allExecutions { @@ -189,13 +227,13 @@ func (h *MonitorHandler) loadExecutionListWithPagination(page, pageSize int, sta } offset := (page - 1) * pageSize - executions, err := h.db.LoadToolExecutionListPage(offset, pageSize, status, toolName) + executions, err := h.db.LoadToolExecutionListPageForAccess(offset, pageSize, status, toolName, access) if err != nil { h.logger.Warn("从数据库加载执行记录列表失败，回退到内存数据", zap.Error(err)) - return h.loadExecutionListWithPaginationFromMemory(page, pageSize, status, toolName) + return h.loadExecutionListWithPaginationFromMemory(page, pageSize, status, toolName, access) } - total, err := h.db.CountToolExecutions(status, toolName) + total, err := h.db.CountToolExecutionsForAccess(status, toolName, access) if err != nil { h.logger.Warn("获取执行记录总数失败", zap.Error(err)) total = offset + len(executions) @@ -207,8 +245,8 @@ func (h *MonitorHandler) loadExecutionListWithPagination(page, pageSize int, sta return executions, total } -func (h *MonitorHandler) loadExecutionListWithPaginationFromMemory(page, pageSize int, status, toolName string) ([]*mcp.ToolExecution, int) { - allExecutions := h.mcpServer.GetAllExecutions() +func (h *MonitorHandler) loadExecutionListWithPaginationFromMemory(page, pageSize int, status, toolName string, access database.RBACListAccess) ([]*mcp.ToolExecution, int) { + allExecutions := filterToolExecutionsForAccess(h.mcpServer.GetAllExecutions(), access, h.db) if status != "" || toolName != "" { filtered := make([]*mcp.ToolExecution, 0) for _, exec := range allExecutions { @@ -260,6 +298,50 @@ func slimToolExecution(exec *mcp.ToolExecution) *mcp.ToolExecution { return slim } +func filterToolExecutionsForAccess(executions []*mcp.ToolExecution, access database.RBACListAccess, db *database.DB) []*mcp.ToolExecution { + if access.Scope == database.RBACScopeAll { + return executions + } + out := make([]*mcp.ToolExecution, 0, len(executions)) + for _, exec := range executions { + if toolExecutionVisible(exec, access, db) { + out = append(out, exec) + } + } + return out +} + +func toolExecutionVisible(exec *mcp.ToolExecution, access database.RBACListAccess, db *database.DB) bool { + if exec == nil || strings.TrimSpace(access.UserID) == "" { + return false + } + if access.Scope == database.RBACScopeAll || strings.TrimSpace(exec.OwnerUserID) == strings.TrimSpace(access.UserID) { + return true + } + conversationID := strings.TrimSpace(exec.ConversationID) + return conversationID != "" && db != nil && db.UserCanAccessResource(access.UserID, access.Scope, "conversation", conversationID) +} + +func (h *MonitorHandler) monitorExecutionAllowed(c *gin.Context, id string) bool { + access := notificationAccessFromContext(c) + if access.Scope == database.RBACScopeAll { + return true + } + id = strings.TrimSpace(id) + if id == "" { + return false + } + if exec, ok := h.mcpServer.GetExecution(id); ok { + return toolExecutionVisible(exec, access, h.db) + } + if h.externalMCPMgr != nil { + if exec, ok := h.externalMCPMgr.GetExecution(id); ok { + return toolExecutionVisible(exec, access, h.db) + } + } + return h.db != nil && h.db.UserCanAccessToolExecution(access.UserID, access.Scope, id) +} + func (h *MonitorHandler) loadExecutionsWithPagination(page, pageSize int, status, toolName string) ([]*mcp.ToolExecution, int) { if h.db == nil { allExecutions := h.mcpServer.GetAllExecutions() @@ -453,6 +535,10 @@ func (h *MonitorHandler) loadStatsMap() map[string]*mcp.ToolStats { // GetExecution 获取特定执行记录 func (h *MonitorHandler) GetExecution(c *gin.Context) { id := c.Param("id") + if !h.monitorExecutionAllowed(c, id) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } // 先从内部MCP服务器查找 exec, exists := h.mcpServer.GetExecution(id) @@ -493,6 +579,10 @@ func (h *MonitorHandler) CancelExecution(c *gin.Context) { c.JSON(http.StatusBadRequest, gin.H{"error": "执行记录ID不能为空"}) return } + if !h.monitorExecutionAllowed(c, id) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } note := "" dec := json.NewDecoder(c.Request.Body) var body struct { @@ -592,6 +682,9 @@ func (h *MonitorHandler) BatchGetToolNames(c *gin.Context) { result := make(map[string]executionSummary, len(req.IDs)) for _, id := range req.IDs { + if !h.monitorExecutionAllowed(c, id) { + continue + } // 先从内部MCP服务器查找 if exec, exists := h.mcpServer.GetExecution(id); exists { result[id] = executionSummary{ToolName: exec.ToolName, Status: exec.Status} @@ -755,6 +848,10 @@ func (h *MonitorHandler) DeleteExecution(c *gin.Context) { c.JSON(http.StatusBadRequest, gin.H{"error": "执行记录ID不能为空"}) return } + if !h.monitorExecutionAllowed(c, id) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } // 如果使用数据库，先获取执行记录信息，然后删除并更新统计 if h.db != nil { @@ -823,6 +920,12 @@ func (h *MonitorHandler) DeleteExecutions(c *gin.Context) { c.JSON(http.StatusBadRequest, gin.H{"error": "执行记录ID列表不能为空"}) return } + for _, id := range request.IDs { + if !h.monitorExecutionAllowed(c, id) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问一个或多个执行记录"}) + return + } + } // 如果使用数据库，先获取执行记录信息，然后删除并更新统计 if h.db != nil { diff --git a/internal/handler/multi_agent.go b/internal/handler/multi_agent.go index f7e0f7ac..0760bcac 100644 --- a/internal/handler/multi_agent.go +++ b/internal/handler/multi_agent.go @@ -133,7 +133,7 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) { "userMessageId": prep.UserMessageID, }) } - if h.runRoleWorkflowStreamIfBound(&req, prep, sendEvent) { + if h.runRoleWorkflowStreamIfBound(c, &req, prep, sendEvent) { return } @@ -163,7 +163,7 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) { var result *multiagent.RunResult var runErr error - baseCtx, cancelWithCause = context.WithCancelCause(context.Background()) + baseCtx, cancelWithCause = context.WithCancelCause(detachedAgentContext(c.Request.Context())) taskCtx, timeoutCancel := context.WithTimeout(baseCtx, 600*time.Minute) if _, err := h.tasks.StartTask(conversationID, req.Message, cancelWithCause); err != nil { @@ -259,7 +259,7 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) { if h.tryContinueOnEinoEmptyResponse(taskCtx, mw, conversationID, result, &emptyResponseContinueAttempt, &curHistory, &curFinalMessage, progressCallback) { mainIterationOffset += segmentMainIterationMax timeoutCancel() - baseCtx, cancelWithCause, taskCtx, timeoutCancel = h.rebindEinoRunningTask(conversationID, timeoutCancel) + baseCtx, cancelWithCause, taskCtx, timeoutCancel = h.rebindEinoRunningTask(taskCtx, conversationID, timeoutCancel) continue } timeoutCancel() @@ -291,7 +291,7 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) { }) mainIterationOffset += segmentMainIterationMax timeoutCancel() - baseCtx, cancelWithCause = context.WithCancelCause(context.Background()) + baseCtx, cancelWithCause = context.WithCancelCause(detachedAgentContext(baseCtx)) h.tasks.BindTaskCancel(conversationID, cancelWithCause) taskCtx, timeoutCancel = context.WithTimeout(baseCtx, 600*time.Minute) h.tasks.UpdateTaskStatus(conversationID, "running") @@ -541,6 +541,8 @@ func formatInterruptContinueUserMessage(note string) string { func multiAgentHTTPErrorStatus(err error) (int, string) { msg := err.Error() switch { + case strings.Contains(msg, "无权访问"): + return http.StatusForbidden, msg case strings.Contains(msg, "对话不存在"): return http.StatusNotFound, msg case strings.Contains(msg, "未找到该 WebShell"): diff --git a/internal/handler/multi_agent_prepare.go b/internal/handler/multi_agent_prepare.go index 8f45919d..4becd55c 100644 --- a/internal/handler/multi_agent_prepare.go +++ b/internal/handler/multi_agent_prepare.go @@ -8,6 +8,7 @@ import ( "cyberstrike-ai/internal/audit" "cyberstrike-ai/internal/database" "cyberstrike-ai/internal/mcp/builtin" + "cyberstrike-ai/internal/security" "github.com/gin-gonic/gin" "go.uber.org/zap" @@ -30,16 +31,34 @@ func (h *AgentHandler) prepareMultiAgentSession(req *ChatRequest, c *gin.Context } conversationID := strings.TrimSpace(req.ConversationID) + projectID := strings.TrimSpace(effectiveProjectID(h.config, req.ProjectID)) + webshellID := strings.TrimSpace(req.WebShellConnectionID) + session, hasSession := security.CurrentSession(c) + if !hasSession || !session.Permissions["chat:write"] { + return nil, fmt.Errorf("无权写入对话") + } + canAccess := func(resourceType, resourceID string) bool { + if !hasSession || h.db == nil || strings.TrimSpace(resourceID) == "" { + return false + } + return h.db.UserCanAccessResource(session.UserID, session.Scope, resourceType, resourceID) + } + if projectID != "" && (!session.Permissions["project:read"] || !canAccess("project", projectID)) { + return nil, fmt.Errorf("无权访问目标项目") + } + if webshellID != "" && (!session.Permissions["webshell:write"] || !canAccess("webshell", webshellID)) { + return nil, fmt.Errorf("无权访问该 WebShell 连接") + } createdNew := false if conversationID == "" { title := safeTruncateString(req.Message, 50) var conv *database.Conversation var err error meta := audit.ConversationCreateMetaFromGin(c, source) - meta.ProjectID = effectiveProjectID(h.config, req.ProjectID) - if strings.TrimSpace(req.WebShellConnectionID) != "" { + meta.ProjectID = projectID + if webshellID != "" { meta.Source = source + "_webshell" - meta.WebShellConnectionID = strings.TrimSpace(req.WebShellConnectionID) + meta.WebShellConnectionID = webshellID conv, err = h.db.CreateConversationWithWebshell(meta.WebShellConnectionID, title, meta) } else { conv, err = h.db.CreateConversation(title, meta) @@ -49,10 +68,17 @@ func (h *AgentHandler) prepareMultiAgentSession(req *ChatRequest, c *gin.Context } conversationID = conv.ID createdNew = true + if hasSession { + _ = h.db.SetResourceOwner("conversation", conversationID, session.UserID) + _ = h.db.AssignResourceToUser(session.UserID, "conversation", conversationID) + } } else { if _, err := h.db.GetConversation(conversationID); err != nil { return nil, fmt.Errorf("对话不存在") } + if !canAccess("conversation", conversationID) { + return nil, fmt.Errorf("无权访问该对话") + } } agentHistoryMessages, err := h.loadHistoryFromAgentTrace(conversationID) @@ -67,8 +93,8 @@ func (h *AgentHandler) prepareMultiAgentSession(req *ChatRequest, c *gin.Context finalMessage := req.Message var roleTools []string - if req.WebShellConnectionID != "" { - conn, errConn := h.db.GetWebshellConnection(strings.TrimSpace(req.WebShellConnectionID)) + if webshellID != "" { + conn, errConn := h.db.GetWebshellConnection(webshellID) if errConn != nil || conn == nil { h.logger.Warn("WebShell AI 助手：未找到连接", zap.String("id", req.WebShellConnectionID), zap.Error(errConn)) return nil, fmt.Errorf("未找到该 WebShell 连接") diff --git a/internal/handler/notification.go b/internal/handler/notification.go index b83c8e93..21f99d9b 100644 --- a/internal/handler/notification.go +++ b/internal/handler/notification.go @@ -494,14 +494,16 @@ func buildPlaceholders(n int) string { return strings.Join(out, ",") } -func (h *NotificationHandler) readStatesByIDs(ids []string) (map[string]bool, error) { +func (h *NotificationHandler) readStatesByIDs(userID string, ids []string) (map[string]bool, error) { result := make(map[string]bool, len(ids)) - if len(ids) == 0 { + userID = strings.TrimSpace(userID) + if len(ids) == 0 || userID == "" { return result, nil } holders := buildPlaceholders(len(ids)) - query := "SELECT event_id FROM notification_reads WHERE event_id IN (" + holders + ")" - args := make([]interface{}, 0, len(ids)) + query := "SELECT event_id FROM notification_reads_by_user WHERE user_id = ? AND event_id IN (" + holders + ")" + args := make([]interface{}, 0, len(ids)+1) + args = append(args, userID) for _, id := range ids { args = append(args, id) } @@ -520,7 +522,7 @@ func (h *NotificationHandler) readStatesByIDs(ids []string) (map[string]bool, er return result, nil } -func (h *NotificationHandler) applyReadStates(items []NotificationSummaryItem) ([]NotificationSummaryItem, error) { +func (h *NotificationHandler) applyReadStates(userID string, items []NotificationSummaryItem) ([]NotificationSummaryItem, error) { markableIDs := make([]string, 0, len(items)) for _, item := range items { if item.Actionable { @@ -528,7 +530,7 @@ func (h *NotificationHandler) applyReadStates(items []NotificationSummaryItem) ( } markableIDs = append(markableIDs, item.ID) } - readMap, err := h.readStatesByIDs(markableIDs) + readMap, err := h.readStatesByIDs(userID, markableIDs) if err != nil { return items, err } @@ -585,34 +587,38 @@ func createNotificationReadTableIfNeeded(db *database.DB) error { return fmt.Errorf("db is nil") } _, err := db.Exec(` - CREATE TABLE IF NOT EXISTS notification_reads ( - event_id TEXT PRIMARY KEY, - read_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP + CREATE TABLE IF NOT EXISTS notification_reads_by_user ( + user_id TEXT NOT NULL, + event_id TEXT NOT NULL, + read_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, + PRIMARY KEY(user_id, event_id) ); `) if err != nil { return err } - _, idxErr := db.Exec(`CREATE INDEX IF NOT EXISTS idx_notification_reads_read_at ON notification_reads(read_at DESC);`) + _, idxErr := db.Exec(`CREATE INDEX IF NOT EXISTS idx_notification_reads_user_read_at ON notification_reads_by_user(user_id, read_at DESC);`) return idxErr } -func pruneNotificationReads(db *database.DB, maxRows int) error { +func pruneNotificationReads(db *database.DB, userID string, maxRows int) error { if db == nil { return fmt.Errorf("db is nil") } - if maxRows <= 0 { + userID = strings.TrimSpace(userID) + if maxRows <= 0 || userID == "" { return nil } _, err := db.Exec(` - DELETE FROM notification_reads - WHERE event_id NOT IN ( + DELETE FROM notification_reads_by_user + WHERE user_id = ? AND event_id NOT IN ( SELECT event_id - FROM notification_reads + FROM notification_reads_by_user + WHERE user_id = ? ORDER BY read_at DESC, rowid DESC LIMIT ? ) - `, maxRows) + `, userID, userID, maxRows) return err } @@ -655,6 +661,11 @@ func (h *NotificationHandler) MarkRead(c *gin.Context) { c.JSON(http.StatusOK, gin.H{"ok": true, "marked": 0}) return } + session, ok := security.CurrentSession(c) + if !ok || strings.TrimSpace(session.UserID) == "" { + c.JSON(http.StatusUnauthorized, gin.H{"error": "missing authenticated user"}) + return + } tx, err := h.db.Begin() if err != nil { c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to begin transaction"}) @@ -664,9 +675,9 @@ func (h *NotificationHandler) MarkRead(c *gin.Context) { _ = tx.Rollback() }() stmt, err := tx.Prepare(` - INSERT INTO notification_reads(event_id, read_at) - VALUES(?, CURRENT_TIMESTAMP) - ON CONFLICT(event_id) DO UPDATE SET read_at = CURRENT_TIMESTAMP + INSERT INTO notification_reads_by_user(user_id, event_id, read_at) + VALUES(?, ?, CURRENT_TIMESTAMP) + ON CONFLICT(user_id, event_id) DO UPDATE SET read_at = CURRENT_TIMESTAMP `) if err != nil { c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to prepare statement"}) @@ -679,7 +690,7 @@ func (h *NotificationHandler) MarkRead(c *gin.Context) { if !ok { continue } - if _, err := stmt.Exec(id); err != nil { + if _, err := stmt.Exec(session.UserID, id); err != nil { c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to mark read"}) return } @@ -689,7 +700,7 @@ func (h *NotificationHandler) MarkRead(c *gin.Context) { c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to commit read marks"}) return } - if err := pruneNotificationReads(h.db, notificationReadMaxRows); err != nil { + if err := pruneNotificationReads(h.db, session.UserID, notificationReadMaxRows); err != nil { h.logger.Warn("裁剪通知已读记录失败", zap.Error(err)) } c.JSON(http.StatusOK, gin.H{"ok": true, "marked": marked}) @@ -776,7 +787,8 @@ func (h *NotificationHandler) GetSummary(c *gin.Context) { items = append(items, longRunningItems...) items = append(items, completedItems...) - items, err := h.applyReadStates(items) + session, _ := security.CurrentSession(c) + items, err := h.applyReadStates(session.UserID, items) if err != nil { h.logger.Warn("加载通知已读状态失败", zap.Error(err)) c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to load notification read states"}) diff --git a/internal/handler/notification_rbac_test.go b/internal/handler/notification_rbac_test.go new file mode 100644 index 00000000..ad05edef --- /dev/null +++ b/internal/handler/notification_rbac_test.go @@ -0,0 +1,45 @@ +package handler + +import ( + "bytes" + "net/http" + "net/http/httptest" + "path/filepath" + "testing" + + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/security" + + "github.com/gin-gonic/gin" + "go.uber.org/zap" +) + +func TestNotificationReadStateIsPerUser(t *testing.T) { + db, err := database.NewDB(filepath.Join(t.TempDir(), "notification-rbac.db"), zap.NewNop()) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = db.Close() }) + h := NewNotificationHandler(db, nil, zap.NewNop()) + router := gin.New() + router.Use(func(c *gin.Context) { + c.Set(security.ContextSessionKey, security.Session{UserID: "u1", Scope: database.RBACScopeAssigned}) + c.Next() + }) + router.POST("/notifications/read", h.MarkRead) + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/notifications/read", bytes.NewBufferString(`{"eventIds":["vuln:v1"]}`)) + req.Header.Set("Content-Type", "application/json") + router.ServeHTTP(w, req) + if w.Code != http.StatusOK { + t.Fatalf("mark read status = %d: %s", w.Code, w.Body.String()) + } + u1, err := h.readStatesByIDs("u1", []string{"vuln:v1"}) + if err != nil || !u1["vuln:v1"] { + t.Fatalf("u1 read state = %#v, err=%v", u1, err) + } + u2, err := h.readStatesByIDs("u2", []string{"vuln:v1"}) + if err != nil || u2["vuln:v1"] { + t.Fatalf("u2 inherited u1 read state = %#v, err=%v", u2, err) + } +} diff --git a/internal/handler/project.go b/internal/handler/project.go index e19e7e38..c671ac65 100644 --- a/internal/handler/project.go +++ b/internal/handler/project.go @@ -642,6 +642,12 @@ func (h *ProjectHandler) DeleteFactEdge(c *gin.Context) { func (h *ProjectHandler) PromoteAttackChain(c *gin.Context) { projectID := c.Param("id") conversationID := c.Param("conversationId") + session, ok := security.CurrentSession(c) + if !ok || !h.db.UserCanAccessResource(session.UserID, session.Scope, "project", projectID) || + !h.db.UserCanAccessResource(session.UserID, session.Scope, "conversation", conversationID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问目标项目或来源对话"}) + return + } result, err := attackchain.PromoteToProject(h.db, projectID, conversationID) if err != nil { c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) diff --git a/internal/handler/rbac.go b/internal/handler/rbac.go index 6dc11d36..4e21ce41 100644 --- a/internal/handler/rbac.go +++ b/internal/handler/rbac.go @@ -1,11 +1,13 @@ package handler import ( + "fmt" "net/http" "strconv" "strings" "cyberstrike-ai/internal/audit" + "cyberstrike-ai/internal/authctx" "cyberstrike-ai/internal/database" "cyberstrike-ai/internal/security" @@ -34,15 +36,22 @@ func (h *RBACHandler) SetAuthManager(m *security.AuthManager) { func (h *RBACHandler) Me(c *gin.Context) { session, _ := security.CurrentSession(c) + resolvedScope := session.Scope + permissionScopes := session.PermissionScopes + if principal, ok := authctx.PrincipalFromContext(c.Request.Context()); ok { + resolvedScope = principal.Scope + permissionScopes = principal.PermissionScopes + } c.JSON(http.StatusOK, gin.H{ "user": gin.H{ "id": session.UserID, "username": session.Username, "display_name": session.DisplayName, }, - "roles": session.Roles, - "permissions": permissionKeys(session.Permissions), - "scope": session.Scope, + "roles": session.Roles, + "permissions": permissionKeys(session.Permissions), + "scope": resolvedScope, + "permission_scopes": permissionScopes, }) } @@ -96,12 +105,29 @@ type upsertRBACRoleRequest struct { Permissions []string `json:"permissions"` } +func validateRBACPermissionKeys(keys []string) error { + for _, key := range keys { + key = strings.TrimSpace(key) + if key == "" { + continue + } + if _, ok := security.PermissionCatalog[key]; !ok { + return fmt.Errorf("未知权限: %s", key) + } + } + return nil +} + func (h *RBACHandler) CreateRole(c *gin.Context) { var req upsertRBACRoleRequest if err := c.ShouldBindJSON(&req); err != nil { c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) return } + if err := validateRBACPermissionKeys(req.Permissions); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } role, err := h.db.UpsertRBACRole("", req.Name, req.Description, req.Scope, req.Permissions) if err != nil { c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) @@ -129,6 +155,10 @@ func (h *RBACHandler) UpdateRole(c *gin.Context) { c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) return } + if err := validateRBACPermissionKeys(req.Permissions); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } role, err := h.db.UpsertRBACRole(id, req.Name, req.Description, req.Scope, req.Permissions) if err != nil { c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) diff --git a/internal/handler/rbac_boundary_test.go b/internal/handler/rbac_boundary_test.go new file mode 100644 index 00000000..c353db2c --- /dev/null +++ b/internal/handler/rbac_boundary_test.go @@ -0,0 +1,206 @@ +package handler + +import ( + "bytes" + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "path/filepath" + "testing" + "time" + + "cyberstrike-ai/internal/authctx" + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/mcp" + "cyberstrike-ai/internal/security" + + "github.com/gin-gonic/gin" + "go.uber.org/zap" +) + +func TestDetachedAgentContextRetainsPrincipalWithoutParentCancellation(t *testing.T) { + parent, cancel := context.WithCancel(context.Background()) + parent = authctx.WithPrincipal(parent, authctx.NewPrincipal("u1", "user", database.RBACScopeAssigned, map[string]bool{"agent:execute": true})) + detached := detachedAgentContext(parent) + cancel() + if err := detached.Err(); err != nil { + t.Fatalf("detached context inherited cancellation: %v", err) + } + principal, ok := authctx.PrincipalFromContext(detached) + if !ok || principal.UserID != "u1" || !principal.HasPermission("agent:execute") { + t.Fatalf("detached context lost principal: %#v, ok=%v", principal, ok) + } +} + +func TestPromoteAttackChainRequiresSourceConversationAccess(t *testing.T) { + db, err := database.NewDB(filepath.Join(t.TempDir(), "promote-rbac.db"), zap.NewNop()) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = db.Close() }) + project, _ := db.CreateProject(&database.Project{Name: "owned"}) + conversation, _ := db.CreateConversation("foreign", database.ConversationCreateMeta{}) + _ = db.SetResourceOwner("project", project.ID, "u1") + _ = db.SetResourceOwner("conversation", conversation.ID, "u2") + h := NewProjectHandler(db, zap.NewNop()) + router := gin.New() + router.Use(func(c *gin.Context) { + c.Set(security.ContextSessionKey, security.Session{UserID: "u1", Scope: database.RBACScopeOwn}) + c.Next() + }) + router.POST("/api/projects/:id/promote-attack-chain/:conversationId", h.PromoteAttackChain) + w := httptest.NewRecorder() + router.ServeHTTP(w, httptest.NewRequest(http.MethodPost, "/api/projects/"+project.ID+"/promote-attack-chain/"+conversation.ID, nil)) + if w.Code != http.StatusForbidden { + t.Fatalf("status = %d, want 403: %s", w.Code, w.Body.String()) + } +} + +func TestVulnerabilityCannotBeReparentedToForeignProject(t *testing.T) { + db, err := database.NewDB(filepath.Join(t.TempDir(), "vuln-reparent-rbac.db"), zap.NewNop()) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = db.Close() }) + owned, _ := db.CreateProject(&database.Project{Name: "owned"}) + foreign, _ := db.CreateProject(&database.Project{Name: "foreign"}) + _ = db.SetResourceOwner("project", owned.ID, "u1") + _ = db.SetResourceOwner("project", foreign.ID, "u2") + vulnerability, err := db.CreateVulnerability(&database.Vulnerability{Title: "v", Severity: "high", ProjectID: owned.ID}) + if err != nil { + t.Fatal(err) + } + _ = db.SetResourceOwner("vulnerability", vulnerability.ID, "u1") + h := NewVulnerabilityHandler(db, zap.NewNop()) + router := gin.New() + router.Use(func(c *gin.Context) { + c.Set(security.ContextSessionKey, security.Session{UserID: "u1", Scope: database.RBACScopeOwn}) + c.Next() + }) + router.PUT("/api/vulnerabilities/:id", h.UpdateVulnerability) + body, _ := json.Marshal(map[string]interface{}{"project_id": foreign.ID}) + w := httptest.NewRecorder() + router.ServeHTTP(w, httptest.NewRequest(http.MethodPut, "/api/vulnerabilities/"+vulnerability.ID, bytes.NewReader(body))) + if w.Code != http.StatusForbidden { + t.Fatalf("status = %d, want 403: %s", w.Code, w.Body.String()) + } +} + +func TestAgentTaskEndpointsFilterAndRejectForeignConversations(t *testing.T) { + gin.SetMode(gin.TestMode) + db, user := setupConversationRBACTest(t) + allowed, _ := db.CreateConversation("allowed", database.ConversationCreateMeta{}) + hidden, _ := db.CreateConversation("hidden", database.ConversationCreateMeta{}) + if err := db.AssignResourceToUser(user.ID, "conversation", allowed.ID); err != nil { + t.Fatal(err) + } + tasks := NewAgentTaskManager() + if _, err := tasks.StartTask(allowed.ID, "visible", func(error) {}); err != nil { + t.Fatal(err) + } + if _, err := tasks.StartTask(hidden.ID, "secret", func(error) {}); err != nil { + t.Fatal(err) + } + h := &AgentHandler{db: db, tasks: tasks, logger: zap.NewNop()} + + w := performAssignedHandler(user, http.MethodGet, "/api/agent-loop/tasks", nil, h.ListAgentTasks) + if w.Code != http.StatusOK { + t.Fatalf("list status = %d: %s", w.Code, w.Body.String()) + } + var response struct { + Tasks []*AgentTask `json:"tasks"` + } + if err := json.Unmarshal(w.Body.Bytes(), &response); err != nil { + t.Fatal(err) + } + if len(response.Tasks) != 1 || response.Tasks[0].ConversationID != allowed.ID { + t.Fatalf("tasks = %#v, want only %s", response.Tasks, allowed.ID) + } + + w = performAssignedHandler(user, http.MethodPost, "/api/agent-loop/cancel", map[string]string{"conversationId": hidden.ID}, h.CancelAgentLoop) + if w.Code != http.StatusForbidden { + t.Fatalf("cancel status = %d, want %d: %s", w.Code, http.StatusForbidden, w.Body.String()) + } +} + +func TestChatUploadPathAuthorizationFollowsConversationAccess(t *testing.T) { + db, user := setupConversationRBACTest(t) + allowed, _ := db.CreateConversation("allowed", database.ConversationCreateMeta{}) + hidden, _ := db.CreateConversation("hidden", database.ConversationCreateMeta{}) + if err := db.AssignResourceToUser(user.ID, "conversation", allowed.ID); err != nil { + t.Fatal(err) + } + h := NewChatUploadsHandler(zap.NewNop(), db) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Set(security.ContextSessionKey, security.Session{UserID: user.ID, Scope: database.RBACScopeAssigned, Permissions: map[string]bool{"chat:write": true}}) + + if !h.pathAllowed(c, filepath.ToSlash(filepath.Join("2026-07-10", allowed.ID, "a.txt"))) { + t.Fatal("assigned conversation attachment should be accessible") + } + if h.pathAllowed(c, filepath.ToSlash(filepath.Join("2026-07-10", hidden.ID, "secret.txt"))) { + t.Fatal("foreign conversation attachment should be denied") + } + if h.pathAllowed(c, "2026-07-10/_manual/secret.txt") { + t.Fatal("unowned manual attachment should fail closed") + } +} + +func TestPrepareMultiAgentSessionRejectsForeignConversation(t *testing.T) { + db, user := setupConversationRBACTest(t) + hidden, _ := db.CreateConversation("hidden", database.ConversationCreateMeta{}) + h := &AgentHandler{db: db, logger: zap.NewNop()} + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Set(security.ContextSessionKey, security.Session{UserID: user.ID, Scope: database.RBACScopeAssigned, Permissions: map[string]bool{"chat:write": true}}) + + _, err := h.prepareMultiAgentSession(&ChatRequest{ConversationID: hidden.ID, Message: "write"}, c, "test") + if err == nil || err.Error() != "无权访问该对话" { + t.Fatalf("err = %v, want unauthorized conversation", err) + } +} + +func TestMonitorExecutionDetailRejectsForeignOwner(t *testing.T) { + db, user := setupConversationRBACTest(t) + for _, exec := range []*mcp.ToolExecution{ + {ID: "exec-allowed", ToolName: "allowed", Status: "completed", StartTime: time.Now(), OwnerUserID: user.ID}, + {ID: "exec-hidden", ToolName: "hidden", Status: "completed", StartTime: time.Now(), OwnerUserID: "another-user"}, + } { + if err := db.SaveToolExecution(exec); err != nil { + t.Fatal(err) + } + } + h := NewMonitorHandler(mcp.NewServerWithStorage(zap.NewNop(), db), nil, db, zap.NewNop()) + + request := func(id string) *httptest.ResponseRecorder { + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest(http.MethodGet, "/api/monitor/execution/"+id, nil) + c.Params = gin.Params{{Key: "id", Value: id}} + c.Set(security.ContextSessionKey, security.Session{UserID: user.ID, Scope: database.RBACScopeAssigned}) + h.GetExecution(c) + return w + } + if w := request("exec-hidden"); w.Code != http.StatusForbidden { + t.Fatalf("hidden status = %d, want %d: %s", w.Code, http.StatusForbidden, w.Body.String()) + } + if w := request("exec-allowed"); w.Code != http.StatusOK { + t.Fatalf("allowed status = %d, want %d: %s", w.Code, http.StatusOK, w.Body.String()) + } +} + +func performAssignedHandler(user *database.RBACUser, method, path string, body interface{}, handler gin.HandlerFunc) *httptest.ResponseRecorder { + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + var req *http.Request + if body == nil { + req = httptest.NewRequest(method, path, nil) + } else { + payload, _ := json.Marshal(body) + req = httptest.NewRequest(method, path, bytes.NewReader(payload)) + req.Header.Set("Content-Type", "application/json") + } + c.Request = req + c.Set(security.ContextSessionKey, security.Session{UserID: user.ID, Scope: database.RBACScopeAssigned}) + handler(c) + return w +} diff --git a/internal/handler/rbac_context.go b/internal/handler/rbac_context.go new file mode 100644 index 00000000..f5b19c2e --- /dev/null +++ b/internal/handler/rbac_context.go @@ -0,0 +1,12 @@ +package handler + +import "context" + +// detachedAgentContext lets a long-running Agent survive an SSE disconnect +// while retaining immutable request values such as the authenticated Principal. +func detachedAgentContext(parent context.Context) context.Context { + if parent == nil { + parent = context.Background() + } + return context.WithoutCancel(parent) +} diff --git a/internal/handler/robot.go b/internal/handler/robot.go index db4c6dab..f67c495b 100644 --- a/internal/handler/robot.go +++ b/internal/handler/robot.go @@ -7,6 +7,7 @@ import ( "crypto/cipher" "crypto/rand" "crypto/sha1" + "crypto/sha256" "encoding/base64" "encoding/binary" "encoding/json" @@ -16,6 +17,7 @@ import ( "io" "net/http" "sort" + "strconv" "strings" "sync" "time" @@ -28,24 +30,24 @@ import ( ) const ( - robotCmdHelp = "帮助" - robotCmdList = "列表" - robotCmdListAlt = "对话列表" - robotCmdSwitch = "切换" - robotCmdContinue = "继续" - robotCmdNew = "新对话" - robotCmdClear = "清空" - robotCmdCurrent = "当前" - robotCmdStop = "停止" - robotCmdRoles = "角色" - robotCmdRolesList = "角色列表" - robotCmdSwitchRole = "切换角色" - robotCmdDelete = "删除" - robotCmdVersion = "版本" - robotCmdProjects = "项目" - robotCmdProjectsList = "项目列表" - robotCmdBindProject = "绑定项目" - robotCmdNewProject = "新建项目" + robotCmdHelp = "帮助" + robotCmdList = "列表" + robotCmdListAlt = "对话列表" + robotCmdSwitch = "切换" + robotCmdContinue = "继续" + robotCmdNew = "新对话" + robotCmdClear = "清空" + robotCmdCurrent = "当前" + robotCmdStop = "停止" + robotCmdRoles = "角色" + robotCmdRolesList = "角色列表" + robotCmdSwitchRole = "切换角色" + robotCmdDelete = "删除" + robotCmdVersion = "版本" + robotCmdProjects = "项目" + robotCmdProjectsList = "项目列表" + robotCmdBindProject = "绑定项目" + robotCmdNewProject = "新建项目" robotCmdUnbindProject = "解除项目" ) @@ -60,6 +62,7 @@ type RobotHandler struct { sessionRoles map[string]string // key: "platform_userID", value: roleName（默认"默认"） cancelMu sync.Mutex // 保护 runningCancels runningCancels map[string]context.CancelFunc // key: "platform_userID", 用于停止命令中断任务 + wecomReplay map[string]time.Time } // NewRobotHandler 创建机器人处理器 @@ -72,14 +75,51 @@ func NewRobotHandler(cfg *config.Config, db *database.DB, agentHandler *AgentHan sessions: make(map[string]string), sessionRoles: make(map[string]string), runningCancels: make(map[string]context.CancelFunc), + wecomReplay: make(map[string]time.Time), } } +func (h *RobotHandler) acceptFreshWecomRequest(timestamp, nonce, signature string) bool { + unixSeconds, err := strconv.ParseInt(strings.TrimSpace(timestamp), 10, 64) + if err != nil { + return false + } + now := time.Now() + requestTime := time.Unix(unixSeconds, 0) + if requestTime.Before(now.Add(-5*time.Minute)) || requestTime.After(now.Add(5*time.Minute)) { + return false + } + key := strings.TrimSpace(timestamp) + "\x00" + strings.TrimSpace(nonce) + "\x00" + strings.TrimSpace(signature) + if strings.TrimSpace(nonce) == "" || strings.TrimSpace(signature) == "" { + return false + } + h.mu.Lock() + defer h.mu.Unlock() + for replayKey, seenAt := range h.wecomReplay { + if now.Sub(seenAt) > 10*time.Minute { + delete(h.wecomReplay, replayKey) + } + } + if _, exists := h.wecomReplay[key]; exists { + return false + } + h.wecomReplay[key] = now + return true +} + // sessionKey 生成会话 key func (h *RobotHandler) sessionKey(platform, userID string) string { return platform + "_" + userID } +// robotOwnerID creates a stable, non-reversible local owner identity for one +// platform user. Robot users are isolated from each other and from Web users +// even though they authenticate through a shared platform callback secret. +func robotOwnerID(platform, userID string) string { + sum := sha256.Sum256([]byte(strings.TrimSpace(platform) + "\x00" + strings.TrimSpace(userID))) + return fmt.Sprintf("robot:%x", sum[:16]) +} + func (h *RobotHandler) loadSessionBinding(sk string) (convID, role string) { if h.db == nil || strings.TrimSpace(sk) == "" { return "", "" @@ -119,18 +159,23 @@ func (h *RobotHandler) getOrCreateConversation(platform, userID, title string) ( h.mu.RLock() convID = h.sessions[sk] h.mu.RUnlock() - if convID != "" { + ownerID := robotOwnerID(platform, userID) + if convID != "" && h.db.UserCanAccessResource(ownerID, database.RBACScopeOwn, "conversation", convID) { return convID, false } if persistedConvID, persistedRole := h.loadSessionBinding(sk); strings.TrimSpace(persistedConvID) != "" { - // 会话绑定持久化：服务重启后也可恢复当前对话和角色。 - h.mu.Lock() - h.sessions[sk] = persistedConvID - if strings.TrimSpace(persistedRole) != "" { - h.sessionRoles[sk] = persistedRole + if !h.db.UserCanAccessResource(ownerID, database.RBACScopeOwn, "conversation", persistedConvID) { + h.deleteSessionBinding(sk) + } else { + // 会话绑定持久化：服务重启后也可恢复当前对话和角色。 + h.mu.Lock() + h.sessions[sk] = persistedConvID + if strings.TrimSpace(persistedRole) != "" { + h.sessionRoles[sk] = persistedRole + } + h.mu.Unlock() + return persistedConvID, false } - h.mu.Unlock() - return persistedConvID, false } t := strings.TrimSpace(title) if t == "" { @@ -140,12 +185,16 @@ func (h *RobotHandler) getOrCreateConversation(platform, userID, title string) ( } meta := database.ConversationCreateMeta{Source: "robot:" + platform} meta.ProjectID = effectiveProjectID(h.config, "") + if meta.ProjectID != "" && !h.db.UserCanAccessResource(ownerID, database.RBACScopeOwn, "project", meta.ProjectID) { + meta.ProjectID = "" + } conv, err := h.db.CreateConversation(t, meta) if err != nil { h.logger.Warn("创建机器人会话失败", zap.Error(err)) return "", false } convID = conv.ID + _ = h.db.SetResourceOwner("conversation", convID, ownerID) h.mu.Lock() role := h.sessionRoles[sk] h.sessions[sk] = convID @@ -197,11 +246,16 @@ func (h *RobotHandler) clearConversation(platform, userID string) (newConvID str title := "新对话 " + time.Now().Format("01-02 15:04") meta := database.ConversationCreateMeta{Source: "robot:" + platform + ":new"} meta.ProjectID = effectiveProjectID(h.config, "") + ownerID := robotOwnerID(platform, userID) + if meta.ProjectID != "" && !h.db.UserCanAccessResource(ownerID, database.RBACScopeOwn, "project", meta.ProjectID) { + meta.ProjectID = "" + } conv, err := h.db.CreateConversation(title, meta) if err != nil { h.logger.Warn("创建新对话失败", zap.Error(err)) return "" } + _ = h.db.SetResourceOwner("conversation", conv.ID, ownerID) h.setConversation(platform, userID, conv.ID) return conv.ID } @@ -251,7 +305,7 @@ func (h *RobotHandler) HandleMessage(platform, userID, text string) (reply strin h.cancelMu.Unlock() }() role := h.getRole(platform, userID) - resp, newConvID, err := h.agentHandler.ProcessMessageForRobot(ctx, platform, convID, text, role) + resp, newConvID, err := h.agentHandler.ProcessMessageForRobot(ctx, platform, robotOwnerID(platform, userID), convID, text, role) if err != nil { h.logger.Warn("机器人 Agent 执行失败", zap.String("platform", platform), zap.String("userID", userID), zap.Error(err)) if errors.Is(err, context.Canceled) { @@ -306,15 +360,19 @@ func (h *RobotHandler) projectsEnabled() bool { return h.config != nil && h.config.Project.Enabled } -func (h *RobotHandler) resolveProjectByIDOrName(idOrName string) (*database.Project, string) { +func (h *RobotHandler) resolveProjectByIDOrName(platform, userID, idOrName string) (*database.Project, string) { idOrName = strings.TrimSpace(idOrName) if idOrName == "" { return nil, "请指定项目 ID 或名称，例如：绑定项目 xxx-xxx" } + ownerID := robotOwnerID(platform, userID) if p, err := h.db.GetProject(idOrName); err == nil { - return p, "" + if h.db.UserCanAccessResource(ownerID, database.RBACScopeOwn, "project", p.ID) { + return p, "" + } + return nil, "项目不存在或无权访问。" } - list, err := h.db.ListProjects("", "", 200, 0) + list, err := h.db.ListProjectsForAccess("", "", 200, 0, ownerID, database.RBACScopeOwn) if err != nil { return nil, "查询项目失败: " + err.Error() } @@ -349,11 +407,11 @@ func (h *RobotHandler) formatProjectLabel(projectID string) string { return projectID } -func (h *RobotHandler) cmdProjects() string { +func (h *RobotHandler) cmdProjects(platform, userID string) string { if !h.projectsEnabled() { return "项目功能未启用（config.project.enabled）。" } - list, err := h.db.ListProjects("", "", 50, 0) + list, err := h.db.ListProjectsForAccess("", "", 50, 0, robotOwnerID(platform, userID), database.RBACScopeOwn) if err != nil { return "获取项目列表失败: " + err.Error() } @@ -380,7 +438,7 @@ func (h *RobotHandler) cmdBindProject(platform, userID, idOrName string) string if !h.projectsEnabled() { return "项目功能未启用（config.project.enabled）。" } - p, errMsg := h.resolveProjectByIDOrName(idOrName) + p, errMsg := h.resolveProjectByIDOrName(platform, userID, idOrName) if p == nil { return errMsg } @@ -407,6 +465,7 @@ func (h *RobotHandler) cmdNewProject(platform, userID, name string) string { if err != nil { return "创建项目失败: " + err.Error() } + _ = h.db.SetResourceOwner("project", created.ID, robotOwnerID(platform, userID)) convID, _ := h.getOrCreateConversation(platform, userID, name) if convID == "" { return fmt.Sprintf("项目已创建：「%s」\nID: %s\n（绑定当前对话失败，请手动发送「绑定项目 %s」）", created.Name, created.ID, created.ID) @@ -430,6 +489,9 @@ func (h *RobotHandler) cmdUnbindProject(platform, userID string) string { convID = persistedConvID } } + if !h.db.UserCanAccessResource(robotOwnerID(platform, userID), database.RBACScopeOwn, "conversation", convID) { + return "当前对话不存在或无权访问。" + } if convID == "" { return "当前没有进行中的对话，无需解除绑定。" } @@ -446,8 +508,8 @@ func (h *RobotHandler) cmdUnbindProject(platform, userID string) string { return "已解除当前对话的项目绑定。" } -func (h *RobotHandler) cmdList() string { - convs, err := h.db.ListConversations(50, 0, "", "", "") +func (h *RobotHandler) cmdList(platform, userID string) string { + convs, err := h.db.ListConversationsForAccess(50, 0, "", "", "", robotOwnerID(platform, userID), database.RBACScopeOwn) if err != nil { return "获取对话列表失败: " + err.Error() } @@ -471,7 +533,7 @@ func (h *RobotHandler) cmdSwitch(platform, userID, convID string) string { return "请指定对话 ID，例如：切换 xxx-xxx-xxx" } conv, err := h.db.GetConversation(convID) - if err != nil { + if err != nil || !h.db.UserCanAccessResource(robotOwnerID(platform, userID), database.RBACScopeOwn, "conversation", convID) { return "对话不存在或 ID 错误。" } h.setConversation(platform, userID, conv.ID) @@ -512,6 +574,9 @@ func (h *RobotHandler) cmdCurrent(platform, userID string) string { if convID == "" { return "当前没有进行中的对话。发送任意内容将创建新对话。" } + if !h.db.UserCanAccessResource(robotOwnerID(platform, userID), database.RBACScopeOwn, "conversation", convID) { + return "当前对话不存在或无权访问。" + } conv, err := h.db.GetConversation(convID) if err != nil { return "当前对话 ID: " + convID + "（获取标题失败）" @@ -582,6 +647,9 @@ func (h *RobotHandler) cmdDelete(platform, userID, convID string) string { if convID == "" { return "请指定对话 ID，例如：删除 xxx-xxx-xxx" } + if !h.db.UserCanAccessResource(robotOwnerID(platform, userID), database.RBACScopeOwn, "conversation", convID) { + return "对话不存在或无权访问。" + } sk := h.sessionKey(platform, userID) h.mu.RLock() currentConvID := h.sessions[sk] @@ -617,7 +685,7 @@ func (h *RobotHandler) handleRobotCommand(platform, userID, text string) (string case text == robotCmdHelp || text == "help" || text == "？" || text == "?": return h.cmdHelp(), true case text == robotCmdList || text == robotCmdListAlt || text == "list": - return h.cmdList(), true + return h.cmdList(platform, userID), true case strings.HasPrefix(text, robotCmdSwitch+" ") || strings.HasPrefix(text, robotCmdContinue+" ") || strings.HasPrefix(text, "switch ") || strings.HasPrefix(text, "continue "): var id string switch { @@ -663,7 +731,7 @@ func (h *RobotHandler) handleRobotCommand(platform, userID, text string) (string case text == robotCmdVersion || text == "version": return h.cmdVersion(), true case text == robotCmdProjects || text == robotCmdProjectsList || text == "projects": - return h.cmdProjects(), true + return h.cmdProjects(platform, userID), true case text == robotCmdUnbindProject || text == "unbind project": return h.cmdUnbindProject(platform, userID), true case strings.HasPrefix(text, robotCmdNewProject+" ") || strings.HasPrefix(text, "new project "): @@ -903,6 +971,11 @@ func (h *RobotHandler) HandleWecomPOST(c *gin.Context) { c.String(http.StatusOK, "") return } + if !h.acceptFreshWecomRequest(timestamp, nonce, msgSignature) { + h.logger.Warn("企业微信 POST 时间戳过期或请求重放，已拒绝") + c.String(http.StatusOK, "") + return + } var body wecomXML if err := xml.Unmarshal(bodyRaw, &body); err != nil { diff --git a/internal/handler/robot_rbac_test.go b/internal/handler/robot_rbac_test.go new file mode 100644 index 00000000..88e88e54 --- /dev/null +++ b/internal/handler/robot_rbac_test.go @@ -0,0 +1,64 @@ +package handler + +import ( + "fmt" + "path/filepath" + "strings" + "testing" + "time" + + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/database" + + "go.uber.org/zap" +) + +func TestRobotUsersAreResourceIsolated(t *testing.T) { + db, err := database.NewDB(filepath.Join(t.TempDir(), "robot-rbac.db"), zap.NewNop()) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = db.Close() }) + cfg := &config.Config{} + cfg.Project.Enabled = true + h := NewRobotHandler(cfg, db, nil, zap.NewNop()) + + conversationID, _ := h.getOrCreateConversation("wecom", "alice", "alice conversation") + if conversationID == "" { + t.Fatal("alice conversation was not created") + } + if got := h.cmdList("wecom", "bob"); strings.Contains(got, conversationID) { + t.Fatalf("bob listed alice conversation: %s", got) + } + if got := h.cmdSwitch("wecom", "bob", conversationID); !strings.Contains(got, "不存在") && !strings.Contains(got, "无权访问") { + t.Fatalf("bob switched to alice conversation: %s", got) + } + if got := h.cmdDelete("wecom", "bob", conversationID); !strings.Contains(got, "无权访问") { + t.Fatalf("bob deleted alice conversation: %s", got) + } + if _, err := db.GetConversation(conversationID); err != nil { + t.Fatalf("alice conversation was deleted: %v", err) + } + + createReply := h.cmdNewProject("wecom", "alice", "alice project") + if !strings.Contains(createReply, "已创建项目") { + t.Fatalf("create project reply: %s", createReply) + } + if got := h.cmdProjects("wecom", "bob"); strings.Contains(got, "alice project") { + t.Fatalf("bob listed alice project: %s", got) + } +} + +func TestWecomReplayGuardRequiresFreshUniqueRequest(t *testing.T) { + h := NewRobotHandler(&config.Config{}, nil, nil, zap.NewNop()) + timestamp := time.Now().Unix() + if !h.acceptFreshWecomRequest(fmt.Sprintf("%d", timestamp), "nonce", "signature") { + t.Fatal("fresh request was rejected") + } + if h.acceptFreshWecomRequest(fmt.Sprintf("%d", timestamp), "nonce", "signature") { + t.Fatal("duplicate request was accepted") + } + if h.acceptFreshWecomRequest(fmt.Sprintf("%d", timestamp-600), "old", "signature") { + t.Fatal("stale request was accepted") + } +} diff --git a/internal/handler/vulnerability.go b/internal/handler/vulnerability.go index 089748a3..d52babf3 100644 --- a/internal/handler/vulnerability.go +++ b/internal/handler/vulnerability.go @@ -254,7 +254,15 @@ func (h *VulnerabilityHandler) UpdateVulnerability(c *gin.Context) { // 更新字段 if req.ProjectID != nil { - existing.ProjectID = strings.TrimSpace(*req.ProjectID) + targetProjectID := strings.TrimSpace(*req.ProjectID) + if targetProjectID != "" { + session, ok := security.CurrentSession(c) + if !ok || !h.db.UserCanAccessResource(session.UserID, session.Scope, "project", targetProjectID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权关联目标项目"}) + return + } + } + existing.ProjectID = targetProjectID } if req.ConversationTag != nil { existing.ConversationTag = *req.ConversationTag diff --git a/internal/handler/webshell.go b/internal/handler/webshell.go index 86113dcc..5b34c368 100644 --- a/internal/handler/webshell.go +++ b/internal/handler/webshell.go @@ -721,10 +721,15 @@ func (h *WebShellHandler) Exec(c *gin.Context) { c.JSON(http.StatusBadRequest, gin.H{"error": "url and command are required"}) return } - if !h.webshellConnectionRequestAllowed(c, req.ConnectionID, req.URL) { + conn, allowed := h.authorizedWebshellConnection(c, req.ConnectionID, req.URL) + if !allowed { c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) return } + // The database record is authoritative. Never let a caller pair an + // authorized ID with attacker-controlled transport credentials or a URL. + req.URL, req.Password, req.Type = conn.URL, conn.Password, conn.Type + req.Method, req.CmdParam, req.Encoding = conn.Method, conn.CmdParam, conn.Encoding parsed, err := url.Parse(req.URL) if err != nil || (parsed.Scheme != "http" && parsed.Scheme != "https") { @@ -818,10 +823,13 @@ func (h *WebShellHandler) FileOp(c *gin.Context) { c.JSON(http.StatusBadRequest, gin.H{"error": "url and action are required"}) return } - if !h.webshellConnectionRequestAllowed(c, req.ConnectionID, req.URL) { + conn, allowed := h.authorizedWebshellConnection(c, req.ConnectionID, req.URL) + if !allowed { c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) return } + req.URL, req.Password, req.Type = conn.URL, conn.Password, conn.Type + req.Method, req.CmdParam, req.Encoding, req.OS = conn.Method, conn.CmdParam, conn.Encoding, conn.OS parsed, err := url.Parse(req.URL) if err != nil || (parsed.Scheme != "http" && parsed.Scheme != "https") { @@ -898,23 +906,26 @@ func (h *WebShellHandler) FileOp(c *gin.Context) { }) } -func (h *WebShellHandler) webshellConnectionRequestAllowed(c *gin.Context, connectionID, requestURL string) bool { +func (h *WebShellHandler) authorizedWebshellConnection(c *gin.Context, connectionID, requestURL string) (*database.WebShellConnection, bool) { connectionID = strings.TrimSpace(connectionID) if connectionID == "" { - return true + return nil, false } if h.db == nil { - return false + return nil, false } session, ok := security.CurrentSession(c) if !ok || !h.db.UserCanAccessResource(session.UserID, session.Scope, "webshell", connectionID) { - return false + return nil, false } conn, err := h.db.GetWebshellConnection(connectionID) if err != nil || conn == nil { - return false + return nil, false } - return strings.TrimSpace(conn.URL) == strings.TrimSpace(requestURL) + if requestURL = strings.TrimSpace(requestURL); requestURL != "" && strings.TrimSpace(conn.URL) != requestURL { + return nil, false + } + return conn, true } // ExecWithConnection 在指定 WebShell 连接上执行命令（供 MCP/Agent 等非 HTTP 调用） diff --git a/internal/handler/webshell_rbac_test.go b/internal/handler/webshell_rbac_test.go index 68d17e63..cea5676b 100644 --- a/internal/handler/webshell_rbac_test.go +++ b/internal/handler/webshell_rbac_test.go @@ -40,6 +40,18 @@ func TestWebshellExecRequiresConnectionAccessWhenConnectionIDProvided(t *testing } } +func TestWebshellExecRejectsAdHocURLWithoutConnectionID(t *testing.T) { + gin.SetMode(gin.TestMode) + _, user, _, _ := setupWebshellRBACTest(t) + handler := NewWebShellHandler(zap.NewNop(), nil) + w := performWebshellJSON(user, http.MethodPost, "/api/webshell/exec", map[string]interface{}{ + "url": "http://127.0.0.1/admin", "command": "id", + }, handler.Exec) + if w.Code != http.StatusForbidden { + t.Fatalf("ad-hoc URL status = %d, want %d: %s", w.Code, http.StatusForbidden, w.Body.String()) + } +} + func TestWebshellFileOpRequiresConnectionAccessWhenConnectionIDProvided(t *testing.T) { gin.SetMode(gin.TestMode) db, user, _, hidden := setupWebshellRBACTest(t) diff --git a/internal/handler/workflow_integration.go b/internal/handler/workflow_integration.go index dee8bd35..3acf5d73 100644 --- a/internal/handler/workflow_integration.go +++ b/internal/handler/workflow_integration.go @@ -36,6 +36,7 @@ func (h *AgentHandler) roleForWorkflow(req *ChatRequest) (config.RoleConfig, boo } func (h *AgentHandler) runRoleWorkflowStreamIfBound( + c *gin.Context, req *ChatRequest, prep *multiAgentPrepared, sendEvent func(eventType, message string, data interface{}), @@ -60,7 +61,10 @@ func (h *AgentHandler) runRoleWorkflowStreamIfBound( } }() - baseCtx, cancelWithCause := context.WithCancelCause(context.Background()) + if c == nil || c.Request == nil { + return false + } + baseCtx, cancelWithCause := context.WithCancelCause(detachedAgentContext(c.Request.Context())) defer cancelWithCause(nil) taskCtx, timeoutCancel := context.WithTimeout(baseCtx, 600*time.Minute) defer timeoutCancel() diff --git a/internal/handler/workflow_run.go b/internal/handler/workflow_run.go index f9e4e6d7..1b18f05a 100644 --- a/internal/handler/workflow_run.go +++ b/internal/handler/workflow_run.go @@ -8,6 +8,8 @@ import ( "cyberstrike-ai/internal/agent" "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/security" workflowrunner "cyberstrike-ai/internal/workflow" "github.com/gin-gonic/gin" @@ -20,6 +22,10 @@ func (h *WorkflowHandler) SetRuntime(agent *agent.Agent, cfg *config.Config) { func (h *WorkflowHandler) GetRun(c *gin.Context) { runID := strings.TrimSpace(c.Param("runId")) + if !h.workflowRunAllowed(c, runID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } run, err := h.db.GetWorkflowRun(runID) if err != nil { c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) @@ -39,6 +45,10 @@ func (h *WorkflowHandler) GetRun(c *gin.Context) { func (h *WorkflowHandler) ReplayRun(c *gin.Context) { runID := strings.TrimSpace(c.Param("runId")) + if !h.workflowRunAllowed(c, runID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } nodeRuns, err := h.db.ListWorkflowNodeRuns(runID) if err != nil { c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) @@ -67,11 +77,18 @@ func (h *WorkflowHandler) ReplayRun(c *gin.Context) { func (h *WorkflowHandler) ListPendingRuns(c *gin.Context) { conversationID := strings.TrimSpace(c.Query("conversationId")) + if conversationID != "" && !h.workflowConversationAllowed(c, conversationID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } runs, err := h.db.ListWorkflowRunsAwaitingHITLFiltered(conversationID, 50) if err != nil { c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) return } + runs = filterSlice(runs, func(run *database.WorkflowRun) bool { + return run != nil && h.workflowConversationAllowed(c, run.ConversationID) + }) c.JSON(http.StatusOK, gin.H{"runs": runs}) } @@ -86,6 +103,10 @@ func (h *WorkflowHandler) ResumeRun(c *gin.Context) { return } runID := strings.TrimSpace(c.Param("runId")) + if !h.workflowRunAllowed(c, runID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } var req workflowResumeRequest if err := c.ShouldBindJSON(&req); err != nil { c.JSON(http.StatusBadRequest, gin.H{"error": "无效的请求参数: " + err.Error()}) @@ -160,3 +181,20 @@ func (h *WorkflowHandler) ResumeRun(c *gin.Context) { "awaitingHitl": result.AwaitingHITL, }) } + +func (h *WorkflowHandler) workflowConversationAllowed(c *gin.Context, conversationID string) bool { + session, ok := security.CurrentSession(c) + return ok && h.db.UserCanAccessResource(session.UserID, session.Scope, "conversation", strings.TrimSpace(conversationID)) +} + +func (h *WorkflowHandler) workflowRunAllowed(c *gin.Context, runID string) bool { + session, ok := security.CurrentSession(c) + if !ok { + return false + } + if session.Scope == database.RBACScopeAll { + return true + } + run, err := h.db.GetWorkflowRun(strings.TrimSpace(runID)) + return err == nil && run != nil && h.workflowConversationAllowed(c, run.ConversationID) +}