xinmi/internal/accesscontrol/middleware_test.go

package accesscontrol

import (
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/Tencent/WeKnora/internal/types"
	"github.com/gin-gonic/gin"
)

func TestFeatureRoutes(t *testing.T) {
	for _, tc := range []struct{ method, path, kind, want string }{
		{"GET", "/api/v1/workbench/resources/:kind", "training-maps", "content.read"},
		{"POST", "/api/v1/workbench/resources/:kind/:id/actions/:action", "content-tasks", "content.write"},
		{"DELETE", "/api/v1/workbench/resources/:kind/:id", "flows", "flows.delete"},
		{"POST", "/api/v1/agent-chat/:session_id", "", "chat.write"},
		{"POST", "/api/v1/knowledge-chat/:session_id", "", "chat.write"},
		{"POST", "/api/v1/agent/tool-approvals/:pending_id", "", "chat.write"},
		{"POST", "/api/v1/agent/mcp-oauth-resolutions/:pending_id", "", "chat.write"},
		{"POST", "/api/v1/agent/mcp-oauth-resolutions/:pending_id/cancel", "", "chat.write"},
		{"POST", "/api/v1/knowledge-search", "", "knowledge.read"},
		{"GET", "/api/v1/initialization/config/:kbId", "", "knowledge.read"},
		{"PUT", "/api/v1/initialization/config/:kbId", "", "knowledge.write"},
		{"POST", "/api/v1/initialization/initialize/:kbId", "", "knowledge.write"},
		{"POST", "/api/v1/skills/instructions/mount", "", "agents.skills"},
		{"GET", "/api/v1/memory/items", "", "agents.read"},
		{"GET", "/api/v1/shared-agents", "", "agents.read"},
		{"GET", "/api/v1/shared-knowledge-bases", "", "knowledge.read"},
		{"GET", "/api/v1/embed-channels", "", "agents.read"},
		{"POST", "/api/v1/wechat/qrcode", "", "digital-humans.wechat"},
		{"GET", "/api/v1/models", "", ""},
		{"POST", "/api/v1/models", "", "settings.write"},
		{"GET", "/api/v1/evaluation", "", "settings.read"},
		{"POST", "/api/v1/evaluation", "", "settings.write"},
	} {
		t.Run(tc.path+tc.method, func(t *testing.T) {
			if got := requiredPermission(tc.method, tc.path, tc.kind); got != tc.want {
				t.Fatalf("got %q want %q", got, tc.want)
			}
		})
	}
}

func TestMiddlewareReloadsPermissionsBeforeAlternateNativeActions(t *testing.T) {
	for _, tc := range []struct {
		method, route, path, permission string
	}{
		{http.MethodGet, "/initialization/config/:kbId", "/initialization/config/owned-kb", "knowledge.read"},
		{http.MethodPut, "/initialization/config/:kbId", "/initialization/config/owned-kb", "knowledge.write"},
		{http.MethodPost, "/initialization/initialize/:kbId", "/initialization/initialize/owned-kb", "knowledge.write"},
		{http.MethodPost, "/agent/tool-approvals/:pending_id", "/agent/tool-approvals/owned-pending", "chat.write"},
		{http.MethodPost, "/agent/mcp-oauth-resolutions/:pending_id", "/agent/mcp-oauth-resolutions/owned-pending", "chat.write"},
		{http.MethodPost, "/agent/mcp-oauth-resolutions/:pending_id/cancel", "/agent/mcp-oauth-resolutions/owned-pending/cancel", "chat.write"},
	} {
		t.Run(tc.method+tc.route, func(t *testing.T) {
			s := versionDB(t)
			role := Role{ID: "native-role", TenantID: 1, Name: "native", Permissions: types.StringArray{"knowledge.read", "knowledge.write", "chat.read", "chat.write", "settings.read"}}
			if err := s.db.Create(&role).Error; err != nil {
				t.Fatal(err)
			}
			if err := s.db.Create(&Assignment{TenantID: 1, UserID: "alice", RoleID: role.ID}).Error; err != nil {
				t.Fatal(err)
			}
			called := 0
			router := gin.New()
			router.Use(s.Middleware())
			router.Handle(tc.method, "/api/v1"+tc.route, func(c *gin.Context) {
				called++
				c.Status(http.StatusNoContent)
			})
			request := func() int {
				req := httptest.NewRequest(tc.method, "/api/v1"+tc.path, nil)
				req = req.WithContext(versionContext(1, "alice", false))
				res := httptest.NewRecorder()
				router.ServeHTTP(res, req)
				return res.Code
			}
			if status := request(); status != http.StatusNoContent || called != 1 {
				t.Fatalf("authorized route blocked: status=%d calls=%d", status, called)
			}
			retained := types.StringArray{}
			for _, permission := range role.Permissions {
				if permission != tc.permission {
					retained = append(retained, permission)
				}
			}
			if err := s.db.Model(&role).Update("permissions", retained).Error; err != nil {
				t.Fatal(err)
			}
			if status := request(); status != http.StatusForbidden || called != 1 {
				t.Fatalf("revoked permission reached native handler: status=%d calls=%d", status, called)
			}
		})
	}
}

func TestPermissionValidation(t *testing.T) {
	if validPermissions([]string{"content.write"}) {
		t.Fatal("write without read accepted")
	}
	if validPermissions([]string{"*"}) {
		t.Fatal("unknown wildcard accepted")
	}
	if !validPermissions([]string{"content.read", "content.write"}) {
		t.Fatal("valid feature permissions rejected")
	}
	if !validPermissions(nil) {
		t.Fatal("deny-all role rejected")
	}
}

func TestContentActionPermissions(t *testing.T) {
	for _, action := range []string{"accept", "skip", "review", "review-settings", "compose"} {
		for _, path := range []string{"/api/v1/workbench/resources/:kind/:id/actions/:action", "/api/v1/workbench/jobs/:id/:action"} {
			if got := requiredPermission("POST", path, "content-tasks", action); got != "content.review" {
				t.Fatalf("%s %s: %s", path, action, got)
			}
		}
	}
	for _, action := range []string{"generate", "revise", "decompose", "import-image"} {
		if got := requiredPermission("POST", "/api/v1/workbench/jobs/:id/:action", "content-tasks", action); got != "content.generate" {
			t.Fatalf("production action %s: %s", action, got)
		}
	}
	if got := requiredPermission("GET", "/api/v1/workbench/resources/:kind/:id/export", "content-tasks"); got != "content.deliver" {
		t.Fatalf("export permission: %s", got)
	}
	for _, permission := range []string{"content.review", "content.deliver"} {
		if validPermissions([]string{permission}) || !validPermissions([]string{"content.read", permission}) {
			t.Fatalf("read prerequisite for %s", permission)
		}
	}
}