package api import ( "encoding/json" "io" "log/slog" "net/http" "net/http/httptest" "net/netip" "os" "path/filepath" "strings" "testing" "time" "git.wntrmute.dev/kyle/sift/internal/auth" "git.wntrmute.dev/kyle/sift/internal/batch" "git.wntrmute.dev/kyle/sift/internal/dcache" "git.wntrmute.dev/kyle/sift/internal/fetch" "git.wntrmute.dev/kyle/sift/internal/guard" ) // newBatchServer builds a Server with the batch runner wired over an in-memory // dcache, exactly as serve does in production. callers are the caller names to // provision (each with a bearer key). func newBatchServer(t *testing.T, g *guard.Guard, callers ...string) (*Server, map[string]string) { t.Helper() dir := t.TempDir() keysPath := filepath.Join(dir, "keys.toml") keys := map[string]string{} for _, name := range callers { out := filepath.Join(dir, name) if err := auth.Add(keysPath, name, "hermes-fleet", out); err != nil { t.Fatal(err) } b, err := os.ReadFile(out) if err != nil { t.Fatal(err) } keys[name] = strings.TrimSpace(string(b)) } store, err := auth.Load(keysPath) if err != nil { t.Fatal(err) } if g == nil { g = &guard.Guard{} } f := &fetch.Fetcher{Guard: g, UA: "sift-test"} s := New("test", g, f, store, slog.New(slog.NewTextHandler(io.Discard, nil))) db, err := dcache.Open("", 1<<20, slog.New(slog.NewTextHandler(io.Discard, nil))) if err != nil { t.Fatal(err) } t.Cleanup(func() { _ = db.Close() }) s.Batch = batch.New(db, s.fetchTarget, batch.Config{}, slog.New(slog.NewTextHandler(io.Discard, nil))) return s, keys } func submitBatch(t *testing.T, s *Server, key string, targets []map[string]string) *httptest.ResponseRecorder { t.Helper() return postJSON(t, s, "/v1/batch", key, map[string]any{"targets": targets}) } func getBatch(t *testing.T, s *Server, key, id string) *httptest.ResponseRecorder { t.Helper() req := httptest.NewRequest(http.MethodGet, "/v1/batch/"+id, nil) req.Header.Set("Authorization", "Bearer "+key) rec := httptest.NewRecorder() s.ServeHTTP(rec, req) return rec } func batchStatus(t *testing.T, s *Server, key, id string) map[string]any { t.Helper() rec := getBatch(t, s, key, id) var env map[string]any if err := json.Unmarshal(rec.Body.Bytes(), &env); err != nil { t.Fatalf("status decode: %v (%s)", err, rec.Body.String()) } return env } // waitBatchTerminal polls until the job reaches a terminal status. func waitBatchTerminal(t *testing.T, s *Server, key, id string) map[string]any { t.Helper() deadline := time.Now().Add(3 * time.Second) for time.Now().Before(deadline) { env := batchStatus(t, s, key, id) if st, _ := env["status"].(string); st == "done" || st == "partial" || st == "failed" || st == "cancelled" { return env } time.Sleep(time.Millisecond) } t.Fatal("batch never reached a terminal status") return nil } func TestBatchSubmitAndCollect(t *testing.T) { g := &guard.Guard{AllowPrefixes: []netip.Prefix{netip.MustParsePrefix("127.0.0.0/8")}} ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { w.Header().Set("Content-Type", "text/html; charset=utf-8") _, _ = w.Write([]byte(`T

hello batch

