easyai-ai-gateway/apps/api/internal/store/access_policy.go

package store

import (
	"context"
	"strings"

	"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
	"github.com/easyai/easyai-ai-gateway/apps/api/internal/modelaccess"
)

type APIKeyAccessRuleDiagnostic struct {
	RuleID       string `json:"ruleId"`
	ResourceType string `json:"resourceType"`
	ResourceID   string `json:"resourceId"`
	ResourceName string `json:"resourceName,omitempty"`
	Effect       string `json:"effect"`
	Effective    bool   `json:"effective"`
	Reason       string `json:"reason,omitempty" enums:"resource_unavailable,owner_access_revoked,scope_not_allowed"`
}

func (s *Store) enabledPlatformModels(ctx context.Context) ([]PlatformModel, []Platform, error) {
	models, err := s.ListModels(ctx)
	if err != nil {
		return nil, nil, err
	}
	platforms, err := s.ListPlatforms(ctx)
	if err != nil {
		return nil, nil, err
	}
	enabledPlatforms := map[string]bool{}
	for _, platform := range platforms {
		if platform.Status == "enabled" {
			enabledPlatforms[platform.ID] = true
		}
	}
	enabled := make([]PlatformModel, 0, len(models))
	for _, model := range models {
		if model.Enabled && enabledPlatforms[model.PlatformID] {
			enabled = append(enabled, model)
		}
	}
	return enabled, platforms, nil
}

func (s *Store) filterPlatformModelsByLayeredAccess(ctx context.Context, user *auth.User, models []PlatformModel) ([]PlatformModel, error) {
	rules, err := s.listActiveAccessRulesForResources(ctx, platformModelAccessResources(models))
	if err != nil {
		return nil, err
	}
	baselineRules, apiKeyRules := splitLayeredAccessRules(rules)
	baseline := filterPlatformModelsByRuleSet(models, baselineRules, baselineAccessRuleSubjects(user), permissionLevel(user))
	if user == nil || strings.TrimSpace(user.APIKeyID) == "" {
		return baseline, nil
	}
	apiKeyUsers, err := s.apiKeyAccessRuleUsers(ctx, apiKeyRules)
	if err != nil {
		return nil, err
	}
	return filterPlatformModelsByLayeredRuleSet(user, baseline, baselineRules, apiKeyRules, apiKeyUsers), nil
}

func filterPlatformModelsByLayeredRuleSet(user *auth.User, baseline []PlatformModel, baselineRules []AccessRule, apiKeyRules []AccessRule, apiKeyUsers map[string]*auth.User) []PlatformModel {
	if user == nil || strings.TrimSpace(user.APIKeyID) == "" {
		return baseline
	}
	keyFiltered := make([]PlatformModel, 0, len(baseline))
	for _, model := range baseline {
		effectiveRules := effectiveAPIKeyRulesForPlatformModel(apiKeyRules, baselineRules, apiKeyUsers, model)
		if platformModelAllowedByAccessRules(model, effectiveRules, apiKeyAccessRuleSubjects(user), permissionLevel(user)) {
			keyFiltered = append(keyFiltered, model)
		}
	}
	return filterPlatformModelsByAPIKeyScopes(keyFiltered, user.APIKeyScopes)
}

func (s *Store) filterPlatformModelsByBaselineAccess(ctx context.Context, user *auth.User, models []PlatformModel) ([]PlatformModel, error) {
	rules, err := s.listActiveAccessRulesForResources(ctx, platformModelAccessResources(models))
	if err != nil {
		return nil, err
	}
	baselineRules, _ := splitLayeredAccessRules(rules)
	return filterPlatformModelsByRuleSet(models, baselineRules, baselineAccessRuleSubjects(user), permissionLevel(user)), nil
}

