diff --git a/cmd/cosift/feedback.go b/cmd/cosift/feedback.go index 16a9901..50946d4 100644 --- a/cmd/cosift/feedback.go +++ b/cmd/cosift/feedback.go @@ -108,11 +108,9 @@ func (s *pebbleHTTP) handleFeedback(w http.ResponseWriter, r *http.Request) { // reward/penalty tallies joined to the query log by qid. This is the "make it // useful" surface — the labeled signal without an external pipeline. func (s *pebbleHTTP) handleFeedbackList(w http.ResponseWriter, r *http.Request) { - if want := s.cluster.PeerAuthToken; want != "" { - if r.Header.Get("Authorization") != "Bearer "+want { - writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") - return - } + if !peerTokenOK(r, s.cluster.PeerAuthToken) { + writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") + return } fp := feedbackLogPath() if fp == "" { diff --git a/cmd/cosift/querylog.go b/cmd/cosift/querylog.go index 6db7163..71147e0 100644 --- a/cmd/cosift/querylog.go +++ b/cmd/cosift/querylog.go @@ -90,12 +90,9 @@ func (s *pebbleHTTP) writeQueryLog(rec queryLogRec) { // handleQueryLog tails the query log. ?n=N (default 100, max 5000) returns the // last N JSON lines as a JSON array; ?raw=1 returns them as raw JSONL. func (s *pebbleHTTP) handleQueryLog(w http.ResponseWriter, r *http.Request) { - if want := s.cluster.PeerAuthToken; want != "" { - got := r.Header.Get("Authorization") - if got != "Bearer "+want { - writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") - return - } + if !peerTokenOK(r, s.cluster.PeerAuthToken) { + writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") + return } path := os.Getenv("COSIFT_QUERY_LOG") if path == "" { diff --git a/cmd/cosift/sec_findings_test.go b/cmd/cosift/sec_findings_test.go new file mode 100644 index 0000000..7486c26 --- /dev/null +++ b/cmd/cosift/sec_findings_test.go @@ -0,0 +1,176 @@ +package main + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "regexp" + "strconv" + "strings" + "testing" + + "github.com/pilot-protocol/cosift/internal/config" + "github.com/pilot-protocol/cosift/internal/index" +) + +// --------------------------------------------------------------------------- +// peer/admin token comparison +// --------------------------------------------------------------------------- + +func TestPeerTokenOK(t *testing.T) { + req := func(auth string) *http.Request { + r := httptest.NewRequest(http.MethodGet, "/admin/x", nil) + if auth != "" { + r.Header.Set("Authorization", auth) + } + return r + } + cases := []struct { + name string + auth string + want string + ok bool + }{ + {"empty want accepts anything", "", "", true}, + {"empty want accepts a bogus token", "Bearer nonsense", "", true}, + {"exact match", "Bearer s3cret", "s3cret", true}, + {"wrong token", "Bearer s3cres", "s3cret", false}, + {"prefix of the real token", "Bearer s3c", "s3cret", false}, + {"real token plus suffix", "Bearer s3cretX", "s3cret", false}, + {"missing header", "", "s3cret", false}, + {"wrong scheme", "Basic s3cret", "s3cret", false}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if got := peerTokenOK(req(tc.auth), tc.want); got != tc.ok { + t.Fatalf("peerTokenOK(%q, %q) = %v, want %v", tc.auth, tc.want, got, tc.ok) + } + }) + } +} + +// bearerCompare matches a direct string comparison of a request-supplied +// bearer credential against the configured token. Every admin gate must route +// through peerTokenOK instead, which compares in constant time. +var bearerCompare = regexp.MustCompile( + `(?:got != want|!= "Bearer "\+want|r\.Header\.Get\("Authorization"\) != )`) + +func TestAdminGatesUseConstantTimeCompare(t *testing.T) { + files, err := filepath.Glob("*.go") + if err != nil { + t.Fatalf("glob: %v", err) + } + var offenders []string + for _, f := range files { + if strings.HasSuffix(f, "_test.go") { + continue + } + src, err := os.ReadFile(f) + if err != nil { + t.Fatalf("read %s: %v", f, err) + } + for i, line := range strings.Split(string(src), "\n") { + if bearerCompare.MatchString(line) { + offenders = append(offenders, f+":"+strconv.Itoa(i+1)+": "+strings.TrimSpace(line)) + } + } + } + if len(offenders) > 0 { + t.Fatalf("token comparisons not routed through peerTokenOK:\n %s", + strings.Join(offenders, "\n ")) + } +} + +// requireAdmin must still reject a wrong token and admit the right one. +func TestRequireAdminGate(t *testing.T) { + s := &pebbleHTTP{} + s.cluster = config.Cluster{PeerAuthToken: "topsecret"} + called := 0 + h := s.requireAdmin(func(w http.ResponseWriter, r *http.Request) { + called++ + w.WriteHeader(http.StatusOK) + }) + + r := httptest.NewRequest(http.MethodPost, "/admin/x", nil) + r.Header.Set("Authorization", "Bearer wrong") + w := httptest.NewRecorder() + h(w, r) + if w.Code != http.StatusUnauthorized { + t.Fatalf("wrong token: status %d, want 401", w.Code) + } + if called != 0 { + t.Fatalf("handler ran for a wrong token") + } + + r = httptest.NewRequest(http.MethodPost, "/admin/x", nil) + r.Header.Set("Authorization", "Bearer topsecret") + w = httptest.NewRecorder() + h(w, r) + if w.Code != http.StatusOK || called != 1 { + t.Fatalf("right token: status %d called %d, want 200/1", w.Code, called) + } +} + +// --------------------------------------------------------------------------- +// embed-backfill durability +// --------------------------------------------------------------------------- + +// TestEmbedBackfillPersistsHNSW checks that vectors added by the backfill +// handler are written to the store, not just to the in-memory graph: after the +// handler returns, loading the graph back from Pebble must yield the nodes. +func TestEmbedBackfillPersistsHNSW(t *testing.T) { + f := populatedPebbleStore(t) + mock := openaiTestServer(t) + srv := f.makeServer(mock) + // Start from an empty graph so every corpus doc counts as missing. + srv.hnsw = index.NewHNSW(f.dim) + + ctx := context.Background() + if _, ok, err := index.LoadHNSWMeta(ctx, f.ps); err != nil { + t.Fatalf("LoadHNSWMeta: %v", err) + } else if ok { + t.Fatalf("precondition: store already holds HNSW meta") + } + + r := httptest.NewRequest(http.MethodPost, "/admin/embed-backfill", + strings.NewReader(`{"workers":2}`)) + w := httptest.NewRecorder() + srv.handleEmbedBackfill(w, r) + if w.Code != http.StatusOK { + t.Fatalf("status %d: %s", w.Code, w.Body.String()) + } + + var resp struct { + Missing int `json:"missing"` + Embedded int `json:"embedded"` + Persisted bool `json:"persisted"` + PersistError string `json:"persist_error"` + } + if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil { + t.Fatalf("decode: %v (%s)", err, w.Body.String()) + } + if resp.Embedded == 0 { + t.Fatalf("nothing embedded: %s", w.Body.String()) + } + if !resp.Persisted { + t.Fatalf("handler reported persisted=false (err=%q)", resp.PersistError) + } + if srv.hnsw.Len() == 0 { + t.Fatalf("in-memory graph is empty after backfill") + } + + loaded, ok, err := index.LoadHNSW(ctx, f.ps) + if err != nil { + t.Fatalf("LoadHNSW: %v", err) + } + if !ok { + t.Fatalf("backfilled vectors were not persisted to the store") + } + if loaded.Len() != srv.hnsw.Len() { + t.Fatalf("persisted graph has %d nodes, in-memory has %d", + loaded.Len(), srv.hnsw.Len()) + } +} diff --git a/cmd/cosift/serve_admin.go b/cmd/cosift/serve_admin.go index 3eabfa8..ddc814d 100644 --- a/cmd/cosift/serve_admin.go +++ b/cmd/cosift/serve_admin.go @@ -26,12 +26,9 @@ import ( // against the codebook already loaded on this serve. Order-of-magnitude // faster than handlePQTrain because it skips k-means; just encode loop. func (s *pebbleHTTP) handlePQEncode(w http.ResponseWriter, r *http.Request) { - if want := s.cluster.PeerAuthToken; want != "" { - got := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") - if got != want { - writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") - return - } + if !peerTokenOK(r, s.cluster.PeerAuthToken) { + writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") + return } if s.hnsw == nil { writeProblem(w, http.StatusBadRequest, "pq-encode: no HNSW loaded") @@ -90,12 +87,9 @@ func (s *pebbleHTTP) handlePQEncode(w http.ResponseWriter, r *http.Request) { // returned path is safe to tar — Pebble's compactor cannot mutate hard-linked // SSTs. Caller is responsible for deleting the dir after consuming it. func (s *pebbleHTTP) handleCheckpoint(w http.ResponseWriter, r *http.Request) { - if want := s.cluster.PeerAuthToken; want != "" { - got := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") - if got != want { - writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") - return - } + if !peerTokenOK(r, s.cluster.PeerAuthToken) { + writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") + return } base := os.Getenv("COSIFT_CHECKPOINT_DIR") if base == "" { @@ -127,12 +121,9 @@ type pqTrainReq struct { } func (s *pebbleHTTP) handlePQTrain(w http.ResponseWriter, r *http.Request) { - if want := s.cluster.PeerAuthToken; want != "" { - got := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") - if got != want { - writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") - return - } + if !peerTokenOK(r, s.cluster.PeerAuthToken) { + writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") + return } if s.hnsw == nil || s.hnsw.Len() == 0 { writeProblem(w, http.StatusBadRequest, "pq-train requires an in-memory HNSW with nodes") @@ -246,12 +237,9 @@ var evalQuickQueries = []string{ // in-process (re-uses the same chat/retrieval stack as a real /answer call), // so the rate that comes back matches what users see. func (s *pebbleHTTP) handleEvalQuick(w http.ResponseWriter, r *http.Request) { - if want := s.cluster.PeerAuthToken; want != "" { - got := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") - if got != want { - writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") - return - } + if !peerTokenOK(r, s.cluster.PeerAuthToken) { + writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") + return } if s.chat == nil { writeProblem(w, http.StatusNotImplemented, "eval-quick requires cfg.Chat.Model") @@ -361,12 +349,9 @@ func (s *pebbleHTTP) handleEvalQuick(w http.ResponseWriter, r *http.Request) { // ResponseController because compacting a multi-million-node graph routinely // runs past 60s. Returns counters so operators can confirm progress. func (s *pebbleHTTP) handleHNSWCompact(w http.ResponseWriter, r *http.Request) { - if want := s.cluster.PeerAuthToken; want != "" { - got := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") - if got != want { - writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") - return - } + if !peerTokenOK(r, s.cluster.PeerAuthToken) { + writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") + return } if s.hnsw == nil { writeProblem(w, http.StatusNotImplemented, "hnsw-compact requires a loaded HNSW graph") @@ -474,12 +459,9 @@ type embedBackfillReq struct { // After it logs "done", set COSIFT_HOST_PARTITION_READ=1 (and =1 for the write // flag to keep it fresh) and restart. func (s *pebbleHTTP) handleHostBackfill(w http.ResponseWriter, r *http.Request) { - if want := s.cluster.PeerAuthToken; want != "" { - got := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") - if got != want { - writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") - return - } + if !peerTokenOK(r, s.cluster.PeerAuthToken) { + writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") + return } if s.store == nil { writeProblem(w, http.StatusNotImplemented, "host-backfill requires a PebbleStore") @@ -510,12 +492,9 @@ func (s *pebbleHTTP) handleHostBackfill(w http.ResponseWriter, r *http.Request) } func (s *pebbleHTTP) handleEmbedBackfill(w http.ResponseWriter, r *http.Request) { - if want := s.cluster.PeerAuthToken; want != "" { - got := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") - if got != want { - writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") - return - } + if !peerTokenOK(r, s.cluster.PeerAuthToken) { + writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") + return } if s.embedder == nil || s.hnsw == nil { writeProblem(w, http.StatusNotImplemented, "embed backfill requires both an embedder and HNSW graph to be configured") @@ -599,14 +578,34 @@ func (s *pebbleHTTP) handleEmbedBackfill(w http.ResponseWriter, r *http.Request) } wg.Wait() - log.Printf("embed-backfill: scanned=%d missing=%d embedded=%d in %s", - scanned.Load(), missing.Load(), embedded.Load(), time.Since(t0).Round(time.Second)) - writeJSON(w, http.StatusOK, map[string]any{ - "scanned": scanned.Load(), - "missing": missing.Load(), - "embedded": embedded.Load(), - "elapsed": time.Since(t0).String(), - }) + // AddPassage only mutates the in-memory graph. Snapshot it to the store + // so the newly-embedded vectors survive a restart; without this the + // backfill's work is lost whenever the process exits. Detached context + // so a client disconnect can't abort the write. + persisted := false + var persistErr string + if s.store != nil && embedded.Load() > 0 { + if err := s.hnsw.Persist(context.Background(), s.store); err != nil { + persistErr = err.Error() + log.Printf("embed-backfill: hnsw persist FAILED: %v", err) + } else { + persisted = true + } + } + + log.Printf("embed-backfill: scanned=%d missing=%d embedded=%d persisted=%v in %s", + scanned.Load(), missing.Load(), embedded.Load(), persisted, time.Since(t0).Round(time.Second)) + resp := map[string]any{ + "scanned": scanned.Load(), + "missing": missing.Load(), + "embedded": embedded.Load(), + "persisted": persisted, + "elapsed": time.Since(t0).String(), + } + if persistErr != "" { + resp["persist_error"] = persistErr + } + writeJSON(w, http.StatusOK, resp) } // truncateForEmbedLite mirrors the crawler's helper without pulling the diff --git a/cmd/cosift/serve_crawl.go b/cmd/cosift/serve_crawl.go index fbfb871..c5ef940 100644 --- a/cmd/cosift/serve_crawl.go +++ b/cmd/cosift/serve_crawl.go @@ -28,12 +28,9 @@ type crawlEnqueueReq struct { func (s *pebbleHTTP) handleCrawlEnqueue(w http.ResponseWriter, r *http.Request) { // Auth - if want := s.cluster.PeerAuthToken; want != "" { - got := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") - if got != want { - writeProblem(w, http.StatusUnauthorized, "missing or invalid peer token") - return - } + if !peerTokenOK(r, s.cluster.PeerAuthToken) { + writeProblem(w, http.StatusUnauthorized, "missing or invalid peer token") + return } if s.crawlSeed == nil { writeProblem(w, http.StatusNotImplemented, "this shard has no in-serve crawler (-crawl-seeds-file not set)") @@ -64,12 +61,9 @@ func (s *pebbleHTTP) handleCrawlEnqueue(w http.ResponseWriter, r *http.Request) // link-target domain recurs above their frequency threshold. Persisted so it // survives restart. Body: {"domain":"example.com"}. func (s *pebbleHTTP) handleAllowDomain(w http.ResponseWriter, r *http.Request) { - if want := s.cluster.PeerAuthToken; want != "" { - got := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") - if got != want { - writeProblem(w, http.StatusUnauthorized, "missing or invalid peer token") - return - } + if !peerTokenOK(r, s.cluster.PeerAuthToken) { + writeProblem(w, http.StatusUnauthorized, "missing or invalid peer token") + return } if s.crawlAllowDomain == nil { writeProblem(w, http.StatusNotImplemented, "this shard has no in-serve crawler") @@ -100,12 +94,9 @@ func (s *pebbleHTTP) handleAllowDomain(w http.ResponseWriter, r *http.Request) { // approximate count of deleted entries (DeleteRange is O(tombstones), // not O(N), so the count is reported as -1 to signal "swept range"). func (s *pebbleHTTP) handleFrontierClear(w http.ResponseWriter, r *http.Request) { - if want := s.cluster.PeerAuthToken; want != "" { - got := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") - if got != want { - writeProblem(w, http.StatusUnauthorized, "missing or invalid peer token") - return - } + if !peerTokenOK(r, s.cluster.PeerAuthToken) { + writeProblem(w, http.StatusUnauthorized, "missing or invalid peer token") + return } if err := s.store.ClearFrontier(r.Context()); err != nil { writeProblem(w, http.StatusInternalServerError, err.Error()) @@ -124,12 +115,9 @@ type frontierPurgeReq struct { } func (s *pebbleHTTP) handleFrontierPurgeHost(w http.ResponseWriter, r *http.Request) { - if want := s.cluster.PeerAuthToken; want != "" { - got := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") - if got != want { - writeProblem(w, http.StatusUnauthorized, "missing or invalid peer token") - return - } + if !peerTokenOK(r, s.cluster.PeerAuthToken) { + writeProblem(w, http.StatusUnauthorized, "missing or invalid peer token") + return } var req frontierPurgeReq body, _ := io.ReadAll(io.LimitReader(r.Body, 64<<10)) @@ -154,12 +142,9 @@ type sitemapImportReq struct { } func (s *pebbleHTTP) handleSitemapImport(w http.ResponseWriter, r *http.Request) { - if want := s.cluster.PeerAuthToken; want != "" { - got := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") - if got != want { - writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") - return - } + if !peerTokenOK(r, s.cluster.PeerAuthToken) { + writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") + return } if s.crawlSeedSitemap == nil { writeProblem(w, http.StatusNotImplemented, "this shard has no in-serve crawler (-crawl-seeds-file not set)") @@ -191,12 +176,9 @@ func (s *pebbleHTTP) handleSitemapImport(w http.ResponseWriter, r *http.Request) // entries are stuck in errored state — sitemap-import's INSERT OR IGNORE won't // re-queue them, but Recrawl will. func (s *pebbleHTTP) handleRecrawlSitemap(w http.ResponseWriter, r *http.Request) { - if want := s.cluster.PeerAuthToken; want != "" { - got := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") - if got != want { - writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") - return - } + if !peerTokenOK(r, s.cluster.PeerAuthToken) { + writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") + return } if s.crawlRecrawl == nil { writeProblem(w, http.StatusNotImplemented, "this shard has no in-serve crawler") @@ -261,12 +243,9 @@ type crawlNowReq struct { } func (s *pebbleHTTP) handleCrawlNow(w http.ResponseWriter, r *http.Request) { - if want := s.cluster.PeerAuthToken; want != "" { - got := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") - if got != want { - writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") - return - } + if !peerTokenOK(r, s.cluster.PeerAuthToken) { + writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") + return } if s.crawlFetchNow == nil { writeProblem(w, http.StatusNotImplemented, "this shard has no in-serve crawler (-crawl-seeds-file not set)") @@ -319,12 +298,9 @@ type sitePackReq struct { } func (s *pebbleHTTP) handleSitePack(w http.ResponseWriter, r *http.Request) { - if want := s.cluster.PeerAuthToken; want != "" { - got := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") - if got != want { - writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") - return - } + if !peerTokenOK(r, s.cluster.PeerAuthToken) { + writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") + return } if s.crawlSeedSitemap == nil || s.crawlSeedRSS == nil { writeProblem(w, http.StatusNotImplemented, "this shard has no in-serve crawler (-crawl-seeds-file not set)") @@ -500,12 +476,9 @@ type siteSubmitReq struct { } func (s *pebbleHTTP) handleSiteSubmit(w http.ResponseWriter, r *http.Request) { - if want := s.cluster.PeerAuthToken; want != "" { - got := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") - if got != want { - writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") - return - } + if !peerTokenOK(r, s.cluster.PeerAuthToken) { + writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") + return } if s.crawlSeedSitemapLane == nil { writeProblem(w, http.StatusNotImplemented, "this shard has no in-serve crawler (-crawl-seeds-file not set)") @@ -581,12 +554,9 @@ type wetImportBulkReq struct { } func (s *pebbleHTTP) handleWETImportBulk(w http.ResponseWriter, r *http.Request) { - if want := s.cluster.PeerAuthToken; want != "" { - got := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") - if got != want { - writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") - return - } + if !peerTokenOK(r, s.cluster.PeerAuthToken) { + writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") + return } if s.crawlSeedWET == nil { writeProblem(w, http.StatusNotImplemented, "this shard has no in-serve crawler (-crawl-seeds-file not set)") @@ -709,12 +679,9 @@ type wetImportReq struct { } func (s *pebbleHTTP) handleWETImport(w http.ResponseWriter, r *http.Request) { - if want := s.cluster.PeerAuthToken; want != "" { - got := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") - if got != want { - writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") - return - } + if !peerTokenOK(r, s.cluster.PeerAuthToken) { + writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") + return } if s.crawlSeedWET == nil { writeProblem(w, http.StatusNotImplemented, "this shard has no in-serve crawler (-crawl-seeds-file not set)") @@ -753,12 +720,9 @@ type frontierDemoteHostReq struct { } func (s *pebbleHTTP) handleFrontierDemoteHost(w http.ResponseWriter, r *http.Request) { - if want := s.cluster.PeerAuthToken; want != "" { - got := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") - if got != want { - writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") - return - } + if !peerTokenOK(r, s.cluster.PeerAuthToken) { + writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") + return } ps, ok := any(s.store).(*store.PebbleStore) if !ok { @@ -797,12 +761,9 @@ func (s *pebbleHTTP) handleFrontierDemoteHost(w http.ResponseWriter, r *http.Req // key. GetLaneStats then reported impossibly-high in_flight counts // (>max_concurrent). Idempotent — re-running is a no-op once clean. func (s *pebbleHTTP) handleFrontierPurgeStaleInFlight(w http.ResponseWriter, r *http.Request) { - if want := s.cluster.PeerAuthToken; want != "" { - got := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") - if got != want { - writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") - return - } + if !peerTokenOK(r, s.cluster.PeerAuthToken) { + writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") + return } ps, ok := any(s.store).(*store.PebbleStore) if !ok { @@ -828,12 +789,9 @@ type rssImportReq struct { } func (s *pebbleHTTP) handleRSSImport(w http.ResponseWriter, r *http.Request) { - if want := s.cluster.PeerAuthToken; want != "" { - got := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") - if got != want { - writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") - return - } + if !peerTokenOK(r, s.cluster.PeerAuthToken) { + writeProblem(w, http.StatusUnauthorized, "missing or invalid admin token") + return } if s.crawlSeedRSS == nil { writeProblem(w, http.StatusNotImplemented, "this shard has no in-serve crawler (-crawl-seeds-file not set)") diff --git a/cmd/cosift/serve_helpers.go b/cmd/cosift/serve_helpers.go index 38f38ac..e7a3747 100644 --- a/cmd/cosift/serve_helpers.go +++ b/cmd/cosift/serve_helpers.go @@ -2,6 +2,7 @@ package main import ( "context" + "crypto/subtle" _ "embed" "encoding/json" "flag" @@ -19,6 +20,25 @@ import ( "github.com/pilot-protocol/cosift/internal/store" ) +// bearerToken extracts the credential from an "Authorization: Bearer " +// request header. Returns "" when the header is absent or carries a different +// scheme. +func bearerToken(r *http.Request) string { + return strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") +} + +// peerTokenOK reports whether the request presents the expected peer/admin +// token. The comparison runs in time independent of how many leading bytes +// match, and returns 1 only when both lengths and all bytes are equal. +// An empty want means the endpoint is unauthenticated (single-node default) +// and every caller is accepted. +func peerTokenOK(r *http.Request, want string) bool { + if want == "" { + return true + } + return subtle.ConstantTimeCompare([]byte(bearerToken(r)), []byte(want)) == 1 +} + // retrievalFilters bundles the four post-retrieval predicates /search, // /answer, /find_similar, and /research all share. Centralized in // so /research's sync + stream paths stay in lockstep. diff --git a/cmd/cosift/serve_setup.go b/cmd/cosift/serve_setup.go index 50286af..bd917a4 100644 --- a/cmd/cosift/serve_setup.go +++ b/cmd/cosift/serve_setup.go @@ -1292,12 +1292,9 @@ func stripPort(remoteAddr string) string { // without auth, even if it forgets the per-handler check. func (s *pebbleHTTP) requireAdmin(h http.HandlerFunc) http.HandlerFunc { return func(w http.ResponseWriter, r *http.Request) { - if want := s.cluster.PeerAuthToken; want != "" { - got := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") - if got != want { - writeProblem(w, http.StatusUnauthorized, "missing or invalid peer token") - return - } + if !peerTokenOK(r, s.cluster.PeerAuthToken) { + writeProblem(w, http.StatusUnauthorized, "missing or invalid peer token") + return } h(w, r) } diff --git a/internal/crawler/crawler.go b/internal/crawler/crawler.go index 2eac2e4..57fe2f2 100644 --- a/internal/crawler/crawler.go +++ b/internal/crawler/crawler.go @@ -65,9 +65,10 @@ type Crawler struct { // bounded diagnostic counter for the zombie-reclaim path. // Logs the first few zero-hit cases so we can confirm the code path - // is actually reached even when no prior passages exist. Cheap (just - // an int read on the hot path; not atomic — diagnostic only). - zombieDebugLogged int + // is actually reached even when no prior passages exist. Written from + // every crawl worker goroutine, so the read-modify-write goes through + // atomic.Int64. + zombieDebugLogged atomic.Int64 // per-host error-rate tracking via sync.Map + atomic // counters. With 512 workers we cannot afford a single write lock @@ -93,6 +94,13 @@ type Crawler struct { embedDone atomic.Int64 embedFailed atomic.Int64 + // embedDrainUntil is a unix-nano wall-clock deadline for the post-crawl + // drain of embedQ. Zero means "no deadline" (crawl still running). Run + // sets it just before closing embedQ; embed workers stop doing work and + // count the remainder as dropped once it passes, so shutdown stays + // bounded even when the embedder is slow or unreachable. + embedDrainUntil atomic.Int64 + // URLs whose frontier transition happened inside processClaimed via // WriteCrawlResult. The worker checks this set before its own // CompleteFrontier call so we don't pay a redundant mu hop. @@ -604,9 +612,15 @@ func (c *Crawler) Run(ctx context.Context) error { if embedWorkers > 0 { bufSize := envIntCrawler("COSIFT_EMBED_DECOUPLE_BUFFER", 4096) c.embedQ = make(chan *embedJob, bufSize) + // runCtx is cancelled the moment the frontier drains (see + // terminator), while jobs already queued still need a live + // context to embed and write. Detach the embed pool from that + // cancellation; its wall-clock bound comes from embedDrainUntil + // plus the per-job timeout instead. + embedCtx := context.WithoutCancel(ctx) for i := 0; i < embedWorkers; i++ { embedWG.Add(1) - go c.embedWorker(runCtx, &embedWG) + go c.embedWorker(embedCtx, &embedWG) } log.Printf("crawler: embed decouple ON (%d workers, %d-buf)", embedWorkers, bufSize) } @@ -637,8 +651,14 @@ func (c *Crawler) Run(ctx context.Context) error { // then wait for embed workers to drain. Otherwise an early close // would race a still-running crawl worker's send and panic. if c.embedQ != nil { + // Cap the drain so a slow or unreachable embedder can't hold the + // crawl open indefinitely. Jobs still queued past the deadline are + // counted as dropped and left for embed-backfill. + drainSecs := envIntCrawler("COSIFT_EMBED_DRAIN_SECONDS", defaultEmbedDrainSeconds) + c.embedDrainUntil.Store(time.Now().Add(time.Duration(drainSecs) * time.Second).UnixNano()) close(c.embedQ) embedWG.Wait() + c.embedDrainUntil.Store(0) log.Printf("crawler: embed pool drained — queued=%d done=%d failed=%d dropped=%d", c.embedQueued.Load(), c.embedDone.Load(), c.embedFailed.Load(), c.embedDropped.Load()) c.embedQ = nil // nil so FetchAndIndexNow doesn't send to closed channel between restarts @@ -661,64 +681,126 @@ func envIntCrawler(key string, def int) int { return n } +// Defaults for the decoupled embed pipeline's wall-clock bounds. Both are +// overridable via env (COSIFT_EMBED_DRAIN_SECONDS, COSIFT_EMBED_JOB_SECONDS, +// COSIFT_EMBED_ENQUEUE_WAIT_MS). +const ( + defaultEmbedDrainSeconds = 300 + defaultEmbedJobSeconds = 120 + defaultEmbedEnqueueWaitMs = 5000 +) + +// enqueueEmbedJob hands a job to the embed pool. It first tries a +// non-blocking send, then waits a bounded amount of time for buffer space +// before giving up. Reports whether the job was accepted; a rejected job's +// passages are left for embed-backfill. +func (c *Crawler) enqueueEmbedJob(ctx context.Context, job *embedJob) bool { + select { + case c.embedQ <- job: + return true + default: + } + wait := time.Duration(envIntCrawler("COSIFT_EMBED_ENQUEUE_WAIT_MS", defaultEmbedEnqueueWaitMs)) * time.Millisecond + if wait <= 0 { + return false + } + t := time.NewTimer(wait) + defer t.Stop() + select { + case c.embedQ <- job: + return true + case <-ctx.Done(): + return false + case <-t.C: + return false + } +} + +// embedJobExpired reports whether the post-crawl drain deadline has passed. +func (c *Crawler) embedJobExpired() bool { + d := c.embedDrainUntil.Load() + return d != 0 && time.Now().UnixNano() > d +} + // embedWorker drains c.embedQ. For each job, embeds the chunk texts and // writes passages to the configured PassageWriter. Uses the optional // batch writer when available for one HNSW lock per doc instead of // one per chunk. +// +// ctx is detached from the crawl's cancellation so queued jobs still +// complete after the frontier drains; each job carries its own timeout and +// the pool stops working once the drain deadline passes. func (c *Crawler) embedWorker(ctx context.Context, wg *sync.WaitGroup) { defer wg.Done() + jobTimeout := time.Duration(envIntCrawler("COSIFT_EMBED_JOB_SECONDS", defaultEmbedJobSeconds)) * time.Second + if jobTimeout <= 0 { + jobTimeout = defaultEmbedJobSeconds * time.Second + } for job := range c.embedQ { - vecs, err := c.embedder.Embed(ctx, job.texts) - if err != nil { - c.embedFailed.Add(1) - log.Printf("embed-decouple %s: %v", job.url, err) - continue - } - if len(vecs) != len(job.chunks) { - c.embedFailed.Add(1) + if c.embedJobExpired() { + c.embedDropped.Add(1) continue } - // Mirror the synchronous path's zombie reclaim so re-crawled - // URLs don't accumulate generations of vectors in HNSW. - if os.Getenv("COSIFT_ZOMBIE_RECLAIM") == "1" { - if inv, ok := c.passageWriter.(URLInvalidator); ok { - _, _ = inv.MarkURLInvalid(ctx, job.url) + c.runEmbedJob(ctx, job, jobTimeout) + } +} + +// runEmbedJob embeds one job's chunk texts and writes the resulting passages. +// The per-job context bounds how long a single document may hold a worker. +func (c *Crawler) runEmbedJob(parent context.Context, job *embedJob, jobTimeout time.Duration) { + ctx, cancel := context.WithTimeout(parent, jobTimeout) + defer cancel() + + vecs, err := c.embedder.Embed(ctx, job.texts) + if err != nil { + c.embedFailed.Add(1) + log.Printf("embed-decouple %s: %v", job.url, err) + return + } + if len(vecs) != len(job.chunks) { + c.embedFailed.Add(1) + return + } + // Mirror the synchronous path's zombie reclaim so re-crawled + // URLs don't accumulate generations of vectors in HNSW. + if os.Getenv("COSIFT_ZOMBIE_RECLAIM") == "1" { + if inv, ok := c.passageWriter.(URLInvalidator); ok { + _, _ = inv.MarkURLInvalid(ctx, job.url) + } + } + // Prefer the batch interface (single HNSW lock for the whole + // doc) over per-chunk writes. + if bw, ok := c.passageWriter.(PassageWriterBatch); ok { + ps := make([]*store.Passage, len(job.chunks)) + for i, ch := range job.chunks { + ps[i] = &store.Passage{ + DocID: job.docID, + Offset: ch.Offset, + Length: ch.Length, + Model: c.embedder.Model(), + Embedding: vecs[i], } } - // Prefer the batch interface (single HNSW lock for the whole - // doc) over per-chunk writes. - if bw, ok := c.passageWriter.(PassageWriterBatch); ok { - ps := make([]*store.Passage, len(job.chunks)) - for i, ch := range job.chunks { - ps[i] = &store.Passage{ - DocID: job.docID, - Offset: ch.Offset, - Length: ch.Length, - Model: c.embedder.Model(), - Embedding: vecs[i], - } - } - if err := bw.UpsertPassageBatch(ctx, ps); err != nil { - c.embedFailed.Add(1) - log.Printf("embed-decouple batch %s: %v", job.url, err) - continue + if err := bw.UpsertPassageBatch(ctx, ps); err != nil { + c.embedFailed.Add(1) + log.Printf("embed-decouple batch %s: %v", job.url, err) + return + } + } else { + for i, ch := range job.chunks { + p := &store.Passage{ + DocID: job.docID, + Offset: ch.Offset, + Length: ch.Length, + Model: c.embedder.Model(), + Embedding: vecs[i], } - } else { - for i, ch := range job.chunks { - p := &store.Passage{ - DocID: job.docID, - Offset: ch.Offset, - Length: ch.Length, - Model: c.embedder.Model(), - Embedding: vecs[i], - } - if err := c.passageWriter.UpsertPassage(ctx, p); err != nil { - log.Printf("embed-decouple passage %s offset=%d: %v", job.url, ch.Offset, err) - } + if err := c.passageWriter.UpsertPassage(ctx, p); err != nil { + log.Printf("embed-decouple passage %s offset=%d: %v", job.url, ch.Offset, err) } } - c.embedDone.Add(1) } + c.embedDone.Add(1) } // statusDumper writes a JSON snapshot of crawl progress every 10s to path. @@ -1375,10 +1457,9 @@ func (c *Crawler) processClaimed(ctx context.Context, item store.FrontierItem, g } if c.embedQ != nil { job := &embedJob{url: item.URL, docID: id, chunks: chunks, texts: texts} - select { - case c.embedQ <- job: + if c.enqueueEmbedJob(ctx, job) { c.embedQueued.Add(1) - default: + } else { c.embedDropped.Add(1) } c.enqueueLinks(ctx, parsed.Links, item.Depth+1) @@ -1417,10 +1498,9 @@ func (c *Crawler) processClaimed(ctx context.Context, item store.FrontierItem, g log.Printf("zombie-reclaim %s: %v", item.URL, err) case n > 0: log.Printf("zombie-reclaim %s: invalidated %d stale passages", item.URL, n) - case c.zombieDebugLogged < 5: + case c.zombieDebugLogged.Add(1) <= 5: // Diagnostic: first few zero-hit cases to confirm code path is reached. log.Printf("zombie-reclaim %s: 0 prior passages (diagnostic)", item.URL) - c.zombieDebugLogged++ } } } diff --git a/internal/crawler/embed_drain_test.go b/internal/crawler/embed_drain_test.go new file mode 100644 index 0000000..32ddc2a --- /dev/null +++ b/internal/crawler/embed_drain_test.go @@ -0,0 +1,259 @@ +package crawler + +import ( + "context" + "fmt" + "net/http" + "net/http/httptest" + "sync" + "testing" + "time" + + "github.com/pilot-protocol/cosift/internal/config" + "github.com/pilot-protocol/cosift/internal/store" +) + +// slowEmbedder takes a fixed amount of time per call and honors context +// cancellation the way a real HTTP-backed embedder does. +type slowEmbedder struct { + dim int + delay time.Duration + + mu sync.Mutex + calls int + failed int +} + +func (s *slowEmbedder) Model() string { return "slow-emb" } +func (s *slowEmbedder) Dim() int { return s.dim } + +func (s *slowEmbedder) Embed(ctx context.Context, texts []string) ([][]float32, error) { + t := time.NewTimer(s.delay) + defer t.Stop() + select { + case <-ctx.Done(): + s.mu.Lock() + s.failed++ + s.mu.Unlock() + return nil, ctx.Err() + case <-t.C: + } + s.mu.Lock() + s.calls++ + s.mu.Unlock() + out := make([][]float32, len(texts)) + for i := range texts { + v := make([]float32, s.dim) + for j := range v { + v[j] = 0.1 + } + out[i] = v + } + return out, nil +} + +// countingWriter records passage writes and reports zero prior passages from +// MarkURLInvalid, which is the branch that drives the zombie-reclaim +// diagnostic counter. +type countingWriter struct { + inner PassageWriter + // invalidateDelay holds each worker inside the reclaim block so that all + // of them reach the diagnostic counter at roughly the same instant. + invalidateDelay time.Duration + + mu sync.Mutex + upserts int + invalids int +} + +func (c *countingWriter) UpsertPassage(ctx context.Context, p *store.Passage) error { + if err := c.inner.UpsertPassage(ctx, p); err != nil { + return err + } + c.mu.Lock() + c.upserts++ + c.mu.Unlock() + return nil +} + +func (c *countingWriter) MarkURLInvalid(ctx context.Context, url string) (int, error) { + if c.invalidateDelay > 0 { + time.Sleep(c.invalidateDelay) + } + c.mu.Lock() + c.invalids++ + c.mu.Unlock() + return 0, nil +} + +func (c *countingWriter) counts() (int, int) { + c.mu.Lock() + defer c.mu.Unlock() + return c.upserts, c.invalids +} + +// pagesServer serves n distinct HTML pages under /p. +func pagesServer(t *testing.T, n int) (*httptest.Server, []string) { + t.Helper() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/html") + fmt.Fprintf(w, `Page %s +

Heading %s

Body text for path %s about distributed indexing and retrieval.

`, + r.URL.Path, r.URL.Path, r.URL.Path) + })) + t.Cleanup(srv.Close) + urls := make([]string, 0, n) + for i := 0; i < n; i++ { + urls = append(urls, fmt.Sprintf("%s/p%d", srv.URL, i)) + } + return srv, urls +} + +func fastCrawlCfg() config.Crawler { + cfg := config.Default().Crawler + cfg.MaxDepth = 0 + cfg.PerHostDelayMs = 0 + cfg.MaxConcurrent = 8 + cfg.RespectRobots = false + return cfg +} + +// TestEmbedPoolFinishesQueuedJobsAfterCrawlContextCancel checks that the +// decoupled embed pool keeps a usable context after the crawl's own context is +// cancelled (the terminator cancels it as soon as the frontier drains). Jobs +// already sitting in the queue must still be embedded and written, not failed. +func TestEmbedPoolFinishesQueuedJobsAfterCrawlContextCancel(t *testing.T) { + const pages = 40 + + t.Setenv("COSIFT_EMBED_DECOUPLE_WORKERS", "1") + t.Setenv("COSIFT_EMBED_DECOUPLE_BUFFER", "256") + + _, urls := pagesServer(t, pages) + s := newStoreT(t) + emb := &slowEmbedder{dim: 8, delay: 100 * time.Millisecond} + cw := &countingWriter{inner: s} + + c := New(fastCrawlCfg(), s).WithEmbedder(emb).WithPassageWriter(cw) + for _, u := range urls { + if err := c.Seed(u); err != nil { + t.Fatalf("seed %s: %v", u, err) + } + } + if err := c.Run(context.Background()); err != nil { + t.Fatalf("run: %v", err) + } + + queued := c.embedQueued.Load() + done := c.embedDone.Load() + failed := c.embedFailed.Load() + dropped := c.embedDropped.Load() + + if queued == 0 { + t.Fatalf("no embed jobs were queued — decoupled path did not engage") + } + if failed != 0 { + t.Errorf("embed jobs failed after the crawl context was cancelled: failed=%d (queued=%d done=%d)", + failed, queued, done) + } + if done != queued { + t.Errorf("embed pool completed %d of %d queued jobs (failed=%d dropped=%d)", + done, queued, failed, dropped) + } + if up, _ := cw.counts(); up == 0 { + t.Errorf("no passages were written") + } +} + +// TestEnqueueEmbedJobWaitsForCapacity checks that a momentarily full queue +// makes the producer wait for space rather than discarding the document's +// passages outright. +func TestEnqueueEmbedJobWaitsForCapacity(t *testing.T) { + c := newBare(config.Default().Crawler) + c.embedQ = make(chan *embedJob, 1) + c.embedQ <- &embedJob{url: "https://x.example/filler"} + + t.Setenv("COSIFT_EMBED_ENQUEUE_WAIT_MS", "2000") + + go func() { + time.Sleep(50 * time.Millisecond) + <-c.embedQ + }() + + start := time.Now() + if !c.enqueueEmbedJob(context.Background(), &embedJob{url: "https://x.example/real"}) { + t.Fatalf("job rejected even though capacity freed up after %s", time.Since(start)) + } + if elapsed := time.Since(start); elapsed < 40*time.Millisecond { + t.Fatalf("returned in %s — did not actually wait for capacity", elapsed) + } +} + +// TestEnqueueEmbedJobGivesUpWhenQueueStaysFull keeps the wall clock bounded: +// a permanently full queue must not block a crawl worker indefinitely. +func TestEnqueueEmbedJobGivesUpWhenQueueStaysFull(t *testing.T) { + c := newBare(config.Default().Crawler) + c.embedQ = make(chan *embedJob, 1) + c.embedQ <- &embedJob{url: "https://x.example/filler"} + + t.Setenv("COSIFT_EMBED_ENQUEUE_WAIT_MS", "50") + + start := time.Now() + if c.enqueueEmbedJob(context.Background(), &embedJob{url: "https://x.example/real"}) { + t.Fatalf("job accepted into a full queue") + } + if elapsed := time.Since(start); elapsed > 2*time.Second { + t.Fatalf("enqueue took %s — wait budget not honored", elapsed) + } +} + +// TestEmbedJobExpiredHonoursDrainDeadline covers the shutdown bound: once the +// drain deadline has passed, remaining queued work is abandoned. +func TestEmbedJobExpiredHonoursDrainDeadline(t *testing.T) { + c := newBare(config.Default().Crawler) + if c.embedJobExpired() { + t.Fatalf("expired with no deadline set") + } + c.embedDrainUntil.Store(time.Now().Add(time.Hour).UnixNano()) + if c.embedJobExpired() { + t.Fatalf("expired with an hour of budget left") + } + c.embedDrainUntil.Store(time.Now().Add(-time.Second).UnixNano()) + if !c.embedJobExpired() { + t.Fatalf("did not expire past the deadline") + } +} + +// TestZombieReclaimDiagnosticCounterConcurrent drives the zombie-reclaim +// diagnostic path from many crawl workers at once. Run under -race. +func TestZombieReclaimDiagnosticCounterConcurrent(t *testing.T) { + const pages = 60 + + t.Setenv("COSIFT_ZOMBIE_RECLAIM", "1") + t.Setenv("COSIFT_EMBED_DECOUPLE_WORKERS", "0") + + _, urls := pagesServer(t, pages) + s := newStoreT(t) + // slowEmbedder with no delay: the package's stubEmbedder keeps an + // unsynchronized call counter and is only safe at MaxConcurrent=1. + emb := &slowEmbedder{dim: 8} + cw := &countingWriter{inner: s, invalidateDelay: 15 * time.Millisecond} + + cfg := fastCrawlCfg() + cfg.MaxConcurrent = 16 + c := New(cfg, s).WithEmbedder(emb).WithPassageWriter(cw) + for _, u := range urls { + if err := c.Seed(u); err != nil { + t.Fatalf("seed %s: %v", u, err) + } + } + if err := c.Run(context.Background()); err != nil { + t.Fatalf("run: %v", err) + } + + if _, inv := cw.counts(); inv == 0 { + t.Fatalf("zombie-reclaim path never ran — counter was not exercised") + } + if got := c.zombieDebugLogged.Load(); got == 0 { + t.Fatalf("diagnostic counter never incremented") + } +}