`)) })) defer ts.Close() s, keys := newBatchServer(t, g, "test") rec := submitBatch(t, s, keys["test"], []map[string]string{{"url": ts.URL}, {"url": ts.URL + "/2"}}) if rec.Code != 202 { t.Fatalf("submit %d %s", rec.Code, rec.Body.String()) } var sub map[string]any if err := json.Unmarshal(rec.Body.Bytes(), &sub); err != nil { t.Fatal(err) } if sub["success"] != true { t.Fatalf("success %v", sub["success"]) } id, _ := sub["job_id"].(string) if id == "" { t.Fatal("missing job_id") } if reused, _ := sub["reused"].(bool); reused { t.Fatal("first submit must not be reused") } env := waitBatchTerminal(t, s, keys["test"], id) if env["status"] != "done" { t.Fatalf("status %v", env["status"]) } targets, _ := env["targets"].([]any) if len(targets) != 2 { t.Fatalf("targets %d", len(targets)) } first := targets[0].(map[string]any) if first["status"] != "done" { t.Fatalf("target0 %v", first["status"]) } res, _ := first["result"].(map[string]any) if res == nil || !strings.Contains(res["content"].(string), "hello batch") { t.Fatalf("target0 result %v", res) } } func TestBatchIdempotentResubmit(t *testing.T) { g := &guard.Guard{AllowPrefixes: []netip.Prefix{netip.MustParsePrefix("127.0.0.0/8")}} ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { w.Header().Set("Content-Type", "text/plain") _, _ = w.Write([]byte("ok")) })) defer ts.Close() s, keys := newBatchServer(t, g, "test") targets := []map[string]string{{"url": ts.URL}} first := submitBatch(t, s, keys["test"], targets) second := submitBatch(t, s, keys["test"], targets) var a, b map[string]any _ = json.Unmarshal(first.Body.Bytes(), &a) _ = json.Unmarshal(second.Body.Bytes(), &b) if a["job_id"] != b["job_id"] { t.Fatalf("ids differ: %v vs %v", a["job_id"], b["job_id"]) } if reused, _ := b["reused"].(bool); !reused { t.Fatal("resubmit must be reused=true") } } func TestBatchGuardRefusedUpFront(t *testing.T) { s, keys := newBatchServer(t, nil, "test") rec := submitBatch(t, s, keys["test"], []map[string]string{{"url": "http://example.com:8080/"}}) if rec.Code != 422 { t.Fatalf("want 422, got %d %s", rec.Code, rec.Body.String()) } if !strings.Contains(rec.Body.String(), "blocked") { t.Fatalf("body %s", rec.Body.String()) } } func TestBatchTooManyTargets(t *testing.T) { s, keys := newBatchServer(t, nil, "test") var targets []map[string]string for i := 0; i < 33; i++ { targets = append(targets, map[string]string{"url": "https://example.com/"}) } rec := submitBatch(t, s, keys["test"], targets) if rec.Code != 422 { t.Fatalf("want 422, got %d %s", rec.Code, rec.Body.String()) } } func TestBatchPartialFailure(t *testing.T) { g := &guard.Guard{AllowPrefixes: []netip.Prefix{netip.MustParsePrefix("127.0.0.0/8")}} ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { w.Header().Set("Content-Type", "text/plain") _, _ = w.Write([]byte("ok")) })) defer ts.Close() s, keys := newBatchServer(t, g, "test") // target 1 is a private address: passes the submit pre-check (valid URL) but // is blocked by the connect-time guard, surfacing as a per-target failure. rec := submitBatch(t, s, keys["test"], []map[string]string{{"url": ts.URL}, {"url": "http://10.0.0.5/"}}) if rec.Code != 202 { t.Fatalf("submit %d %s", rec.Code, rec.Body.String()) } var sub map[string]any _ = json.Unmarshal(rec.Body.Bytes(), &sub) id := sub["job_id"].(string) env := waitBatchTerminal(t, s, keys["test"], id) if env["status"] != "partial" { t.Fatalf("status %v", env["status"]) } targets := env["targets"].([]any) second := targets[1].(map[string]any) if second["status"] != "failed" { t.Fatalf("target1 status %v", second["status"]) } if err, _ := second["error"].(string); !strings.Contains(err, "blocked") { t.Fatalf("target1 error %q", err) } } func TestBatchOwnership404(t *testing.T) { s, keys := newBatchServer(t, nil, "test", "other") rec := submitBatch(t, s, keys["test"], []map[string]string{{"url": "https://example.com/"}}) if rec.Code != 202 { t.Fatalf("submit %d", rec.Code) } var sub map[string]any _ = json.Unmarshal(rec.Body.Bytes(), &sub) id := sub["job_id"].(string) if got := getBatch(t, s, keys["other"], id); got.Code != 404 { t.Fatalf("other caller GET: %d (must not reveal the job)", got.Code) } if got := getBatch(t, s, keys["test"], id); got.Code != 200 { t.Fatalf("owner GET: %d", got.Code) } } func TestBatchCancel(t *testing.T) { g := &guard.Guard{AllowPrefixes: []netip.Prefix{netip.MustParsePrefix("127.0.0.0/8")}} s, keys := newBatchServer(t, g, "test") rec := submitBatch(t, s, keys["test"], []map[string]string{{"url": "https://example.com/"}}) var sub map[string]any _ = json.Unmarshal(rec.Body.Bytes(), &sub) id := sub["job_id"].(string) req := httptest.NewRequest(http.MethodDelete, "/v1/batch/"+id, nil) req.Header.Set("Authorization", "Bearer "+keys["test"]) drec := httptest.NewRecorder() s.ServeHTTP(drec, req) if drec.Code != 200 { t.Fatalf("cancel %d %s", drec.Code, drec.Body.String()) } env := waitBatchTerminal(t, s, keys["test"], id) if env["status"] != "cancelled" { t.Fatalf("status %v", env["status"]) } }