xinmi/internal/application/repository/resource_access.go

package repository

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"strings"
	"time"

	"github.com/Tencent/WeKnora/internal/accesspolicy"
	"github.com/Tencent/WeKnora/internal/types"
	"github.com/Tencent/WeKnora/internal/types/interfaces"
	"gorm.io/gorm"
)

var errFileAccessDenied = errors.New("file is not attached to an accessible resource")

// A storage locator is never an authorization capability. Resolve it back to
// live domain owners, then apply the same personal/knowledge scope as their APIs.
func (r *resourceRepository) AuthorizeFileRead(ctx context.Context, value string, scope interfaces.ResourceFileScope) error {
	if scope.TenantID == 0 {
		return errFileAccessDenied
	}
	var resource *types.StoredResource
	var err error
	if handle, ok := types.ParseResourcePath(value); ok {
		resource, err = r.GetByHandle(ctx, handle)
	} else {
		hash := sha256.Sum256([]byte(value))
		resource, err = r.GetByTenantLocation(ctx, scope.TenantID, hex.EncodeToString(hash[:]))
	}
	if err != nil {
		return err
	}
	var bindings []types.ResourceBinding
	if resource != nil {
		if resource.TenantID != scope.TenantID {
			return errFileAccessDenied
		}
		if err := r.db.WithContext(ctx).Where("resource_id = ? AND tenant_id = ?", resource.ID, scope.TenantID).Find(&bindings).Error; err != nil {
			return err
		}
		// Older registrations may lack a source_file relation. An attachment
		// alias for the same original must not bypass the download restriction.
		var originals []types.Knowledge
		if err := r.db.WithContext(ctx).Select("id").Where("tenant_id = ? AND file_path IN ?", scope.TenantID,
			[]string{value, resource.PhysicalPath, types.BuildResourcePath(resource.Handle)}).Find(&originals).Error; err != nil {
			return err
		}
		for _, doc := range originals {
			bindings = append(bindings, types.ResourceBinding{OwnerType: types.ResourceOwnerKnowledge, OwnerID: doc.ID, Relation: types.ResourceRelationSourceFile})
		}
	} else if _, isHandle := types.ParseResourcePath(value); !isHandle {
		// Legacy source files predate resource bindings but still have an exact
		// knowledge owner. Do not accept an arbitrary same-tenant path.
		var documents []types.Knowledge
		if err := r.db.WithContext(ctx).Select("id").Where("tenant_id = ? AND file_path = ?", scope.TenantID, value).Find(&documents).Error; err != nil {
			return err
		}
		for _, doc := range documents {
			bindings = append(bindings, types.ResourceBinding{OwnerType: types.ResourceOwnerKnowledge, OwnerID: doc.ID, Relation: types.ResourceRelationSourceFile})
		}
	}
	var sourceBindings []types.ResourceBinding
	for _, binding := range bindings {
		if binding.OwnerType == types.ResourceOwnerKnowledge && binding.Relation == types.ResourceRelationSourceFile {
			sourceBindings = append(sourceBindings, binding)
		}
	}
	if len(sourceBindings) > 0 {
		bindings = sourceBindings
	}
	for _, binding := range bindings {
		allowed, err := r.readableFileBinding(ctx, scope, binding)
		if err != nil {
			return err
		}
		if allowed {
			return nil
		}
	}
	return errFileAccessDenied
}

