From 21d437babd04077df644ad84d75f7747070dcd6f Mon Sep 17 00:00:00 2001 From: Jonathanchang31 <55106972+jonathanchang31@users.noreply.github.com> Date: Thu, 4 Jun 2026 17:31:21 +0200 Subject: [PATCH 1/4] fix: exposing private knowledge base --- internal/dao/kb.go | 9 ++-- internal/dao/kb_test.go | 96 ++++++++++++++++++++++++++++++++++++----- 2 files changed, 92 insertions(+), 13 deletions(-) diff --git a/internal/dao/kb.go b/internal/dao/kb.go index d74c98b7b0d..59cf1b9bf8d 100644 --- a/internal/dao/kb.go +++ b/internal/dao/kb.go @@ -245,9 +245,12 @@ func (dao *KnowledgebaseDAO) GetDetail(kbID string) (*entity.KnowledgebaseDetail func (dao *KnowledgebaseDAO) Accessible(kbID, userID string) bool { var count int64 err := DB.Table("knowledgebase"). - Joins("JOIN user_tenant ON user_tenant.tenant_id = knowledgebase.tenant_id"). - Where("knowledgebase.id = ? AND user_tenant.user_id = ? AND knowledgebase.status = ?", - kbID, userID, string(entity.StatusValid)). + Joins("LEFT JOIN user_tenant ON user_tenant.tenant_id = knowledgebase.tenant_id AND user_tenant.user_id = ?", userID). + Where(`knowledgebase.id = ? AND knowledgebase.status = ? AND ( + knowledgebase.created_by = ? OR + knowledgebase.tenant_id = ? OR + (knowledgebase.permission = ? AND user_tenant.user_id IS NOT NULL) + )`, kbID, string(entity.StatusValid), userID, userID, string(entity.TenantPermissionTeam)). Count(&count).Error if err != nil { diff --git a/internal/dao/kb_test.go b/internal/dao/kb_test.go index 47b0704e250..dca86849ad3 100644 --- a/internal/dao/kb_test.go +++ b/internal/dao/kb_test.go @@ -36,8 +36,8 @@ func setupKBTestDB(t *testing.T) *gorm.DB { t.Fatalf("failed to open sqlite: %v", err) } - // Migrate knowledgebase table - if err := db.AutoMigrate(&entity.Knowledgebase{}); err != nil { + // Migrate tables needed by KB access checks. + if err := db.AutoMigrate(&entity.Knowledgebase{}, &entity.UserTenant{}); err != nil { t.Fatalf("failed to migrate: %v", err) } @@ -47,14 +47,16 @@ func setupKBTestDB(t *testing.T) *gorm.DB { func testKnowledgebase(t *testing.T, db *gorm.DB, id string, docNum, tokenNum, chunkNum int64) *entity.Knowledgebase { t.Helper() kb := &entity.Knowledgebase{ - ID: id, - TenantID: "tenant-1", - Name: "test-kb-" + id, - EmbdID: "embd-1", - DocNum: docNum, - TokenNum: tokenNum, - ChunkNum: chunkNum, - Status: stringPtr(string(entity.StatusValid)), + ID: id, + TenantID: "tenant-1", + Name: "test-kb-" + id, + EmbdID: "embd-1", + CreatedBy: "tenant-1", + Permission: string(entity.TenantPermissionMe), + DocNum: docNum, + TokenNum: tokenNum, + ChunkNum: chunkNum, + Status: stringPtr(string(entity.StatusValid)), } if err := db.Create(kb).Error; err != nil { t.Fatalf("failed to create test kb: %v", err) @@ -62,6 +64,22 @@ func testKnowledgebase(t *testing.T, db *gorm.DB, id string, docNum, tokenNum, c return kb } +func testUserTenant(t *testing.T, db *gorm.DB, id, userID, tenantID string) { + t.Helper() + status := string(entity.StatusValid) + userTenant := &entity.UserTenant{ + ID: id, + UserID: userID, + TenantID: tenantID, + Role: "normal", + InvitedBy: tenantID, + Status: &status, + } + if err := db.Create(userTenant).Error; err != nil { + t.Fatalf("failed to create test user_tenant: %v", err) + } +} + func stringPtr(s string) *string { return &s } @@ -119,3 +137,61 @@ func TestKnowledgebaseDAO_DecreaseDocumentNum_ZeroDecrement(t *testing.T) { t.Fatalf("chunk_num should be unchanged: expected 15, got %d", kb.ChunkNum) } } + +func TestKnowledgebaseDAO_AccessibleAllowsOwnerForPrivateKnowledgebase(t *testing.T) { + db := setupKBTestDB(t) + pushDB(t, db) + dao := NewKnowledgebaseDAO() + + testKnowledgebase(t, db, "kb-private", 0, 0, 0) + + if !dao.Accessible("kb-private", "tenant-1") { + t.Fatal("owner should access a private knowledgebase") + } +} + +func TestKnowledgebaseDAO_AccessibleRejectsTenantMemberForPrivateKnowledgebase(t *testing.T) { + db := setupKBTestDB(t) + pushDB(t, db) + dao := NewKnowledgebaseDAO() + + testKnowledgebase(t, db, "kb-private", 0, 0, 0) + testUserTenant(t, db, "ut-1", "member-1", "tenant-1") + + if dao.Accessible("kb-private", "member-1") { + t.Fatal("tenant member should not access a private knowledgebase") + } +} + +func TestKnowledgebaseDAO_AccessibleAllowsTenantMemberForTeamKnowledgebase(t *testing.T) { + db := setupKBTestDB(t) + pushDB(t, db) + dao := NewKnowledgebaseDAO() + + kb := testKnowledgebase(t, db, "kb-team", 0, 0, 0) + kb.Permission = string(entity.TenantPermissionTeam) + if err := db.Save(kb).Error; err != nil { + t.Fatalf("failed to update test kb permission: %v", err) + } + testUserTenant(t, db, "ut-1", "member-1", "tenant-1") + + if !dao.Accessible("kb-team", "member-1") { + t.Fatal("tenant member should access a team knowledgebase") + } +} + +func TestKnowledgebaseDAO_AccessibleRejectsUnrelatedUserForTeamKnowledgebase(t *testing.T) { + db := setupKBTestDB(t) + pushDB(t, db) + dao := NewKnowledgebaseDAO() + + kb := testKnowledgebase(t, db, "kb-team", 0, 0, 0) + kb.Permission = string(entity.TenantPermissionTeam) + if err := db.Save(kb).Error; err != nil { + t.Fatalf("failed to update test kb permission: %v", err) + } + + if dao.Accessible("kb-team", "stranger-1") { + t.Fatal("unrelated user should not access a team knowledgebase") + } +} From 7ed6854648c06d10d2a902019098a0c2998e7d88 Mon Sep 17 00:00:00 2001 From: Jonathanchang31 <55106972+jonathanchang31@users.noreply.github.com> Date: Thu, 4 Jun 2026 20:10:06 +0200 Subject: [PATCH 2/4] fix: Bot reviews code --- internal/dao/kb.go | 11 +++++------ internal/dao/kb_test.go | 17 +++++++++++++++++ 2 files changed, 22 insertions(+), 6 deletions(-) diff --git a/internal/dao/kb.go b/internal/dao/kb.go index 59cf1b9bf8d..010d1221cc3 100644 --- a/internal/dao/kb.go +++ b/internal/dao/kb.go @@ -248,9 +248,8 @@ func (dao *KnowledgebaseDAO) Accessible(kbID, userID string) bool { Joins("LEFT JOIN user_tenant ON user_tenant.tenant_id = knowledgebase.tenant_id AND user_tenant.user_id = ?", userID). Where(`knowledgebase.id = ? AND knowledgebase.status = ? AND ( knowledgebase.created_by = ? OR - knowledgebase.tenant_id = ? OR (knowledgebase.permission = ? AND user_tenant.user_id IS NOT NULL) - )`, kbID, string(entity.StatusValid), userID, userID, string(entity.TenantPermissionTeam)). + )`, kbID, string(entity.StatusValid), userID, string(entity.TenantPermissionTeam)). Count(&count).Error if err != nil { @@ -327,7 +326,7 @@ func (dao *KnowledgebaseDAO) AtomicIncreaseDocNumByID(kbID string) error { return DB.Model(&entity.Knowledgebase{}). Where("id = ?", kbID). Updates(map[string]interface{}{ - "doc_num": DB.Raw("doc_num + 1"), + "doc_num": DB.Raw("doc_num + 1"), }).Error } @@ -337,9 +336,9 @@ func (dao *KnowledgebaseDAO) DecreaseDocumentNum(kbID string, docNum, chunkNum, return DB.Model(&entity.Knowledgebase{}). Where("id = ?", kbID). Updates(map[string]interface{}{ - "doc_num": DB.Raw("doc_num - ?", docNum), - "chunk_num": DB.Raw("chunk_num - ?", chunkNum), - "token_num": DB.Raw("token_num - ?", tokenNum), + "doc_num": DB.Raw("doc_num - ?", docNum), + "chunk_num": DB.Raw("chunk_num - ?", chunkNum), + "token_num": DB.Raw("token_num - ?", tokenNum), }).Error } diff --git a/internal/dao/kb_test.go b/internal/dao/kb_test.go index dca86849ad3..c04532fbcb6 100644 --- a/internal/dao/kb_test.go +++ b/internal/dao/kb_test.go @@ -163,6 +163,23 @@ func TestKnowledgebaseDAO_AccessibleRejectsTenantMemberForPrivateKnowledgebase(t } } +func TestKnowledgebaseDAO_AccessibleRejectsTenantIDMatchWhenNotCreatorForPrivateKnowledgebase(t *testing.T) { + db := setupKBTestDB(t) + pushDB(t, db) + dao := NewKnowledgebaseDAO() + + kb := testKnowledgebase(t, db, "kb-private-tenant-match", 0, 0, 0) + kb.CreatedBy = "member-creator" + kb.Permission = string(entity.TenantPermissionMe) + if err := db.Save(kb).Error; err != nil { + t.Fatalf("failed to update test kb creator: %v", err) + } + + if dao.Accessible("kb-private-tenant-match", "tenant-1") { + t.Fatal("non-creator should not access private knowledgebase even when tenant_id matches requester") + } +} + func TestKnowledgebaseDAO_AccessibleAllowsTenantMemberForTeamKnowledgebase(t *testing.T) { db := setupKBTestDB(t) pushDB(t, db) From c53353e3ce872a990e41ba89e063d3176d351968 Mon Sep 17 00:00:00 2001 From: Jonathanchang31 <55106972+jonathanchang31@users.noreply.github.com> Date: Fri, 5 Jun 2026 06:23:07 +0200 Subject: [PATCH 3/4] fix: serveral minor errors --- internal/dao/kb.go | 2 +- internal/dao/kb_test.go | 29 +++++++++++++++++++++++++++++ 2 files changed, 30 insertions(+), 1 deletion(-) diff --git a/internal/dao/kb.go b/internal/dao/kb.go index 010d1221cc3..f0f2f59718d 100644 --- a/internal/dao/kb.go +++ b/internal/dao/kb.go @@ -245,7 +245,7 @@ func (dao *KnowledgebaseDAO) GetDetail(kbID string) (*entity.KnowledgebaseDetail func (dao *KnowledgebaseDAO) Accessible(kbID, userID string) bool { var count int64 err := DB.Table("knowledgebase"). - Joins("LEFT JOIN user_tenant ON user_tenant.tenant_id = knowledgebase.tenant_id AND user_tenant.user_id = ?", userID). + Joins("LEFT JOIN user_tenant ON user_tenant.tenant_id = knowledgebase.tenant_id AND user_tenant.user_id = ? AND user_tenant.status = ?", userID, string(entity.StatusValid)). Where(`knowledgebase.id = ? AND knowledgebase.status = ? AND ( knowledgebase.created_by = ? OR (knowledgebase.permission = ? AND user_tenant.user_id IS NOT NULL) diff --git a/internal/dao/kb_test.go b/internal/dao/kb_test.go index c04532fbcb6..fbc3df4018f 100644 --- a/internal/dao/kb_test.go +++ b/internal/dao/kb_test.go @@ -197,6 +197,35 @@ func TestKnowledgebaseDAO_AccessibleAllowsTenantMemberForTeamKnowledgebase(t *te } } +func TestKnowledgebaseDAO_AccessibleRejectsRevokedTenantMemberForTeamKnowledgebase(t *testing.T) { + db := setupKBTestDB(t) + pushDB(t, db) + dao := NewKnowledgebaseDAO() + + kb := testKnowledgebase(t, db, "kb-team", 0, 0, 0) + kb.Permission = string(entity.TenantPermissionTeam) + if err := db.Save(kb).Error; err != nil { + t.Fatalf("failed to update test kb permission: %v", err) + } + + invalidStatus := string(entity.StatusInvalid) + revokedMembership := &entity.UserTenant{ + ID: "ut-revoked", + UserID: "revoked-member", + TenantID: "tenant-1", + Role: "normal", + InvitedBy: "tenant-1", + Status: &invalidStatus, + } + if err := db.Create(revokedMembership).Error; err != nil { + t.Fatalf("failed to create revoked user_tenant: %v", err) + } + + if dao.Accessible("kb-team", "revoked-member") { + t.Fatal("revoked tenant member should not access a team knowledgebase") + } +} + func TestKnowledgebaseDAO_AccessibleRejectsUnrelatedUserForTeamKnowledgebase(t *testing.T) { db := setupKBTestDB(t) pushDB(t, db) From 4b79ce9c2cb36a2ebb04b800b76c502a9d92f71d Mon Sep 17 00:00:00 2001 From: Jonathanchang31 <55106972+jonathanchang31@users.noreply.github.com> Date: Fri, 5 Jun 2026 06:51:28 +0200 Subject: [PATCH 4/4] fix: final error before merge --- internal/dao/kb.go | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/internal/dao/kb.go b/internal/dao/kb.go index f0f2f59718d..36e6078cda4 100644 --- a/internal/dao/kb.go +++ b/internal/dao/kb.go @@ -326,7 +326,7 @@ func (dao *KnowledgebaseDAO) AtomicIncreaseDocNumByID(kbID string) error { return DB.Model(&entity.Knowledgebase{}). Where("id = ?", kbID). Updates(map[string]interface{}{ - "doc_num": DB.Raw("doc_num + 1"), + "doc_num": DB.Raw("doc_num + 1"), }).Error } @@ -336,9 +336,9 @@ func (dao *KnowledgebaseDAO) DecreaseDocumentNum(kbID string, docNum, chunkNum, return DB.Model(&entity.Knowledgebase{}). Where("id = ?", kbID). Updates(map[string]interface{}{ - "doc_num": DB.Raw("doc_num - ?", docNum), - "chunk_num": DB.Raw("chunk_num - ?", chunkNum), - "token_num": DB.Raw("token_num - ?", tokenNum), + "doc_num": DB.Raw("doc_num - ?", docNum), + "chunk_num": DB.Raw("chunk_num - ?", chunkNum), + "token_num": DB.Raw("token_num - ?", tokenNum), }).Error }