Add files via upload

2026-07-14 07:57:26 +02:00 · 2026-07-10 18:54:04 +08:00
parent a145687508
commit 1b64f5d8a0
28 changed files with 1040 additions and 125 deletions
@@ -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
+}