func (r *resourceRepository) readableFileBinding(ctx context.Context, scope interfaces.ResourceFileScope, binding types.ResourceBinding) (bool, error) {
	q := r.db.WithContext(ctx)
	callerTenant, _ := types.TenantIDFromContext(ctx)
	userID, _ := types.UserIDFromContext(ctx)
	switch binding.OwnerType {
	case types.ResourceOwnerKnowledge:
		if binding.Relation == types.ResourceRelationSourceFile {
			if callerTenant != scope.TenantID || userID == "" || scope.EmbedChannelID != "" || !accesspolicy.Allowed(ctx, "knowledge.read") {
				return false, nil
			}
			var source types.Knowledge
			if err := q.Select("id,knowledge_base_id").Where("id = ? AND tenant_id = ?", binding.OwnerID, scope.TenantID).Take(&source).Error; err != nil {
				return false, ignoreMissingFileOwner(err)
			}
			var base types.KnowledgeBase
			if err := q.Select("id,creator_id").Where("id = ? AND tenant_id = ?", source.KnowledgeBaseID, scope.TenantID).Take(&base).Error; err != nil {
				return false, ignoreMissingFileOwner(err)
			}
			if base.CreatorID != userID && !types.IsSystemAdminFromContext(ctx) && !types.TenantRoleFromContext(ctx).HasPermission(types.TenantRoleAdmin) {
				return false, nil
			}
		}
		if scope.MessageID != "" {
			var message types.Message
			if err := q.Select("id,session_id,knowledge_references").Where("id = ?", scope.MessageID).Take(&message).Error; err != nil {
				return false, ignoreMissingFileOwner(err)
			}
			for _, reference := range message.KnowledgeReferences {
				if reference != nil && reference.KnowledgeID == binding.OwnerID {
					var count int64
					if err := q.Model(&types.Knowledge{}).Where("id = ? AND tenant_id = ?", binding.OwnerID, scope.TenantID).Count(&count).Error; err != nil {
						return false, err
					}
					if count == 0 {
						return false, nil
					}
					return r.readableFileSession(ctx, scope, message.SessionID)
				}
			}
			return false, nil
		}
		var knowledge types.Knowledge
		if err := q.Select("id,knowledge_base_id").Where("id = ? AND tenant_id = ?", binding.OwnerID, scope.TenantID).Take(&knowledge).Error; err != nil {
			return false, ignoreMissingFileOwner(err)
		}
		if scope.KnowledgeBaseID != "" {
			// KBAccessRead has already authorized this exact knowledge base.
			return knowledge.KnowledgeBaseID == scope.KnowledgeBaseID, nil
		}
		var base types.KnowledgeBase
		if err := q.Select("id,tenant_id,creator_id,is_temporary,version_state").Where("id = ? AND tenant_id = ?", knowledge.KnowledgeBaseID, scope.TenantID).Take(&base).Error; err != nil {
			return false, ignoreMissingFileOwner(err)
		}
		if base.IsTemporary || base.VersionState != "" {
			return false, nil
		}
		if scope.EmbedChannelID != "" {
			var channel types.EmbedChannel
			if err := q.Where("id = ? AND tenant_id = ? AND enabled = ?", scope.EmbedChannelID, scope.TenantID, true).Take(&channel).Error; err != nil {
				return false, ignoreMissingFileOwner(err)
			}
			var agent types.CustomAgent
			if err := q.Where("id = ? AND tenant_id = ? AND enabled = ?", channel.AgentID, scope.TenantID, true).Take(&agent).Error; err != nil {
				return false, ignoreMissingFileOwner(err)
			}
			switch agent.Config.KBSelectionMode {
			case "none":
				return false, nil
			case "selected":
				for _, id := range agent.Config.KnowledgeBases {
					if id == base.ID {
						return true, nil
					}
				}
				return false, nil
			default:
				return true, nil
			}
		}
		if callerTenant != scope.TenantID || userID == "" || !accesspolicy.Allowed(ctx, "knowledge.read") {
			return false, nil
		}
		if types.IsSystemAdminFromContext(ctx) || types.TenantRoleFromContext(ctx).HasPermission(types.TenantRoleAdmin) {
			return true, nil
		}
		return accesspolicy.Readable(ctx, "knowledge", base.ID, base.CreatorID == userID, accesspolicy.FromContext(ctx) == nil), nil
	case types.ResourceOwnerMessage, types.ResourceOwnerTemporaryDocument:
		userID = types.SessionOwnerIDFromContext(ctx)
		if scope.KnowledgeBaseID != "" || userID == "" {
			return false, nil
		}
		var sessionID string
		if binding.OwnerType == types.ResourceOwnerMessage {
			if scope.MessageID != "" && binding.OwnerID != scope.MessageID {
				return false, nil
			}
			var message types.Message
			if err := q.Select("id,session_id").Where("id = ?", binding.OwnerID).Take(&message).Error; err != nil {
				return false, ignoreMissingFileOwner(err)
			}
			sessionID = message.SessionID
		} else {
			if scope.MessageID != "" {
				return false, nil
			}
			var document types.TemporaryDocument
			if err := q.Select("session_id").Where("id = ? AND tenant_id = ? AND expires_at > ?", binding.OwnerID, scope.TenantID, time.Now()).Take(&document).Error; err != nil {
				return false, ignoreMissingFileOwner(err)
			}
			sessionID = document.SessionID
		}
		return r.readableFileSession(ctx, scope, sessionID)
	}
	return false, nil
}

func (r *resourceRepository) readableFileSession(ctx context.Context, scope interfaces.ResourceFileScope, sessionID string) (bool, error) {
	tenantID, _ := types.TenantIDFromContext(ctx)
	userID := types.SessionOwnerIDFromContext(ctx)
	if userID == "" || scope.SessionID != "" && scope.SessionID != sessionID {
		return false, nil
	}
	var session types.Session
	if err := r.db.WithContext(ctx).Select("id,tenant_id,user_id").Where("id = ? AND tenant_id = ? AND user_id = ?", sessionID, tenantID, userID).Take(&session).Error; err != nil {
		return false, ignoreMissingFileOwner(err)
	}
	if scope.EmbedChannelID != "" {
		return scope.SessionID != "" && strings.HasPrefix(userID, types.PrincipalEmbedSession+":"), nil
	}
	return true, nil
}

func ignoreMissingFileOwner(err error) error {
	if errors.Is(err, gorm.ErrRecordNotFound) {
		return nil
	}
	return err
}