func (s *Store) filterRuntimeCandidatesByLayeredAccess(ctx context.Context, user *auth.User, candidates []RuntimeModelCandidate) ([]RuntimeModelCandidate, error) {
	if len(candidates) == 0 {
		return candidates, nil
	}
	accessUser, err := s.resolveCurrentAccessUser(ctx, user)
	if err != nil {
		return nil, err
	}
	rules, err := s.listActiveAccessRulesForResources(ctx, candidateAccessResources(candidates))
	if err != nil {
		return nil, err
	}
	baselineRules, apiKeyRules := splitLayeredAccessRules(rules)
	baseline := filterCandidatesByRuleSet(candidates, baselineRules, baselineAccessRuleSubjects(accessUser), permissionLevel(accessUser))
	if accessUser == nil || strings.TrimSpace(accessUser.APIKeyID) == "" {
		return baseline, nil
	}
	apiKeyUsers, err := s.apiKeyAccessRuleUsers(ctx, apiKeyRules)
	if err != nil {
		return nil, err
	}
	keyFiltered := make([]RuntimeModelCandidate, 0, len(baseline))
	for _, candidate := range baseline {
		effectiveRules := effectiveAPIKeyRulesForCandidate(apiKeyRules, baselineRules, apiKeyUsers, candidate)
		if candidateAllowedByAccessRules(candidate, effectiveRules, apiKeyAccessRuleSubjects(accessUser), permissionLevel(accessUser)) {
			keyFiltered = append(keyFiltered, candidate)
		}
	}
	filtered := make([]RuntimeModelCandidate, 0, len(keyFiltered))
	for _, candidate := range keyFiltered {
		if modelaccess.ScopeAllowsModelType(accessUser.APIKeyScopes, candidate.ModelType) {
			filtered = append(filtered, candidate)
		}
	}
	return filtered, nil
}

func (s *Store) ListAPIKeyAssignablePlatformModelsForKey(ctx context.Context, user *auth.User, apiKeyID string) ([]PlatformModel, []APIKeyAccessRuleDiagnostic, error) {
	if localGatewayUserID(user) == "" {
		return nil, nil, ErrLocalUserRequired
	}
	accessUser, err := s.resolveOwnedAPIKeyAccessUser(ctx, user, apiKeyID)
	if err != nil {
		return nil, nil, err
	}
	allModels, err := s.ListModels(ctx)
	if err != nil {
		return nil, nil, err
	}
	enabledModels, platforms, err := s.enabledPlatformModels(ctx)
	if err != nil {
		return nil, nil, err
	}
	rules, err := s.listActiveAccessRulesForResources(ctx, platformModelAccessResources(enabledModels))
	if err != nil {
		return nil, nil, err
	}
	baselineRules, _ := splitLayeredAccessRules(rules)
	baseline := filterPlatformModelsByRuleSet(enabledModels, baselineRules, baselineAccessRuleSubjects(accessUser), permissionLevel(accessUser))
	scoped := filterPlatformModelsByAPIKeyScopes(baseline, accessUser.APIKeyScopes)
	ownedRules, err := s.ListAPIKeyAccessRules(ctx, user)
	if err != nil {
		return nil, nil, err
	}
	diagnostics := diagnoseAPIKeyRules(apiKeyID, ownedRules, allModels, enabledModels, baseline, scoped, platforms)
	return scoped, diagnostics, nil
}

func (s *Store) resolveOwnedAPIKeyAccessUser(ctx context.Context, user *auth.User, apiKeyID string) (*auth.User, error) {
	if user == nil {
		return nil, ErrLocalUserRequired
	}
	gatewayUserID := localGatewayUserID(user)
	if gatewayUserID == "" {
		return nil, ErrLocalUserRequired
	}
	var scopesBytes []byte
	var userGroupID string
	var userGroupKey string
	err := s.pool.QueryRow(ctx, `
SELECT k.scopes, COALESCE(k.user_group_id::text, u.default_user_group_id::text, ''), COALESCE(g.group_key, '')
FROM gateway_api_keys k
JOIN gateway_users u ON u.id = k.gateway_user_id
LEFT JOIN gateway_user_groups g ON g.id = COALESCE(k.user_group_id, u.default_user_group_id)
WHERE k.id = $1::uuid
  AND k.gateway_user_id = $2::uuid
  AND k.deleted_at IS NULL
  AND u.deleted_at IS NULL`, strings.TrimSpace(apiKeyID), gatewayUserID).Scan(&scopesBytes, &userGroupID, &userGroupKey)
	if err != nil {
		return nil, err
	}
	next := *user
	next.APIKeyID = strings.TrimSpace(apiKeyID)
	next.APIKeyScopes = decodeStringArray(scopesBytes)
	next.UserGroupID = userGroupID
	next.UserGroupKey = userGroupKey
	next.UserGroupKeys = nil
	if userGroupKey != "" {
		next.UserGroupKeys = []string{userGroupKey}
	}
	return &next, nil
}

