163 lines
12 KiB
Go
163 lines
12 KiB
Go
|
package repository
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"testing"
|
||
|
|
|
||
|
|
"github.com/Tencent/WeKnora/internal/accesspolicy"
|
||
|
|
"github.com/Tencent/WeKnora/internal/types"
|
||
|
|
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||
|
|
"github.com/stretchr/testify/require"
|
||
|
|
"gorm.io/driver/sqlite"
|
||
|
|
"gorm.io/gorm"
|
||
|
|
)
|
||
|
|
|
||
|
|
func fileAccessFixture(t *testing.T) (*resourceRepository, *gorm.DB, string) {
|
||
|
|
t.Helper()
|
||
|
|
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
||
|
|
require.NoError(t, err)
|
||
|
|
sqlDB, err := db.DB()
|
||
|
|
require.NoError(t, err)
|
||
|
|
sqlDB.SetMaxOpenConns(1)
|
||
|
|
t.Cleanup(func() { _ = sqlDB.Close() })
|
||
|
|
require.NoError(t, db.AutoMigrate(&types.StoredResource{}, &types.ResourceBinding{}))
|
||
|
|
for _, schema := range []string{
|
||
|
|
"CREATE TABLE knowledges (id TEXT PRIMARY KEY, tenant_id INTEGER, knowledge_base_id TEXT, file_path TEXT, deleted_at DATETIME)",
|
||
|
|
"CREATE TABLE knowledge_bases (id TEXT PRIMARY KEY, tenant_id INTEGER, creator_id TEXT, is_temporary BOOLEAN DEFAULT FALSE, version_state TEXT DEFAULT '', deleted_at DATETIME)",
|
||
|
|
"CREATE TABLE sessions (id TEXT PRIMARY KEY, tenant_id INTEGER, user_id TEXT, description TEXT, deleted_at DATETIME)",
|
||
|
|
"CREATE TABLE messages (id TEXT PRIMARY KEY, session_id TEXT, knowledge_references TEXT, deleted_at DATETIME)",
|
||
|
|
"CREATE TABLE temporary_documents (id TEXT PRIMARY KEY, tenant_id INTEGER, session_id TEXT, expires_at DATETIME, deleted_at DATETIME)",
|
||
|
|
"CREATE TABLE custom_agents (id TEXT PRIMARY KEY, tenant_id INTEGER, enabled BOOLEAN, config TEXT, deleted_at DATETIME)",
|
||
|
|
"CREATE TABLE embed_channels (id TEXT PRIMARY KEY, tenant_id INTEGER, agent_id TEXT, enabled BOOLEAN, deleted_at DATETIME)",
|
||
|
|
} {
|
||
|
|
require.NoError(t, db.Exec(schema).Error)
|
||
|
|
}
|
||
|
|
resource := &types.StoredResource{ID: "file-1", Handle: "AbCdEfGhIjKlMnOpQrStUv", TenantID: 7, PhysicalPath: "local://7/exports/private.png", Provider: "local", LocationHash: "unused"}
|
||
|
|
require.NoError(t, db.Create(resource).Error)
|
||
|
|
return &resourceRepository{db: db}, db, types.BuildResourcePath(resource.Handle)
|
||
|
|
}
|
||
|
|
|
||
|
|
func fileUserContext(user string) context.Context {
|
||
|
|
ctx := context.WithValue(context.Background(), types.TenantIDContextKey, uint64(7))
|
||
|
|
ctx = context.WithValue(ctx, types.UserIDContextKey, user)
|
||
|
|
return accesspolicy.WithSnapshot(ctx, &accesspolicy.Snapshot{Publications: map[string]bool{"knowledge:kb-a": false}})
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestResourceFileAccessKnowledgeBinding(t *testing.T) {
|
||
|
|
r, db, ref := fileAccessFixture(t)
|
||
|
|
require.NoError(t, db.Exec("INSERT INTO knowledge_bases (id,tenant_id,creator_id) VALUES ('kb-a',7,'alice'),('kb-b',7,'bob')").Error)
|
||
|
|
require.NoError(t, db.Exec("INSERT INTO knowledges (id,tenant_id,knowledge_base_id,file_path) VALUES ('doc-a',7,'kb-a','local://7/legacy.doc')").Error)
|
||
|
|
require.NoError(t, r.CreateBinding(context.Background(), &types.ResourceBinding{ResourceID: "file-1", TenantID: 7, OwnerType: types.ResourceOwnerKnowledge, OwnerID: "doc-a"}))
|
||
|
|
scope := interfaces.ResourceFileScope{TenantID: 7}
|
||
|
|
require.NoError(t, r.AuthorizeFileRead(fileUserContext("alice"), ref, scope))
|
||
|
|
require.Error(t, r.AuthorizeFileRead(fileUserContext("bob"), ref, scope))
|
||
|
|
require.Error(t, r.AuthorizeFileRead(fileUserContext("bob"), "local://7/legacy.doc", scope))
|
||
|
|
require.NoError(t, r.AuthorizeFileRead(fileUserContext("alice"), "local://7/legacy.doc", scope))
|
||
|
|
require.Error(t, r.AuthorizeFileRead(fileUserContext("alice"), "local://7/unregistered.png", scope))
|
||
|
|
require.Error(t, r.AuthorizeFileRead(fileUserContext("alice"), ref, interfaces.ResourceFileScope{TenantID: 8}))
|
||
|
|
require.Error(t, r.AuthorizeFileRead(fileUserContext("bob"), ref, interfaces.ResourceFileScope{TenantID: 7, KnowledgeBaseID: "kb-b"}))
|
||
|
|
require.NoError(t, r.AuthorizeFileRead(fileUserContext("bob"), ref, interfaces.ResourceFileScope{TenantID: 7, KnowledgeBaseID: "kb-a"}))
|
||
|
|
shared := accesspolicy.WithSnapshot(fileUserContext("bob"), &accesspolicy.Snapshot{Publications: map[string]bool{"knowledge:kb-a": false}, Grants: map[string]bool{"knowledge:kb-a": true}})
|
||
|
|
require.NoError(t, r.AuthorizeFileRead(shared, ref, scope))
|
||
|
|
denied := accesspolicy.WithSnapshot(fileUserContext("alice"), &accesspolicy.Snapshot{Restricted: true, Permissions: map[string]bool{}})
|
||
|
|
require.Error(t, r.AuthorizeFileRead(denied, ref, scope))
|
||
|
|
require.NoError(t, db.Exec("UPDATE knowledges SET deleted_at=CURRENT_TIMESTAMP WHERE id='doc-a'").Error)
|
||
|
|
require.Error(t, r.AuthorizeFileRead(fileUserContext("alice"), ref, scope))
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestResourceFileAccessSharedOriginalCannotBeDownloaded(t *testing.T) {
|
||
|
|
r, db, ref := fileAccessFixture(t)
|
||
|
|
require.NoError(t, db.Exec("INSERT INTO knowledge_bases (id,tenant_id,creator_id) VALUES ('kb-a',7,'alice')").Error)
|
||
|
|
require.NoError(t, db.Exec("INSERT INTO knowledges (id,tenant_id,knowledge_base_id,file_path) VALUES ('doc-a',7,'kb-a','local://7/legacy.doc')").Error)
|
||
|
|
require.NoError(t, r.CreateBinding(context.Background(), &types.ResourceBinding{ResourceID: "file-1", TenantID: 7, OwnerType: types.ResourceOwnerKnowledge, OwnerID: "doc-a", Relation: types.ResourceRelationSourceFile}))
|
||
|
|
require.NoError(t, r.CreateBinding(context.Background(), &types.ResourceBinding{ResourceID: "file-1", TenantID: 7, OwnerType: types.ResourceOwnerKnowledge, OwnerID: "doc-a", Relation: types.ResourceRelationAttachment}))
|
||
|
|
require.NoError(t, db.Exec("INSERT INTO sessions (id,tenant_id,user_id) VALUES ('s-b',7,'bob')").Error)
|
||
|
|
require.NoError(t, db.Exec("INSERT INTO messages (id,session_id,knowledge_references) VALUES ('msg-b','s-b',?)", []byte(`[{"knowledge_id":"doc-a"}]`)).Error)
|
||
|
|
shared := accesspolicy.WithSnapshot(fileUserContext("bob"), &accesspolicy.Snapshot{Publications: map[string]bool{"knowledge:kb-a": true}, Grants: map[string]bool{"knowledge:kb-a": true}})
|
||
|
|
for _, scope := range []interfaces.ResourceFileScope{
|
||
|
|
{TenantID: 7},
|
||
|
|
{TenantID: 7, KnowledgeBaseID: "kb-a"},
|
||
|
|
{TenantID: 7, MessageID: "msg-b", SessionID: "s-b"},
|
||
|
|
{TenantID: 7, EmbedChannelID: "channel-a"},
|
||
|
|
} {
|
||
|
|
require.ErrorIs(t, r.AuthorizeFileRead(shared, ref, scope), errFileAccessDenied)
|
||
|
|
require.ErrorIs(t, r.AuthorizeFileRead(shared, "local://7/legacy.doc", scope), errFileAccessDenied)
|
||
|
|
}
|
||
|
|
scope := interfaces.ResourceFileScope{TenantID: 7}
|
||
|
|
require.NoError(t, r.AuthorizeFileRead(fileUserContext("alice"), ref, scope))
|
||
|
|
admin := context.WithValue(fileUserContext("admin"), types.TenantRoleContextKey, types.TenantRoleAdmin)
|
||
|
|
require.NoError(t, r.AuthorizeFileRead(admin, ref, scope))
|
||
|
|
systemAdmin := context.WithValue(fileUserContext("system"), types.SystemAdminContextKey, true)
|
||
|
|
require.NoError(t, r.AuthorizeFileRead(systemAdmin, ref, scope))
|
||
|
|
crossTenant := context.WithValue(systemAdmin, types.TenantIDContextKey, uint64(8))
|
||
|
|
require.ErrorIs(t, r.AuthorizeFileRead(crossTenant, ref, scope), errFileAccessDenied)
|
||
|
|
denied := accesspolicy.WithSnapshot(fileUserContext("alice"), &accesspolicy.Snapshot{Restricted: true, Permissions: map[string]bool{}})
|
||
|
|
require.ErrorIs(t, r.AuthorizeFileRead(denied, ref, interfaces.ResourceFileScope{TenantID: 7, KnowledgeBaseID: "kb-a"}), errFileAccessDenied)
|
||
|
|
|
||
|
|
// A legacy registration without a source relation is still an original.
|
||
|
|
require.NoError(t, db.Exec("DELETE FROM resource_bindings WHERE relation=?", types.ResourceRelationSourceFile).Error)
|
||
|
|
for _, location := range []string{ref, "local://7/exports/private.png"} {
|
||
|
|
require.NoError(t, db.Exec("UPDATE knowledges SET file_path=? WHERE id='doc-a'", location).Error)
|
||
|
|
require.ErrorIs(t, r.AuthorizeFileRead(shared, ref, scope), errFileAccessDenied)
|
||
|
|
require.NoError(t, r.AuthorizeFileRead(fileUserContext("alice"), ref, scope))
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestResourceFileAccessPrivateMessages(t *testing.T) {
|
||
|
|
r, db, ref := fileAccessFixture(t)
|
||
|
|
require.NoError(t, db.Exec("INSERT INTO sessions (id,tenant_id,user_id) VALUES ('s-a',7,'alice'),('s-b',7,'bob'),('s-legacy',7,'')").Error)
|
||
|
|
require.NoError(t, db.Exec("INSERT INTO messages (id,session_id) VALUES ('msg-a','s-a'),('msg-b','s-b'),('msg-legacy','s-legacy')").Error)
|
||
|
|
binding := &types.ResourceBinding{ResourceID: "file-1", TenantID: 7, OwnerType: types.ResourceOwnerMessage, OwnerID: "msg-a"}
|
||
|
|
require.NoError(t, r.CreateBinding(context.Background(), binding))
|
||
|
|
scope := interfaces.ResourceFileScope{TenantID: 7}
|
||
|
|
require.NoError(t, r.AuthorizeFileRead(fileUserContext("alice"), ref, scope))
|
||
|
|
require.Error(t, r.AuthorizeFileRead(fileUserContext("bob"), ref, scope))
|
||
|
|
require.Error(t, r.AuthorizeFileRead(fileUserContext("alice"), ref, interfaces.ResourceFileScope{TenantID: 7, MessageID: "msg-b", SessionID: "s-b"}))
|
||
|
|
require.NoError(t, r.AuthorizeFileRead(fileUserContext("alice"), ref, interfaces.ResourceFileScope{TenantID: 7, MessageID: "msg-a", SessionID: "s-a"}))
|
||
|
|
require.Error(t, r.AuthorizeFileRead(fileUserContext("alice"), ref, interfaces.ResourceFileScope{TenantID: 7, KnowledgeBaseID: "kb-a"}))
|
||
|
|
require.NoError(t, db.Model(&types.ResourceBinding{}).Where("id = ?", binding.ID).Update("owner_id", "msg-legacy").Error)
|
||
|
|
require.Error(t, r.AuthorizeFileRead(fileUserContext("bob"), ref, scope))
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestResourceFileAccessEmbedKnowledgeScope(t *testing.T) {
|
||
|
|
r, db, ref := fileAccessFixture(t)
|
||
|
|
require.NoError(t, db.Exec("INSERT INTO knowledge_bases (id,tenant_id,creator_id) VALUES ('kb-a',7,'alice')").Error)
|
||
|
|
require.NoError(t, db.Exec("INSERT INTO knowledges (id,tenant_id,knowledge_base_id) VALUES ('doc-a',7,'kb-a')").Error)
|
||
|
|
require.NoError(t, r.CreateBinding(context.Background(), &types.ResourceBinding{ResourceID: "file-1", TenantID: 7, OwnerType: types.ResourceOwnerKnowledge, OwnerID: "doc-a"}))
|
||
|
|
require.NoError(t, db.Exec("INSERT INTO custom_agents (id,tenant_id,enabled,config) VALUES ('agent-a',7,TRUE,?)", `{"kb_selection_mode":"selected","knowledge_bases":["kb-a"]}`).Error)
|
||
|
|
require.NoError(t, db.Exec("INSERT INTO embed_channels (id,tenant_id,agent_id,enabled) VALUES ('channel-a',7,'agent-a',TRUE)").Error)
|
||
|
|
scope := interfaces.ResourceFileScope{TenantID: 7, EmbedChannelID: "channel-a"}
|
||
|
|
require.NoError(t, r.AuthorizeFileRead(fileUserContext("visitor"), ref, scope))
|
||
|
|
for _, config := range []string{`{"kb_selection_mode":"none"}`, `{"kb_selection_mode":"selected","knowledge_bases":["kb-other"]}`, `{"kb_selection_mode":"selected","knowledge_bases":[]}`} {
|
||
|
|
require.NoError(t, db.Exec("UPDATE custom_agents SET config=?", config).Error)
|
||
|
|
require.Error(t, r.AuthorizeFileRead(fileUserContext("visitor"), ref, scope))
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestResourceFileAccessMessageKnowledgeReferences(t *testing.T) {
|
||
|
|
r, db, ref := fileAccessFixture(t)
|
||
|
|
require.NoError(t, db.Exec("INSERT INTO knowledges (id,tenant_id,knowledge_base_id) VALUES ('doc-a',7,'kb-a')").Error)
|
||
|
|
require.NoError(t, r.CreateBinding(context.Background(), &types.ResourceBinding{ResourceID: "file-1", TenantID: 7, OwnerType: types.ResourceOwnerKnowledge, OwnerID: "doc-a"}))
|
||
|
|
require.NoError(t, db.Exec("INSERT INTO sessions (id,tenant_id,user_id) VALUES ('s-a',7,'alice')").Error)
|
||
|
|
require.NoError(t, db.Exec("INSERT INTO messages (id,session_id,knowledge_references) VALUES ('msg-a','s-a',?)", []byte(`[{"knowledge_id":"doc-a"}]`)).Error)
|
||
|
|
scope := interfaces.ResourceFileScope{TenantID: 7, MessageID: "msg-a", SessionID: "s-a"}
|
||
|
|
require.NoError(t, r.AuthorizeFileRead(fileUserContext("alice"), ref, scope))
|
||
|
|
require.Error(t, r.AuthorizeFileRead(fileUserContext("bob"), ref, scope))
|
||
|
|
require.NoError(t, db.Exec("UPDATE messages SET knowledge_references=?", []byte(`[{"knowledge_id":"doc-unrelated"}]`)).Error)
|
||
|
|
require.Error(t, r.AuthorizeFileRead(fileUserContext("alice"), ref, scope))
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestResourceFileAccessEmbedPrivateSession(t *testing.T) {
|
||
|
|
r, db, ref := fileAccessFixture(t)
|
||
|
|
principal := types.EmbedSessionPrincipal(7, "channel-a", "s-a")
|
||
|
|
require.NoError(t, db.Exec("INSERT INTO sessions (id,tenant_id,user_id) VALUES ('s-a',7,?)", principal.StorageID()).Error)
|
||
|
|
require.NoError(t, db.Exec("INSERT INTO messages (id,session_id) VALUES ('msg-a','s-a')").Error)
|
||
|
|
require.NoError(t, r.CreateBinding(context.Background(), &types.ResourceBinding{ResourceID: "file-1", TenantID: 7, OwnerType: types.ResourceOwnerMessage, OwnerID: "msg-a"}))
|
||
|
|
ctx := types.WithPrincipal(fileUserContext("embed"), principal)
|
||
|
|
scope := interfaces.ResourceFileScope{TenantID: 7, EmbedChannelID: "channel-a", SessionID: "s-a"}
|
||
|
|
require.NoError(t, r.AuthorizeFileRead(ctx, ref, scope))
|
||
|
|
require.Error(t, r.AuthorizeFileRead(fileUserContext("embed"), ref, scope))
|
||
|
|
scope.SessionID = "s-other"
|
||
|
|
require.Error(t, r.AuthorizeFileRead(ctx, ref, scope))
|
||
|
|
}
|