xinmi/internal/accesscontrol/middleware.go

package accesscontrol

import (
	"net/http"
	"strings"

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

func requiredPermission(method, path, kind string, actions ...string) string {
	path = strings.TrimPrefix(path, "/api/v1/")
	mutating := method != http.MethodGet && method != http.MethodHead
	if strings.HasPrefix(path, "workbench/resources/") && strings.HasSuffix(path, "/export") {
		return "content.deliver"
	}
	if method != http.MethodGet && method != http.MethodHead && len(actions) > 0 &&
		(strings.Contains(path, "/actions/") || strings.HasPrefix(path, "workbench/jobs/")) {
		return accesspolicy.WorkbenchActionPermission(kind, actions[0])
	}
	module := ""
	switch {
	case strings.HasPrefix(path, "access/"):
		return ""
	case strings.HasPrefix(path, "workbench/resources/"):
		if method == http.MethodDelete {
			if kind == "flows" {
				return "flows.delete"
			}
			return "content.delete"
		}
		if mutating && kind == "training-maps" {
			return "content.maps"
		}
		if kind == "flows" {
			module = "flows"
		} else {
			module = "content"
		}
	case strings.HasPrefix(path, "workbench/jobs/"):
		if kind == "flows" {
			module = "flows"
		} else {
			module = "content"
		}
	case strings.HasPrefix(path, "digital-humans"), strings.HasPrefix(path, "im/"), strings.HasPrefix(path, "im-channels"), strings.HasPrefix(path, "wechat"):
		if strings.Contains(path, "/wechat/") || strings.HasPrefix(path, "wechat") || strings.HasPrefix(path, "im/") || strings.HasPrefix(path, "im-channels") {
			return "digital-humans.wechat"
		}
		if mutating && (strings.HasSuffix(path, "/distill") || strings.HasSuffix(path, "/trial")) {
			return "digital-humans.distill"
		}
		if mutating && strings.HasSuffix(path, "/review") {
			return "digital-humans.review"
		}
		if method == http.MethodDelete && !strings.Contains(path, "/bindings/") {
			return "digital-humans.delete"
		}
		module = "digital-humans"
	case path == "workbench/operations/consumption":
		return "knowledge.consumption"
	case strings.HasPrefix(path, "workbench/operations/memories/"):
		return "agents.memory"
	case strings.HasPrefix(path, "workbench/operations/accounts"), strings.HasPrefix(path, "workbench/operations/grants"), strings.HasPrefix(path, "workbench/operations/quotas"):
		module = "balance"
	case strings.HasPrefix(path, "knowledge-chat"), strings.HasPrefix(path, "agent-chat"), strings.HasPrefix(path, "agent/"):
		module = "chat"
	case strings.HasPrefix(path, "knowledge"), strings.HasPrefix(path, "shared-knowledge-bases"), strings.HasPrefix(path, "chunks"), strings.HasPrefix(path, "tags"), strings.HasPrefix(path, "faq"), strings.HasPrefix(path, "initialization/config/"), strings.HasPrefix(path, "initialization/initialize/"):
		module = "knowledge"
	case strings.HasPrefix(path, "agents"), strings.HasPrefix(path, "shared-agents"), strings.HasPrefix(path, "embed-channels"), strings.HasPrefix(path, "custom-agents"), strings.HasPrefix(path, "skills"), strings.HasPrefix(path, "instruction-skills"), strings.HasPrefix(path, "memory"):
		if mutating {
			if strings.HasPrefix(path, "skills/instructions") || strings.HasPrefix(path, "instruction-skills") {
				return "agents.skills"
			}
			if strings.HasPrefix(path, "memory") {
				return "agents.memory"
			}
			if method == http.MethodDelete && path == "agents/:id" {
				return "agents.delete"
			}
		}
		module = "agents"
	case strings.HasPrefix(path, "sessions"), strings.HasPrefix(path, "messages"), strings.HasPrefix(path, "chat"):
		module = "chat"
	case strings.HasPrefix(path, "organizations"):
		module = "sharing"
	case path == "evaluation", strings.HasPrefix(path, "tenants/kv/") && mutating, strings.HasPrefix(path, "tenants/") && (strings.Contains(path, "/members") || strings.Contains(path, "/invitations")), strings.HasPrefix(path, "me/env-vars"):
		module = "settings"
	case strings.HasPrefix(path, "models"), strings.HasPrefix(path, "mcp-services"), strings.HasPrefix(path, "sandbox-configs"), strings.HasPrefix(path, "vector-stores"), strings.HasPrefix(path, "storage-backends"), strings.HasPrefix(path, "web-search-providers"), strings.HasPrefix(path, "datasource"):
		// Read-only selectors are used by content/Agent setup as well as settings.
		if method == http.MethodGet || method == http.MethodHead {
			return ""
		}
		module = "settings"
	}
	if module == "" {
		return ""
	}
	action := "read"
	if method != http.MethodGet && method != http.MethodHead && !strings.HasSuffix(path, "/search") && !strings.HasSuffix(path, "/hybrid-search") && path != "knowledge-search" {
		action = "write"
	}
	return module + "." + action
}

func (s *Service) Middleware() gin.HandlerFunc {
	return func(c *gin.Context) {
		t, u := identity(c.Request.Context())
		if t == 0 || u == "" || types.IsSyntheticUserID(u) {
			c.Next()
			return
		}
		path := c.FullPath()
		if strings.Contains(path, "/auth/") || path == "/api/v1/access/me" {
			c.Next()
			return
		}
		snapshot, err := s.Snapshot(c.Request.Context())
		if err != nil {
			fail(c, 503, "权限信息暂不可用，请稍后重试")
			return
		}
		ctx := accesspolicy.WithSnapshot(c.Request.Context(), snapshot)
		c.Request = c.Request.WithContext(ctx)
		kind := c.Param("kind")
		action := c.Param("action")
		if strings.Contains(path, "/workbench/jobs/") {
			var row struct{ Kind, Action string }
			if e := s.db.Table("workbench_jobs").Select("kind, action").Where("tenant_id=? AND owner_id=? AND id=?", t, u, c.Param("id")).Take(&row).Error; e == nil {
				kind = row.Kind
				action = row.Action
			}
		}
		permission := requiredPermission(c.Request.Method, path, kind, action)
		if strings.HasSuffix(path, "/operations/records/:recordKind") {
			switch c.Param("recordKind") {
			case "transactions":
				permission = "balance.read"
			case "usage":
				permission = "usage.read"
			case "audit":
				permission = "audit.read"
			case "system":
				permission = "system.read"
			case "consumption":
				permission = "knowledge.read"
			case "memories":
				permission = "agents.read"
			}
		}
		if permission != "" && ((!administrator(ctx) && adminOnlyPermission(permission)) || !accesspolicy.Allowed(ctx, permission)) {
			label := permission
			for _, m := range Catalog {
				for _, p := range m.Permissions {
					if p.ID == permission {
						label = m.Name + " / " + p.Name
					}
				}
			}
			fail(c, 403, "没有权限："+label)
			c.Set("required_permission_label", label)
			return
		}
		c.Next()
	}
}