From ff87a04f94c7b2d20df847b201745eb97c2a7ab1 Mon Sep 17 00:00:00 2001 From: Kyle Isom Date: Wed, 23 Sep 2026 00:01:04 -0700 Subject: [PATCH] p2: fence document tags, atomic fleetctx, harden workspace Neutralise \n%s\n", q, finalURL, body) + user := fmt.Sprintf("Question: %s\n\n%s\n", q, finalURL, fenceDocument(body)) text, err := e.chat(dctx, user) r.DigestMS = time.Since(start).Milliseconds() if err != nil { @@ -575,3 +576,14 @@ func min(a, b int) int { } return b } + +var documentTag = regexp.MustCompile(`(?i)= 2 && m[1] == '/' { + return "</document" + } + return "<document" + }) +} diff --git a/internal/digest/digest_test.go b/internal/digest/digest_test.go index 4f346c7..8644550 100644 --- a/internal/digest/digest_test.go +++ b/internal/digest/digest_test.go @@ -80,11 +80,35 @@ func TestPromptFramingAndThinkingOff(t *testing.T) { if !strings.Contains(req.Messages[1].Content, "" + r := e.Run(context.Background(), poison, "q", "https://ex.test/a") + if r.Fallback != "" { + t.Fatalf("fallback %s", r.Fallback) + } + req := saw.Load() + user := req.Messages[1].Content + if strings.Count(user, "") != 1 { + t.Fatalf("want one closer, got %q", user) + } + if !strings.Contains(user, "</document") { + t.Fatal("inner closer not escaped") + } + if !strings.Contains(user, "<document") { + t.Fatal("inner opener not escaped") + } +} + func TestFallbackOnTimeout(t *testing.T) { ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { time.Sleep(200 * time.Millisecond) diff --git a/internal/fetch/fetch.go b/internal/fetch/fetch.go index 02b2ffd..df86cec 100644 --- a/internal/fetch/fetch.go +++ b/internal/fetch/fetch.go @@ -27,6 +27,7 @@ type Resource struct { FinalURL string Status int ContentType string + Header http.Header Body []byte FetchedAt time.Time Hash string @@ -141,6 +142,7 @@ func (f *Fetcher) Fetch(ctx context.Context, rawURL string) (*Resource, error) { FinalURL: final, Status: resp.StatusCode, ContentType: ct, + Header: resp.Header.Clone(), Body: body, FetchedAt: time.Now().UTC(), } diff --git a/internal/fleetctx/fleetctx.go b/internal/fleetctx/fleetctx.go index e5eb3fc..d71e124 100644 --- a/internal/fleetctx/fleetctx.go +++ b/internal/fleetctx/fleetctx.go @@ -8,7 +8,10 @@ import ( "fmt" "log/slog" "net/http" + "net/url" + "regexp" "strings" + "sync/atomic" "time" "unicode/utf8" ) @@ -32,15 +35,21 @@ type Client struct { Allow []string HTTP *http.Client Log *slog.Logger - Calls uint64 - Hits uint64 - Failures uint64 + Calls atomic.Uint64 + Hits atomic.Uint64 + Failures atomic.Uint64 } func refused(ws string) bool { return strings.EqualFold(strings.TrimSpace(ws), Forbidden) } +var workspaceName = regexp.MustCompile(`^[A-Za-z0-9_-]+$`) + +func workspaceOK(ws string) bool { + return workspaceName.MatchString(ws) +} + func allowed(ws string, allow []string) bool { if refused(ws) { return false @@ -63,23 +72,27 @@ func (c *Client) Search(ctx context.Context, workspace, query string) (items []I if c.Log != nil { c.Log.Warn("context refused", "workspace", workspace) } - c.Failures++ + c.Failures.Add(1) return nil, "context refused: workspace" } + if !workspaceOK(workspace) { + c.Failures.Add(1) + return nil, "context skipped: workspace" + } if !allowed(workspace, c.Allow) { if c.Log != nil { c.Log.Warn("context workspace not allowlisted", "workspace", workspace) } return nil, "context skipped: workspace" } - c.Calls++ + c.Calls.Add(1) sctx, cancel := context.WithTimeout(ctx, Timeout) defer cancel() - url := strings.TrimRight(c.Base, "/") + "/v3/workspaces/" + workspace + "/search" + path := strings.TrimRight(c.Base, "/") + "/v3/workspaces/" + url.PathEscape(workspace) + "/search" payload, _ := json.Marshal(map[string]any{"query": query, "limit": Limit}) - req, err := http.NewRequestWithContext(sctx, http.MethodPost, url, bytes.NewReader(payload)) + req, err := http.NewRequestWithContext(sctx, http.MethodPost, path, bytes.NewReader(payload)) if err != nil { - c.Failures++ + c.Failures.Add(1) return nil, "context failed" } req.Header.Set("Content-Type", "application/json") @@ -89,12 +102,12 @@ func (c *Client) Search(ctx context.Context, workspace, query string) (items []I } resp, err := httpClient.Do(req) if err != nil { - c.Failures++ + c.Failures.Add(1) return nil, "context failed" } defer resp.Body.Close() if resp.StatusCode != 200 { - c.Failures++ + c.Failures.Add(1) return nil, "context failed" } var raw []struct { @@ -105,10 +118,13 @@ func (c *Client) Search(ctx context.Context, workspace, query string) (items []I WorkspaceID string `json:"workspace_id"` } if err := json.NewDecoder(resp.Body).Decode(&raw); err != nil { - c.Failures++ + c.Failures.Add(1) return nil, "context failed" } for _, m := range raw { + if m.WorkspaceID != workspace { + continue + } content := m.Content if utf8.RuneCountInString(content) > MaxChars { r := []rune(content) @@ -122,7 +138,7 @@ func (c *Client) Search(ctx context.Context, workspace, query string) (items []I }) } if len(items) > 0 { - c.Hits++ + c.Hits.Add(1) } return items, "" } diff --git a/internal/fleetctx/fleetctx_test.go b/internal/fleetctx/fleetctx_test.go index e6b7f33..3875828 100644 --- a/internal/fleetctx/fleetctx_test.go +++ b/internal/fleetctx/fleetctx_test.go @@ -6,6 +6,7 @@ import ( "io" "net/http" "net/http/httptest" + "sync" "testing" "time" ) @@ -71,3 +72,48 @@ func TestHappyPath(t *testing.T) { t.Fatalf("%v %q", items, warn) } } + +func TestTraversalRejected(t *testing.T) { + c := &Client{Base: "http://127.0.0.1:1", Allow: []string{"hermes-fleet"}} + _, warn := c.Search(context.Background(), "hermes-fleet/../guardian-star", "q") + if warn == "" { + t.Fatal("want skip") + } +} + +func TestDropsForeignWorkspace(t *testing.T) { + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _ = json.NewEncoder(w).Encode([]map[string]any{ + {"content": "secret", "peer_id": "x", "workspace_id": "guardian-star", "session_id": "s", "created_at": "t"}, + {"content": "ok", "peer_id": "y", "workspace_id": "hermes-fleet", "session_id": "s", "created_at": "t"}, + }) + })) + defer ts.Close() + c := &Client{Base: ts.URL, Allow: []string{"hermes-fleet"}, HTTP: ts.Client()} + items, warn := c.Search(context.Background(), "hermes-fleet", "q") + if warn != "" || len(items) != 1 || items[0].Content != "ok" { + t.Fatalf("%v %q", items, warn) + } +} + +func TestConcurrentSearch(t *testing.T) { + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _ = json.NewEncoder(w).Encode([]map[string]any{{ + "content": "ok", "peer_id": "y", "workspace_id": "hermes-fleet", "session_id": "s", "created_at": "t", + }}) + })) + defer ts.Close() + c := &Client{Base: ts.URL, Allow: []string{"hermes-fleet"}, HTTP: ts.Client()} + var wg sync.WaitGroup + for i := 0; i < 32; i++ { + wg.Add(1) + go func() { + defer wg.Done() + _, _ = c.Search(context.Background(), "hermes-fleet", "q") + }() + } + wg.Wait() + if c.Calls.Load() != 32 { + t.Fatalf("calls %d", c.Calls.Load()) + } +} diff --git a/internal/recover/recover.go b/internal/recover/recover.go index f27c159..40d9595 100644 --- a/internal/recover/recover.go +++ b/internal/recover/recover.go @@ -9,7 +9,6 @@ import ( "net/url" "regexp" "strings" - "time" "git.wntrmute.dev/kyle/sift/internal/fetch" ) @@ -75,6 +74,10 @@ func wayback(ctx context.Context, f Fetcher, original string) (*Result, error) { if !c.Available || c.URL == "" { return nil, fmt.Errorf("wayback: unavailable") } + u, err := url.Parse(c.URL) + if err != nil || u.Scheme != "https" || strings.ToLower(u.Hostname()) != "web.archive.org" { + return nil, fmt.Errorf("wayback: bad snapshot url") + } snapURL := insertID(c.URL) page, err := f.Fetch(ctx, snapURL) if err != nil { @@ -122,7 +125,8 @@ func archiveToday(ctx context.Context, f Fetcher, original string) (*Result, err last = fmt.Errorf("archive.today: fake success") continue } - return &Result{Res: res, Via: ViaToday, SnapshotAt: res.FetchedAt.UTC().Format(time.RFC3339), OriginalURL: original}, nil + at := mementoDatetime(res) + return &Result{Res: res, Via: ViaToday, SnapshotAt: at, OriginalURL: original}, nil } if last == nil { last = fmt.Errorf("archive.today: exhausted") @@ -162,6 +166,13 @@ func bytesTrim(b []byte) []byte { return []byte(strings.TrimSpace(string(b))) } +func mementoDatetime(res *fetch.Resource) string { + if res == nil || res.Header == nil { + return "" + } + return strings.TrimSpace(res.Header.Get("Memento-Datetime")) +} + func hostOf(raw string) string { u, err := url.Parse(raw) if err != nil { diff --git a/internal/recover/recover_test.go b/internal/recover/recover_test.go index 8f9f585..f3e7de0 100644 --- a/internal/recover/recover_test.go +++ b/internal/recover/recover_test.go @@ -69,6 +69,22 @@ func TestWaybackAvailable(t *testing.T) { } } +func TestWaybackRejectsForeignHost(t *testing.T) { + avail, _ := json.Marshal(map[string]any{ + "archived_snapshots": map[string]any{ + "closest": map[string]any{"available": true, "url": "https://evil.example/x", "timestamp": "20240101120000"}, + }, + }) + f := &fakeFetch{ + body: map[string][]byte{"https://archive.org/wayback/available": avail}, + ct: map[string]string{"https://archive.org/wayback/available": "application/json"}, + } + _, err := Try(context.Background(), f, "https://example.com/x") + if err == nil { + t.Fatal("want reject") + } +} + func TestWaybackUnavailableGoesToArchiveToday(t *testing.T) { avail, _ := json.Marshal(map[string]any{"archived_snapshots": map[string]any{}}) page := []byte("Saved" + strings.Repeat("ok ", 200) + "")