diff --git a/internal/agent/tools/database_query.go b/internal/agent/tools/database_query.go index 118f854d..b3aaa972 100644 --- a/internal/agent/tools/database_query.go +++ b/internal/agent/tools/database_query.go @@ -255,21 +255,13 @@ func (t *DatabaseQueryTool) Execute(ctx context.Context, args json.RawMessage) ( // validateAndSecureSQL validates the SQL query and injects tenant_id conditions func (t *DatabaseQueryTool) validateAndSecureSQL(sqlQuery string, tenantID uint64) (string, error) { - kbIDs := t.searchTargets.GetAllKnowledgeBaseIDs() - var knowledgeIDs []string - for _, target := range t.searchTargets { - if target.Type == types.SearchTargetTypeKnowledge && len(target.KnowledgeIDs) > 0 { - knowledgeIDs = append(knowledgeIDs, target.KnowledgeIDs...) - } - } - securedSQL, validationResult, err := utils.ValidateAndSecureSQL( sqlQuery, utils.WithSecurityDefaults(tenantID), utils.WithSoftDeleteFilter("knowledge_bases", "knowledges", "chunks"), utils.WithHiddenKBFilter(), utils.WithInjectionRiskCheck(), - utils.WithSearchScopeFilter(kbIDs, knowledgeIDs), + utils.WithSearchScopes(searchScopesFromTargets(t.searchTargets)), ) if err != nil { return "", err @@ -286,6 +278,21 @@ func (t *DatabaseQueryTool) validateAndSecureSQL(sqlQuery string, tenantID uint6 return securedSQL, nil } +func searchScopesFromTargets(searchTargets types.SearchTargets) []utils.SearchScope { + scopes := make([]utils.SearchScope, 0, len(searchTargets)) + for _, target := range searchTargets { + if target == nil || target.KnowledgeBaseID == "" { + continue + } + scopes = append(scopes, utils.SearchScope{ + KnowledgeBaseID: target.KnowledgeBaseID, + KnowledgeIDs: append([]string(nil), target.KnowledgeIDs...), + TagIDs: append([]string(nil), target.TagIDs...), + }) + } + return scopes +} + // formatQueryResults formats query results into readable text func (t *DatabaseQueryTool) formatQueryResults( columns []string, diff --git a/internal/agent/tools/get_document_info.go b/internal/agent/tools/get_document_info.go index da8cc5ed..a28cf584 100644 --- a/internal/agent/tools/get_document_info.go +++ b/internal/agent/tools/get_document_info.go @@ -144,6 +144,17 @@ func (t *GetDocumentInfoTool) Execute(ctx context.Context, args json.RawMessage) mu.Unlock() return } + allowed, scopeErr := searchTargetsAllowKnowledgeID(ctx, t.searchTargets, chunk.KnowledgeID, chunk.KnowledgeBaseID, t.knowledgeService) + if scopeErr != nil || !allowed { + mu.Lock() + if scopeErr != nil { + results["faq:"+id] = &docInfo{err: fmt.Errorf("failed to validate FAQ scope: %v", scopeErr)} + } else { + results["faq:"+id] = &docInfo{err: fmt.Errorf("FAQ entry %s is not within the current @mention scope", id)} + } + mu.Unlock() + return + } var meta *types.FAQChunkMetadata if chunk.ChunkType == types.ChunkTypeFAQ { meta, _ = chunk.FAQMetadata() @@ -179,6 +190,17 @@ func (t *GetDocumentInfoTool) Execute(ctx context.Context, args json.RawMessage) mu.Unlock() return } + allowed, scopeErr := searchTargetsAllowKnowledgeID(ctx, t.searchTargets, knowledge.ID, knowledge.KnowledgeBaseID, t.knowledgeService) + if scopeErr != nil || !allowed { + mu.Lock() + if scopeErr != nil { + results[id] = &docInfo{err: fmt.Errorf("failed to validate document scope: %v", scopeErr)} + } else { + results[id] = &docInfo{err: fmt.Errorf("document %s is not within the current @mention scope", knowledge.ID)} + } + mu.Unlock() + return + } // Use knowledge's actual tenant_id for chunk query (supports cross-tenant shared KB). // Keep chunk-type filter aligned with list_knowledge_chunks so the diff --git a/internal/agent/tools/list_knowledge_chunks.go b/internal/agent/tools/list_knowledge_chunks.go index 5bbee4aa..6f6bd55a 100644 --- a/internal/agent/tools/list_knowledge_chunks.go +++ b/internal/agent/tools/list_knowledge_chunks.go @@ -133,6 +133,19 @@ func (t *ListKnowledgeChunksTool) Execute(ctx context.Context, args json.RawMess Error: fmt.Sprintf("Knowledge base %s is not accessible", knowledge.KnowledgeBaseID), }, fmt.Errorf("knowledge base not in search targets") } + allowed, err := searchTargetsAllowKnowledgeID(ctx, t.searchTargets, knowledge.ID, knowledge.KnowledgeBaseID, t.knowledgeService) + if err != nil { + return &types.ToolResult{ + Success: false, + Error: fmt.Sprintf("failed to validate knowledge scope: %v", err), + }, err + } + if !allowed { + return &types.ToolResult{ + Success: false, + Error: fmt.Sprintf("Knowledge %s is not within the current @mention scope", knowledge.ID), + }, fmt.Errorf("knowledge not in search target scope") + } // Use the knowledge's actual tenant_id for chunk query (supports cross-tenant shared KB) effectiveTenantID := knowledge.TenantID @@ -296,6 +309,19 @@ func (t *ListKnowledgeChunksTool) executeByChunkID(ctx context.Context, chunkID Error: fmt.Sprintf("knowledge base %s is not accessible", chunk.KnowledgeBaseID), }, fmt.Errorf("knowledge base not in search targets") } + allowed, scopeErr := searchTargetsAllowKnowledgeID(ctx, t.searchTargets, chunk.KnowledgeID, chunk.KnowledgeBaseID, t.knowledgeService) + if scopeErr != nil { + return &types.ToolResult{ + Success: false, + Error: fmt.Sprintf("failed to validate chunk scope: %v", scopeErr), + }, scopeErr + } + if !allowed { + return &types.ToolResult{ + Success: false, + Error: fmt.Sprintf("chunk %s is not within the current @mention scope", chunk.ID), + }, fmt.Errorf("chunk not in search target scope") + } chunks := []*types.Chunk{chunk} if chunk.ImageInfo == "" { diff --git a/internal/agent/tools/scope_authorization.go b/internal/agent/tools/scope_authorization.go new file mode 100644 index 00000000..8937c84f --- /dev/null +++ b/internal/agent/tools/scope_authorization.go @@ -0,0 +1,86 @@ +package tools + +import ( + "context" + + "github.com/Tencent/WeKnora/internal/types" + "github.com/Tencent/WeKnora/internal/types/interfaces" +) + +type knowledgeTagsFetcher func(context.Context, []string) (map[string][]*types.KnowledgeTag, error) + +func searchTargetsAllowKnowledgeID( + ctx context.Context, + searchTargets types.SearchTargets, + knowledgeID string, + kbID string, + knowledgeService interfaces.KnowledgeService, +) (bool, error) { + if knowledgeID == "" || kbID == "" { + return false, nil + } + + var tagIDs []string + matchedKB := false + for _, target := range searchTargets { + if target == nil || target.KnowledgeBaseID != kbID { + continue + } + matchedKB = true + if target.Type == types.SearchTargetTypeKnowledgeBase && len(target.TagIDs) == 0 { + return true, nil + } + for _, allowedID := range target.KnowledgeIDs { + if allowedID == knowledgeID { + return true, nil + } + } + tagIDs = append(tagIDs, target.TagIDs...) + } + if !matchedKB || len(tagIDs) == 0 || knowledgeService == nil { + return false, nil + } + + matches, err := knowledgeIDsMatchingAnyTag(ctx, []string{knowledgeID}, tagIDs, knowledgeService.GetKnowledgeTags) + if err != nil { + return false, err + } + return matches[knowledgeID], nil +} + +func knowledgeIDsMatchingAnyTag( + ctx context.Context, + knowledgeIDs []string, + tagIDs []string, + fetchTags knowledgeTagsFetcher, +) (map[string]bool, error) { + result := make(map[string]bool) + if len(knowledgeIDs) == 0 || len(tagIDs) == 0 || fetchTags == nil { + return result, nil + } + + uniqueKnowledgeIDs := dedupNonEmptyStrings(knowledgeIDs) + uniqueTagIDs := dedupNonEmptyStrings(tagIDs) + if len(uniqueKnowledgeIDs) == 0 || len(uniqueTagIDs) == 0 { + return result, nil + } + + tagSet := make(map[string]bool, len(uniqueTagIDs)) + for _, tagID := range uniqueTagIDs { + tagSet[tagID] = true + } + + tagMap, err := fetchTags(ctx, uniqueKnowledgeIDs) + if err != nil { + return nil, err + } + for _, knowledgeID := range uniqueKnowledgeIDs { + for _, tag := range tagMap[knowledgeID] { + if tag != nil && tagSet[tag.ID] { + result[knowledgeID] = true + break + } + } + } + return result, nil +} diff --git a/internal/agent/tools/scope_authorization_test.go b/internal/agent/tools/scope_authorization_test.go new file mode 100644 index 00000000..8e8a9325 --- /dev/null +++ b/internal/agent/tools/scope_authorization_test.go @@ -0,0 +1,78 @@ +package tools + +import ( + "context" + "testing" + + "github.com/Tencent/WeKnora/internal/types" +) + +func testKnowledgeTag(id string) *types.KnowledgeTag { + return &types.KnowledgeTag{ID: id} +} + +func TestKnowledgeIDsMatchingAnyTag(t *testing.T) { + matches, err := knowledgeIDsMatchingAnyTag( + context.Background(), + []string{"doc-1", "doc-2"}, + []string{"tag-a"}, + func(_ context.Context, ids []string) (map[string][]*types.KnowledgeTag, error) { + return map[string][]*types.KnowledgeTag{ + "doc-1": []*types.KnowledgeTag{testKnowledgeTag("tag-z")}, + "doc-2": []*types.KnowledgeTag{testKnowledgeTag("tag-a")}, + }, nil + }, + ) + if err != nil { + t.Fatalf("knowledgeIDsMatchingAnyTag() error = %v", err) + } + if matches["doc-1"] { + t.Fatalf("doc-1 should not match tag-a") + } + if !matches["doc-2"] { + t.Fatalf("doc-2 should match tag-a") + } +} + +func TestPagePassesWikiScope_TagScope(t *testing.T) { + page := &types.WikiPage{ + Slug: "entity/acme", + SourceRefs: []string{"doc-1|Acme intro", "doc-2"}, + } + + passes, err := pagePassesWikiScope( + context.Background(), + page, + WikiScope{KnowledgeBaseID: "kb-1", TagIDs: []string{"tag-a"}}, + func(_ context.Context, ids []string) (map[string][]*types.KnowledgeTag, error) { + return map[string][]*types.KnowledgeTag{ + "doc-1": []*types.KnowledgeTag{testKnowledgeTag("tag-z")}, + "doc-2": []*types.KnowledgeTag{testKnowledgeTag("tag-a")}, + }, nil + }, + ) + if err != nil { + t.Fatalf("pagePassesWikiScope() error = %v", err) + } + if !passes { + t.Fatalf("page should pass when one source document has the mentioned tag") + } + + passes, err = pagePassesWikiScope( + context.Background(), + page, + WikiScope{KnowledgeBaseID: "kb-1", TagIDs: []string{"tag-missing"}}, + func(_ context.Context, ids []string) (map[string][]*types.KnowledgeTag, error) { + return map[string][]*types.KnowledgeTag{ + "doc-1": []*types.KnowledgeTag{testKnowledgeTag("tag-z")}, + "doc-2": []*types.KnowledgeTag{testKnowledgeTag("tag-a")}, + }, nil + }, + ) + if err != nil { + t.Fatalf("pagePassesWikiScope() error = %v", err) + } + if passes { + t.Fatalf("page should be filtered when no source document has the mentioned tag") + } +} diff --git a/internal/agent/tools/wiki_tools.go b/internal/agent/tools/wiki_tools.go index d7a988f9..d1ef8c0d 100644 --- a/internal/agent/tools/wiki_tools.go +++ b/internal/agent/tools/wiki_tools.go @@ -99,9 +99,12 @@ func renderIndexOverviewForAgent(resp *types.WikiIndexResponse) string { // When non-empty, a wiki page is only surfaced if its SourceRefs // intersect this set. Used to honour user-level @mentions that pin the // search to specific documents inside a KB. +// - TagIDs: OPTIONAL whitelist of document tags. When non-empty, a wiki page +// is only surfaced if at least one of its SourceRefs belongs to any tag. type WikiScope struct { KnowledgeBaseID string KnowledgeIDs []string + TagIDs []string } // NewWikiScopesFromKBIDs is a convenience constructor for callers that only @@ -232,15 +235,50 @@ func pageIntersectsKnowledgeIDs(page *types.WikiPage, allowed map[string]bool) b return false } -type wikiReadPageTool struct { - BaseTool - wikiService interfaces.WikiPageService - scopes []WikiScope - seenLinks map[string]bool - mu sync.Mutex +func pagePassesWikiScope( + ctx context.Context, + page *types.WikiPage, + scope WikiScope, + fetchTags knowledgeTagsFetcher, +) (bool, error) { + if allowed, has := scopeKnowledgeFilter(scope); has { + if !pageIntersectsKnowledgeIDs(page, allowed) { + return false, nil + } + } + + tagIDs := dedupNonEmptyStrings(scope.TagIDs) + if len(tagIDs) == 0 { + return true, nil + } + if isStructuralPage(page) { + return true, nil + } + sourceKnowledgeIDs := extractSourceKnowledgeIDs(page) + if len(sourceKnowledgeIDs) == 0 { + return true, nil + } + matches, err := knowledgeIDsMatchingAnyTag(ctx, sourceKnowledgeIDs, tagIDs, fetchTags) + if err != nil { + return false, err + } + return len(matches) > 0, nil } -func NewWikiReadPageTool(wikiService interfaces.WikiPageService, scopes []WikiScope) types.Tool { +type wikiReadPageTool struct { + BaseTool + wikiService interfaces.WikiPageService + knowledgeService interfaces.KnowledgeService + scopes []WikiScope + seenLinks map[string]bool + mu sync.Mutex +} + +func NewWikiReadPageTool( + wikiService interfaces.WikiPageService, + knowledgeService interfaces.KnowledgeService, + scopes []WikiScope, +) types.Tool { return &wikiReadPageTool{ BaseTool: NewBaseTool( ToolWikiReadPage, @@ -263,9 +301,10 @@ When the same slug exists in multiple knowledge bases, all matching pages are re "required": ["slugs"] }`), ), - wikiService: wikiService, - scopes: scopes, - seenLinks: make(map[string]bool), + wikiService: wikiService, + knowledgeService: knowledgeService, + scopes: scopes, + seenLinks: make(map[string]bool), } } @@ -419,6 +458,10 @@ func (t *wikiReadPageTool) Execute(ctx context.Context, args json.RawMessage) (* // Track slugs that were found in the raw lookup but filtered out by a // knowledge_ids whitelist, so we can surface a clearer error. filteredOut := make(map[string][]string) // slug -> list of KB IDs where filtered + var fetchTags knowledgeTagsFetcher + if t.knowledgeService != nil { + fetchTags = t.knowledgeService.GetKnowledgeTags + } for _, slug := range slugsToFetch { var hits []struct { page *types.WikiPage @@ -438,13 +481,16 @@ func (t *wikiReadPageTool) Execute(ctx context.Context, args json.RawMessage) (* actualKBID = page.KnowledgeBaseID } - // Apply server-enforced knowledge-ID scope (silent; never exposed - // to the model as a tool argument). - if allowed, has := scopeKnowledgeFilter(sc); has { - if !pageIntersectsKnowledgeIDs(page, allowed) { - filteredOut[slug] = append(filteredOut[slug], actualKBID) - continue - } + // Apply server-enforced source-document / tag scope (silent; never + // exposed to the model as a tool argument). + passesScope, scopeErr := pagePassesWikiScope(ctx, page, sc, fetchTags) + if scopeErr != nil { + errs = append(errs, fmt.Sprintf("Failed to validate wiki scope for '%s' in KB %s: %v", slug, actualKBID, scopeErr)) + continue + } + if !passesScope { + filteredOut[slug] = append(filteredOut[slug], actualKBID) + continue } hits = append(hits, struct { @@ -513,13 +559,18 @@ func (t *wikiReadPageTool) Execute(ctx context.Context, args json.RawMessage) (* type wikiSearchTool struct { BaseTool - wikiService interfaces.WikiPageService - scopes []WikiScope - seenSlugs map[string]bool - mu sync.Mutex + wikiService interfaces.WikiPageService + knowledgeService interfaces.KnowledgeService + scopes []WikiScope + seenSlugs map[string]bool + mu sync.Mutex } -func NewWikiSearchTool(wikiService interfaces.WikiPageService, scopes []WikiScope) types.Tool { +func NewWikiSearchTool( + wikiService interfaces.WikiPageService, + knowledgeService interfaces.KnowledgeService, + scopes []WikiScope, +) types.Tool { return &wikiSearchTool{ BaseTool: NewBaseTool( ToolWikiSearch, @@ -553,9 +604,10 @@ Use this to find relevant wiki pages when you don't know the exact slug.`, "required": ["queries"] }`), ), - wikiService: wikiService, - scopes: scopes, - seenSlugs: make(map[string]bool), + wikiService: wikiService, + knowledgeService: knowledgeService, + scopes: scopes, + seenSlugs: make(map[string]bool), } } @@ -607,6 +659,10 @@ func (t *wikiSearchTool) Execute(ctx context.Context, args json.RawMessage) (*ty page *types.WikiPage kbID string } + var fetchTags knowledgeTagsFetcher + if t.knowledgeService != nil { + fetchTags = t.knowledgeService.GetKnowledgeTags + } for _, query := range queriesToRun { var allHits []searchHit @@ -616,8 +672,6 @@ func (t *wikiSearchTool) Execute(ctx context.Context, args json.RawMessage) (*ty if kbID == "" { continue } - allowed, hasFilter := scopeKnowledgeFilter(sc) - pages, err := t.wikiService.SearchPages(ctx, kbID, query, params.Limit) if err != nil { continue @@ -626,7 +680,8 @@ func (t *wikiSearchTool) Execute(ctx context.Context, args json.RawMessage) (*ty if p == nil { continue } - if hasFilter && !pageIntersectsKnowledgeIDs(p, allowed) { + passesScope, scopeErr := pagePassesWikiScope(ctx, p, sc, fetchTags) + if scopeErr != nil || !passesScope { filteredCount++ continue } diff --git a/internal/application/service/agent_service.go b/internal/application/service/agent_service.go index 9afd3769..4460195f 100644 --- a/internal/application/service/agent_service.go +++ b/internal/application/service/agent_service.go @@ -440,6 +440,9 @@ func (s *agentService) registerTools( if target.Type == types.SearchTargetTypeKnowledge && len(target.KnowledgeIDs) > 0 { scope.KnowledgeIDs = append([]string(nil), target.KnowledgeIDs...) } + if len(target.TagIDs) > 0 { + scope.TagIDs = append([]string(nil), target.TagIDs...) + } wikiScopes = append(wikiScopes, scope) } } @@ -611,9 +614,9 @@ func (s *agentService) registerTools( // Wiki tools — only registered when wiki KBs are detected case tools.ToolWikiReadPage: - toolToRegister = tools.NewWikiReadPageTool(s.wikiPageService, wikiScopes) + toolToRegister = tools.NewWikiReadPageTool(s.wikiPageService, s.knowledgeService, wikiScopes) case tools.ToolWikiSearch: - toolToRegister = tools.NewWikiSearchTool(s.wikiPageService, wikiScopes) + toolToRegister = tools.NewWikiSearchTool(s.wikiPageService, s.knowledgeService, wikiScopes) case tools.ToolWikiReadSourceDoc: toolToRegister = tools.NewWikiReadSourceDocTool(s.knowledgeService, s.chunkService) case tools.ToolWikiFlagIssue: diff --git a/internal/utils/inject.go b/internal/utils/inject.go index 8a31f83e..a2bee90d 100644 --- a/internal/utils/inject.go +++ b/internal/utils/inject.go @@ -144,6 +144,15 @@ type sqlValidator struct { enableSearchScopeFilter bool searchScopeKBIDs []string searchScopeKnowledgeIDs []string + searchScopes []SearchScope +} + +// SearchScope describes one allowed knowledge scope for SQL query injection. +// Empty KnowledgeIDs and TagIDs means the whole KB is in scope. +type SearchScope struct { + KnowledgeBaseID string + KnowledgeIDs []string + TagIDs []string } // ParseSQL parses a SQL statement using pg_query_go and extracts table names, select fields, and where fields @@ -630,6 +639,18 @@ func WithSearchScopeFilter(kbIDs []string, knowledgeIDs []string) SQLValidationO } } +// WithSearchScopes restricts queries using structured OR scopes. Each scope can +// represent a full KB, specific documents, or a tag-constrained KB. +func WithSearchScopes(scopes []SearchScope) SQLValidationOption { + return func(v *sqlValidator) { + if len(scopes) == 0 { + return + } + v.enableSearchScopeFilter = true + v.searchScopes = append([]SearchScope(nil), scopes...) + } +} + // WithSecurityDefaults applies a comprehensive set of security validations func WithSecurityDefaults(tenantID uint64) SQLValidationOption { return func(v *sqlValidator) { @@ -1015,7 +1036,13 @@ func (v *sqlValidator) injectHiddenKBFilter(sql string, tablesInQuery map[string // and (optionally) specific knowledge documents. func (v *sqlValidator) injectSearchScopeConditions(sql string, tablesInQuery map[string]string) string { if !v.enableSearchScopeFilter || len(v.searchScopeKBIDs) == 0 { - return sql + if len(v.searchScopes) == 0 { + return sql + } + } + + if len(v.searchScopes) > 0 { + return v.injectStructuredSearchScopeConditions(sql, tablesInQuery) } quotedKBIDs := quoteStringSlice(v.searchScopeKBIDs) @@ -1050,6 +1077,116 @@ func (v *sqlValidator) injectSearchScopeConditions(sql string, tablesInQuery map return InjectAndConditions(sql, strings.Join(conditions, " AND ")) } +func (v *sqlValidator) injectStructuredSearchScopeConditions(sql string, tablesInQuery map[string]string) string { + var conditions []string + + if alias, ok := tablesInQuery["knowledge_bases"]; ok { + if cond := buildKnowledgeBaseScopeCondition(alias, v.searchScopes); cond != "" { + conditions = append(conditions, cond) + } + } + if alias, ok := tablesInQuery["knowledges"]; ok { + if cond := buildKnowledgeScopeCondition(alias, v.searchScopes); cond != "" { + conditions = append(conditions, cond) + } + } + if alias, ok := tablesInQuery["chunks"]; ok { + if cond := buildChunkScopeCondition(alias, v.searchScopes); cond != "" { + conditions = append(conditions, cond) + } + } + + if len(conditions) == 0 { + return sql + } + return InjectAndConditions(sql, strings.Join(conditions, " AND ")) +} + +func buildKnowledgeBaseScopeCondition(alias string, scopes []SearchScope) string { + kbIDs := uniqueScopeKBIDs(scopes) + if len(kbIDs) == 0 { + return "" + } + return fmt.Sprintf("%s.id IN (%s)", alias, strings.Join(quoteStringSlice(kbIDs), ", ")) +} + +func buildKnowledgeScopeCondition(alias string, scopes []SearchScope) string { + clauses := make([]string, 0, len(scopes)) + for _, scope := range scopes { + if scope.KnowledgeBaseID == "" { + continue + } + kbID := quoteString(scope.KnowledgeBaseID) + switch { + case len(scope.KnowledgeIDs) > 0: + clauses = append(clauses, fmt.Sprintf( + "(%s.knowledge_base_id = %s AND %s.id IN (%s))", + alias, kbID, alias, strings.Join(quoteStringSlice(scope.KnowledgeIDs), ", "), + )) + case len(scope.TagIDs) > 0: + clauses = append(clauses, fmt.Sprintf( + "(%s.knowledge_base_id = %s AND EXISTS (SELECT 1 FROM knowledge_tag_relations ktr WHERE ktr.knowledge_id = %s.id AND ktr.tag_id IN (%s)))", + alias, kbID, alias, strings.Join(quoteStringSlice(scope.TagIDs), ", "), + )) + default: + clauses = append(clauses, fmt.Sprintf("%s.knowledge_base_id = %s", alias, kbID)) + } + } + return joinOrClauses(clauses) +} + +func buildChunkScopeCondition(alias string, scopes []SearchScope) string { + clauses := make([]string, 0, len(scopes)) + for _, scope := range scopes { + if scope.KnowledgeBaseID == "" { + continue + } + kbID := quoteString(scope.KnowledgeBaseID) + switch { + case len(scope.KnowledgeIDs) > 0: + clauses = append(clauses, fmt.Sprintf( + "(%s.knowledge_base_id = %s AND %s.knowledge_id IN (%s))", + alias, kbID, alias, strings.Join(quoteStringSlice(scope.KnowledgeIDs), ", "), + )) + case len(scope.TagIDs) > 0: + clauses = append(clauses, fmt.Sprintf( + "(%s.knowledge_base_id = %s AND EXISTS (SELECT 1 FROM knowledge_tag_relations ktr WHERE ktr.knowledge_id = %s.knowledge_id AND ktr.tag_id IN (%s)))", + alias, kbID, alias, strings.Join(quoteStringSlice(scope.TagIDs), ", "), + )) + default: + clauses = append(clauses, fmt.Sprintf("%s.knowledge_base_id = %s", alias, kbID)) + } + } + return joinOrClauses(clauses) +} + +func uniqueScopeKBIDs(scopes []SearchScope) []string { + seen := make(map[string]bool, len(scopes)) + out := make([]string, 0, len(scopes)) + for _, scope := range scopes { + if scope.KnowledgeBaseID == "" || seen[scope.KnowledgeBaseID] { + continue + } + seen[scope.KnowledgeBaseID] = true + out = append(out, scope.KnowledgeBaseID) + } + return out +} + +func joinOrClauses(clauses []string) string { + if len(clauses) == 0 { + return "" + } + if len(clauses) == 1 { + return clauses[0] + } + return "(" + strings.Join(clauses, " OR ") + ")" +} + +func quoteString(s string) string { + return quoteStringSlice([]string{s})[0] +} + func quoteStringSlice(ss []string) []string { quoted := make([]string, len(ss)) for i, s := range ss { diff --git a/internal/utils/inject_test.go b/internal/utils/inject_test.go index 730c2fe0..f5a515a9 100644 --- a/internal/utils/inject_test.go +++ b/internal/utils/inject_test.go @@ -3,6 +3,7 @@ package utils import ( "encoding/json" "fmt" + "strings" "testing" ) @@ -449,6 +450,37 @@ func TestInjectAndConditions(t *testing.T) { } } +func TestValidateAndSecureSQL_WithStructuredSearchScopes(t *testing.T) { + securedSQL, validation, err := ValidateAndSecureSQL( + "SELECT id FROM chunks", + WithSearchScopes([]SearchScope{ + {KnowledgeBaseID: "kb-full"}, + {KnowledgeBaseID: "kb-doc", KnowledgeIDs: []string{"doc-1"}}, + {KnowledgeBaseID: "kb-tag", TagIDs: []string{"tag-a", "tag-b"}}, + }), + ) + if err != nil { + t.Fatalf("ValidateAndSecureSQL() error = %v", err) + } + if !validation.Valid { + t.Fatalf("expected validation to pass, got %#v", validation.Errors) + } + + for _, want := range []string{ + "chunks.knowledge_base_id = 'kb-full'", + "chunks.knowledge_base_id = 'kb-doc' AND chunks.knowledge_id IN ('doc-1')", + "chunks.knowledge_base_id = 'kb-tag' AND EXISTS", + "knowledge_tag_relations", + "ktr.knowledge_id = chunks.knowledge_id", + "ktr.tag_id IN ('tag-a', 'tag-b')", + " OR ", + } { + if !strings.Contains(securedSQL, want) { + t.Fatalf("secured SQL missing %q:\n%s", want, securedSQL) + } + } +} + func BenchmarkInjectAndConditions(b *testing.B) { const sql = "SELECT id, title FROM docs WHERE status = 'active' ORDER BY created_at LIMIT 50" for i := 0; i < b.N; i++ {