func effectiveAPIKeyRulesForPlatformModel(rules []AccessRule, baselineRules []AccessRule, users map[string]*auth.User, model PlatformModel) []AccessRule {
	effective := make([]AccessRule, 0, len(rules))
	for _, rule := range rules {
		accessUser := users[rule.SubjectID]
		if accessUser == nil || !accessRuleMatchesPlatformModel(rule, model) ||
			!platformModelAllowedByAccessRules(model, baselineRules, baselineAccessRuleSubjects(accessUser), permissionLevel(accessUser)) ||
			len(modelaccess.FilterModelTypes(accessUser.APIKeyScopes, model.ModelType)) == 0 {
			continue
		}
		effective = append(effective, rule)
	}
	return effective
}

func effectiveAPIKeyRulesForCandidate(rules []AccessRule, baselineRules []AccessRule, users map[string]*auth.User, candidate RuntimeModelCandidate) []AccessRule {
	effective := make([]AccessRule, 0, len(rules))
	for _, rule := range rules {
		accessUser := users[rule.SubjectID]
		if accessUser == nil || !accessRuleMatchesCandidate(rule, candidate) ||
			!candidateAllowedByAccessRules(candidate, baselineRules, baselineAccessRuleSubjects(accessUser), permissionLevel(accessUser)) ||
			!modelaccess.ScopeAllowsModelType(accessUser.APIKeyScopes, candidate.ModelType) {
			continue
		}
		effective = append(effective, rule)
	}
	return effective
}

func (s *Store) apiKeyAccessRuleUsers(ctx context.Context, rules []AccessRule) (map[string]*auth.User, error) {
	ids := make([]string, 0, len(rules))
	seen := map[string]bool{}
	for _, rule := range rules {
		if rule.SubjectType != "api_key" || rule.SubjectID == "" || seen[rule.SubjectID] {
			continue
		}
		seen[rule.SubjectID] = true
		ids = append(ids, rule.SubjectID)
	}
	if len(ids) == 0 {
		return map[string]*auth.User{}, nil
	}
	rows, err := s.pool.Query(ctx, `
SELECT k.id::text, k.scopes, u.id::text,
       COALESCE(u.gateway_tenant_id::text, ''), COALESCE(u.tenant_id, ''), COALESCE(u.tenant_key, ''),
       u.roles, COALESCE(k.user_group_id::text, u.default_user_group_id::text, ''), COALESCE(g.group_key, '')
FROM gateway_api_keys k
JOIN gateway_users u ON u.id = k.gateway_user_id
LEFT JOIN gateway_user_groups g ON g.id = COALESCE(k.user_group_id, u.default_user_group_id)
WHERE k.id = ANY($1::uuid[])
  AND k.status = 'active'
  AND k.deleted_at IS NULL
  AND (k.expires_at IS NULL OR k.expires_at > now())
  AND u.status = 'active'
  AND u.deleted_at IS NULL`, ids)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	users := map[string]*auth.User{}
	for rows.Next() {
		var apiKeyID string
		var scopesBytes []byte
		var rolesBytes []byte
		var gatewayUserID string
		var gatewayTenantID string
		var tenantID string
		var tenantKey string
		var userGroupID string
		var userGroupKey string
		if err := rows.Scan(&apiKeyID, &scopesBytes, &gatewayUserID, &gatewayTenantID, &tenantID, &tenantKey, &rolesBytes, &userGroupID, &userGroupKey); err != nil {
			return nil, err
		}
		groupKeys := []string(nil)
		if userGroupKey != "" {
			groupKeys = []string{userGroupKey}
		}
		users[apiKeyID] = &auth.User{
			GatewayUserID:   gatewayUserID,
			GatewayTenantID: gatewayTenantID,
			TenantID:        tenantID,
			TenantKey:       tenantKey,
			Roles:           decodeStringArray(rolesBytes),
			UserGroupID:     userGroupID,
			UserGroupKey:    userGroupKey,
			UserGroupKeys:   groupKeys,
			APIKeyID:        apiKeyID,
			APIKeyScopes:    decodeStringArray(scopesBytes),
		}
	}
	return users, rows.Err()
}

