First Commit
ci / go (push) Waiting to run
ci / go-db (agent) (push) Waiting to run
ci / go-db (config) (push) Waiting to run
ci / go-db (db) (push) Waiting to run
ci / go-db (evidence) (push) Waiting to run
ci / go-db (llmrec) (push) Waiting to run
ci / go-db (server) (push) Waiting to run
detections / detections (push) Waiting to run
web / web (push) Waiting to run
docs / links (push) Canceled after 0s
ci / go (push) Waiting to run
ci / go-db (agent) (push) Waiting to run
ci / go-db (config) (push) Waiting to run
ci / go-db (db) (push) Waiting to run
ci / go-db (evidence) (push) Waiting to run
ci / go-db (llmrec) (push) Waiting to run
ci / go-db (server) (push) Waiting to run
detections / detections (push) Waiting to run
web / web (push) Waiting to run
docs / links (push) Canceled after 0s
This commit is contained in:
@@ -0,0 +1,419 @@
|
||||
// Package evidence preserves finding evidence independently of disposable traffic.
|
||||
package evidence
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"database/sql"
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"io/fs"
|
||||
"log"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/artex/db"
|
||||
"github.com/Autumn-27/artex/traffic"
|
||||
)
|
||||
|
||||
type Store struct {
|
||||
DB *db.DB
|
||||
Traffic *traffic.Traffic
|
||||
Dir string
|
||||
}
|
||||
|
||||
func New(pg *db.DB, tr *traffic.Traffic, dir string) *Store {
|
||||
return &Store{DB: pg, Traffic: tr, Dir: dir}
|
||||
}
|
||||
|
||||
func hashPath(dir, hash string) (string, error) {
|
||||
if len(hash) != 64 {
|
||||
return "", errors.New("invalid evidence hash")
|
||||
}
|
||||
if _, err := hex.DecodeString(hash); err != nil || strings.ToLower(hash) != hash {
|
||||
return "", errors.New("invalid evidence hash")
|
||||
}
|
||||
return filepath.Join(dir, "blobs", hash[:2], hash+".bin"), nil
|
||||
}
|
||||
|
||||
func verifyFile(path, hash string, length int64) error {
|
||||
f, err := os.Open(path)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer f.Close()
|
||||
h := sha256.New()
|
||||
n, err := io.Copy(h, f)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if n != length || hex.EncodeToString(h.Sum(nil)) != hash {
|
||||
return fmt.Errorf("证据正文校验失败: %s", hash)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// A new body becomes visible only after a durable write. Failed SQL commits may
|
||||
// leave unreferenced files; GC reaps those after a full day's grace period.
|
||||
func (s *Store) writeBody(r io.Reader, expectedLength int64, expectedHash string) (hash string, err error) {
|
||||
stage := filepath.Join(s.Dir, ".staging")
|
||||
if err = os.MkdirAll(stage, 0o700); err != nil {
|
||||
return
|
||||
}
|
||||
f, err := os.CreateTemp(stage, "body-")
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
defer func() { f.Close(); os.Remove(f.Name()) }()
|
||||
h := sha256.New()
|
||||
n, err := io.Copy(io.MultiWriter(f, h), r)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
hash = hex.EncodeToString(h.Sum(nil))
|
||||
if n != expectedLength {
|
||||
return "", fmt.Errorf("正文不完整: 预期 %d 字节,读取 %d 字节", expectedLength, n)
|
||||
}
|
||||
if expectedHash != "" && expectedHash != hash {
|
||||
return "", errors.New("原始流量正文哈希不匹配")
|
||||
}
|
||||
if err = f.Sync(); err != nil {
|
||||
return "", err
|
||||
}
|
||||
if err = f.Close(); err != nil {
|
||||
return "", err
|
||||
}
|
||||
path, err := hashPath(s.Dir, hash)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if _, err = os.Stat(path); err == nil {
|
||||
if err = verifyFile(path, hash, n); err != nil {
|
||||
return "", err
|
||||
}
|
||||
// Refresh the grace period for a restored but not-yet-committed body.
|
||||
now := time.Now()
|
||||
return hash, os.Chtimes(path, now, now)
|
||||
} else if !os.IsNotExist(err) {
|
||||
return "", err
|
||||
}
|
||||
if err = os.MkdirAll(filepath.Dir(path), 0o700); err != nil {
|
||||
return "", err
|
||||
}
|
||||
if err = os.Rename(f.Name(), path); err != nil {
|
||||
return "", err
|
||||
}
|
||||
d, err := os.Open(filepath.Dir(path))
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
defer d.Close()
|
||||
return hash, d.Sync()
|
||||
}
|
||||
|
||||
// StageFindingsExport freezes bindings/report versions and makes private body
|
||||
// copies before the HTTP response is started. The caller owns and removes dest.
|
||||
func (s *Store) StageFindingsExport(ctx context.Context, findings []*db.DBFinding, dest string, copyBodies bool) error {
|
||||
return s.DB.WithEvidenceTx(ctx, func(tx *sql.Tx) error {
|
||||
var snapshots []db.TrafficEvidenceSnapshot
|
||||
for _, f := range findings {
|
||||
list, err := db.FindingTrafficTx(tx, f.ID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
f.TrafficBindings = list.Bindings
|
||||
f.TrafficCount = len(list.Bindings)
|
||||
f.EvidenceVersion = list.Version
|
||||
f.ReportEvidenceVersion = list.ReportVersion
|
||||
if err = tx.QueryRow(`SELECT report FROM findings WHERE id=$1`, f.ID).Scan(&f.Report); err != nil {
|
||||
return err
|
||||
}
|
||||
for _, b := range list.Bindings {
|
||||
snapshots = append(snapshots, b.Snapshot)
|
||||
}
|
||||
}
|
||||
if copyBodies {
|
||||
return s.copySnapshots(snapshots, dest)
|
||||
}
|
||||
return nil
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Store) prepare(ctx context.Context, refs []db.TrafficRef) ([]db.PreparedTrafficEvidence, error) {
|
||||
refs, err := db.NormalizeTrafficRefs(refs)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out := make([]db.PreparedTrafficEvidence, 0, len(refs))
|
||||
if len(refs) == 0 {
|
||||
return out, nil
|
||||
}
|
||||
ids := make([]string, len(refs))
|
||||
byID := map[string]db.TrafficRef{}
|
||||
for i, ref := range refs {
|
||||
ids[i] = ref.TrafficID
|
||||
byID[ref.TrafficID] = ref
|
||||
}
|
||||
err = s.Traffic.ReadEvidence(ctx, ids, func(e traffic.EvidenceExchange) error {
|
||||
rh, err := s.writeBody(e.Request, e.ReqLen, e.ReqHash)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
ph, err := s.writeBody(e.Response, e.RespLen, e.RespHash)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
v := db.TrafficEvidenceSnapshot{SourceTrafficID: e.ID, CapturedAt: e.TS, URL: e.URL, Method: e.Method, Status: e.Status, ContentType: e.ContentType,
|
||||
ReqHead: e.ReqHead, RespHead: e.RespHead, ReqHash: rh, RespHash: ph, ReqLen: e.ReqLen, RespLen: e.RespLen}
|
||||
// Raw wire bytes: normalize once here so the ID, the stored row and every
|
||||
// downstream consumer (archive, API responses) all see the same text.
|
||||
v = v.Normalize()
|
||||
v.ID = db.TrafficSnapshotID(v)
|
||||
out = append(out, db.PreparedTrafficEvidence{Ref: byID[e.ID], Snapshot: v})
|
||||
return nil
|
||||
})
|
||||
return out, err
|
||||
}
|
||||
|
||||
func (s *Store) Record(ctx context.Context, in db.RecordFindingInput, refs []db.TrafficRef) (out *db.RecordedFinding, err error) {
|
||||
err = s.DB.WithEvidenceTx(ctx, func(tx *sql.Tx) error {
|
||||
if err := db.LockTaskEvidenceTx(tx, in.TaskID); err != nil {
|
||||
return err
|
||||
}
|
||||
prepared, err := s.prepare(ctx, refs)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
out, err = db.RecordFindingTx(ctx, tx, in, prepared)
|
||||
return err
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
func (s *Store) Bind(ctx context.Context, findingID int64, refs []db.TrafficRef) (out *db.FindingTraffic, err error) {
|
||||
err = s.DB.WithEvidenceTx(ctx, func(tx *sql.Tx) error {
|
||||
if err := db.LockFindingEvidenceTx(tx, findingID, nil); err != nil {
|
||||
return err
|
||||
}
|
||||
prepared, err := s.prepare(ctx, refs)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err = db.AddFindingTrafficTx(tx, findingID, prepared); err != nil {
|
||||
return err
|
||||
}
|
||||
out, err = db.FindingTrafficTx(tx, findingID)
|
||||
return err
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
func (s *Store) WithBinding(ctx context.Context, findingID, bindingID int64, fn func(db.FindingTrafficBinding) error) error {
|
||||
return s.DB.WithEvidenceTx(ctx, func(tx *sql.Tx) error {
|
||||
list, err := db.FindingTrafficTx(tx, findingID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
for _, b := range list.Bindings {
|
||||
if b.ID == bindingID {
|
||||
return fn(b)
|
||||
}
|
||||
}
|
||||
return db.ErrEvidenceNotFound
|
||||
})
|
||||
}
|
||||
|
||||
// Binding resolves one binding's metadata under the evidence lock and releases
|
||||
// the lock before returning. Callers that then stream a body to a client must
|
||||
// use this instead of WithBinding: verifyFile+io.Copy is O(body size), so a
|
||||
// large download (or a slow client) holding WithEvidenceTx would block every
|
||||
// evidence write process-wide. Reading the blob afterwards is safe — blobs are
|
||||
// content-addressed and GC only reaps unreferenced files after a 24h grace
|
||||
// period, and an already-open fd survives an unlink regardless.
|
||||
func (s *Store) Binding(ctx context.Context, findingID, bindingID int64) (db.FindingTrafficBinding, error) {
|
||||
var out db.FindingTrafficBinding
|
||||
err := s.WithBinding(ctx, findingID, bindingID, func(b db.FindingTrafficBinding) error {
|
||||
out = b
|
||||
return nil
|
||||
})
|
||||
return out, err
|
||||
}
|
||||
|
||||
// OpenBody may be called without holding the evidence lock; see Binding.
|
||||
func (s *Store) OpenBody(snapshot db.TrafficEvidenceSnapshot, side string) (*os.File, int64, error) {
|
||||
hash, length := snapshot.ReqHash, snapshot.ReqLen
|
||||
if side == "response" {
|
||||
hash, length = snapshot.RespHash, snapshot.RespLen
|
||||
} else if side != "request" {
|
||||
return nil, 0, errors.New("side 必须为 request 或 response")
|
||||
}
|
||||
path, err := hashPath(s.Dir, hash)
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
if err = verifyFile(path, hash, length); err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
f, err := os.Open(path)
|
||||
return f, length, err
|
||||
}
|
||||
|
||||
// CopySnapshots is used by both report downloads and portable task archives.
|
||||
// The destination owns real copies, never links into either disposable store.
|
||||
func (s *Store) CopySnapshots(ctx context.Context, snapshots []db.TrafficEvidenceSnapshot, dest string) error {
|
||||
return s.DB.WithEvidenceTx(ctx, func(*sql.Tx) error { return s.copySnapshots(snapshots, dest) })
|
||||
}
|
||||
|
||||
func (s *Store) copySnapshots(snapshots []db.TrafficEvidenceSnapshot, dest string) error {
|
||||
for _, v := range snapshots {
|
||||
for _, side := range []string{"request", "response"} {
|
||||
if err := func() error {
|
||||
f, length, err := s.OpenBody(v, side)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer f.Close()
|
||||
hash := v.ReqHash
|
||||
if side == "response" {
|
||||
hash = v.RespHash
|
||||
}
|
||||
target, err := hashPath(dest, hash)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err = os.MkdirAll(filepath.Dir(target), 0o700); err != nil {
|
||||
return err
|
||||
}
|
||||
out, err := os.OpenFile(target, os.O_CREATE|os.O_WRONLY|os.O_EXCL, 0o600)
|
||||
if os.IsExist(err) {
|
||||
return verifyFile(target, hash, length)
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
_, copyErr := io.Copy(out, f)
|
||||
syncErr := out.Sync()
|
||||
closeErr := out.Close()
|
||||
return errors.Join(copyErr, syncErr, closeErr)
|
||||
}(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Store) InstallSnapshots(ctx context.Context, snapshots []db.TrafficEvidenceSnapshot, source string) error {
|
||||
return s.WithInstalledSnapshots(ctx, snapshots, source, func() error { return nil })
|
||||
}
|
||||
|
||||
// WithInstalledSnapshots pins installed bodies until the metadata restore finishes.
|
||||
// The callback must not acquire another evidence advisory lock.
|
||||
func (s *Store) WithInstalledSnapshots(ctx context.Context, snapshots []db.TrafficEvidenceSnapshot, source string, restore func() error) error {
|
||||
return s.DB.WithEvidenceTx(ctx, func(*sql.Tx) error {
|
||||
for _, v := range snapshots {
|
||||
if v.ID != db.TrafficSnapshotID(v) {
|
||||
return errors.New("归档证据快照元数据哈希不匹配")
|
||||
}
|
||||
for _, body := range []struct {
|
||||
hash string
|
||||
length int64
|
||||
}{{v.ReqHash, v.ReqLen}, {v.RespHash, v.RespLen}} {
|
||||
if err := func() error {
|
||||
path, err := hashPath(source, body.hash)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
f, err := os.Open(path)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer f.Close()
|
||||
_, err = s.writeBody(f, body.length, body.hash)
|
||||
return err
|
||||
}(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
return restore()
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Store) Collect(ctx context.Context, now time.Time) error {
|
||||
return s.DB.WithEvidenceTx(ctx, func(tx *sql.Tx) error {
|
||||
if _, err := tx.Exec(`UPDATE traffic_evidence_snapshots s SET unreferenced_at=$1
|
||||
WHERE unreferenced_at IS NULL AND NOT EXISTS(SELECT 1 FROM finding_traffic_bindings b WHERE b.snapshot_id=s.id)`, now); err != nil {
|
||||
return err
|
||||
}
|
||||
if _, err := tx.Exec(`DELETE FROM traffic_evidence_snapshots s WHERE unreferenced_at<$1
|
||||
AND NOT EXISTS(SELECT 1 FROM finding_traffic_bindings b WHERE b.snapshot_id=s.id)`, now.Add(-24*time.Hour)); err != nil {
|
||||
return err
|
||||
}
|
||||
rows, err := tx.Query(`SELECT req_hash FROM traffic_evidence_snapshots UNION SELECT resp_hash FROM traffic_evidence_snapshots`)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
refs := map[string]bool{}
|
||||
for rows.Next() {
|
||||
var hash string
|
||||
if err := rows.Scan(&hash); err != nil {
|
||||
rows.Close()
|
||||
return err
|
||||
}
|
||||
refs[hash] = true
|
||||
}
|
||||
err = rows.Err()
|
||||
rows.Close()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
for _, root := range []string{filepath.Join(s.Dir, "blobs"), filepath.Join(s.Dir, ".staging")} {
|
||||
err := filepath.WalkDir(root, func(path string, entry fs.DirEntry, err error) error {
|
||||
if os.IsNotExist(err) {
|
||||
return nil
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if entry.IsDir() {
|
||||
return nil
|
||||
}
|
||||
if refs[strings.TrimSuffix(entry.Name(), ".bin")] {
|
||||
return nil
|
||||
}
|
||||
info, err := entry.Info()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if now.Sub(info.ModTime()) < 24*time.Hour {
|
||||
return nil
|
||||
}
|
||||
return os.Remove(path)
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Store) RunGC(ctx context.Context) {
|
||||
timer := time.NewTicker(time.Hour)
|
||||
defer timer.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case now := <-timer.C:
|
||||
if err := s.Collect(ctx, now); err != nil && !errors.Is(err, context.Canceled) {
|
||||
log.Printf("[evidence] cleanup failed: %v", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,333 @@
|
||||
package evidence
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/artex/db"
|
||||
"github.com/Autumn-27/artex/traffic"
|
||||
)
|
||||
|
||||
func evidenceFixture(t *testing.T) (*Store, db.RecordFindingInput, string) {
|
||||
t.Helper()
|
||||
dsn := os.Getenv("ARTEX_PG_DSN")
|
||||
if dsn == "" {
|
||||
t.Skip("ARTEX_PG_DSN is required for evidence integration tests")
|
||||
}
|
||||
pg, err := db.Open(dsn)
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) — skipping", err)
|
||||
}
|
||||
t.Cleanup(func() { pg.Close() })
|
||||
task, err := pg.CreateTask("evidence integration "+t.Name(), "local fixtures", nil, 0, 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { pg.DeleteFindingsByTask(task.ID); pg.DeleteTask(task.ID) })
|
||||
dir := t.TempDir()
|
||||
tr, err := traffic.Open(filepath.Join(dir, "traffic"), ":0")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { tr.Close() })
|
||||
store := New(pg, tr, filepath.Join(dir, "evidence"))
|
||||
in := db.RecordFindingInput{TaskID: task.ID, ExplorationID: task.ExplorationID, Worker: "test", VulnClass: "TEST", Name: "Evidence fixture", Severity: "low", Summary: "local test"}
|
||||
return store, in, dir
|
||||
}
|
||||
|
||||
func seedExchange(t *testing.T, s *Store, id string, body []byte, spill bool) {
|
||||
t.Helper()
|
||||
var blob any
|
||||
inline := body
|
||||
if spill {
|
||||
sum := sha256.Sum256(body)
|
||||
hash := hex.EncodeToString(sum[:])
|
||||
blob = hash
|
||||
path := filepath.Join(filepath.Dir(s.Dir), "traffic", "_blobs", "sha256", hash[:2], hash+".bin")
|
||||
if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := os.WriteFile(path, body, 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
inline = []byte("TRUNCATED PREVIEW")
|
||||
}
|
||||
_, err := s.Traffic.DB().Exec(`INSERT INTO exchanges(id,ts,host,method,url_template,url,status,content_type,req_len,resp_len,path) VALUES(?,?,'fixture.local','POST','/test',?,200,'application/octet-stream',2,?,'')`, id, time.Now().Unix(), "https://fixture.local/"+id, len(body))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_, err = s.Traffic.DB().Exec(`INSERT INTO exchange_bodies(id,req_head,req_body,resp_head,resp_body,resp_blob) VALUES(?,?,?, ?,?,?)`, id, "POST /test HTTP/1.1\nHost: fixture.local\n", []byte("{}"), "HTTP 200\nContent-Type: application/octet-stream\n", inline, blob)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEvidenceBindingLifecycle(t *testing.T) {
|
||||
s, in, _ := evidenceFixture(t)
|
||||
ctx := context.Background()
|
||||
large := bytes.Repeat([]byte{0, 1, 2, 255, 'a', 'b'}, 180000)
|
||||
for i := 0; i < 3; i++ {
|
||||
seedExchange(t, s, fmt.Sprint(i), large, i == 2)
|
||||
}
|
||||
refs := []db.TrafficRef{{TrafficID: "0", Role: "baseline", Note: "normal"}, {TrafficID: "1", Role: "proof", Note: "proof"}, {TrafficID: "2", Role: "verification", Note: "binary"}}
|
||||
r, err := s.Record(ctx, in, refs)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(r.Traffic.Bindings) != 3 || r.Traffic.Version != 1 {
|
||||
t.Fatalf("record=%+v", r)
|
||||
}
|
||||
again, err := s.Bind(ctx, r.FindingID, []db.TrafficRef{{TrafficID: "1", Note: "must not overwrite"}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(again.Bindings) != 3 || again.Version != 1 || again.Bindings[1].Note != "proof" {
|
||||
t.Fatalf("duplicate changed bindings: %+v", again)
|
||||
}
|
||||
other, err := s.Record(ctx, in, []db.TrafficRef{refs[2]})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if other.Traffic.Bindings[0].SnapshotID != r.Traffic.Bindings[2].SnapshotID {
|
||||
t.Fatal("snapshot was not shared")
|
||||
}
|
||||
if _, err = s.Traffic.DeleteHost("fixture.local"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, binding := range r.Traffic.Bindings {
|
||||
if err = s.WithBinding(ctx, r.FindingID, binding.ID, func(b db.FindingTrafficBinding) error {
|
||||
f, _, err := s.OpenBody(b.Snapshot, "response")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer f.Close()
|
||||
got, err := io.ReadAll(f)
|
||||
if !bytes.Equal(got, large) {
|
||||
t.Fatal("body was truncated/changed")
|
||||
}
|
||||
return err
|
||||
}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
v := again.Version
|
||||
if _, err = s.DB.SetFindingReportVersionByNodeID(ctx, r.NodeID, "report", &v); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
note := "updated"
|
||||
if err = s.DB.EditFindingTraffic(ctx, r.FindingID, again.Bindings[0].ID, v, nil, ¬e, false, nil); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err = s.DB.SetFindingReportVersionByNodeID(ctx, r.NodeID, "stale report", &v); !errors.Is(err, db.ErrEvidenceConflict) {
|
||||
t.Fatalf("stale report: %v", err)
|
||||
}
|
||||
f, err := s.DB.GetFinding(r.FindingID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if f.Report != "report" || f.EvidenceVersion == f.ReportEvidenceVersion || f.TrafficCount != 3 {
|
||||
t.Fatalf("finding versions: %+v", f)
|
||||
}
|
||||
if err = s.DB.EditFindingTraffic(ctx, r.FindingID, again.Bindings[0].ID, v, nil, nil, true, nil); !errors.Is(err, db.ErrEvidenceConflict) {
|
||||
t.Fatalf("stale delete: %v", err)
|
||||
}
|
||||
if _, err = s.DB.SetFindingReportByNodeID(r.NodeID, "legacy report"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
f, err = s.DB.GetFinding(r.FindingID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if f.ReportEvidenceVersion != -1 {
|
||||
t.Fatal("legacy write claimed current evidence")
|
||||
}
|
||||
if _, err = s.DB.DeleteFinding(r.FindingID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err = s.Collect(ctx, time.Now().Add(48*time.Hour)); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err = s.WithBinding(ctx, other.FindingID, other.Traffic.Bindings[0].ID, func(b db.FindingTrafficBinding) error {
|
||||
f, _, err := s.OpenBody(b.Snapshot, "response")
|
||||
if f != nil {
|
||||
f.Close()
|
||||
}
|
||||
return err
|
||||
}); err != nil {
|
||||
t.Fatal("GC removed shared body:", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEvidenceAtomicFailuresAndLegacy(t *testing.T) {
|
||||
s, in, dir := evidenceFixture(t)
|
||||
ctx := context.Background()
|
||||
seedExchange(t, s, "valid", []byte("OK"), false)
|
||||
counts := func() (findings, nodes int) {
|
||||
s.DB.QueryRow(`SELECT count(*) FROM findings WHERE task_id=$1`, in.TaskID).Scan(&findings)
|
||||
s.DB.QueryRow(`SELECT count(*) FROM exploration_nodes WHERE exploration_id=$1 AND kind='finding'`, in.ExplorationID).Scan(&nodes)
|
||||
return
|
||||
}
|
||||
for _, kind := range []string{"missing_id", "missing_body", "disk_failure", "database_failure"} {
|
||||
t.Run(kind, func(t *testing.T) {
|
||||
input := in
|
||||
refs := []db.TrafficRef{{TrafficID: "valid"}, {TrafficID: "absent"}}
|
||||
local := *s
|
||||
switch kind {
|
||||
case "missing_body":
|
||||
seedExchange(t, s, "broken", []byte("more than preview"), false)
|
||||
s.Traffic.DB().Exec(`UPDATE exchange_bodies SET resp_body=? WHERE id='broken'`, []byte("x"))
|
||||
refs = []db.TrafficRef{{TrafficID: "valid"}, {TrafficID: "broken"}}
|
||||
case "disk_failure":
|
||||
local.Dir = filepath.Join(dir, "file-not-dir")
|
||||
if err := os.WriteFile(local.Dir, []byte("x"), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
refs = refs[:1]
|
||||
case "database_failure":
|
||||
input.AssetIDs = []int64{9223372036854775807}
|
||||
refs = refs[:1]
|
||||
}
|
||||
if _, err := local.Record(ctx, input, refs); err == nil {
|
||||
t.Fatal("expected failure")
|
||||
}
|
||||
if f, n := counts(); f != 0 || n != 0 {
|
||||
t.Fatalf("partial record: findings=%d nodes=%d", f, n)
|
||||
}
|
||||
})
|
||||
}
|
||||
legacy := filepath.Join(dir, "traffic", "legacy")
|
||||
os.MkdirAll(legacy, 0o700)
|
||||
os.WriteFile(filepath.Join(legacy, "request.http"), []byte("GET / HTTP/1.1\nHost: old.local\n"), 0o600)
|
||||
os.WriteFile(filepath.Join(legacy, "response.http"), []byte("HTTP 200\n\nlegacy body"), 0o600)
|
||||
_, err := s.Traffic.DB().Exec(`INSERT INTO exchanges(id,ts,host,method,url_template,url,status,content_type,req_len,resp_len,path) VALUES('old',1,'old.local','GET','/','http://old.local/',200,'text/plain',0,11,'legacy')`)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err = s.Record(ctx, in, []db.TrafficRef{{TrafficID: "old"}}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err = s.Record(ctx, in, nil); err != nil {
|
||||
t.Fatal("legacy no-traffic report failed:", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEvidenceConcurrentBindAndDelete(t *testing.T) {
|
||||
s, in, _ := evidenceFixture(t)
|
||||
ctx := context.Background()
|
||||
seedExchange(t, s, "race", bytes.Repeat([]byte("large"), 90000), true)
|
||||
r, err := s.Record(ctx, in, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var wg sync.WaitGroup
|
||||
errs := make(chan error, 10)
|
||||
for i := 0; i < 8; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
_, err := s.Bind(ctx, r.FindingID, []db.TrafficRef{{TrafficID: "race"}})
|
||||
errs <- err
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
close(errs)
|
||||
for err := range errs {
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
list, err := s.DB.GetFindingTraffic(ctx, r.FindingID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(list.Bindings) != 1 || list.Version != 1 {
|
||||
t.Fatalf("concurrent duplicates: %+v", list)
|
||||
}
|
||||
// Race a new capture copy with source deletion: either the full copy commits
|
||||
// or nothing does. An already-bound snapshot remains readable in both cases.
|
||||
wg.Add(2)
|
||||
go func() { defer wg.Done(); s.Bind(ctx, r.FindingID, []db.TrafficRef{{TrafficID: "race"}}) }()
|
||||
go func() { defer wg.Done(); s.Traffic.DeleteHost("fixture.local") }()
|
||||
wg.Wait()
|
||||
if err = s.WithBinding(ctx, r.FindingID, list.Bindings[0].ID, func(b db.FindingTrafficBinding) error {
|
||||
f, _, err := s.OpenBody(b.Snapshot, "response")
|
||||
if f != nil {
|
||||
f.Close()
|
||||
}
|
||||
return err
|
||||
}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEvidenceGCGraceAndActiveRestore(t *testing.T) {
|
||||
s, in, _ := evidenceFixture(t)
|
||||
ctx := context.Background()
|
||||
seedExchange(t, s, "gc", []byte("unique unreferenced gc body"), false)
|
||||
f, err := s.Record(ctx, in, []db.TrafficRef{{TrafficID: "gc"}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
snap := f.Traffic.Bindings[0].Snapshot
|
||||
archive := t.TempDir()
|
||||
if err = s.CopySnapshots(ctx, []db.TrafficEvidenceSnapshot{snap}, archive); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err = s.DB.DeleteFinding(f.FindingID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
now := time.Now()
|
||||
if err = s.Collect(ctx, now); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
path, err := hashPath(s.Dir, snap.RespHash)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err = s.Collect(ctx, now.Add(23*time.Hour)); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err = os.Stat(path); err != nil {
|
||||
t.Fatal("GC grace ignored", err)
|
||||
}
|
||||
if err = s.Collect(ctx, now.Add(25*time.Hour)); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err = os.Stat(path); !os.IsNotExist(err) {
|
||||
t.Fatal("orphan was not collected", err)
|
||||
}
|
||||
entered, release, finished := make(chan struct{}), make(chan struct{}), make(chan error, 1)
|
||||
go func() {
|
||||
finished <- s.WithInstalledSnapshots(ctx, []db.TrafficEvidenceSnapshot{snap}, archive, func() error { close(entered); <-release; return nil })
|
||||
}()
|
||||
<-entered
|
||||
short, cancel := context.WithTimeout(ctx, 100*time.Millisecond)
|
||||
err = s.Collect(short, now.Add(72*time.Hour))
|
||||
cancel()
|
||||
close(release)
|
||||
if err == nil {
|
||||
t.Fatal("GC entered during active restore")
|
||||
}
|
||||
if err = <-finished; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err = os.Stat(path); err != nil {
|
||||
t.Fatal("active restore lost body", err)
|
||||
}
|
||||
// A portable package cannot substitute metadata while retaining its old ID.
|
||||
bad := snap
|
||||
bad.URL = "http://tampered.local/"
|
||||
if err = s.InstallSnapshots(ctx, []db.TrafficEvidenceSnapshot{bad}, archive); err == nil {
|
||||
t.Fatal("accepted tampered snapshot")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,35 @@
|
||||
package evidence
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"testing"
|
||||
|
||||
"github.com/Autumn-27/artex/db"
|
||||
)
|
||||
|
||||
// Initialize an explicitly configured fresh database before taking the same
|
||||
// suite lock as db, agent and server. Hold it on one pinned connection.
|
||||
func TestMain(m *testing.M) {
|
||||
if os.Getenv("ARTEX_PG_DSN") == "" {
|
||||
os.Exit(m.Run())
|
||||
}
|
||||
os.Exit(runEvidenceSuite(m))
|
||||
}
|
||||
func runEvidenceSuite(m *testing.M) int {
|
||||
pg, err := db.Open(os.Getenv("ARTEX_PG_DSN"))
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
defer pg.Close()
|
||||
conn, err := pg.Conn(context.Background())
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
defer conn.Close()
|
||||
if _, err = conn.ExecContext(context.Background(), `SELECT pg_advisory_lock(7337741002)`); err != nil {
|
||||
panic(err)
|
||||
}
|
||||
defer conn.ExecContext(context.Background(), `SELECT pg_advisory_unlock(7337741002)`)
|
||||
return m.Run()
|
||||
}
|
||||
Reference in New Issue
Block a user