sec: admin auth, embed-backfill durability, and embed-queue hardening (M23, L20, L21, L22)

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 == "" {

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 == "" {

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())
+	}
+}

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