func splitLayeredAccessRules(rules []AccessRule) ([]AccessRule, []AccessRule) {
	baseline := make([]AccessRule, 0, len(rules))
	apiKeys := make([]AccessRule, 0, len(rules))
	for _, rule := range rules {
		if rule.SubjectType == "api_key" {
			apiKeys = append(apiKeys, rule)
		} else {
			baseline = append(baseline, rule)
		}
	}
	return baseline, apiKeys
}

func filterPlatformModelsByRuleSet(models []PlatformModel, rules []AccessRule, subjects map[string]bool, level int) []PlatformModel {
	filtered := make([]PlatformModel, 0, len(models))
	for _, model := range models {
		if platformModelAllowedByAccessRules(model, rules, subjects, level) {
			filtered = append(filtered, model)
		}
	}
	return filtered
}

func filterCandidatesByRuleSet(candidates []RuntimeModelCandidate, rules []AccessRule, subjects map[string]bool, level int) []RuntimeModelCandidate {
	filtered := make([]RuntimeModelCandidate, 0, len(candidates))
	for _, candidate := range candidates {
		if candidateAllowedByAccessRules(candidate, rules, subjects, level) {
			filtered = append(filtered, candidate)
		}
	}
	return filtered
}

func filterPlatformModelsByAPIKeyScopes(models []PlatformModel, scopes []string) []PlatformModel {
	filtered := make([]PlatformModel, 0, len(models))
	for _, model := range models {
		declaredTypes := append(StringList(nil), model.ModelType...)
		allowedTypes := modelaccess.FilterModelTypes(scopes, model.ModelType)
		if len(allowedTypes) == 0 {
			continue
		}
		model.ModelType = StringList(allowedTypes)
		model.Capabilities = filterPlatformModelTypeConfig(model.Capabilities, allowedTypes, declaredTypes)
		model.BaseCapabilities = filterPlatformModelTypeConfig(model.BaseCapabilities, allowedTypes, declaredTypes)
		model.CapabilityOverride = filterPlatformModelTypeConfig(model.CapabilityOverride, allowedTypes, declaredTypes)
		filtered = append(filtered, model)
	}
	return filtered
}

func filterPlatformModelTypeConfig(config map[string]any, allowedTypes []string, declaredTypes []string) map[string]any {
	if len(config) == 0 {
		return config
	}
	allowed := map[string]bool{}
	for _, modelType := range allowedTypes {
		allowed[modelType] = true
	}
	declared := map[string]bool{}
	for _, modelType := range declaredTypes {
		declared[modelType] = true
	}
	out := make(map[string]any, len(config))
	for key, value := range config {
		if key == "originalTypes" {
			original := stringValues(value)
			kept := make([]string, 0, len(original))
			for _, modelType := range original {
				if allowed[modelType] {
					kept = append(kept, modelType)
				}
			}
			if len(kept) > 0 {
				out[key] = kept
			}
			continue
		}
		if (declared[key] || knownPlatformModelType(key)) && !allowed[key] {
			continue
		}
		out[key] = value
	}
	return out
}

func knownPlatformModelType(value string) bool {
	switch value {
	case "text_generate", "text_embedding", "text_rerank", "tools_call",
		"image_generate", "image_edit", "image_analysis", "image_vectorize",
		"video_generate", "video_enhance", "image_to_video", "text_to_video", "video_edit", "video_reference", "video_first_last_frame", "video_understanding", "omni_video", "omni",
		"audio_generate", "music_generate", "audio_understanding", "text_to_speech", "voice_clone":
		return true
	default:
		return false
	}
}

func stringValues(value any) []string {
	switch values := value.(type) {
	case []string:
		return values
	case StringList:
		return []string(values)
	case []any:
		out := make([]string, 0, len(values))
		for _, value := range values {
			if text, ok := value.(string); ok && text != "" {
				out = append(out, text)
			}
		}
		return out
	default:
		return nil
	}
}

func baselineAccessRuleSubjects(user *auth.User) map[string]bool {
	subjects := accessRuleSubjects(user)
	if user != nil && user.APIKeyID != "" {
		delete(subjects, "api_key:"+user.APIKeyID)
	}
	return subjects
}

