diff --git a/internal/api/api.go b/internal/api/api.go
index 3779477..c5b5a98 100644
--- a/internal/api/api.go
+++ b/internal/api/api.go
@@ -353,9 +353,14 @@ func (s *Server) v1fetch(w http.ResponseWriter, r *http.Request) {
}
if snap != nil {
src["kind"] = "snapshot"
- src["snapshot_at"] = snap.SnapshotAt
src["via"] = string(snap.Via)
src["original_url"] = snap.OriginalURL
+ if snap.SnapshotAt != "" {
+ src["snapshot_at"] = snap.SnapshotAt
+ } else {
+ src["snapshot_at"] = nil
+ warn = append(warn, "snapshot date unknown")
+ }
}
out := map[string]any{
"content": content,
diff --git a/internal/digest/digest.go b/internal/digest/digest.go
index ac4df03..c55f9f8 100644
--- a/internal/digest/digest.go
+++ b/internal/digest/digest.go
@@ -17,6 +17,7 @@ import (
"log/slog"
"math"
"net/http"
+ "regexp"
"strings"
"sync"
"time"
@@ -260,7 +261,7 @@ func (e *Engine) Run(ctx context.Context, cleaned, question, finalURL string) Re
if q == "" {
q = "Summarize the key facts."
}
- user := fmt.Sprintf("Question: %s\n\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)?document`)
+
+func fenceDocument(s string) string {
+ return documentTag.ReplaceAllStringFunc(s, func(m string) string {
+ if len(m) >= 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) + "")