From ad2feaf810faf90d10f0a00e42efda6530a0fc98 Mon Sep 17 00:00:00 2001 From: afjcjsbx Date: Sun, 8 Mar 2026 16:15:08 +0100 Subject: [PATCH] fix encapsulation bypass on registry internals --- pkg/tools/registry.go | 39 ++++++++++++++++++ pkg/tools/search_tool.go | 88 ++++++++++++++++------------------------ 2 files changed, 75 insertions(+), 52 deletions(-) diff --git a/pkg/tools/registry.go b/pkg/tools/registry.go index a1107039d..91af0d6f9 100644 --- a/pkg/tools/registry.go +++ b/pkg/tools/registry.go @@ -88,6 +88,45 @@ func (r *ToolRegistry) TickTTL() { } } +// Version returns the current registry version (atomically). +func (r *ToolRegistry) Version() uint64 { + return r.version.Load() +} + +// HiddenToolSnapshot holds a consistent snapshot of hidden tools and the +// registry version at which it was taken. Used by BM25SearchTool cache. +type HiddenToolSnapshot struct { + Docs []HiddenToolDoc + Version uint64 +} + +// HiddenToolDoc is a lightweight representation of a hidden tool for search indexing. +type HiddenToolDoc struct { + Name string + Description string +} + +// SnapshotHiddenTools returns all non-core tools and the current registry +// version under a single read-lock, guaranteeing consistency between the +// two values. +func (r *ToolRegistry) SnapshotHiddenTools() HiddenToolSnapshot { + r.mu.RLock() + defer r.mu.RUnlock() + docs := make([]HiddenToolDoc, 0, len(r.tools)) + for name, entry := range r.tools { + if !entry.IsCore { + docs = append(docs, HiddenToolDoc{ + Name: name, + Description: entry.Tool.Description(), + }) + } + } + return HiddenToolSnapshot{ + Docs: docs, + Version: r.version.Load(), + } +} + func (r *ToolRegistry) Get(name string) (Tool, bool) { r.mu.RLock() defer r.mu.RUnlock() diff --git a/pkg/tools/search_tool.go b/pkg/tools/search_tool.go index 338c6b919..3a59a4f12 100644 --- a/pkg/tools/search_tool.go +++ b/pkg/tools/search_tool.go @@ -214,85 +214,69 @@ type bm25CachedEngine struct { engine *utils.BM25Engine[searchDoc] } +// snapshotToSearchDocs converts a HiddenToolSnapshot to BM25 searchDoc slice. +func snapshotToSearchDocs(snap HiddenToolSnapshot) []searchDoc { + docs := make([]searchDoc, len(snap.Docs)) + for i, d := range snap.Docs { + docs[i] = searchDoc{Name: d.Name, Description: d.Description} + } + return docs +} + +// buildBM25Engine creates a BM25Engine from a slice of searchDocs. +func buildBM25Engine(docs []searchDoc) *utils.BM25Engine[searchDoc] { + return utils.NewBM25Engine( + docs, + func(doc searchDoc) string { + return doc.Name + " " + doc.Description + }, + ) +} + // getOrBuildEngine returns a cached BM25 engine, rebuilding it only when // the registry version has changed (new tools registered). func (t *BM25SearchTool) getOrBuildEngine() *bm25CachedEngine { - // Fast path: optimistic check without locking the registry. - // If the version hasn't changed, the cache is still valid. - if t.cachedEngine != nil && t.cacheVersion == t.registry.version.Load() { + // Fast path: optimistic check without locking. + if t.cachedEngine != nil && t.cacheVersion == t.registry.Version() { return t.cachedEngine } t.cacheMu.Lock() defer t.cacheMu.Unlock() - // Read version inside the registry RLock so the snapshot and version - // are guaranteed to be consistent (no TOCTOU between Load and RLock). - t.registry.mu.RLock() - snapshotVersion := t.registry.version.Load() - snapshot := make([]searchDoc, 0, len(t.registry.tools)) - for name, entry := range t.registry.tools { - if !entry.IsCore { - snapshot = append(snapshot, searchDoc{ - Name: name, - Description: entry.Tool.Description(), - }) - } - } - t.registry.mu.RUnlock() + // Snapshot + version are read under a single registry RLock, + // guaranteeing consistency (no TOCTOU). + snap := t.registry.SnapshotHiddenTools() // Re-check: another goroutine may have rebuilt while we waited for cacheMu. - if t.cachedEngine != nil && t.cacheVersion == snapshotVersion { + if t.cachedEngine != nil && t.cacheVersion == snap.Version { return t.cachedEngine } - if len(snapshot) == 0 { + docs := snapshotToSearchDocs(snap) + if len(docs) == 0 { t.cachedEngine = nil - t.cacheVersion = snapshotVersion + t.cacheVersion = snap.Version return nil } - engine := utils.NewBM25Engine( - snapshot, - func(doc searchDoc) string { - return doc.Name + " " + doc.Description - }, - ) - - cached := &bm25CachedEngine{engine: engine} + cached := &bm25CachedEngine{engine: buildBM25Engine(docs)} t.cachedEngine = cached - t.cacheVersion = snapshotVersion + t.cacheVersion = snap.Version return cached } // SearchBM25 ranks hidden tools against query using BM25 via utils.BM25Engine. -// The corpus snapshot is built under the registry read-lock, then released -// before scoring so the lock is not held during CPU-intensive work. +// This non-cached variant rebuilds the engine on every call. Used by tests +// and any code that doesn't hold a BM25SearchTool instance. func (r *ToolRegistry) SearchBM25(query string, maxSearchResults int) []ToolSearchResult { - r.mu.RLock() - snapshot := make([]searchDoc, 0, len(r.tools)) - for name, entry := range r.tools { - if !entry.IsCore { - snapshot = append(snapshot, searchDoc{ - Name: name, - Description: entry.Tool.Description(), - }) - } - } - r.mu.RUnlock() - - if len(snapshot) == 0 { + snap := r.SnapshotHiddenTools() + docs := snapshotToSearchDocs(snap) + if len(docs) == 0 { return nil } - engine := utils.NewBM25Engine( - snapshot, - func(doc searchDoc) string { - return doc.Name + " " + doc.Description - }, - ) - - ranked := engine.Search(query, maxSearchResults) + ranked := buildBM25Engine(docs).Search(query, maxSearchResults) if len(ranked) == 0 { return nil }