func apiKeyAccessRuleSubjects(user *auth.User) map[string]bool {
	subjects := map[string]bool{}
	if user != nil && strings.TrimSpace(user.APIKeyID) != "" {
		subjects["api_key:"+strings.TrimSpace(user.APIKeyID)] = true
	}
	return subjects
}

func permissionLevel(user *auth.User) int {
	if user == nil {
		return 0
	}
	return auth.PermissionLevel(user.Roles)
}

func accessRuleMatchesPlatformModel(rule AccessRule, model PlatformModel) bool {
	switch rule.ResourceType {
	case "platform":
		return rule.ResourceID == model.PlatformID
	case "platform_model":
		return rule.ResourceID == model.ID
	case "base_model":
		return rule.ResourceID != "" && rule.ResourceID == model.BaseModelID
	default:
		return false
	}
}

func accessRuleMatchesCandidate(rule AccessRule, candidate RuntimeModelCandidate) bool {
	switch rule.ResourceType {
	case "platform":
		return rule.ResourceID == candidate.PlatformID
	case "platform_model":
		return rule.ResourceID == candidate.PlatformModelID
	case "base_model":
		return rule.ResourceID != "" && rule.ResourceID == candidate.BaseModelID
	default:
		return false
	}
}

func diagnoseAPIKeyRules(apiKeyID string, rules []AccessRule, allModels []PlatformModel, enabledModels []PlatformModel, baselineModels []PlatformModel, scopedModels []PlatformModel, platforms []Platform) []APIKeyAccessRuleDiagnostic {
	platformNames := map[string]string{}
	for _, platform := range platforms {
		platformNames[platform.ID] = firstNonEmpty(platform.InternalName, platform.Name, platform.PlatformKey)
	}
	modelNames := map[string]string{}
	baseModelNames := map[string]string{}
	for _, model := range allModels {
		modelNames[model.ID] = firstNonEmpty(model.DisplayName, model.ModelName, model.ID)
		if model.BaseModelID != "" && baseModelNames[model.BaseModelID] == "" {
			baseModelNames[model.BaseModelID] = firstNonEmpty(model.DisplayName, model.ModelName, model.BaseModelID)
		}
	}
	diagnostics := make([]APIKeyAccessRuleDiagnostic, 0)
	for _, rule := range rules {
		if rule.SubjectType != "api_key" || rule.SubjectID != apiKeyID || rule.Status != "active" {
			continue
		}
		diagnostic := APIKeyAccessRuleDiagnostic{
			RuleID:       rule.ID,
			ResourceType: rule.ResourceType,
			ResourceID:   rule.ResourceID,
			ResourceName: accessRuleResourceName(rule, platformNames, modelNames, baseModelNames),
			Effect:       rule.Effect,
			Effective:    true,
		}
		switch {
		case !accessRuleMatchesAnyPlatformModel(rule, allModels) || !accessRuleMatchesAnyPlatformModel(rule, enabledModels):
			diagnostic.Effective = false
			diagnostic.Reason = "resource_unavailable"
		case !accessRuleMatchesAnyPlatformModel(rule, baselineModels):
			diagnostic.Effective = false
			diagnostic.Reason = "owner_access_revoked"
		case !accessRuleMatchesAnyPlatformModel(rule, scopedModels):
			diagnostic.Effective = false
			diagnostic.Reason = "scope_not_allowed"
		}
		diagnostics = append(diagnostics, diagnostic)
	}
	return diagnostics
}

func accessRuleMatchesAnyPlatformModel(rule AccessRule, models []PlatformModel) bool {
	for _, model := range models {
		if accessRuleMatchesPlatformModel(rule, model) {
			return true
		}
	}
	return false
}

func accessRuleResourceName(rule AccessRule, platformNames map[string]string, modelNames map[string]string, baseModelNames map[string]string) string {
	switch rule.ResourceType {
	case "platform":
		return firstNonEmpty(platformNames[rule.ResourceID], rule.ResourceID)
	case "platform_model":
		return firstNonEmpty(modelNames[rule.ResourceID], rule.ResourceID)
	case "base_model":
		return firstNonEmpty(baseModelNames[rule.ResourceID], rule.ResourceID)
	default:
		return rule.ResourceID
	}
}