xinmi/internal/application/service/knowledgebase_version_access_test.go

package service

import (
	"context"
	"testing"

	"github.com/Tencent/WeKnora/internal/accesspolicy"
	"github.com/Tencent/WeKnora/internal/knowledgeversion"
	"github.com/Tencent/WeKnora/internal/types"
	"github.com/Tencent/WeKnora/internal/types/interfaces"
	"github.com/stretchr/testify/require"
)

type versionAccessKBRepository struct {
	interfaces.KnowledgeBaseRepository
	bases map[string]*types.KnowledgeBase
}

func (r versionAccessKBRepository) GetKnowledgeBaseByID(_ context.Context, id string) (*types.KnowledgeBase, error) {
	return r.bases[id], nil
}

func (r versionAccessKBRepository) GetKnowledgeBaseByIDs(_ context.Context, ids []string) ([]*types.KnowledgeBase, error) {
	result := make([]*types.KnowledgeBase, 0, len(ids))
	for _, id := range ids {
		result = append(result, r.bases[id])
	}
	return result, nil
}

func TestKnowledgeBaseVersionStateAccessIsolation(t *testing.T) {
	owner := context.WithValue(context.Background(), types.TenantIDContextKey, uint64(7))
	owner = context.WithValue(owner, types.UserIDContextKey, "owner")
	owner = context.WithValue(owner, types.TenantRoleContextKey, types.TenantRoleContributor)
	owner = accesspolicy.WithSnapshot(owner, &accesspolicy.Snapshot{Restricted: true, Permissions: map[string]bool{"knowledge.read": true}})
	for _, state := range []string{"copying", "snapshot"} {
		for _, actor := range []struct {
			name    string
			ctx     context.Context
			allowed bool
		}{
			{"owner", owner, false},
			{"workspace_admin", context.WithValue(owner, types.TenantRoleContextKey, types.TenantRoleAdmin), false},
			{"system_admin", context.WithValue(owner, types.SystemAdminContextKey, true), false},
			{"cross_workspace", context.WithValue(owner, types.TenantIDContextKey, uint64(8)), false},
			{"worker", knowledgeversion.WithWorker(owner), true},
		} {
			t.Run(state+"/"+actor.name, func(t *testing.T) {
				staging := &types.KnowledgeBase{ID: "staging", TenantID: 7, CreatorID: "owner", VersionState: state}
				ready := &types.KnowledgeBase{ID: "ready", TenantID: 7, CreatorID: "owner"}
				svc := &knowledgeBaseService{repo: versionAccessKBRepository{bases: map[string]*types.KnowledgeBase{staging.ID: staging, ready.ID: ready}}}
				require.Equal(t, actor.allowed, knowledgeBaseReadable(actor.ctx, staging))
				require.True(t, knowledgeBaseReadable(actor.ctx, ready))
				for _, get := range []func(context.Context, string) (*types.KnowledgeBase, error){svc.GetKnowledgeBaseByID, svc.GetKnowledgeBaseByIDOnly} {
					kb, err := get(actor.ctx, staging.ID)
					if actor.allowed {
						require.NoError(t, err)
						require.Equal(t, staging.ID, kb.ID)
					} else {
						require.ErrorContains(t, err, "没有此知识库")
						require.Nil(t, kb)
					}
				}
				listed, err := svc.GetKnowledgeBasesByIDsOnly(actor.ctx, []string{staging.ID, ready.ID})
				require.NoError(t, err)
				if actor.allowed {
					require.Len(t, listed, 2)
				} else {
					require.Len(t, listed, 1)
					require.Equal(t, ready.ID, listed[0].ID)
				}
			})
		}
	}
}

func TestKnowledgeBaseRequiresExplicitSharing(t *testing.T) {
	base := context.WithValue(context.Background(), types.TenantIDContextKey, uint64(7))
	base = context.WithValue(base, types.UserIDContextKey, "other-user")
	base = context.WithValue(base, types.TenantRoleContextKey, types.TenantRoleContributor)
	kb := &types.KnowledgeBase{ID: "private-kb", TenantID: 7, CreatorID: "owner"}
	for _, tt := range []struct {
		name     string
		owner    string
		role     types.TenantRole
		snapshot *accesspolicy.Snapshot
		allowed  bool
	}{
		{"unmanaged_private", "other-user", types.TenantRoleContributor, &accesspolicy.Snapshot{}, false},
		{"owner", "owner", types.TenantRoleContributor, &accesspolicy.Snapshot{}, true},
		{"workspace_admin", "admin", types.TenantRoleAdmin, &accesspolicy.Snapshot{}, true},
		{"published", "other-user", types.TenantRoleContributor, &accesspolicy.Snapshot{Publications: map[string]bool{"knowledge:private-kb": true}}, true},
		{"granted", "other-user", types.TenantRoleContributor, &accesspolicy.Snapshot{Grants: map[string]bool{"knowledge:private-kb": true}}, true},
		{"revoked", "other-user", types.TenantRoleContributor, &accesspolicy.Snapshot{Publications: map[string]bool{"knowledge:private-kb": false}}, false},
		{"module_denied", "owner", types.TenantRoleContributor, &accesspolicy.Snapshot{Restricted: true}, false},
		{"background", "other-user", types.TenantRoleContributor, nil, true},
	} {
		t.Run(tt.name, func(t *testing.T) {
			ctx := context.WithValue(base, types.UserIDContextKey, tt.owner)
			ctx = context.WithValue(ctx, types.TenantRoleContextKey, tt.role)
			if tt.snapshot != nil {
				ctx = accesspolicy.WithSnapshot(ctx, tt.snapshot)
			}
			require.Equal(t, tt.allowed, knowledgeBaseReadable(ctx, kb))
			listed := filterReadableKnowledgeBases(ctx, []*types.KnowledgeBase{kb})
			require.Equal(t, tt.allowed, len(listed) == 1)
		})
	}
}