diff --git a/internal/api/signals.go b/internal/api/signals.go index 01080c4..7b2591d 100644 --- a/internal/api/signals.go +++ b/internal/api/signals.go @@ -40,6 +40,7 @@ func (s *Server) handleSignal(w http.ResponseWriter, r *http.Request) { ctx := ctxOrBackground(r) metadata := normalizeSignalMetadata(req.Metadata) metadata = s.enrichIncidentMetadata(ctx, req.Chain, req.TxHash, metadata) + metadata = withCandidateTxHashes(metadata, req.CandidateTxHashes) // /signals does not honour force_rerun (hack-detector cannot opt into reruns). c, dedup, err := s.Store.SubmitCase(ctx, req.Chain, req.TxHash, req.Source, req.DetectedAt, metadata, false) if err != nil { @@ -147,6 +148,25 @@ func normalizeSignalMetadata(raw json.RawMessage) json.RawMessage { return data } +// withCandidateTxHashes persists hack-detector's multi-tx candidate list inside +// the case metadata blob (no schema change) so the worker can run drain +// selection later. Single-tx submissions (<2 candidates) are left untouched. +func withCandidateTxHashes(raw json.RawMessage, candidates []string) json.RawMessage { + if len(candidates) < 2 { + return raw + } + m := map[string]any{} + if len(raw) > 0 { + _ = json.Unmarshal(raw, &m) + } + m["candidate_tx_hashes"] = candidates + data, err := json.Marshal(m) + if err != nil { + return raw + } + return data +} + func signalMetadataMergeFields(meta signalMetadataContract) map[string]any { return map[string]any{ "schema_version": signalMetadataSchemaVersion, diff --git a/internal/api/types.go b/internal/api/types.go index e07540c..f6873be 100644 --- a/internal/api/types.go +++ b/internal/api/types.go @@ -15,6 +15,10 @@ type SubmissionRequest struct { DetectedAt *string `json:"detected_at,omitempty"` ForceRerun *bool `json:"force_rerun,omitempty"` Metadata json.RawMessage `json:"metadata,omitempty"` + // CandidateTxHashes carries up to 3 exploit tx candidates when a single + // alert names more than one tx (e.g. authority-setup + asset-drain). The + // worker runs cheap economic selection to pick the real drain before RCA. + CandidateTxHashes []string `json:"candidate_tx_hashes,omitempty"` } // SubmissionResponse is the 202 body for POST /signals and POST /cases. diff --git a/internal/api/validation.go b/internal/api/validation.go index 74bf124..7837ca3 100644 --- a/internal/api/validation.go +++ b/internal/api/validation.go @@ -1,6 +1,7 @@ package api import ( + "fmt" "net/http" "regexp" "time" @@ -13,6 +14,10 @@ var ( txHashRegex = regexp.MustCompile(`^(0x)?[0-9a-fA-F]{64}$`) ) +// maxCandidateTxHashes bounds how many multi-tx exploit candidates a single +// submission may carry (mirrors hack-detector's extractor cap). +const maxCandidateTxHashes = 3 + type validationError struct { Field string `json:"field"` Message string `json:"message"` @@ -28,6 +33,14 @@ func validateSubmission(req *SubmissionRequest) []validationError { } else if !txHashRegex.MatchString(req.TxHash) { errs = append(errs, validationError{Field: "tx_hash", Message: "tx_hash must be a 64-character hex string with optional 0x prefix"}) } + if len(req.CandidateTxHashes) > maxCandidateTxHashes { + errs = append(errs, validationError{Field: "candidate_tx_hashes", Message: fmt.Sprintf("at most %d candidate_tx_hashes are allowed", maxCandidateTxHashes)}) + } + for i, h := range req.CandidateTxHashes { + if !txHashRegex.MatchString(h) { + errs = append(errs, validationError{Field: fmt.Sprintf("candidate_tx_hashes[%d]", i), Message: "candidate_tx_hashes must be a 64-character hex string with optional 0x prefix"}) + } + } if req.DetectedAt != nil && *req.DetectedAt != "" { if _, err := time.Parse(time.RFC3339Nano, *req.DetectedAt); err != nil { if _, err2 := time.Parse(time.RFC3339, *req.DetectedAt); err2 != nil { diff --git a/internal/lumoskit/runner.go b/internal/lumoskit/runner.go index 8604745..8a580f2 100644 --- a/internal/lumoskit/runner.go +++ b/internal/lumoskit/runner.go @@ -43,6 +43,10 @@ func (o RunOptions) stages() []string { return []string{"agent_poc", "rca"} case "agent_poc_repair", "poc_repair", "economic_proof_repair": return []string{"agent_poc_repair", "rca"} + case "flow_context_select": + // Cheap deterministic prefix (no agent/LLM) used by the multi-tx drain + // selector: cefg→localize→lift→flow_context produces fund_flows.json. + return []string{"cefg", "localize", "lift", "flow_context"} default: return []string{o.Stage} } diff --git a/internal/store/cases.go b/internal/store/cases.go index 44ffd7c..876a914 100644 --- a/internal/store/cases.go +++ b/internal/store/cases.go @@ -311,6 +311,17 @@ func (s *Store) AppendCaseEvent(ctx context.Context, caseID, eventType string, e }) } +// UpdateCaseTxHash re-points a case at a different tx hash. The multi-tx drain +// selector uses this when the real asset-drain tx differs from the primary tx +// the case was created with. The (chain, tx_hash) index is non-unique and dedup +// already ran at submission time, so re-pointing is safe. +func (s *Store) UpdateCaseTxHash(ctx context.Context, caseID, txHash string) error { + _, err := s.db.ExecContext(ctx, + `UPDATE cases SET tx_hash = ?, updated_at = ? WHERE case_id = ?`, + txHash, nowUTC(), caseID) + return err +} + func (s *Store) GetCase(ctx context.Context, caseID string) (*Case, error) { row := s.db.QueryRowContext(ctx, ` SELECT case_id, chain, tx_hash, source, detected_at, metadata, diff --git a/internal/worker/selector.go b/internal/worker/selector.go new file mode 100644 index 0000000..2e90c23 --- /dev/null +++ b/internal/worker/selector.go @@ -0,0 +1,186 @@ +package worker + +import ( + "context" + "encoding/json" + "log/slog" + "math/big" + "os" + "path/filepath" + "strings" + "sync" + + "github.com/UPside-Lumos-V2/helios/internal/lumoskit" + "github.com/UPside-Lumos-V2/helios/internal/store" +) + +// drainSelectSubdir holds per-candidate cheap-prefix artifacts under the case +// output root. It is scratch for the selection decision, separate from the +// winner's full-RCA artifacts at the output-root top level. +const drainSelectSubdir = "_drain_select" + +// candidateScore is the per-candidate economic summary recorded on the +// drain_tx_selected event for auditability. +type candidateScore struct { + TxHash string `json:"tx_hash"` + ExitCode int `json:"exit_code"` + NetFlowCount int `json:"net_flow_count"` + MaxAbsDelta string `json:"max_abs_delta"` +} + +type drainSelection struct { + PrimaryTxHash string + WinnerTxHash string + Ranking []candidateScore +} + +// candidateTxHashesFromMetadata pulls the optional candidate_tx_hashes list that +// hack-detector attaches when a single alert names more than one exploit tx. +func candidateTxHashesFromMetadata(metadata json.RawMessage) []string { + var root map[string]json.RawMessage + if len(metadata) == 0 || json.Unmarshal(metadata, &root) != nil { + return nil + } + raw, ok := root["candidate_tx_hashes"] + if !ok || len(raw) == 0 { + return nil + } + var list []string + if json.Unmarshal(raw, &list) != nil { + return nil + } + return list +} + +// selectDrainTxIfNeeded runs drain selection for fresh multi-tx cases and, when +// the winner differs from the case's primary tx, re-points the case at it. It is +// a no-op for stage resumes (the tx is already chosen) and single-tx cases. +func (w *Worker) selectDrainTxIfNeeded(ctx context.Context, c *store.Case, runOptions lumoskit.RunOptions, log *slog.Logger) { + if runOptions.Stage != "" { + return // stage resume reuses the already-selected tx + } + if c.OutputRoot == nil { + return + } + candidates := candidateTxHashesFromMetadata(c.Metadata) + if len(candidates) < 2 { + return + } + + sel := w.selectDrainTx(ctx, c, candidates) + if sel.WinnerTxHash != "" && sel.WinnerTxHash != c.TxHash { + if err := w.Store.UpdateCaseTxHash(ctx, c.CaseID, sel.WinnerTxHash); err != nil { + log.Warn("update case tx_hash to selected drain failed", "err", err, "selected", sel.WinnerTxHash) + } else { + log.Info("drain tx re-pointed", + "from", sel.PrimaryTxHash, "to", sel.WinnerTxHash, "candidates", len(candidates)) + c.TxHash = sel.WinnerTxHash + } + } + if err := w.Store.AppendCaseEvent(ctx, c.CaseID, "drain_tx_selected", map[string]any{ + "primary_tx_hash": sel.PrimaryTxHash, + "selected_tx_hash": sel.WinnerTxHash, + "candidate_count": len(candidates), + "ranking": sel.Ranking, + }); err != nil { + log.Warn("record drain_tx_selected event failed", "err", err) + } +} + +// selectDrainTx runs the cheap deterministic flow_context prefix +// (cefg→localize→lift→flow_context, no agent/LLM) on each candidate tx and picks +// the one whose net fund-flow movement is largest — the real asset-drain. A +// setup/authority tx moves ~no funds, so its net_flows are empty. +// +// Per ADR-0018, helios runs lumoskit only as a subprocess; it reads the +// resulting flow_context artifact and makes the routing decision itself (it does +// not perform analysis). The winner defaults to the primary tx, so a failed or +// movement-free selection never makes things worse than today's behaviour. +func (w *Worker) selectDrainTx(ctx context.Context, c *store.Case, candidates []string) drainSelection { + sel := drainSelection{PrimaryTxHash: c.TxHash, WinnerTxHash: c.TxHash} + + type candResult struct { + score candidateScore + maxAbs *big.Float + ok bool // prefix exited cleanly (ExitCode == 0) + } + results := make([]candResult, len(candidates)) + + // Fan out the per-candidate cheap prefix concurrently. Each candidate writes + // to an isolated output subdir (drainSelectSubdir/), so there is no + // contention, and Runner holds no per-call state. Candidate count is capped + // at maxCandidateTxHashes (3), so unbounded fan-out is bounded in practice. + // Wall-clock drops from N×prefix to ~1×prefix. + var wg sync.WaitGroup + for i, tx := range candidates { + wg.Add(1) + go func(i int, tx string) { + defer wg.Done() + score := candidateScore{TxHash: tx, MaxAbsDelta: "0"} + root := filepath.Join(*c.OutputRoot, drainSelectSubdir, txDirToken(tx)) + res := w.Runner.RunWithOptions(ctx, c.Chain, tx, root, lumoskit.RunOptions{Stage: "flow_context_select"}) + score.ExitCode = res.ExitCode + maxAbs := new(big.Float) + if res.ExitCode == 0 { + count, m := readNetFlowMagnitude(root) + score.NetFlowCount = count + score.MaxAbsDelta = m.Text('f', 0) + maxAbs = m + } + results[i] = candResult{score: score, maxAbs: maxAbs, ok: res.ExitCode == 0} + }(i, tx) + } + wg.Wait() + + // Deterministic selection in candidate order (independent of completion + // order): largest net movement wins, ties keep the earlier candidate. + bestDelta := new(big.Float) + bestFound := false + for i := range candidates { + r := results[i] + sel.Ranking = append(sel.Ranking, r.score) + if r.ok && r.maxAbs.Sign() > 0 && (!bestFound || r.maxAbs.Cmp(bestDelta) > 0) { + bestFound = true + bestDelta = r.maxAbs + sel.WinnerTxHash = candidates[i] + } + } + return sel +} + +func txDirToken(tx string) string { + return strings.TrimPrefix(strings.ToLower(strings.TrimSpace(tx)), "0x") +} + +// readNetFlowMagnitude reads /artifacts/flow_context/fund_flows.json and +// returns the number of net-flow entries and the largest absolute per-holder +// delta. This is a coarse drain-vs-setup discriminator, not a USD figure — the +// precise loss comes from the full RCA on the winner. +func readNetFlowMagnitude(outputRoot string) (count int, maxAbsDelta *big.Float) { + maxAbsDelta = new(big.Float) + path := filepath.Join(outputRoot, "artifacts", "flow_context", "fund_flows.json") + data, err := os.ReadFile(path) + if err != nil { + return 0, maxAbsDelta + } + var doc struct { + NetFlows []struct { + Delta string `json:"delta"` + } `json:"net_flows"` + } + if json.Unmarshal(data, &doc) != nil { + return 0, maxAbsDelta + } + for _, nf := range doc.NetFlows { + abs := strings.TrimPrefix(strings.TrimSpace(nf.Delta), "-") + v, ok := new(big.Float).SetString(abs) + if !ok { + continue + } + count++ + if v.Cmp(maxAbsDelta) > 0 { + maxAbsDelta = v + } + } + return count, maxAbsDelta +} diff --git a/internal/worker/selector_test.go b/internal/worker/selector_test.go new file mode 100644 index 0000000..36f2b66 --- /dev/null +++ b/internal/worker/selector_test.go @@ -0,0 +1,177 @@ +package worker + +import ( + "context" + "encoding/json" + "fmt" + "os" + "path/filepath" + "reflect" + "strings" + "testing" + + "github.com/UPside-Lumos-V2/helios/internal/lumoskit" + "github.com/UPside-Lumos-V2/helios/internal/store" +) + +func TestCandidateTxHashesFromMetadata(t *testing.T) { + cases := []struct { + name string + meta string + want []string + }{ + {"present", `{"candidate_tx_hashes":["0xaa","0xbb"]}`, []string{"0xaa", "0xbb"}}, + {"absent", `{"protocol_name":"X"}`, nil}, + {"empty", ``, nil}, + {"malformed", `not json`, nil}, + {"wrong_type", `{"candidate_tx_hashes":"oops"}`, nil}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + got := candidateTxHashesFromMetadata(json.RawMessage(tc.meta)) + if !reflect.DeepEqual(got, tc.want) { + t.Fatalf("got %v want %v", got, tc.want) + } + }) + } +} + +func TestReadNetFlowMagnitude(t *testing.T) { + root := t.TempDir() + dir := filepath.Join(root, "artifacts", "flow_context") + if err := os.MkdirAll(dir, 0o755); err != nil { + t.Fatal(err) + } + body := `{"net_flows":[{"delta":"-27000000"},{"delta":"15"},{"delta":"100"}]}` + if err := os.WriteFile(filepath.Join(dir, "fund_flows.json"), []byte(body), 0o644); err != nil { + t.Fatal(err) + } + + count, maxAbs := readNetFlowMagnitude(root) + if count != 3 { + t.Fatalf("count=%d want 3", count) + } + if got := maxAbs.Text('f', 0); got != "27000000" { + t.Fatalf("maxAbs=%s want 27000000", got) + } + + // Missing file (e.g. a setup tx whose prefix produced no flows) → zero. + count2, maxAbs2 := readNetFlowMagnitude(t.TempDir()) + if count2 != 0 || maxAbs2.Sign() != 0 { + t.Fatalf("expected zero for missing file, got count=%d max=%s", count2, maxAbs2.Text('f', 0)) + } +} + +func TestSelectDrainTxPicksLargestNetMovement(t *testing.T) { + setup := "0x" + strings.Repeat("a", 64) + drain := "0x" + strings.Repeat("b", 64) + bin := writeSelectorLumoskit(t, t.TempDir(), drain) + + w := &Worker{Runner: &lumoskit.Runner{Binary: bin}} + root := t.TempDir() + c := &store.Case{CaseID: "c1", Chain: "base", TxHash: setup, OutputRoot: &root} + + sel := w.selectDrainTx(context.Background(), c, []string{setup, drain}) + + if sel.WinnerTxHash != drain { + t.Fatalf("expected winner %s, got %s (ranking=%+v)", drain, sel.WinnerTxHash, sel.Ranking) + } + if sel.PrimaryTxHash != setup { + t.Fatalf("expected primary %s, got %s", setup, sel.PrimaryTxHash) + } +} + +func TestSelectDrainTxKeepsPrimaryWhenNoMovement(t *testing.T) { + setup := "0x" + strings.Repeat("a", 64) + other := "0x" + strings.Repeat("c", 64) + // No tx matches the drain marker → both produce empty net_flows. + bin := writeSelectorLumoskit(t, t.TempDir(), "0x"+strings.Repeat("f", 64)) + + w := &Worker{Runner: &lumoskit.Runner{Binary: bin}} + root := t.TempDir() + c := &store.Case{CaseID: "c1", Chain: "base", TxHash: setup, OutputRoot: &root} + + sel := w.selectDrainTx(context.Background(), c, []string{setup, other}) + + if sel.WinnerTxHash != setup { + t.Fatalf("expected winner to stay primary %s, got %s", setup, sel.WinnerTxHash) + } +} + +func TestSelectDrainTxKeepsPrimaryWhenNoMovementAndPrimaryIsNotFirstCandidate(t *testing.T) { + setup := "0x" + strings.Repeat("a", 64) + other := "0x" + strings.Repeat("c", 64) + // No tx matches the drain marker → both produce empty net_flows. + bin := writeSelectorLumoskit(t, t.TempDir(), "0x"+strings.Repeat("f", 64)) + + w := &Worker{Runner: &lumoskit.Runner{Binary: bin}} + root := t.TempDir() + c := &store.Case{CaseID: "c1", Chain: "base", TxHash: setup, OutputRoot: &root} + + sel := w.selectDrainTx(context.Background(), c, []string{other, setup}) + + if sel.WinnerTxHash != setup { + t.Fatalf("expected winner to stay primary %s, got %s", setup, sel.WinnerTxHash) + } +} + +func TestSelectDrainTxThreeCandidatesParallelDeterministic(t *testing.T) { + a := "0x" + strings.Repeat("a", 64) + drain := "0x" + strings.Repeat("b", 64) + c3 := "0x" + strings.Repeat("c", 64) + bin := writeSelectorLumoskit(t, t.TempDir(), drain) + + w := &Worker{Runner: &lumoskit.Runner{Binary: bin}} + root := t.TempDir() + c := &store.Case{CaseID: "c1", Chain: "base", TxHash: a, OutputRoot: &root} + + // Candidates fan out concurrently; ranking must still come back in candidate + // order and the winner must be the drain regardless of completion order. + sel := w.selectDrainTx(context.Background(), c, []string{a, drain, c3}) + + if sel.WinnerTxHash != drain { + t.Fatalf("expected winner %s, got %s", drain, sel.WinnerTxHash) + } + if len(sel.Ranking) != 3 { + t.Fatalf("expected 3 ranking entries, got %d", len(sel.Ranking)) + } + for i, want := range []string{a, drain, c3} { + if sel.Ranking[i].TxHash != want { + t.Fatalf("ranking[%d]=%s want %s (order not preserved)", i, sel.Ranking[i].TxHash, want) + } + } +} + +// writeSelectorLumoskit emits a fake lumoskit that writes a large net_flows +// fund_flows.json only for drainTx, and an empty one for every other tx. +func writeSelectorLumoskit(t *testing.T, dir, drainTx string) string { + t.Helper() + path := filepath.Join(dir, "fake-selector-lumoskit") + script := fmt.Sprintf(`#!/bin/sh +set -eu +out=""; tx="" +while [ "$#" -gt 0 ]; do + case "$1" in + --output-root) out="$2"; shift 2 ;; + --tx) tx="$2"; shift 2 ;; + --chain) shift 2 ;; + --stage) shift 2 ;; + *) shift ;; + esac +done +mkdir -p "$out/artifacts/flow_context" +if [ "$tx" = "%s" ]; then + cat > "$out/artifacts/flow_context/fund_flows.json" <<'JSON' +{"net_flows":[{"holder":"0xvictim","asset":"TRAC","delta":"-27000000","flow_ids":["f1"]},{"holder":"0xattacker","asset":"TRAC","delta":"27000000","flow_ids":["f1"]}]} +JSON +else + cat > "$out/artifacts/flow_context/fund_flows.json" <<'JSON' +{"net_flows":[]} +JSON +fi +`, drainTx) + if err := os.WriteFile(path, []byte(script), 0o755); err != nil { + t.Fatal(err) + } + return path +} diff --git a/internal/worker/worker.go b/internal/worker/worker.go index 24fba23..e0bf2e3 100644 --- a/internal/worker/worker.go +++ b/internal/worker/worker.go @@ -159,6 +159,11 @@ func (w *Worker) process(ctx context.Context, c *store.Case) { log.Info("stage resume prepared", "stage", runOptions.Stage, "source_case_id", resumePayload["resume_source_case_id"]) } + // Multi-tx incident: if the alert carried several candidate txs, pick the + // real asset-drain (cheap, deterministic, no LLM) before the expensive RCA. + // This may re-point c.TxHash, so it runs before signal-context is written. + w.selectDrainTxIfNeeded(ctx, c, runOptions, log) + if err := writeSignalContext(c); err != nil { log.Warn("write signal context failed", "err", err) }