diff --git a/README.md b/README.md index 39fec353..ebe5deef 100644 --- a/README.md +++ b/README.md @@ -77,21 +77,30 @@ current modes are `postgres_sql`, `local_traversal`, and `neo4j`; AGE is referen comparison mode yet. The command can emit JSONL records plus Markdown and JSON summaries, and can compare current timings against a previous JSONL baseline. -`go run ./cmd/retriever` dumps and loads live Dawgs graph databases as -manifest-based collections of compressed JSONL fragments. It supports -PostgreSQL and Neo4j, uncompressed, gzip, and zstd fragments, bounded keyset -scans, resumable dump checkpoints, checksum validation before load, optional -deterministic property scrubbing, and a read-throughput benchmark mode. It can -also package dumps as single HPKE/ML-KEM encrypted TAR archives. -See [cmd/retriever/README.md](cmd/retriever/README.md) for dump, encrypted -archive, load, scrubbed dump, metrics verification, and benchmark examples. -The same import/export functionality is available to library consumers from -`github.com/specterops/dawgs/retriever`; callers provide an already-open -`graph.Database`, and archive helpers support both path-based and stream-based -APIs. The package exposes CLI-matching default option constructors, structured -progress callbacks, manifest/metrics helpers, HPKE key envelope reader/writer -helpers, and typed errors for validation, compatibility, checksum, metrics, and -count mismatches. +`go run ./cmd/retriever` exports live Dawgs graphs to local +`ret-collection-v1` collections. JSONL with zstd is the default output; +Parquet with unshredded VARIANT properties is enabled independently, and +JSONL-only, Parquet-only, and dual-output collections are supported. Only +complete JSONL output is loadable. Parquet-only collections remain valid, +verifiable analytical exports. + +The command surface separates `dump`, `load`, `verify-collection`, +`verify-database`, `keygen`, `pack`, `unpack`, and `bench`. Archive creation is +a separate post-dump operation, and an encrypted archive must be unpacked to a +verified local collection before loading. Dump checkpoints are resumable only +while source node and relationship counts remain unchanged; same-count content +mutation is outside that guarantee. Load requires empty target graphs, is not +resumable, and may leave partial state after a write failure, so clear an +affected graph before retrying. + +The new library surface is the small operation-oriented +`github.com/specterops/dawgs/ret` facade over concrete component packages. +Legacy `github.com/specterops/dawgs/retriever` remains temporarily alongside it +for side-by-side review, but the CLI no longer uses the legacy package. Paths +are local-filesystem-only. See +[cmd/retriever/README.md](cmd/retriever/README.md) for exact dump, load, +verification, encrypted archive, scrubbing, force-replacement, and concrete +benchmark examples. PostgreSQL translates exact string property equality with a JSON string type guard and `properties ->>` extraction, so indexes created on expressions such as `properties ->> 'objectid'` and `properties ->> 'name'` can be used for selective diff --git a/cmd/retriever/README.md b/cmd/retriever/README.md index 67aba113..f03b0764 100644 --- a/cmd/retriever/README.md +++ b/cmd/retriever/README.md @@ -1,23 +1,9 @@ # retriever -`retriever` dumps and loads live Dawgs graph databases as manifest-based -collections of compact JSONL fragments. - -The v1 collection format is intended for idle databases. Dumps are deterministic -by entity ID order and cap each graph scan at the entity counts observed when -counting starts, but they do not provide a transactional cross-fragment snapshot. - -The collection manifest format is `retriever-jsonl-collection-v1`. Node and -relationship fragments contain one JSON object per line and use paths such as -`graphs/default/nodes-000001.jsonl.zst` and -`graphs/default/edges-000001.jsonl.zst`. The manifest identifies each file's -phase, record count, compressed and uncompressed sizes, and SHA-256 checksum. -The writer always terminates records with a newline; the loader also accepts a -final record without one. The loader starts with a 64 KiB scan buffer, grows it -only when needed, and rejects JSONL lines larger than 10 MiB. -New dumps default to zstd level 3. `-compression` accepts `zstd`, `gzip`, or -`none`; higher zstd levels remain available through `-zstd-level` for operators -who have measured an artifact-size benefit worth the additional workspace. +`retriever` exports Dawgs graphs to local `ret-collection-v1` collections, +loads JSONL collections, verifies artifacts or database contents, and manages +encrypted collection archives. Collection and archive paths must be on a local +filesystem; object stores and remote blob stores are not supported. ## Dump @@ -25,170 +11,78 @@ who have measured an artifact-size benefit worth the additional workspace. retriever dump \ -connection "$CONNECTION_STRING" \ -out ./dumpdir \ - -graph default \ - -scrub none \ - -compression zstd \ - -zstd-level 3 \ - -shard-size 100000 + -graph default ``` -Use repeated `-graph` flags to dump multiple named graphs. For PostgreSQL, -`-all-graphs` discovers graph names from Dawgs' `graph` metadata table and -validates that expected node and edge partitions exist. For Neo4j, `-all-graphs` -means the selected Neo4j database only. - -Existing non-empty output directories are refused unless `-force` is supplied. -The manifest is written last as `manifest.json`; if a dump fails before that -point, the directory is intentionally left for inspection without a success -manifest. - -Every database read is an ascending keyset query with a server-side -`-batch-size` limit. Retriever also rejects record B+1 at the cursor boundary, -checks every ID for strict increase, and processes records directly from the -cursor instead of retaining a result-sized slice. Dump, verify, and bench use -this same bounded scan path. - -Interrupted dumps can resume from the last atomically committed fragment by -repeating the original command with `-resume`: +JSONL is enabled by default with zstd at the JSONL package's default level +(`-jsonl-level 0`). Parquet is independent and disabled by default: ```bash retriever dump \ -connection "$CONNECTION_STRING" \ -out ./dumpdir \ - -graph default \ - -scrub none \ - -compression zstd \ - -zstd-level 3 \ - -shard-size 100000 \ - -batch-size 10000 \ - -resume + -graph first \ + -graph second \ + -jsonl=true \ + -jsonl-compression zstd \ + -jsonl-level 0 \ + -parquet ``` -The hidden checkpoint binds the driver, ordered graph targets, batch and shard -sizes, compression settings, scrub rules/configuration, and a SHA-256 identity -of the salt without storing the salt. Resume verifies every committed fragment, -rejects unexpected files, recounts the source, reconstructs compact metrics -state, and continues after the last committed source ID. `-resume` and `-force` -are mutually exclusive. The source must remain quiescent for the entire original -and resumed dump: count changes are detected, but a same-count property or -topology replacement cannot be identified without a database snapshot token. - -New dumps include a `retriever-metrics-v1` manifest section with graph metrics -computed from the same node and relationship streams written to the fragments. -The metrics include entity counts, kind histograms, degree histograms, endpoint -kind-shape histograms, and a canonical SHA-256 fingerprint. They intentionally -exclude IDs, property keys, property values, source identifiers, examples, and -sampled paths. - -Rows inserted after the initial graph count are ignored once the counted entity -total has been scanned. Deletes or other concurrent source mutations can still -make the final dumped counts differ from the initial counts, in which case dump -fails rather than writing a misleading manifest. - -Dump progress is emitted with `log/slog` on stderr. Notices mark output -directory preparation, graph counting, node and relationship phase boundaries, -periodic entity progress, checkpoint publication, manifest writing, and -completion. Phase and periodic progress includes `heap_alloc_bytes`, -`heap_inuse_bytes`, `sys_bytes`, `gc_count`, and `rss_bytes`; routine telemetry -does not force a garbage collection. Library `ProgressEvent` callbacks receive -the same memory fields, plus fragment pass and compressed/decompressed byte -counters where applicable. - -`GOMEMLIMIT` is useful defense in depth after choosing a measured batch size; -it is not a substitute for the server and client cursor bounds. - -## Runtime Profiling - -The long-running `dump`, `load`, `verify`, and `bench` commands can expose Go's -standard pprof HTTP endpoints with `-pprof-listen`. Profiling is disabled when -the flag is omitted. For example: +For a Parquet-only analytical export, disable JSONL explicitly: ```bash retriever dump \ -connection "$CONNECTION_STRING" \ - -out ./dumpdir \ + -out ./parquet-dump \ -graph default \ - -scrub full \ - -salt "$RETRIEVER_SCRUB_SALT" \ - -pprof-listen 127.0.0.1:6060 + -jsonl=false \ + -parquet=true ``` -The command writes the active endpoint to stderr after the listener starts. A -heap profile can then be captured and inspected while the command is running: +JSONL and Parquet are distinct first-class shard outputs. JSONL compression +flags do not configure Parquet; Parquet uses its own unshredded VARIANT +representation. -```bash -curl -o retriever-heap.pb.gz \ - 'http://127.0.0.1:6060/debug/pprof/heap?gc=1' -go tool pprof -http=127.0.0.1:0 retriever-heap.pb.gz -``` - -The `gc=1` query runs a garbage collection before capturing the heap, making the -profile useful for distinguishing live retained state from allocation churn. -Use `/debug/pprof/allocs` to inspect cumulative allocations and -`/debug/pprof/goroutine?debug=1` for a text goroutine dump. Captures from two -phases can be compared with `go tool pprof -base earlier.pb.gz later.pb.gz`. - -The pprof server has no authentication and accepts loopback listen addresses -only. Its dedicated handler deliberately does not register the command-line -endpoint because arguments may contain connection strings, salts, or key paths. -The listener shuts down when the Retriever command exits, and failure to bind -the requested address fails the command before database work starts. Treat all -captured profiles as sensitive operational artifacts. - -## Encrypted Archives +Repeated `-graph` values are preserved in command order. PostgreSQL +`-all-graphs` discovers graphs from Dawgs metadata; Neo4j selects its one +effective graph target. The database must remain quiescent while dumping. +Retriever snapshots node and relationship counts, rechecks them before +completion, and rejects a resume when either total has changed. This +count-stability guarantee cannot detect a same-count content replacement or +mutation because Dawgs does not expose a cross-transaction snapshot token. -Generate a recipient key pair before creating encrypted archives: - -```bash -retriever keygen \ - -private retriever-private.key \ - -public retriever-public.key -``` - -The private key is written as an unencrypted JSON key envelope with restrictive -file permissions. Store it separately from archives and treat it as sensitive. -Key generation refuses to replace existing key files. - -Pass `-archive-out` and `-recipient` to create a single encrypted TAR archive -after a dump succeeds: +Interrupted dumps can resume with the same configuration: ```bash retriever dump \ -connection "$CONNECTION_STRING" \ -out ./dumpdir \ -graph default \ - -scrub full \ - -salt "$RETRIEVER_SCRUB_SALT" \ - -archive-out ./dump.tar.pq \ - -recipient ./retriever-public.key + -resume ``` -The archive layer uses an uncompressed TAR stream encrypted with HPKE using -ML-KEM-1024, HKDF-SHA512, and AES-256-GCM. Fragment compression inside the dump -directory remains controlled by `-compression`; the archive itself is not -compressed. If archive creation fails, the command fails and removes partial -archive output while leaving the completed dump directory for inspection. +`-force` is available only for a fresh dump, is mutually exclusive with +`-resume`, and currently requires Linux or Darwin. Normal fresh and resumed +dumps remain portable. Force rejects destination or intermediate symlinks, +pins the physical parent and destination, and atomically quarantines the exact +approved directory with a handle-relative no-replace rename. It does not +enumerate or mutate anything inside the prior collection. The complete prior +collection remains alongside the new dump as a `.ret-force-*.preserved` +tombstone, and the command reports its exact path. Force therefore does not +reclaim the prior collection's disk space. Removing a preserved collection is +a separate manual action that requires a quiescent filesystem and accepts a +weaker concurrent-substitution contract than this command. Filesystem roots, +home or repository-wide targets, and their physical ancestors are rejected +before mutation. If another object appears at the original destination before +handoff, force stops before profiling, database access, or dumping; the prior +collection is restored when possible, otherwise both objects are preserved and +the error reports their absolute paths. -Unpack an encrypted archive with the private key: - -```bash -retriever unpack \ - -archive ./dump.tar.pq \ - -identity ./retriever-private.key \ - -out ./dumpdir -``` +Force replacement is CLI policy. The `/ret` library facade has no force option +and never replaces a dump destination. -Existing non-empty unpack output directories are refused unless `-force` is -supplied. Unpacked collections retain the normal `manifest.json` and fragment -layout and can be loaded with `retriever load -in ./dumpdir`. Before promoting -the staged output directory, unpack validates the manifest, expected paths, -compressed sizes and checksums, and the absence of unexpected files. Unpack does -not reopen every extracted fragment for a standalone checksum pass: it hashes -the decrypted TAR payload while writing it. Unpack does not decompress -fragments or validate JSONL records, source IDs, or record counts; that semantic -validation occurs when the collection is passed to `retriever load`. - -## Scrubbed Dumps +Full scrubbing uses the existing policy: ```bash retriever dump \ @@ -199,17 +93,12 @@ retriever dump \ -salt "$RETRIEVER_SCRUB_SALT" ``` -`-scrub full` fails closed unless a salt is supplied by `-salt` or -`RETRIEVER_SCRUB_SALT`. The legacy `RETRIEVR_SCRUB_SALT` name is accepted as a -fallback for existing scripts. Scrubbing preserves topology and source database -IDs while deterministically transforming sensitive property values. Action -counts are recorded per file, per graph, and globally in the manifest. -Pseudonyms are derived directly from the stable salt, so full scrub does not -retain a graph-wide identifier registry or perform a node observation pre-pass. -Scrub state is isolated per selected graph. - -Classifier and graph-identifier settings can be overridden with `-config`; see -`../../retriever/defaults.toml` for the supported TOML shape. +`RETRIEVR_SCRUB_SALT` remains a fallback spelling. `-config` reads the direct +scrub TOML policy shape illustrated by +[`ret/scrub/example.toml`](../../ret/scrub/example.toml): scalar policy fields +are top-level, with `[graph_rules]` and `[classifier]` sections. Salt is +runtime-only and must come from `-salt` or the environment; TOML cannot set it. +Command-line mode controls whether the policy runs. ## Load @@ -219,154 +108,141 @@ retriever load \ -in ./dumpdir ``` -Encrypted archives produced by `dump -archive-out` can be loaded directly: +Load accepts a local collection directory only. It requires complete JSONL +output, validates the JSONL representation before database writes, and does +not open or validate Parquet artifacts. Parquet-only collections are valid +exports and pass `verify-collection`, but `load` intentionally rejects them. +Targets must be empty; load never clears or replaces a graph. Load is not +resumable, and a failed write can leave a partial graph. Clear that graph +before retrying. + +Optional database verification is a distinct operation after a successful +load: ```bash retriever load \ -connection "$CONNECTION_STRING" \ - -archive ./dump.tar.pq \ - -identity ./retriever-private.key + -in ./dumpdir \ + -verify-database ``` -Load reads and validates `manifest.json`, then verifies every fragment checksum, -JSONL record, and record count before performing database writes. It asserts -destination graph schemas from the manifest metadata and loads all nodes before -relationships. Graph names from the dump are preserved; load does not support -overriding graph names. Direct archive loads decrypt and perform integrity-only -unpack validation in a temporary collection directory, then run the same -semantic preflight as directory loads before schema assertion or database -writes. - -`-batch-size` bounds both node and relationship writes. PostgreSQL uses a -correlated bulk insert and Neo4j uses grouped `UNWIND` creates; generated node -IDs are returned in stable input-ordinal order before relationship endpoints -are resolved. Retriever-produced decimal source IDs use a compact numeric map, -while arbitrary valid archive IDs retain a compatibility fallback. - -The preflight combines compressed byte counting and SHA-256 calculation with -JSONL decoding in one read of each fragment. Load then reopens and decodes the -fragments while writing to the database. This intentional second decode ensures -that every fragment is semantically valid before any database mutation. Measure -the complete preflight-plus-decode fragment path with: +## Verification + +Collection verification needs no database and validates every declared JSONL +and Parquet artifact: ```bash -go test ./retriever -run '^$' -bench '^BenchmarkLoadFragmentPath$' -benchmem +retriever verify-collection -in ./dumpdir ``` -Load refuses to write into target graphs that already contain nodes or -relationships; clear the destination graph before restoring a collection. +Database verification opens the selected backend and compares current graph +metrics with `manifest.json`: -Load progress is emitted with `log/slog` on stderr. Notices mark manifest -reading, checksum verification, schema assertion, graph boundaries, node and -relationship phase boundaries, periodic entity progress, and completion. +```bash +retriever verify-database \ + -connection "$CONNECTION_STRING" \ + -in ./dumpdir +``` + +## Encrypted archives -Pass `-verify-metrics` to scan the loaded destination graph after load and -compare it against the metrics stored in the manifest: +Archive creation is independent from dumping. Generate keys with: ```bash -retriever load \ - -connection "$CONNECTION_STRING" \ +retriever keygen \ + -private-key ./retriever-private.key \ + -public-key ./retriever-public.key +``` + +Pack a fully verified collection: + +```bash +retriever pack \ -in ./dumpdir \ - -verify-metrics + -archive ./dump.tar.enc \ + -recipient ./retriever-public.key +``` + +Unpack, authenticate, fully verify, and atomically publish a collection: + +```bash +retriever unpack \ + -archive ./dump.tar.enc \ + -out ./restored-dump \ + -identity ./retriever-private.key ``` -Metrics verification adds a full post-load node and relationship scan. +Archive publication currently requires Linux or Darwin. Pack and unpack never +replace their destinations. -## Verify +An encrypted archive is not a load input. Unpack it to a verified local +collection first, then load the unpacked directory: ```bash -retriever verify \ +retriever unpack \ + -archive ./dump.tar.enc \ + -out ./restored-dump \ + -identity ./retriever-private.key + +retriever load \ -connection "$CONNECTION_STRING" \ - -in ./dumpdir + -in ./restored-dump ``` -Verification reads expected metrics from `manifest.json`, computes actual -metrics from the destination graph database, and compares them strictly. A -successful run prints the verified graph, node, and relationship counts. A -mismatch exits non-zero and prints deterministic differences, for example: +## Runtime profiling and progress -```text -retriever: graph metrics mismatch: - graph "default" node_count: expected 884868, actual 884867 -``` +`dump`, `load`, `verify-database`, and `bench` accept +`-pprof-listen` with a loopback-only address such as `127.0.0.1:6060`. +Profiling is disabled when omitted. The dedicated server omits the command-line +endpoint because arguments may contain connection strings, salts, or key +paths. -Use standalone verification when the load already happened or when you want the -proof step to be a separate operational checkpoint. Older dumps without a -metrics section cannot be verified with this command. +Root operations emit typed events. The CLI translates them to structured +`slog` records and owns progress sampling, elapsed-rate calculation, Go runtime +memory statistics, and RSS sampling. ## Bench +`bench` measures database reads and each selected concrete artifact format: + ```bash retriever bench \ -connection "$CONNECTION_STRING" \ -graph default \ -workers 1 \ -batch-size 10000 \ - -sample-size 1000000 + -sample-size 1000000 \ + -jsonl=true \ + -jsonl-compression zstd \ + -jsonl-level 0 \ + -parquet=true ``` -Benchmark mode performs read-only scans and reports node and relationship -throughput. By default, each phase scans at most 1,000,000 nodes or -relationships so large graphs do not require a full read to produce a throughput -estimate. Pass `-sample-size 0` to scan the full graph. Use `-all-graphs` to -benchmark every graph discoverable by the selected driver, mirroring the dump -command's discovery behavior. Use `-workers` to compare concurrent batch -processing counts; database keyset scans stay sequential, while JSON -encode/compression work runs across the requested workers. The benchmark keeps -database read timing separate from optional JSON encode/compression timing: +JSONL and Parquet can be selected independently, but at least one must be +enabled. When both are selected, the report contains separate `jsonl` and +`parquet` results for each worker count. JSONL compression flags apply only to +JSONL; Parquet always uses its own unshredded VARIANT writer configuration. +`-json` emits the same report shape as JSON with a `format` field on each +result. Worker counts control benchmark-only concrete write processing; each +format phase owns its own ordered database source read. Benchmark-only worker +concurrency does not affect normal dump or load operations. -```bash -retriever bench -connection "$CONNECTION_STRING" -graph default -compression zstd -json -``` +## Go library -Benchmark progress is emitted with `log/slog` on stderr. Notices mark benchmark -start, graph counting, each worker run, node and edge phase boundaries, and -completion. Text or JSON benchmark reports remain on stdout. +New library consumers should use the small operation-oriented facade at +`github.com/specterops/dawgs/ret`, with concrete component packages such as +`ret/jsonl`, `ret/parquet`, `ret/scrub`, and `ret/archive` when their owned +types are needed. The legacy `github.com/specterops/dawgs/retriever` package +remains temporarily alongside `/ret` for review, but the CLI no longer imports +or delegates to it. -## Testing Policy +## Testing -Unit tests focus on collection format validation, compression/checksum behavior, -scrub transformations, schema planning, edge resolution, CLI flag validation, -and other pure helpers. Full database dump/load behavior is covered by -integration tests for the backend selected by `CONNECTION_STRING`. The round-trip -case crosses scan-batch and shard boundaries and verifies restored properties and -topology. `BenchmarkLoadFragmentPath` measures preflight plus the second fragment -decode used during database loading. - -### Retained-heap baseline and harness - -The incident baseline from the 1,845,833-node/44-million-edge run was 5.26 GB -forced-GC `HeapAlloc` in nodes and 16.70 GB early in edges, with 4.41 GB and -14.88 GB respectively retained by graph-wide result materialization. The last -controlled sample reached 27.31 GB `HeapAlloc` and 29.68 GB MaxRSS. The compact -metrics baseline was 401.7 MiB, the scrub registry baseline was 294.6 MiB, and -a level-11 zstd writer retained about 51 MiB versus 18.25 MiB at level 3. - -Run the retained-heap harness with one measured construction per case: +Run command unit tests with: ```bash -go test ./retriever -run '^$' \ - -bench '^BenchmarkRetainedHeap' \ - -benchmem -benchtime=1x -count=1 - -go test ./retriever -run '^$' \ - -bench '^(BenchmarkEdgeMetricsAllocation|BenchmarkEdgeScrubAllocation)$' \ - -benchmem -benchtime=100000x -count=1 +go test ./cmd/retriever ``` -The 2026-07-21 local post-change sample reported 112,591,760 retained bytes for -metrics at 1,845,833 nodes (about 73% below the 401.7 MiB baseline), 178,152 -retained bytes after pseudonymizing the same number of unique scrub identifiers, -and 0 B/op with 0 allocs/op on the warmed steady-state edge metrics path. The -hydration cases at B, 10B, and 100B remain in the harness to make raw result -materialization growth visible; `BenchmarkScanEntityBatchesFixedBatch` verifies -the production scan while holding B constant. Numbers are machine-specific, so -save benchmark output with the backend, batch, scrub, metrics, and compression -settings when evaluating the 2 GiB HeapAlloc/4 GiB RSS scale gate. - -Database validation requires separate runs for both schemes: - -```bash -CONNECTION_STRING='postgresql://...' make test_all -CONNECTION_STRING='neo4j://...' make test_all -``` +Repository-wide integration validation uses `make test_all` and requires +`CONNECTION_STRING` for the selected backend. diff --git a/cmd/retriever/bench.go b/cmd/retriever/bench.go index 8658b570..ebed879e 100644 --- a/cmd/retriever/bench.go +++ b/cmd/retriever/bench.go @@ -2,15 +2,21 @@ package main import ( "context" + "errors" "fmt" "io" "log/slog" - "sort" + "os" + "path/filepath" "sync" + "sync/atomic" "time" "github.com/specterops/dawgs/graph" - "github.com/specterops/dawgs/retriever" + "github.com/specterops/dawgs/ret/dawgs" + "github.com/specterops/dawgs/ret/entity" + "github.com/specterops/dawgs/ret/jsonl" + "github.com/specterops/dawgs/ret/parquet" ) type benchReport struct { @@ -25,6 +31,7 @@ type benchGraphReport struct { } type benchResult struct { + Format string `json:"format"` Workers int `json:"workers"` BatchSize int `json:"batch_size"` SampleSize int `json:"sample_size,omitempty"` @@ -55,7 +62,14 @@ type benchPhaseResult struct { CompressedByteSize int64 } -func Bench(ctx context.Context, db graph.Database, driverName string, targets []retriever.GraphTarget, options benchOptions) (benchReport, error) { +func benchJSONLCodec(config *jsonl.Config) string { + if config == nil { + return "" + } + return string(config.Codec) +} + +func Bench(ctx context.Context, db graph.Database, driverName string, graphNames []string, options benchOptions) (benchReport, error) { if err := options.validate(); err != nil { return benchReport{}, err } @@ -63,160 +77,189 @@ func Bench(ctx context.Context, db graph.Database, driverName string, targets [] startedAt := time.Now() slog.Info("retriever bench started", slog.String("driver", driverName), - slog.Int("graph_count", len(targets)), + slog.Int("graph_count", len(graphNames)), slog.Int("batch_size", options.BatchSize), slog.Int("sample_size", options.SampleSize), slog.Any("workers", options.Workers), - slog.String("compression", string(options.Compression)), + slog.Bool("jsonl", options.JSONL != nil), + slog.String("jsonl_codec", benchJSONLCodec(options.JSONL)), + slog.Bool("parquet", options.Parquet != nil), ) report := benchReport{ Driver: driverName, GeneratedAt: time.Now().UTC(), - Graphs: make([]benchGraphReport, 0, len(targets)), + Graphs: make([]benchGraphReport, 0, len(graphNames)), } - for targetIndex, target := range targets { + for targetIndex, graphName := range graphNames { graphStartedAt := time.Now() slog.Info("retriever bench graph started", - slog.String("graph", target.Name), + slog.String("graph", graphName), slog.Int("graph_index", targetIndex+1), - slog.Int("graph_count", len(targets)), + slog.Int("graph_count", len(graphNames)), ) - targetGraph := graph.Graph{ - Name: target.Name, - } - slog.Info("retriever bench counting graph entities", - slog.String("graph", target.Name), + slog.String("graph", graphName), ) - nodeCount, edgeCount, err := countGraphEntities(ctx, db, targetGraph) + source, err := dawgs.NewSource(db, graphName, options.BatchSize) + if err != nil { + return benchReport{}, err + } + snapshot, err := source.Snapshot(ctx) if err != nil { return benchReport{}, err } slog.Info("retriever bench graph counts ready", - slog.String("graph", target.Name), - slog.Int64("node_count", nodeCount), - slog.Int64("edge_count", edgeCount), + slog.String("graph", graphName), + slog.Int64("node_count", snapshot.NodeCount), + slog.Int64("edge_count", snapshot.RelationshipCount), ) graphReport := benchGraphReport{ - Name: target.Name, + Name: graphName, } for workerIndex, workerCount := range options.Workers { - workerStartedAt := time.Now() - - slog.Info("retriever bench worker run started", - slog.String("graph", target.Name), - slog.Int("worker_count", workerCount), - slog.Int("worker_index", workerIndex+1), - slog.Int("worker_runs", len(options.Workers)), - slog.Int("batch_size", options.BatchSize), - slog.Int("sample_size", options.SampleSize), - ) - - plannedNodes := benchPlannedCount(nodeCount, options.SampleSize) - - slog.Info("retriever bench node phase started", - slog.String("graph", target.Name), - slog.Int("worker_count", workerCount), - slog.Int64("node_count", nodeCount), - slog.Int64("planned_count", plannedNodes), - ) - - nodeResult, err := benchNodes(ctx, db, targetGraph, nodeCount, workerCount, options) - if err != nil { - return benchReport{}, err + if options.JSONL != nil { + result, err := benchJSONLRun(ctx, db, graphName, snapshot, workerCount, workerIndex, options) + if err != nil { + return benchReport{}, err + } + graphReport.Results = append(graphReport.Results, result) } - - slog.Info("retriever bench node phase completed", - slog.String("graph", target.Name), - slog.Int("worker_count", workerCount), - slog.Int64("processed", nodeResult.Count), - slog.Duration("wall_elapsed", nodeResult.WallElapsed), - slog.Duration("db_read_elapsed", nodeResult.DBReadElapsed), - slog.Duration("encode_compress_elapsed", nodeResult.EncodeCompressTime), - slog.Float64("entities_per_second", perSecond(nodeResult.Count, nodeResult.WallElapsed)), - ) - - plannedEdges := benchPlannedCount(edgeCount, options.SampleSize) - - slog.Info("retriever bench edge phase started", - slog.String("graph", target.Name), - slog.Int("worker_count", workerCount), - slog.Int64("edge_count", edgeCount), - slog.Int64("planned_count", plannedEdges), - ) - - edgeResult, err := benchEdges(ctx, db, targetGraph, edgeCount, workerCount, options) - if err != nil { - return benchReport{}, err + if options.Parquet != nil { + result, err := benchParquetRun(ctx, db, graphName, snapshot, workerCount, workerIndex, options) + if err != nil { + return benchReport{}, err + } + graphReport.Results = append(graphReport.Results, result) } - - slog.Info("retriever bench edge phase completed", - slog.String("graph", target.Name), - slog.Int("worker_count", workerCount), - slog.Int64("processed", edgeResult.Count), - slog.Duration("wall_elapsed", edgeResult.WallElapsed), - slog.Duration("db_read_elapsed", edgeResult.DBReadElapsed), - slog.Duration("encode_compress_elapsed", edgeResult.EncodeCompressTime), - slog.Float64("entities_per_second", perSecond(edgeResult.Count, edgeResult.WallElapsed)), - ) - - totalWall := nodeResult.WallElapsed + edgeResult.WallElapsed - graphReport.Results = append(graphReport.Results, benchResult{ - Workers: workerCount, - BatchSize: options.BatchSize, - SampleSize: options.SampleSize, - NodeCount: nodeCount, - EdgeCount: edgeCount, - NodeProcessed: nodeResult.Count, - EdgeProcessed: edgeResult.Count, - NodeWallMillis: nodeResult.WallElapsed.Milliseconds(), - EdgeWallMillis: edgeResult.WallElapsed.Milliseconds(), - NodeDBReadMillis: nodeResult.DBReadElapsed.Milliseconds(), - EdgeDBReadMillis: edgeResult.DBReadElapsed.Milliseconds(), - NodeEncodeCompressMillis: nodeResult.EncodeCompressTime.Milliseconds(), - EdgeEncodeCompressMillis: edgeResult.EncodeCompressTime.Milliseconds(), - TotalWallMillis: totalWall.Milliseconds(), - NodesPerSecond: perSecond(nodeResult.Count, nodeResult.WallElapsed), - EdgesPerSecond: perSecond(edgeResult.Count, edgeResult.WallElapsed), - EntitiesPerSecond: perSecond(nodeResult.Count+edgeResult.Count, totalWall), - UncompressedBytes: nodeResult.UncompressedByteSize + edgeResult.UncompressedByteSize, - CompressedBytes: nodeResult.CompressedByteSize + edgeResult.CompressedByteSize, - }) - - slog.Info("retriever bench worker run completed", - slog.String("graph", target.Name), - slog.Int("worker_count", workerCount), - slog.Duration("wall_elapsed", time.Since(workerStartedAt)), - slog.Float64("entities_per_second", perSecond(nodeResult.Count+edgeResult.Count, totalWall)), - ) } report.Graphs = append(report.Graphs, graphReport) slog.Info("retriever bench graph completed", - slog.String("graph", target.Name), + slog.String("graph", graphName), slog.Duration("wall_elapsed", time.Since(graphStartedAt)), ) } slog.Info("retriever bench completed", slog.String("driver", driverName), - slog.Int("graph_count", len(targets)), + slog.Int("graph_count", len(graphNames)), slog.Duration("wall_elapsed", time.Since(startedAt)), ) return report, nil } +func benchJSONLRun(ctx context.Context, db graph.Database, graphName string, snapshot dawgs.Snapshot, workers, workerIndex int, options benchOptions) (benchResult, error) { + runStartedAt := time.Now() + logBenchRunStarted(graphName, "jsonl", workers, workerIndex, options) + nodeResult, err := benchNodes( + ctx, db, graphName, "jsonl", snapshot.NodeCount, workers, options, + func(path string, nodes []entity.Node) (benchPhaseResult, error) { + return benchJSONLNodeBatch(path, nodes, *options.JSONL) + }, + ) + if err != nil { + return benchResult{}, err + } + relationshipResult, err := benchRelationships( + ctx, db, graphName, "jsonl", snapshot.RelationshipCount, workers, options, + func(path string, relationships []entity.Relationship) (benchPhaseResult, error) { + return benchJSONLRelationshipBatch(path, relationships, *options.JSONL) + }, + ) + if err != nil { + return benchResult{}, err + } + result := newBenchResult("jsonl", snapshot, workers, options, nodeResult, relationshipResult) + logBenchRunCompleted(graphName, result, time.Since(runStartedAt)) + return result, nil +} + +func benchParquetRun(ctx context.Context, db graph.Database, graphName string, snapshot dawgs.Snapshot, workers, workerIndex int, options benchOptions) (benchResult, error) { + runStartedAt := time.Now() + logBenchRunStarted(graphName, "parquet", workers, workerIndex, options) + nodeResult, err := benchNodes( + ctx, db, graphName, "parquet", snapshot.NodeCount, workers, options, + func(path string, nodes []entity.Node) (benchPhaseResult, error) { + return benchParquetNodeBatch(path, nodes, *options.Parquet) + }, + ) + if err != nil { + return benchResult{}, err + } + relationshipResult, err := benchRelationships( + ctx, db, graphName, "parquet", snapshot.RelationshipCount, workers, options, + func(path string, relationships []entity.Relationship) (benchPhaseResult, error) { + return benchParquetRelationshipBatch(path, relationships, *options.Parquet) + }, + ) + if err != nil { + return benchResult{}, err + } + result := newBenchResult("parquet", snapshot, workers, options, nodeResult, relationshipResult) + logBenchRunCompleted(graphName, result, time.Since(runStartedAt)) + return result, nil +} + +func logBenchRunStarted(graphName, format string, workers, workerIndex int, options benchOptions) { + slog.Info("retriever bench worker run started", + slog.String("graph", graphName), + slog.String("format", format), + slog.Int("worker_count", workers), + slog.Int("worker_index", workerIndex+1), + slog.Int("worker_runs", len(options.Workers)), + slog.Int("batch_size", options.BatchSize), + slog.Int("sample_size", options.SampleSize), + ) +} + +func newBenchResult(format string, snapshot dawgs.Snapshot, workers int, options benchOptions, nodeResult, relationshipResult benchPhaseResult) benchResult { + totalWall := nodeResult.WallElapsed + relationshipResult.WallElapsed + return benchResult{ + Format: format, + Workers: workers, + BatchSize: options.BatchSize, + SampleSize: options.SampleSize, + NodeCount: snapshot.NodeCount, + EdgeCount: snapshot.RelationshipCount, + NodeProcessed: nodeResult.Count, + EdgeProcessed: relationshipResult.Count, + NodeWallMillis: nodeResult.WallElapsed.Milliseconds(), + EdgeWallMillis: relationshipResult.WallElapsed.Milliseconds(), + NodeDBReadMillis: nodeResult.DBReadElapsed.Milliseconds(), + EdgeDBReadMillis: relationshipResult.DBReadElapsed.Milliseconds(), + NodeEncodeCompressMillis: nodeResult.EncodeCompressTime.Milliseconds(), + EdgeEncodeCompressMillis: relationshipResult.EncodeCompressTime.Milliseconds(), + TotalWallMillis: totalWall.Milliseconds(), + NodesPerSecond: perSecond(nodeResult.Count, nodeResult.WallElapsed), + EdgesPerSecond: perSecond(relationshipResult.Count, relationshipResult.WallElapsed), + EntitiesPerSecond: perSecond(nodeResult.Count+relationshipResult.Count, totalWall), + UncompressedBytes: nodeResult.UncompressedByteSize + relationshipResult.UncompressedByteSize, + CompressedBytes: nodeResult.CompressedByteSize + relationshipResult.CompressedByteSize, + } +} + +func logBenchRunCompleted(graphName string, result benchResult, elapsed time.Duration) { + slog.Info("retriever bench worker run completed", + slog.String("graph", graphName), + slog.String("format", result.Format), + slog.Int("worker_count", result.Workers), + slog.Duration("wall_elapsed", elapsed), + slog.Float64("entities_per_second", result.EntitiesPerSecond), + ) +} + type benchBatchProcessor[T any] struct { + parent context.Context ctx context.Context cancel context.CancelFunc process func([]T) (benchPhaseResult, error) @@ -239,6 +282,7 @@ func newBenchBatchProcessor[T any](ctx context.Context, workers int, process fun var ( scanCtx, cancel = context.WithCancel(ctx) processor = &benchBatchProcessor[T]{ + parent: ctx, ctx: scanCtx, cancel: cancel, process: process, @@ -322,6 +366,9 @@ func (s *benchBatchProcessor[T]) closeAndWait() (benchPhaseResult, error) { if err := s.currentError(); err != nil { return benchPhaseResult{}, err } + if err := s.parent.Err(); err != nil { + return benchPhaseResult{}, err + } return s.snapshot(), nil } @@ -364,157 +411,347 @@ func (s *benchBatchProcessor[T]) snapshot() benchPhaseResult { return s.result } -func benchNodes(ctx context.Context, db graph.Database, targetGraph graph.Graph, total int64, workers int, options benchOptions) (benchPhaseResult, error) { - startedAt := time.Now() - planned := benchPlannedCount(total, options.SampleSize) +type benchArtifactFilesystem struct { + mkdirTemp func(string, string) (string, error) + removeAll func(string) error +} - processor, scanCtx, err := newBenchBatchProcessor(ctx, workers, func(nodes []*graph.Node) (benchPhaseResult, error) { - return benchNodeBatch(nodes, options) +func (s benchArtifactFilesystem) withDefaults() benchArtifactFilesystem { + if s.mkdirTemp == nil { + s.mkdirTemp = os.MkdirTemp + } + if s.removeAll == nil { + s.removeAll = os.RemoveAll + } + return s +} + +func benchNodes( + ctx context.Context, + db graph.Database, + graphName string, + format string, + total int64, + workers int, + options benchOptions, + write func(string, []entity.Node) (benchPhaseResult, error), +) (benchPhaseResult, error) { + return benchNodesWithFilesystem(ctx, db, graphName, format, total, workers, options, write, benchArtifactFilesystem{}) +} + +func benchNodesWithFilesystem( + ctx context.Context, + db graph.Database, + graphName string, + format string, + total int64, + workers int, + options benchOptions, + write func(string, []entity.Node) (benchPhaseResult, error), + filesystem benchArtifactFilesystem, +) (phaseResult benchPhaseResult, resultErr error) { + filesystem = filesystem.withDefaults() + tempDir, err := filesystem.mkdirTemp("", "retriever-bench-"+format+"-nodes-") + if err != nil { + return benchPhaseResult{}, fmt.Errorf("create %s node benchmark directory: %w", format, err) + } + defer func() { + if err := filesystem.removeAll(tempDir); err != nil { + resultErr = errors.Join(resultErr, fmt.Errorf("remove %s node benchmark directory %q: %w", format, tempDir, err)) + } + }() + + source, err := dawgs.NewSource(db, graphName, options.BatchSize) + if err != nil { + return benchPhaseResult{}, err + } + var batchNumber atomic.Int64 + processor, scanCtx, err := newBenchBatchProcessor(ctx, workers, func(nodes []entity.Node) (benchPhaseResult, error) { + path := filepath.Join(tempDir, fmt.Sprintf("worker-batch-%06d", batchNumber.Add(1))) + return write(path, nodes) }) if err != nil { return benchPhaseResult{}, err } - var ( - nodes = make([]*graph.Node, 0, options.BatchSize) - nextProgressAt = retrieverInitialProgressAt(planned) + startedAt := time.Now() + planned := benchPlannedCount(total, options.SampleSize) + processed := int64(0) + nextProgressAt := retrieverInitialProgressAt(planned) + slog.Info("retriever bench node phase started", + slog.String("graph", graphName), + slog.String("format", format), + slog.Int("worker_count", workers), + slog.Int64("node_count", total), + slog.Int64("planned_count", planned), ) - _, scanErr := retriever.ScanDatabaseNodes(scanCtx, db, targetGraph, planned, options.BatchSize, func(node *graph.Node) error { - nodes = append(nodes, node) - return nil - }, func(event retriever.ScanBatchEvent) error { - processor.addDBReadElapsed(event.ReadElapsed) - if err := processor.handle(nodes); err != nil { - return err - } - nodes = make([]*graph.Node, 0, options.BatchSize) - - progressResult := processor.snapshot() - progressResult.Count = event.Processed - nextProgressAt = logBenchPhaseProgress(targetGraph.Name, retriever.PhaseNodes, workers, progressResult, planned, startedAt, nextProgressAt) - return nil - }) + for processed < planned { + readStartedAt := time.Now() + batch, readErr := source.NextNodes(scanCtx) + processor.addDBReadElapsed(time.Since(readStartedAt)) + if readErr != nil { + _, closeErr := processor.closeAndWait() + return benchPhaseResult{}, errors.Join(readErr, closeErr) + } + if len(batch.Entities) == 0 { + _, closeErr := processor.closeAndWait() + return benchPhaseResult{}, errors.Join( + fmt.Errorf("node benchmark scan ended after %d of %d entities", processed, planned), + closeErr, + ) + } - result, processErr := processor.closeAndWait() - result.WallElapsed = time.Since(startedAt) + remaining := planned - processed + if int64(len(batch.Entities)) > remaining { + batch.Entities = batch.Entities[:remaining] + } + if err := processor.handle(batch.Entities); err != nil { + _, closeErr := processor.closeAndWait() + return benchPhaseResult{}, errors.Join(err, closeErr) + } + processed += int64(len(batch.Entities)) - if processErr != nil { - return benchPhaseResult{}, processErr + progress := processor.snapshot() + progress.Count = processed + nextProgressAt = logBenchPhaseProgress(graphName, "nodes", workers, progress, planned, startedAt, nextProgressAt) } - if scanErr != nil { - return benchPhaseResult{}, scanErr + phaseResult, err = processor.closeAndWait() + phaseResult.WallElapsed = time.Since(startedAt) + if err != nil { + return benchPhaseResult{}, err } + if phaseResult.Count != planned { + return benchPhaseResult{}, fmt.Errorf("node benchmark wrote %d of %d planned entities", phaseResult.Count, planned) + } + slog.Info("retriever bench node phase completed", + slog.String("graph", graphName), + slog.String("format", format), + slog.Int("worker_count", workers), + slog.Int64("processed", phaseResult.Count), + slog.Duration("wall_elapsed", phaseResult.WallElapsed), + slog.Duration("db_read_elapsed", phaseResult.DBReadElapsed), + slog.Duration("encode_compress_elapsed", phaseResult.EncodeCompressTime), + slog.Float64("entities_per_second", perSecond(phaseResult.Count, phaseResult.WallElapsed)), + ) + return phaseResult, nil +} - return result, nil +func benchRelationships( + ctx context.Context, + db graph.Database, + graphName string, + format string, + total int64, + workers int, + options benchOptions, + write func(string, []entity.Relationship) (benchPhaseResult, error), +) (benchPhaseResult, error) { + return benchRelationshipsWithFilesystem(ctx, db, graphName, format, total, workers, options, write, benchArtifactFilesystem{}) } -func benchEdges(ctx context.Context, db graph.Database, targetGraph graph.Graph, total int64, workers int, options benchOptions) (benchPhaseResult, error) { - startedAt := time.Now() - planned := benchPlannedCount(total, options.SampleSize) +func benchRelationshipsWithFilesystem( + ctx context.Context, + db graph.Database, + graphName string, + format string, + total int64, + workers int, + options benchOptions, + write func(string, []entity.Relationship) (benchPhaseResult, error), + filesystem benchArtifactFilesystem, +) (phaseResult benchPhaseResult, resultErr error) { + filesystem = filesystem.withDefaults() + tempDir, err := filesystem.mkdirTemp("", "retriever-bench-"+format+"-relationships-") + if err != nil { + return benchPhaseResult{}, fmt.Errorf("create %s relationship benchmark directory: %w", format, err) + } + defer func() { + if err := filesystem.removeAll(tempDir); err != nil { + resultErr = errors.Join(resultErr, fmt.Errorf("remove %s relationship benchmark directory %q: %w", format, tempDir, err)) + } + }() - processor, scanCtx, err := newBenchBatchProcessor(ctx, workers, func(relationships []*graph.Relationship) (benchPhaseResult, error) { - return benchRelationshipBatch(relationships, options) + source, err := dawgs.NewSource(db, graphName, options.BatchSize) + if err != nil { + return benchPhaseResult{}, err + } + var batchNumber atomic.Int64 + processor, scanCtx, err := newBenchBatchProcessor(ctx, workers, func(relationships []entity.Relationship) (benchPhaseResult, error) { + path := filepath.Join(tempDir, fmt.Sprintf("worker-batch-%06d", batchNumber.Add(1))) + return write(path, relationships) }) if err != nil { return benchPhaseResult{}, err } - var ( - relationships = make([]*graph.Relationship, 0, options.BatchSize) - nextProgressAt = retrieverInitialProgressAt(planned) + startedAt := time.Now() + planned := benchPlannedCount(total, options.SampleSize) + processed := int64(0) + nextProgressAt := retrieverInitialProgressAt(planned) + slog.Info("retriever bench relationship phase started", + slog.String("graph", graphName), + slog.String("format", format), + slog.Int("worker_count", workers), + slog.Int64("relationship_count", total), + slog.Int64("planned_count", planned), ) - _, scanErr := retriever.ScanDatabaseRelationships(scanCtx, db, targetGraph, planned, options.BatchSize, func(relationship *graph.Relationship) error { - relationships = append(relationships, relationship) - return nil - }, func(event retriever.ScanBatchEvent) error { - processor.addDBReadElapsed(event.ReadElapsed) - if err := processor.handle(relationships); err != nil { - return err - } - relationships = make([]*graph.Relationship, 0, options.BatchSize) - - progressResult := processor.snapshot() - progressResult.Count = event.Processed - nextProgressAt = logBenchPhaseProgress(targetGraph.Name, retriever.PhaseEdges, workers, progressResult, planned, startedAt, nextProgressAt) - return nil - }) + for processed < planned { + readStartedAt := time.Now() + batch, readErr := source.NextRelationships(scanCtx) + processor.addDBReadElapsed(time.Since(readStartedAt)) + if readErr != nil { + _, closeErr := processor.closeAndWait() + return benchPhaseResult{}, errors.Join(readErr, closeErr) + } + if len(batch.Entities) == 0 { + _, closeErr := processor.closeAndWait() + return benchPhaseResult{}, errors.Join( + fmt.Errorf("relationship benchmark scan ended after %d of %d entities", processed, planned), + closeErr, + ) + } - result, processErr := processor.closeAndWait() - result.WallElapsed = time.Since(startedAt) + remaining := planned - processed + if int64(len(batch.Entities)) > remaining { + batch.Entities = batch.Entities[:remaining] + } + if err := processor.handle(batch.Entities); err != nil { + _, closeErr := processor.closeAndWait() + return benchPhaseResult{}, errors.Join(err, closeErr) + } + processed += int64(len(batch.Entities)) - if processErr != nil { - return benchPhaseResult{}, processErr + progress := processor.snapshot() + progress.Count = processed + nextProgressAt = logBenchPhaseProgress(graphName, "relationships", workers, progress, planned, startedAt, nextProgressAt) } - if scanErr != nil { - return benchPhaseResult{}, scanErr + phaseResult, err = processor.closeAndWait() + phaseResult.WallElapsed = time.Since(startedAt) + if err != nil { + return benchPhaseResult{}, err } - - return result, nil + if phaseResult.Count != planned { + return benchPhaseResult{}, fmt.Errorf("relationship benchmark wrote %d of %d planned entities", phaseResult.Count, planned) + } + slog.Info("retriever bench relationship phase completed", + slog.String("graph", graphName), + slog.String("format", format), + slog.Int("worker_count", workers), + slog.Int64("processed", phaseResult.Count), + slog.Duration("wall_elapsed", phaseResult.WallElapsed), + slog.Duration("db_read_elapsed", phaseResult.DBReadElapsed), + slog.Duration("encode_compress_elapsed", phaseResult.EncodeCompressTime), + slog.Float64("entities_per_second", perSecond(phaseResult.Count, phaseResult.WallElapsed)), + ) + return phaseResult, nil } -func benchNodeBatch(nodes []*graph.Node, options benchOptions) (benchPhaseResult, error) { - return benchCompressedBatch(len(nodes), options, func() []retriever.FragmentNode { - items := make([]retriever.FragmentNode, 0, len(nodes)) - for _, node := range nodes { - kinds := node.Kinds.Strings() - sort.Strings(kinds) - - items = append(items, retriever.FragmentNode{ - ID: node.ID.String(), - Kinds: kinds, - Properties: node.Properties.MapOrEmpty(), - }) - } - - return items +func benchJSONLNodeBatch(path string, nodes []entity.Node, config jsonl.Config) (benchPhaseResult, error) { + startedAt := time.Now() + artifact, err := benchWriteArtifactFile(path, nodes, func(output io.Writer) (benchArtifactWriter[entity.Node, jsonl.Artifact], error) { + writer, err := jsonl.NewNodeWriter(output, config) + return &writer, err }) + elapsed := time.Since(startedAt) + if err != nil { + return benchPhaseResult{}, err + } + return benchPhaseResult{ + Count: artifact.Count, + EncodeCompressTime: elapsed, + UncompressedByteSize: artifact.UncompressedBytes, + CompressedByteSize: artifact.StoredBytes, + }, nil } -func benchRelationshipBatch(relationships []*graph.Relationship, options benchOptions) (benchPhaseResult, error) { - return benchCompressedBatch(len(relationships), options, func() []retriever.FragmentEdge { - items := make([]retriever.FragmentEdge, 0, len(relationships)) - for _, relationship := range relationships { - kind := "" - if relationship.Kind != nil { - kind = relationship.Kind.String() - } - - items = append(items, retriever.FragmentEdge{ - StartID: relationship.StartID.String(), - EndID: relationship.EndID.String(), - Kind: kind, - Properties: relationship.Properties.MapOrEmpty(), - }) - } - - return items +func benchJSONLRelationshipBatch(path string, relationships []entity.Relationship, config jsonl.Config) (benchPhaseResult, error) { + startedAt := time.Now() + artifact, err := benchWriteArtifactFile(path, relationships, func(output io.Writer) (benchArtifactWriter[entity.Relationship, jsonl.Artifact], error) { + writer, err := jsonl.NewRelationshipWriter(output, config) + return &writer, err }) + elapsed := time.Since(startedAt) + if err != nil { + return benchPhaseResult{}, err + } + return benchPhaseResult{ + Count: artifact.Count, + EncodeCompressTime: elapsed, + UncompressedByteSize: artifact.UncompressedBytes, + CompressedByteSize: artifact.StoredBytes, + }, nil } -func benchCompressedBatch[T any](count int, options benchOptions, buildRecords func() []T) (benchPhaseResult, error) { - result := benchPhaseResult{ - Count: int64(count), - } - if options.Compression == retriever.CompressionDisabled || count == 0 { - return result, nil +func benchParquetNodeBatch(path string, nodes []entity.Node, config parquet.Config) (benchPhaseResult, error) { + startedAt := time.Now() + artifact, err := benchWriteArtifactFile(path, nodes, func(output io.Writer) (benchArtifactWriter[entity.Node, parquet.Artifact], error) { + writer, err := parquet.NewNodeWriter(output, config) + return &writer, err + }) + elapsed := time.Since(startedAt) + if err != nil { + return benchPhaseResult{}, err } + return benchPhaseResult{ + Count: artifact.Count, + EncodeCompressTime: elapsed, + CompressedByteSize: artifact.StoredBytes, + }, nil +} - encodeStarted := time.Now() - uncompressedBytes, compressedBytes, err := retriever.CompressedJSONLinesSize(options.Compression, options.ZstdLevel, buildRecords()) - result.EncodeCompressTime = time.Since(encodeStarted) - +func benchParquetRelationshipBatch(path string, relationships []entity.Relationship, config parquet.Config) (benchPhaseResult, error) { + startedAt := time.Now() + artifact, err := benchWriteArtifactFile(path, relationships, func(output io.Writer) (benchArtifactWriter[entity.Relationship, parquet.Artifact], error) { + writer, err := parquet.NewRelationshipWriter(output, config) + return &writer, err + }) + elapsed := time.Since(startedAt) if err != nil { return benchPhaseResult{}, err } + return benchPhaseResult{ + Count: artifact.Count, + EncodeCompressTime: elapsed, + CompressedByteSize: artifact.StoredBytes, + }, nil +} - result.UncompressedByteSize = uncompressedBytes - result.CompressedByteSize = compressedBytes +type benchArtifactWriter[E, A any] interface { + Push([]E) error + Close() error + Result() (A, error) +} - return result, nil +func benchWriteArtifactFile[E, A any]( + path string, + values []E, + newWriter func(io.Writer) (benchArtifactWriter[E, A], error), +) (A, error) { + var zero A + file, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0o600) + if err != nil { + return zero, fmt.Errorf("open benchmark artifact: %w", err) + } + writer, err := newWriter(file) + if err != nil { + return zero, errors.Join(err, file.Close()) + } + pushErr := writer.Push(values) + closeWriterErr := writer.Close() + var artifact A + var resultErr error + if pushErr == nil && closeWriterErr == nil { + artifact, resultErr = writer.Result() + } + closeFileErr := file.Close() + if err := errors.Join(pushErr, closeWriterErr, resultErr, closeFileErr); err != nil { + return zero, err + } + return artifact, nil } func benchPlannedCount(total int64, sampleSize int) int64 { @@ -534,14 +771,14 @@ func benchPlannedCount(total int64, sampleSize int) int64 { return sampleCount } -func logBenchPhaseProgress(graphName string, phaseName retriever.Phase, workers int, result benchPhaseResult, planned int64, startedAt time.Time, nextProgressAt int64) int64 { +func logBenchPhaseProgress(graphName string, phaseName string, workers int, result benchPhaseResult, planned int64, startedAt time.Time, nextProgressAt int64) int64 { if nextProgressAt == 0 || result.Count < nextProgressAt || result.Count >= planned { return nextProgressAt } slog.Info("retriever bench phase progress", slog.String("graph", graphName), - slog.String("phase", string(phaseName)), + slog.String("phase", phaseName), slog.Int("worker_count", workers), slog.Int64("processed", result.Count), slog.Int64("planned_count", planned), @@ -561,7 +798,8 @@ func writeBenchReport(writer io.Writer, report benchReport) { for _, result := range graphReport.Results { fmt.Fprintf( writer, - " workers=%d batch=%d sample_size=%d nodes=%d/%d edges=%d/%d total_ms=%d entities_per_sec=%.2f db_read_ms=%d encode_compress_ms=%d\n", + " format=%s workers=%d batch=%d sample_size=%d nodes=%d/%d edges=%d/%d total_ms=%d entities_per_sec=%.2f db_read_ms=%d encode_compress_ms=%d\n", + result.Format, result.Workers, result.BatchSize, result.SampleSize, diff --git a/cmd/retriever/bench_test.go b/cmd/retriever/bench_test.go index e902d3ba..6f0ab56c 100644 --- a/cmd/retriever/bench_test.go +++ b/cmd/retriever/bench_test.go @@ -2,13 +2,24 @@ package main import ( "bytes" + "compress/gzip" "context" + "encoding/json" + "errors" + "os" + "path/filepath" + "reflect" + "sort" "strings" + "sync" "testing" "time" + cypherModel "github.com/specterops/dawgs/cypher/models/cypher" "github.com/specterops/dawgs/graph" - "github.com/specterops/dawgs/retriever" + "github.com/specterops/dawgs/ret/entity" + "github.com/specterops/dawgs/ret/jsonl" + "github.com/specterops/dawgs/ret/parquet" ) func TestBenchSamplingHelpers(t *testing.T) { @@ -51,6 +62,7 @@ func TestBenchFormattingHelpers(t *testing.T) { Graphs: []benchGraphReport{{ Name: "default", Results: []benchResult{{ + Format: "parquet", Workers: 2, BatchSize: 100, SampleSize: 2, @@ -66,7 +78,7 @@ func TestBenchFormattingHelpers(t *testing.T) { }}, }) output := buffer.String() - for _, expected := range []string{"graph: default", "workers=2", "sample_size=2", "nodes=2/3", "edges=2/4", "entities_per_sec=140.00", "db_read_ms=30"} { + for _, expected := range []string{"graph: default", "format=parquet", "workers=2", "sample_size=2", "nodes=2/3", "edges=2/4", "entities_per_sec=140.00", "db_read_ms=30"} { if !strings.Contains(output, expected) { t.Fatalf("bench report missing %q in %q", expected, output) } @@ -78,27 +90,27 @@ func TestLogBenchPhaseProgressThresholds(t *testing.T) { nextProgressAt := retrieverProgressEntityInterval startedAt := time.Now().Add(-time.Second) - if got := logBenchPhaseProgress("default", retriever.PhaseNodes, 1, benchPhaseResult{ + if got := logBenchPhaseProgress("default", "nodes", 1, benchPhaseResult{ Count: nextProgressAt - 1, }, planned, startedAt, nextProgressAt); got != nextProgressAt { t.Fatalf("progress before threshold advanced to %d", got) } - if got := logBenchPhaseProgress("default", retriever.PhaseNodes, 1, benchPhaseResult{ + if got := logBenchPhaseProgress("default", "nodes", 1, benchPhaseResult{ Count: nextProgressAt, }, planned, startedAt, nextProgressAt); got != nextProgressAt*2 { t.Fatalf("progress at threshold advanced to %d", got) } - if got := logBenchPhaseProgress("default", retriever.PhaseNodes, 1, benchPhaseResult{ + if got := logBenchPhaseProgress("default", "nodes", 1, benchPhaseResult{ Count: nextProgressAt*2 + 1, }, planned*2, startedAt, nextProgressAt); got != nextProgressAt*3 { t.Fatalf("progress after large jump advanced to %d", got) } - if got := logBenchPhaseProgress("default", retriever.PhaseNodes, 1, benchPhaseResult{ + if got := logBenchPhaseProgress("default", "nodes", 1, benchPhaseResult{ Count: planned, }, planned, startedAt, nextProgressAt); got != nextProgressAt { t.Fatalf("completed progress advanced to %d", got) } - if got := logBenchPhaseProgress("default", retriever.PhaseNodes, 1, benchPhaseResult{ + if got := logBenchPhaseProgress("default", "nodes", 1, benchPhaseResult{ Count: nextProgressAt, }, planned, startedAt, 0); got != 0 { t.Fatalf("disabled progress advanced to %d", got) @@ -143,44 +155,466 @@ func TestBenchBatchProcessorAggregatesConcurrentResults(t *testing.T) { } } -func TestBenchBatchCompression(t *testing.T) { - options := benchOptions{ - Compression: retriever.CompressionGzip, - ZstdLevel: retriever.DefaultZstdLevel, +func TestBenchBatchProcessorReturnsParentCancellationWithQueuedJobs(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + started := make(chan struct{}, 2) + release := make(chan struct{}) + processor, _, err := newBenchBatchProcessor(ctx, 2, func(values []int) (benchPhaseResult, error) { + started <- struct{}{} + <-release + return benchPhaseResult{Count: int64(len(values))}, nil + }) + if err != nil { + t.Fatalf("create bench batch processor: %v", err) + } + + if err := processor.handle([]int{1}); err != nil { + t.Fatalf("handle first active batch: %v", err) + } + if err := processor.handle([]int{2}); err != nil { + t.Fatalf("handle second active batch: %v", err) + } + <-started + <-started + if err := processor.handle([]int{3}); err != nil { + t.Fatalf("queue third batch: %v", err) } - nodes := []*graph.Node{ - graph.NewNode(2, graph.AsProperties(map[string]any{"name": "alice"}), graph.StringKind("User"), graph.StringKind("Admin")), - graph.NewNode(1, nil, graph.StringKind("Computer")), + if err := processor.handle([]int{4}); err != nil { + t.Fatalf("queue fourth batch: %v", err) } - nodeResult, err := benchNodeBatch(nodes, options) + cancel() + close(release) + if _, err := processor.closeAndWait(); !errors.Is(err, context.Canceled) { + t.Fatalf("wait error = %v, want caller cancellation", err) + } +} + +func TestBenchBatchProcessorPrefersWriterErrorOverParentCancellation(t *testing.T) { + writerFailure := errors.New("writer failed") + ctx, cancel := context.WithCancel(context.Background()) + started := make(chan struct{}) + release := make(chan struct{}) + processor, _, err := newBenchBatchProcessor(ctx, 2, func([]int) (benchPhaseResult, error) { + close(started) + <-release + return benchPhaseResult{}, writerFailure + }) if err != nil { - t.Fatalf("bench node batch: %v", err) + t.Fatalf("create bench batch processor: %v", err) } - if nodeResult.Count != 2 || nodeResult.UncompressedByteSize <= 0 || nodeResult.CompressedByteSize <= 0 { - t.Fatalf("unexpected node batch result: %+v", nodeResult) + if err := processor.handle([]int{1}); err != nil { + t.Fatalf("handle active batch: %v", err) } + <-started + cancel() + close(release) - noCompressionResult, err := benchNodeBatch(nodes, benchOptions{ - Compression: retriever.CompressionDisabled, - ZstdLevel: retriever.DefaultZstdLevel, - }) + if _, err := processor.closeAndWait(); !errors.Is(err, writerFailure) { + t.Fatalf("wait error = %v, want writer failure", err) + } +} + +func TestBenchSourcePhasesPreserveSamplePrefixAndAggregateWorkers(t *testing.T) { + database := &benchSourceDatabase{ + nodes: []*graph.Node{ + graph.NewNode(1, graph.AsProperties(map[string]any{"name": "one"}), graph.StringKind("User"), graph.StringKind("Admin"), graph.StringKind("User")), + graph.NewNode(2, graph.AsProperties(map[string]any{"name": "two"}), graph.StringKind("Computer")), + graph.NewNode(3, graph.AsProperties(map[string]any{"name": "three"}), graph.StringKind("Group")), + graph.NewNode(4, graph.AsProperties(map[string]any{"name": "four"}), graph.StringKind("Role")), + }, + relationships: []*graph.Relationship{ + graph.NewRelationship(10, 1, 2, graph.AsProperties(map[string]any{"position": 1}), graph.StringKind("First")), + graph.NewRelationship(11, 2, 3, graph.AsProperties(map[string]any{"position": 2}), graph.StringKind("Second")), + graph.NewRelationship(12, 3, 4, graph.AsProperties(map[string]any{"position": 3}), graph.StringKind("Third")), + graph.NewRelationship(13, 4, 1, graph.AsProperties(map[string]any{"position": 4}), graph.StringKind("Fourth")), + }, + } + options := benchOptions{BatchSize: 2, SampleSize: 3} + + var ( + mu sync.Mutex + nodes []entity.Node + nodeBatchSizes []int + relationships []entity.Relationship + relationshipBatchSizes []int + ) + nodeResult, err := benchNodes( + context.Background(), database, "example", "jsonl", int64(len(database.nodes)), 2, options, + func(_ string, batch []entity.Node) (benchPhaseResult, error) { + mu.Lock() + nodes = append(nodes, batch...) + nodeBatchSizes = append(nodeBatchSizes, len(batch)) + mu.Unlock() + return benchPhaseResult{Count: int64(len(batch))}, nil + }, + ) + if err != nil { + t.Fatalf("bench nodes: %v", err) + } + relationshipResult, err := benchRelationships( + context.Background(), database, "example", "parquet", int64(len(database.relationships)), 2, options, + func(_ string, batch []entity.Relationship) (benchPhaseResult, error) { + mu.Lock() + relationships = append(relationships, batch...) + relationshipBatchSizes = append(relationshipBatchSizes, len(batch)) + mu.Unlock() + return benchPhaseResult{Count: int64(len(batch))}, nil + }, + ) if err != nil { - t.Fatalf("bench uncompressed node batch: %v", err) + t.Fatalf("bench relationships: %v", err) + } + + if nodeResult.Count != 3 || relationshipResult.Count != 3 { + t.Fatalf("phase counts = nodes %d relationships %d, want 3 each", nodeResult.Count, relationshipResult.Count) + } + sort.Slice(nodes, func(i, j int) bool { return nodes[i].SourceID < nodes[j].SourceID }) + if got := []string{nodes[0].SourceID, nodes[1].SourceID, nodes[2].SourceID}; !reflect.DeepEqual(got, []string{"1", "2", "3"}) { + t.Fatalf("node prefix = %v, want [1 2 3]", got) + } + if got := nodes[0].Kinds; !reflect.DeepEqual(got, []string{"User", "Admin", "User"}) { + t.Fatalf("first node kinds = %v", got) } - if noCompressionResult.Count != 2 || noCompressionResult.UncompressedByteSize != 0 || noCompressionResult.CompressedByteSize != 0 { - t.Fatalf("unexpected no-compression node result: %+v", noCompressionResult) + sort.Ints(nodeBatchSizes) + if !reflect.DeepEqual(nodeBatchSizes, []int{1, 2}) { + t.Fatalf("node batch sizes = %v, want final partial batch", nodeBatchSizes) } - relationships := []*graph.Relationship{ - graph.NewRelationship(10, 1, 2, graph.AsProperties(map[string]any{"source": "test"}), graph.StringKind("AdminTo")), - graph.NewRelationship(11, 2, 3, nil, nil), + sort.Slice(relationships, func(i, j int) bool { return relationships[i].SourceID < relationships[j].SourceID }) + if got := []string{relationships[0].SourceID, relationships[1].SourceID, relationships[2].SourceID}; !reflect.DeepEqual(got, []string{"10", "11", "12"}) { + t.Fatalf("relationship prefix/source IDs = %v, want [10 11 12]", got) } - relationshipResult, err := benchRelationshipBatch(relationships, options) + sort.Ints(relationshipBatchSizes) + if !reflect.DeepEqual(relationshipBatchSizes, []int{1, 2}) { + t.Fatalf("relationship batch sizes = %v, want final partial batch", relationshipBatchSizes) + } +} + +func TestBenchSourcePhaseErrorsAndCountMismatches(t *testing.T) { + sourceFailure := errors.New("source failed") + database := &benchSourceDatabase{ + nodes: []*graph.Node{graph.NewNode(1, graph.NewProperties(), graph.StringKind("User"))}, + nodeFetchErr: sourceFailure, + } + options := benchOptions{BatchSize: 1, SampleSize: 1} + if _, err := benchNodes( + context.Background(), database, "example", "jsonl", 1, 1, options, + func(_ string, batch []entity.Node) (benchPhaseResult, error) { + return benchPhaseResult{Count: int64(len(batch))}, nil + }, + ); !errors.Is(err, sourceFailure) { + t.Fatalf("node source error = %v, want source failure", err) + } + + database.nodeFetchErr = nil + if _, err := benchNodes( + context.Background(), database, "example", "jsonl", 1, 1, options, + func(string, []entity.Node) (benchPhaseResult, error) { + return benchPhaseResult{}, nil + }, + ); err == nil || !strings.Contains(err.Error(), "wrote 0 of 1") { + t.Fatalf("node count mismatch error = %v", err) + } + + database.relationships = []*graph.Relationship{ + graph.NewRelationship(2, 1, 3, graph.NewProperties(), graph.StringKind("MemberOf")), + } + database.relationshipFetchErr = sourceFailure + if _, err := benchRelationships( + context.Background(), database, "example", "parquet", 1, 1, options, + func(_ string, batch []entity.Relationship) (benchPhaseResult, error) { + return benchPhaseResult{Count: int64(len(batch))}, nil + }, + ); !errors.Is(err, sourceFailure) { + t.Fatalf("relationship source error = %v, want source failure", err) + } + + database.relationshipFetchErr = nil + if _, err := benchRelationships( + context.Background(), database, "example", "parquet", 1, 1, options, + func(string, []entity.Relationship) (benchPhaseResult, error) { + return benchPhaseResult{}, nil + }, + ); err == nil || !strings.Contains(err.Error(), "wrote 0 of 1") { + t.Fatalf("relationship count mismatch error = %v", err) + } +} + +func TestBenchPhaseReportsArtifactCleanupErrors(t *testing.T) { + cleanupFailure := errors.New("cleanup failed") + phaseFailure := errors.New("phase failed") + filesystem := benchArtifactFilesystem{ + removeAll: func(path string) error { + if err := os.RemoveAll(path); err != nil { + t.Fatalf("remove benchmark test directory: %v", err) + } + return cleanupFailure + }, + } + options := benchOptions{BatchSize: 1} + nodeDatabase := &benchSourceDatabase{ + nodes: []*graph.Node{graph.NewNode(1, graph.NewProperties(), graph.StringKind("User"))}, + nodeFetchErr: phaseFailure, + } + + if _, err := benchNodesWithFilesystem( + context.Background(), nodeDatabase, "example", "jsonl", 1, 1, options, + func(string, []entity.Node) (benchPhaseResult, error) { return benchPhaseResult{}, nil }, + filesystem, + ); !errors.Is(err, phaseFailure) || !errors.Is(err, cleanupFailure) { + t.Fatalf("node cleanup error = %v, want joined phase and cleanup failures", err) + } + if _, err := benchRelationshipsWithFilesystem( + context.Background(), emptyBenchDatabase{}, "example", "parquet", 0, 1, options, + func(string, []entity.Relationship) (benchPhaseResult, error) { return benchPhaseResult{}, nil }, + filesystem, + ); !errors.Is(err, cleanupFailure) { + t.Fatalf("relationship cleanup error = %v, want cleanup failure", err) + } +} + +func TestBenchConcreteFormatBatches(t *testing.T) { + jsonlDirectory := t.TempDir() + jsonlPath := filepath.Join(jsonlDirectory, "nodes.jsonl.gz") + nodes := []entity.Node{ + {SourceID: "2", Kinds: []string{"User", "Admin", "User"}, Properties: map[string]any{"name": "alice"}}, + {SourceID: "1", Kinds: []string{"Computer"}}, + } + jsonlResult, err := benchJSONLNodeBatch( + jsonlPath, + nodes, + jsonl.Config{Codec: jsonl.CodecGzip}, + ) + if err != nil { + t.Fatalf("bench JSONL node batch: %v", err) + } + if jsonlResult.Count != 2 || jsonlResult.UncompressedByteSize <= 0 || jsonlResult.CompressedByteSize <= 0 { + t.Fatalf("unexpected JSONL node batch result: %+v", jsonlResult) + } + file, err := openGzipJSONL(jsonlPath) + if err != nil { + t.Fatalf("open benchmark JSONL: %v", err) + } + defer file.Close() + var first struct { + Kinds []string `json:"kinds"` + } + if err := json.NewDecoder(file).Decode(&first); err != nil { + t.Fatalf("decode benchmark JSONL: %v", err) + } + if got := strings.Join(first.Kinds, ","); got != "User,Admin,User" { + t.Fatalf("node kind order = %q", got) + } + parquetResult, err := benchParquetNodeBatch( + filepath.Join(t.TempDir(), "nodes.parquet"), + nodes, + parquet.Config{}, + ) + if err != nil { + t.Fatalf("bench Parquet node batch: %v", err) + } + if parquetResult.Count != 2 || parquetResult.UncompressedByteSize != 0 || parquetResult.CompressedByteSize <= 0 { + t.Fatalf("unexpected Parquet node batch result: %+v", parquetResult) + } + + relationships := []entity.Relationship{ + {SourceID: "10", StartID: "1", EndID: "2", Kind: "AdminTo", Properties: map[string]any{"source": "test"}}, + {SourceID: "11", StartID: "2", EndID: "3", Kind: "MemberOf"}, + } + relationshipResult, err := benchParquetRelationshipBatch( + filepath.Join(t.TempDir(), "relationships.parquet"), + relationships, + parquet.Config{}, + ) if err != nil { - t.Fatalf("bench relationship batch: %v", err) + t.Fatalf("bench Parquet relationship batch: %v", err) } - if relationshipResult.Count != 2 || relationshipResult.UncompressedByteSize <= 0 || relationshipResult.CompressedByteSize <= 0 { + if relationshipResult.Count != 2 || relationshipResult.CompressedByteSize <= 0 { t.Fatalf("unexpected relationship batch result: %+v", relationshipResult) } } + +func TestBenchReportsBothConcreteFormatsIndependently(t *testing.T) { + report, err := Bench(context.Background(), emptyBenchDatabase{}, "pg", []string{"default"}, benchOptions{ + Workers: []int{1}, + BatchSize: 10, + SampleSize: 10, + JSONL: &jsonl.Config{Codec: jsonl.CodecZstd}, + Parquet: &parquet.Config{}, + }) + if err != nil { + t.Fatalf("bench: %v", err) + } + if len(report.Graphs) != 1 || len(report.Graphs[0].Results) != 2 { + t.Fatalf("report = %+v", report) + } + if got := report.Graphs[0].Results[0].Format; got != "jsonl" { + t.Fatalf("first format = %q", got) + } + if got := report.Graphs[0].Results[1].Format; got != "parquet" { + t.Fatalf("second format = %q", got) + } +} + +func openGzipJSONL(path string) (*gzip.Reader, error) { + stored, err := os.ReadFile(path) + if err != nil { + return nil, err + } + return gzip.NewReader(bytes.NewReader(stored)) +} + +type emptyBenchDatabase struct { + graph.Database +} + +func (emptyBenchDatabase) ReadTransaction(_ context.Context, delegate graph.TransactionDelegate, _ ...graph.TransactionOption) error { + return delegate(emptyBenchTransaction{}) +} + +type emptyBenchTransaction struct { + graph.Transaction +} + +func (emptyBenchTransaction) WithGraph(graph.Graph) graph.Transaction { + return emptyBenchTransaction{} +} + +func (emptyBenchTransaction) Nodes() graph.NodeQuery { + return emptyBenchNodeQuery{} +} + +func (emptyBenchTransaction) Relationships() graph.RelationshipQuery { + return emptyBenchRelationshipQuery{} +} + +type emptyBenchNodeQuery struct { + graph.NodeQuery +} + +func (emptyBenchNodeQuery) Count() (int64, error) { + return 0, nil +} + +type emptyBenchRelationshipQuery struct { + graph.RelationshipQuery +} + +func (emptyBenchRelationshipQuery) Count() (int64, error) { + return 0, nil +} + +type benchTestCursor[T any] struct { + values chan T + err error +} + +func newBenchTestCursor[T any](values []T) *benchTestCursor[T] { + channel := make(chan T, len(values)) + for _, value := range values { + channel <- value + } + close(channel) + return &benchTestCursor[T]{values: channel} +} + +func (s *benchTestCursor[T]) Error() error { return s.err } +func (s *benchTestCursor[T]) Close() {} +func (s *benchTestCursor[T]) Chan() chan T { return s.values } + +type benchSourceDatabase struct { + graph.Database + nodes []*graph.Node + relationships []*graph.Relationship + nodeFetchErr error + relationshipFetchErr error +} + +func (s *benchSourceDatabase) ReadTransaction(_ context.Context, delegate graph.TransactionDelegate, _ ...graph.TransactionOption) error { + return delegate(&benchSourceTransaction{database: s}) +} + +type benchSourceTransaction struct { + graph.Transaction + database *benchSourceDatabase +} + +func (s *benchSourceTransaction) WithGraph(graph.Graph) graph.Transaction { return s } +func (s *benchSourceTransaction) Nodes() graph.NodeQuery { + return &benchSourceNodeQuery{database: s.database} +} +func (s *benchSourceTransaction) Relationships() graph.RelationshipQuery { + return &benchSourceRelationshipQuery{database: s.database} +} + +type benchSourceNodeQuery struct { + graph.NodeQuery + database *benchSourceDatabase + afterID graph.ID + limit int +} + +func (s *benchSourceNodeQuery) OrderBy(...graph.Criteria) graph.NodeQuery { return s } +func (s *benchSourceNodeQuery) Filter(criteria graph.Criteria) graph.NodeQuery { + s.afterID = benchSourceAfterID(criteria) + return s +} +func (s *benchSourceNodeQuery) Limit(limit int) graph.NodeQuery { + s.limit = limit + return s +} +func (s *benchSourceNodeQuery) Count() (int64, error) { + return int64(len(s.database.nodes)), nil +} +func (s *benchSourceNodeQuery) Fetch(delegate func(graph.Cursor[*graph.Node]) error, _ ...graph.Criteria) error { + if s.database.nodeFetchErr != nil { + return s.database.nodeFetchErr + } + values := make([]*graph.Node, 0, s.limit) + for _, node := range s.database.nodes { + if node.ID > s.afterID && len(values) < s.limit { + values = append(values, node) + } + } + return delegate(newBenchTestCursor(values)) +} + +type benchSourceRelationshipQuery struct { + graph.RelationshipQuery + database *benchSourceDatabase + afterID graph.ID + limit int +} + +func (s *benchSourceRelationshipQuery) OrderBy(...graph.Criteria) graph.RelationshipQuery { + return s +} +func (s *benchSourceRelationshipQuery) Filter(criteria graph.Criteria) graph.RelationshipQuery { + s.afterID = benchSourceAfterID(criteria) + return s +} +func (s *benchSourceRelationshipQuery) Limit(limit int) graph.RelationshipQuery { + s.limit = limit + return s +} +func (s *benchSourceRelationshipQuery) Count() (int64, error) { + return int64(len(s.database.relationships)), nil +} +func (s *benchSourceRelationshipQuery) Fetch(delegate func(graph.Cursor[*graph.Relationship]) error) error { + if s.database.relationshipFetchErr != nil { + return s.database.relationshipFetchErr + } + values := make([]*graph.Relationship, 0, s.limit) + for _, relationship := range s.database.relationships { + if relationship.ID > s.afterID && len(values) < s.limit { + values = append(values, relationship) + } + } + return delegate(newBenchTestCursor(values)) +} + +func benchSourceAfterID(criteria graph.Criteria) graph.ID { + comparison := criteria.(*cypherModel.Comparison) + return comparison.Partials[0].Right.(*cypherModel.Parameter).Value.(graph.ID) +} diff --git a/cmd/retriever/config.go b/cmd/retriever/config.go index a9aac300..e7d527b6 100644 --- a/cmd/retriever/config.go +++ b/cmd/retriever/config.go @@ -3,14 +3,151 @@ package main import ( "flag" "fmt" + "io" "os" "strconv" "strings" - "github.com/specterops/dawgs/retriever" + "github.com/specterops/dawgs/ret" + "github.com/specterops/dawgs/ret/jsonl" + "github.com/specterops/dawgs/ret/parquet" + "github.com/specterops/dawgs/ret/scrub" ) -const defaultBenchSampleSize = 1_000_000 +const ( + defaultGraphName = "default" + defaultEntityBatchSize = 10_000 + defaultShardSize = 100_000 + defaultBenchSampleSize = 1_000_000 +) + +type dumpCommandConfig struct { + database databaseConfig + dump ret.DumpConfig + graphs []string + allGraphs bool + force bool + pprof string +} + +func parseDumpCommand(args []string, output io.Writer) (dumpCommandConfig, error) { + jsonlConfig := jsonl.Config{ + Codec: jsonl.CodecZstd, + Level: 0, + } + parquetConfig := parquet.Config{} + config := dumpCommandConfig{ + dump: ret.DumpConfig{ + EntityBatchSize: defaultEntityBatchSize, + ShardSize: defaultShardSize, + JSONL: &jsonlConfig, + }, + } + var ( + graphs stringList + scrubMode string + scrubSalt string + scrubConfigPath string + jsonlCompression = string(jsonlConfig.Codec) + jsonlEnabled = true + parquetEnabled bool + scrubConfig = scrub.DefaultConfig() + ) + flags := flag.NewFlagSet("retriever dump", flag.ContinueOnError) + flags.SetOutput(output) + commonDatabaseFlags(flags, &config.database) + flags.Var(&graphs, "graph", "Graph target. May be repeated.") + flags.BoolVar(&config.allGraphs, "all-graphs", false, "Dump every graph discoverable by the selected driver.") + flags.StringVar(&config.dump.Directory, "out", "", "Output collection directory.") + flags.BoolVar(&config.force, "force", false, "Replace the exact output directory before a fresh dump.") + flags.BoolVar(&config.dump.Resume, "resume", false, "Resume an interrupted dump from its validated checkpoint.") + flags.BoolVar(&jsonlEnabled, "jsonl", jsonlEnabled, "Write JSONL artifacts.") + flags.StringVar(&jsonlCompression, "jsonl-compression", jsonlCompression, "JSONL compression codec: zstd, gzip, or none.") + flags.IntVar(&jsonlConfig.Level, "jsonl-level", jsonlConfig.Level, "JSONL compression level; 0 selects the package default.") + flags.BoolVar(&parquetEnabled, "parquet", parquetEnabled, "Write Parquet artifacts.") + flags.StringVar(&scrubMode, "scrub", "none", "Scrub mode: none or full.") + flags.StringVar(&scrubSalt, "salt", "", "Scrub salt. Overrides RETRIEVER_SCRUB_SALT and is never written.") + flags.StringVar(&scrubConfigPath, "config", "", "Optional retriever TOML scrub configuration.") + flags.IntVar(&config.dump.ShardSize, "shard-size", config.dump.ShardSize, "Maximum entities per shard.") + flags.IntVar(&config.dump.EntityBatchSize, "batch-size", config.dump.EntityBatchSize, "Database read batch size.") + commonPprofFlag(flags, &config.pprof) + if err := flags.Parse(args); err != nil { + return dumpCommandConfig{}, err + } + + fillConnectionFromEnv(&config.database) + config.graphs = append([]string(nil), graphs...) + config.dump.Directory = strings.TrimSpace(config.dump.Directory) + jsonlConfig.Codec = jsonl.Codec(strings.TrimSpace(jsonlCompression)) + if jsonlEnabled { + config.dump.JSONL = &jsonlConfig + } else { + config.dump.JSONL = nil + } + if parquetEnabled { + config.dump.Parquet = &parquetConfig + } else { + config.dump.Parquet = nil + } + + if path := strings.TrimSpace(scrubConfigPath); path != "" { + loaded, err := scrub.ReadConfig(path) + if err != nil { + return dumpCommandConfig{}, err + } + scrubConfig = loaded + } + if strings.TrimSpace(scrubSalt) == "" { + scrubSalt = strings.TrimSpace(os.Getenv("RETRIEVER_SCRUB_SALT")) + if scrubSalt == "" { + scrubSalt = strings.TrimSpace(os.Getenv("RETRIEVR_SCRUB_SALT")) + } + } + scrubConfig.Salt = strings.TrimSpace(scrubSalt) + switch strings.TrimSpace(scrubMode) { + case "none": + config.dump.Scrub = nil + case "full": + if scrubConfig.Salt == "" { + return dumpCommandConfig{}, fmt.Errorf("-scrub full requires -salt, RETRIEVER_SCRUB_SALT, or legacy RETRIEVR_SCRUB_SALT") + } + config.dump.Scrub = &scrubConfig + default: + return dumpCommandConfig{}, fmt.Errorf("unsupported scrub mode %q", scrubMode) + } + + if config.dump.Directory == "" { + return dumpCommandConfig{}, fmt.Errorf("output directory is required; pass -out") + } + if config.force && config.dump.Resume { + return dumpCommandConfig{}, fmt.Errorf("-force and -resume are mutually exclusive") + } + if config.dump.EntityBatchSize <= 0 { + return dumpCommandConfig{}, fmt.Errorf("batch-size must be > 0") + } + if config.dump.ShardSize <= 0 { + return dumpCommandConfig{}, fmt.Errorf("shard-size must be > 0") + } + if config.dump.JSONL == nil && config.dump.Parquet == nil { + return dumpCommandConfig{}, fmt.Errorf("at least one of -jsonl or -parquet must be enabled") + } + if config.dump.JSONL != nil { + if err := config.dump.JSONL.Validate(); err != nil { + return dumpCommandConfig{}, fmt.Errorf("JSONL configuration: %w", err) + } + } + if config.dump.Parquet != nil { + if err := config.dump.Parquet.Validate(); err != nil { + return dumpCommandConfig{}, fmt.Errorf("Parquet configuration: %w", err) + } + } + if config.dump.Scrub != nil { + if err := config.dump.Scrub.Validate(); err != nil { + return dumpCommandConfig{}, fmt.Errorf("scrub configuration: %w", err) + } + } + return config, nil +} type stringList []string @@ -99,12 +236,12 @@ func parseWorkerList(value string) ([]int, error) { } type benchOptions struct { - Workers []int - BatchSize int - SampleSize int - Compression retriever.CompressionCodec - ZstdLevel int - JSONOutput bool + Workers []int + BatchSize int + SampleSize int + JSONL *jsonl.Config + Parquet *parquet.Config + JSONOutput bool } func (s benchOptions) validate() error { @@ -126,19 +263,90 @@ func (s benchOptions) validate() error { return fmt.Errorf("sample-size must be >= 0") } - if s.ZstdLevel <= 0 { - return fmt.Errorf("zstd-level must be > 0") + if s.JSONL == nil && s.Parquet == nil { + return fmt.Errorf("at least one of -jsonl or -parquet must be enabled") } - if s.Compression != retriever.CompressionDisabled { - if err := retriever.ValidateCompression(s.Compression); err != nil { - return err + if s.JSONL != nil { + if err := s.JSONL.Validate(); err != nil { + return fmt.Errorf("JSONL configuration: %w", err) + } + } + if s.Parquet != nil { + if err := s.Parquet.Validate(); err != nil { + return fmt.Errorf("Parquet configuration: %w", err) } } return nil } +type benchCommandConfig struct { + database databaseConfig + bench benchOptions + graphs []string + allGraphs bool + pprof string +} + +func parseBenchCommand(args []string, output io.Writer) (benchCommandConfig, error) { + jsonlConfig := jsonl.Config{Codec: jsonl.CodecZstd} + parquetConfig := parquet.Config{} + config := benchCommandConfig{ + bench: benchOptions{ + Workers: []int{1}, + BatchSize: defaultEntityBatchSize, + SampleSize: defaultBenchSampleSize, + JSONL: &jsonlConfig, + }, + } + var ( + graphs stringList + workers workerList + jsonlCompression = string(jsonlConfig.Codec) + jsonlEnabled = true + parquetEnabled bool + ) + flags := flag.NewFlagSet("retriever bench", flag.ContinueOnError) + flags.SetOutput(output) + commonDatabaseFlags(flags, &config.database) + flags.Var(&graphs, "graph", "Graph target. May be repeated.") + flags.BoolVar(&config.allGraphs, "all-graphs", false, "Benchmark every graph discoverable by the selected driver.") + flags.Var(&workers, "workers", "Comma-separated worker counts.") + flags.IntVar(&config.bench.BatchSize, "batch-size", config.bench.BatchSize, "Database read batch size.") + flags.IntVar(&config.bench.SampleSize, "sample-size", config.bench.SampleSize, "Maximum nodes and relationships to scan per phase; 0 scans the full graph.") + flags.BoolVar(&jsonlEnabled, "jsonl", jsonlEnabled, "Benchmark JSONL artifacts.") + flags.StringVar(&jsonlCompression, "jsonl-compression", jsonlCompression, "JSONL compression codec: zstd, gzip, or none.") + flags.IntVar(&jsonlConfig.Level, "jsonl-level", jsonlConfig.Level, "JSONL compression level; 0 selects the package default.") + flags.BoolVar(&parquetEnabled, "parquet", parquetEnabled, "Benchmark Parquet artifacts.") + flags.BoolVar(&config.bench.JSONOutput, "json", false, "Emit machine-readable JSON.") + commonPprofFlag(flags, &config.pprof) + if err := flags.Parse(args); err != nil { + return benchCommandConfig{}, err + } + + fillConnectionFromEnv(&config.database) + config.graphs = append([]string(nil), graphs...) + if len(workers) > 0 { + config.bench.Workers = append([]int(nil), workers...) + } + jsonlConfig.Codec = jsonl.Codec(strings.TrimSpace(jsonlCompression)) + if jsonlEnabled { + config.bench.JSONL = &jsonlConfig + } else { + config.bench.JSONL = nil + } + if parquetEnabled { + config.bench.Parquet = &parquetConfig + } else { + config.bench.Parquet = nil + } + if err := config.bench.validate(); err != nil { + return benchCommandConfig{}, err + } + return config, nil +} + func commonDatabaseFlags(flags *flag.FlagSet, cfg *databaseConfig) { flags.StringVar(&cfg.Driver, "driver", "", "Graph database driver. Inferred from -connection when omitted.") flags.StringVar(&cfg.Connection, "connection", "", "Graph database connection string. Falls back to CONNECTION_STRING.") diff --git a/cmd/retriever/config_test.go b/cmd/retriever/config_test.go index 041d0c07..462357d0 100644 --- a/cmd/retriever/config_test.go +++ b/cmd/retriever/config_test.go @@ -1,27 +1,139 @@ package main import ( + "os" + "path/filepath" + "reflect" "testing" - "github.com/specterops/dawgs/retriever" + "github.com/specterops/dawgs/ret/jsonl" + "github.com/specterops/dawgs/ret/parquet" + "github.com/specterops/dawgs/ret/scrub" ) +func TestParseDumpDefaultsToJSONLZstdWithoutParquet(t *testing.T) { + config, err := parseDumpCommand([]string{"-out", t.TempDir()}, testFlagOutput(t)) + if err != nil { + t.Fatalf("parse dump: %v", err) + } + if config.dump.JSONL == nil { + t.Fatal("JSONL is disabled by default") + } + if config.dump.JSONL.Codec != jsonl.CodecZstd { + t.Fatalf("JSONL codec = %q, want %q", config.dump.JSONL.Codec, jsonl.CodecZstd) + } + if config.dump.JSONL.Level != 0 { + t.Fatalf("JSONL level = %d, want package default 0", config.dump.JSONL.Level) + } + if config.dump.Parquet != nil { + t.Fatal("Parquet is enabled by default") + } + if config.dump.Scrub != nil { + t.Fatalf("scrubbing is enabled by default: %+v", config.dump.Scrub) + } +} + +func TestParseDumpKeepsJSONLAndParquetIndependent(t *testing.T) { + config, err := parseDumpCommand([]string{ + "-out", t.TempDir(), + "-jsonl=false", + "-parquet", + }, testFlagOutput(t)) + if err != nil { + t.Fatalf("parse dump: %v", err) + } + if config.dump.JSONL != nil { + t.Fatal("JSONL remained enabled") + } + if config.dump.Parquet == nil { + t.Fatal("Parquet was not enabled") + } +} + +func TestParseDumpRejectsInvalidIndependentJSONLConfig(t *testing.T) { + _, err := parseDumpCommand([]string{ + "-out", t.TempDir(), + "-jsonl-compression", "zip", + }, testFlagOutput(t)) + if err == nil { + t.Fatal("expected invalid JSONL codec") + } +} + +func TestParseDumpEnablesScrubbingWithRuntimeSalt(t *testing.T) { + config, err := parseDumpCommand([]string{ + "-out", t.TempDir(), + "-scrub", "full", + "-salt", "runtime-salt", + }, testFlagOutput(t)) + if err != nil { + t.Fatalf("parse dump: %v", err) + } + if config.dump.Scrub == nil { + t.Fatal("scrubbing remained disabled") + } + if config.dump.Scrub.Salt != "runtime-salt" { + t.Fatalf("scrub salt = %q", config.dump.Scrub.Salt) + } + if !reflect.DeepEqual(config.dump.Scrub.Rules, scrub.DefaultConfig().Rules) { + t.Fatalf("scrub rules differ from defaults: %+v", config.dump.Scrub.Rules) + } +} + +func TestParseDumpReadsDirectScrubPolicyAndOverridesFileSaltAtRuntime(t *testing.T) { + path := filepath.Join(t.TempDir(), "scrub.toml") + if err := os.WriteFile(path, []byte(` +salt = "must-not-load" +fake_domain = "scrub.example" + +[graph_rules] +domain_kind = "CustomDomain" + +[classifier] +long_text_threshold = 8 +`), 0o600); err != nil { + t.Fatalf("write scrub config: %v", err) + } + + config, err := parseDumpCommand([]string{ + "-out", t.TempDir(), + "-scrub", "full", + "-salt", "runtime-salt", + "-config", path, + }, testFlagOutput(t)) + if err != nil { + t.Fatalf("parse dump: %v", err) + } + if config.dump.Scrub == nil { + t.Fatal("scrubbing remained disabled") + } + if config.dump.Scrub.Salt != "runtime-salt" { + t.Fatalf("scrub salt = %q", config.dump.Scrub.Salt) + } + if config.dump.Scrub.Rules.FakeDomain != "scrub.example" { + t.Fatalf("fake domain = %q", config.dump.Scrub.Rules.FakeDomain) + } + if config.dump.Scrub.Rules.GraphRules.DomainKind != "CustomDomain" { + t.Fatalf("domain kind = %q", config.dump.Scrub.Rules.GraphRules.DomainKind) + } + if config.dump.Scrub.Rules.Classifier.LongTextThreshold != 8 { + t.Fatalf("long text threshold = %d", config.dump.Scrub.Rules.Classifier.LongTextThreshold) + } +} + func TestParseWorkerList(t *testing.T) { workers, err := parseWorkerList("1,2,4,2") if err != nil { t.Fatalf("parse worker list: %v", err) } - if got, want := len(workers), 3; got != want { t.Fatalf("worker count length = %d, want %d", got, want) } - if workers[0] != 1 || workers[1] != 2 || workers[2] != 4 { t.Fatalf("unexpected workers: %v", workers) } - if _, err := parseWorkerList("0"); err == nil { - t.Fatalf("expected invalid worker count error") + t.Fatal("expected invalid worker count error") } } @@ -30,21 +142,17 @@ func TestFlagListTypes(t *testing.T) { if err := graphs.Set(" default "); err != nil { t.Fatalf("set graph: %v", err) } - if err := graphs.Set(""); err == nil { - t.Fatalf("expected empty graph error") + t.Fatal("expected empty graph error") } - if graphs.String() != "default" { t.Fatalf("graph list string = %q", graphs.String()) } var workers workerList - if err := workers.Set("2,4"); err != nil { t.Fatalf("set workers: %v", err) } - if workers.String() != "2,4" { t.Fatalf("worker list string = %q", workers.String()) } @@ -52,23 +160,18 @@ func TestFlagListTypes(t *testing.T) { func TestWorkerListAppendsRepeatedFlags(t *testing.T) { var workers workerList - if err := workers.Set("1,2"); err != nil { t.Fatalf("set initial workers: %v", err) } - if err := workers.Set("2,4"); err != nil { t.Fatalf("set repeated workers: %v", err) } - if workers.String() != "1,2,4" { t.Fatalf("worker list string = %q", workers.String()) } - if err := workers.Set("bad"); err == nil { - t.Fatalf("expected invalid worker count") + t.Fatal("expected invalid worker count") } - if workers.String() != "1,2,4" { t.Fatalf("invalid worker update changed list to %q", workers.String()) } @@ -79,34 +182,95 @@ func TestBenchOptionsValidate(t *testing.T) { Workers: []int{1}, BatchSize: 1, SampleSize: 1, - ZstdLevel: retriever.DefaultZstdLevel, + JSONL: &jsonl.Config{Codec: jsonl.CodecZstd}, } if err := bench.validate(); err != nil { t.Fatalf("valid bench options: %v", err) } - bench.Workers = nil - if err := bench.validate(); err == nil { - t.Fatalf("expected missing workers") + t.Fatal("expected missing workers") } - bench.Workers = []int{2} - if err := bench.validate(); err != nil { t.Fatalf("valid parallel bench workers: %v", err) } - bench.Workers = []int{0} - if err := bench.validate(); err == nil { - t.Fatalf("expected invalid worker count") + t.Fatal("expected invalid worker count") } - bench.Workers = []int{1} bench.SampleSize = -1 - if err := bench.validate(); err == nil { - t.Fatalf("expected invalid sample size") + t.Fatal("expected invalid sample size") + } +} + +func TestParseBenchAllowsIndependentFormatSelection(t *testing.T) { + for name, args := range map[string][]string{ + "jsonl": {"-jsonl=true", "-parquet=false"}, + "parquet": {"-jsonl=false", "-parquet=true"}, + "both": {"-jsonl=true", "-parquet=true"}, + } { + t.Run(name, func(t *testing.T) { + config, err := parseBenchCommand(args, testFlagOutput(t)) + if err != nil { + t.Fatalf("parse bench: %v", err) + } + if (config.bench.JSONL != nil) != (name != "parquet") { + t.Fatalf("JSONL enabled = %t", config.bench.JSONL != nil) + } + if (config.bench.Parquet != nil) != (name != "jsonl") { + t.Fatalf("Parquet enabled = %t", config.bench.Parquet != nil) + } + }) } } + +func TestParseBenchDefaultsToConcreteJSONLZstd(t *testing.T) { + config, err := parseBenchCommand(nil, testFlagOutput(t)) + if err != nil { + t.Fatalf("parse bench: %v", err) + } + if config.bench.JSONL == nil || config.bench.JSONL.Codec != jsonl.CodecZstd || config.bench.JSONL.Level != 0 { + t.Fatalf("JSONL config = %+v", config.bench.JSONL) + } + if config.bench.Parquet != nil { + t.Fatal("Parquet enabled by default") + } + if len(config.bench.Workers) != 1 || config.bench.Workers[0] != 1 { + t.Fatalf("workers = %v", config.bench.Workers) + } +} + +func TestParseBenchRejectsNeitherFormatAndInvalidJSONLConfig(t *testing.T) { + if _, err := parseBenchCommand([]string{"-jsonl=false", "-parquet=false"}, testFlagOutput(t)); err == nil { + t.Fatal("expected neither-format error") + } + if _, err := parseBenchCommand([]string{"-jsonl-compression", "zip"}, testFlagOutput(t)); err == nil { + t.Fatal("expected invalid JSONL codec error") + } + + config, err := parseBenchCommand([]string{ + "-jsonl=false", + "-parquet=true", + "-jsonl-compression", "zip", + }, testFlagOutput(t)) + if err != nil { + t.Fatalf("disabled JSONL config affected Parquet-only benchmark: %v", err) + } + if config.bench.Parquet == nil || *config.bench.Parquet != (parquet.Config{}) { + t.Fatalf("Parquet config = %+v", config.bench.Parquet) + } +} + +func testFlagOutput(t *testing.T) *discardWriter { + t.Helper() + return &discardWriter{} +} + +type discardWriter struct{} + +func (*discardWriter) Write(value []byte) (int, error) { + return len(value), nil +} diff --git a/cmd/retriever/database.go b/cmd/retriever/database.go index 5c85760d..e5d2d5c6 100644 --- a/cmd/retriever/database.go +++ b/cmd/retriever/database.go @@ -2,16 +2,18 @@ package main import ( "context" + "errors" "fmt" "net/url" + "path" "strings" + "time" "github.com/jackc/pgx/v5/pgxpool" "github.com/specterops/dawgs" "github.com/specterops/dawgs/drivers/neo4j" "github.com/specterops/dawgs/drivers/pg" "github.com/specterops/dawgs/graph" - "github.com/specterops/dawgs/retriever" "github.com/specterops/dawgs/util/size" ) @@ -21,7 +23,17 @@ type databaseConfig struct { Graph string } +const productDatabaseCloseTimeout = 30 * time.Second + +type databaseOpenOperations struct { + open func(context.Context, string, dawgs.Config) (graph.Database, error) +} + func openDatabase(ctx context.Context, cfg databaseConfig) (graph.Database, string, error) { + return openDatabaseWith(ctx, cfg, databaseOpenOperations{}) +} + +func openDatabaseWith(ctx context.Context, cfg databaseConfig, operations databaseOpenOperations) (result graph.Database, driver string, resultErr error) { connection := strings.TrimSpace(cfg.Connection) if connection == "" { return nil, "", fmt.Errorf("database connection is required; pass -connection or set CONNECTION_STRING") @@ -65,17 +77,20 @@ func openDatabase(ctx context.Context, cfg databaseConfig) (graph.Database, stri return nil, "", fmt.Errorf("unsupported driver %q; expected %s or %s", driverName, pg.DriverName, neo4j.DriverName) } - db, err := dawgs.Open(ctx, driverName, openConfig) + open := operations.open + if open == nil { + open = dawgs.Open + } + db, err := open(ctx, driverName, openConfig) if err != nil { return nil, "", fmt.Errorf("open %s database: %w", driverName, err) } poolOwnedByDriver = true - openSuccess := false defer func() { - if !openSuccess { - _ = db.Close(ctx) + if resultErr != nil { + resultErr = errors.Join(resultErr, closeProductDatabase(db)) } }() @@ -87,11 +102,21 @@ func openDatabase(ctx context.Context, cfg databaseConfig) (graph.Database, stri } } - openSuccess = true - return db, driverName, nil } +func closeProductDatabase(database graph.Database) error { + if database == nil { + return nil + } + ctx, cancel := context.WithTimeout(context.Background(), productDatabaseCloseTimeout) + defer cancel() + if err := database.Close(ctx); err != nil { + return fmt.Errorf("close database: %w", err) + } + return nil +} + func driverFromConnectionString(connection string) (string, error) { parsedURL, err := url.Parse(connection) if err != nil { @@ -108,9 +133,9 @@ func driverFromConnectionString(connection string) (string, error) { } } -func resolveGraphTargets(ctx context.Context, db graph.Database, driverName string, requested []string, allGraphs bool) ([]retriever.GraphTarget, error) { - if allGraphs && len(requested) > 0 { - return nil, fmt.Errorf("-all-graphs cannot be combined with -graph") +func resolveGraphNames(ctx context.Context, db graph.Database, driverName string, requested []string, allGraphs bool) ([]string, error) { + if err := validateGraphSelection(requested, allGraphs); err != nil { + return nil, err } if allGraphs { @@ -118,47 +143,51 @@ func resolveGraphTargets(ctx context.Context, db graph.Database, driverName stri case pg.DriverName: return discoverPostgresGraphs(ctx, db) case neo4j.DriverName: - return []retriever.GraphTarget{{ - Name: retriever.DefaultGraphName, - }}, nil + return []string{defaultGraphName}, nil default: return nil, fmt.Errorf("all-graphs is not supported for driver %q", driverName) } } if len(requested) == 0 { - return []retriever.GraphTarget{{ - Name: retriever.DefaultGraphName, - }}, nil + return []string{defaultGraphName}, nil } if driverName == neo4j.DriverName && len(requested) > 1 { return nil, fmt.Errorf("neo4j supports one retriever graph target because Dawgs graph names are no-ops for that driver") } - targets := make([]retriever.GraphTarget, 0, len(requested)) - seen := map[string]struct{}{} + targets := make([]string, 0, len(requested)) + for _, name := range requested { + trimmed := strings.TrimSpace(name) + targets = append(targets, trimmed) + } + + return targets, nil +} +func validateGraphSelection(requested []string, allGraphs bool) error { + if allGraphs && len(requested) > 0 { + return fmt.Errorf("-all-graphs cannot be combined with -graph") + } + seen := make(map[string]struct{}, len(requested)) for _, name := range requested { trimmed := strings.TrimSpace(name) - if trimmed == "" { - return nil, fmt.Errorf("graph name cannot be empty") + if trimmed == "" || trimmed == "." || trimmed == ".." || + path.Clean(trimmed) != trimmed || + strings.ContainsAny(trimmed, "/\\") || + strings.ContainsRune(trimmed, '\x00') { + return fmt.Errorf("graph name %q is not a safe path segment", name) } - - if _, ok := seen[trimmed]; ok { - return nil, fmt.Errorf("duplicate graph target %q", trimmed) + if _, found := seen[trimmed]; found { + return fmt.Errorf("duplicate graph target %q", trimmed) } - seen[trimmed] = struct{}{} - targets = append(targets, retriever.GraphTarget{ - Name: trimmed, - }) } - - return targets, nil + return nil } -func discoverPostgresGraphs(ctx context.Context, db graph.Database) ([]retriever.GraphTarget, error) { +func discoverPostgresGraphs(ctx context.Context, db graph.Database) ([]string, error) { const graphQuery = ` select g.name, @@ -167,7 +196,7 @@ select from graph g order by g.name` - var targets []retriever.GraphTarget + var targets []string if err := db.ReadTransaction(ctx, func(tx graph.Transaction) error { result := tx.Raw(graphQuery, nil) defer result.Close() @@ -187,9 +216,7 @@ order by g.name` return fmt.Errorf("PostgreSQL graph %q is missing expected node/edge partitions", name) } - targets = append(targets, retriever.GraphTarget{ - Name: name, - }) + targets = append(targets, name) } return result.Error() @@ -203,12 +230,3 @@ order by g.name` return targets, nil } - -func graphDirectoryName(name string) string { - escaped := url.PathEscape(name) - if escaped == "" { - return retriever.DefaultGraphName - } - - return escaped -} diff --git a/cmd/retriever/database_test.go b/cmd/retriever/database_test.go index 2b0358fe..0eba46fe 100644 --- a/cmd/retriever/database_test.go +++ b/cmd/retriever/database_test.go @@ -2,11 +2,14 @@ package main import ( "context" + "errors" + "reflect" "testing" + "github.com/specterops/dawgs" "github.com/specterops/dawgs/drivers/neo4j" "github.com/specterops/dawgs/drivers/pg" - "github.com/specterops/dawgs/retriever" + "github.com/specterops/dawgs/graph" ) func TestDriverFromConnectionString(t *testing.T) { @@ -22,68 +25,84 @@ func TestDriverFromConnectionString(t *testing.T) { if err != nil { t.Fatalf("driverFromConnectionString(%q): %v", connection, err) } - if actual != expected { t.Fatalf("driverFromConnectionString(%q) = %q, want %q", connection, actual, expected) } } - if _, err := driverFromConnectionString("mysql://example"); err == nil { - t.Fatalf("expected unsupported scheme error") + t.Fatal("expected unsupported scheme error") } } -func TestResolveGraphTargets(t *testing.T) { - targets, err := resolveGraphTargets(context.Background(), nil, pg.DriverName, nil, false) +func TestResolveGraphNamesPreservesOrderAndRejectsDuplicates(t *testing.T) { + names, err := resolveGraphNames(context.Background(), nil, pg.DriverName, nil, false) if err != nil { t.Fatalf("resolve default graph: %v", err) } - - if len(targets) != 1 || targets[0].Name != retriever.DefaultGraphName { - t.Fatalf("unexpected default targets: %+v", targets) + if !reflect.DeepEqual(names, []string{defaultGraphName}) { + t.Fatalf("default graph names = %v", names) } - targets, err = resolveGraphTargets(context.Background(), nil, pg.DriverName, []string{"a", "b"}, false) + names, err = resolveGraphNames(context.Background(), nil, pg.DriverName, []string{"b", "a"}, false) if err != nil { t.Fatalf("resolve explicit graphs: %v", err) } - - if len(targets) != 2 || targets[0].Name != "a" || targets[1].Name != "b" { - t.Fatalf("unexpected explicit targets: %+v", targets) + if !reflect.DeepEqual(names, []string{"b", "a"}) { + t.Fatalf("explicit graph names = %v", names) } - - if _, err := resolveGraphTargets(context.Background(), nil, pg.DriverName, []string{"a", "a"}, false); err == nil { - t.Fatalf("expected duplicate graph error") + if _, err := resolveGraphNames(context.Background(), nil, pg.DriverName, []string{"a", "a"}, false); err == nil { + t.Fatal("expected duplicate graph error") } - - if _, err := resolveGraphTargets(context.Background(), nil, pg.DriverName, []string{"a"}, true); err == nil { - t.Fatalf("expected all-graphs and graph conflict") + if _, err := resolveGraphNames(context.Background(), nil, pg.DriverName, []string{"a"}, true); err == nil { + t.Fatal("expected all-graphs and graph conflict") } - - if _, err := resolveGraphTargets(context.Background(), nil, neo4j.DriverName, []string{"a", "b"}, false); err == nil { - t.Fatalf("expected neo4j multi-graph error") + if _, err := resolveGraphNames(context.Background(), nil, neo4j.DriverName, []string{"a", "b"}, false); err == nil { + t.Fatal("expected neo4j multi-graph error") } - targets, err = resolveGraphTargets(context.Background(), nil, neo4j.DriverName, nil, true) + names, err = resolveGraphNames(context.Background(), nil, neo4j.DriverName, nil, true) if err != nil { t.Fatalf("resolve neo4j all-graphs: %v", err) } - - if len(targets) != 1 || targets[0].Name != retriever.DefaultGraphName { - t.Fatalf("unexpected neo4j all-graphs target: %+v", targets) + if !reflect.DeepEqual(names, []string{defaultGraphName}) { + t.Fatalf("neo4j all-graphs names = %v", names) } } -func TestGraphDirectoryName(t *testing.T) { - if got := graphDirectoryName("default"); got != "default" { - t.Fatalf("graphDirectoryName(default) = %q", got) +func TestOpenDatabaseJoinsSetDefaultGraphAndCloseFailures(t *testing.T) { + setFailure := errors.New("set default graph failed") + closeFailure := errors.New("close failed") + ctx, cancel := context.WithCancel(context.Background()) + database := &defaultGraphFailureDatabase{ + closingTestDatabase: closingTestDatabase{closeErr: closeFailure}, + setErr: setFailure, + cancel: cancel, } - if got := graphDirectoryName("graph/name"); got != "graph%2Fname" { - t.Fatalf("graphDirectoryName(graph/name) = %q", got) + _, _, err := openDatabaseWith(ctx, databaseConfig{ + Driver: neo4j.DriverName, + Connection: "neo4j://example", + Graph: "asset", + }, databaseOpenOperations{ + open: func(context.Context, string, dawgs.Config) (graph.Database, error) { + return database, nil + }, + }) + if !errors.Is(err, setFailure) || !errors.Is(err, closeFailure) { + t.Fatalf("open database error = %v, want joined set-default and close failures", err) } - - if got := graphDirectoryName(""); got != retriever.DefaultGraphName { - t.Fatalf("graphDirectoryName(empty) = %q", got) + if len(database.closeContextErrors) != 1 || database.closeContextErrors[0] != nil { + t.Fatalf("close context errors = %v, want one non-canceled cleanup context", database.closeContextErrors) } } + +type defaultGraphFailureDatabase struct { + closingTestDatabase + setErr error + cancel context.CancelFunc +} + +func (s *defaultGraphFailureDatabase) SetDefaultGraph(context.Context, graph.Graph) error { + s.cancel() + return s.setErr +} diff --git a/cmd/retriever/force.go b/cmd/retriever/force.go new file mode 100644 index 00000000..cb57e647 --- /dev/null +++ b/cmd/retriever/force.go @@ -0,0 +1,430 @@ +package main + +import ( + "crypto/rand" + "encoding/hex" + "errors" + "fmt" + "io/fs" + "os" + "path/filepath" + "runtime" + "strings" +) + +type forceReplaceOperations struct { + afterParentPinned func(parentPath string, parent *os.Root) error + beforeQuarantine func(parent *os.Root, name string, target *os.Root) error + afterQuarantine func(parent *os.Root, originalName, quarantineName string) error + closeRoot func(role string, root *os.Root) error + closeFile func(role string, file *os.File) error +} + +type forceReplacement struct { + destination string + tombstone string +} + +func replaceDumpDestination(target string, operations forceReplaceOperations) (result forceReplacement, resultErr error) { + if err := validateForcePlatform(runtime.GOOS); err != nil { + return forceReplacement{}, err + } + absolute, err := cleanAbsoluteForceTarget(target) + if err != nil { + return forceReplacement{}, err + } + parentPath, name := filepath.Split(absolute) + parentPath = filepath.Clean(parentPath) + if name == "" || name == "." || name == string(os.PathSeparator) { + return forceReplacement{}, fmt.Errorf("unsafe dump replacement target %q", absolute) + } + + parent, parentInfo, err := openPinnedAbsoluteDirectory(parentPath, operations) + if err != nil { + return forceReplacement{}, fmt.Errorf("pin dump replacement parent: %w", err) + } + parentClosed := false + defer func() { + if !parentClosed { + resultErr = errors.Join(resultErr, wrapForceCloseError( + "parent", + operations.closeRootHandle("parent", parent), + )) + } + }() + if operations.afterParentPinned != nil { + if err := operations.afterParentPinned(parentPath, parent); err != nil { + return forceReplacement{}, fmt.Errorf("force replacement after parent pin: %w", err) + } + } + if err := provePinnedAbsoluteDirectory(parentPath, parentInfo, operations); err != nil { + return forceReplacement{}, fmt.Errorf("dump replacement parent changed after validation: %w", err) + } + + targetInfo, err := parent.Lstat(name) + if errors.Is(err, os.ErrNotExist) { + return forceReplacement{destination: absolute}, nil + } + if err != nil { + return forceReplacement{}, fmt.Errorf("inspect dump replacement target: %w", err) + } + if targetInfo.Mode()&os.ModeSymlink != 0 { + return forceReplacement{}, fmt.Errorf("dump replacement target %q is a symbolic link", absolute) + } + if !targetInfo.IsDir() { + return forceReplacement{}, fmt.Errorf("dump replacement target %q is not a directory", absolute) + } + if err := rejectProtectedForceIdentity(targetInfo); err != nil { + return forceReplacement{}, err + } + + targetRoot, err := parent.OpenRoot(name) + if err != nil { + return forceReplacement{}, fmt.Errorf("pin dump replacement target: %w", err) + } + targetClosed := false + defer func() { + if !targetClosed { + resultErr = errors.Join(resultErr, wrapForceCloseError( + "target", + operations.closeRootHandle("target", targetRoot), + )) + } + }() + pinnedTargetInfo, err := targetRoot.Stat(".") + if err != nil { + return forceReplacement{}, fmt.Errorf("inspect pinned dump replacement target: %w", err) + } + if !os.SameFile(targetInfo, pinnedTargetInfo) { + return forceReplacement{}, fmt.Errorf("dump replacement target changed while being pinned; preserving pathname") + } + + directory, err := parent.Open(".") + if err != nil { + return forceReplacement{}, fmt.Errorf("open pinned dump replacement parent: %w", err) + } + directoryClosed := false + defer func() { + if !directoryClosed { + resultErr = errors.Join(resultErr, wrapForceCloseError( + "parent directory handle", + operations.closeFileHandle("parent directory handle", directory), + )) + } + }() + + current, err := parent.Lstat(name) + if err != nil || !os.SameFile(targetInfo, current) { + return forceReplacement{}, fmt.Errorf("dump replacement target changed before quarantine; preserving replacement") + } + if operations.beforeQuarantine != nil { + if err := operations.beforeQuarantine(parent, name, targetRoot); err != nil { + return forceReplacement{}, fmt.Errorf("force replacement before quarantine: %w", err) + } + } + + quarantineName, err := quarantineForceTarget(directory, name) + if err != nil { + return forceReplacement{}, fmt.Errorf("atomically quarantine dump replacement target: %w", err) + } + tombstone := filepath.Join(parentPath, quarantineName) + restoreOnFailure := true + defer func() { + if restoreOnFailure { + resultErr = errors.Join( + resultErr, + restoreForceQuarantine(directory, parentPath, name, quarantineName), + ) + } + }() + if operations.afterQuarantine != nil { + if err := operations.afterQuarantine(parent, name, quarantineName); err != nil { + return forceReplacement{}, fmt.Errorf("force replacement after quarantine: %w", err) + } + } + quarantinedInfo, err := parent.Lstat(quarantineName) + if err != nil || !os.SameFile(targetInfo, quarantinedInfo) { + return forceReplacement{}, fmt.Errorf("dump replacement target was substituted at quarantine boundary; replacement will be restored or preserved") + } + if err := provePinnedAbsoluteDirectory(parentPath, parentInfo, operations); err != nil { + return forceReplacement{}, fmt.Errorf("dump replacement parent changed after quarantine: %w", err) + } + if _, err := parent.Lstat(name); err == nil { + return forceReplacement{}, fmt.Errorf( + "original destination %q was recreated after quarantine; refusing dump handoff", + absolute, + ) + } else if !errors.Is(err, os.ErrNotExist) { + return forceReplacement{}, fmt.Errorf( + "inspect original destination %q after quarantine: %w", + absolute, + err, + ) + } + targetClosed = true + if err := operations.closeRootHandle("target", targetRoot); err != nil { + return forceReplacement{}, fmt.Errorf("close preserved prior dump destination: %w", err) + } + + parentClosed = true + if err := operations.closeRootHandle("parent", parent); err != nil { + return forceReplacement{}, wrapForceCloseError("parent", err) + } + + directoryClosed = true + if err := operations.closeFileHandle("parent directory handle", directory); err != nil { + restoreOnFailure = false + return forceReplacement{}, fmt.Errorf( + "%w; approved destination preserved at tombstone %q", + wrapForceCloseError("parent directory handle", err), + tombstone, + ) + } + + restoreOnFailure = false + return forceReplacement{ + destination: absolute, + tombstone: tombstone, + }, nil +} + +func validateForcePlatform(platform string) error { + switch platform { + case "linux", "darwin": + return nil + default: + return fmt.Errorf( + "dump -force is unsupported on platform %q; supported platforms are linux and darwin", + platform, + ) + } +} + +func cleanAbsoluteForceTarget(target string) (string, error) { + if strings.TrimSpace(target) == "" { + return "", fmt.Errorf("unsafe dump replacement target: path is empty") + } + absolute, err := filepath.Abs(target) + if err != nil { + return "", fmt.Errorf("resolve dump replacement target: %w", err) + } + absolute = filepath.Clean(absolute) + volumeRoot := filepath.Clean(filepath.VolumeName(absolute) + string(os.PathSeparator)) + if absolute == volumeRoot { + return "", fmt.Errorf("unsafe dump replacement target %q: filesystem root is protected", absolute) + } + return absolute, nil +} + +func openPinnedAbsoluteDirectory(path string, operations forceReplaceOperations) (*os.Root, fs.FileInfo, error) { + absolute, err := filepath.Abs(path) + if err != nil { + return nil, nil, err + } + absolute = filepath.Clean(absolute) + volumeRoot := filepath.Clean(filepath.VolumeName(absolute) + string(os.PathSeparator)) + root, err := os.OpenRoot(volumeRoot) + if err != nil { + return nil, nil, err + } + rootInfo, err := root.Stat(".") + if err != nil { + return nil, nil, errors.Join( + err, + wrapForceCloseError("absolute traversal root", operations.closeRootHandle("absolute traversal root", root)), + ) + } + if absolute == volumeRoot { + return root, rootInfo, nil + } + relative, err := filepath.Rel(volumeRoot, absolute) + if err != nil { + return nil, nil, errors.Join( + err, + wrapForceCloseError("absolute traversal root", operations.closeRootHandle("absolute traversal root", root)), + ) + } + for _, component := range strings.Split(relative, string(os.PathSeparator)) { + if component == "" || component == "." || component == ".." { + return nil, nil, errors.Join( + fmt.Errorf("unsafe path component %q", component), + wrapForceCloseError("absolute traversal root", operations.closeRootHandle("absolute traversal root", root)), + ) + } + info, err := root.Lstat(component) + if err != nil { + return nil, nil, errors.Join( + err, + wrapForceCloseError("absolute traversal root", operations.closeRootHandle("absolute traversal root", root)), + ) + } + if info.Mode()&os.ModeSymlink != 0 { + return nil, nil, errors.Join( + fmt.Errorf("path component %q is a symbolic link", component), + wrapForceCloseError("absolute traversal root", operations.closeRootHandle("absolute traversal root", root)), + ) + } + if !info.IsDir() { + return nil, nil, errors.Join( + fmt.Errorf("path component %q is not a directory", component), + wrapForceCloseError("absolute traversal root", operations.closeRootHandle("absolute traversal root", root)), + ) + } + child, err := root.OpenRoot(component) + if err != nil { + return nil, nil, errors.Join( + err, + wrapForceCloseError("absolute traversal root", operations.closeRootHandle("absolute traversal root", root)), + ) + } + childInfo, err := child.Stat(".") + if err != nil || !os.SameFile(info, childInfo) { + identityErr := err + if identityErr == nil { + identityErr = fmt.Errorf("physical directory identity changed") + } + return nil, nil, errors.Join( + fmt.Errorf("path component %q changed while being pinned: %w", component, identityErr), + wrapForceCloseError("absolute traversal child", operations.closeRootHandle("absolute traversal child", child)), + wrapForceCloseError("absolute traversal root", operations.closeRootHandle("absolute traversal root", root)), + ) + } + if err := operations.closeRootHandle("absolute traversal root", root); err != nil { + return nil, nil, errors.Join( + wrapForceCloseError("absolute traversal root", err), + wrapForceCloseError("absolute traversal child", operations.closeRootHandle("absolute traversal child", child)), + ) + } + root = child + rootInfo = childInfo + } + return root, rootInfo, nil +} + +func provePinnedAbsoluteDirectory( + path string, + expected fs.FileInfo, + operations forceReplaceOperations, +) error { + current, currentInfo, err := openPinnedAbsoluteDirectory(path, operations) + if err != nil { + return err + } + var proofErr error + if !os.SameFile(expected, currentInfo) { + proofErr = fmt.Errorf("physical directory identity changed") + } + return errors.Join( + proofErr, + wrapForceCloseError("absolute directory proof", operations.closeRootHandle("absolute directory proof", current)), + ) +} + +func rejectProtectedForceIdentity(candidate fs.FileInfo) error { + protectedPaths := []string{string(os.PathSeparator)} + if home, err := os.UserHomeDir(); err == nil { + protectedPaths = append(protectedPaths, home) + } + if repository, err := findRepositoryRoot(); err == nil { + protectedPaths = append(protectedPaths, repository) + } + for _, protected := range protectedPaths { + resolved := protected + if physical, err := filepath.EvalSymlinks(protected); err == nil { + resolved = physical + } + for { + info, err := os.Stat(resolved) + if err == nil && os.SameFile(candidate, info) { + return fmt.Errorf("unsafe dump replacement target: physical root, home, or repository ancestor is protected") + } + parent := filepath.Dir(resolved) + if parent == resolved { + break + } + resolved = parent + } + } + return nil +} + +func quarantineForceTarget(directory *os.File, name string) (string, error) { + for range 100 { + var random [16]byte + if _, err := rand.Read(random[:]); err != nil { + return "", fmt.Errorf("generate force quarantine name: %w", err) + } + quarantine := ".ret-force-" + hex.EncodeToString(random[:]) + ".preserved" + if err := forceRenameNoReplace(directory, name, quarantine); err == nil { + return quarantine, nil + } else if !errors.Is(err, os.ErrExist) { + return "", err + } + } + return "", fmt.Errorf("force quarantine name attempts exhausted") +} + +func restoreForceQuarantine( + directory *os.File, + parentPath string, + original string, + quarantine string, +) error { + if err := forceRenameNoReplace(directory, quarantine, original); err != nil { + originalPath := filepath.Join(parentPath, original) + quarantinePath := filepath.Join(parentPath, quarantine) + return fmt.Errorf( + "restore preserved prior collection %q to %q: %w; preserving both pathnames: prior collection at %q; competing destination at %q", + quarantinePath, + originalPath, + err, + quarantinePath, + originalPath, + ) + } + return nil +} + +func (s forceReplaceOperations) closeRootHandle(role string, root *os.Root) error { + if s.closeRoot != nil { + return s.closeRoot(role, root) + } + return root.Close() +} + +func (s forceReplaceOperations) closeFileHandle(role string, file *os.File) error { + if s.closeFile != nil { + return s.closeFile(role, file) + } + return file.Close() +} + +func wrapForceCloseError(description string, err error) error { + if err == nil { + return nil + } + return fmt.Errorf("close force replacement %s: %w", description, err) +} + +func findRepositoryRoot() (string, error) { + current, err := os.Getwd() + if err != nil { + return "", fmt.Errorf("get working directory: %w", err) + } + current, err = filepath.Abs(current) + if err != nil { + return "", fmt.Errorf("resolve working directory: %w", err) + } + for { + if _, err := os.Lstat(filepath.Join(current, ".git")); err == nil { + return current, nil + } else if !errors.Is(err, os.ErrNotExist) { + return "", fmt.Errorf("inspect repository root: %w", err) + } + parent := filepath.Dir(current) + if parent == current { + return "", fmt.Errorf("repository root not found") + } + current = parent + } +} diff --git a/cmd/retriever/force_rename_darwin.go b/cmd/retriever/force_rename_darwin.go new file mode 100644 index 00000000..806923ed --- /dev/null +++ b/cmd/retriever/force_rename_darwin.go @@ -0,0 +1,19 @@ +//go:build darwin + +package main + +import ( + "os" + + "golang.org/x/sys/unix" +) + +func forceRenameNoReplace(directory *os.File, oldName, newName string) error { + return unix.RenameatxNp( + int(directory.Fd()), + oldName, + int(directory.Fd()), + newName, + unix.RENAME_EXCL, + ) +} diff --git a/cmd/retriever/force_rename_linux.go b/cmd/retriever/force_rename_linux.go new file mode 100644 index 00000000..256e9d31 --- /dev/null +++ b/cmd/retriever/force_rename_linux.go @@ -0,0 +1,19 @@ +//go:build linux + +package main + +import ( + "os" + + "golang.org/x/sys/unix" +) + +func forceRenameNoReplace(directory *os.File, oldName, newName string) error { + return unix.Renameat2( + int(directory.Fd()), + oldName, + int(directory.Fd()), + newName, + unix.RENAME_NOREPLACE, + ) +} diff --git a/cmd/retriever/force_rename_other.go b/cmd/retriever/force_rename_other.go new file mode 100644 index 00000000..cf1824a2 --- /dev/null +++ b/cmd/retriever/force_rename_other.go @@ -0,0 +1,11 @@ +//go:build !linux && !darwin + +package main + +import ( + "os" +) + +func forceRenameNoReplace(_ *os.File, _, _ string) error { + return validateForcePlatform("unsupported") +} diff --git a/cmd/retriever/force_test.go b/cmd/retriever/force_test.go new file mode 100644 index 00000000..0f47fd06 --- /dev/null +++ b/cmd/retriever/force_test.go @@ -0,0 +1,339 @@ +package main + +import ( + "bytes" + "context" + "errors" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/specterops/dawgs/graph" + "github.com/specterops/dawgs/ret" +) + +func TestDumpForcePreservesCompletePriorCollectionAndCallsDumpAtOriginalPath(t *testing.T) { + parent := t.TempDir() + destination := filepath.Join(parent, "dump") + nested := filepath.Join(destination, "nested") + simulatedMount := filepath.Join(destination, "simulated-mount") + external := t.TempDir() + externalMarker := filepath.Join(external, "external") + sibling := filepath.Join(parent, "sibling") + if err := os.MkdirAll(nested, 0o755); err != nil { + t.Fatalf("mkdir nested prior collection: %v", err) + } + if err := os.Mkdir(simulatedMount, 0o755); err != nil { + t.Fatalf("mkdir simulated mount entry: %v", err) + } + regular := filepath.Join(nested, "regular") + if err := os.WriteFile(regular, []byte{0, 1, 2, 3, 255}, 0o640); err != nil { + t.Fatalf("write regular file: %v", err) + } + hardlinkSupported := true + if err := os.Link(regular, filepath.Join(nested, "hardlink")); err != nil { + hardlinkSupported = false + } + if err := os.WriteFile(filepath.Join(simulatedMount, "mounted-data"), []byte("mounted"), 0o600); err != nil { + t.Fatalf("write simulated mount data: %v", err) + } + if err := os.WriteFile(externalMarker, []byte("external"), 0o600); err != nil { + t.Fatalf("write external marker: %v", err) + } + if err := os.Symlink(external, filepath.Join(destination, "external-link")); err != nil { + t.Fatalf("create external symlink: %v", err) + } + if err := os.WriteFile(sibling, []byte("sibling"), 0o600); err != nil { + t.Fatalf("write sibling: %v", err) + } + + dumpCalled := false + runtime := newTestCommandRuntime(commandOperations{ + openDatabase: successfulTestDatabaseOpen, + dump: func(_ context.Context, _ graph.Database, config ret.DumpConfig) (ret.DumpResult, error) { + dumpCalled = true + if config.Directory != destination { + t.Fatalf("dump directory = %q, want original path %q", config.Directory, destination) + } + if _, err := os.Lstat(config.Directory); !errors.Is(err, os.ErrNotExist) { + t.Fatalf("original destination exists before Dump: %v", err) + } + if err := os.Mkdir(config.Directory, 0o755); err != nil { + return ret.DumpResult{}, err + } + return ret.DumpResult{}, os.WriteFile( + filepath.Join(config.Directory, "fresh"), + []byte("fresh"), + 0o600, + ) + }, + }) + + if err := runtime.run(context.Background(), []string{ + "dump", "-out", destination, "-force", "-graph", "asset", + }); err != nil { + t.Fatalf("dump force: %v", err) + } + if !dumpCalled { + t.Fatal("ret.Dump was not called") + } + + tombstone := requireSinglePreservedForceTombstone(t, parent) + if output := runtime.stderr.(*bytes.Buffer).String(); !strings.Contains(output, tombstone) { + t.Fatalf("force output %q does not report tombstone %q", output, tombstone) + } + requirePreservedPriorCollection(t, tombstone, external, hardlinkSupported) + if contents, err := os.ReadFile(externalMarker); err != nil || string(contents) != "external" { + t.Fatalf("external target changed: contents=%q err=%v", contents, err) + } + if contents, err := os.ReadFile(sibling); err != nil || string(contents) != "sibling" { + t.Fatalf("sibling changed: contents=%q err=%v", contents, err) + } + if contents, err := os.ReadFile(filepath.Join(destination, "fresh")); err != nil || string(contents) != "fresh" { + t.Fatalf("fresh dump output missing: contents=%q err=%v", contents, err) + } +} + +func TestDumpForceRejectsPostQuarantineReplacementBeforeExternalWork(t *testing.T) { + cases := []struct { + name string + create func(parent *os.Root, original string) error + require func(t *testing.T, destination string) + }{ + { + name: "regular file", + create: func(parent *os.Root, original string) error { + return parent.WriteFile(original, []byte("replacement-file"), 0o600) + }, + require: func(t *testing.T, destination string) { + t.Helper() + if contents, err := os.ReadFile(destination); err != nil || + string(contents) != "replacement-file" { + t.Fatalf("replacement file changed: contents=%q err=%v", contents, err) + } + }, + }, + { + name: "directory", + create: func(parent *os.Root, original string) error { + if err := parent.Mkdir(original, 0o755); err != nil { + return err + } + return parent.WriteFile(filepath.Join(original, "replacement"), []byte("replacement-dir"), 0o600) + }, + require: func(t *testing.T, destination string) { + t.Helper() + if contents, err := os.ReadFile(filepath.Join(destination, "replacement")); err != nil || + string(contents) != "replacement-dir" { + t.Fatalf("replacement directory changed: contents=%q err=%v", contents, err) + } + }, + }, + { + name: "symlink", + create: func(parent *os.Root, original string) error { + return parent.Symlink("replacement-target", original) + }, + require: func(t *testing.T, destination string) { + t.Helper() + if target, err := os.Readlink(destination); err != nil || target != "replacement-target" { + t.Fatalf("replacement symlink changed: target=%q err=%v", target, err) + } + }, + }, + } + + for _, test := range cases { + t.Run(test.name, func(t *testing.T) { + parent := t.TempDir() + destination := filepath.Join(parent, "dump") + if err := os.Mkdir(destination, 0o755); err != nil { + t.Fatalf("mkdir prior collection: %v", err) + } + if err := os.WriteFile(filepath.Join(destination, "prior"), []byte("prior"), 0o600); err != nil { + t.Fatalf("write prior marker: %v", err) + } + + databaseCalled := false + dumpCalled := false + runtime := newTestCommandRuntime(commandOperations{ + openDatabase: func(context.Context, databaseConfig) (graph.Database, string, error) { + databaseCalled = true + return nil, "", errors.New("database must not be opened") + }, + dump: func(context.Context, graph.Database, ret.DumpConfig) (ret.DumpResult, error) { + dumpCalled = true + return ret.DumpResult{}, errors.New("dump must not be called") + }, + }) + runtime.force.afterQuarantine = func(parent *os.Root, original, _ string) error { + return test.create(parent, original) + } + + err := runtime.run(context.Background(), []string{ + "dump", + "-out", destination, + "-force", + "-graph", "asset", + "-pprof-listen", "127.0.0.1:0", + }) + if err == nil || !strings.Contains(err.Error(), "original destination") { + t.Fatalf("dump error = %v, want original destination occupancy failure", err) + } + if databaseCalled { + t.Fatal("database was opened after post-quarantine replacement") + } + if dumpCalled { + t.Fatal("ret.Dump was called after post-quarantine replacement") + } + if output := runtime.stderr.(*bytes.Buffer).String(); strings.Contains(output, "pprof:") { + t.Fatalf("pprof started after post-quarantine replacement: %q", output) + } + + test.require(t, destination) + tombstone := requireSinglePreservedForceTombstone(t, parent) + if !strings.Contains(err.Error(), tombstone) { + t.Fatalf("dump error %q does not report preserved prior collection %q", err, tombstone) + } + if contents, readErr := os.ReadFile(filepath.Join(tombstone, "prior")); readErr != nil || + string(contents) != "prior" { + t.Fatalf("prior collection changed: contents=%q err=%v", contents, readErr) + } + }) + } +} + +func TestReplaceDumpDestinationRestoresAfterTargetRootCloseFailure(t *testing.T) { + testForceCloseFailureRestoration(t, "target", false) +} + +func TestReplaceDumpDestinationRestoresAfterParentRootCloseFailure(t *testing.T) { + testForceCloseFailureRestoration(t, "parent", false) +} + +func TestReplaceDumpDestinationReportsTombstoneAfterParentDirectoryCloseFailure(t *testing.T) { + testForceCloseFailureRestoration(t, "parent directory handle", true) +} + +func testForceCloseFailureRestoration(t *testing.T, failingRole string, expectTombstone bool) { + t.Helper() + parent := t.TempDir() + destination := filepath.Join(parent, "dump") + if err := os.Mkdir(destination, 0o755); err != nil { + t.Fatalf("mkdir destination: %v", err) + } + if err := os.WriteFile(filepath.Join(destination, "old"), []byte("old"), 0o600); err != nil { + t.Fatalf("write destination marker: %v", err) + } + + closeFailure := errors.New("injected close failure") + dumpCalled := false + runtime := newTestCommandRuntime(commandOperations{ + openDatabase: successfulTestDatabaseOpen, + dump: func(context.Context, graph.Database, ret.DumpConfig) (ret.DumpResult, error) { + dumpCalled = true + return ret.DumpResult{}, nil + }, + }) + runtime.force = forceReplaceOperations{ + closeRoot: func(role string, root *os.Root) error { + closeErr := root.Close() + if role == failingRole { + return errors.Join(closeErr, closeFailure) + } + return closeErr + }, + closeFile: func(role string, file *os.File) error { + closeErr := file.Close() + if role == failingRole { + return errors.Join(closeErr, closeFailure) + } + return closeErr + }, + } + err := runtime.run(context.Background(), []string{ + "dump", "-out", destination, "-force", "-graph", "asset", + }) + if !errors.Is(err, closeFailure) { + t.Fatalf("replace error = %v, want injected close failure", err) + } + if dumpCalled { + t.Fatal("ret.Dump was called after force close failure") + } + + tombstones, globErr := filepath.Glob(filepath.Join(parent, ".ret-force-*.preserved")) + if globErr != nil { + t.Fatalf("glob tombstones: %v", globErr) + } + if expectTombstone { + if len(tombstones) != 1 { + t.Fatalf("tombstones = %v, want one preserved tombstone", tombstones) + } + if !strings.Contains(err.Error(), tombstones[0]) { + t.Fatalf("close error %q does not report tombstone %q", err, tombstones[0]) + } + if contents, readErr := os.ReadFile(filepath.Join(tombstones[0], "old")); readErr != nil || + string(contents) != "old" { + t.Fatalf("preserved tombstone changed: contents=%q err=%v", contents, readErr) + } + if _, statErr := os.Lstat(destination); !errors.Is(statErr, os.ErrNotExist) { + t.Fatalf("destination exists after unrecoverable close failure: %v", statErr) + } + return + } + + if len(tombstones) != 0 { + t.Fatalf("unexpected tombstones after restoration: %v", tombstones) + } + if contents, readErr := os.ReadFile(filepath.Join(destination, "old")); readErr != nil || + string(contents) != "old" { + t.Fatalf("restored prior collection changed: contents=%q err=%v", contents, readErr) + } +} + +func requireSinglePreservedForceTombstone(t *testing.T, parent string) string { + t.Helper() + tombstones, err := filepath.Glob(filepath.Join(parent, ".ret-force-*.preserved")) + if err != nil { + t.Fatalf("glob preserved tombstones: %v", err) + } + if len(tombstones) != 1 { + t.Fatalf("preserved tombstones = %v, want one", tombstones) + } + return tombstones[0] +} + +func requirePreservedPriorCollection( + t *testing.T, + tombstone string, + external string, + hardlinkSupported bool, +) { + t.Helper() + regular := filepath.Join(tombstone, "nested", "regular") + if contents, err := os.ReadFile(regular); err != nil || + string(contents) != string([]byte{0, 1, 2, 3, 255}) { + t.Fatalf("regular file changed: contents=%v err=%v", contents, err) + } + if hardlinkSupported { + regularInfo, err := os.Stat(regular) + if err != nil { + t.Fatalf("stat regular file: %v", err) + } + hardlinkInfo, err := os.Stat(filepath.Join(tombstone, "nested", "hardlink")) + if err != nil { + t.Fatalf("stat hardlink: %v", err) + } + if !os.SameFile(regularInfo, hardlinkInfo) { + t.Fatal("hardlink identity was not preserved") + } + } + if contents, err := os.ReadFile(filepath.Join(tombstone, "simulated-mount", "mounted-data")); err != nil || + string(contents) != "mounted" { + t.Fatalf("simulated mount data changed: contents=%q err=%v", contents, err) + } + if target, err := os.Readlink(filepath.Join(tombstone, "external-link")); err != nil || target != external { + t.Fatalf("external symlink changed: target=%q err=%v", target, err) + } +} diff --git a/cmd/retriever/force_unsupported_test.go b/cmd/retriever/force_unsupported_test.go new file mode 100644 index 00000000..378a7fcf --- /dev/null +++ b/cmd/retriever/force_unsupported_test.go @@ -0,0 +1,33 @@ +//go:build !linux && !darwin + +package main + +import ( + "context" + "os" + "path/filepath" + "strings" + "testing" +) + +func TestDumpForceUnsupportedPlatformFailsBeforeMutation(t *testing.T) { + destination := filepath.Join(t.TempDir(), "dump") + if err := os.Mkdir(destination, 0o755); err != nil { + t.Fatalf("mkdir destination: %v", err) + } + marker := filepath.Join(destination, "keep") + if err := os.WriteFile(marker, []byte("keep"), 0o600); err != nil { + t.Fatalf("write marker: %v", err) + } + + runtime := newTestCommandRuntime(commandOperations{}) + err := runtime.run(context.Background(), []string{ + "dump", "-out", destination, "-force", "-graph", "asset", + }) + if err == nil || !strings.Contains(err.Error(), "unsupported") { + t.Fatalf("dump error = %v, want unsupported-platform failure", err) + } + if contents, err := os.ReadFile(marker); err != nil || string(contents) != "keep" { + t.Fatalf("unsupported force changed marker: contents=%q err=%v", contents, err) + } +} diff --git a/cmd/retriever/main.go b/cmd/retriever/main.go index 1b5ca896..89dfd954 100644 --- a/cmd/retriever/main.go +++ b/cmd/retriever/main.go @@ -2,31 +2,79 @@ package main import ( "context" - "crypto/hpke" "encoding/json" + "errors" "flag" "fmt" "io" + "log/slog" "os" "strings" - "github.com/specterops/dawgs/retriever" + "github.com/specterops/dawgs/graph" + "github.com/specterops/dawgs/ret" + "github.com/specterops/dawgs/ret/archive" + "github.com/specterops/dawgs/ret/observe" ) const usage = `usage: retriever [options] Commands: - keygen Generate an HPKE recipient key pair for encrypted archives. - dump Dump live Dawgs graph data into a manifest-based collection. - unpack Decrypt and unpack an encrypted retriever archive. - load Load a manifest-based collection into a Dawgs graph database. - verify Verify loaded graph metrics against a dump manifest. - bench Benchmark read throughput for dump planning. + dump Dump live Dawgs graph data into a local collection. + load Load a local JSONL collection into a Dawgs database. + verify-collection Verify every artifact in a local collection. + verify-database Verify database metrics against a collection manifest. + pack Create an encrypted archive from a verified collection. + unpack Decrypt, verify, and publish a collection. + keygen Generate an archive recipient key pair. + bench Benchmark read throughput for dump planning. ` +type commandOperations struct { + openDatabase func(context.Context, databaseConfig) (graph.Database, string, error) + dump func(context.Context, graph.Database, ret.DumpConfig) (ret.DumpResult, error) + load func(context.Context, graph.Database, ret.LoadConfig) (ret.LoadResult, error) + verifyCollection func(context.Context, ret.VerifyCollectionConfig) (ret.VerifyCollectionResult, error) + verifyDatabase func(context.Context, graph.Database, ret.VerifyDatabaseConfig) (ret.VerifyDatabaseResult, error) + pack func(context.Context, ret.PackConfig) error + unpack func(context.Context, ret.UnpackConfig) error + keygen func(ret.KeygenConfig) error +} + +func (s commandOperations) withDefaults() commandOperations { + if s.openDatabase == nil { + s.openDatabase = openDatabase + } + if s.dump == nil { + s.dump = ret.Dump + } + if s.load == nil { + s.load = ret.Load + } + if s.verifyCollection == nil { + s.verifyCollection = ret.VerifyCollection + } + if s.verifyDatabase == nil { + s.verifyDatabase = ret.VerifyDatabase + } + if s.pack == nil { + s.pack = ret.Pack + } + if s.unpack == nil { + s.unpack = ret.Unpack + } + if s.keygen == nil { + s.keygen = ret.Keygen + } + return s +} + type commandRuntime struct { - stdout io.Writer - stderr io.Writer + stdout io.Writer + stderr io.Writer + operations commandOperations + observer observe.Observer + force forceReplaceOperations } func main() { @@ -34,7 +82,6 @@ func main() { stdout: os.Stdout, stderr: os.Stderr, } - if err := runtime.run(context.Background(), os.Args[1:]); err != nil { fmt.Fprintf(os.Stderr, "retriever: %v\n", err) os.Exit(1) @@ -51,16 +98,20 @@ func (s commandRuntime) run(ctx context.Context, args []string) error { case "help", "-h", "--help": fmt.Fprint(s.stdout, usage) return nil - case "keygen": - return s.runKeygen(args[1:]) case "dump": return s.runDump(ctx, args[1:]) - case "unpack": - return s.runUnpack(args[1:]) case "load": return s.runLoad(ctx, args[1:]) - case "verify": - return s.runVerify(ctx, args[1:]) + case "verify-collection": + return s.runVerifyCollection(ctx, args[1:]) + case "verify-database": + return s.runVerifyDatabase(ctx, args[1:]) + case "pack": + return s.runPack(ctx, args[1:]) + case "unpack": + return s.runUnpack(ctx, args[1:]) + case "keygen": + return s.runKeygen(args[1:]) case "bench": return s.runBench(ctx, args[1:]) default: @@ -70,381 +121,382 @@ func (s commandRuntime) run(ctx context.Context, args []string) error { } func (s commandRuntime) runDump(ctx context.Context, args []string) error { - var ( - dbCfg databaseConfig - cfg = retriever.DefaultDumpOptions("") - - graphs stringList - scrubValue string - compressionVal string - archiveOut string - recipientPath string - scrubConfig string - pprofListen string - ) - - flags := flag.NewFlagSet("retriever dump", flag.ContinueOnError) - flags.SetOutput(s.stderr) - commonDatabaseFlags(flags, &dbCfg) - flags.Var(&graphs, "graph", "Graph target. May be repeated.") - allGraphs := flags.Bool("all-graphs", false, "Dump every graph discoverable by the selected driver.") - flags.StringVar(&cfg.OutputDir, "out", "", "Output collection directory.") - flags.BoolVar(&cfg.Force, "force", false, "Replace an existing non-empty output directory.") - flags.BoolVar(&cfg.Resume, "resume", false, "Resume an interrupted dump from its validated checkpoint.") - flags.StringVar(&archiveOut, "archive-out", "", "Optional encrypted archive output path.") - flags.StringVar(&recipientPath, "recipient", "", "Recipient public key for -archive-out.") - flags.StringVar(&scrubValue, "scrub", string(cfg.Scrub), "Scrub mode: none or full.") - flags.StringVar(&cfg.Salt, "salt", "", "Scrub salt. Overrides RETRIEVER_SCRUB_SALT and is never written.") - flags.StringVar(&scrubConfig, "config", "", "Optional retriever TOML config for scrub classifier settings.") - flags.StringVar(&compressionVal, "compression", string(cfg.Compression), "Compression codec: zstd, gzip, or none.") - flags.IntVar(&cfg.ZstdLevel, "zstd-level", cfg.ZstdLevel, "zstd compression level.") - flags.IntVar(&cfg.ShardSize, "shard-size", cfg.ShardSize, "Maximum entities per fragment.") - flags.IntVar(&cfg.BatchSize, "batch-size", cfg.BatchSize, "Database read batch size.") - commonPprofFlag(flags, &pprofListen) - if err := flags.Parse(args); err != nil { + command, err := parseDumpCommand(args, s.stderr) + if err != nil { return err } - - fillConnectionFromEnv(&dbCfg) - - cfg.Scrub = retriever.ScrubMode(strings.TrimSpace(scrubValue)) - cfg.Compression = retriever.CompressionCodec(strings.TrimSpace(compressionVal)) - - if strings.TrimSpace(cfg.Salt) == "" { - cfg.Salt = strings.TrimSpace(os.Getenv("RETRIEVER_SCRUB_SALT")) - if cfg.Salt == "" { - cfg.Salt = strings.TrimSpace(os.Getenv("RETRIEVR_SCRUB_SALT")) + if err := validateGraphSelection(command.graphs, command.allGraphs); err != nil { + return err + } + if address := strings.TrimSpace(command.pprof); address != "" { + if err := validatePprofListenAddress(address); err != nil { + return err } } - - if strings.TrimSpace(scrubConfig) != "" { - file, err := os.Open(scrubConfig) + if command.force { + replacement, err := replaceDumpDestination(command.dump.Directory, s.force) if err != nil { - return fmt.Errorf("open scrub config: %w", err) + return err + } + command.dump.Directory = replacement.destination + if replacement.tombstone != "" { + fmt.Fprintf(s.stderr, "force: previous destination preserved intact at %s\n", replacement.tombstone) } - defer file.Close() - - cfg.ScrubConfig = file } - if err := cfg.Validate(); err != nil { + profileServer, err := startPprofServer(command.pprof, s.stderr) + if err != nil { return err } + defer stopPprofServer(profileServer, s.stderr) - var archiveRecipient hpke.PublicKey - if strings.TrimSpace(archiveOut) != "" { - if strings.TrimSpace(recipientPath) == "" { - return fmt.Errorf("-archive-out requires -recipient") + operations := s.operations.withDefaults() + var result ret.DumpResult + if err := func() (resultErr error) { + database, driverName, err := operations.openDatabase(ctx, command.database) + if err != nil { + return err } + defer func() { + resultErr = errors.Join(resultErr, closeProductDatabase(database)) + }() - var err error - archiveRecipient, err = retriever.LoadArchivePublicKey(recipientPath) + graphs, err := resolveGraphNames(ctx, database, driverName, command.graphs, command.allGraphs) if err != nil { return err } - - if err := retriever.PreflightArchiveOutputPath(archiveOut); err != nil { + command.dump.Graphs = graphs + command.dump.Observer = s.commandObserver() + if err := command.dump.Validate(); err != nil { return err } - } else if strings.TrimSpace(recipientPath) != "" { - return fmt.Errorf("-recipient requires -archive-out") - } - - profileServer, err := startPprofServer(pprofListen, s.stderr) - if err != nil { + result, err = operations.dump(ctx, database, command.dump) + return err + }(); err != nil { return err } - defer stopPprofServer(profileServer, s.stderr) + fmt.Fprintf(s.stdout, + "dumped %d graph(s)\nmanifest: %s\nnodes: %d\nrelationships: %d\n", + result.GraphCount, + result.ManifestPath, + result.NodeCount, + result.RelationshipCount, + ) + return nil +} - db, driverName, err := openDatabase(ctx, dbCfg) - if err != nil { +func (s commandRuntime) runLoad(ctx context.Context, args []string) error { + var ( + databaseConfig databaseConfig + config = ret.LoadConfig{BatchSize: defaultEntityBatchSize} + verify bool + pprofListen string + ) + flags := flag.NewFlagSet("retriever load", flag.ContinueOnError) + flags.SetOutput(s.stderr) + commonDatabaseFlags(flags, &databaseConfig) + flags.StringVar(&config.Directory, "in", "", "Input collection directory.") + flags.IntVar(&config.BatchSize, "batch-size", config.BatchSize, "Database write batch size.") + flags.BoolVar(&verify, "verify-database", false, "Verify database metrics after a successful load.") + commonPprofFlag(flags, &pprofListen) + if err := flags.Parse(args); err != nil { return err } - defer db.Close(ctx) - - targets, err := resolveGraphTargets(ctx, db, driverName, []string(graphs), *allGraphs) - if err != nil { + fillConnectionFromEnv(&databaseConfig) + config.Directory = strings.TrimSpace(config.Directory) + config.Observer = s.commandObserver() + if err := config.Validate(); err != nil { return err } - result, err := retriever.Dump(ctx, db, driverName, targets, cfg) + profileServer, err := startPprofServer(pprofListen, s.stderr) if err != nil { return err } + defer stopPprofServer(profileServer, s.stderr) - var archiveLine string - if strings.TrimSpace(archiveOut) != "" { - if err := retriever.WriteEncryptedCollectionArchiveFile(cfg.OutputDir, archiveOut, archiveRecipient); err != nil { + operations := s.operations.withDefaults() + var ( + result ret.LoadResult + verifyResult ret.VerifyDatabaseResult + ) + if err := func() (resultErr error) { + database, _, err := operations.openDatabase(ctx, databaseConfig) + if err != nil { return err } + defer func() { + resultErr = errors.Join(resultErr, closeProductDatabase(database)) + }() - archiveLine = fmt.Sprintf("archive: %s\n", archiveOut) + result, err = operations.load(ctx, database, config) + if err != nil { + return fmt.Errorf("load failed; if a graph was partially loaded, clear it before retry: %w", err) + } + if verify { + verifyResult, err = operations.verifyDatabase(ctx, database, ret.VerifyDatabaseConfig{ + Directory: config.Directory, + BatchSize: config.BatchSize, + Observer: config.Observer, + }) + if err != nil { + return err + } + } + return nil + }(); err != nil { + return err } + fmt.Fprintf(s.stdout, + "loaded %d graph(s)\nnodes: %d\nrelationships: %d\n", + result.GraphCount, + result.NodeCount, + result.RelationshipCount, + ) - fmt.Fprintf(s.stdout, "dumped %d graph(s)\nmanifest: %s\n%snodes: %d\nrelationships: %d\n", len(result.Manifest.Graphs), result.ManifestPath, archiveLine, result.NodeCount, result.EdgeCount) + if verify { + fmt.Fprintf(s.stdout, + "verified database: %d graph(s)\nnodes: %d\nrelationships: %d\n", + verifyResult.GraphCount, + verifyResult.NodeCount, + verifyResult.RelationshipCount, + ) + } return nil } -func (s commandRuntime) runKeygen(args []string) error { - var cfg retriever.KeygenOptions - - flags := flag.NewFlagSet("retriever keygen", flag.ContinueOnError) +func (s commandRuntime) runVerifyCollection(ctx context.Context, args []string) error { + var config ret.VerifyCollectionConfig + flags := flag.NewFlagSet("retriever verify-collection", flag.ContinueOnError) flags.SetOutput(s.stderr) - flags.StringVar(&cfg.PrivatePath, "private", "", "Private key output path.") - flags.StringVar(&cfg.PublicPath, "public", "", "Public key output path.") + flags.StringVar(&config.Directory, "in", "", "Input collection directory.") if err := flags.Parse(args); err != nil { return err } - - if err := cfg.Validate(); err != nil { + config.Directory = strings.TrimSpace(config.Directory) + config.Observer = s.commandObserver() + if err := config.Validate(); err != nil { return err } - if err := retriever.Keygen(cfg); err != nil { + result, err := s.operations.withDefaults().verifyCollection(ctx, config) + if err != nil { return err } - - fmt.Fprintf(s.stdout, "private key: %s\npublic key: %s\n", cfg.PrivatePath, cfg.PublicPath) + fmt.Fprintf(s.stdout, + "verified collection: %d graph(s)\nnodes: %d\nrelationships: %d\n", + result.GraphCount, + result.NodeCount, + result.RelationshipCount, + ) return nil } -func (s commandRuntime) runUnpack(args []string) error { +func (s commandRuntime) runVerifyDatabase(ctx context.Context, args []string) error { var ( - archivePath string - identityPath string - outputDir string - force bool + databaseConfig databaseConfig + config = ret.VerifyDatabaseConfig{BatchSize: defaultEntityBatchSize} + pprofListen string ) - - flags := flag.NewFlagSet("retriever unpack", flag.ContinueOnError) + flags := flag.NewFlagSet("retriever verify-database", flag.ContinueOnError) flags.SetOutput(s.stderr) - flags.StringVar(&archivePath, "archive", "", "Encrypted archive input path.") - flags.StringVar(&identityPath, "identity", "", "Recipient private key path.") - flags.StringVar(&outputDir, "out", "", "Output collection directory.") - flags.BoolVar(&force, "force", false, "Replace an existing non-empty output directory.") + commonDatabaseFlags(flags, &databaseConfig) + flags.StringVar(&config.Directory, "in", "", "Input collection directory.") + flags.IntVar(&config.BatchSize, "batch-size", config.BatchSize, "Database read batch size.") + commonPprofFlag(flags, &pprofListen) if err := flags.Parse(args); err != nil { return err } - - if strings.TrimSpace(archivePath) == "" { - return fmt.Errorf("archive path is required; pass -archive") - } - - if strings.TrimSpace(identityPath) == "" { - return fmt.Errorf("identity key path is required; pass -identity") - } - - if strings.TrimSpace(outputDir) == "" { - return fmt.Errorf("output directory is required; pass -out") + fillConnectionFromEnv(&databaseConfig) + config.Directory = strings.TrimSpace(config.Directory) + config.Observer = s.commandObserver() + if err := config.Validate(); err != nil { + return err } - identity, err := retriever.LoadArchivePrivateKey(identityPath) + profileServer, err := startPprofServer(pprofListen, s.stderr) if err != nil { return err } + defer stopPprofServer(profileServer, s.stderr) - if err := retriever.UnpackEncryptedCollectionArchiveFile(archivePath, outputDir, force, identity); err != nil { + operations := s.operations.withDefaults() + var result ret.VerifyDatabaseResult + if err := func() (resultErr error) { + database, _, err := operations.openDatabase(ctx, databaseConfig) + if err != nil { + return err + } + defer func() { + resultErr = errors.Join(resultErr, closeProductDatabase(database)) + }() + result, err = operations.verifyDatabase(ctx, database, config) + return err + }(); err != nil { return err } - - fmt.Fprintf(s.stdout, "unpacked archive: %s\noutput: %s\n", archivePath, outputDir) + fmt.Fprintf(s.stdout, + "verified database: %d graph(s)\nnodes: %d\nrelationships: %d\n", + result.GraphCount, + result.NodeCount, + result.RelationshipCount, + ) return nil } -func (s commandRuntime) runLoad(ctx context.Context, args []string) error { +func (s commandRuntime) runPack(ctx context.Context, args []string) error { var ( - dbCfg databaseConfig - cfg = retriever.DefaultLoadOptions("") - inputDir string - archivePath string - identityPath string - pprofListen string + collectionDirectory string + archivePath string + recipientPath string ) - - flags := flag.NewFlagSet("retriever load", flag.ContinueOnError) + flags := flag.NewFlagSet("retriever pack", flag.ContinueOnError) flags.SetOutput(s.stderr) - commonDatabaseFlags(flags, &dbCfg) - flags.StringVar(&inputDir, "in", "", "Input collection directory.") - flags.StringVar(&archivePath, "archive", "", "Encrypted archive input path.") - flags.StringVar(&identityPath, "identity", "", "Recipient private key path for -archive.") - flags.IntVar(&cfg.BatchSize, "batch-size", cfg.BatchSize, "Database write batch size.") - flags.BoolVar(&cfg.VerifyMetrics, "verify-metrics", false, "Verify loaded graph metrics against the dump manifest after load.") - commonPprofFlag(flags, &pprofListen) + flags.StringVar(&collectionDirectory, "in", "", "Input collection directory.") + flags.StringVar(&archivePath, "archive", "", "Encrypted archive output path.") + flags.StringVar(&recipientPath, "recipient", "", "Recipient public key path.") if err := flags.Parse(args); err != nil { return err } - - fillConnectionFromEnv(&dbCfg) - cfg.InputDir = strings.TrimSpace(inputDir) - - if strings.TrimSpace(archivePath) != "" { - if cfg.InputDir != "" { - return fmt.Errorf("load accepts either -in or -archive, not both") - } - - if strings.TrimSpace(identityPath) == "" { - return fmt.Errorf("-archive requires -identity") - } - - identity, err := retriever.LoadArchivePrivateKey(identityPath) - if err != nil { - return err - } - - archiveFile, err := os.Open(archivePath) - if err != nil { - return fmt.Errorf("open archive: %w", err) - } - defer archiveFile.Close() - - cfg.ArchiveReader = archiveFile - cfg.ArchiveIdentity = identity - } else if strings.TrimSpace(identityPath) != "" { - return fmt.Errorf("-identity requires -archive") + if strings.TrimSpace(collectionDirectory) == "" { + return fmt.Errorf("collection directory is required; pass -in") } - - if err := cfg.Validate(); err != nil { - return err + if strings.TrimSpace(archivePath) == "" { + return fmt.Errorf("archive path is required; pass -archive") } - - profileServer, err := startPprofServer(pprofListen, s.stderr) + if strings.TrimSpace(recipientPath) == "" { + return fmt.Errorf("recipient key path is required; pass -recipient") + } + recipient, err := archive.ReadPublicKey(recipientPath) if err != nil { return err } - defer stopPprofServer(profileServer, s.stderr) - - db, driverName, err := openDatabase(ctx, dbCfg) - if err != nil { + config := ret.PackConfig{ + CollectionDirectory: strings.TrimSpace(collectionDirectory), + ArchivePath: strings.TrimSpace(archivePath), + Recipient: recipient, + Observer: s.commandObserver(), + } + if err := config.Validate(); err != nil { return err } - defer db.Close(ctx) - - result, err := retriever.Load(ctx, db, driverName, cfg) - if err != nil { + if err := s.operations.withDefaults().pack(ctx, config); err != nil { return err } - - fmt.Fprintf(s.stdout, "loaded %d graph(s)\nnodes: %d\nrelationships: %d\n", result.GraphCount, result.NodeCount, result.EdgeCount) + fmt.Fprintf(s.stdout, "archive: %s\n", config.ArchivePath) return nil } -func (s commandRuntime) runVerify(ctx context.Context, args []string) error { +func (s commandRuntime) runUnpack(ctx context.Context, args []string) error { var ( - dbCfg databaseConfig - cfg = retriever.DefaultVerifyOptions("") - pprofListen string + archivePath string + outputDir string + identityPath string ) - - flags := flag.NewFlagSet("retriever verify", flag.ContinueOnError) + flags := flag.NewFlagSet("retriever unpack", flag.ContinueOnError) flags.SetOutput(s.stderr) - commonDatabaseFlags(flags, &dbCfg) - flags.StringVar(&cfg.InputDir, "in", "", "Input collection directory.") - flags.IntVar(&cfg.BatchSize, "batch-size", cfg.BatchSize, "Database read batch size.") - commonPprofFlag(flags, &pprofListen) + flags.StringVar(&archivePath, "archive", "", "Encrypted archive input path.") + flags.StringVar(&outputDir, "out", "", "Output collection directory.") + flags.StringVar(&identityPath, "identity", "", "Recipient private key path.") if err := flags.Parse(args); err != nil { return err } - - fillConnectionFromEnv(&dbCfg) - - if err := cfg.Validate(); err != nil { - return err + if strings.TrimSpace(archivePath) == "" { + return fmt.Errorf("archive path is required; pass -archive") } - - profileServer, err := startPprofServer(pprofListen, s.stderr) + if strings.TrimSpace(outputDir) == "" { + return fmt.Errorf("output directory is required; pass -out") + } + if strings.TrimSpace(identityPath) == "" { + return fmt.Errorf("identity key path is required; pass -identity") + } + identity, err := archive.ReadPrivateKey(identityPath) if err != nil { return err } - defer stopPprofServer(profileServer, s.stderr) - - db, driverName, err := openDatabase(ctx, dbCfg) - if err != nil { + config := ret.UnpackConfig{ + ArchivePath: strings.TrimSpace(archivePath), + OutputDirectory: strings.TrimSpace(outputDir), + Identity: identity, + Observer: s.commandObserver(), + } + if err := config.Validate(); err != nil { return err } - defer db.Close(ctx) - - result, err := retriever.Verify(ctx, db, driverName, cfg) - if err != nil { + if err := s.operations.withDefaults().unpack(ctx, config); err != nil { return err } - - fmt.Fprintf(s.stdout, "verified %d graph(s)\nnodes: %d\nrelationships: %d\n", result.GraphCount, result.NodeCount, result.EdgeCount) + fmt.Fprintf(s.stdout, "unpacked archive: %s\noutput: %s\n", config.ArchivePath, config.OutputDirectory) return nil } -func (s commandRuntime) runBench(ctx context.Context, args []string) error { - var ( - dbCfg databaseConfig - cfg benchOptions - - graphs stringList - workers workerList - compressionVal string - pprofListen string - ) - - cfg.BatchSize = retriever.DefaultBatchSize - cfg.SampleSize = defaultBenchSampleSize - cfg.ZstdLevel = retriever.DefaultZstdLevel - - flags := flag.NewFlagSet("retriever bench", flag.ContinueOnError) +func (s commandRuntime) runKeygen(args []string) error { + var config ret.KeygenConfig + flags := flag.NewFlagSet("retriever keygen", flag.ContinueOnError) flags.SetOutput(s.stderr) - commonDatabaseFlags(flags, &dbCfg) - flags.Var(&graphs, "graph", "Graph target. May be repeated.") - allGraphs := flags.Bool("all-graphs", false, "Benchmark every graph discoverable by the selected driver.") - flags.Var(&workers, "workers", "Comma-separated worker counts.") - flags.IntVar(&cfg.BatchSize, "batch-size", cfg.BatchSize, "Database read batch size.") - flags.IntVar(&cfg.SampleSize, "sample-size", cfg.SampleSize, "Maximum nodes and relationships to scan per phase; 0 scans the full graph.") - flags.StringVar(&compressionVal, "compression", "", "Optional compression codec to include encode/compress timing: zstd, gzip, or none.") - flags.IntVar(&cfg.ZstdLevel, "zstd-level", cfg.ZstdLevel, "zstd compression level.") - flags.BoolVar(&cfg.JSONOutput, "json", false, "Emit machine-readable JSON.") - commonPprofFlag(flags, &pprofListen) + flags.StringVar(&config.PrivateKeyPath, "private-key", "", "Private key output path.") + flags.StringVar(&config.PublicKeyPath, "public-key", "", "Public key output path.") if err := flags.Parse(args); err != nil { return err } - - fillConnectionFromEnv(&dbCfg) - - if len(workers) == 0 { - workers = workerList{1} - } - - cfg.Workers = []int(workers) - cfg.Compression = retriever.CompressionCodec(strings.TrimSpace(compressionVal)) - - if err := cfg.validate(); err != nil { + if err := config.Validate(); err != nil { return err } - - profileServer, err := startPprofServer(pprofListen, s.stderr) - if err != nil { + if err := s.operations.withDefaults().keygen(config); err != nil { return err } - defer stopPprofServer(profileServer, s.stderr) + fmt.Fprintf(s.stdout, "private key: %s\npublic key: %s\n", config.PrivateKeyPath, config.PublicKeyPath) + return nil +} - db, driverName, err := openDatabase(ctx, dbCfg) +func (s commandRuntime) runBench(ctx context.Context, args []string) error { + config, err := parseBenchCommand(args, s.stderr) if err != nil { return err } - defer db.Close(ctx) - targets, err := resolveGraphTargets(ctx, db, driverName, []string(graphs), *allGraphs) + profileServer, err := startPprofServer(config.pprof, s.stderr) if err != nil { return err } + defer stopPprofServer(profileServer, s.stderr) - report, err := Bench(ctx, db, driverName, targets, cfg) - if err != nil { + operations := s.operations.withDefaults() + var report benchReport + if err := func() (resultErr error) { + database, driverName, err := operations.openDatabase(ctx, config.database) + if err != nil { + return err + } + defer func() { + resultErr = errors.Join(resultErr, closeProductDatabase(database)) + }() + + graphNames, err := resolveGraphNames(ctx, database, driverName, config.graphs, config.allGraphs) + if err != nil { + return err + } + + report, err = Bench(ctx, database, driverName, graphNames, config.bench) + if err != nil { + return err + } + return nil + }(); err != nil { return err } - - if cfg.JSONOutput { + if config.bench.JSONOutput { encoder := json.NewEncoder(s.stdout) encoder.SetIndent("", " ") - return encoder.Encode(report) } - writeBenchReport(s.stdout, report) return nil } + +func (s commandRuntime) commandObserver() observe.Observer { + if s.observer != nil { + return s.observer + } + return newCommandObserver(slog.Default()) +} diff --git a/cmd/retriever/main_test.go b/cmd/retriever/main_test.go index 7e265278..a9b6f27b 100644 --- a/cmd/retriever/main_test.go +++ b/cmd/retriever/main_test.go @@ -3,15 +3,18 @@ package main import ( "bytes" "context" + "errors" "os" "path/filepath" + "reflect" "strings" "testing" - "github.com/specterops/dawgs/retriever" + "github.com/specterops/dawgs/graph" + "github.com/specterops/dawgs/ret" ) -func TestCommandRuntimeHelpAndValidation(t *testing.T) { +func TestCommandRuntimeHelpListsProductCommands(t *testing.T) { runtime := commandRuntime{ stdout: &bytes.Buffer{}, stderr: &bytes.Buffer{}, @@ -20,74 +23,899 @@ func TestCommandRuntimeHelpAndValidation(t *testing.T) { t.Fatalf("help: %v", err) } - helpOutput := runtime.stdout.(*bytes.Buffer).String() - - for _, command := range []string{"keygen", "dump", "unpack", "load", "verify", "bench"} { - if !strings.Contains(helpOutput, command) { + output := runtime.stdout.(*bytes.Buffer).String() + for _, command := range []string{ + "dump", + "load", + "verify-collection", + "verify-database", + "pack", + "unpack", + "keygen", + "bench", + } { + if !strings.Contains(output, command) { t.Fatalf("help output missing %s command", command) } } + if strings.Contains(output, "\n verify ") { + t.Fatalf("help output retained ambiguous verify command:\n%s", output) + } +} - err := runtime.run(context.Background(), []string{"unknown"}) - if err == nil || !strings.Contains(err.Error(), "unknown command") { - t.Fatalf("expected unknown command error, got %v", err) +func TestProductCommandsValidateRequiredPathsBeforeExternalWork(t *testing.T) { + cases := []struct { + command string + want string + }{ + {command: "dump", want: "output directory"}, + {command: "load", want: "load directory"}, + {command: "verify-collection", want: "collection directory"}, + {command: "verify-database", want: "database verification directory"}, + {command: "pack", want: "collection directory"}, + {command: "unpack", want: "archive path"}, + {command: "keygen", want: "private"}, } + for _, test := range cases { + t.Run(test.command, func(t *testing.T) { + runtime := newTestCommandRuntime(commandOperations{}) + err := runtime.run(context.Background(), []string{test.command}) + if err == nil || !strings.Contains(err.Error(), test.want) { + t.Fatalf("%s error = %v, want containing %q", test.command, err, test.want) + } + }) + } +} - err = runtime.run(context.Background(), []string{"dump", "-out", t.TempDir(), "-scrub", "full"}) - if err == nil || !strings.Contains(err.Error(), "-scrub full requires") { - t.Fatalf("expected scrub salt validation error, got %v", err) +func TestProductCommandsRejectRemovedFlagBleed(t *testing.T) { + cases := [][]string{ + {"dump", "-archive-out", "dump.tar.enc"}, + {"dump", "-recipient", "public.key"}, + {"dump", "-workers", "2"}, + {"load", "-archive", "dump.tar.enc"}, + {"load", "-identity", "private.key"}, + {"load", "-workers", "2"}, + {"unpack", "-force"}, + } + for _, args := range cases { + t.Run(strings.Join(args, "_"), func(t *testing.T) { + runtime := newTestCommandRuntime(commandOperations{}) + err := runtime.run(context.Background(), args) + if err == nil || !strings.Contains(err.Error(), "flag provided but not defined") { + t.Fatalf("run(%v) error = %v, want undefined flag", args, err) + } + }) } +} + +func TestVerifyCollectionDoesNotOpenDatabase(t *testing.T) { + var opened int + runtime := newTestCommandRuntime(commandOperations{ + openDatabase: func(context.Context, databaseConfig) (graph.Database, string, error) { + opened++ + return nil, "", errors.New("database must not open") + }, + verifyCollection: func(_ context.Context, config ret.VerifyCollectionConfig) (ret.VerifyCollectionResult, error) { + if config.Directory != "collection" { + t.Fatalf("verify collection directory = %q", config.Directory) + } + return ret.VerifyCollectionResult{GraphCount: 2, NodeCount: 3, RelationshipCount: 4}, nil + }, + }) - err = runtime.run(context.Background(), []string{"load"}) - if err == nil || !strings.Contains(err.Error(), "input directory or archive reader is required") { - t.Fatalf("expected load input validation error, got %v", err) + if err := runtime.run(context.Background(), []string{"verify-collection", "-in", "collection"}); err != nil { + t.Fatalf("verify collection: %v", err) + } + if opened != 0 { + t.Fatalf("database open calls = %d, want 0", opened) } +} + +func TestVerifyDatabaseOpensDatabaseAndCallsFacade(t *testing.T) { + var opened, verified int + runtime := newTestCommandRuntime(commandOperations{ + openDatabase: func(context.Context, databaseConfig) (graph.Database, string, error) { + opened++ + return nil, "pg", nil + }, + verifyDatabase: func(_ context.Context, database graph.Database, config ret.VerifyDatabaseConfig) (ret.VerifyDatabaseResult, error) { + verified++ + if database != nil { + t.Fatal("test opener should supply nil database") + } + if config.Directory != "collection" || config.BatchSize != defaultEntityBatchSize { + t.Fatalf("verify database config = %+v", config) + } + return ret.VerifyDatabaseResult{}, nil + }, + }) - err = runtime.run(context.Background(), []string{"keygen"}) - if err == nil || !strings.Contains(err.Error(), "private key path is required") { - t.Fatalf("expected keygen validation error, got %v", err) + if err := runtime.run(context.Background(), []string{ + "verify-database", + "-in", "collection", + "-connection", "postgresql://example/database", + }); err != nil { + t.Fatalf("verify database: %v", err) } + if opened != 1 || verified != 1 { + t.Fatalf("open calls = %d, verify calls = %d; want 1 each", opened, verified) + } +} + +func TestLoadVerifyDatabaseIsASecondFacadeCall(t *testing.T) { + var calls []string + runtime := newTestCommandRuntime(commandOperations{ + openDatabase: successfulTestDatabaseOpen, + load: func(_ context.Context, _ graph.Database, config ret.LoadConfig) (ret.LoadResult, error) { + calls = append(calls, "load") + if config.Directory != "collection" || config.BatchSize != 77 { + t.Fatalf("load config = %+v", config) + } + return ret.LoadResult{}, nil + }, + verifyDatabase: func(_ context.Context, _ graph.Database, config ret.VerifyDatabaseConfig) (ret.VerifyDatabaseResult, error) { + calls = append(calls, "verify-database") + if config.Directory != "collection" || config.BatchSize != 77 { + t.Fatalf("verify database config = %+v", config) + } + return ret.VerifyDatabaseResult{}, nil + }, + }) - err = runtime.run(context.Background(), []string{"unpack"}) - if err == nil || !strings.Contains(err.Error(), "archive path is required") { - t.Fatalf("expected unpack validation error, got %v", err) + if err := runtime.run(context.Background(), []string{ + "load", + "-in", "collection", + "-batch-size", "77", + "-verify-database", + "-connection", "postgresql://example/database", + }); err != nil { + t.Fatalf("load: %v", err) + } + if !reflect.DeepEqual(calls, []string{"load", "verify-database"}) { + t.Fatalf("operation calls = %v", calls) + } +} + +func TestLoadFailureWarnsToClearBeforeRetryWithoutContinuing(t *testing.T) { + var calls []string + runtime := newTestCommandRuntime(commandOperations{ + openDatabase: successfulTestDatabaseOpen, + load: func(context.Context, graph.Database, ret.LoadConfig) (ret.LoadResult, error) { + calls = append(calls, "load") + return ret.LoadResult{}, errors.New("write transaction failed") + }, + verifyDatabase: func(context.Context, graph.Database, ret.VerifyDatabaseConfig) (ret.VerifyDatabaseResult, error) { + calls = append(calls, "verify-database") + return ret.VerifyDatabaseResult{}, nil + }, + }) + + err := runtime.run(context.Background(), []string{ + "load", + "-in", "collection", + "-verify-database", + "-connection", "postgresql://example/database", + }) + if err == nil || !strings.Contains(err.Error(), "clear") || !strings.Contains(err.Error(), "retry") { + t.Fatalf("load error = %v, want clear-then-retry guidance", err) + } + if !reflect.DeepEqual(calls, []string{"load"}) { + t.Fatalf("operation calls = %v, want only load", calls) } +} - err = runtime.run(context.Background(), []string{"verify"}) - if err == nil || !strings.Contains(err.Error(), "input directory is required") { - t.Fatalf("expected verify input validation error, got %v", err) +func TestPackUnpackAndKeygenCallIndependentFacadeOperations(t *testing.T) { + keyDir := t.TempDir() + privatePath := filepath.Join(keyDir, "private.key") + publicPath := filepath.Join(keyDir, "public.key") + if err := ret.Keygen(ret.KeygenConfig{ + PrivateKeyPath: privatePath, + PublicKeyPath: publicPath, + }); err != nil { + t.Fatalf("generate test keys: %v", err) } - err = runtime.run(context.Background(), []string{"bench", "-workers", "0"}) - if err == nil || !strings.Contains(err.Error(), "worker counts must be > 0") { - t.Fatalf("expected worker validation error, got %v", err) + var calls []string + runtime := newTestCommandRuntime(commandOperations{ + pack: func(_ context.Context, config ret.PackConfig) error { + calls = append(calls, "pack") + if config.CollectionDirectory != "collection" || config.ArchivePath != "collection.tar.enc" { + t.Fatalf("pack config = %+v", config) + } + return nil + }, + unpack: func(_ context.Context, config ret.UnpackConfig) error { + calls = append(calls, "unpack") + if config.ArchivePath != "collection.tar.enc" || config.OutputDirectory != "restored" { + t.Fatalf("unpack config = %+v", config) + } + return nil + }, + keygen: func(config ret.KeygenConfig) error { + calls = append(calls, "keygen") + if config.PrivateKeyPath != "new-private.key" || config.PublicKeyPath != "new-public.key" { + t.Fatalf("keygen config = %+v", config) + } + return nil + }, + }) + + if err := runtime.run(context.Background(), []string{ + "pack", + "-in", "collection", + "-archive", "collection.tar.enc", + "-recipient", publicPath, + }); err != nil { + t.Fatalf("pack: %v", err) + } + if err := runtime.run(context.Background(), []string{ + "unpack", + "-archive", "collection.tar.enc", + "-out", "restored", + "-identity", privatePath, + }); err != nil { + t.Fatalf("unpack: %v", err) + } + if err := runtime.run(context.Background(), []string{ + "keygen", + "-private-key", "new-private.key", + "-public-key", "new-public.key", + }); err != nil { + t.Fatalf("keygen: %v", err) + } + if !reflect.DeepEqual(calls, []string{"pack", "unpack", "keygen"}) { + t.Fatalf("operation calls = %v", calls) } } -func TestDumpArchiveOutputPreflightBeforeDatabase(t *testing.T) { - runtime := commandRuntime{ - stdout: &bytes.Buffer{}, - stderr: &bytes.Buffer{}, +func TestKeygenRejectsLegacyKeyFlagNames(t *testing.T) { + runtime := newTestCommandRuntime(commandOperations{}) + err := runtime.run(context.Background(), []string{ + "keygen", + "-private", "private.key", + "-public", "public.key", + }) + if err == nil || !strings.Contains(err.Error(), "flag provided but not defined") { + t.Fatalf("keygen error = %v, want undefined legacy flag", err) } - dir := t.TempDir() - privatePath := filepath.Join(dir, "private.key") - publicPath := filepath.Join(dir, "public.key") - if err := retriever.GenerateArchiveKeyFiles(privatePath, publicPath); err != nil { - t.Fatalf("generate archive keys: %v", err) +} + +func TestDumpPassesGraphOrderToFacade(t *testing.T) { + runtime := newTestCommandRuntime(commandOperations{ + openDatabase: successfulTestDatabaseOpen, + dump: func(_ context.Context, _ graph.Database, config ret.DumpConfig) (ret.DumpResult, error) { + if !reflect.DeepEqual(config.Graphs, []string{"second", "first"}) { + t.Fatalf("dump graphs = %v", config.Graphs) + } + return ret.DumpResult{}, nil + }, + }) + + if err := runtime.run(context.Background(), []string{ + "dump", + "-out", filepath.Join(t.TempDir(), "dump"), + "-graph", "second", + "-graph", "first", + "-connection", "postgresql://example/database", + }); err != nil { + t.Fatalf("dump: %v", err) } +} - archivePath := filepath.Join(dir, "dump.tar.pq") - if err := os.WriteFile(archivePath, []byte("exists"), 0o600); err != nil { - t.Fatalf("write existing archive: %v", err) +func TestDumpForceResumeFailsBeforeDeletionOrDatabaseOpen(t *testing.T) { + destination := t.TempDir() + marker := filepath.Join(destination, "keep") + if err := os.WriteFile(marker, []byte("keep"), 0o600); err != nil { + t.Fatalf("write marker: %v", err) } + var opened int + runtime := newTestCommandRuntime(commandOperations{ + openDatabase: func(context.Context, databaseConfig) (graph.Database, string, error) { + opened++ + return nil, "", nil + }, + }) err := runtime.run(context.Background(), []string{ "dump", - "-out", filepath.Join(dir, "dump"), - "-archive-out", archivePath, - "-recipient", publicPath, + "-out", destination, + "-force", + "-resume", + }) + if err == nil || !strings.Contains(err.Error(), "mutually exclusive") { + t.Fatalf("dump error = %v, want force/resume conflict", err) + } + if opened != 0 { + t.Fatalf("database open calls = %d, want 0", opened) + } + if _, err := os.Stat(marker); err != nil { + t.Fatalf("force/resume removed marker: %v", err) + } +} + +func TestDumpForcePureValidationFailuresPreserveDestination(t *testing.T) { + cases := []struct { + name string + args []string + want string + }{ + { + name: "all graphs conflict", + args: []string{"-graph", "asset", "-all-graphs"}, + want: "-all-graphs cannot be combined", + }, + { + name: "duplicate graph", + args: []string{"-graph", "asset", "-graph", "asset"}, + want: "duplicate graph", + }, + { + name: "unsafe graph path", + args: []string{"-graph", "../asset"}, + want: "safe path segment", + }, + { + name: "invalid pprof address", + args: []string{"-graph", "asset", "-pprof-listen", "not-an-address"}, + want: "invalid pprof listen address", + }, + { + name: "non-loopback pprof address", + args: []string{"-graph", "asset", "-pprof-listen", "0.0.0.0:6060"}, + want: "not loopback", + }, + } + + for _, test := range cases { + t.Run(test.name, func(t *testing.T) { + destination := filepath.Join(t.TempDir(), "dump") + if err := os.Mkdir(destination, 0o755); err != nil { + t.Fatalf("mkdir destination: %v", err) + } + marker := filepath.Join(destination, "keep") + if err := os.WriteFile(marker, []byte("keep"), 0o600); err != nil { + t.Fatalf("write marker: %v", err) + } + + runtime := newTestCommandRuntime(commandOperations{}) + args := append([]string{"dump", "-out", destination, "-force"}, test.args...) + err := runtime.run(context.Background(), args) + if err == nil || !strings.Contains(err.Error(), test.want) { + t.Fatalf("dump error = %v, want containing %q", err, test.want) + } + if contents, err := os.ReadFile(marker); err != nil || string(contents) != "keep" { + t.Fatalf("pure validation failure changed marker: contents=%q err=%v", contents, err) + } + }) + } +} + +func TestDumpForceRejectsRepositoryBroadTargetBeforeDeletion(t *testing.T) { + repositoryRoot, err := findRepositoryRoot() + if err != nil { + t.Fatalf("find repository root: %v", err) + } + marker := filepath.Join(repositoryRoot, "go.mod") + + runtime := newTestCommandRuntime(commandOperations{}) + err = runtime.run(context.Background(), []string{ + "dump", + "-out", repositoryRoot, + "-force", + }) + if err == nil || !strings.Contains(err.Error(), "unsafe") { + t.Fatalf("dump error = %v, want unsafe target", err) + } + if _, err := os.Stat(marker); err != nil { + t.Fatalf("repository marker missing after rejection: %v", err) + } +} + +func TestDumpForceRejectsStaticIntermediateSymlinkWithoutDeletingTarget(t *testing.T) { + physicalParent := t.TempDir() + physicalDestination := filepath.Join(physicalParent, "dump") + if err := os.Mkdir(physicalDestination, 0o755); err != nil { + t.Fatalf("mkdir physical destination: %v", err) + } + marker := filepath.Join(physicalDestination, "keep") + if err := os.WriteFile(marker, []byte("keep"), 0o600); err != nil { + t.Fatalf("write marker: %v", err) + } + + aliasParent := t.TempDir() + alias := filepath.Join(aliasParent, "alias") + if err := os.Symlink(physicalParent, alias); err != nil { + t.Fatalf("create intermediate symlink: %v", err) + } + runtime := newTestCommandRuntime(commandOperations{}) + err := runtime.run(context.Background(), []string{ + "dump", + "-out", filepath.Join(alias, "dump"), + "-force", + "-graph", "asset", + }) + if err == nil || !strings.Contains(err.Error(), "symbolic link") { + t.Fatalf("dump error = %v, want intermediate symlink rejection", err) + } + if contents, err := os.ReadFile(marker); err != nil || string(contents) != "keep" { + t.Fatalf("symlink target changed: contents=%q err=%v", contents, err) + } +} + +func TestDumpForceParentSubstitutionAfterPinPreservesBothTrees(t *testing.T) { + base := t.TempDir() + parent := filepath.Join(base, "parent") + movedParent := filepath.Join(base, "approved-parent") + destination := filepath.Join(parent, "dump") + if err := os.MkdirAll(destination, 0o755); err != nil { + t.Fatalf("mkdir destination: %v", err) + } + if err := os.WriteFile(filepath.Join(destination, "approved"), []byte("approved"), 0o600); err != nil { + t.Fatalf("write approved marker: %v", err) + } + + runtime := newTestCommandRuntime(commandOperations{}) + runtime.force.afterParentPinned = func(parentPath string, _ *os.Root) error { + if err := os.Rename(parentPath, movedParent); err != nil { + return err + } + if err := os.Mkdir(parentPath, 0o755); err != nil { + return err + } + if err := os.Mkdir(filepath.Join(parentPath, "dump"), 0o755); err != nil { + return err + } + return os.WriteFile(filepath.Join(parentPath, "dump", "replacement"), []byte("replacement"), 0o600) + } + err := runtime.run(context.Background(), []string{ + "dump", "-out", destination, "-force", "-graph", "asset", }) - if err == nil || !strings.Contains(err.Error(), "already exists") { - t.Fatalf("expected archive preflight error before database open, got %v", err) + if err == nil || !strings.Contains(err.Error(), "parent changed") { + t.Fatalf("dump error = %v, want pinned parent substitution failure", err) + } + if _, err := os.Stat(filepath.Join(movedParent, "dump", "approved")); err != nil { + t.Fatalf("approved tree was not preserved: %v", err) + } + if _, err := os.Stat(filepath.Join(parent, "dump", "replacement")); err != nil { + t.Fatalf("replacement tree was not preserved: %v", err) } } + +func TestDumpForceTargetSubstitutionAtQuarantineBoundaryRestoresReplacement(t *testing.T) { + parentPath := t.TempDir() + destination := filepath.Join(parentPath, "dump") + approvedMoved := filepath.Join(parentPath, "approved-moved") + if err := os.Mkdir(destination, 0o755); err != nil { + t.Fatalf("mkdir destination: %v", err) + } + if err := os.WriteFile(filepath.Join(destination, "approved"), []byte("approved"), 0o600); err != nil { + t.Fatalf("write approved marker: %v", err) + } + + runtime := newTestCommandRuntime(commandOperations{}) + runtime.force.beforeQuarantine = func(parent *os.Root, name string, _ *os.Root) error { + if err := parent.Rename(name, "approved-moved"); err != nil { + return err + } + if err := parent.Mkdir(name, 0o755); err != nil { + return err + } + return parent.WriteFile(filepath.Join(name, "replacement"), []byte("replacement"), 0o600) + } + err := runtime.run(context.Background(), []string{ + "dump", "-out", destination, "-force", "-graph", "asset", + }) + if err == nil || !strings.Contains(err.Error(), "substituted") { + t.Fatalf("dump error = %v, want target substitution failure", err) + } + if _, err := os.Stat(filepath.Join(approvedMoved, "approved")); err != nil { + t.Fatalf("approved tree was not preserved: %v", err) + } + if _, err := os.Stat(filepath.Join(destination, "replacement")); err != nil { + t.Fatalf("replacement tree was not restored: %v", err) + } +} + +func TestDumpForcePostQuarantineSubstitutionRestoresReplacement(t *testing.T) { + parentPath := t.TempDir() + destination := filepath.Join(parentPath, "dump") + if err := os.Mkdir(destination, 0o755); err != nil { + t.Fatalf("mkdir destination: %v", err) + } + if err := os.WriteFile(filepath.Join(destination, "approved"), []byte("approved"), 0o600); err != nil { + t.Fatalf("write approved marker: %v", err) + } + + var approvedMoved string + runtime := newTestCommandRuntime(commandOperations{}) + runtime.force.afterQuarantine = func(parent *os.Root, _, quarantine string) error { + approvedMoved = quarantine + ".approved" + if err := parent.Rename(quarantine, approvedMoved); err != nil { + return err + } + if err := parent.Mkdir(quarantine, 0o755); err != nil { + return err + } + return parent.WriteFile(filepath.Join(quarantine, "replacement"), []byte("replacement"), 0o600) + } + err := runtime.run(context.Background(), []string{ + "dump", "-out", destination, "-force", "-graph", "asset", + }) + if err == nil || !strings.Contains(err.Error(), "substituted") { + t.Fatalf("dump error = %v, want post-quarantine substitution failure", err) + } + if _, err := os.Stat(filepath.Join(parentPath, approvedMoved, "approved")); err != nil { + t.Fatalf("approved tree was not preserved: %v", err) + } + if _, err := os.Stat(filepath.Join(destination, "replacement")); err != nil { + t.Fatalf("replacement tree was not restored: %v", err) + } +} + +func TestDumpForceBlockedSubstitutionRestorePreservesEveryObject(t *testing.T) { + parentPath := t.TempDir() + destination := filepath.Join(parentPath, "dump") + if err := os.Mkdir(destination, 0o755); err != nil { + t.Fatalf("mkdir destination: %v", err) + } + if err := os.WriteFile(filepath.Join(destination, "approved"), []byte("approved"), 0o600); err != nil { + t.Fatalf("write approved marker: %v", err) + } + + runtime := newTestCommandRuntime(commandOperations{}) + runtime.force.beforeQuarantine = func(parent *os.Root, name string, _ *os.Root) error { + if err := parent.Rename(name, "approved-moved"); err != nil { + return err + } + if err := parent.Mkdir(name, 0o755); err != nil { + return err + } + return parent.WriteFile(filepath.Join(name, "replacement"), []byte("replacement"), 0o600) + } + runtime.force.afterQuarantine = func(parent *os.Root, original, _ string) error { + if err := parent.Mkdir(original, 0o755); err != nil { + return err + } + return parent.WriteFile(filepath.Join(original, "blocker"), []byte("blocker"), 0o600) + } + err := runtime.run(context.Background(), []string{ + "dump", "-out", destination, "-force", "-graph", "asset", + }) + if err == nil || !strings.Contains(err.Error(), "preserving both pathnames") { + t.Fatalf("dump error = %v, want blocked restoration report", err) + } + if _, err := os.Stat(filepath.Join(parentPath, "approved-moved", "approved")); err != nil { + t.Fatalf("approved tree was not preserved: %v", err) + } + if _, err := os.Stat(filepath.Join(destination, "blocker")); err != nil { + t.Fatalf("blocking tree was not preserved: %v", err) + } + quarantines, err := filepath.Glob(filepath.Join(parentPath, ".ret-force-*.preserved")) + if err != nil { + t.Fatalf("glob quarantines: %v", err) + } + if len(quarantines) != 1 { + t.Fatalf("quarantines = %v, want preserved replacement", quarantines) + } + if _, err := os.Stat(filepath.Join(quarantines[0], "replacement")); err != nil { + t.Fatalf("quarantined replacement was not preserved: %v", err) + } +} + +func TestDumpForceRejectsExistingDestinationSymlink(t *testing.T) { + parentPath := t.TempDir() + physical := filepath.Join(parentPath, "physical") + if err := os.Mkdir(physical, 0o755); err != nil { + t.Fatalf("mkdir physical target: %v", err) + } + marker := filepath.Join(physical, "keep") + if err := os.WriteFile(marker, []byte("keep"), 0o600); err != nil { + t.Fatalf("write marker: %v", err) + } + destination := filepath.Join(parentPath, "dump") + if err := os.Symlink(physical, destination); err != nil { + t.Fatalf("symlink destination: %v", err) + } + + runtime := newTestCommandRuntime(commandOperations{}) + err := runtime.run(context.Background(), []string{ + "dump", "-out", destination, "-force", "-graph", "asset", + }) + if err == nil || !strings.Contains(err.Error(), "symbolic link") { + t.Fatalf("dump error = %v, want destination symlink rejection", err) + } + if _, err := os.Stat(marker); err != nil { + t.Fatalf("symlink target changed: %v", err) + } +} + +func TestDumpForceAbsentTargetCreatesNoTombstoneAndCallsDump(t *testing.T) { + parent := t.TempDir() + destination := filepath.Join(parent, "dump") + var called bool + runtime := newTestCommandRuntime(commandOperations{ + openDatabase: successfulTestDatabaseOpen, + dump: func(_ context.Context, _ graph.Database, config ret.DumpConfig) (ret.DumpResult, error) { + called = true + if config.Directory != destination { + t.Fatalf("dump directory = %q, want %q", config.Directory, destination) + } + return ret.DumpResult{}, nil + }, + }) + if err := runtime.run(context.Background(), []string{ + "dump", "-out", destination, "-force", "-graph", "asset", + }); err != nil { + t.Fatalf("dump: %v", err) + } + if !called { + t.Fatal("dump operation was not called") + } + tombstones, err := filepath.Glob(filepath.Join(parent, ".ret-force-*.preserved")) + if err != nil { + t.Fatalf("glob tombstones: %v", err) + } + if len(tombstones) != 0 { + t.Fatalf("absent target created tombstones: %v", tombstones) + } +} + +func TestForcePlatformValidationRejectsUnsupportedName(t *testing.T) { + err := validateForcePlatform("windows") + if err == nil || !strings.Contains(err.Error(), "unsupported") { + t.Fatalf("platform validation error = %v", err) + } +} + +func TestReplaceDumpDestinationRejectsRootHomeAndRepositoryAncestors(t *testing.T) { + repositoryRoot, err := findRepositoryRoot() + if err != nil { + t.Fatalf("find repository root: %v", err) + } + home, err := os.UserHomeDir() + if err != nil { + t.Fatalf("user home: %v", err) + } + root := filepath.Clean(filepath.VolumeName(repositoryRoot) + string(os.PathSeparator)) + + for _, target := range []string{ + root, + home, + repositoryRoot, + filepath.Dir(repositoryRoot), + } { + t.Run(target, func(t *testing.T) { + if _, err := replaceDumpDestination(target, forceReplaceOperations{}); err == nil || !strings.Contains(err.Error(), "unsafe") { + t.Fatalf("replaceDumpDestination(%q) error = %v, want unsafe", target, err) + } + }) + } +} + +func TestReplaceDumpDestinationReturnsCleanAbsoluteChild(t *testing.T) { + target := filepath.Join(t.TempDir(), "parent", "..", "dump") + replacement, err := replaceDumpDestination(target, forceReplaceOperations{}) + if err != nil { + t.Fatalf("validate replace target: %v", err) + } + if !filepath.IsAbs(replacement.destination) || replacement.destination != filepath.Clean(target) { + t.Fatalf("resolved target = %q, want clean absolute %q", replacement.destination, filepath.Clean(target)) + } +} + +func TestDumpForceMovesAsideOnlyExactDestinationBeforeDump(t *testing.T) { + parent := t.TempDir() + destination := filepath.Join(parent, "replace") + sibling := filepath.Join(parent, "keep") + if err := os.Mkdir(destination, 0o755); err != nil { + t.Fatalf("mkdir destination: %v", err) + } + if err := os.WriteFile(filepath.Join(destination, "old"), []byte("old"), 0o600); err != nil { + t.Fatalf("write old file: %v", err) + } + if err := os.WriteFile(sibling, []byte("keep"), 0o600); err != nil { + t.Fatalf("write sibling: %v", err) + } + + runtime := newTestCommandRuntime(commandOperations{ + openDatabase: successfulTestDatabaseOpen, + dump: func(_ context.Context, _ graph.Database, config ret.DumpConfig) (ret.DumpResult, error) { + if _, err := os.Lstat(config.Directory); !errors.Is(err, os.ErrNotExist) { + t.Fatalf("dump destination still exists before facade call: %v", err) + } + return ret.DumpResult{}, nil + }, + }) + if err := runtime.run(context.Background(), []string{ + "dump", + "-out", destination, + "-force", + "-connection", "postgresql://example/database", + }); err != nil { + t.Fatalf("dump: %v", err) + } + if contents, err := os.ReadFile(sibling); err != nil || string(contents) != "keep" { + t.Fatalf("sibling changed: contents=%q err=%v", contents, err) + } + tombstones, err := filepath.Glob(filepath.Join(parent, ".ret-force-*.preserved")) + if err != nil { + t.Fatalf("glob tombstones: %v", err) + } + if len(tombstones) != 1 { + t.Fatalf("tombstones = %v, want one preserved prior collection", tombstones) + } + if contents, err := os.ReadFile(filepath.Join(tombstones[0], "old")); err != nil || string(contents) != "old" { + t.Fatalf("prior destination changed: contents=%q err=%v", contents, err) + } +} + +func TestProductCommandsReportDatabaseCloseErrors(t *testing.T) { + closeFailure := errors.New("close failed") + cases := []struct { + name string + args []string + ops func(error) commandOperations + wantNoOutput bool + }{ + { + name: "dump", + args: []string{"dump", "-out", filepath.Join(t.TempDir(), "dump"), "-graph", "asset"}, + ops: func(closeErr error) commandOperations { + return commandOperations{ + dump: func(context.Context, graph.Database, ret.DumpConfig) (ret.DumpResult, error) { + return ret.DumpResult{}, nil + }, + } + }, + }, + { + name: "load", + args: []string{"load", "-in", "collection"}, + ops: func(closeErr error) commandOperations { + return commandOperations{ + load: func(context.Context, graph.Database, ret.LoadConfig) (ret.LoadResult, error) { + return ret.LoadResult{}, nil + }, + } + }, + }, + { + name: "verify database", + args: []string{"verify-database", "-in", "collection"}, + ops: func(closeErr error) commandOperations { + return commandOperations{ + verifyDatabase: func(context.Context, graph.Database, ret.VerifyDatabaseConfig) (ret.VerifyDatabaseResult, error) { + return ret.VerifyDatabaseResult{}, nil + }, + } + }, + }, + { + name: "bench", + args: []string{"bench", "-graph", "asset"}, + wantNoOutput: true, + ops: func(closeErr error) commandOperations { + return commandOperations{} + }, + }, + } + + for _, test := range cases { + t.Run(test.name, func(t *testing.T) { + database := &closingTestDatabase{closeErr: closeFailure} + operations := test.ops(closeFailure) + operations.openDatabase = func(context.Context, databaseConfig) (graph.Database, string, error) { + return database, "pg", nil + } + runtime := newTestCommandRuntime(operations) + err := runtime.run(context.Background(), test.args) + if !errors.Is(err, closeFailure) { + t.Fatalf("command error = %v, want close failure", err) + } + if test.wantNoOutput && runtime.stdout.(*bytes.Buffer).Len() != 0 { + t.Fatalf("command emitted success report before close failure: %q", runtime.stdout) + } + }) + } +} + +func TestProductCommandJoinsPrimaryAndDatabaseCloseErrors(t *testing.T) { + primaryFailure := errors.New("dump failed") + closeFailure := errors.New("close failed") + database := &closingTestDatabase{closeErr: closeFailure} + runtime := newTestCommandRuntime(commandOperations{ + openDatabase: func(context.Context, databaseConfig) (graph.Database, string, error) { + return database, "pg", nil + }, + dump: func(context.Context, graph.Database, ret.DumpConfig) (ret.DumpResult, error) { + return ret.DumpResult{}, primaryFailure + }, + }) + + err := runtime.run(context.Background(), []string{ + "dump", + "-out", filepath.Join(t.TempDir(), "dump"), + "-graph", "asset", + }) + if !errors.Is(err, primaryFailure) || !errors.Is(err, closeFailure) { + t.Fatalf("command error = %v, want joined primary and close failures", err) + } +} + +func TestProductDatabaseCloseUsesNonCanceledContext(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + database := &closingTestDatabase{} + runtime := newTestCommandRuntime(commandOperations{ + openDatabase: func(context.Context, databaseConfig) (graph.Database, string, error) { + return database, "pg", nil + }, + dump: func(context.Context, graph.Database, ret.DumpConfig) (ret.DumpResult, error) { + cancel() + return ret.DumpResult{}, context.Canceled + }, + }) + + err := runtime.run(ctx, []string{ + "dump", + "-out", filepath.Join(t.TempDir(), "dump"), + "-graph", "asset", + }) + if !errors.Is(err, context.Canceled) { + t.Fatalf("command error = %v, want context cancellation", err) + } + if len(database.closeContextErrors) != 1 { + t.Fatalf("close contexts = %d, want 1", len(database.closeContextErrors)) + } + if err := database.closeContextErrors[0]; err != nil { + t.Fatalf("database close context was already canceled: %v", err) + } +} + +func TestBenchDatabaseCloseUsesNonCanceledContextAfterCancellation(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + database := &closingTestDatabase{} + runtime := newTestCommandRuntime(commandOperations{ + openDatabase: func(context.Context, databaseConfig) (graph.Database, string, error) { + cancel() + return database, "pg", nil + }, + }) + + err := runtime.run(ctx, []string{"bench", "-graph", "asset"}) + if !errors.Is(err, context.Canceled) { + t.Fatalf("bench error = %v, want context cancellation", err) + } + if len(database.closeContextErrors) != 1 { + t.Fatalf("close contexts = %d, want 1", len(database.closeContextErrors)) + } + if err := database.closeContextErrors[0]; err != nil { + t.Fatalf("database close context was already canceled: %v", err) + } +} + +func newTestCommandRuntime(operations commandOperations) commandRuntime { + return commandRuntime{ + stdout: &bytes.Buffer{}, + stderr: &bytes.Buffer{}, + operations: operations, + } +} + +func successfulTestDatabaseOpen(context.Context, databaseConfig) (graph.Database, string, error) { + return nil, "pg", nil +} + +type closingTestDatabase struct { + graph.Database + closeErr error + closeContextErrors []error +} + +func (s *closingTestDatabase) Close(ctx context.Context) error { + s.closeContextErrors = append(s.closeContextErrors, ctx.Err()) + return s.closeErr +} + +func (s *closingTestDatabase) ReadTransaction(_ context.Context, delegate graph.TransactionDelegate, _ ...graph.TransactionOption) error { + return delegate(emptyBenchTransaction{}) +} diff --git a/cmd/retriever/observer.go b/cmd/retriever/observer.go new file mode 100644 index 00000000..0281ea3c --- /dev/null +++ b/cmd/retriever/observer.go @@ -0,0 +1,234 @@ +package main + +import ( + "context" + "fmt" + "log/slog" + "os" + "runtime" + "strconv" + "strings" + "sync" + "time" + + "github.com/specterops/dawgs/ret/observe" +) + +const commandProgressInterval int64 = 250_000 + +type commandObserver struct { + logger *slog.Logger + now func() time.Time + progressInterval int64 + + mu sync.Mutex + phases map[commandPhaseKey]commandPhaseState +} + +type commandPhaseKey struct { + operation string + graph string + phase string +} + +type commandPhaseState struct { + started time.Time + baseline int64 + next int64 +} + +type commandRuntimeTelemetry struct { + heapAlloc uint64 + heapInuse uint64 + sys uint64 + numGC uint32 + rss uint64 +} + +func newCommandObserver(logger *slog.Logger) *commandObserver { + if logger == nil { + logger = slog.Default() + } + return &commandObserver{ + logger: logger, + now: time.Now, + progressInterval: commandProgressInterval, + phases: map[commandPhaseKey]commandPhaseState{}, + } +} + +func (s *commandObserver) Observe(ctx context.Context, event observe.Event) { + switch value := event.(type) { + case observe.OperationStarted: + s.logger.InfoContext(ctx, "retriever operation started", + slog.String("operation", value.Operation)) + case observe.OperationCompleted: + attributes := []any{ + slog.String("operation", value.Operation), + slog.Duration("duration", value.Duration), + } + if value.Err != nil { + attributes = append(attributes, slog.Any("error", value.Err)) + s.logger.ErrorContext(ctx, "retriever operation completed", attributes...) + } else { + s.logger.InfoContext(ctx, "retriever operation completed", attributes...) + } + case observe.GraphStarted: + s.logger.InfoContext(ctx, "retriever graph started", + slog.String("operation", value.Operation), + slog.String("graph", value.Graph)) + case observe.GraphCompleted: + s.logger.InfoContext(ctx, "retriever graph completed", + slog.String("operation", value.Operation), + slog.String("graph", value.Graph), + slog.Int64("nodes", value.Nodes), + slog.Int64("relationships", value.Relationships), + slog.Duration("duration", value.Duration)) + case observe.PhaseStarted: + s.phaseStarted(ctx, value) + case observe.PhaseProgress: + s.phaseProgress(ctx, value) + case observe.PhaseCompleted: + s.phaseCompleted(ctx, value) + case observe.ShardCommitted: + s.logger.InfoContext(ctx, "retriever shard committed", + slog.String("graph", value.Graph), + slog.String("entity_type", value.EntityType), + slog.Int("shard", value.Index), + slog.Int64("count", value.Count), + slog.String("jsonl_path", value.JSONLPath), + slog.Int64("jsonl_bytes", value.JSONLBytes), + slog.String("parquet_path", value.ParquetPath), + slog.Int64("parquet_bytes", value.ParquetBytes)) + case observe.ArtifactVerified: + s.logger.InfoContext(ctx, "retriever artifact verified", + slog.String("graph", value.Graph), + slog.String("entity_type", value.EntityType), + slog.String("format", value.Format), + slog.String("path", value.Path), + slog.Int64("count", value.Count), + slog.Int64("bytes", value.Bytes)) + case observe.ArchiveEntryProcessed: + s.logger.InfoContext(ctx, "retriever archive entry processed", + slog.String("operation", value.Operation), + slog.String("path", value.Path), + slog.Int64("bytes", value.Size)) + default: + s.logger.DebugContext(ctx, "retriever event", + slog.String("type", fmt.Sprintf("%T", event))) + } +} + +func (s *commandObserver) phaseStarted(ctx context.Context, event observe.PhaseStarted) { + key := commandPhaseKey{operation: event.Operation, graph: event.Graph, phase: event.Phase} + interval := s.normalizedProgressInterval() + next := event.Completed + interval + if event.Total > event.Completed && event.Total-event.Completed < interval { + next = event.Total + } + s.mu.Lock() + s.phases[key] = commandPhaseState{ + started: s.now(), + baseline: event.Completed, + next: next, + } + s.mu.Unlock() + + s.logger.InfoContext(ctx, "retriever phase started", + slog.String("operation", event.Operation), + slog.String("graph", event.Graph), + slog.String("phase", event.Phase), + slog.Int64("completed", event.Completed), + slog.Int64("total", event.Total)) +} + +func (s *commandObserver) phaseProgress(ctx context.Context, event observe.PhaseProgress) { + key := commandPhaseKey{operation: event.Operation, graph: event.Graph, phase: event.Phase} + now := s.now() + interval := s.normalizedProgressInterval() + + s.mu.Lock() + state, found := s.phases[key] + if !found { + state = commandPhaseState{ + started: now, + next: interval, + } + } + report := event.Completed >= state.next || event.Completed >= event.Total + if report { + for state.next <= event.Completed { + state.next += interval + } + s.phases[key] = state + } + s.mu.Unlock() + if !report { + return + } + + elapsed := now.Sub(state.started) + telemetry := sampleCommandRuntimeTelemetry() + s.logger.InfoContext(ctx, "retriever phase progress", + slog.String("operation", event.Operation), + slog.String("graph", event.Graph), + slog.String("phase", event.Phase), + slog.Int64("completed", event.Completed), + slog.Int64("total", event.Total), + slog.Duration("elapsed", elapsed), + slog.Float64("entities_per_second", commandPerSecond(event.Completed-state.baseline, elapsed)), + slog.Uint64("heap_alloc_bytes", telemetry.heapAlloc), + slog.Uint64("heap_inuse_bytes", telemetry.heapInuse), + slog.Uint64("sys_bytes", telemetry.sys), + slog.Uint64("gc_count", uint64(telemetry.numGC)), + slog.Uint64("rss_bytes", telemetry.rss)) +} + +func (s *commandObserver) phaseCompleted(ctx context.Context, event observe.PhaseCompleted) { + key := commandPhaseKey{operation: event.Operation, graph: event.Graph, phase: event.Phase} + s.mu.Lock() + delete(s.phases, key) + s.mu.Unlock() + + s.logger.InfoContext(ctx, "retriever phase completed", + slog.String("operation", event.Operation), + slog.String("graph", event.Graph), + slog.String("phase", event.Phase), + slog.Int64("completed", event.Completed), + slog.Duration("duration", event.Duration)) +} + +func (s *commandObserver) normalizedProgressInterval() int64 { + if s.progressInterval <= 0 { + return commandProgressInterval + } + return s.progressInterval +} + +func sampleCommandRuntimeTelemetry() commandRuntimeTelemetry { + var stats runtime.MemStats + runtime.ReadMemStats(&stats) + return commandRuntimeTelemetry{ + heapAlloc: stats.HeapAlloc, + heapInuse: stats.HeapInuse, + sys: stats.Sys, + numGC: stats.NumGC, + rss: commandCurrentRSS(), + } +} + +func commandCurrentRSS() uint64 { + contents, err := os.ReadFile("/proc/self/statm") + if err != nil { + return 0 + } + fields := strings.Fields(string(contents)) + if len(fields) < 2 { + return 0 + } + pages, err := strconv.ParseUint(fields[1], 10, 64) + if err != nil { + return 0 + } + return pages * uint64(os.Getpagesize()) +} diff --git a/cmd/retriever/observer_test.go b/cmd/retriever/observer_test.go new file mode 100644 index 00000000..13f2e19b --- /dev/null +++ b/cmd/retriever/observer_test.go @@ -0,0 +1,325 @@ +package main + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "log/slog" + "strings" + "sync" + "testing" + "time" + + "github.com/specterops/dawgs/ret/observe" +) + +func TestCommandObserverTranslatesEveryConcreteEvent(t *testing.T) { + cases := []struct { + name string + event observe.Event + want string + }{ + {name: "operation started", event: observe.OperationStarted{Operation: "dump"}, want: "operation started"}, + {name: "operation completed", event: observe.OperationCompleted{Operation: "dump", Duration: time.Second}, want: "operation completed"}, + {name: "graph started", event: observe.GraphStarted{Operation: "dump", Graph: "asset"}, want: "graph started"}, + {name: "graph completed", event: observe.GraphCompleted{Operation: "dump", Graph: "asset", Nodes: 2, Relationships: 1, Duration: time.Second}, want: "graph completed"}, + {name: "phase started", event: observe.PhaseStarted{Operation: "dump", Graph: "asset", Phase: "nodes", Total: 2}, want: "phase started"}, + {name: "phase progress", event: observe.PhaseProgress{Operation: "dump", Graph: "asset", Phase: "nodes", Completed: 2, Total: 2}, want: "phase progress"}, + {name: "phase completed", event: observe.PhaseCompleted{Operation: "dump", Graph: "asset", Phase: "nodes", Completed: 2, Duration: time.Second}, want: "phase completed"}, + {name: "shard committed", event: observe.ShardCommitted{Graph: "asset", EntityType: "nodes", Index: 1, Count: 2, JSONLPath: "nodes.jsonl"}, want: "shard committed"}, + {name: "artifact verified", event: observe.ArtifactVerified{Graph: "asset", EntityType: "nodes", Format: "jsonl", Path: "nodes.jsonl", Count: 2}, want: "artifact verified"}, + {name: "archive entry", event: observe.ArchiveEntryProcessed{Operation: "pack", Path: "manifest.json", Size: 100}, want: "archive entry processed"}, + } + + for _, test := range cases { + t.Run(test.name, func(t *testing.T) { + var output bytes.Buffer + observer := newCommandObserver(slog.New(slog.NewTextHandler(&output, &slog.HandlerOptions{ + Level: slog.LevelDebug, + }))) + observer.Observe(context.Background(), test.event) + if !strings.Contains(output.String(), test.want) { + t.Fatalf("log output %q does not contain %q", output.String(), test.want) + } + }) + } +} + +func TestCommandObserverLogsOperationFailureAsError(t *testing.T) { + var output bytes.Buffer + observer := newCommandObserver(slog.New(slog.NewTextHandler(&output, nil))) + observer.Observe(context.Background(), observe.OperationCompleted{ + Operation: "load", + Duration: time.Second, + Err: errors.New("failed"), + }) + + if got := output.String(); !strings.Contains(got, "level=ERROR") || !strings.Contains(got, "error=failed") { + t.Fatalf("failure log = %q", got) + } +} + +func TestCommandObserverWritesStructuredEventAttributes(t *testing.T) { + cases := []struct { + name string + event observe.Event + want map[string]any + }{ + { + name: "operation", + event: observe.OperationStarted{Operation: "dump"}, + want: map[string]any{"operation": "dump"}, + }, + { + name: "graph", + event: observe.GraphCompleted{ + Operation: "load", + Graph: "asset", + Nodes: 4, + Relationships: 3, + }, + want: map[string]any{ + "operation": "load", + "graph": "asset", + "nodes": float64(4), + "relationships": float64(3), + }, + }, + { + name: "shard", + event: observe.ShardCommitted{ + Graph: "asset", + EntityType: "nodes", + Index: 2, + Count: 5, + JSONLPath: "nodes.jsonl", + JSONLBytes: 10, + ParquetPath: "nodes.parquet", + ParquetBytes: 20, + }, + want: map[string]any{ + "graph": "asset", + "entity_type": "nodes", + "shard": float64(2), + "count": float64(5), + "jsonl_path": "nodes.jsonl", + "jsonl_bytes": float64(10), + "parquet_path": "nodes.parquet", + "parquet_bytes": float64(20), + }, + }, + { + name: "artifact", + event: observe.ArtifactVerified{ + Graph: "asset", + EntityType: "relationships", + Format: "parquet", + Path: "relationships.parquet", + Count: 6, + Bytes: 30, + }, + want: map[string]any{ + "graph": "asset", + "entity_type": "relationships", + "format": "parquet", + "path": "relationships.parquet", + "count": float64(6), + "bytes": float64(30), + }, + }, + { + name: "archive", + event: observe.ArchiveEntryProcessed{Operation: "unpack", Path: "manifest.json", Size: 40}, + want: map[string]any{"operation": "unpack", "path": "manifest.json", "bytes": float64(40)}, + }, + } + + for _, test := range cases { + t.Run(test.name, func(t *testing.T) { + var output bytes.Buffer + observer := newCommandObserver(slog.New(slog.NewJSONHandler(&output, nil))) + observer.Observe(context.Background(), test.event) + var record map[string]any + if err := json.Unmarshal(output.Bytes(), &record); err != nil { + t.Fatalf("decode structured log %q: %v", output.String(), err) + } + for key, expected := range test.want { + if actual := record[key]; actual != expected { + t.Fatalf("structured attribute %q = %#v, want %#v; record=%v", key, actual, expected, record) + } + } + }) + } +} + +func TestCommandObserverSamplesProgressAndAddsRuntimeTelemetry(t *testing.T) { + var output bytes.Buffer + observer := newCommandObserver(slog.New(slog.NewTextHandler(&output, nil))) + observer.progressInterval = 100 + started := time.Unix(100, 0) + observer.now = func() time.Time { return started } + observer.Observe(context.Background(), observe.PhaseStarted{ + Operation: "dump", + Graph: "asset", + Phase: "nodes", + Total: 250, + }) + + output.Reset() + observer.now = func() time.Time { return started.Add(2 * time.Second) } + observer.Observe(context.Background(), observe.PhaseProgress{ + Operation: "dump", + Graph: "asset", + Phase: "nodes", + Completed: 50, + Total: 250, + }) + if output.Len() != 0 { + t.Fatalf("unsampled progress was logged: %q", output.String()) + } + + observer.now = func() time.Time { return started.Add(4 * time.Second) } + observer.Observe(context.Background(), observe.PhaseProgress{ + Operation: "dump", + Graph: "asset", + Phase: "nodes", + Completed: 150, + Total: 250, + }) + got := output.String() + for _, value := range []string{ + "phase progress", + "entities_per_second=37.5", + "heap_alloc_bytes=", + "heap_inuse_bytes=", + "sys_bytes=", + "gc_count=", + "rss_bytes=", + } { + if !strings.Contains(got, value) { + t.Fatalf("progress log %q does not contain %q", got, value) + } + } +} + +func TestCommandObserverUsesFirstProgressAsResumedBaseline(t *testing.T) { + var output bytes.Buffer + observer := newCommandObserver(slog.New(slog.NewTextHandler(&output, nil))) + observer.progressInterval = 100 + started := time.Unix(100, 0) + observer.now = func() time.Time { return started } + observer.Observe(context.Background(), observe.PhaseStarted{ + Operation: "dump", + Graph: "asset", + Phase: "nodes", + Completed: 500, + Total: 1_000, + }) + + output.Reset() + observer.now = func() time.Time { return started.Add(2 * time.Second) } + observer.Observe(context.Background(), observe.PhaseProgress{ + Operation: "dump", + Graph: "asset", + Phase: "nodes", + Completed: 600, + Total: 1_000, + }) + if got := output.String(); !strings.Contains(got, "entities_per_second=50") { + t.Fatalf("resumed rate log = %q, want 100/2 entities per second", got) + } +} + +func TestCommandObserverReportsFreshOneBatchProgress(t *testing.T) { + var output bytes.Buffer + observer := newCommandObserver(slog.New(slog.NewTextHandler(&output, nil))) + observer.progressInterval = 100 + started := time.Unix(100, 0) + observer.now = func() time.Time { return started } + observer.Observe(context.Background(), observe.PhaseStarted{ + Operation: "load", + Graph: "asset", + Phase: "nodes", + Completed: 0, + Total: 1, + }) + + output.Reset() + observer.now = func() time.Time { return started.Add(time.Second) } + observer.Observe(context.Background(), observe.PhaseProgress{ + Operation: "load", + Graph: "asset", + Phase: "nodes", + Completed: 1, + Total: 1, + }) + if got := output.String(); !strings.Contains(got, "phase progress") || + !strings.Contains(got, "entities_per_second=1") || + !strings.Contains(got, "heap_alloc_bytes=") { + t.Fatalf("fresh one-batch progress log = %q, want progress, rate, and telemetry", got) + } +} + +func TestCommandObserverCompletionRemovesPhaseState(t *testing.T) { + observer := newCommandObserver(slog.New(slog.NewTextHandler(&bytes.Buffer{}, nil))) + event := observe.PhaseStarted{Operation: "load", Graph: "asset", Phase: "nodes", Total: 10} + observer.Observe(context.Background(), event) + if len(observer.phases) != 1 { + t.Fatalf("phase state count = %d, want 1", len(observer.phases)) + } + observer.Observe(context.Background(), observe.PhaseCompleted{ + Operation: event.Operation, + Graph: event.Graph, + Phase: event.Phase, + Completed: 10, + }) + if len(observer.phases) != 0 { + t.Fatalf("phase state count after completion = %d, want 0", len(observer.phases)) + } +} + +func TestCommandObserverHandlesConcurrentPhasesAndCanceledContexts(t *testing.T) { + var output bytes.Buffer + observer := newCommandObserver(slog.New(slog.NewTextHandler(&output, nil))) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + const phases = 32 + var wait sync.WaitGroup + wait.Add(phases) + for index := range phases { + go func() { + defer wait.Done() + graphName := fmt.Sprintf("graph-%d", index) + observer.Observe(ctx, observe.PhaseStarted{ + Operation: "verify_database", + Graph: graphName, + Phase: "nodes", + Total: 1, + }) + observer.Observe(ctx, observe.PhaseProgress{ + Operation: "verify_database", + Graph: graphName, + Phase: "nodes", + Completed: 1, + Total: 1, + }) + observer.Observe(ctx, observe.PhaseCompleted{ + Operation: "verify_database", + Graph: graphName, + Phase: "nodes", + Completed: 1, + }) + }() + } + wait.Wait() + + if len(observer.phases) != 0 { + t.Fatalf("phase state count = %d, want 0", len(observer.phases)) + } + if !strings.Contains(output.String(), "operation=verify_database") { + t.Fatalf("canceled-context events were not structured: %q", output.String()) + } +} diff --git a/cmd/retriever/pprof_test.go b/cmd/retriever/pprof_test.go index de1e0b9f..63e6a9d9 100644 --- a/cmd/retriever/pprof_test.go +++ b/cmd/retriever/pprof_test.go @@ -143,7 +143,7 @@ func TestPprofServerRejectsNonLoopbackAddress(t *testing.T) { } func TestLongRunningCommandsExposePprofFlag(t *testing.T) { - for _, command := range []string{"dump", "load", "verify", "bench"} { + for _, command := range []string{"dump", "load", "verify-database", "bench"} { var ( stdout bytes.Buffer stderr bytes.Buffer diff --git a/cmd/retriever/rate.go b/cmd/retriever/rate.go new file mode 100644 index 00000000..d300724e --- /dev/null +++ b/cmd/retriever/rate.go @@ -0,0 +1,10 @@ +package main + +import "time" + +func commandPerSecond(count int64, elapsed time.Duration) float64 { + if count <= 0 || elapsed <= 0 { + return 0 + } + return float64(count) / elapsed.Seconds() +} diff --git a/cmd/retriever/retriever_integration_test.go b/cmd/retriever/retriever_integration_test.go index 1009e8f2..baed0bc4 100644 --- a/cmd/retriever/retriever_integration_test.go +++ b/cmd/retriever/retriever_integration_test.go @@ -19,374 +19,865 @@ package main import ( + "bytes" "context" + "encoding/json" "errors" "fmt" "os" + "path/filepath" "reflect" - "strings" + "runtime" + "sync" + "sync/atomic" "testing" "time" "github.com/specterops/dawgs/graph" "github.com/specterops/dawgs/ops" - "github.com/specterops/dawgs/query" - "github.com/specterops/dawgs/retriever" + "github.com/specterops/dawgs/ret" + "github.com/specterops/dawgs/ret/archive" + "github.com/specterops/dawgs/ret/collection" + "github.com/specterops/dawgs/ret/dawgs" + "github.com/specterops/dawgs/ret/entity" + "github.com/specterops/dawgs/ret/jsonl" + "github.com/specterops/dawgs/ret/observe" + "github.com/specterops/dawgs/ret/parquet" + "github.com/specterops/dawgs/ret/scrub" ) -func TestDumpLoadRoundTrip(t *testing.T) { - connection := os.Getenv("CONNECTION_STRING") - if connection == "" { - t.Skip("CONNECTION_STRING not set") +func TestRetFacadeCollectionMatrix(t *testing.T) { + for _, testCase := range []struct { + name string + jsonl *jsonl.Config + parquet *parquet.Config + loadable bool + }{ + { + name: "jsonl", + jsonl: &jsonl.Config{Codec: jsonl.CodecZstd}, + loadable: true, + }, + { + name: "parquet", + parquet: &parquet.Config{}, + loadable: false, + }, + { + name: "dual", + jsonl: &jsonl.Config{Codec: jsonl.CodecZstd}, + parquet: &parquet.Config{}, + loadable: true, + }, + } { + t.Run(testCase.name, func(t *testing.T) { + harness := newRetIntegrationHarness(t) + fixture := harness.seedStandardGraph(t) + config := harness.dumpConfig(fixture.name, testCase.jsonl, testCase.parquet) + + dumpResult, err := ret.Dump(harness.ctx, harness.database, config) + if err != nil { + t.Fatalf("dump: %v", err) + } + assertOperationCounts(t, dumpResult.GraphCount, dumpResult.NodeCount, dumpResult.RelationshipCount, 1, 4, 3) + + verifyResult, err := ret.VerifyCollection(harness.ctx, ret.VerifyCollectionConfig{Directory: config.Directory}) + if err != nil { + t.Fatalf("verify collection: %v", err) + } + assertOperationCounts(t, verifyResult.GraphCount, verifyResult.NodeCount, verifyResult.RelationshipCount, 1, 4, 3) + assertConcreteOutputs(t, config.Directory, testCase.jsonl != nil, testCase.parquet != nil) + assertArchiveRoundTrip(t, config.Directory) + if testCase.jsonl != nil && testCase.parquet != nil { + damageFirstParquetArtifact(t, config.Directory) + } + + before := harness.snapshot(t, fixture.name) + loadResult, err := ret.Load(harness.ctx, harness.database, ret.LoadConfig{ + Directory: config.Directory, + BatchSize: 2, + }) + if !testCase.loadable { + if !errors.Is(err, ret.ErrCollectionNotLoadable) { + t.Fatalf("load parquet-only collection error = %v, want %v", err, ret.ErrCollectionNotLoadable) + } + if after := harness.snapshot(t, fixture.name); after != before { + t.Fatalf("parquet-only load mutated target: before=%+v after=%+v", before, after) + } + return + } + if !errors.Is(err, ret.ErrNonEmptyTarget) { + t.Fatalf("load nonempty target error = %v, want %v", err, ret.ErrNonEmptyTarget) + } + if after := harness.snapshot(t, fixture.name); after != before { + t.Fatalf("nonempty-target rejection mutated target: before=%+v after=%+v", before, after) + } + harness.assertStandardGraph(t, fixture) + + harness.clearGraph(t, fixture.name) + loadResult, err = ret.Load(harness.ctx, harness.database, ret.LoadConfig{ + Directory: config.Directory, + BatchSize: 2, + }) + if err != nil { + t.Fatalf("load: %v", err) + } + assertOperationCounts(t, loadResult.GraphCount, loadResult.NodeCount, loadResult.RelationshipCount, 1, 4, 3) + + databaseResult, err := ret.VerifyDatabase(harness.ctx, harness.database, ret.VerifyDatabaseConfig{ + Directory: config.Directory, + BatchSize: 2, + }) + if err != nil { + t.Fatalf("verify database: %v", err) + } + assertOperationCounts(t, databaseResult.GraphCount, databaseResult.NodeCount, databaseResult.RelationshipCount, 1, 4, 3) + harness.assertStandardGraph(t, fixture) + }) } - driverName, err := driverFromConnectionString(connection) - if err != nil { - t.Fatalf("infer driver: %v", err) +} + +func TestRetFacadeDumpResume(t *testing.T) { + harness := newRetIntegrationHarness(t) + fixture := harness.seedStandardGraph(t) + config := harness.dumpConfig( + fixture.name, + &jsonl.Config{Codec: jsonl.CodecZstd}, + &parquet.Config{}, + ) + config.ShardSize = 2 + + interruptedConfig, err := interruptDumpAfterFirstNodeShard(harness.ctx, harness.database, config) + if !errors.Is(err, context.Canceled) { + t.Fatalf("interrupted dump error = %v, want %v", err, context.Canceled) } - ctx := context.Background() - db, _, err := openDatabase(ctx, databaseConfig{ - Connection: connection, - }) + + interruptedConfig.Resume = true + result, err := ret.Dump(harness.ctx, harness.database, interruptedConfig) if err != nil { - t.Fatalf("open database: %v", err) + t.Fatalf("resume dump: %v", err) + } + assertOperationCounts(t, result.GraphCount, result.NodeCount, result.RelationshipCount, 1, 4, 3) + if _, err := ret.VerifyCollection(harness.ctx, ret.VerifyCollectionConfig{Directory: config.Directory}); err != nil { + t.Fatalf("verify resumed collection: %v", err) + } +} + +func TestRetFacadeResumeRejectsChangedCounts(t *testing.T) { + for _, testCase := range []struct { + name string + mutate func(*testing.T, *retIntegrationHarness, seededGraph) + }{ + {name: "node total", mutate: mutateNodeTotal}, + {name: "relationship total", mutate: mutateRelationshipTotal}, + } { + t.Run(testCase.name, func(t *testing.T) { + harness := newRetIntegrationHarness(t) + fixture := harness.seedStandardGraph(t) + config := harness.dumpConfig( + fixture.name, + &jsonl.Config{Codec: jsonl.CodecZstd}, + nil, + ) + config.ShardSize = 2 + + interruptedConfig, err := interruptDumpAfterFirstNodeShard(harness.ctx, harness.database, config) + if !errors.Is(err, context.Canceled) { + t.Fatalf("interrupted dump error = %v, want %v", err, context.Canceled) + } + testCase.mutate(t, harness, fixture) + + interruptedConfig.Resume = true + if _, err := ret.Dump(harness.ctx, harness.database, interruptedConfig); !errors.Is(err, ret.ErrSourceCountChanged) { + t.Fatalf("resume after %s change error = %v, want %v", testCase.name, err, ret.ErrSourceCountChanged) + } + }) + } +} + +func TestRetFacadeScrubbedDualOutput(t *testing.T) { + harness := newRetIntegrationHarness(t) + fixture := harness.seedScrubGraph(t) + scrubConfig := scrub.DefaultConfig() + scrubConfig.Salt = "ret-integration-salt" + config := harness.dumpConfig( + fixture.name, + &jsonl.Config{Codec: jsonl.CodecZstd}, + &parquet.Config{}, + ) + config.Scrub = &scrubConfig + + if _, err := ret.Dump(harness.ctx, harness.database, config); err != nil { + t.Fatalf("dump scrubbed collection: %v", err) + } + if _, err := ret.VerifyCollection(harness.ctx, ret.VerifyCollectionConfig{Directory: config.Directory}); err != nil { + t.Fatalf("verify scrubbed collection: %v", err) } - defer db.Close(ctx) - graphName := fmt.Sprintf("retriever_it_%d", time.Now().UTC().UnixNano()) - userKind := graph.StringKind("RetrieverUser") - systemKind := graph.StringKind("RetrieverSystem") - adminKind := graph.StringKind("RetrieverAdminTo") - graphSchema := graph.Graph{ - Name: graphName, - Nodes: graph.Kinds{userKind, systemKind}, - Edges: graph.Kinds{adminKind}, + artifacts := readConcreteArtifacts(t, config.Directory) + if !reflect.DeepEqual(normalizeNodes(t, artifacts.jsonlNodes), normalizeNodes(t, artifacts.parquetNodes)) { + t.Fatalf("JSONL and Parquet node values differ:\nJSONL: %#v\nParquet: %#v", artifacts.jsonlNodes, artifacts.parquetNodes) } - if err := db.AssertSchema(ctx, graph.Schema{ - Graphs: []graph.Graph{graphSchema}, - DefaultGraph: graphSchema, + if !reflect.DeepEqual( + normalizeRelationships(t, artifacts.jsonlRelationships), + normalizeRelationships(t, artifacts.parquetRelationships), + ) { + t.Fatalf("JSONL and Parquet relationship values differ:\nJSONL: %#v\nParquet: %#v", artifacts.jsonlRelationships, artifacts.parquetRelationships) + } + assertScrubFixture(t, artifacts, fixture.kinds) + + harness.clearGraph(t, fixture.name) + if _, err := ret.Load(harness.ctx, harness.database, ret.LoadConfig{ + Directory: config.Directory, + BatchSize: 2, }); err != nil { - t.Fatalf("assert schema: %v", err) + t.Fatalf("load scrubbed JSONL: %v", err) } - clearGraph := func() error { - return db.WriteTransaction(ctx, func(tx graph.Transaction) error { - return tx.WithGraph(graph.Graph{ - Name: graphName, - }).Nodes().Delete() - }) + if _, err := ret.VerifyDatabase(harness.ctx, harness.database, ret.VerifyDatabaseConfig{ + Directory: config.Directory, + BatchSize: 2, + }); err != nil { + t.Fatalf("verify scrubbed database: %v", err) } - if err := clearGraph(); err != nil { - t.Fatalf("clear graph before seed: %v", err) + harness.assertScrubGraph(t, fixture, artifacts) +} + +func assertArchiveRoundTrip(t *testing.T, collectionDirectory string) { + t.Helper() + t.Run("archive", func(t *testing.T) { + if runtime.GOOS != "linux" && runtime.GOOS != "darwin" { + t.Skip("ret archive publication is supported only on Linux and Darwin") + } + + archivePath, unpackedDirectory, recipient, identity := newArchiveFixture(t) + if err := ret.Pack(context.Background(), ret.PackConfig{ + CollectionDirectory: collectionDirectory, + ArchivePath: archivePath, + Recipient: recipient, + }); err != nil { + t.Fatalf("pack collection: %v", err) + } + if err := ret.Unpack(context.Background(), ret.UnpackConfig{ + ArchivePath: archivePath, + OutputDirectory: unpackedDirectory, + Identity: identity, + }); err != nil { + t.Fatalf("unpack collection: %v", err) + } + if _, err := ret.VerifyCollection(context.Background(), ret.VerifyCollectionConfig{Directory: unpackedDirectory}); err != nil { + t.Fatalf("verify unpacked collection: %v", err) + } + }) +} + +var retIntegrationSequence atomic.Uint64 + +type retIntegrationHarness struct { + ctx context.Context + database graph.Database + root string +} + +type seededGraph struct { + name string + nodeIDs []graph.ID + nodeKinds map[string][]string + nodeRoles map[string]string + relationshipKind string + kinds []string +} + +type concreteArtifacts struct { + jsonlNodes []entity.Node + parquetNodes []entity.Node + jsonlRelationships []entity.Relationship + parquetRelationships []entity.Relationship + parquetRelationshipIDs []string + scrubCounts scrub.ActionCounts + scrubMetadata collection.ScrubMetadata +} + +func newRetIntegrationHarness(t *testing.T) *retIntegrationHarness { + t.Helper() + connection := os.Getenv("CONNECTION_STRING") + if connection == "" { + t.Skip("CONNECTION_STRING not set") + } + + ctx := context.Background() + database, _, err := openDatabase(ctx, databaseConfig{Connection: connection}) + if err != nil { + t.Fatalf("open database: %v", err) } t.Cleanup(func() { - _ = clearGraph() + closeCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + if err := database.Close(closeCtx); err != nil { + t.Errorf("close database: %v", err) + } }) - const ( - nodeCount = 7 - edgeCount = 5 - batchSize = 2 - shardSize = 3 - ) - var seededNodeIDs, seededEdgeIDs []graph.ID - if err := db.WriteTransaction(ctx, func(tx graph.Transaction) error { - tx = tx.WithGraph(graph.Graph{ - Name: graphName, - }) + return &retIntegrationHarness{ + ctx: ctx, + database: database, + root: t.TempDir(), + } +} + +func (s *retIntegrationHarness) graphName() string { + return fmt.Sprintf("ret_it_%d_%d", time.Now().UTC().UnixNano(), retIntegrationSequence.Add(1)) +} + +func (s *retIntegrationHarness) dumpConfig(graphName string, jsonlConfig *jsonl.Config, parquetConfig *parquet.Config) ret.DumpConfig { + return ret.DumpConfig{ + Directory: filepath.Join(s.root, fmt.Sprintf("collection-%d", retIntegrationSequence.Add(1))), + Graphs: []string{graphName}, + EntityBatchSize: 2, + ShardSize: 2, + JSONL: jsonlConfig, + Parquet: parquetConfig, + } +} + +func (s *retIntegrationHarness) seedStandardGraph(t *testing.T) seededGraph { + t.Helper() + graphName := s.graphName() + userKind := graph.StringKind("RetIntegrationUser") + systemKind := graph.StringKind("RetIntegrationSystem") + relationshipKind := graph.StringKind("RetIntegrationLink") + s.assertSchema(t, graphName, graph.Kinds{userKind, systemKind}, graph.Kinds{relationshipKind}) + s.clearGraph(t, graphName) + t.Cleanup(func() { + cleanupCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + if err := s.clearGraphError(cleanupCtx, graphName); err != nil { + t.Errorf("clean up graph %q: %v", graphName, err) + } + }) - nodes := make([]*graph.Node, 0, nodeCount) - for index := range nodeCount { - kind := userKind + fixture := seededGraph{ + name: graphName, + nodeKinds: make(map[string][]string, 4), + nodeRoles: make(map[string]string, 4), + relationshipKind: relationshipKind.String(), + } + if err := s.database.WriteTransaction(s.ctx, func(tx graph.Transaction) error { + tx = tx.WithGraph(graph.Graph{Name: graphName}) + nodes := make([]*graph.Node, 0, 4) + for index := range 4 { + kinds := graph.Kinds{userKind} if index%2 == 1 { - kind = systemKind + kinds = graph.Kinds{systemKind} } - node, err := tx.CreateNode(graph.AsProperties(map[string]any{ "name": fmt.Sprintf("node-%d", index), - "role": fmt.Sprintf("role-%d", index%3), - }), kind) + "role": fmt.Sprintf("role-%d", index%2), + }), kinds...) if err != nil { return err } nodes = append(nodes, node) - seededNodeIDs = append(seededNodeIDs, node.ID) + fixture.nodeIDs = append(fixture.nodeIDs, node.ID) + name := fmt.Sprintf("node-%d", index) + fixture.nodeKinds[name] = kinds.Strings() + fixture.nodeRoles[name] = fmt.Sprintf("role-%d", index%2) } - - for index := range edgeCount { - relationship, err := tx.CreateRelationshipByIDs(nodes[index].ID, nodes[index+1].ID, adminKind, graph.AsProperties(map[string]any{ - "route": fmt.Sprintf("route-%d", index), - })) - if err != nil { + for index := range 3 { + if _, err := tx.CreateRelationshipByIDs( + nodes[index].ID, + nodes[index+1].ID, + relationshipKind, + graph.AsProperties(map[string]any{"route": fmt.Sprintf("route-%d", index)}), + ); err != nil { return err } - seededEdgeIDs = append(seededEdgeIDs, relationship.ID) } - return nil }); err != nil { - t.Fatalf("seed graph: %v", err) + t.Fatalf("seed standard graph: %v", err) } + return fixture +} - entitySnapshot, err := countGraphEntitySnapshot(ctx, db, graph.Graph{ - Name: graphName, +func (s *retIntegrationHarness) seedScrubGraph(t *testing.T) seededGraph { + t.Helper() + graphName := s.graphName() + firstKind := graph.StringKind("RetIntegrationFirst") + secondKind := graph.StringKind("RetIntegrationSecond") + relationshipKind := graph.StringKind("RetIntegrationScrubLink") + s.assertSchema(t, graphName, graph.Kinds{firstKind, secondKind}, graph.Kinds{relationshipKind}) + s.clearGraph(t, graphName) + t.Cleanup(func() { + cleanupCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + if err := s.clearGraphError(cleanupCtx, graphName); err != nil { + t.Errorf("clean up graph %q: %v", graphName, err) + } }) - if err != nil { - t.Fatalf("count graph entities: %v", err) - } - if entitySnapshot.NodeCount != nodeCount || entitySnapshot.EdgeCount != edgeCount { - t.Fatalf("unexpected seeded graph counts: nodes=%d edges=%d", entitySnapshot.NodeCount, entitySnapshot.EdgeCount) - } - - assertBoundedRetrieverScans(t, ctx, db, graph.Graph{Name: graphName}, entitySnapshot, batchSize) - assertSkipAndLimit(t, ctx, db, graph.Graph{Name: graphName}, seededNodeIDs, seededEdgeIDs) - dumpDir := t.TempDir() - dumpResult, err := retriever.Dump(ctx, db, driverName, []retriever.GraphTarget{{ - Name: graphName, - }}, retriever.DumpOptions{ - OutputDir: dumpDir, - Scrub: retriever.ScrubNone, - Compression: retriever.CompressionGzip, - ZstdLevel: retriever.DefaultZstdLevel, - ShardSize: shardSize, - BatchSize: batchSize, - }) - if err != nil { - t.Fatalf("dump: %v", err) - } - if dumpResult.NodeCount != nodeCount || dumpResult.EdgeCount != edgeCount { - t.Fatalf("unexpected dump counts: nodes=%d edges=%d", dumpResult.NodeCount, dumpResult.EdgeCount) - } - if got := len(dumpResult.Manifest.Graphs); got != 1 { - t.Fatalf("dump manifest graph count = %d", got) - } - graphEntry := dumpResult.Manifest.Graphs[0] - if graphEntry.NodeCount != nodeCount || graphEntry.EdgeCount != edgeCount { - t.Fatalf("unexpected manifest graph counts: nodes=%d edges=%d", graphEntry.NodeCount, graphEntry.EdgeCount) + orderedDuplicateKinds := graph.Kinds{secondKind, firstKind, secondKind} + graphNodes := []*graph.Node{ + graph.NewNode(0, graph.AsProperties(map[string]any{ + "enabled": true, + "preserved_count": int64(42), + "email": "alice@example.com", + "password": "super-secret", + "created_at": "2026-01-01T00:00:00Z", + }), orderedDuplicateKinds...), + graph.NewNode(0, graph.AsProperties(map[string]any{}), firstKind), } - var nodeFiles, edgeFiles int - for _, fileEntry := range graphEntry.Files { - switch fileEntry.Phase { - case retriever.PhaseNodes: - nodeFiles++ - case retriever.PhaseEdges: - edgeFiles++ - } - if fileEntry.Count > shardSize { - t.Fatalf("fragment %s count %d exceeds shard size %d", fileEntry.Path, fileEntry.Count, shardSize) + var nodeIDs []graph.ID + if err := s.database.BatchOperation(s.ctx, func(batch graph.Batch) error { + creator, ok := batch.WithGraph(graph.Graph{Name: graphName}).(graph.NodeBatchCreator) + if !ok { + return errors.New("database batch does not support correlated node creation") } + var err error + nodeIDs, err = creator.CreateNodes(graphNodes) + return err + }, graph.WithBatchSize(2)); err != nil { + t.Fatalf("seed scrub nodes: %v", err) + } + if len(nodeIDs) != 2 { + t.Fatalf("seed scrub node IDs = %d, want 2", len(nodeIDs)) + } + if err := s.database.WriteTransaction(s.ctx, func(tx graph.Transaction) error { + _, err := tx.WithGraph(graph.Graph{Name: graphName}).CreateRelationshipByIDs( + nodeIDs[0], + nodeIDs[1], + relationshipKind, + graph.AsProperties(map[string]any{}), + ) + return err + }); err != nil { + t.Fatalf("seed scrub relationship: %v", err) } - if nodeFiles != 3 || edgeFiles != 2 { - t.Fatalf("unexpected manifest shards: node files=%d edge files=%d", nodeFiles, edgeFiles) - } - if dumpResult.Manifest.Metrics == nil { - t.Fatalf("dump manifest is missing metrics") - } - if got := len(dumpResult.Manifest.Metrics.Graphs); got != 1 { - t.Fatalf("dump metrics graph count = %d", got) + + nodes, _ := s.fetchGraph(t, graphName) + var observedKinds []string + for _, node := range nodes { + if node.Properties != nil && node.Properties.MapOrEmpty()["enabled"] == true { + observedKinds = node.Kinds.Strings() + break + } } - if dumpResult.Manifest.Metrics.Graphs[0].NodeCount != nodeCount || dumpResult.Manifest.Metrics.Graphs[0].EdgeCount != edgeCount { - t.Fatalf("unexpected dump metrics counts: %+v", dumpResult.Manifest.Metrics.Graphs[0]) + if observedKinds == nil { + t.Fatal("seeded scrub node was not returned by the database") } + return seededGraph{name: graphName, nodeIDs: nodeIDs, kinds: observedKinds} +} - if _, err := retriever.Load(ctx, db, driverName, retriever.LoadOptions{ - InputDir: dumpDir, - BatchSize: batchSize, - VerifyMetrics: true, - }); err == nil || !strings.Contains(err.Error(), "is not empty") { - t.Fatalf("expected non-empty target load error, got %v", err) +func (s *retIntegrationHarness) assertSchema(t *testing.T, graphName string, nodeKinds, relationshipKinds graph.Kinds) { + t.Helper() + target := graph.Graph{Name: graphName, Nodes: nodeKinds, Edges: relationshipKinds} + if err := s.database.AssertSchema(s.ctx, graph.Schema{ + Graphs: []graph.Graph{target}, + DefaultGraph: target, + }); err != nil { + t.Fatalf("assert graph schema: %v", err) } +} - if err := clearGraph(); err != nil { - t.Fatalf("clear graph before load: %v", err) +func (s *retIntegrationHarness) clearGraph(t *testing.T, graphName string) { + t.Helper() + if err := s.clearGraphError(s.ctx, graphName); err != nil { + t.Fatalf("clear graph %q: %v", graphName, err) } +} - loadResult, err := retriever.Load(ctx, db, driverName, retriever.LoadOptions{ - InputDir: dumpDir, - BatchSize: batchSize, - VerifyMetrics: true, +func (s *retIntegrationHarness) clearGraphError(ctx context.Context, graphName string) error { + return s.database.WriteTransaction(ctx, func(tx graph.Transaction) error { + return tx.WithGraph(graph.Graph{Name: graphName}).Nodes().Delete() }) +} + +func (s *retIntegrationHarness) snapshot(t *testing.T, graphName string) dawgs.Snapshot { + t.Helper() + source, err := dawgs.NewSource(s.database, graphName, 2) if err != nil { - t.Fatalf("load: %v", err) + t.Fatalf("create graph source: %v", err) } - if loadResult.NodeCount != nodeCount || loadResult.EdgeCount != edgeCount { - t.Fatalf("unexpected load counts: nodes=%d edges=%d", loadResult.NodeCount, loadResult.EdgeCount) + snapshot, err := source.Snapshot(s.ctx) + if err != nil { + t.Fatalf("snapshot graph: %v", err) } + return snapshot +} - var loadedNodes []*graph.Node - var loadedEdges []*graph.Relationship - if err := db.ReadTransaction(ctx, func(tx graph.Transaction) error { +func (s *retIntegrationHarness) fetchGraph(t *testing.T, graphName string) ([]*graph.Node, []*graph.Relationship) { + t.Helper() + var nodes []*graph.Node + var relationships []*graph.Relationship + if err := s.database.ReadTransaction(s.ctx, func(tx graph.Transaction) error { tx = tx.WithGraph(graph.Graph{Name: graphName}) var err error - if loadedNodes, err = ops.FetchNodes(tx.Nodes()); err != nil { + if nodes, err = ops.FetchNodes(tx.Nodes()); err != nil { return err } - loadedEdges, err = ops.FetchRelationships(tx.Relationships()) + relationships, err = ops.FetchRelationships(tx.Relationships()) return err }); err != nil { - t.Fatalf("read restored topology: %v", err) - } - if len(loadedNodes) != nodeCount || len(loadedEdges) != edgeCount { - t.Fatalf("unexpected restored topology counts: nodes=%d edges=%d", len(loadedNodes), len(loadedEdges)) + t.Fatalf("fetch graph %q: %v", graphName, err) } + return nodes, relationships +} - nodeNames := make(map[graph.ID]string, len(loadedNodes)) - expectedRoles := make(map[string]string, nodeCount) - expectedKinds := make(map[string]string, nodeCount) - for index := range nodeCount { - name := fmt.Sprintf("node-%d", index) - expectedRoles[name] = fmt.Sprintf("role-%d", index%3) - expectedKinds[name] = userKind.String() - if index%2 == 1 { - expectedKinds[name] = systemKind.String() - } +func (s *retIntegrationHarness) assertStandardGraph(t *testing.T, fixture seededGraph) { + t.Helper() + nodes, relationships := s.fetchGraph(t, fixture.name) + if len(nodes) != 4 || len(relationships) != 3 { + t.Fatalf("loaded graph counts: nodes=%d relationships=%d, want 4 and 3", len(nodes), len(relationships)) } - for _, node := range loadedNodes { - if node.Properties == nil { - t.Fatalf("restored node %d is missing properties", node.ID) - } - name := fmt.Sprint(node.Properties.Get("name").Any()) - role := fmt.Sprint(node.Properties.Get("role").Any()) - nodeNames[node.ID] = name - if expectedRole, ok := expectedRoles[name]; !ok || role != expectedRole { - t.Fatalf("unexpected restored node properties: %+v", node.Properties.MapOrEmpty()) + + namesByID := make(map[graph.ID]string, len(nodes)) + for _, node := range nodes { + properties := node.Properties.MapOrEmpty() + name := fmt.Sprint(properties["name"]) + namesByID[node.ID] = name + if got, want := node.Kinds.Strings(), fixture.nodeKinds[name]; !reflect.DeepEqual(got, want) { + t.Fatalf("node %q kinds = %v, want %v", name, got, want) } - if len(node.Kinds) != 1 || node.Kinds[0].String() != expectedKinds[name] { - t.Fatalf("unexpected restored node kinds for %s: %v", name, node.Kinds.Strings()) + wantRole := fixture.nodeRoles[name] + if got := fmt.Sprint(properties["role"]); got != wantRole { + t.Fatalf("node %q role = %q, want %q", name, got, wantRole) } } - restoredRoutes := map[string]string{} - for _, relationship := range loadedEdges { - startName, startOK := nodeNames[relationship.StartID] - endName, endOK := nodeNames[relationship.EndID] - if !startOK || !endOK { - t.Fatalf("restored relationship has unresolved endpoint IDs: %+v", relationship) - } - if relationship.Kind == nil || relationship.Kind.String() != adminKind.String() { - t.Fatalf("unexpected restored relationship kind: %+v", relationship.Kind) + routes := make(map[string]string, len(relationships)) + for _, relationship := range relationships { + if relationship.Kind == nil || relationship.Kind.String() != fixture.relationshipKind { + t.Fatalf("relationship kind = %v, want %q", relationship.Kind, fixture.relationshipKind) } - if relationship.Properties == nil { - t.Fatalf("restored relationship is missing properties: %+v", relationship) - } - restoredRoutes[startName+"->"+endName] = fmt.Sprint(relationship.Properties.Get("route").Any()) + key := namesByID[relationship.StartID] + "->" + namesByID[relationship.EndID] + routes[key] = fmt.Sprint(relationship.Properties.MapOrEmpty()["route"]) } - for index := range edgeCount { - path := fmt.Sprintf("node-%d->node-%d", index, index+1) - if route := restoredRoutes[path]; route != fmt.Sprintf("route-%d", index) { - t.Fatalf("restored route %s = %q", path, route) + for index := range 3 { + key := fmt.Sprintf("node-%d->node-%d", index, index+1) + if got, want := routes[key], fmt.Sprintf("route-%d", index); got != want { + t.Fatalf("relationship %q route = %q, want %q", key, got, want) } } +} - verifyResult, err := retriever.Verify(ctx, db, driverName, retriever.VerifyOptions{ - InputDir: dumpDir, - BatchSize: batchSize, - }) - if err != nil { - t.Fatalf("verify: %v", err) +func (s *retIntegrationHarness) assertScrubGraph(t *testing.T, fixture seededGraph, artifacts concreteArtifacts) { + t.Helper() + nodes, relationships := s.fetchGraph(t, fixture.name) + if len(nodes) != 2 || len(relationships) != 1 { + t.Fatalf("loaded scrub graph counts: nodes=%d relationships=%d, want 2 and 1", len(nodes), len(relationships)) + } + var loaded *graph.Node + for _, node := range nodes { + if node.Properties != nil && node.Properties.MapOrEmpty()["enabled"] == true { + loaded = node + break + } } - if verifyResult.NodeCount != nodeCount || verifyResult.EdgeCount != edgeCount { - t.Fatalf("unexpected verify counts: nodes=%d edges=%d", verifyResult.NodeCount, verifyResult.EdgeCount) + if loaded == nil { + t.Fatal("loaded scrub graph is missing the primary node") } + if got := loaded.Kinds.Strings(); !reflect.DeepEqual(got, fixture.kinds) { + t.Fatalf("loaded ordered kinds = %v, want %v", got, fixture.kinds) + } + artifactNode := primaryArtifactNode(t, artifacts.jsonlNodes) + if got, want := normalizeProperties(t, loaded.Properties.MapOrEmpty()), normalizeProperties(t, artifactNode.Properties); !reflect.DeepEqual(got, want) { + t.Fatalf("loaded scrub properties = %#v, want %#v", got, want) + } +} - if err := db.WriteTransaction(ctx, func(tx graph.Transaction) error { - tx = tx.WithGraph(graph.Graph{ - Name: graphName, - }) - _, err := tx.CreateNode(graph.AsProperties(map[string]any{"name": "extra"}), userKind) +func interruptDumpAfterFirstNodeShard( + parent context.Context, + database graph.Database, + config ret.DumpConfig, +) (ret.DumpConfig, error) { + ctx, cancel := context.WithCancel(parent) + defer cancel() + config.Observer = &cancelOnNodeShardObserver{cancel: cancel} + _, err := ret.Dump(ctx, database, config) + config.Observer = nil + return config, err +} + +type cancelOnNodeShardObserver struct { + cancel context.CancelFunc + once sync.Once +} + +func (s *cancelOnNodeShardObserver) Observe(_ context.Context, event observe.Event) { + if committed, ok := event.(observe.ShardCommitted); ok && committed.EntityType == "node" { + s.once.Do(s.cancel) + } +} + +func mutateNodeTotal(t *testing.T, harness *retIntegrationHarness, fixture seededGraph) { + t.Helper() + if err := harness.database.WriteTransaction(harness.ctx, func(tx graph.Transaction) error { + _, err := tx.WithGraph(graph.Graph{Name: fixture.name}).CreateNode( + graph.AsProperties(map[string]any{"name": "count-change"}), + graph.StringKind("RetIntegrationUser"), + ) return err }); err != nil { - t.Fatalf("mutate loaded graph: %v", err) + t.Fatalf("mutate node total: %v", err) } +} - _, err = retriever.Verify(ctx, db, driverName, retriever.VerifyOptions{ - InputDir: dumpDir, - BatchSize: batchSize, - }) - var mismatch retriever.MetricsMismatchError - if !errors.As(err, &mismatch) { - t.Fatalf("expected metrics mismatch error, got %v", err) - } - if !strings.Contains(err.Error(), "node_count") { - t.Fatalf("expected mismatch to include node_count, got %v", err) +func mutateRelationshipTotal(t *testing.T, harness *retIntegrationHarness, fixture seededGraph) { + t.Helper() + if err := harness.database.WriteTransaction(harness.ctx, func(tx graph.Transaction) error { + _, err := tx.WithGraph(graph.Graph{Name: fixture.name}).CreateRelationshipByIDs( + fixture.nodeIDs[0], + fixture.nodeIDs[len(fixture.nodeIDs)-1], + graph.StringKind("RetIntegrationLink"), + graph.AsProperties(map[string]any{"route": "count-change"}), + ) + return err + }); err != nil { + t.Fatalf("mutate relationship total: %v", err) } } -func assertBoundedRetrieverScans(t *testing.T, ctx context.Context, db graph.Database, targetGraph graph.Graph, snapshot graphEntitySnapshot, batchSize int) { +func assertOperationCounts( + t *testing.T, + graphs int, + nodes, relationships int64, + wantGraphs int, + wantNodes, wantRelationships int64, +) { t.Helper() - - var nodeIDs []graph.ID - processed, err := retriever.ScanDatabaseNodes(ctx, db, targetGraph, snapshot.NodeCount, batchSize, func(node *graph.Node) error { - nodeIDs = append(nodeIDs, node.ID) - return nil - }, func(event retriever.ScanBatchEvent) error { - if event.Count > batchSize { - return fmt.Errorf("node cursor callback count %d exceeds batch size %d", event.Count, batchSize) - } - return nil - }) - if err != nil || processed != snapshot.NodeCount { - t.Fatalf("bounded node scan processed=%d err=%v", processed, err) + if graphs != wantGraphs || nodes != wantNodes || relationships != wantRelationships { + t.Fatalf( + "operation counts: graphs=%d nodes=%d relationships=%d, want %d %d %d", + graphs, + nodes, + relationships, + wantGraphs, + wantNodes, + wantRelationships, + ) } - assertStrictIDs(t, "node", nodeIDs) +} - var edgeIDs []graph.ID - processed, err = retriever.ScanDatabaseRelationships(ctx, db, targetGraph, snapshot.EdgeCount, batchSize, func(relationship *graph.Relationship) error { - edgeIDs = append(edgeIDs, relationship.ID) - return nil - }, func(event retriever.ScanBatchEvent) error { - if event.Count > batchSize { - return fmt.Errorf("relationship cursor callback count %d exceeds batch size %d", event.Count, batchSize) +func assertConcreteOutputs(t *testing.T, root string, wantJSONL, wantParquet bool) { + t.Helper() + manifest, err := collection.Read(root) + if err != nil { + t.Fatalf("read collection manifest: %v", err) + } + if got := manifest.Outputs.JSONL != nil; got != wantJSONL { + t.Fatalf("manifest JSONL enabled = %t, want %t", got, wantJSONL) + } + if got := manifest.Outputs.Parquet != nil; got != wantParquet { + t.Fatalf("manifest Parquet enabled = %t, want %t", got, wantParquet) + } + for _, graphEntry := range manifest.Graphs { + for _, shard := range graphEntry.NodeShards { + if (shard.JSONL != nil) != wantJSONL || (shard.Parquet != nil) != wantParquet { + t.Fatalf("node shard %d concrete outputs do not match collection capabilities", shard.Index) + } + } + for _, shard := range graphEntry.RelationshipShards { + if (shard.JSONL != nil) != wantJSONL || (shard.Parquet != nil) != wantParquet { + t.Fatalf("relationship shard %d concrete outputs do not match collection capabilities", shard.Index) + } } - return nil - }) - if err != nil || processed != snapshot.EdgeCount { - t.Fatalf("bounded relationship scan processed=%d err=%v", processed, err) } - assertStrictIDs(t, "relationship", edgeIDs) } -func assertStrictIDs(t *testing.T, entityName string, ids []graph.ID) { +func damageFirstParquetArtifact(t *testing.T, root string) { t.Helper() - for index := 1; index < len(ids); index++ { - if ids[index] <= ids[index-1] { - t.Fatalf("%s IDs are not strictly increasing: %v", entityName, ids) + manifest, err := collection.Read(root) + if err != nil { + t.Fatalf("read dual-output manifest before damaging Parquet: %v", err) + } + for _, graphEntry := range manifest.Graphs { + for _, shard := range graphEntry.NodeShards { + if shard.Parquet != nil { + path := filepath.Join(root, filepath.FromSlash(shard.Parquet.Path)) + if err := os.WriteFile(path, []byte("intentionally damaged Parquet"), 0o600); err != nil { + t.Fatalf("damage Parquet artifact: %v", err) + } + return + } } } + t.Fatal("dual-output collection has no Parquet artifact to damage") } -func assertSkipAndLimit(t *testing.T, ctx context.Context, db graph.Database, targetGraph graph.Graph, nodeIDs, edgeIDs []graph.ID) { +func newArchiveFixture(t *testing.T) (string, string, archive.PublicKey, archive.PrivateKey) { t.Helper() + root := t.TempDir() + privatePath := filepath.Join(root, "identity.key") + publicPath := filepath.Join(root, "recipient.key") + if err := ret.Keygen(ret.KeygenConfig{ + PrivateKeyPath: privatePath, + PublicKeyPath: publicPath, + }); err != nil { + t.Fatalf("generate archive keys: %v", err) + } + recipient, err := archive.ReadPublicKey(publicPath) + if err != nil { + t.Fatalf("read archive recipient: %v", err) + } + identity, err := archive.ReadPrivateKey(privatePath) + if err != nil { + t.Fatalf("read archive identity: %v", err) + } + return filepath.Join(root, "collection.ret.enc"), filepath.Join(root, "unpacked"), recipient, identity +} - if err := db.ReadTransaction(ctx, func(tx graph.Transaction) error { - tx = tx.WithGraph(targetGraph) - var actualNodeIDs []graph.ID - if err := tx.Nodes().OrderBy(query.NodeID()).Offset(1).Limit(2).FetchIDs(func(cursor graph.Cursor[graph.ID]) error { - for id := range cursor.Chan() { - actualNodeIDs = append(actualNodeIDs, id) - } - return cursor.Error() +func readConcreteArtifacts(t *testing.T, root string) concreteArtifacts { + t.Helper() + manifest, err := collection.Read(root) + if err != nil { + t.Fatalf("read scrub manifest: %v", err) + } + if len(manifest.Graphs) != 1 { + t.Fatalf("scrub manifest graph count = %d, want 1", len(manifest.Graphs)) + } + graphEntry := manifest.Graphs[0] + result := concreteArtifacts{ + scrubCounts: scrub.ActionCounts{}, + scrubMetadata: manifest.Scrub, + } + for _, shard := range graphEntry.NodeShards { + result.scrubCounts.Add(shard.ScrubCounts) + if shard.JSONL == nil || shard.Parquet == nil { + t.Fatalf("dual node shard %d is missing a concrete artifact", shard.Index) + } + var nodes []entity.Node + err := collection.ReadJSONLNodes(root, *shard.JSONL, func(node entity.Node) error { + nodes = append(nodes, node) + return nil + }) + if err != nil { + t.Fatalf("read JSONL node shard %d: %v", shard.Index, err) + } + result.jsonlNodes = append(result.jsonlNodes, nodes...) + if err := collection.ReadParquetNodes(root, *shard.Parquet, func(node entity.Node) error { + result.parquetNodes = append(result.parquetNodes, node) + return nil }); err != nil { - return err + t.Fatalf("read Parquet node shard %d: %v", shard.Index, err) + } + } + for _, shard := range graphEntry.RelationshipShards { + result.scrubCounts.Add(shard.ScrubCounts) + if shard.JSONL == nil || shard.Parquet == nil { + t.Fatalf("dual relationship shard %d is missing a concrete artifact", shard.Index) } - if !reflect.DeepEqual(actualNodeIDs, nodeIDs[1:3]) { - return fmt.Errorf("node skip/limit IDs = %v, want %v", actualNodeIDs, nodeIDs[1:3]) + var relationships []entity.Relationship + err := collection.ReadJSONLRelationships(root, *shard.JSONL, func(relationship entity.Relationship) error { + relationships = append(relationships, relationship) + return nil + }) + if err != nil { + t.Fatalf("read JSONL relationship shard %d: %v", shard.Index, err) } - - var actualEdgeIDs []graph.ID - if err := tx.Relationships().OrderBy(query.RelationshipID()).Offset(1).Limit(2).FetchIDs(func(cursor graph.Cursor[graph.ID]) error { - for id := range cursor.Chan() { - actualEdgeIDs = append(actualEdgeIDs, id) - } - return cursor.Error() + result.jsonlRelationships = append(result.jsonlRelationships, relationships...) + if err := collection.ReadParquetRelationships(root, *shard.Parquet, func(relationship entity.Relationship) error { + result.parquetRelationshipIDs = append(result.parquetRelationshipIDs, relationship.SourceID) + relationship.SourceID = "" + result.parquetRelationships = append(result.parquetRelationships, relationship) + return nil }); err != nil { - return err + t.Fatalf("read Parquet relationship shard %d: %v", shard.Index, err) } - if !reflect.DeepEqual(actualEdgeIDs, edgeIDs[1:3]) { - return fmt.Errorf("relationship skip/limit IDs = %v, want %v", actualEdgeIDs, edgeIDs[1:3]) + } + return result +} + +func assertScrubFixture(t *testing.T, artifacts concreteArtifacts, wantKinds []string) { + t.Helper() + if !artifacts.scrubMetadata.Enabled { + t.Fatal("scrub metadata is not enabled") + } + wantCounts := scrub.ActionCounts{ + Preserve: 2, + Pseudonymize: 1, + Redact: 1, + ShiftTimestamp: 1, + } + if !reflect.DeepEqual(artifacts.scrubCounts, wantCounts) { + t.Fatalf("scrub action counts = %#v, want %#v", artifacts.scrubCounts, wantCounts) + } + for index, relationship := range artifacts.jsonlRelationships { + if relationship.SourceID != "" { + t.Fatalf("JSONL relationship %d retained source ID %q", index, relationship.SourceID) + } + } + for index, sourceID := range artifacts.parquetRelationshipIDs { + if sourceID == "" { + t.Fatalf("Parquet relationship %d omitted its source ID", index) } + } - return nil - }); err != nil { - t.Fatalf("assert skip and limit: %v", err) + node := primaryArtifactNode(t, artifacts.jsonlNodes) + if !reflect.DeepEqual(node.Kinds, wantKinds) { + t.Fatalf("artifact ordered kinds = %v, want database-observed %v", node.Kinds, wantKinds) + } + if node.Properties["enabled"] != true { + t.Fatalf("preserved property = %#v, want true", node.Properties["enabled"]) + } + normalized := normalizeProperties(t, node.Properties) + if got, ok := normalized["preserved_count"].(json.Number); !ok || got.String() != "42" { + t.Fatalf( + "normalized preserved count = %#v (%T), want json.Number(%q)", + normalized["preserved_count"], + normalized["preserved_count"], + "42", + ) + } + if node.Properties["password"] != "[REDACTED]" { + t.Fatalf("redacted property = %#v, want [REDACTED]", node.Properties["password"]) + } + if node.Properties["created_at"] != "2026-01-18T00:00:00Z" { + t.Fatalf("shifted timestamp = %#v, want 2026-01-18T00:00:00Z", node.Properties["created_at"]) + } + if got := fmt.Sprint(node.Properties["email"]); got != "user-d47583d80c3e@example.invalid" { + t.Fatalf("pseudonymized email = %q, want user-d47583d80c3e@example.invalid", got) + } +} + +func primaryArtifactNode(t *testing.T, nodes []entity.Node) entity.Node { + t.Helper() + for _, node := range nodes { + if node.Properties["enabled"] == true { + return node + } + } + t.Fatal("artifact is missing the primary scrub node") + return entity.Node{} +} + +func normalizeNodes(t *testing.T, nodes []entity.Node) []entity.Node { + t.Helper() + normalized := append([]entity.Node(nil), nodes...) + for index := range normalized { + normalized[index].Properties = normalizeProperties(t, normalized[index].Properties) + } + return normalized +} + +func normalizeRelationships(t *testing.T, relationships []entity.Relationship) []entity.Relationship { + t.Helper() + normalized := append([]entity.Relationship(nil), relationships...) + for index := range normalized { + normalized[index].Properties = normalizeProperties(t, normalized[index].Properties) + } + return normalized +} + +func normalizeProperties(t *testing.T, properties map[string]any) map[string]any { + t.Helper() + encoded, err := json.Marshal(properties) + if err != nil { + t.Fatalf("encode properties for logical normalization: %v", err) + } + decoder := json.NewDecoder(bytes.NewReader(encoded)) + decoder.UseNumber() + var normalized map[string]any + if err := decoder.Decode(&normalized); err != nil { + t.Fatalf("decode properties for logical normalization: %v", err) } + return normalized } diff --git a/go.mod b/go.mod index 1f380c05..98bf8dc5 100644 --- a/go.mod +++ b/go.mod @@ -15,9 +15,11 @@ require ( github.com/jackc/pgx/v5 v5.10.0 github.com/klauspost/compress v1.19.0 github.com/neo4j/neo4j-go-driver/v5 v5.28.4 + github.com/parquet-go/parquet-go v0.30.1 github.com/pashagolub/pgxmock/v5 v5.1.0 github.com/pelletier/go-toml/v2 v2.4.3 github.com/stretchr/testify v1.11.1 + golang.org/x/sys v0.46.0 golang.org/x/tools v0.47.0 ) @@ -59,6 +61,7 @@ require ( github.com/alfatraining/structtag v1.0.0 // indirect github.com/alingse/asasalint v0.0.11 // indirect github.com/alingse/nilnesserr v0.2.0 // indirect + github.com/andybalholm/brotli v1.1.1 // indirect github.com/ashanbrown/forbidigo/v2 v2.3.0 // indirect github.com/ashanbrown/makezero/v2 v2.1.0 // indirect github.com/aymanbagabas/go-osc52/v2 v2.0.1 // indirect @@ -120,6 +123,7 @@ require ( github.com/golangci/swaggoswag v0.0.0-20250504205917-77f2aca3143e // indirect github.com/golangci/unconvert v0.0.0-20250410112200-a129a6e6413e // indirect github.com/google/go-cmp v0.7.0 // indirect + github.com/google/uuid v1.6.0 // indirect github.com/gordonklaus/ineffassign v0.2.0 // indirect github.com/gostaticanalysis/analysisutil v0.7.1 // indirect github.com/gostaticanalysis/comment v1.5.0 // indirect @@ -172,6 +176,9 @@ require ( github.com/nishanths/exhaustive v0.12.0 // indirect github.com/nishanths/predeclared v0.2.2 // indirect github.com/nunnatsa/ginkgolinter v0.23.0 // indirect + github.com/parquet-go/bitpack v1.0.0 // indirect + github.com/parquet-go/jsonlite v1.0.0 // indirect + github.com/pierrec/lz4/v4 v4.1.21 // indirect github.com/pkg/errors v0.9.1 // indirect github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect github.com/prometheus/client_golang v1.22.0 // indirect @@ -212,6 +219,7 @@ require ( github.com/timonwong/loggercheck v0.11.0 // indirect github.com/tomarrell/wrapcheck/v2 v2.12.0 // indirect github.com/tommy-muehle/go-mnd/v2 v2.5.1 // indirect + github.com/twpayne/go-geom v1.6.1 // indirect github.com/ultraware/funlen v0.2.0 // indirect github.com/ultraware/whitespace v0.2.0 // indirect github.com/uudashr/gocognit v1.2.1 // indirect @@ -234,7 +242,6 @@ require ( golang.org/x/exp/typeparams v0.0.0-20260209203927-2842357ff358 // indirect golang.org/x/mod v0.37.0 // indirect golang.org/x/sync v0.21.0 // indirect - golang.org/x/sys v0.46.0 // indirect golang.org/x/text v0.39.0 // indirect google.golang.org/protobuf v1.36.11 // indirect gopkg.in/yaml.v3 v3.0.1 // indirect diff --git a/go.sum b/go.sum index bfc4bb5c..2c101706 100644 --- a/go.sum +++ b/go.sum @@ -60,6 +60,8 @@ github.com/alingse/asasalint v0.0.11 h1:SFwnQXJ49Kx/1GghOFz1XGqHYKp21Kq1nHad/0WQ github.com/alingse/asasalint v0.0.11/go.mod h1:nCaoMhw7a9kSJObvQyVzNTPBDbNpdocqrSP7t/cW5+I= github.com/alingse/nilnesserr v0.2.0 h1:raLem5KG7EFVb4UIDAXgrv3N2JIaffeKNtcEXkEWd/w= github.com/alingse/nilnesserr v0.2.0/go.mod h1:1xJPrXonEtX7wyTq8Dytns5P2hNzoWymVUIaKm4HNFg= +github.com/andybalholm/brotli v1.1.1 h1:PR2pgnyFznKEugtsUo0xLdDop5SKXd5Qf5ysW+7XdTA= +github.com/andybalholm/brotli v1.1.1/go.mod h1:05ib4cKhjx3OQYUY22hTVd34Bc8upXjOLL2rKwwZBoA= github.com/antlr4-go/antlr/v4 v4.13.1 h1:SqQKkuVZ+zWkMMNkjy5FZe5mr5WURWnlpmOuzYWrPrQ= github.com/antlr4-go/antlr/v4 v4.13.1/go.mod h1:GKmUxMtwp6ZgGwZSva4eWPC5mS6vUAmOABFgjdkM7Nw= github.com/ashanbrown/forbidigo/v2 v2.3.0 h1:OZZDOchCgsX5gvToVtEBoV2UWbFfI6RKQTir2UZzSxo= @@ -443,10 +445,18 @@ github.com/otiai10/curr v0.0.0-20150429015615-9b4961190c95/go.mod h1:9qAhocn7zKJ github.com/otiai10/curr v1.0.0/go.mod h1:LskTG5wDwr8Rs+nNQ+1LlxRjAtTZZjtJW4rMXl6j4vs= github.com/otiai10/mint v1.3.0/go.mod h1:F5AjcsTsWUqX+Na9fpHb52P8pcRX2CI6A3ctIT91xUo= github.com/otiai10/mint v1.3.1/go.mod h1:/yxELlJQ0ufhjUwhshSj+wFjZ78CnZ48/1wtmBH1OTc= +github.com/parquet-go/bitpack v1.0.0 h1:AUqzlKzPPXf2bCdjfj4sTeacrUwsT7NlcYDMUQxPcQA= +github.com/parquet-go/bitpack v1.0.0/go.mod h1:XnVk9TH+O40eOOmvpAVZ7K2ocQFrQwysLMnc6M/8lgs= +github.com/parquet-go/jsonlite v1.0.0 h1:87QNdi56wOfsE5bdgas0vRzHPxfJgzrXGml1zZdd7VU= +github.com/parquet-go/jsonlite v1.0.0/go.mod h1:nDjpkpL4EOtqs6NQugUsi0Rleq9sW/OtC1NnZEnxzF0= +github.com/parquet-go/parquet-go v0.30.1 h1:Oy6ganNrAdFiVwy7wNmWagfPTWA2X9Z3tVHBc7JtuX8= +github.com/parquet-go/parquet-go v0.30.1/go.mod h1:navtkAYr2LGoJVp141oXPlO/sxLvaOe3la2JEoD8+rg= github.com/pashagolub/pgxmock/v5 v5.1.0 h1:NZ4pl82b335sEGIbD/+tk2fVIgVs3yNWr1R42ukpUvU= github.com/pashagolub/pgxmock/v5 v5.1.0/go.mod h1:8IJct22b7+EuqecVmYb9aKiENJLLqTsbjFHXH/znAEg= github.com/pelletier/go-toml/v2 v2.4.3 h1:GTRvJQutkOSftxIFD5xw9aepkYNuPWmVJpffdDPYVpY= github.com/pelletier/go-toml/v2 v2.4.3/go.mod h1:2gIqNv+qfxSVS7cM2xJQKtLSTLUE9V8t9Stt+h56mCY= +github.com/pierrec/lz4/v4 v4.1.21 h1:yOVMLb6qSIDP67pl/5F7RepeKYu/VmTyEXvuMI5d9mQ= +github.com/pierrec/lz4/v4 v4.1.21/go.mod h1:gZWDp/Ze/IJXGXf23ltt2EXimqmTUXEy0GFuRQyBid4= github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= @@ -568,6 +578,8 @@ github.com/tomarrell/wrapcheck/v2 v2.12.0 h1:H/qQ1aNWz/eeIhxKAFvkfIA+N7YDvq6TWVF github.com/tomarrell/wrapcheck/v2 v2.12.0/go.mod h1:AQhQuZd0p7b6rfW+vUwHm5OMCGgp63moQ9Qr/0BpIWo= github.com/tommy-muehle/go-mnd/v2 v2.5.1 h1:NowYhSdyE/1zwK9QCLeRb6USWdoif80Ie+v+yU8u1Zw= github.com/tommy-muehle/go-mnd/v2 v2.5.1/go.mod h1:WsUAkMJMYww6l/ufffCD3m+P7LEvr8TnZn9lwVDlgzw= +github.com/twpayne/go-geom v1.6.1 h1:iLE+Opv0Ihm/ABIcvQFGIiFBXd76oBIar9drAwHFhR4= +github.com/twpayne/go-geom v1.6.1/go.mod h1:Kr+Nly6BswFsKM5sd31YaoWS5PeDDH2NftJTK7Gd028= github.com/ultraware/funlen v0.2.0 h1:gCHmCn+d2/1SemTdYMiKLAHFYxTYz7z9VIDRaTGyLkI= github.com/ultraware/funlen v0.2.0/go.mod h1:ZE0q4TsJ8T1SQcjmkhN/w+MceuatI6pBFSxxyteHIJA= github.com/ultraware/whitespace v0.2.0 h1:TYowo2m9Nfj1baEQBjuHzvMRbp19i+RCcRYrSWoFa+g= @@ -580,6 +592,7 @@ github.com/xen0n/gosmopolitan v1.3.0 h1:zAZI1zefvo7gcpbCOrPSHJZJYA9ZgLfJqtKzZ5pH github.com/xen0n/gosmopolitan v1.3.0/go.mod h1:rckfr5T6o4lBtM1ga7mLGKZmLxswUoH1zxHgNXOsEt4= github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e h1:JVG44RsyaB9T2KIHavMF/ppJZNG9ZpyihvCd0w101no= github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e/go.mod h1:RbqR21r5mrJuqunuUZ/Dhy/avygyECGrLceyNeo4LiM= +github.com/xyproto/randomstring v1.0.5/go.mod h1:rgmS5DeNXLivK7YprL0pY+lTuhNQW3iGxZ18UQApw/E= github.com/yagipy/maintidx v1.0.0 h1:h5NvIsCz+nRDapQ0exNv4aJ0yXSI0420omVANTv3GJM= github.com/yagipy/maintidx v1.0.0/go.mod h1:0qNf/I/CCZXSMhsRsrEPDZ+DkekpKLXAJfsTACwgXLk= github.com/yeya24/promlinter v0.3.0 h1:JVDbMp08lVCP7Y6NP3qHroGAO6z2yGKQtS5JsjqtoFs= diff --git a/ret/archive.go b/ret/archive.go new file mode 100644 index 00000000..0440d274 --- /dev/null +++ b/ret/archive.go @@ -0,0 +1,143 @@ +package ret + +import ( + "context" + "fmt" + "path/filepath" + "strings" + "time" + + "github.com/specterops/dawgs/ret/archive" + "github.com/specterops/dawgs/ret/observe" +) + +const ( + packOperationName = "pack" + unpackOperationName = "unpack" +) + +type PackConfig struct { + CollectionDirectory string + ArchivePath string + Recipient archive.PublicKey + Observer observe.Observer +} + +type UnpackConfig struct { + ArchivePath string + OutputDirectory string + Identity archive.PrivateKey + Observer observe.Observer +} + +type KeygenConfig struct { + PrivateKeyPath string + PublicKeyPath string +} + +func (s PackConfig) Validate() error { + if strings.TrimSpace(s.CollectionDirectory) == "" || + strings.TrimSpace(s.ArchivePath) == "" { + return fmt.Errorf("%w: collection directory and archive path are required", ErrInvalidConfig) + } + if s.Recipient == (archive.PublicKey{}) { + return fmt.Errorf("%w: archive recipient is required", ErrInvalidConfig) + } + return nil +} + +func (s UnpackConfig) Validate() error { + if strings.TrimSpace(s.ArchivePath) == "" || + strings.TrimSpace(s.OutputDirectory) == "" { + return fmt.Errorf("%w: archive path and output directory are required", ErrInvalidConfig) + } + if s.Identity == (archive.PrivateKey{}) { + return fmt.Errorf("%w: archive identity is required", ErrInvalidConfig) + } + return nil +} + +func (s KeygenConfig) Validate() error { + if strings.TrimSpace(s.PrivateKeyPath) == "" || + strings.TrimSpace(s.PublicKeyPath) == "" { + return fmt.Errorf("%w: private and public key paths are required", ErrInvalidConfig) + } + privatePath, err := filepath.Abs(s.PrivateKeyPath) + if err != nil { + return fmt.Errorf("%w: resolve private key path: %w", ErrInvalidConfig, err) + } + publicPath, err := filepath.Abs(s.PublicKeyPath) + if err != nil { + return fmt.Errorf("%w: resolve public key path: %w", ErrInvalidConfig, err) + } + if filepath.Clean(privatePath) == filepath.Clean(publicPath) { + return fmt.Errorf("%w: private and public key paths must differ", ErrInvalidConfig) + } + return nil +} + +func Pack(ctx context.Context, config PackConfig) (resultErr error) { + started := time.Now() + observe.Emit(ctx, config.Observer, observe.OperationStarted{Operation: packOperationName}) + defer func() { + observe.Emit(ctx, config.Observer, observe.OperationCompleted{ + Operation: packOperationName, + Duration: time.Since(started), + Err: resultErr, + }) + }() + + if err := ctx.Err(); err != nil { + return fmt.Errorf("pack: %w", err) + } + if err := config.Validate(); err != nil { + return err + } + return archive.Create(ctx, archive.CreateConfig{ + CollectionDirectory: config.CollectionDirectory, + ArchivePath: config.ArchivePath, + Recipient: config.Recipient, + Observer: config.Observer, + }) +} + +func Unpack(ctx context.Context, config UnpackConfig) (resultErr error) { + started := time.Now() + observe.Emit(ctx, config.Observer, observe.OperationStarted{Operation: unpackOperationName}) + defer func() { + observe.Emit(ctx, config.Observer, observe.OperationCompleted{ + Operation: unpackOperationName, + Duration: time.Since(started), + Err: resultErr, + }) + }() + + if err := ctx.Err(); err != nil { + return fmt.Errorf("unpack: %w", err) + } + if err := config.Validate(); err != nil { + return err + } + return archive.Extract(ctx, archive.ExtractConfig{ + ArchivePath: config.ArchivePath, + OutputDirectory: config.OutputDirectory, + Identity: config.Identity, + Observer: config.Observer, + }) +} + +func Keygen(config KeygenConfig) error { + if err := config.Validate(); err != nil { + return err + } + public, private, err := archive.GenerateKeyPair() + if err != nil { + return err + } + return archive.WriteKeyPair( + config.PrivateKeyPath, + private, + config.PublicKeyPath, + public, + ) +} diff --git a/ret/archive/cleanup.go b/ret/archive/cleanup.go new file mode 100644 index 00000000..5f747a09 --- /dev/null +++ b/ret/archive/cleanup.go @@ -0,0 +1,405 @@ +package archive + +import ( + "crypto/rand" + "encoding/hex" + "errors" + "fmt" + "io/fs" + "os" +) + +func removeOwnedEntry( + parent *os.Root, + name string, + expected fs.FileInfo, + retainedStat func() (fs.FileInfo, error), + description string, + operations archiveOperations, +) (resultErr error) { + if err := proveOwnedEntry(parent, name, expected, retainedStat, description); err != nil { + return err + } + if err := operations.runBeforeOwnedEntryQuarantine(); err != nil { + return fmt.Errorf( + "ownership cleanup for %s: before-quarantine operation: %w; preserving pathname", + description, + err, + ) + } + + directory, err := parent.Open(".") + if err != nil { + return fmt.Errorf( + "ownership cleanup for %s: open pinned parent directory: %w; preserving pathname", + description, + err, + ) + } + defer func() { + if err := directory.Close(); err != nil { + resultErr = errors.Join( + resultErr, + fmt.Errorf("ownership cleanup for %s: close pinned parent directory: %w", description, err), + ) + } + }() + + quarantineName, err := quarantineOwnedName(directory, name) + if err != nil { + return fmt.Errorf( + "ownership cleanup for %s: atomically quarantine pathname: %w; preserving pathname", + description, + err, + ) + } + hookErr := operations.runAfterOwnedEntryQuarantine(quarantineName) + quarantined, inspectErr := parent.Lstat(quarantineName) + if inspectErr != nil { + return errors.Join( + fmt.Errorf( + "ownership cleanup for %s: inspect quarantine %q: %w; preserving unproven quarantine", + description, + quarantineName, + inspectErr, + ), + hookErr, + ) + } + if quarantined.Mode().Type() == expected.Mode().Type() && + os.SameFile(expected, quarantined) { + return errors.Join( + fmt.Errorf( + "ownership cleanup for %s: created object preserved in quarantine %q because conditional unlink by identity is unavailable", + description, + quarantineName, + ), + hookErr, + ) + } + + restoreErr := renameNoReplace(directory, "", quarantineName, name) + if restoreErr == nil { + return errors.Join( + fmt.Errorf( + "ownership cleanup for %s: pathname was substituted at the cleanup boundary; replacement restored and preserved", + description, + ), + hookErr, + ) + } + return errors.Join( + fmt.Errorf( + "ownership cleanup for %s: pathname was substituted at the cleanup boundary; replacement preserved in quarantine %q", + description, + quarantineName, + ), + fmt.Errorf("restore quarantined replacement: %w", restoreErr), + hookErr, + ) +} + +func proveOwnedEntry( + parent *os.Root, + name string, + expected fs.FileInfo, + retainedStat func() (fs.FileInfo, error), + description string, +) error { + if expected == nil { + return fmt.Errorf( + "ownership cleanup for %s: created identity is unavailable; preserving pathname", + description, + ) + } + if retainedStat == nil { + return fmt.Errorf( + "ownership cleanup for %s: retained identity handle is unavailable; preserving pathname", + description, + ) + } + retained, err := retainedStat() + if err != nil { + return fmt.Errorf( + "ownership cleanup for %s: inspect retained identity handle: %w; preserving pathname", + description, + err, + ) + } + if retained.Mode().Type() != expected.Mode().Type() || !os.SameFile(expected, retained) { + return fmt.Errorf( + "ownership cleanup for %s: retained identity handle changed; preserving pathname", + description, + ) + } + current, err := parent.Lstat(name) + if err != nil { + if errors.Is(err, os.ErrNotExist) { + return fmt.Errorf( + "ownership cleanup for %s: created pathname disappeared", + description, + ) + } + return fmt.Errorf("ownership cleanup for %s: inspect pathname: %w", description, err) + } + if current.Mode().Type() != expected.Mode().Type() || !os.SameFile(expected, current) { + return fmt.Errorf( + "ownership cleanup for %s: pathname no longer identifies the created object; preserving replacement", + description, + ) + } + return nil +} + +func quarantineOwnedName(directory *os.File, name string) (string, error) { + for range 100 { + var random [16]byte + if _, err := rand.Read(random[:]); err != nil { + return "", fmt.Errorf("generate quarantine name: %w", err) + } + quarantineName := ".ret-cleanup-" + hex.EncodeToString(random[:]) + ".quarantine" + if err := renameNoReplace(directory, "", name, quarantineName); err == nil { + return quarantineName, nil + } else if !errors.Is(err, os.ErrExist) { + return "", err + } + } + return "", fmt.Errorf("quarantine name attempts exhausted") +} + +func removeOwnedDirectory( + parent *os.Root, + name string, + root *os.Root, + expected fs.FileInfo, + description string, + operations archiveOperations, +) error { + if root == nil { + return fmt.Errorf( + "ownership cleanup for %s: pinned directory is unavailable; preserving pathname", + description, + ) + } + pinned, err := root.Stat(".") + if err != nil { + return errors.Join( + fmt.Errorf("ownership cleanup for %s: inspect pinned directory: %w", description, err), + wrapExtractCloseError(description, root.Close()), + ) + } + if expected == nil || !pinned.IsDir() || !os.SameFile(expected, pinned) { + return errors.Join( + fmt.Errorf( + "ownership cleanup for %s: pinned directory identity changed; preserving pathname", + description, + ), + wrapExtractCloseError(description, root.Close()), + ) + } + + var cleanupErr error + entries, err := operations.runReadOwnedDirectory(root) + if err != nil { + cleanupErr = fmt.Errorf( + "ownership cleanup for %s: read pinned directory: %w", + description, + err, + ) + } else { + for _, entry := range entries { + if err := operations.runRemoveOwnedDirectoryEntry(root, entry.Name()); err != nil { + cleanupErr = errors.Join( + cleanupErr, + fmt.Errorf( + "ownership cleanup for %s: remove pinned entry %q: %w", + description, + entry.Name(), + err, + ), + ) + } + } + } + if cleanupErr == nil { + remaining, err := operations.runReadOwnedDirectory(root) + if err != nil { + cleanupErr = fmt.Errorf( + "ownership cleanup for %s: prove pinned directory empty: %w", + description, + err, + ) + } else if len(remaining) != 0 { + cleanupErr = fmt.Errorf( + "ownership cleanup for %s: pinned directory is not empty after rooted cleanup", + description, + ) + } + } + retainedStat := func() (fs.FileInfo, error) { + return root.Stat(".") + } + var removeErr error + if cleanupErr == nil { + removeErr = removeOwnedEntry( + parent, + name, + expected, + retainedStat, + description, + operations, + ) + } else { + removeErr = preserveOwnedEntry( + parent, + name, + expected, + retainedStat, + description, + "rooted directory emptying or proof did not complete", + ) + } + closeErr := wrapExtractCloseError(description, root.Close()) + return errors.Join(cleanupErr, closeErr, removeErr) +} + +type ownedPath struct { + parent *os.Root + handle *os.File + name string + info fs.FileInfo + description string +} + +func preserveOwnedEntry( + parent *os.Root, + name string, + expected fs.FileInfo, + retainedStat func() (fs.FileInfo, error), + description string, + reason string, +) error { + proofErr := proveOwnedEntry(parent, name, expected, retainedStat, description) + if proofErr != nil { + return errors.Join( + fmt.Errorf( + "ownership cleanup for %s: %s; preserving pathname", + description, + reason, + ), + proofErr, + ) + } + return fmt.Errorf( + "ownership cleanup for %s: %s; preserving proven original pathname", + description, + reason, + ) +} + +func sanitizeOwnedFile( + handle *os.File, + description string, + operations archiveOperations, +) error { + if handle == nil { + return fmt.Errorf("sanitize %s before quarantine: retained writable handle is unavailable", description) + } + if err := operations.runTruncateOwnedFile(handle, 0); err != nil { + return fmt.Errorf("truncate %s before quarantine: %w", description, err) + } + if err := operations.runSyncOwnedFile(handle); err != nil { + return fmt.Errorf("sync %s before quarantine: %w", description, err) + } + return nil +} + +func sanitizeAndRemoveOwnedEntry( + parent *os.Root, + name string, + expected fs.FileInfo, + handle *os.File, + description string, + operations archiveOperations, +) error { + if err := sanitizeOwnedFile(handle, description, operations); err != nil { + return errors.Join( + err, + preserveOwnedEntry( + parent, + name, + expected, + func() (fs.FileInfo, error) { + if handle == nil { + return nil, fmt.Errorf("retained writable handle is unavailable") + } + return handle.Stat() + }, + description, + "sanitization did not complete", + ), + ) + } + return removeOwnedEntry( + parent, + name, + expected, + handle.Stat, + description, + operations, + ) +} + +func (s *ownedPath) remove(operations archiveOperations) error { + if s == nil || s.parent == nil { + return fmt.Errorf("ownership cleanup: owned pathname is unavailable") + } + if s.handle == nil { + return errors.Join( + fmt.Errorf("ownership cleanup: owned object handle is unavailable"), + s.release(), + ) + } + removeErr := sanitizeAndRemoveOwnedEntry( + s.parent, + s.name, + s.info, + s.handle, + s.description, + operations, + ) + var handleCloseErr error + if s.handle != nil { + handleCloseErr = s.handle.Close() + s.handle = nil + if handleCloseErr != nil { + handleCloseErr = fmt.Errorf("close owned object: %w", handleCloseErr) + } + } + closeErr := s.parent.Close() + s.parent = nil + if closeErr != nil { + closeErr = fmt.Errorf("close owned pathname parent: %w", closeErr) + } + return errors.Join(removeErr, handleCloseErr, closeErr) +} + +func (s *ownedPath) release() error { + if s == nil || s.parent == nil { + return nil + } + var handleCloseErr error + if s.handle != nil { + handleCloseErr = s.handle.Close() + s.handle = nil + if handleCloseErr != nil { + handleCloseErr = fmt.Errorf("close owned object: %w", handleCloseErr) + } + } + parentCloseErr := s.parent.Close() + s.parent = nil + if parentCloseErr != nil { + parentCloseErr = fmt.Errorf("close owned pathname parent: %w", parentCloseErr) + } + return errors.Join(handleCloseErr, parentCloseErr) +} diff --git a/ret/archive/cleanup_test.go b/ret/archive/cleanup_test.go new file mode 100644 index 00000000..03122336 --- /dev/null +++ b/ret/archive/cleanup_test.go @@ -0,0 +1,264 @@ +package archive + +import ( + "errors" + "io/fs" + "os" + "path/filepath" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestRemoveOwnedEntryPreservesSubstitutionAtCheckRemovalBoundary(t *testing.T) { + // Break caught: Lstat/SameFile succeeds, a replacement is installed, and a + // later pathname Remove deletes the replacement rather than the owned file. + parentPath := t.TempDir() + parent, err := os.OpenRoot(parentPath) + require.NoError(t, err) + t.Cleanup(func() { + require.NoError(t, parent.Close()) + }) + const name = "owned.tmp" + owned, err := parent.OpenFile(name, os.O_RDWR|os.O_CREATE|os.O_EXCL, 0o600) + require.NoError(t, err) + t.Cleanup(func() { + require.NoError(t, owned.Close()) + }) + expected, err := owned.Stat() + require.NoError(t, err) + replaced := false + operations := archiveOperations{ + beforeOwnedEntryQuarantine: func() error { + require.NoError(t, parent.Remove(name)) + replacement, err := parent.OpenFile( + name, + os.O_WRONLY|os.O_CREATE|os.O_EXCL, + 0o600, + ) + require.NoError(t, err) + _, err = replacement.Write([]byte("preserve replacement")) + require.NoError(t, err) + require.NoError(t, replacement.Close()) + replaced = true + return nil + }, + } + + err = removeOwnedEntry( + parent, + name, + expected, + owned.Stat, + "test object", + operations, + ) + + require.ErrorContains(t, err, "ownership cleanup") + require.True(t, replaced) + payload, readErr := os.ReadFile(filepath.Join(parentPath, name)) + require.NoError(t, readErr) + require.Equal(t, []byte("preserve replacement"), payload) + quarantines, globErr := filepath.Glob(filepath.Join(parentPath, ".ret-cleanup-*.quarantine")) + require.NoError(t, globErr) + require.Empty(t, quarantines) + require.False(t, errors.Is(err, os.ErrNotExist)) +} + +func TestRemoveOwnedEntryPreservesQuarantinedSubstitutionWhenRestoreIsBlocked(t *testing.T) { + // Break caught: deleting the quarantined substitution, or replacing a + // second object that occupies the original name before restoration. + parentPath := t.TempDir() + parent, err := os.OpenRoot(parentPath) + require.NoError(t, err) + t.Cleanup(func() { + require.NoError(t, parent.Close()) + }) + const name = "owned.tmp" + owned, err := parent.OpenFile(name, os.O_RDWR|os.O_CREATE|os.O_EXCL, 0o600) + require.NoError(t, err) + t.Cleanup(func() { + require.NoError(t, owned.Close()) + }) + expected, err := owned.Stat() + require.NoError(t, err) + operations := archiveOperations{ + beforeOwnedEntryQuarantine: func() error { + require.NoError(t, parent.Remove(name)) + replacement, err := parent.OpenFile( + name, + os.O_WRONLY|os.O_CREATE|os.O_EXCL, + 0o600, + ) + require.NoError(t, err) + _, err = replacement.Write([]byte("first replacement")) + require.NoError(t, err) + return replacement.Close() + }, + afterOwnedEntryQuarantine: func(_ string) error { + blocker, err := parent.OpenFile( + name, + os.O_WRONLY|os.O_CREATE|os.O_EXCL, + 0o600, + ) + require.NoError(t, err) + _, err = blocker.Write([]byte("second replacement")) + require.NoError(t, err) + return blocker.Close() + }, + } + + err = removeOwnedEntry( + parent, + name, + expected, + owned.Stat, + "test object", + operations, + ) + + require.ErrorContains(t, err, "ownership cleanup") + require.ErrorContains(t, err, "quarantine") + payload, readErr := os.ReadFile(filepath.Join(parentPath, name)) + require.NoError(t, readErr) + require.Equal(t, []byte("second replacement"), payload) + quarantines, globErr := filepath.Glob(filepath.Join(parentPath, ".ret-cleanup-*.quarantine")) + require.NoError(t, globErr) + require.Len(t, quarantines, 1) + quarantined, readErr := os.ReadFile(quarantines[0]) + require.NoError(t, readErr) + require.Equal(t, []byte("first replacement"), quarantined) +} + +func TestRemoveOwnedDirectoryPreservesOriginalWhenRootedCleanupIsUnproven(t *testing.T) { + // Break caught: moving a stage into quarantine after rooted enumeration, + // entry removal, or final empty-inventory proof failed. + tests := []struct { + name string + operation func(t *testing.T) archiveOperations + match string + wantEntry string + }{ + { + name: "initial read directory", + operation: func(_ *testing.T) archiveOperations { + return archiveOperations{ + readOwnedDirectory: func(_ *os.Root) ([]fs.DirEntry, error) { + return nil, errors.New("injected read directory failure") + }, + } + }, + match: "injected read directory failure", + wantEntry: "marker", + }, + { + name: "remove entry", + operation: func(_ *testing.T) archiveOperations { + return archiveOperations{ + removeOwnedDirectoryEntry: func(_ *os.Root, _ string) error { + return errors.New("injected remove entry failure") + }, + } + }, + match: "injected remove entry failure", + wantEntry: "marker", + }, + { + name: "final read directory", + operation: func(_ *testing.T) archiveOperations { + readCalls := 0 + return archiveOperations{ + readOwnedDirectory: func(root *os.Root) ([]fs.DirEntry, error) { + readCalls++ + if readCalls == 2 { + return nil, errors.New("injected final read directory failure") + } + return fs.ReadDir(root.FS(), ".") + }, + } + }, + match: "injected final read directory failure", + }, + { + name: "final inventory not empty", + operation: func(t *testing.T) archiveOperations { + readCalls := 0 + return archiveOperations{ + readOwnedDirectory: func(root *os.Root) ([]fs.DirEntry, error) { + readCalls++ + if readCalls == 2 { + late, err := root.OpenFile( + "late-entry", + os.O_WRONLY|os.O_CREATE|os.O_EXCL, + 0o600, + ) + require.NoError(t, err) + _, err = late.Write([]byte("late payload")) + require.NoError(t, err) + require.NoError(t, late.Close()) + } + return fs.ReadDir(root.FS(), ".") + }, + } + }, + match: "not empty", + wantEntry: "late-entry", + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + parentPath := t.TempDir() + parent, err := os.OpenRoot(parentPath) + require.NoError(t, err) + t.Cleanup(func() { + require.NoError(t, parent.Close()) + }) + const stageName = "stage.tmp" + require.NoError(t, parent.Mkdir(stageName, 0o700)) + stage, err := parent.OpenRoot(stageName) + require.NoError(t, err) + info, err := stage.Stat(".") + require.NoError(t, err) + marker, err := stage.OpenFile( + "marker", + os.O_WRONLY|os.O_CREATE|os.O_EXCL, + 0o600, + ) + require.NoError(t, err) + _, err = marker.Write([]byte("stage payload")) + require.NoError(t, err) + require.NoError(t, marker.Close()) + + err = removeOwnedDirectory( + parent, + stageName, + stage, + info, + "test stage", + test.operation(t), + ) + + require.ErrorContains(t, err, test.match) + require.ErrorContains(t, err, "ownership cleanup") + stagePath := filepath.Join(parentPath, stageName) + require.DirExists(t, stagePath) + require.Empty(t, archiveCleanupQuarantinePaths(t, parentPath)) + entries, readErr := os.ReadDir(stagePath) + require.NoError(t, readErr) + if test.wantEntry == "" { + require.Empty(t, entries) + } else { + require.Len(t, entries, 1) + require.Equal(t, test.wantEntry, entries[0].Name()) + } + }) + } +} + +func archiveCleanupQuarantinePaths(t *testing.T, parent string) []string { + t.Helper() + quarantines, err := filepath.Glob(filepath.Join(parent, ".ret-cleanup-*.quarantine")) + require.NoError(t, err) + return quarantines +} diff --git a/ret/archive/envelope.go b/ret/archive/envelope.go new file mode 100644 index 00000000..a83b7681 --- /dev/null +++ b/ret/archive/envelope.go @@ -0,0 +1,479 @@ +package archive + +import ( + "bytes" + "crypto/hpke" + "crypto/sha256" + "encoding/base64" + "encoding/binary" + "encoding/json" + "errors" + "fmt" + "io" + "math" +) + +const ( + envelopeMagic = "RET-PQ-ARCHIVE-v1" + envelopeHPKEInfo = "ret/archive/hpke/v1" + envelopeFrameSize = 1024 * 1024 + maxEnvelopeHeaderSize = 8 * 1024 + + frameHeaderSize = 13 + aeadOverhead = 16 + + dataFrame byte = 0 + finalFrame byte = 1 + + maxEnvelopeFrameSequence = math.MaxUint64 - 1 +) + +type envelopeHeader struct { + Format string `json:"format"` + KEM string `json:"kem"` + KDF string `json:"kdf"` + AEAD string `json:"aead"` + EncapsulatedKey string `json:"encapsulated_key"` + FrameSize int `json:"frame_size"` +} + +type encryptWriter struct { + destination io.Writer + sender *hpke.Sender + headerDigest [sha256.Size]byte + sequence uint64 + writeErr error + closeErr error + closed bool +} + +type decryptReader struct { + source io.Reader + recipient *hpke.Recipient + headerDigest [sha256.Size]byte + sequence uint64 + plaintext []byte + terminalErr error + closeErr error + final bool + closed bool +} + +func newEncryptWriter(destination io.Writer, recipient PublicKey) (io.WriteCloser, error) { + if destination == nil { + return nil, fmt.Errorf("encryption destination is required") + } + if !recipient.valid { + return nil, fmt.Errorf("recipient public key is required") + } + + publicHPKEKey, err := archiveHPKESuite.kem.NewPublicKey(recipient.material[:]) + if err != nil { + return nil, fmt.Errorf("parse recipient public key: %w", err) + } + encapsulatedKey, sender, err := hpke.NewSender( + publicHPKEKey, + archiveHPKESuite.kdf, + archiveHPKESuite.aead, + []byte(envelopeHPKEInfo), + ) + if err != nil { + return nil, fmt.Errorf("create archive sender: %w", err) + } + if len(encapsulatedKey) != publicKeyMaterialSize { + return nil, fmt.Errorf( + "generated encapsulated key has length %d, want %d", + len(encapsulatedKey), + publicKeyMaterialSize, + ) + } + + header := envelopeHeader{ + Format: EnvelopeFormat, + KEM: kemName, + KDF: kdfName, + AEAD: aeadName, + EncapsulatedKey: base64.StdEncoding.EncodeToString(encapsulatedKey), + FrameSize: envelopeFrameSize, + } + headerBytes, err := json.Marshal(header) + if err != nil { + return nil, fmt.Errorf("encode archive header: %w", err) + } + if len(headerBytes) == 0 || len(headerBytes) > maxEnvelopeHeaderSize { + return nil, fmt.Errorf("archive header size %d is invalid", len(headerBytes)) + } + + if err := writeAll(destination, []byte(envelopeMagic)); err != nil { + return nil, fmt.Errorf("write archive magic: %w", err) + } + var headerSize [4]byte + binary.BigEndian.PutUint32(headerSize[:], uint32(len(headerBytes))) + if err := writeAll(destination, headerSize[:]); err != nil { + return nil, fmt.Errorf("write archive header size: %w", err) + } + if err := writeAll(destination, headerBytes); err != nil { + return nil, fmt.Errorf("write archive header: %w", err) + } + + return &encryptWriter{ + destination: destination, + sender: sender, + headerDigest: sha256.Sum256(headerBytes), + }, nil +} + +func (s *encryptWriter) Write(p []byte) (int, error) { + if s.closed { + return 0, fmt.Errorf("encrypted archive writer is closed") + } + if s.writeErr != nil { + return 0, s.writeErr + } + + written := 0 + for len(p) > 0 { + chunkSize := min(len(p), envelopeFrameSize) + if err := s.writeFrame(dataFrame, p[:chunkSize]); err != nil { + s.writeErr = err + return written, err + } + written += chunkSize + p = p[chunkSize:] + } + return written, nil +} + +func (s *encryptWriter) Close() error { + if s.closed { + if s.closeErr != nil { + return s.closeErr + } + return fmt.Errorf("encrypted archive writer is already closed") + } + s.closed = true + + if s.writeErr != nil { + s.closeErr = s.writeErr + return s.closeErr + } + if err := s.writeFrame(finalFrame, nil); err != nil { + s.closeErr = err + return err + } + return nil +} + +func (s *encryptWriter) writeFrame(frameType byte, plaintext []byte) error { + final := frameType == finalFrame + if err := validateFrameSequence(s.sequence, final); err != nil { + return err + } + ciphertext, err := s.sender.Seal(frameAAD(s.headerDigest[:], s.sequence, final), plaintext) + if err != nil { + return fmt.Errorf("encrypt archive frame %d: %w", s.sequence, err) + } + if len(ciphertext) > envelopeFrameSize+aeadOverhead { + return fmt.Errorf("encrypted archive frame %d is too large", s.sequence) + } + + var header [frameHeaderSize]byte + header[0] = frameType + binary.BigEndian.PutUint64(header[1:9], s.sequence) + binary.BigEndian.PutUint32(header[9:13], uint32(len(ciphertext))) + if err := writeAll(s.destination, header[:]); err != nil { + return fmt.Errorf("write archive frame %d header: %w", s.sequence, err) + } + if err := writeAll(s.destination, ciphertext); err != nil { + return fmt.Errorf("write archive frame %d ciphertext: %w", s.sequence, err) + } + + s.sequence++ + return nil +} + +func newDecryptReader(source io.Reader, identity PrivateKey) (io.ReadCloser, error) { + if source == nil { + return nil, fmt.Errorf("encrypted archive source is required") + } + if !identity.valid { + return nil, fmt.Errorf("identity private key is required") + } + privateHPKEKey, err := archiveHPKESuite.kem.NewPrivateKey(identity.material[:]) + if err != nil { + return nil, fmt.Errorf("parse identity private key: %w", err) + } + + headerBytes, _, encapsulatedKey, err := readEnvelopeHeader(source) + if err != nil { + return nil, err + } + recipient, err := hpke.NewRecipient( + encapsulatedKey, + privateHPKEKey, + archiveHPKESuite.kdf, + archiveHPKESuite.aead, + []byte(envelopeHPKEInfo), + ) + if err != nil { + return nil, fmt.Errorf("create archive recipient: %w", err) + } + + return &decryptReader{ + source: source, + recipient: recipient, + headerDigest: sha256.Sum256(headerBytes), + }, nil +} + +func (s *decryptReader) Read(p []byte) (int, error) { + if s.closed { + return 0, io.ErrClosedPipe + } + if len(p) == 0 { + return 0, nil + } + if s.terminalErr != nil { + return 0, s.terminalErr + } + + for len(s.plaintext) == 0 { + if s.final { + return 0, io.EOF + } + if err := s.readNextFrame(); err != nil { + s.terminalErr = err + return 0, err + } + } + + n := copy(p, s.plaintext) + s.plaintext = s.plaintext[n:] + return n, nil +} + +func (s *decryptReader) Close() error { + if s.closed { + return s.closeErr + } + + if s.terminalErr == nil { + s.plaintext = nil + for !s.final { + if err := s.readNextFrame(); err != nil { + s.terminalErr = err + break + } + s.plaintext = nil + } + } + s.closed = true + s.closeErr = s.terminalErr + return s.closeErr +} + +func (s *decryptReader) readNextFrame() error { + var header [frameHeaderSize]byte + if _, err := io.ReadFull(s.source, header[:]); err != nil { + if errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) { + return fmt.Errorf("encrypted archive is missing its final frame") + } + return fmt.Errorf("read archive frame %d header: %w", s.sequence, err) + } + + frameType := header[0] + if frameType != dataFrame && frameType != finalFrame { + return fmt.Errorf("archive frame %d has unsupported type %d", s.sequence, frameType) + } + sequence := binary.BigEndian.Uint64(header[1:9]) + if sequence != s.sequence { + return fmt.Errorf("archive frame sequence is %d, want %d", sequence, s.sequence) + } + final := frameType == finalFrame + if err := validateFrameSequence(s.sequence, final); err != nil { + return err + } + ciphertextSize := binary.BigEndian.Uint32(header[9:13]) + if ciphertextSize < aeadOverhead { + return fmt.Errorf("archive frame %d ciphertext is too short", s.sequence) + } + if ciphertextSize > envelopeFrameSize+aeadOverhead { + return fmt.Errorf("archive frame %d is too large", s.sequence) + } + + ciphertext := make([]byte, int(ciphertextSize)) + if _, err := io.ReadFull(s.source, ciphertext); err != nil { + return fmt.Errorf("read archive frame %d ciphertext: %w", s.sequence, err) + } + plaintext, err := s.recipient.Open(frameAAD(s.headerDigest[:], s.sequence, final), ciphertext) + if err != nil { + return fmt.Errorf("authenticate archive frame %d: %w", s.sequence, err) + } + if len(plaintext) > envelopeFrameSize { + return fmt.Errorf("archive frame %d plaintext is too large", s.sequence) + } + s.sequence++ + + if final { + if len(plaintext) != 0 { + return fmt.Errorf("authenticated final frame contained plaintext") + } + if err := requireEnvelopeEOF(s.source); err != nil { + return err + } + s.final = true + return nil + } + if len(plaintext) == 0 { + return fmt.Errorf("authenticated data frame %d contained no plaintext", sequence) + } + + s.plaintext = plaintext + return nil +} + +func validateFrameSequence(sequence uint64, final bool) error { + if sequence > maxEnvelopeFrameSequence { + return fmt.Errorf("archive frame sequence exhausted at %d", sequence) + } + if !final && sequence == maxEnvelopeFrameSequence { + return fmt.Errorf( + "archive frame sequence exhausted at %d: sequence is reserved for the final frame", + sequence, + ) + } + return nil +} + +func readEnvelopeHeader(source io.Reader) ([]byte, envelopeHeader, []byte, error) { + var header envelopeHeader + magic := make([]byte, len(envelopeMagic)) + if _, err := io.ReadFull(source, magic); err != nil { + return nil, header, nil, fmt.Errorf("read archive magic: %w", err) + } + if string(magic) != envelopeMagic { + return nil, header, nil, fmt.Errorf("archive magic must be %q, got %q", envelopeMagic, string(magic)) + } + + var sizeBytes [4]byte + if _, err := io.ReadFull(source, sizeBytes[:]); err != nil { + return nil, header, nil, fmt.Errorf("read archive header size: %w", err) + } + headerSize := binary.BigEndian.Uint32(sizeBytes[:]) + if headerSize == 0 || headerSize > maxEnvelopeHeaderSize { + return nil, header, nil, fmt.Errorf("archive header size %d is invalid", headerSize) + } + + headerBytes := make([]byte, int(headerSize)) + if _, err := io.ReadFull(source, headerBytes); err != nil { + return nil, header, nil, fmt.Errorf("read archive header: %w", err) + } + if err := validateExactJSONObject( + headerBytes, + "format", + "kem", + "kdf", + "aead", + "encapsulated_key", + "frame_size", + ); err != nil { + return nil, header, nil, fmt.Errorf("validate archive header JSON fields: %w", err) + } + decoder := json.NewDecoder(bytes.NewReader(headerBytes)) + decoder.DisallowUnknownFields() + if err := decoder.Decode(&header); err != nil { + return nil, header, nil, fmt.Errorf("decode archive header JSON: %w", err) + } + var trailing any + if err := decoder.Decode(&trailing); !errors.Is(err, io.EOF) { + if err == nil { + return nil, header, nil, fmt.Errorf("archive header JSON contains a trailing value") + } + return nil, header, nil, fmt.Errorf("decode trailing archive header JSON: %w", err) + } + + encapsulatedKey, err := validateEnvelopeHeader(header) + if err != nil { + return nil, header, nil, err + } + return headerBytes, header, encapsulatedKey, nil +} + +func validateEnvelopeHeader(header envelopeHeader) ([]byte, error) { + if header.Format != EnvelopeFormat { + return nil, fmt.Errorf("envelope format must be %q, got %q", EnvelopeFormat, header.Format) + } + if header.KEM != kemName { + return nil, fmt.Errorf("envelope KEM must be %q, got %q", kemName, header.KEM) + } + if header.KDF != kdfName { + return nil, fmt.Errorf("envelope KDF must be %q, got %q", kdfName, header.KDF) + } + if header.AEAD != aeadName { + return nil, fmt.Errorf("envelope AEAD must be %q, got %q", aeadName, header.AEAD) + } + if header.FrameSize != envelopeFrameSize { + return nil, fmt.Errorf("envelope frame size must be %d, got %d", envelopeFrameSize, header.FrameSize) + } + if header.EncapsulatedKey == "" { + return nil, fmt.Errorf("envelope encapsulated key is required") + } + + encapsulatedKey, err := base64.StdEncoding.Strict().DecodeString(header.EncapsulatedKey) + if err != nil || base64.StdEncoding.EncodeToString(encapsulatedKey) != header.EncapsulatedKey { + if err == nil { + err = fmt.Errorf("encoding is not canonical") + } + return nil, fmt.Errorf("decode encapsulated key as canonical base64: %w", err) + } + if len(encapsulatedKey) != publicKeyMaterialSize { + return nil, fmt.Errorf( + "encapsulated key has length %d, want %d", + len(encapsulatedKey), + publicKeyMaterialSize, + ) + } + return encapsulatedKey, nil +} + +func requireEnvelopeEOF(source io.Reader) error { + var extra [1]byte + n, err := source.Read(extra[:]) + if n > 0 { + return fmt.Errorf("encrypted archive has data after final frame") + } + if err == nil { + return fmt.Errorf("encrypted archive source did not end after final frame") + } + if !errors.Is(err, io.EOF) { + return fmt.Errorf("read encrypted archive trailer: %w", err) + } + return nil +} + +func frameAAD(headerDigest []byte, sequence uint64, final bool) []byte { + aad := make([]byte, 0, len(headerDigest)+9) + aad = append(aad, headerDigest...) + aad = binary.BigEndian.AppendUint64(aad, sequence) + if final { + return append(aad, 1) + } + return append(aad, 0) +} + +func writeAll(destination io.Writer, payload []byte) error { + for len(payload) > 0 { + n, err := destination.Write(payload) + if n < 0 || n > len(payload) { + return fmt.Errorf("writer returned invalid byte count %d", n) + } + payload = payload[n:] + if err != nil { + return err + } + if n == 0 { + return io.ErrShortWrite + } + } + return nil +} diff --git a/ret/archive/envelope_test.go b/ret/archive/envelope_test.go new file mode 100644 index 00000000..2559c52e --- /dev/null +++ b/ret/archive/envelope_test.go @@ -0,0 +1,822 @@ +package archive + +import ( + "bytes" + "crypto/hpke" + "crypto/sha256" + "encoding/base64" + "encoding/binary" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "io" + "math" + "strings" + "testing" + + "github.com/stretchr/testify/require" +) + +const ( + testEnvelopeMagic = "RET-PQ-ARCHIVE-v1" + testEnvelopeHPKEInfo = "ret/archive/hpke/v1" + testEnvelopeFrameSize = 1024 * 1024 + testEnvelopeFrameHeaderSize = 13 + testEnvelopeAEADOverhead = 16 + + testEnvelopeDataFrame byte = 0 + testEnvelopeFinalFrame byte = 1 +) + +type testEnvelopeHeader struct { + Format string `json:"format"` + KEM string `json:"kem"` + KDF string `json:"kdf"` + AEAD string `json:"aead"` + EncapsulatedKey string `json:"encapsulated_key"` + FrameSize int `json:"frame_size"` +} + +type testEnvelopeFrame struct { + start int + end int + ciphertextStart int + frameType byte + sequence uint64 + ciphertextSize uint32 +} + +func TestEnvelopeRoundTripAcrossFrames(t *testing.T) { + recipient, identity, err := GenerateKeyPair() + require.NoError(t, err) + + for _, size := range []int{ + 0, + 1, + testEnvelopeFrameSize - 1, + testEnvelopeFrameSize, + testEnvelopeFrameSize + 1, + 2*testEnvelopeFrameSize + 31, + } { + t.Run(fmt.Sprintf("%d_bytes", size), func(t *testing.T) { + plaintext := testPlaintext(size) + ciphertext := encryptTestEnvelope(t, recipient, plaintext) + + decrypted, err := decryptTestEnvelope(identity, ciphertext) + require.NoError(t, err) + require.Equal(t, plaintext, decrypted) + }) + } +} + +func TestEnvelopeHeaderUsesExactFormatSuiteAndBound(t *testing.T) { + recipient, _, err := GenerateKeyPair() + require.NoError(t, err) + secret := []byte("header-must-not-contain-this-plaintext") + ciphertext := encryptTestEnvelope(t, recipient, secret) + + headerBytes, header, frames := parseTestEnvelope(t, ciphertext) + require.Equal(t, EnvelopeFormat, header.Format) + require.Equal(t, testKEMName, header.KEM) + require.Equal(t, testKDFName, header.KDF) + require.Equal(t, testAEADName, header.AEAD) + require.Equal(t, testEnvelopeFrameSize, header.FrameSize) + require.NotContains(t, string(headerBytes), string(secret)) + require.Len(t, decodeTestMaterial(t, header.EncapsulatedKey), testPublicMaterialSize) + require.Len(t, frames, 2) + require.Equal(t, testEnvelopeDataFrame, frames[0].frameType) + require.Equal(t, uint64(0), frames[0].sequence) + require.LessOrEqual(t, frames[0].ciphertextSize, uint32(testEnvelopeFrameSize+testEnvelopeAEADOverhead)) + require.Equal(t, testEnvelopeFinalFrame, frames[1].frameType) + require.Equal(t, uint64(1), frames[1].sequence) + require.Equal(t, uint32(testEnvelopeAEADOverhead), frames[1].ciphertextSize) +} + +func TestEnvelopeRejectsLegacyMagicAndFormat(t *testing.T) { + _, identity, err := GenerateKeyPair() + require.NoError(t, err) + + _, err = newDecryptReader(bytes.NewReader([]byte("RTRV-PQ-ARCHIVE-v1")), identity) + require.ErrorContains(t, err, "magic") + + recipient, _, err := GenerateKeyPair() + require.NoError(t, err) + ciphertext := encryptTestEnvelope(t, recipient, []byte("payload")) + legacyFormat := rewriteTestEnvelopeHeader(t, ciphertext, func(header *testEnvelopeHeader) { + header.Format = "retriever-encrypted-tar-v1" + }) + _, err = decryptTestEnvelope(identity, legacyFormat) + require.ErrorContains(t, err, EnvelopeFormat) +} + +func TestEnvelopeRejectsHeaderValidationFailures(t *testing.T) { + recipient, identity, err := GenerateKeyPair() + require.NoError(t, err) + ciphertext := encryptTestEnvelope(t, recipient, []byte("payload")) + + tests := []struct { + name string + mutate func(*testEnvelopeHeader) + match string + }{ + {name: "wrong KEM", mutate: func(header *testEnvelopeHeader) { + header.KEM = "ML-KEM-768" + }, match: testKEMName}, + {name: "wrong KDF", mutate: func(header *testEnvelopeHeader) { + header.KDF = "HKDF-SHA256" + }, match: testKDFName}, + {name: "wrong AEAD", mutate: func(header *testEnvelopeHeader) { + header.AEAD = "AES-128-GCM" + }, match: testAEADName}, + {name: "wrong frame size", mutate: func(header *testEnvelopeHeader) { + header.FrameSize++ + }, match: "frame size"}, + {name: "empty encapsulation", mutate: func(header *testEnvelopeHeader) { + header.EncapsulatedKey = "" + }, match: "encapsulated key"}, + {name: "invalid encapsulation base64", mutate: func(header *testEnvelopeHeader) { + header.EncapsulatedKey = "not base64!" + }, match: "base64"}, + {name: "noncanonical encapsulation base64", mutate: func(header *testEnvelopeHeader) { + header.EncapsulatedKey = header.EncapsulatedKey[:8] + "\n" + header.EncapsulatedKey[8:] + }, match: "base64"}, + {name: "short encapsulation", mutate: func(header *testEnvelopeHeader) { + header.EncapsulatedKey = base64.StdEncoding.EncodeToString(make([]byte, testPublicMaterialSize-1)) + }, match: "1568"}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + tampered := rewriteTestEnvelopeHeader(t, ciphertext, test.mutate) + _, err := decryptTestEnvelope(identity, tampered) + require.ErrorContains(t, err, test.match) + }) + } +} + +func TestEnvelopeRejectsUnknownAndTrailingHeaderJSON(t *testing.T) { + recipient, identity, err := GenerateKeyPair() + require.NoError(t, err) + ciphertext := encryptTestEnvelope(t, recipient, []byte("payload")) + headerBytes, _, _ := parseTestEnvelope(t, ciphertext) + + unknown := append(append([]byte(nil), bytes.TrimSuffix(headerBytes, []byte("}"))...), []byte(`,"extra":true}`)...) + _, err = decryptTestEnvelope(identity, replaceTestEnvelopeHeader(t, ciphertext, unknown)) + require.ErrorContains(t, err, "unknown field") + + trailing := append(append([]byte(nil), headerBytes...), []byte(`{}`)...) + _, err = decryptTestEnvelope(identity, replaceTestEnvelopeHeader(t, ciphertext, trailing)) + require.ErrorContains(t, err, "trailing") +} + +func TestEnvelopeHeaderRequiresEachExactLowercaseFieldOnce(t *testing.T) { + recipient, identity, err := GenerateKeyPair() + require.NoError(t, err) + ciphertext := encryptTestEnvelope(t, recipient, []byte("payload")) + validHeader, _, _ := parseTestEnvelope(t, ciphertext) + fields := []string{ + "format", + "kem", + "kdf", + "aead", + "encapsulated_key", + "frame_size", + } + + for _, field := range fields { + t.Run(field+"/case_variant", func(t *testing.T) { + header := envelopeJSONWithCaseVariantField(t, validHeader, field) + _, err := newDecryptReader( + bytes.NewReader(replaceTestEnvelopeHeader(t, ciphertext, header)), + identity, + ) + require.ErrorContains(t, err, "exact lowercase") + }) + t.Run(field+"/attacker_duplicate_first", func(t *testing.T) { + header := envelopeJSONWithDuplicateField(t, validHeader, field, true) + _, err := newDecryptReader( + bytes.NewReader(replaceTestEnvelopeHeader(t, ciphertext, header)), + identity, + ) + require.ErrorContains(t, err, "duplicate") + }) + t.Run(field+"/attacker_duplicate_last", func(t *testing.T) { + header := envelopeJSONWithDuplicateField(t, validHeader, field, false) + _, err := newDecryptReader( + bytes.NewReader(replaceTestEnvelopeHeader(t, ciphertext, header)), + identity, + ) + require.ErrorContains(t, err, "duplicate") + }) + } +} + +func TestEnvelopeAuthenticatesExactHeaderBytesAndEncapsulation(t *testing.T) { + recipient, identity, err := GenerateKeyPair() + require.NoError(t, err) + ciphertext := encryptTestEnvelope(t, recipient, []byte("payload")) + headerBytes, _, _ := parseTestEnvelope(t, ciphertext) + + whitespaceTamper := append(append([]byte(nil), headerBytes...), ' ') + plaintext, err := decryptTestEnvelope(identity, replaceTestEnvelopeHeader(t, ciphertext, whitespaceTamper)) + require.Error(t, err) + require.Empty(t, plaintext) + + encapsulationTamper := rewriteTestEnvelopeHeader(t, ciphertext, func(header *testEnvelopeHeader) { + replacement := byte('A') + if header.EncapsulatedKey[0] == replacement { + replacement = 'B' + } + header.EncapsulatedKey = string(replacement) + header.EncapsulatedKey[1:] + }) + plaintext, err = decryptTestEnvelope(identity, encapsulationTamper) + require.Error(t, err) + require.Empty(t, plaintext) +} + +func TestEnvelopeRejectsWrongIdentity(t *testing.T) { + recipient, _, err := GenerateKeyPair() + require.NoError(t, err) + _, wrongIdentity, err := GenerateKeyPair() + require.NoError(t, err) + + plaintext, err := decryptTestEnvelope(wrongIdentity, encryptTestEnvelope(t, recipient, []byte("payload"))) + require.Error(t, err) + require.Empty(t, plaintext) +} + +func TestEnvelopeRejectsCiphertextTamperWithoutReportingCurrentFrame(t *testing.T) { + recipient, identity, err := GenerateKeyPair() + require.NoError(t, err) + plaintext := testPlaintext(2*testEnvelopeFrameSize + 17) + ciphertext := encryptTestEnvelope(t, recipient, plaintext) + _, _, frames := parseTestEnvelope(t, ciphertext) + require.Len(t, frames, 4) + + firstTampered := append([]byte(nil), ciphertext...) + firstTampered[frames[0].ciphertextStart] ^= 0x01 + decrypted, err := decryptTestEnvelope(identity, firstTampered) + require.Error(t, err) + require.Empty(t, decrypted) + + secondTampered := append([]byte(nil), ciphertext...) + secondTampered[frames[1].ciphertextStart+11] ^= 0x01 + decrypted, err = decryptTestEnvelope(identity, secondTampered) + require.Error(t, err) + require.Equal(t, plaintext[:testEnvelopeFrameSize], decrypted) +} + +func TestEnvelopeRejectsReorderedAndRepeatedFrameNumbers(t *testing.T) { + recipient, identity, err := GenerateKeyPair() + require.NoError(t, err) + ciphertext := encryptTestEnvelope(t, recipient, testPlaintext(2*testEnvelopeFrameSize+1)) + _, _, frames := parseTestEnvelope(t, ciphertext) + require.Len(t, frames, 4) + + reordered := append([]byte(nil), ciphertext[:frames[0].start]...) + reordered = append(reordered, ciphertext[frames[1].start:frames[1].end]...) + reordered = append(reordered, ciphertext[frames[0].start:frames[0].end]...) + reordered = append(reordered, ciphertext[frames[2].start:]...) + decrypted, err := decryptTestEnvelope(identity, reordered) + require.ErrorContains(t, err, "sequence") + require.Empty(t, decrypted) + + repeated := append([]byte(nil), ciphertext...) + binary.BigEndian.PutUint64(repeated[frames[1].start+1:], frames[0].sequence) + decrypted, err = decryptTestEnvelope(identity, repeated) + require.ErrorContains(t, err, "sequence") + require.Equal(t, testPlaintext(testEnvelopeFrameSize), decrypted) +} + +func TestEncryptWriterReservesLastSequenceForFinalFrame(t *testing.T) { + recipient, identity, err := GenerateKeyPair() + require.NoError(t, err) + + var destination bytes.Buffer + writeCloser, err := newEncryptWriter(&destination, recipient) + require.NoError(t, err) + writer := writeCloser.(*encryptWriter) + writer.sequence = math.MaxUint64 - 2 + + plaintext := []byte("last-data-frame") + n, err := writer.Write(plaintext) + require.NoError(t, err) + require.Equal(t, len(plaintext), n) + require.Equal(t, uint64(math.MaxUint64-1), writer.sequence) + require.NoError(t, writer.Close()) + require.Equal(t, uint64(math.MaxUint64), writer.sequence) + + _, _, frames := parseTestEnvelope(t, destination.Bytes()) + require.Len(t, frames, 2) + require.Equal(t, uint64(math.MaxUint64-2), frames[0].sequence) + require.Equal(t, testEnvelopeDataFrame, frames[0].frameType) + require.Equal(t, uint64(math.MaxUint64-1), frames[1].sequence) + require.Equal(t, testEnvelopeFinalFrame, frames[1].frameType) + + readCloser, err := newDecryptReader(bytes.NewReader(destination.Bytes()), identity) + require.NoError(t, err) + reader := readCloser.(*decryptReader) + reader.sequence = math.MaxUint64 - 2 + decrypted, err := io.ReadAll(reader) + require.NoError(t, err) + require.Equal(t, plaintext, decrypted) + require.NoError(t, reader.Close()) + require.Equal(t, uint64(math.MaxUint64), reader.sequence) +} + +func TestEncryptWriterRejectsDataWhenOnlyFinalSequenceRemainsBeforeSeal(t *testing.T) { + recipient, identity, err := GenerateKeyPair() + require.NoError(t, err) + + var destination bytes.Buffer + writeCloser, err := newEncryptWriter(&destination, recipient) + require.NoError(t, err) + writer := writeCloser.(*encryptWriter) + writer.sequence = math.MaxUint64 - 1 + before := append([]byte(nil), destination.Bytes()...) + + n, err := writer.Write([]byte("must-not-be-consumed")) + require.ErrorContains(t, err, "sequence exhausted") + require.Zero(t, n) + require.Equal(t, before, destination.Bytes()) + require.Equal(t, uint64(math.MaxUint64-1), writer.sequence) + + // A retry after clearing only the sticky API error proves the rejected call + // did not consume an HPKE nonce before the guard fired. + writer.writeErr = nil + writer.sequence = math.MaxUint64 - 2 + plaintext := []byte("nonce-zero-remains-unused") + n, err = writer.Write(plaintext) + require.NoError(t, err) + require.Equal(t, len(plaintext), n) + require.NoError(t, writer.Close()) + + readCloser, err := newDecryptReader(bytes.NewReader(destination.Bytes()), identity) + require.NoError(t, err) + reader := readCloser.(*decryptReader) + reader.sequence = math.MaxUint64 - 2 + decrypted, err := io.ReadAll(reader) + require.NoError(t, err) + require.Equal(t, plaintext, decrypted) + require.NoError(t, reader.Close()) +} + +func TestEncryptWriterRejectsFinalAtExhaustedSequenceBeforeSeal(t *testing.T) { + recipient, identity, err := GenerateKeyPair() + require.NoError(t, err) + + var destination bytes.Buffer + writeCloser, err := newEncryptWriter(&destination, recipient) + require.NoError(t, err) + writer := writeCloser.(*encryptWriter) + writer.sequence = math.MaxUint64 + before := append([]byte(nil), destination.Bytes()...) + + require.ErrorContains(t, writer.Close(), "sequence exhausted") + require.Equal(t, before, destination.Bytes()) + require.Equal(t, uint64(math.MaxUint64), writer.sequence) + + // Re-running only the close state at the last safe final sequence proves + // the rejected close did not consume the HPKE nonce. + writer.closed = false + writer.closeErr = nil + writer.sequence = math.MaxUint64 - 1 + require.NoError(t, writer.Close()) + + readCloser, err := newDecryptReader(bytes.NewReader(destination.Bytes()), identity) + require.NoError(t, err) + reader := readCloser.(*decryptReader) + reader.sequence = math.MaxUint64 - 1 + decrypted, err := io.ReadAll(reader) + require.NoError(t, err) + require.Empty(t, decrypted) + require.NoError(t, reader.Close()) +} + +func TestDecryptReaderRejectsDataWhenOnlyFinalSequenceRemainsBeforeOpen(t *testing.T) { + wire, identity := buildTestEnvelopeWithSingleFrame( + t, + math.MaxUint64-1, + testEnvelopeDataFrame, + []byte("must-not-be-reported"), + ) + source := &countingReader{reader: bytes.NewReader(wire)} + readCloser, err := newDecryptReader(source, identity) + require.NoError(t, err) + reader := readCloser.(*decryptReader) + reader.sequence = math.MaxUint64 - 1 + beforeFrame := source.bytesRead + output := []byte{0xa5} + + n, err := reader.Read(output) + require.ErrorContains(t, err, "sequence exhausted") + require.Zero(t, n) + require.Equal(t, []byte{0xa5}, output) + require.Equal(t, testEnvelopeFrameHeaderSize, source.bytesRead-beforeFrame) + require.Equal(t, uint64(math.MaxUint64-1), reader.sequence) + require.Nil(t, reader.plaintext) + afterRead := source.bytesRead + require.ErrorContains(t, reader.Close(), "sequence exhausted") + require.Equal(t, afterRead, source.bytesRead) +} + +func TestDecryptReaderRejectsFinalAtExhaustedSequenceBeforeOpen(t *testing.T) { + wire, identity := buildTestEnvelopeWithSingleFrame( + t, + math.MaxUint64, + testEnvelopeFinalFrame, + nil, + ) + source := &countingReader{reader: bytes.NewReader(wire)} + readCloser, err := newDecryptReader(source, identity) + require.NoError(t, err) + reader := readCloser.(*decryptReader) + reader.sequence = math.MaxUint64 + beforeFrame := source.bytesRead + output := []byte{0xa5} + + n, err := reader.Read(output) + require.ErrorContains(t, err, "sequence exhausted") + require.Zero(t, n) + require.Equal(t, []byte{0xa5}, output) + require.Equal(t, testEnvelopeFrameHeaderSize, source.bytesRead-beforeFrame) + require.Equal(t, uint64(math.MaxUint64), reader.sequence) + require.Nil(t, reader.plaintext) +} + +func TestEnvelopeRequiresOneTerminalAuthenticatedFinalFrame(t *testing.T) { + recipient, identity, err := GenerateKeyPair() + require.NoError(t, err) + plaintext := []byte("payload") + ciphertext := encryptTestEnvelope(t, recipient, plaintext) + _, _, frames := parseTestEnvelope(t, ciphertext) + require.Len(t, frames, 2) + + missingFinal := append([]byte(nil), ciphertext[:frames[1].start]...) + decrypted, err := decryptTestEnvelope(identity, missingFinal) + require.ErrorContains(t, err, "final frame") + require.Equal(t, plaintext, decrypted) + + duplicateFinal := append(append([]byte(nil), ciphertext...), ciphertext[frames[1].start:frames[1].end]...) + decrypted, err = decryptTestEnvelope(identity, duplicateFinal) + require.ErrorContains(t, err, "after final frame") + require.Equal(t, plaintext, decrypted) + + dataAfterFinal := append(append([]byte(nil), ciphertext...), ciphertext[frames[0].start:frames[0].end]...) + decrypted, err = decryptTestEnvelope(identity, dataAfterFinal) + require.ErrorContains(t, err, "after final frame") + require.Equal(t, plaintext, decrypted) +} + +func TestEnvelopeRejectsAuthenticatedFinalFrameContainingPlaintext(t *testing.T) { + wire, identity := buildTestEnvelopeWithFinalPlaintext(t, []byte("not empty")) + + plaintext, err := decryptTestEnvelope(identity, wire) + require.ErrorContains(t, err, "final frame contained plaintext") + require.Empty(t, plaintext) +} + +func TestEnvelopeRejectsUnsupportedTypeAndOversizedFrameBeforeAllocation(t *testing.T) { + recipient, identity, err := GenerateKeyPair() + require.NoError(t, err) + ciphertext := encryptTestEnvelope(t, recipient, []byte("payload")) + _, _, frames := parseTestEnvelope(t, ciphertext) + + unsupported := append([]byte(nil), ciphertext...) + unsupported[frames[0].start] = 2 + decrypted, err := decryptTestEnvelope(identity, unsupported) + require.ErrorContains(t, err, "type") + require.Empty(t, decrypted) + + oversized := append([]byte(nil), ciphertext...) + binary.BigEndian.PutUint32( + oversized[frames[0].start+9:], + uint32(testEnvelopeFrameSize+testEnvelopeAEADOverhead+1), + ) + decrypted, err = decryptTestEnvelope(identity, oversized) + require.ErrorContains(t, err, "too large") + require.Empty(t, decrypted) +} + +func TestEnvelopeRejectsEveryTruncationOffset(t *testing.T) { + recipient, identity, err := GenerateKeyPair() + require.NoError(t, err) + ciphertext := encryptTestEnvelope(t, recipient, nil) + + for cut := range len(ciphertext) { + plaintext, err := decryptTestEnvelope(identity, ciphertext[:cut]) + if err == nil { + t.Fatalf("truncation at byte %d of %d was accepted with plaintext %x", cut, len(ciphertext), plaintext) + } + require.Emptyf(t, plaintext, "truncation at byte %d reported plaintext", cut) + } +} + +func TestDecryptReaderCloseAuthenticatesUnreadRemainder(t *testing.T) { + recipient, identity, err := GenerateKeyPair() + require.NoError(t, err) + ciphertext := encryptTestEnvelope(t, recipient, testPlaintext(testEnvelopeFrameSize+31)) + + reader, err := newDecryptReader(bytes.NewReader(ciphertext), identity) + require.NoError(t, err) + var one [1]byte + n, err := reader.Read(one[:]) + require.NoError(t, err) + require.Equal(t, 1, n) + require.NoError(t, reader.Close()) + require.ErrorIs(t, readOnce(reader), io.ErrClosedPipe) + require.NoError(t, reader.Close()) + + _, _, frames := parseTestEnvelope(t, ciphertext) + truncated := ciphertext[:frames[len(frames)-1].start] + reader, err = newDecryptReader(bytes.NewReader(truncated), identity) + require.NoError(t, err) + n, err = reader.Read(one[:]) + require.NoError(t, err) + require.Equal(t, 1, n) + require.ErrorContains(t, reader.Close(), "final frame") + require.ErrorContains(t, reader.Close(), "final frame") +} + +func TestEncryptWriterCloseFinalizesOnceAndPreservesWriteFailure(t *testing.T) { + recipient, _, err := GenerateKeyPair() + require.NoError(t, err) + + var destination bytes.Buffer + writer, err := newEncryptWriter(&destination, recipient) + require.NoError(t, err) + _, err = writer.Write([]byte("payload")) + require.NoError(t, err) + require.NoError(t, writer.Close()) + sizeAfterClose := destination.Len() + require.ErrorContains(t, writer.Close(), "closed") + require.Equal(t, sizeAfterClose, destination.Len()) + _, err = writer.Write([]byte("more")) + require.ErrorContains(t, err, "closed") + + emptyEnvelope := encryptTestEnvelope(t, recipient, nil) + finalFailure := errors.New("final write failed") + finalDestination := "aWriter{ + remaining: len(emptyEnvelope) - 1, + failure: finalFailure, + } + writer, err = newEncryptWriter(finalDestination, recipient) + require.NoError(t, err) + require.ErrorIs(t, writer.Close(), finalFailure) + require.ErrorIs(t, writer.Close(), finalFailure) + + frameStart := testEnvelopeFrameStart(t, emptyEnvelope) + dataFailure := errors.New("data write failed") + dataDestination := "aWriter{ + remaining: frameStart + testEnvelopeFrameHeaderSize + 1, + failure: dataFailure, + } + writer, err = newEncryptWriter(dataDestination, recipient) + require.NoError(t, err) + n, err := writer.Write([]byte("payload")) + require.ErrorIs(t, err, dataFailure) + require.Zero(t, n) + require.ErrorIs(t, writer.Close(), dataFailure) +} + +func TestEnvelopeConstructorsRejectMissingInputs(t *testing.T) { + recipient, identity, err := GenerateKeyPair() + require.NoError(t, err) + + _, err = newEncryptWriter(nil, recipient) + require.ErrorContains(t, err, "destination") + _, err = newEncryptWriter(io.Discard, PublicKey{}) + require.ErrorContains(t, err, "recipient") + _, err = newDecryptReader(nil, identity) + require.ErrorContains(t, err, "source") + _, err = newDecryptReader(bytes.NewReader(nil), PrivateKey{}) + require.ErrorContains(t, err, "identity") +} + +func encryptTestEnvelope(t *testing.T, recipient PublicKey, plaintext []byte) []byte { + t.Helper() + var destination bytes.Buffer + writer, err := newEncryptWriter(&destination, recipient) + require.NoError(t, err) + n, err := writer.Write(plaintext) + require.NoError(t, err) + require.Equal(t, len(plaintext), n) + require.NoError(t, writer.Close()) + return append([]byte(nil), destination.Bytes()...) +} + +func decryptTestEnvelope(identity PrivateKey, ciphertext []byte) ([]byte, error) { + reader, err := newDecryptReader(bytes.NewReader(ciphertext), identity) + if err != nil { + return nil, err + } + plaintext, readErr := io.ReadAll(reader) + closeErr := reader.Close() + return plaintext, errors.Join(readErr, closeErr) +} + +func parseTestEnvelope(t *testing.T, ciphertext []byte) ([]byte, testEnvelopeHeader, []testEnvelopeFrame) { + t.Helper() + require.GreaterOrEqual(t, len(ciphertext), len(testEnvelopeMagic)+4) + require.Equal(t, testEnvelopeMagic, string(ciphertext[:len(testEnvelopeMagic)])) + headerSize := int(binary.BigEndian.Uint32(ciphertext[len(testEnvelopeMagic):])) + headerStart := len(testEnvelopeMagic) + 4 + frameStart := headerStart + headerSize + require.LessOrEqual(t, frameStart, len(ciphertext)) + headerBytes := append([]byte(nil), ciphertext[headerStart:frameStart]...) + var header testEnvelopeHeader + require.NoError(t, json.Unmarshal(headerBytes, &header)) + + var frames []testEnvelopeFrame + for offset := frameStart; offset < len(ciphertext); { + require.GreaterOrEqual(t, len(ciphertext)-offset, testEnvelopeFrameHeaderSize) + frame := testEnvelopeFrame{ + start: offset, + frameType: ciphertext[offset], + sequence: binary.BigEndian.Uint64(ciphertext[offset+1:]), + ciphertextSize: binary.BigEndian.Uint32(ciphertext[offset+9:]), + } + frame.ciphertextStart = offset + testEnvelopeFrameHeaderSize + frame.end = frame.ciphertextStart + int(frame.ciphertextSize) + require.LessOrEqual(t, frame.end, len(ciphertext)) + frames = append(frames, frame) + offset = frame.end + } + return headerBytes, header, frames +} + +func testEnvelopeFrameStart(t *testing.T, ciphertext []byte) int { + t.Helper() + require.GreaterOrEqual(t, len(ciphertext), len(testEnvelopeMagic)+4) + headerSize := int(binary.BigEndian.Uint32(ciphertext[len(testEnvelopeMagic):])) + return len(testEnvelopeMagic) + 4 + headerSize +} + +func envelopeJSONWithCaseVariantField(t *testing.T, valid []byte, field string) []byte { + t.Helper() + needle := []byte(fmt.Sprintf("%q:", field)) + replacement := []byte(fmt.Sprintf("%q:", strings.ToUpper(field))) + payload := bytes.Replace(valid, needle, replacement, 1) + require.NotEqual(t, valid, payload) + return payload +} + +func envelopeJSONWithDuplicateField(t *testing.T, valid []byte, field string, attackerFirst bool) []byte { + t.Helper() + require.True(t, bytes.HasPrefix(valid, []byte("{"))) + require.True(t, bytes.HasSuffix(valid, []byte("}"))) + attacker := []byte(fmt.Sprintf("%q:%q", field, "attacker-controlled")) + if attackerFirst { + payload := make([]byte, 0, len(valid)+len(attacker)+1) + payload = append(payload, '{') + payload = append(payload, attacker...) + payload = append(payload, ',') + return append(payload, valid[1:]...) + } + + payload := append([]byte(nil), valid[:len(valid)-1]...) + payload = append(payload, ',') + payload = append(payload, attacker...) + return append(payload, '}') +} + +func rewriteTestEnvelopeHeader( + t *testing.T, + ciphertext []byte, + mutate func(*testEnvelopeHeader), +) []byte { + t.Helper() + headerBytes, header, _ := parseTestEnvelope(t, ciphertext) + mutate(&header) + replacement, err := json.Marshal(header) + require.NoError(t, err) + require.NotEqual(t, headerBytes, replacement) + return replaceTestEnvelopeHeader(t, ciphertext, replacement) +} + +func replaceTestEnvelopeHeader(t *testing.T, ciphertext, replacement []byte) []byte { + t.Helper() + frameStart := testEnvelopeFrameStart(t, ciphertext) + require.LessOrEqual(t, len(replacement), int(^uint32(0))) + result := make([]byte, 0, len(testEnvelopeMagic)+4+len(replacement)+len(ciphertext)-frameStart) + result = append(result, testEnvelopeMagic...) + result = binary.BigEndian.AppendUint32(result, uint32(len(replacement))) + result = append(result, replacement...) + result = append(result, ciphertext[frameStart:]...) + return result +} + +func buildTestEnvelopeWithFinalPlaintext(t *testing.T, finalPlaintext []byte) ([]byte, PrivateKey) { + return buildTestEnvelopeWithSingleFrame(t, 0, testEnvelopeFinalFrame, finalPlaintext) +} + +func buildTestEnvelopeWithSingleFrame( + t *testing.T, + sequence uint64, + frameType byte, + plaintext []byte, +) ([]byte, PrivateKey) { + t.Helper() + rawPrivate, err := hpke.MLKEM1024().GenerateKey() + require.NoError(t, err) + rawPrivateBytes, err := rawPrivate.Bytes() + require.NoError(t, err) + rawPublic := rawPrivate.PublicKey() + digest := sha256.Sum256(rawPublic.Bytes()) + privateDocument := testKeyDocument{ + Format: KeyFormat, + Role: "private", + KEM: testKEMName, + KDF: testKDFName, + AEAD: testAEADName, + Material: base64.StdEncoding.EncodeToString(rawPrivateBytes), + Fingerprint: hex.EncodeToString(digest[:]), + } + identity, err := ReadPrivateKey(writeTestKeyDocument(t, privateDocument)) + require.NoError(t, err) + + encapsulation, sender, err := hpke.NewSender( + rawPublic, + hpke.HKDFSHA512(), + hpke.AES256GCM(), + []byte(testEnvelopeHPKEInfo), + ) + require.NoError(t, err) + header := testEnvelopeHeader{ + Format: EnvelopeFormat, + KEM: testKEMName, + KDF: testKDFName, + AEAD: testAEADName, + EncapsulatedKey: base64.StdEncoding.EncodeToString(encapsulation), + FrameSize: testEnvelopeFrameSize, + } + headerBytes, err := json.Marshal(header) + require.NoError(t, err) + headerDigest := sha256.Sum256(headerBytes) + final := frameType == testEnvelopeFinalFrame + ciphertext, err := sender.Seal(testFrameAAD(headerDigest[:], sequence, final), plaintext) + require.NoError(t, err) + + wire := make([]byte, 0, len(testEnvelopeMagic)+4+len(headerBytes)+testEnvelopeFrameHeaderSize+len(ciphertext)) + wire = append(wire, testEnvelopeMagic...) + wire = binary.BigEndian.AppendUint32(wire, uint32(len(headerBytes))) + wire = append(wire, headerBytes...) + wire = append(wire, frameType) + wire = binary.BigEndian.AppendUint64(wire, sequence) + wire = binary.BigEndian.AppendUint32(wire, uint32(len(ciphertext))) + wire = append(wire, ciphertext...) + return wire, identity +} + +func testFrameAAD(headerDigest []byte, sequence uint64, final bool) []byte { + aad := append([]byte(nil), headerDigest...) + aad = binary.BigEndian.AppendUint64(aad, sequence) + if final { + return append(aad, 1) + } + return append(aad, 0) +} + +func testPlaintext(size int) []byte { + pattern := []byte("frame-data-0123456789abcdef") + plaintext := bytes.Repeat(pattern, (size+len(pattern)-1)/len(pattern)) + return plaintext[:size] +} + +func readOnce(reader io.Reader) error { + var value [1]byte + _, err := reader.Read(value[:]) + return err +} + +type quotaWriter struct { + buffer bytes.Buffer + remaining int + failure error +} + +type countingReader struct { + reader io.Reader + bytesRead int +} + +func (s *countingReader) Read(p []byte) (int, error) { + n, err := s.reader.Read(p) + s.bytesRead += n + return n, err +} + +func (s *quotaWriter) Write(p []byte) (int, error) { + if s.remaining == 0 { + return 0, s.failure + } + if len(p) <= s.remaining { + s.remaining -= len(p) + return s.buffer.Write(p) + } + + n, _ := s.buffer.Write(p[:s.remaining]) + s.remaining = 0 + return n, s.failure +} diff --git a/ret/archive/key.go b/ret/archive/key.go new file mode 100644 index 00000000..50fbe722 --- /dev/null +++ b/ret/archive/key.go @@ -0,0 +1,486 @@ +package archive + +import ( + "bytes" + "crypto/hpke" + "crypto/sha256" + "encoding/base64" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "io" + "io/fs" + "os" + "path/filepath" + "strings" +) + +const ( + KeyFormat = "ret-hpke-key-v1" + EnvelopeFormat = "ret-encrypted-tar-v1" + + kemName = "ML-KEM-1024" + kdfName = "HKDF-SHA512" + aeadName = "AES-256-GCM" + + publicRole = "public" + privateRole = "private" + + publicKeyMaterialSize = 1568 + privateKeyMaterialSize = 64 + maxKeyDocumentSize = 8 * 1024 +) + +type hpkeSuite struct { + kem hpke.KEM + kdf hpke.KDF + aead hpke.AEAD +} + +var archiveHPKESuite = hpkeSuite{ + kem: hpke.MLKEM1024(), + kdf: hpke.HKDFSHA512(), + aead: hpke.AES256GCM(), +} + +// PublicKey is an opaque archive recipient key. +type PublicKey struct { + material [publicKeyMaterialSize]byte + valid bool +} + +// PrivateKey is an opaque archive identity key. +type PrivateKey struct { + material [privateKeyMaterialSize]byte + valid bool +} + +type keyDocument struct { + Format string `json:"format"` + Role string `json:"role"` + KEM string `json:"kem"` + KDF string `json:"kdf"` + AEAD string `json:"aead"` + Material string `json:"material"` + Fingerprint string `json:"public_fingerprint"` +} + +// GenerateKeyPair creates an ML-KEM-1024 archive recipient and identity. +func GenerateKeyPair() (PublicKey, PrivateKey, error) { + privateHPKEKey, err := archiveHPKESuite.kem.GenerateKey() + if err != nil { + return PublicKey{}, PrivateKey{}, fmt.Errorf("generate archive key pair: %w", err) + } + + privateMaterial, err := privateHPKEKey.Bytes() + if err != nil { + return PublicKey{}, PrivateKey{}, fmt.Errorf("serialize archive private key: %w", err) + } + if len(privateMaterial) != privateKeyMaterialSize { + return PublicKey{}, PrivateKey{}, fmt.Errorf( + "generated private key material has length %d, want %d", + len(privateMaterial), + privateKeyMaterialSize, + ) + } + + publicMaterial := privateHPKEKey.PublicKey().Bytes() + if len(publicMaterial) != publicKeyMaterialSize { + return PublicKey{}, PrivateKey{}, fmt.Errorf( + "generated public key material has length %d, want %d", + len(publicMaterial), + publicKeyMaterialSize, + ) + } + + var publicKey PublicKey + copy(publicKey.material[:], publicMaterial) + publicKey.valid = true + var privateKey PrivateKey + copy(privateKey.material[:], privateMaterial) + privateKey.valid = true + return publicKey, privateKey, nil +} + +// ReadPublicKey reads and validates a public archive key document. +func ReadPublicKey(path string) (PublicKey, error) { + document, err := readKeyDocument(path) + if err != nil { + return PublicKey{}, fmt.Errorf("read public key: %w", err) + } + material, err := validateKeyDocument(document, publicRole, publicKeyMaterialSize) + if err != nil { + return PublicKey{}, err + } + + publicHPKEKey, err := archiveHPKESuite.kem.NewPublicKey(material) + if err != nil { + return PublicKey{}, fmt.Errorf("parse public key material: %w", err) + } + if err := validatePublicFingerprint(document.Fingerprint, publicHPKEKey.Bytes()); err != nil { + return PublicKey{}, err + } + + var key PublicKey + copy(key.material[:], material) + key.valid = true + return key, nil +} + +// ReadPrivateKey reads and validates a private archive key document. +func ReadPrivateKey(path string) (PrivateKey, error) { + document, err := readKeyDocument(path) + if err != nil { + return PrivateKey{}, fmt.Errorf("read private key: %w", err) + } + material, err := validateKeyDocument(document, privateRole, privateKeyMaterialSize) + if err != nil { + return PrivateKey{}, err + } + + privateHPKEKey, err := archiveHPKESuite.kem.NewPrivateKey(material) + if err != nil { + return PrivateKey{}, fmt.Errorf("parse private key material: %w", err) + } + if err := validatePublicFingerprint(document.Fingerprint, privateHPKEKey.PublicKey().Bytes()); err != nil { + return PrivateKey{}, err + } + + var key PrivateKey + copy(key.material[:], material) + key.valid = true + return key, nil +} + +// WritePublicKey writes a public archive key document without replacing a path. +func WritePublicKey(path string, key PublicKey) error { + if !key.valid { + return fmt.Errorf("public key is required") + } + document := keyDocument{ + Format: KeyFormat, + Role: publicRole, + KEM: kemName, + KDF: kdfName, + AEAD: aeadName, + Material: base64.StdEncoding.EncodeToString(key.material[:]), + Fingerprint: fingerprint(key.material[:]), + } + return writeKeyDocument(path, 0o644, document) +} + +// WritePrivateKey writes a private archive key document without replacing a path. +func WritePrivateKey(path string, key PrivateKey) error { + owned, err := writePrivateKeyOwned(path, key) + if err != nil { + return err + } + return owned.release() +} + +func writePrivateKeyOwned(path string, key PrivateKey) (*ownedPath, error) { + if !key.valid { + return nil, fmt.Errorf("private key is required") + } + privateHPKEKey, err := archiveHPKESuite.kem.NewPrivateKey(key.material[:]) + if err != nil { + return nil, fmt.Errorf("parse private key material: %w", err) + } + document := keyDocument{ + Format: KeyFormat, + Role: privateRole, + KEM: kemName, + KDF: kdfName, + AEAD: aeadName, + Material: base64.StdEncoding.EncodeToString(key.material[:]), + Fingerprint: fingerprint(privateHPKEKey.PublicKey().Bytes()), + } + return writeKeyDocumentOwned(path, 0o600, document) +} + +// WriteKeyPair writes a private key followed by its public key without replacing +// either destination. A public-key failure removes only the private file created +// by this call when its recorded identity still owns the private pathname. +func WriteKeyPair( + privatePath string, + privateKey PrivateKey, + publicPath string, + publicKey PublicKey, +) error { + return writeKeyPair( + privatePath, + privateKey, + publicPath, + publicKey, + archiveOperations{}, + ) +} + +func writeKeyPair( + privatePath string, + privateKey PrivateKey, + publicPath string, + publicKey PublicKey, + operations archiveOperations, +) error { + privateOwned, err := writePrivateKeyOwned(privatePath, privateKey) + if err != nil { + return err + } + if err := WritePublicKey(publicPath, publicKey); err != nil { + return errors.Join(err, privateOwned.remove(operations)) + } + return privateOwned.release() +} + +// PublicFingerprint returns the lowercase SHA-256 fingerprint of the public key. +func PublicFingerprint(key PublicKey) string { + if !key.valid { + return "" + } + return fingerprint(key.material[:]) +} + +func readKeyDocument(path string) (document keyDocument, returnErr error) { + if strings.TrimSpace(path) == "" { + return document, fmt.Errorf("key path is required") + } + file, err := os.Open(path) + if err != nil { + return document, err + } + defer func() { + if err := file.Close(); err != nil { + returnErr = errors.Join(returnErr, fmt.Errorf("close key document: %w", err)) + } + }() + + payload, err := io.ReadAll(io.LimitReader(file, maxKeyDocumentSize+1)) + if err != nil { + return document, fmt.Errorf("read key JSON: %w", err) + } + if len(payload) > maxKeyDocumentSize { + return document, fmt.Errorf("key JSON exceeds %d bytes", maxKeyDocumentSize) + } + if err := validateExactJSONObject( + payload, + "format", + "role", + "kem", + "kdf", + "aead", + "material", + "public_fingerprint", + ); err != nil { + return document, fmt.Errorf("validate key JSON fields: %w", err) + } + + decoder := json.NewDecoder(bytes.NewReader(payload)) + decoder.DisallowUnknownFields() + if err := decoder.Decode(&document); err != nil { + return document, fmt.Errorf("decode key JSON: %w", err) + } + var trailing any + if err := decoder.Decode(&trailing); !errors.Is(err, io.EOF) { + if err == nil { + return document, fmt.Errorf("key JSON contains a trailing value") + } + return document, fmt.Errorf("decode trailing key JSON: %w", err) + } + return document, nil +} + +func validateExactJSONObject(payload []byte, expectedFields ...string) error { + expected := make(map[string]struct{}, len(expectedFields)) + for _, field := range expectedFields { + expected[field] = struct{}{} + } + + decoder := json.NewDecoder(bytes.NewReader(payload)) + token, err := decoder.Token() + if err != nil { + return fmt.Errorf("decode JSON object start: %w", err) + } + opening, ok := token.(json.Delim) + if !ok || opening != '{' { + return fmt.Errorf("JSON value must be an object") + } + + seen := make(map[string]struct{}, len(expectedFields)) + for decoder.More() { + token, err := decoder.Token() + if err != nil { + return fmt.Errorf("decode JSON field name: %w", err) + } + field, ok := token.(string) + if !ok { + return fmt.Errorf("JSON object field name must be a string") + } + if _, ok := expected[field]; !ok { + return fmt.Errorf( + "unknown field %q; top-level field names must use the exact lowercase schema", + field, + ) + } + if _, ok := seen[field]; ok { + return fmt.Errorf("duplicate top-level field %q", field) + } + seen[field] = struct{}{} + + var value json.RawMessage + if err := decoder.Decode(&value); err != nil { + return fmt.Errorf("decode JSON field %q: %w", field, err) + } + } + + token, err = decoder.Token() + if err != nil { + return fmt.Errorf("decode JSON object end: %w", err) + } + closing, ok := token.(json.Delim) + if !ok || closing != '}' { + return fmt.Errorf("JSON object is not terminated") + } + for _, field := range expectedFields { + if _, ok := seen[field]; !ok { + return fmt.Errorf("missing required top-level field %q", field) + } + } + return nil +} + +func validateKeyDocument(document keyDocument, expectedRole string, expectedMaterialSize int) ([]byte, error) { + if document.Format != KeyFormat { + return nil, fmt.Errorf("key format must be %q, got %q", KeyFormat, document.Format) + } + if document.Role != expectedRole { + return nil, fmt.Errorf("key role must be %q, got %q", expectedRole, document.Role) + } + if document.KEM != kemName { + return nil, fmt.Errorf("key KEM must be %q, got %q", kemName, document.KEM) + } + if document.KDF != kdfName { + return nil, fmt.Errorf("key KDF must be %q, got %q", kdfName, document.KDF) + } + if document.AEAD != aeadName { + return nil, fmt.Errorf("key AEAD must be %q, got %q", aeadName, document.AEAD) + } + if document.Material == "" { + return nil, fmt.Errorf("%s key material is required", expectedRole) + } + + material, err := base64.StdEncoding.Strict().DecodeString(document.Material) + if err != nil || base64.StdEncoding.EncodeToString(material) != document.Material { + if err == nil { + err = fmt.Errorf("encoding is not canonical") + } + return nil, fmt.Errorf("decode %s key material as canonical base64: %w", expectedRole, err) + } + if len(material) != expectedMaterialSize { + return nil, fmt.Errorf( + "%s key material has length %d, want %d", + expectedRole, + len(material), + expectedMaterialSize, + ) + } + return material, nil +} + +func validatePublicFingerprint(got string, publicMaterial []byte) error { + want := fingerprint(publicMaterial) + if len(got) != sha256.Size*2 { + return fmt.Errorf("public fingerprint must be %d lowercase hexadecimal characters", sha256.Size*2) + } + if _, err := hex.DecodeString(got); err != nil || strings.ToLower(got) != got { + return fmt.Errorf("public fingerprint must be %d lowercase hexadecimal characters", sha256.Size*2) + } + if got != want { + return fmt.Errorf("public fingerprint does not match public key material") + } + return nil +} + +func fingerprint(publicMaterial []byte) string { + digest := sha256.Sum256(publicMaterial) + return hex.EncodeToString(digest[:]) +} + +func writeKeyDocument(path string, mode fs.FileMode, document keyDocument) error { + owned, err := writeKeyDocumentOwned(path, mode, document) + if err != nil { + return err + } + return owned.release() +} + +func writeKeyDocumentOwned( + path string, + mode fs.FileMode, + document keyDocument, +) (*ownedPath, error) { + if strings.TrimSpace(path) == "" { + return nil, fmt.Errorf("key path is required") + } + absolute, err := filepath.Abs(path) + if err != nil { + return nil, fmt.Errorf("resolve key path: %w", err) + } + parentPath, err := filepath.EvalSymlinks(filepath.Dir(absolute)) + if err != nil { + return nil, fmt.Errorf("resolve key parent: %w", err) + } + parent, err := os.OpenRoot(parentPath) + if err != nil { + return nil, fmt.Errorf("open key parent: %w", err) + } + name := filepath.Base(absolute) + file, err := parent.OpenFile(name, os.O_WRONLY|os.O_CREATE|os.O_EXCL, mode) + if err != nil { + return nil, errors.Join( + fmt.Errorf("create key document: %w", err), + wrapKeyCloseError("key parent", parent.Close()), + ) + } + info, err := file.Stat() + if err != nil { + return nil, errors.Join( + fmt.Errorf("inspect created key document: %w", err), + fmt.Errorf("close key document: %w", file.Close()), + fmt.Errorf( + "ownership cleanup for key document: created identity is unavailable; preserving pathname", + ), + wrapKeyCloseError("key parent", parent.Close()), + ) + } + owned := &ownedPath{ + parent: parent, + handle: file, + name: name, + info: info, + description: "key document", + } + + var primaryErr error + if err := file.Chmod(mode); err != nil { + primaryErr = fmt.Errorf("set key document mode: %w", err) + } else { + encoder := json.NewEncoder(file) + encoder.SetIndent("", " ") + if err := encoder.Encode(document); err != nil { + primaryErr = fmt.Errorf("encode key document: %w", err) + } + } + if primaryErr == nil { + return owned, nil + } + return nil, errors.Join(primaryErr, owned.remove(archiveOperations{})) +} + +func wrapKeyCloseError(description string, err error) error { + if err == nil { + return nil + } + return fmt.Errorf("close %s: %w", description, err) +} diff --git a/ret/archive/key_test.go b/ret/archive/key_test.go new file mode 100644 index 00000000..4dcc40e0 --- /dev/null +++ b/ret/archive/key_test.go @@ -0,0 +1,495 @@ +package archive + +import ( + "bytes" + "crypto/hpke" + "crypto/sha256" + "encoding/base64" + "encoding/hex" + "encoding/json" + "fmt" + "os" + "path/filepath" + "reflect" + "strings" + "testing" + + "github.com/stretchr/testify/require" +) + +const ( + testKEMName = "ML-KEM-1024" + testKDFName = "HKDF-SHA512" + testAEADName = "AES-256-GCM" + testPublicMaterialSize = 1568 + testPrivateMaterialSize = 64 +) + +type testKeyDocument struct { + Format string `json:"format"` + Role string `json:"role"` + KEM string `json:"kem"` + KDF string `json:"kdf"` + AEAD string `json:"aead"` + Material string `json:"material"` + Fingerprint string `json:"public_fingerprint"` +} + +func TestGenerateKeyPairProducesOpaqueUsableKeys(t *testing.T) { + publicKey, privateKey, err := GenerateKeyPair() + require.NoError(t, err) + + dir := t.TempDir() + publicPath := filepath.Join(dir, "public.json") + privatePath := filepath.Join(dir, "private.json") + require.NoError(t, WritePublicKey(publicPath, publicKey)) + require.NoError(t, WritePrivateKey(privatePath, privateKey)) + + publicDocument := readTestKeyDocument(t, publicPath) + privateDocument := readTestKeyDocument(t, privatePath) + require.Equal(t, KeyFormat, publicDocument.Format) + require.Equal(t, KeyFormat, privateDocument.Format) + require.Equal(t, "public", publicDocument.Role) + require.Equal(t, "private", privateDocument.Role) + require.Equal(t, testKEMName, publicDocument.KEM) + require.Equal(t, testKEMName, privateDocument.KEM) + require.Equal(t, testKDFName, publicDocument.KDF) + require.Equal(t, testKDFName, privateDocument.KDF) + require.Equal(t, testAEADName, publicDocument.AEAD) + require.Equal(t, testAEADName, privateDocument.AEAD) + require.Equal(t, privateDocument.Fingerprint, publicDocument.Fingerprint) + require.Regexp(t, `^[0-9a-f]{64}$`, publicDocument.Fingerprint) + require.Len(t, decodeTestMaterial(t, publicDocument.Material), testPublicMaterialSize) + require.Len(t, decodeTestMaterial(t, privateDocument.Material), testPrivateMaterialSize) + require.Equal(t, publicDocument.Fingerprint, PublicFingerprint(publicKey)) + + readPublic, err := ReadPublicKey(publicPath) + require.NoError(t, err) + require.Equal(t, publicDocument.Fingerprint, PublicFingerprint(readPublic)) + _, err = ReadPrivateKey(privatePath) + require.NoError(t, err) +} + +func TestKeyWrappersExposeNoRawHPKEKeys(t *testing.T) { + for _, wrapper := range []reflect.Type{ + reflect.TypeFor[PublicKey](), + reflect.TypeFor[PrivateKey](), + } { + for fieldIndex := range wrapper.NumField() { + field := wrapper.Field(fieldIndex) + require.Falsef(t, field.IsExported(), "%s field %s is exported", wrapper, field.Name) + require.NotContains(t, field.Type.String(), "hpke") + require.NotEqual(t, "crypto/hpke", field.Type.PkgPath()) + } + for methodIndex := range wrapper.NumMethod() { + method := wrapper.Method(methodIndex) + require.NotContains(t, method.Type.String(), "hpke") + } + } +} + +func TestPublicFingerprintIsLowercaseSHA256OfPublicMaterial(t *testing.T) { + document, _, publicBytes := newTestKeyDocuments(t) + path := writeTestKeyDocument(t, document) + + key, err := ReadPublicKey(path) + require.NoError(t, err) + + digest := sha256.Sum256(publicBytes) + require.Equal(t, hex.EncodeToString(digest[:]), PublicFingerprint(key)) + require.Equal(t, PublicFingerprint(key), PublicFingerprint(key)) + require.Empty(t, PublicFingerprint(PublicKey{})) +} + +func TestReadKeysRejectMalformedDocuments(t *testing.T) { + publicDocument, privateDocument, _ := newTestKeyDocuments(t) + tests := []struct { + name string + private bool + mutate func(*testKeyDocument) + match string + }{ + {name: "retriever format", mutate: func(document *testKeyDocument) { + document.Format = "retriever-hpke-key-v1" + }, match: KeyFormat}, + {name: "empty format", mutate: func(document *testKeyDocument) { + document.Format = "" + }, match: KeyFormat}, + {name: "wrong public role", mutate: func(document *testKeyDocument) { + document.Role = "private" + }, match: "public"}, + {name: "wrong private role", private: true, mutate: func(document *testKeyDocument) { + document.Role = "public" + }, match: "private"}, + {name: "wrong KEM", mutate: func(document *testKeyDocument) { + document.KEM = "ML-KEM-768" + }, match: testKEMName}, + {name: "wrong KDF", mutate: func(document *testKeyDocument) { + document.KDF = "HKDF-SHA256" + }, match: testKDFName}, + {name: "wrong AEAD", mutate: func(document *testKeyDocument) { + document.AEAD = "AES-128-GCM" + }, match: testAEADName}, + {name: "empty material", mutate: func(document *testKeyDocument) { + document.Material = "" + }, match: "material"}, + {name: "invalid base64", mutate: func(document *testKeyDocument) { + document.Material = "not base64!" + }, match: "base64"}, + {name: "noncanonical base64", mutate: func(document *testKeyDocument) { + document.Material = document.Material[:8] + "\n" + document.Material[8:] + }, match: "base64"}, + {name: "short public material", mutate: func(document *testKeyDocument) { + document.Material = base64.StdEncoding.EncodeToString(make([]byte, testPublicMaterialSize-1)) + }, match: "1568"}, + {name: "long public material", mutate: func(document *testKeyDocument) { + document.Material = base64.StdEncoding.EncodeToString(make([]byte, testPublicMaterialSize+1)) + }, match: "1568"}, + {name: "short private material", private: true, mutate: func(document *testKeyDocument) { + document.Material = base64.StdEncoding.EncodeToString(make([]byte, testPrivateMaterialSize-1)) + }, match: "64"}, + {name: "long private material", private: true, mutate: func(document *testKeyDocument) { + document.Material = base64.StdEncoding.EncodeToString(make([]byte, testPrivateMaterialSize+1)) + }, match: "64"}, + {name: "empty fingerprint", mutate: func(document *testKeyDocument) { + document.Fingerprint = "" + }, match: "fingerprint"}, + {name: "uppercase fingerprint", mutate: func(document *testKeyDocument) { + document.Fingerprint = strings.ToUpper(document.Fingerprint) + }, match: "fingerprint"}, + {name: "wrong fingerprint", mutate: func(document *testKeyDocument) { + document.Fingerprint = strings.Repeat("0", 64) + }, match: "fingerprint"}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + document := publicDocument + if test.private { + document = privateDocument + } + test.mutate(&document) + path := writeTestKeyDocument(t, document) + + var err error + if test.private { + _, err = ReadPrivateKey(path) + } else { + _, err = ReadPublicKey(path) + } + require.ErrorContains(t, err, test.match) + }) + } +} + +func TestReadKeysRejectInvalidJSONUnknownFieldsAndTrailingValues(t *testing.T) { + publicDocument, _, _ := newTestKeyDocuments(t) + valid, err := json.Marshal(publicDocument) + require.NoError(t, err) + + tests := []struct { + name string + payload string + match string + }{ + {name: "invalid JSON", payload: `{`, match: "JSON"}, + {name: "unknown field", payload: strings.TrimSuffix(string(valid), "}") + `,"extra":true}`, match: "unknown field"}, + {name: "trailing value", payload: string(valid) + "\n{}", match: "trailing"}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + path := filepath.Join(t.TempDir(), "key.json") + require.NoError(t, os.WriteFile(path, []byte(test.payload), 0o600)) + _, err := ReadPublicKey(path) + require.ErrorContains(t, err, test.match) + }) + } +} + +func TestReadPublicKeyRequiresEachExactLowercaseFieldOnce(t *testing.T) { + publicDocument, _, _ := newTestKeyDocuments(t) + valid, err := json.Marshal(publicDocument) + require.NoError(t, err) + fields := []string{ + "format", + "role", + "kem", + "kdf", + "aead", + "material", + "public_fingerprint", + } + + for _, field := range fields { + t.Run(field+"/case_variant", func(t *testing.T) { + payload := keyJSONWithCaseVariantField(t, valid, field) + _, err := ReadPublicKey(writeTestKeyPayload(t, payload)) + require.ErrorContains(t, err, "exact lowercase") + }) + t.Run(field+"/attacker_duplicate_first", func(t *testing.T) { + payload := keyJSONWithDuplicateField(t, valid, field, true) + _, err := ReadPublicKey(writeTestKeyPayload(t, payload)) + require.ErrorContains(t, err, "duplicate") + }) + t.Run(field+"/attacker_duplicate_last", func(t *testing.T) { + payload := keyJSONWithDuplicateField(t, valid, field, false) + _, err := ReadPublicKey(writeTestKeyPayload(t, payload)) + require.ErrorContains(t, err, "duplicate") + }) + } +} + +func TestReadPublicKeyRejectsStructurallyInvalidPublicMaterial(t *testing.T) { + document, _, _ := newTestKeyDocuments(t) + material := make([]byte, testPublicMaterialSize) + for index := range material { + material[index] = 0xff + } + document.Material = base64.StdEncoding.EncodeToString(material) + digest := sha256.Sum256(material) + document.Fingerprint = hex.EncodeToString(digest[:]) + + _, err := ReadPublicKey(writeTestKeyDocument(t, document)) + require.ErrorContains(t, err, "public key material") +} + +func TestKeyWritersUseRequiredModesAndRefuseOverwrite(t *testing.T) { + publicDocument, privateDocument, _ := newTestKeyDocuments(t) + inputPublic, err := ReadPublicKey(writeTestKeyDocument(t, publicDocument)) + require.NoError(t, err) + inputPrivate, err := ReadPrivateKey(writeTestKeyDocument(t, privateDocument)) + require.NoError(t, err) + + dir := t.TempDir() + publicPath := filepath.Join(dir, "nested", "public.json") + privatePath := filepath.Join(dir, "nested", "private.json") + require.NoError(t, os.MkdirAll(filepath.Dir(publicPath), 0o755)) + require.NoError(t, WritePublicKey(publicPath, inputPublic)) + require.NoError(t, WritePrivateKey(privatePath, inputPrivate)) + + publicInfo, err := os.Stat(publicPath) + require.NoError(t, err) + require.Equal(t, os.FileMode(0o644), publicInfo.Mode().Perm()) + privateInfo, err := os.Stat(privatePath) + require.NoError(t, err) + require.Equal(t, os.FileMode(0o600), privateInfo.Mode().Perm()) + + publicBefore, err := os.ReadFile(publicPath) + require.NoError(t, err) + privateBefore, err := os.ReadFile(privatePath) + require.NoError(t, err) + require.Error(t, WritePublicKey(publicPath, inputPublic)) + require.Error(t, WritePrivateKey(privatePath, inputPrivate)) + publicAfter, err := os.ReadFile(publicPath) + require.NoError(t, err) + privateAfter, err := os.ReadFile(privatePath) + require.NoError(t, err) + require.Equal(t, publicBefore, publicAfter) + require.Equal(t, privateBefore, privateAfter) +} + +func TestKeyReadersAndWritersRejectEmptyInputs(t *testing.T) { + _, err := ReadPublicKey("") + require.ErrorContains(t, err, "path") + _, err = ReadPrivateKey("") + require.ErrorContains(t, err, "path") + + publicPath := filepath.Join(t.TempDir(), "public.json") + require.ErrorContains(t, WritePublicKey(publicPath, PublicKey{}), "public key") + _, err = os.Stat(publicPath) + require.ErrorIs(t, err, os.ErrNotExist) + + privatePath := filepath.Join(t.TempDir(), "private.json") + require.ErrorContains(t, WritePrivateKey(privatePath, PrivateKey{}), "private key") + _, err = os.Stat(privatePath) + require.ErrorIs(t, err, os.ErrNotExist) +} + +func TestPrivateKeyRollbackPreservesReplacement(t *testing.T) { + // Break caught: deleting a replacement installed at the private-key path + // when public-key publication later requires rollback. + _, private, err := GenerateKeyPair() + require.NoError(t, err) + privatePath := filepath.Join(t.TempDir(), "private.json") + owned, err := writePrivateKeyOwned(privatePath, private) + require.NoError(t, err) + require.NoError(t, os.Remove(privatePath)) + require.NoError(t, os.WriteFile(privatePath, []byte("preserve replacement"), 0o600)) + + err = owned.remove(archiveOperations{}) + + require.ErrorContains(t, err, "ownership cleanup") + require.Equal(t, []byte("preserve replacement"), mustReadArchiveTestFile(t, privatePath)) +} + +func TestPrivateKeyRollbackSanitizesQuarantinedShell(t *testing.T) { + public, private, err := GenerateKeyPair() + require.NoError(t, err) + parent := t.TempDir() + privatePath := filepath.Join(parent, "private.json") + publicPath := filepath.Join(parent, "public.json") + require.NoError(t, os.WriteFile(publicPath, []byte("preserve public"), 0o600)) + + err = WriteKeyPair(privatePath, private, publicPath, public) + + require.ErrorContains(t, err, "create key document") + require.ErrorContains(t, err, "ownership cleanup") + require.NoFileExists(t, privatePath) + require.Equal(t, []byte("preserve public"), mustReadArchiveTestFile(t, publicPath)) + quarantines := archiveCleanupQuarantinePaths(t, parent) + require.Len(t, quarantines, 1) + info, statErr := os.Stat(quarantines[0]) + require.NoError(t, statErr) + require.True(t, info.Mode().IsRegular()) + require.Zero(t, info.Size()) +} + +func TestPrivateKeyRollbackPreservesOriginalWhenSanitizationFails(t *testing.T) { + // Break caught: moving a private-key file to quarantine after its retained + // handle could not be truncated or synced successfully. + tests := []struct { + name string + operations archiveOperations + wantPayload bool + match string + }{ + { + name: "truncate", + operations: archiveOperations{ + truncateOwnedFile: func(_ *os.File, _ int64) error { + return fmt.Errorf("injected truncate failure") + }, + }, + wantPayload: true, + match: "injected truncate failure", + }, + { + name: "sync", + operations: archiveOperations{ + syncOwnedFile: func(_ *os.File) error { + return fmt.Errorf("injected sync failure") + }, + }, + match: "injected sync failure", + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + public, private, err := GenerateKeyPair() + require.NoError(t, err) + parent := t.TempDir() + privatePath := filepath.Join(parent, "private.json") + publicPath := filepath.Join(parent, "public.json") + require.NoError(t, os.WriteFile(publicPath, []byte("preserve public"), 0o600)) + + err = writeKeyPair( + privatePath, + private, + publicPath, + public, + test.operations, + ) + + require.ErrorContains(t, err, test.match) + require.ErrorContains(t, err, "ownership cleanup") + privatePayload := mustReadArchiveTestFile(t, privatePath) + if test.wantPayload { + require.NotEmpty(t, privatePayload) + } else { + require.Empty(t, privatePayload) + } + require.Equal(t, []byte("preserve public"), mustReadArchiveTestFile(t, publicPath)) + require.Empty(t, archiveCleanupQuarantinePaths(t, parent)) + }) + } +} + +func newTestKeyDocuments(t *testing.T) (testKeyDocument, testKeyDocument, []byte) { + t.Helper() + privateKey, err := hpke.MLKEM1024().GenerateKey() + require.NoError(t, err) + privateBytes, err := privateKey.Bytes() + require.NoError(t, err) + publicBytes := privateKey.PublicKey().Bytes() + require.Len(t, publicBytes, testPublicMaterialSize) + require.Len(t, privateBytes, testPrivateMaterialSize) + + digest := sha256.Sum256(publicBytes) + fingerprint := hex.EncodeToString(digest[:]) + base := testKeyDocument{ + Format: KeyFormat, + KEM: testKEMName, + KDF: testKDFName, + AEAD: testAEADName, + Fingerprint: fingerprint, + } + publicDocument := base + publicDocument.Role = "public" + publicDocument.Material = base64.StdEncoding.EncodeToString(publicBytes) + privateDocument := base + privateDocument.Role = "private" + privateDocument.Material = base64.StdEncoding.EncodeToString(privateBytes) + return publicDocument, privateDocument, publicBytes +} + +func writeTestKeyDocument(t *testing.T, document testKeyDocument) string { + t.Helper() + path := filepath.Join(t.TempDir(), "key.json") + payload, err := json.Marshal(document) + require.NoError(t, err) + require.NoError(t, os.WriteFile(path, payload, 0o600)) + return path +} + +func readTestKeyDocument(t *testing.T, path string) testKeyDocument { + t.Helper() + payload, err := os.ReadFile(path) + require.NoError(t, err) + var document testKeyDocument + require.NoError(t, json.Unmarshal(payload, &document)) + return document +} + +func decodeTestMaterial(t *testing.T, value string) []byte { + t.Helper() + material, err := base64.StdEncoding.DecodeString(value) + require.NoError(t, err) + return material +} + +func writeTestKeyPayload(t *testing.T, payload []byte) string { + t.Helper() + path := filepath.Join(t.TempDir(), "key.json") + require.NoError(t, os.WriteFile(path, payload, 0o600)) + return path +} + +func keyJSONWithCaseVariantField(t *testing.T, valid []byte, field string) []byte { + t.Helper() + needle := []byte(fmt.Sprintf("%q:", field)) + replacement := []byte(fmt.Sprintf("%q:", strings.ToUpper(field))) + payload := bytes.Replace(valid, needle, replacement, 1) + require.NotEqual(t, valid, payload) + return payload +} + +func keyJSONWithDuplicateField(t *testing.T, valid []byte, field string, attackerFirst bool) []byte { + t.Helper() + require.True(t, bytes.HasPrefix(valid, []byte("{"))) + require.True(t, bytes.HasSuffix(valid, []byte("}"))) + attacker := []byte(fmt.Sprintf("%q:%q", field, "attacker-controlled")) + if attackerFirst { + payload := make([]byte, 0, len(valid)+len(attacker)+1) + payload = append(payload, '{') + payload = append(payload, attacker...) + payload = append(payload, ',') + return append(payload, valid[1:]...) + } + + payload := append([]byte(nil), valid[:len(valid)-1]...) + payload = append(payload, ',') + payload = append(payload, attacker...) + return append(payload, '}') +} diff --git a/ret/archive/operations.go b/ret/archive/operations.go new file mode 100644 index 00000000..f6b8758e --- /dev/null +++ b/ret/archive/operations.go @@ -0,0 +1,81 @@ +package archive + +import ( + "io/fs" + "os" +) + +type archiveOperations struct { + beforeOwnedEntryQuarantine func() error + afterOwnedEntryQuarantine func(quarantineName string) error + truncateOwnedFile func(file *os.File, size int64) error + syncOwnedFile func(file *os.File) error + readOwnedDirectory func(root *os.Root) ([]fs.DirEntry, error) + removeOwnedDirectoryEntry func(root *os.Root, name string) error + openRoot func(parent *os.Root, name string) (*os.Root, error) + statRoot func(root *os.Root) (fs.FileInfo, error) + afterExtractFileSync func(root *os.Root, relative string) error +} + +func (s archiveOperations) runBeforeOwnedEntryQuarantine() error { + if s.beforeOwnedEntryQuarantine == nil { + return nil + } + return s.beforeOwnedEntryQuarantine() +} + +func (s archiveOperations) runAfterOwnedEntryQuarantine(quarantineName string) error { + if s.afterOwnedEntryQuarantine == nil { + return nil + } + return s.afterOwnedEntryQuarantine(quarantineName) +} + +func (s archiveOperations) runTruncateOwnedFile(file *os.File, size int64) error { + if s.truncateOwnedFile == nil { + return file.Truncate(size) + } + return s.truncateOwnedFile(file, size) +} + +func (s archiveOperations) runSyncOwnedFile(file *os.File) error { + if s.syncOwnedFile == nil { + return file.Sync() + } + return s.syncOwnedFile(file) +} + +func (s archiveOperations) runReadOwnedDirectory(root *os.Root) ([]fs.DirEntry, error) { + if s.readOwnedDirectory == nil { + return fs.ReadDir(root.FS(), ".") + } + return s.readOwnedDirectory(root) +} + +func (s archiveOperations) runRemoveOwnedDirectoryEntry(root *os.Root, name string) error { + if s.removeOwnedDirectoryEntry == nil { + return root.RemoveAll(name) + } + return s.removeOwnedDirectoryEntry(root, name) +} + +func (s archiveOperations) runOpenRoot(parent *os.Root, name string) (*os.Root, error) { + if s.openRoot == nil { + return parent.OpenRoot(name) + } + return s.openRoot(parent, name) +} + +func (s archiveOperations) runStatRoot(root *os.Root) (fs.FileInfo, error) { + if s.statRoot == nil { + return root.Stat(".") + } + return s.statRoot(root) +} + +func (s archiveOperations) runAfterExtractFileSync(root *os.Root, relative string) error { + if s.afterExtractFileSync == nil { + return nil + } + return s.afterExtractFileSync(root, relative) +} diff --git a/ret/archive/pack.go b/ret/archive/pack.go new file mode 100644 index 00000000..5283fe84 --- /dev/null +++ b/ret/archive/pack.go @@ -0,0 +1,426 @@ +package archive + +import ( + "context" + "crypto/rand" + "encoding/hex" + "errors" + "fmt" + "io" + "os" + "path/filepath" + "runtime" + "strings" + + "github.com/specterops/dawgs/ret/observe" +) + +type CreateConfig struct { + CollectionDirectory string + ArchivePath string + Recipient PublicKey + Observer observe.Observer +} + +type archiveDestination struct { + parentPath string + base string +} + +func Create(ctx context.Context, config CreateConfig) (resultErr error) { + return create(ctx, config, runtime.GOOS, archiveOperations{}) +} + +func create( + ctx context.Context, + config CreateConfig, + platform string, + operations archiveOperations, +) (resultErr error) { + if err := requirePlatformSupport(platform); err != nil { + return err + } + if err := validateCreateConfig(config); err != nil { + return err + } + if err := ctx.Err(); err != nil { + return fmt.Errorf("create archive: %w", err) + } + + destination, err := prepareArchiveDestination(config.CollectionDirectory, config.ArchivePath) + if err != nil { + return err + } + plan, err := collectionTarPaths(ctx, config.CollectionDirectory, config.Observer) + if err != nil { + return err + } + if err := ctx.Err(); err != nil { + return fmt.Errorf("create archive after collection verification: %w", err) + } + + parent, err := os.OpenRoot(destination.parentPath) + if err != nil { + return fmt.Errorf("open archive destination directory: %w", err) + } + defer func() { + resultErr = errors.Join(resultErr, wrapCreateCloseError("archive destination directory", parent.Close())) + }() + directory, err := openDirectoryMatchingRoot(destination.parentPath, parent) + if err != nil { + return err + } + defer func() { + resultErr = errors.Join(resultErr, wrapCreateCloseError("archive destination directory file", directory.Close())) + }() + if _, err := parent.Lstat(destination.base); err == nil { + return fmt.Errorf("archive destination exists: %s", config.ArchivePath) + } else if !errors.Is(err, os.ErrNotExist) { + return fmt.Errorf("inspect archive destination: %w", err) + } + + temporaryName, file, temporaryHandle, createdTemporaryInfo, err := createArchiveTemporary( + parent, + destination.base, + operations, + ) + if err != nil { + return err + } + var encrypted io.WriteCloser + published := false + defer func() { + if encrypted != nil { + resultErr = errors.Join(resultErr, wrapCreateCloseError("archive encryption writer", encrypted.Close())) + } + if !published { + resultErr = errors.Join( + resultErr, + sanitizeAndRemoveOwnedEntry( + parent, + temporaryName, + createdTemporaryInfo, + temporaryHandle, + "temporary archive", + operations, + ), + ) + } + if file != nil { + resultErr = errors.Join(resultErr, wrapCreateCloseError("temporary archive", file.Close())) + } + if temporaryHandle != nil { + resultErr = errors.Join( + resultErr, + wrapCreateCloseError("temporary archive identity handle", temporaryHandle.Close()), + ) + } + }() + + encrypted, err = newEncryptWriter(file, config.Recipient) + if err != nil { + return fmt.Errorf("start archive encryption: %w", err) + } + if err := writeCollectionTar( + ctx, + encrypted, + config.CollectionDirectory, + plan, + config.Observer, + ); err != nil { + return fmt.Errorf("create collection TAR: %w", err) + } + + closeEncryptionErr := encrypted.Close() + encrypted = nil + syncErr := file.Sync() + temporaryInfo, statErr := file.Stat() + closeFileErr := file.Close() + file = nil + if err := errors.Join( + wrapCreateCloseError("archive encryption writer", closeEncryptionErr), + wrapCreateSyncError("temporary archive", syncErr), + wrapCreateStatError("temporary archive", statErr), + wrapCreateCloseError("temporary archive", closeFileErr), + ); err != nil { + return err + } + if err := requireDirectoryPathIdentity(destination.parentPath, directory); err != nil { + return fmt.Errorf("validate archive destination directory before publication: %w", err) + } + if err := requireRootAndPhysicalEntryIdentity( + parent, + destination.parentPath, + temporaryName, + "temporary archive", + temporaryInfo, + ); err != nil { + return err + } + if err := ctx.Err(); err != nil { + return fmt.Errorf("publish archive: %w", err) + } + if _, err := parent.Lstat(destination.base); err == nil { + return fmt.Errorf("archive destination exists: %s", config.ArchivePath) + } else if !errors.Is(err, os.ErrNotExist) { + return fmt.Errorf("inspect archive destination before publication: %w", err) + } + if err := renameNoReplace(directory, destination.parentPath, temporaryName, destination.base); err != nil { + return fmt.Errorf("publish archive: %w", err) + } + published = true + if err := requireRootAndPhysicalEntryIdentity( + parent, + destination.parentPath, + destination.base, + "published archive", + temporaryInfo, + ); err != nil { + return err + } + if err := directory.Sync(); err != nil { + return fmt.Errorf("sync archive destination directory: %w", err) + } + return nil +} + +func validateCreateConfig(config CreateConfig) error { + if strings.TrimSpace(config.CollectionDirectory) == "" { + return fmt.Errorf("collection directory is required") + } + if strings.TrimSpace(config.ArchivePath) == "" { + return fmt.Errorf("archive path is required") + } + if !config.Recipient.valid { + return fmt.Errorf("recipient public key is required") + } + return nil +} + +func prepareArchiveDestination(collectionRoot, archivePath string) (archiveDestination, error) { + if _, err := os.Lstat(archivePath); err == nil { + return archiveDestination{}, fmt.Errorf("archive destination exists: %s", archivePath) + } else if !errors.Is(err, os.ErrNotExist) { + return archiveDestination{}, fmt.Errorf("inspect archive destination: %w", err) + } + + absoluteCollection, err := filepath.Abs(collectionRoot) + if err != nil { + return archiveDestination{}, fmt.Errorf("resolve collection directory: %w", err) + } + absoluteArchive, err := filepath.Abs(archivePath) + if err != nil { + return archiveDestination{}, fmt.Errorf("resolve archive path: %w", err) + } + if pathWithin(absoluteCollection, absoluteArchive) { + return archiveDestination{}, fmt.Errorf("archive path must not be inside the collection") + } + + physicalCollection, err := filepath.EvalSymlinks(absoluteCollection) + if err != nil { + return archiveDestination{}, fmt.Errorf("resolve physical collection directory: %w", err) + } + physicalParent, err := filepath.EvalSymlinks(filepath.Dir(absoluteArchive)) + if err != nil { + return archiveDestination{}, fmt.Errorf("resolve physical archive parent: %w", err) + } + physicalArchive := filepath.Join(physicalParent, filepath.Base(absoluteArchive)) + if pathWithin(physicalCollection, physicalArchive) { + return archiveDestination{}, fmt.Errorf("archive path must not be physically inside the collection") + } + if _, err := os.Lstat(physicalArchive); err == nil { + return archiveDestination{}, fmt.Errorf("archive destination exists: %s", archivePath) + } else if !errors.Is(err, os.ErrNotExist) { + return archiveDestination{}, fmt.Errorf("inspect physical archive destination: %w", err) + } + parentInfo, err := os.Lstat(physicalParent) + if err != nil { + return archiveDestination{}, fmt.Errorf("inspect physical archive parent: %w", err) + } + if parentInfo.Mode()&os.ModeSymlink != 0 || !parentInfo.IsDir() { + return archiveDestination{}, fmt.Errorf("archive parent is not a directory") + } + return archiveDestination{parentPath: physicalParent, base: filepath.Base(physicalArchive)}, nil +} + +func pathWithin(root, candidate string) bool { + relative, err := filepath.Rel(filepath.Clean(root), filepath.Clean(candidate)) + if err != nil { + return false + } + return relative == "." || + (relative != ".." && !strings.HasPrefix(relative, ".."+string(filepath.Separator))) +} + +func openDirectoryMatchingRoot(path string, root *os.Root) (*os.File, error) { + rootInfo, err := root.Stat(".") + if err != nil { + return nil, fmt.Errorf("inspect pinned directory root: %w", err) + } + directory, err := os.Open(path) + if err != nil { + return nil, fmt.Errorf("open pinned directory file: %w", err) + } + directoryInfo, err := directory.Stat() + if err != nil { + return nil, errors.Join( + fmt.Errorf("inspect pinned directory file: %w", err), + wrapCreateCloseError("pinned directory file", directory.Close()), + ) + } + if !rootInfo.IsDir() || !directoryInfo.IsDir() || !os.SameFile(rootInfo, directoryInfo) { + return nil, errors.Join( + fmt.Errorf("directory changed while pinning"), + wrapCreateCloseError("pinned directory file", directory.Close()), + ) + } + return directory, nil +} + +func requireDirectoryPathIdentity(path string, directory *os.File) error { + pathInfo, err := os.Lstat(path) + if err != nil { + return fmt.Errorf("inspect directory path: %w", err) + } + if pathInfo.Mode()&os.ModeSymlink != 0 || !pathInfo.IsDir() { + return fmt.Errorf("directory path is not a non-symlink directory") + } + pinnedInfo, err := directory.Stat() + if err != nil { + return fmt.Errorf("inspect pinned directory: %w", err) + } + if !pinnedInfo.IsDir() || !os.SameFile(pathInfo, pinnedInfo) { + return fmt.Errorf("directory path changed while processing") + } + return nil +} + +func requireRootAndPhysicalEntryIdentity( + root *os.Root, + parentPath string, + name string, + description string, + expected os.FileInfo, +) error { + rootedInfo, err := root.Lstat(name) + if err != nil { + return fmt.Errorf("inspect rooted %s: %w", description, err) + } + physicalInfo, err := os.Lstat(filepath.Join(parentPath, name)) + if err != nil { + return fmt.Errorf("inspect physical %s: %w", description, err) + } + if rootedInfo.Mode()&os.ModeSymlink != 0 || + physicalInfo.Mode()&os.ModeSymlink != 0 || + rootedInfo.Mode().Type() != expected.Mode().Type() || + physicalInfo.Mode().Type() != expected.Mode().Type() || + !os.SameFile(expected, rootedInfo) || + !os.SameFile(expected, physicalInfo) { + return fmt.Errorf("%s changed while processing", description) + } + return nil +} + +func createArchiveTemporary( + parent *os.Root, + archiveBase string, + operations archiveOperations, +) (string, *os.File, *os.File, os.FileInfo, error) { + for range 100 { + var random [16]byte + if _, err := rand.Read(random[:]); err != nil { + return "", nil, nil, nil, fmt.Errorf("generate temporary archive name: %w", err) + } + name := "." + archiveBase + ".create-" + hex.EncodeToString(random[:]) + ".tmp" + file, err := parent.OpenFile(name, os.O_RDWR|os.O_CREATE|os.O_EXCL, 0o600) + if err == nil { + info, statErr := file.Stat() + if statErr != nil { + return "", nil, nil, nil, errors.Join( + fmt.Errorf("inspect temporary archive: %w", statErr), + wrapCreateCloseError("temporary archive", file.Close()), + fmt.Errorf( + "ownership cleanup for temporary archive: created identity is unavailable; preserving pathname", + ), + ) + } + identityHandle, openErr := parent.OpenFile(name, os.O_RDWR, 0) + if openErr != nil { + return "", nil, nil, nil, errors.Join( + fmt.Errorf("pin temporary archive identity: %w", openErr), + sanitizeAndRemoveOwnedEntry( + parent, + name, + info, + file, + "temporary archive", + operations, + ), + wrapCreateCloseError("temporary archive", file.Close()), + ) + } + identityInfo, identityStatErr := identityHandle.Stat() + if identityStatErr != nil || + !identityInfo.Mode().IsRegular() || + !os.SameFile(info, identityInfo) { + if identityStatErr == nil { + identityStatErr = fmt.Errorf("temporary archive identity changed while pinning") + } + return "", nil, nil, nil, errors.Join( + fmt.Errorf("inspect temporary archive identity handle: %w", identityStatErr), + sanitizeAndRemoveOwnedEntry( + parent, + name, + info, + file, + "temporary archive", + operations, + ), + wrapCreateCloseError("temporary archive", file.Close()), + wrapCreateCloseError("temporary archive identity handle", identityHandle.Close()), + ) + } + if err := file.Chmod(0o600); err != nil { + return "", nil, nil, nil, errors.Join( + fmt.Errorf("set temporary archive mode: %w", err), + sanitizeAndRemoveOwnedEntry( + parent, + name, + info, + identityHandle, + "temporary archive", + operations, + ), + wrapCreateCloseError("temporary archive", file.Close()), + wrapCreateCloseError("temporary archive identity handle", identityHandle.Close()), + ) + } + return name, file, identityHandle, info, nil + } + if !errors.Is(err, os.ErrExist) { + return "", nil, nil, nil, fmt.Errorf("create temporary archive: %w", err) + } + } + return "", nil, nil, nil, fmt.Errorf("create unique temporary archive: name attempts exhausted") +} + +func wrapCreateCloseError(name string, err error) error { + if err == nil { + return nil + } + return fmt.Errorf("close %s: %w", name, err) +} + +func wrapCreateSyncError(name string, err error) error { + if err == nil { + return nil + } + return fmt.Errorf("sync %s: %w", name, err) +} + +func wrapCreateStatError(name string, err error) error { + if err == nil { + return nil + } + return fmt.Errorf("inspect %s: %w", name, err) +} diff --git a/ret/archive/pack_test.go b/ret/archive/pack_test.go new file mode 100644 index 00000000..dc226603 --- /dev/null +++ b/ret/archive/pack_test.go @@ -0,0 +1,377 @@ +package archive + +import ( + "archive/tar" + "bytes" + "context" + "errors" + "io" + "os" + "path/filepath" + "runtime" + "testing" + + "github.com/specterops/dawgs/ret/observe" + "github.com/stretchr/testify/require" +) + +func TestCreatePublishesAuthenticatedArchiveExclusively(t *testing.T) { + // Break caught: publishing an incomplete envelope, using a predictable + // temporary name, or replacing a destination that appeared before Create. + root := writeArchiveTestCollection(t, true, true) + recipient, identity, err := GenerateKeyPair() + require.NoError(t, err) + destinationParent := t.TempDir() + archivePath := filepath.Join(destinationParent, "collection.ret") + collision := filepath.Join(destinationParent, ".collection.ret.create-collision.tmp") + require.NoError(t, os.WriteFile(collision, []byte("preserve"), 0o600)) + + require.NoError(t, Create(context.Background(), CreateConfig{ + CollectionDirectory: root, + ArchivePath: archivePath, + Recipient: recipient, + })) + + info, err := os.Stat(archivePath) + require.NoError(t, err) + require.Equal(t, os.FileMode(0o600), info.Mode().Perm()) + require.Equal(t, []byte("preserve"), mustReadArchiveTestFile(t, collision)) + require.Equal(t, []string{collision}, archiveCreateTemporaryPaths(t, archivePath)) + + file, err := os.Open(archivePath) + require.NoError(t, err) + decrypted, err := newDecryptReader(file, identity) + require.NoError(t, err) + contents, err := io.ReadAll(decrypted) + require.NoError(t, err) + require.NoError(t, decrypted.Close()) + require.NoError(t, file.Close()) + + reader := tar.NewReader(bytes.NewReader(contents)) + var names []string + for { + header, err := reader.Next() + if errors.Is(err, io.EOF) { + break + } + require.NoError(t, err) + names = append(names, header.Name) + source := filepath.Join(root, filepath.FromSlash(header.Name)) + require.Equal(t, mustReadArchiveTestFile(t, source), mustReadAllArchiveTest(t, reader)) + } + require.Equal(t, []string{ + "graphs/example/nodes/000001.jsonl", + "graphs/example/nodes/000001.parquet", + "graphs/example/relationships/000001.jsonl", + "graphs/example/relationships/000001.parquet", + "manifest.json", + }, names) + + before := mustReadArchiveTestFile(t, archivePath) + err = Create(context.Background(), CreateConfig{ + CollectionDirectory: root, + ArchivePath: archivePath, + Recipient: recipient, + }) + require.ErrorContains(t, err, "exists") + require.Equal(t, before, mustReadArchiveTestFile(t, archivePath)) + require.Equal(t, []string{collision}, archiveCreateTemporaryPaths(t, archivePath)) +} + +func TestCreateRejectsArchiveLexicallyOrPhysicallyInsideCollection(t *testing.T) { + recipient, _, err := GenerateKeyPair() + require.NoError(t, err) + + t.Run("lexical", func(t *testing.T) { + root := writeArchiveTestCollection(t, true, false) + archivePath := filepath.Join(root, "archive.ret") + + err := Create(context.Background(), CreateConfig{ + CollectionDirectory: root, + ArchivePath: archivePath, + Recipient: recipient, + }) + + require.ErrorContains(t, err, "inside") + require.NoFileExists(t, archivePath) + require.Empty(t, archiveCreateTemporaryPaths(t, archivePath)) + }) + + t.Run("physical through symlinked parent", func(t *testing.T) { + root := writeArchiveTestCollection(t, true, false) + physicalParent := filepath.Join(root, "physical-output") + require.NoError(t, os.Mkdir(physicalParent, 0o700)) + linkParent := filepath.Join(t.TempDir(), "linked-output") + if err := os.Symlink(physicalParent, linkParent); err != nil { + t.Skipf("symlinks unavailable: %v", err) + } + archivePath := filepath.Join(linkParent, "archive.ret") + + err := Create(context.Background(), CreateConfig{ + CollectionDirectory: root, + ArchivePath: archivePath, + Recipient: recipient, + }) + + require.ErrorContains(t, err, "inside") + require.NoFileExists(t, filepath.Join(physicalParent, "archive.ret")) + require.Empty(t, archiveCreateTemporaryPaths(t, archivePath)) + }) +} + +func TestCreateFailureLeavesNoArchiveOrTemporaryFile(t *testing.T) { + recipient, _, err := GenerateKeyPair() + require.NoError(t, err) + + tests := []struct { + name string + prepare func(*testing.T, string) context.Context + }{ + { + name: "full collection verification failure", + prepare: func(t *testing.T, root string) context.Context { + manifest := readArchiveTestManifest(t, root) + parquetPath := manifest.Graphs[0].NodeShards[0].Parquet.Path + require.NoError(t, os.WriteFile( + filepath.Join(root, filepath.FromSlash(parquetPath)), + []byte("corrupt"), + 0o600, + )) + return context.Background() + }, + }, + { + name: "context cancelled before start", + prepare: func(_ *testing.T, _ string) context.Context { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + return ctx + }, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + root := writeArchiveTestCollection(t, true, true) + ctx := test.prepare(t, root) + archivePath := filepath.Join(t.TempDir(), "collection.ret") + + err := Create(ctx, CreateConfig{ + CollectionDirectory: root, + ArchivePath: archivePath, + Recipient: recipient, + }) + + require.Error(t, err) + require.NoFileExists(t, archivePath) + require.Empty(t, archiveCreateTemporaryPaths(t, archivePath)) + }) + } +} + +func TestCreateObserverCancellationSanitizesTemporaryBeforeQuarantine(t *testing.T) { + root := writeArchiveTestCollection(t, true, true) + recipient, _, err := GenerateKeyPair() + require.NoError(t, err) + archiveParent := t.TempDir() + archivePath := filepath.Join(archiveParent, "collection.ret") + ctx, cancel := context.WithCancel(context.Background()) + observer := observe.ObserverFunc(func(_ context.Context, event observe.Event) { + if _, ok := event.(observe.ArchiveEntryProcessed); ok { + cancel() + } + }) + + err = Create(ctx, CreateConfig{ + CollectionDirectory: root, + ArchivePath: archivePath, + Recipient: recipient, + Observer: observer, + }) + + require.ErrorIs(t, err, context.Canceled) + require.ErrorContains(t, err, "ownership cleanup") + require.NoFileExists(t, archivePath) + require.Empty(t, archiveCreateTemporaryPaths(t, archivePath)) + quarantines := archiveCleanupQuarantinePaths(t, archiveParent) + require.Len(t, quarantines, 1) + info, statErr := os.Stat(quarantines[0]) + require.NoError(t, statErr) + require.True(t, info.Mode().IsRegular()) + require.Zero(t, info.Size()) +} + +func TestCreatePreservesTemporaryWhenSanitizationFails(t *testing.T) { + // Break caught: quarantining a partially written archive when retained- + // handle truncate or sync reports failure. + tests := []struct { + name string + operations archiveOperations + wantPayload bool + match string + }{ + { + name: "truncate", + operations: archiveOperations{ + truncateOwnedFile: func(_ *os.File, _ int64) error { + return errors.New("injected truncate failure") + }, + }, + wantPayload: true, + match: "injected truncate failure", + }, + { + name: "sync", + operations: archiveOperations{ + syncOwnedFile: func(_ *os.File) error { + return errors.New("injected sync failure") + }, + }, + match: "injected sync failure", + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + root := writeArchiveTestCollection(t, true, true) + recipient, _, err := GenerateKeyPair() + require.NoError(t, err) + archiveParent := t.TempDir() + archivePath := filepath.Join(archiveParent, "collection.ret") + ctx, cancel := context.WithCancel(context.Background()) + observer := observe.ObserverFunc(func(_ context.Context, event observe.Event) { + if _, ok := event.(observe.ArchiveEntryProcessed); ok { + cancel() + } + }) + + err = create(ctx, CreateConfig{ + CollectionDirectory: root, + ArchivePath: archivePath, + Recipient: recipient, + Observer: observer, + }, runtime.GOOS, test.operations) + + require.ErrorIs(t, err, context.Canceled) + require.ErrorContains(t, err, test.match) + require.ErrorContains(t, err, "ownership cleanup") + require.NoFileExists(t, archivePath) + temporary := archiveCreateTemporaryPaths(t, archivePath) + require.Len(t, temporary, 1) + payload := mustReadArchiveTestFile(t, temporary[0]) + if test.wantPayload { + require.NotEmpty(t, payload) + } else { + require.Empty(t, payload) + } + require.Empty(t, archiveCleanupQuarantinePaths(t, archiveParent)) + }) + } +} + +func TestCreatePreservesDestinationCreatedWhileStreaming(t *testing.T) { + // Break caught: replacing a destination created after the initial + // exclusivity check but before archive publication. + root := writeArchiveTestCollection(t, true, true) + recipient, _, err := GenerateKeyPair() + require.NoError(t, err) + archivePath := filepath.Join(t.TempDir(), "collection.ret") + created := false + observer := observe.ObserverFunc(func(_ context.Context, event observe.Event) { + if _, ok := event.(observe.ArchiveEntryProcessed); ok && !created { + created = true + require.NoError(t, os.WriteFile(archivePath, []byte("concurrent owner"), 0o600)) + } + }) + + err = Create(context.Background(), CreateConfig{ + CollectionDirectory: root, + ArchivePath: archivePath, + Recipient: recipient, + Observer: observer, + }) + + require.Error(t, err) + require.True(t, created) + require.Equal(t, []byte("concurrent owner"), mustReadArchiveTestFile(t, archivePath)) + require.Empty(t, archiveCreateTemporaryPaths(t, archivePath)) +} + +func TestCreateRejectsTemporaryFileReplacedWhileStreaming(t *testing.T) { + // Break caught: publishing a different file substituted at the temporary + // pathname while Create still holds the original temporary file open. + root := writeArchiveTestCollection(t, true, true) + recipient, _, err := GenerateKeyPair() + require.NoError(t, err) + archivePath := filepath.Join(t.TempDir(), "collection.ret") + replaced := false + observer := observe.ObserverFunc(func(_ context.Context, event observe.Event) { + entry, ok := event.(observe.ArchiveEntryProcessed) + if !ok || replaced || entry.Path != "manifest.json" { + return + } + temporary := archiveCreateTemporaryPaths(t, archivePath) + require.Len(t, temporary, 1) + require.NoError(t, os.Remove(temporary[0])) + require.NoError(t, os.WriteFile(temporary[0], []byte("substituted"), 0o600)) + replaced = true + }) + + err = Create(context.Background(), CreateConfig{ + CollectionDirectory: root, + ArchivePath: archivePath, + Recipient: recipient, + Observer: observer, + }) + + require.ErrorContains(t, err, "temporary archive changed") + require.ErrorContains(t, err, "ownership cleanup") + require.True(t, replaced) + require.NoFileExists(t, archivePath) + temporary := archiveCreateTemporaryPaths(t, archivePath) + require.Len(t, temporary, 1) + require.Equal(t, []byte("substituted"), mustReadArchiveTestFile(t, temporary[0])) +} + +func TestRequireDirectoryPathIdentityRejectsReplacedPath(t *testing.T) { + parent := filepath.Join(t.TempDir(), "parent") + require.NoError(t, os.Mkdir(parent, 0o700)) + pinned, err := os.Open(parent) + require.NoError(t, err) + t.Cleanup(func() { + require.NoError(t, pinned.Close()) + }) + moved := parent + ".moved" + require.NoError(t, os.Rename(parent, moved)) + require.NoError(t, os.Mkdir(parent, 0o700)) + + require.ErrorContains( + t, + requireDirectoryPathIdentity(parent, pinned), + "changed", + ) +} + +func archiveCreateTemporaryPaths(t *testing.T, archivePath string) []string { + t.Helper() + matches, err := filepath.Glob(filepath.Join( + filepath.Dir(archivePath), + "."+filepath.Base(archivePath)+".create-*.tmp", + )) + require.NoError(t, err) + return matches +} + +func mustReadArchiveTestFile(t *testing.T, path string) []byte { + t.Helper() + value, err := os.ReadFile(path) + require.NoError(t, err) + return value +} + +func mustReadAllArchiveTest(t *testing.T, reader io.Reader) []byte { + t.Helper() + value, err := io.ReadAll(reader) + require.NoError(t, err) + return value +} diff --git a/ret/archive/platform.go b/ret/archive/platform.go new file mode 100644 index 00000000..f2d049b4 --- /dev/null +++ b/ret/archive/platform.go @@ -0,0 +1,15 @@ +package archive + +import "fmt" + +func requirePlatformSupport(platform string) error { + switch platform { + case "linux", "darwin": + return nil + default: + return fmt.Errorf( + "archive publication is unsupported on platform %q; supported platforms are linux and darwin", + platform, + ) + } +} diff --git a/ret/archive/platform_test.go b/ret/archive/platform_test.go new file mode 100644 index 00000000..d9c6f1cc --- /dev/null +++ b/ret/archive/platform_test.go @@ -0,0 +1,46 @@ +package archive + +import ( + "context" + "path/filepath" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestArchivePlatformSupportIsLimitedToHandleRelativePublication(t *testing.T) { + require.NoError(t, requirePlatformSupport("linux")) + require.NoError(t, requirePlatformSupport("darwin")) + for _, platform := range []string{"windows", "freebsd", "plan9"} { + t.Run(platform, func(t *testing.T) { + require.ErrorContains(t, requirePlatformSupport(platform), "unsupported") + require.ErrorContains(t, requirePlatformSupport(platform), platform) + }) + } +} + +func TestCreateAndExtractGateUnsupportedPlatformsBeforeFilesystemWork(t *testing.T) { + recipient, identity, err := GenerateKeyPair() + require.NoError(t, err) + parent := t.TempDir() + archivePath := filepath.Join(parent, "collection.ret") + output := filepath.Join(parent, "collection") + + err = create(context.Background(), CreateConfig{ + CollectionDirectory: filepath.Join(parent, "missing-collection"), + ArchivePath: archivePath, + Recipient: recipient, + }, "windows", archiveOperations{}) + require.ErrorContains(t, err, "unsupported") + require.NoFileExists(t, archivePath) + require.Empty(t, archiveCreateTemporaryPaths(t, archivePath)) + + err = extract(context.Background(), ExtractConfig{ + ArchivePath: filepath.Join(parent, "missing-archive.ret"), + OutputDirectory: output, + Identity: identity, + }, "windows", archiveOperations{}) + require.ErrorContains(t, err, "unsupported") + require.NoDirExists(t, output) + require.Empty(t, archiveExtractStagePaths(t, output)) +} diff --git a/ret/archive/platform_unsupported_test.go b/ret/archive/platform_unsupported_test.go new file mode 100644 index 00000000..f361e6f7 --- /dev/null +++ b/ret/archive/platform_unsupported_test.go @@ -0,0 +1,71 @@ +//go:build !linux && !darwin + +package archive + +import ( + "context" + "os" + "path/filepath" + "runtime" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestArchiveWorkflowsRejectUnsupportedRuntimeBeforeFilesystemWork(t *testing.T) { + recipient, identity, err := GenerateKeyPair() + require.NoError(t, err) + parent := t.TempDir() + archivePath := filepath.Join(parent, "collection.ret") + output := filepath.Join(parent, "collection") + + err = Create(context.Background(), CreateConfig{ + CollectionDirectory: filepath.Join(parent, "missing-collection"), + ArchivePath: archivePath, + Recipient: recipient, + }) + require.ErrorContains(t, err, "unsupported") + require.ErrorContains(t, err, runtime.GOOS) + require.NoFileExists(t, archivePath) + + err = Extract(context.Background(), ExtractConfig{ + ArchivePath: filepath.Join(parent, "missing-archive.ret"), + OutputDirectory: output, + Identity: identity, + }) + require.ErrorContains(t, err, "unsupported") + require.ErrorContains(t, err, runtime.GOOS) + require.NoDirExists(t, output) +} + +func TestOwnedCleanupOnUnsupportedRuntimePreservesPath(t *testing.T) { + parentPath := t.TempDir() + parent, err := os.OpenRoot(parentPath) + require.NoError(t, err) + t.Cleanup(func() { + require.NoError(t, parent.Close()) + }) + const name = "owned.tmp" + owned, err := parent.OpenFile(name, os.O_RDWR|os.O_CREATE|os.O_EXCL, 0o600) + require.NoError(t, err) + t.Cleanup(func() { + require.NoError(t, owned.Close()) + }) + _, err = owned.Write([]byte("preserve")) + require.NoError(t, err) + expected, err := owned.Stat() + require.NoError(t, err) + + err = removeOwnedEntry( + parent, + name, + expected, + owned.Stat, + "unsupported cleanup test", + archiveOperations{}, + ) + + require.ErrorContains(t, err, "unsupported") + require.Equal(t, []byte("preserve"), mustReadArchiveTestFile(t, filepath.Join(parentPath, name))) + require.Empty(t, archiveCleanupQuarantinePaths(t, parentPath)) +} diff --git a/ret/archive/rename_noreplace_darwin.go b/ret/archive/rename_noreplace_darwin.go new file mode 100644 index 00000000..386f88c0 --- /dev/null +++ b/ret/archive/rename_noreplace_darwin.go @@ -0,0 +1,19 @@ +//go:build darwin + +package archive + +import ( + "os" + + "golang.org/x/sys/unix" +) + +func renameNoReplace(directory *os.File, _ string, oldName, newName string) error { + return unix.RenameatxNp( + int(directory.Fd()), + oldName, + int(directory.Fd()), + newName, + unix.RENAME_EXCL, + ) +} diff --git a/ret/archive/rename_noreplace_linux.go b/ret/archive/rename_noreplace_linux.go new file mode 100644 index 00000000..792ac220 --- /dev/null +++ b/ret/archive/rename_noreplace_linux.go @@ -0,0 +1,19 @@ +//go:build linux + +package archive + +import ( + "os" + + "golang.org/x/sys/unix" +) + +func renameNoReplace(directory *os.File, _ string, oldName, newName string) error { + return unix.Renameat2( + int(directory.Fd()), + oldName, + int(directory.Fd()), + newName, + unix.RENAME_NOREPLACE, + ) +} diff --git a/ret/archive/rename_noreplace_other.go b/ret/archive/rename_noreplace_other.go new file mode 100644 index 00000000..569a43f3 --- /dev/null +++ b/ret/archive/rename_noreplace_other.go @@ -0,0 +1,12 @@ +//go:build !linux && !darwin && !windows + +package archive + +import ( + "os" + "runtime" +) + +func renameNoReplace(_ *os.File, _ string, _, _ string) error { + return requirePlatformSupport(runtime.GOOS) +} diff --git a/ret/archive/rename_noreplace_windows.go b/ret/archive/rename_noreplace_windows.go new file mode 100644 index 00000000..136ee99b --- /dev/null +++ b/ret/archive/rename_noreplace_windows.go @@ -0,0 +1,12 @@ +//go:build windows + +package archive + +import ( + "os" + "runtime" +) + +func renameNoReplace(_ *os.File, _ string, _, _ string) error { + return requirePlatformSupport(runtime.GOOS) +} diff --git a/ret/archive/tar.go b/ret/archive/tar.go new file mode 100644 index 00000000..8a4595b4 --- /dev/null +++ b/ret/archive/tar.go @@ -0,0 +1,514 @@ +package archive + +import ( + "archive/tar" + "context" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "io" + "io/fs" + "os" + "path" + "path/filepath" + "reflect" + "sort" + "strings" + "time" + + "github.com/specterops/dawgs/ret/collection" + "github.com/specterops/dawgs/ret/observe" +) + +const collectionTarFileMode int64 = 0o600 + +type collectionTarFile struct { + path string + info fs.FileInfo +} + +type collectionTarPlan struct { + rootInfo fs.FileInfo + files []collectionTarFile + manifest collection.Manifest + artifactSHA256s map[string]string +} + +func collectionTarPaths( + ctx context.Context, + root string, + observer observe.Observer, +) (collectionTarPlan, error) { + if err := ctx.Err(); err != nil { + return collectionTarPlan{}, fmt.Errorf("prepare collection TAR: %w", err) + } + verification, err := collection.Verify(ctx, root, observer) + if err != nil { + return collectionTarPlan{}, fmt.Errorf("verify collection before TAR creation: %w", err) + } + declared, err := manifestTarPaths(verification.Manifest) + if err != nil { + return collectionTarPlan{}, err + } + plan, err := inventoryCollectionTarFiles(root, declared) + if err != nil { + return collectionTarPlan{}, err + } + plan.manifest = verification.Manifest + plan.artifactSHA256s = manifestArtifactSHA256s(verification.Manifest) + return plan, nil +} + +func manifestArtifactSHA256s(manifest collection.Manifest) map[string]string { + result := make(map[string]string) + for _, graph := range manifest.Graphs { + for _, shard := range graph.NodeShards { + if shard.JSONL != nil { + result[shard.JSONL.Path] = shard.JSONL.SHA256 + } + if shard.Parquet != nil { + result[shard.Parquet.Path] = shard.Parquet.SHA256 + } + } + for _, shard := range graph.RelationshipShards { + if shard.JSONL != nil { + result[shard.JSONL.Path] = shard.JSONL.SHA256 + } + if shard.Parquet != nil { + result[shard.Parquet.Path] = shard.Parquet.SHA256 + } + } + } + return result +} + +func manifestTarPaths(manifest collection.Manifest) ([]string, error) { + declared := []string{collection.ManifestName} + for _, graph := range manifest.Graphs { + for _, shard := range graph.NodeShards { + if shard.JSONL != nil { + declared = append(declared, shard.JSONL.Path) + } + if shard.Parquet != nil { + declared = append(declared, shard.Parquet.Path) + } + } + for _, shard := range graph.RelationshipShards { + if shard.JSONL != nil { + declared = append(declared, shard.JSONL.Path) + } + if shard.Parquet != nil { + declared = append(declared, shard.Parquet.Path) + } + } + } + + seen := make(map[string]struct{}, len(declared)) + for _, relative := range declared { + if _, err := collection.SafeJoin(".", relative); err != nil { + return nil, fmt.Errorf("validate declared collection path %q: %w", relative, err) + } + if _, found := seen[relative]; found { + return nil, fmt.Errorf("duplicate declared collection path %q", relative) + } + seen[relative] = struct{}{} + } + sort.Strings(declared) + return declared, nil +} + +func inventoryCollectionTarFiles(root string, declared []string) (collectionTarPlan, error) { + rootInfo, err := os.Lstat(root) + if err != nil { + return collectionTarPlan{}, fmt.Errorf("inspect collection root: %w", err) + } + if rootInfo.Mode()&os.ModeSymlink != 0 || !rootInfo.IsDir() { + return collectionTarPlan{}, fmt.Errorf("collection root is not a non-symlink directory: %q", root) + } + + declaredFiles := make(map[string]struct{}, len(declared)) + declaredDirectories := map[string]struct{}{".": {}} + for _, relative := range declared { + declaredFiles[relative] = struct{}{} + for directory := path.Dir(relative); directory != "."; directory = path.Dir(directory) { + declaredDirectories[directory] = struct{}{} + } + } + + found := make(map[string]fs.FileInfo, len(declared)) + err = filepath.WalkDir(root, func(candidate string, entry fs.DirEntry, walkErr error) error { + if walkErr != nil { + return walkErr + } + relative, err := filepath.Rel(root, candidate) + if err != nil { + return fmt.Errorf("resolve collection entry %q: %w", candidate, err) + } + relative = filepath.ToSlash(relative) + info, err := entry.Info() + if err != nil { + return fmt.Errorf("inspect collection entry %q: %w", relative, err) + } + if info.Mode()&os.ModeSymlink != 0 { + return fmt.Errorf("collection entry %q is a symlink", relative) + } + if info.IsDir() { + if _, ok := declaredDirectories[relative]; !ok { + return fmt.Errorf("unexpected collection directory %q", relative) + } + return nil + } + if !info.Mode().IsRegular() { + return fmt.Errorf("collection entry %q is not a regular file", relative) + } + if _, ok := declaredFiles[relative]; !ok { + return fmt.Errorf("unexpected collection file %q", relative) + } + found[relative] = info + return nil + }) + if err != nil { + return collectionTarPlan{}, fmt.Errorf("inventory collection TAR inputs: %w", err) + } + + files := make([]collectionTarFile, 0, len(declared)) + for _, relative := range declared { + info, ok := found[relative] + if !ok { + return collectionTarPlan{}, fmt.Errorf("declared collection file %q is missing", relative) + } + files = append(files, collectionTarFile{path: relative, info: info}) + } + return collectionTarPlan{rootInfo: rootInfo, files: files}, nil +} + +func inventoryPinnedCollectionTarFiles(root *os.Root, declared []string) (collectionTarPlan, error) { + rootInfo, err := root.Stat(".") + if err != nil { + return collectionTarPlan{}, fmt.Errorf("inspect pinned collection root: %w", err) + } + if !rootInfo.IsDir() { + return collectionTarPlan{}, fmt.Errorf("pinned collection root is not a directory") + } + + declaredFiles := make(map[string]struct{}, len(declared)) + declaredDirectories := map[string]struct{}{".": {}} + for _, relative := range declared { + declaredFiles[relative] = struct{}{} + for directory := path.Dir(relative); directory != "."; directory = path.Dir(directory) { + declaredDirectories[directory] = struct{}{} + } + } + + found := make(map[string]fs.FileInfo, len(declared)) + err = fs.WalkDir(root.FS(), ".", func(relative string, entry fs.DirEntry, walkErr error) error { + if walkErr != nil { + return walkErr + } + info, err := entry.Info() + if err != nil { + return fmt.Errorf("inspect pinned collection entry %q: %w", relative, err) + } + if info.Mode()&os.ModeSymlink != 0 { + return fmt.Errorf("pinned collection entry %q is a symlink", relative) + } + if info.IsDir() { + if _, ok := declaredDirectories[relative]; !ok { + return fmt.Errorf("unexpected pinned collection directory %q", relative) + } + return nil + } + if !info.Mode().IsRegular() { + return fmt.Errorf("pinned collection entry %q is not a regular file", relative) + } + if _, ok := declaredFiles[relative]; !ok { + return fmt.Errorf("unexpected pinned collection file %q", relative) + } + found[relative] = info + return nil + }) + if err != nil { + return collectionTarPlan{}, fmt.Errorf("inventory pinned collection TAR inputs: %w", err) + } + + files := make([]collectionTarFile, 0, len(declared)) + for _, relative := range declared { + info, ok := found[relative] + if !ok { + return collectionTarPlan{}, fmt.Errorf("declared pinned collection file %q is missing", relative) + } + files = append(files, collectionTarFile{path: relative, info: info}) + } + return collectionTarPlan{rootInfo: rootInfo, files: files}, nil +} + +func readPinnedCollectionManifest(root *os.Root) ( + manifest collection.Manifest, + digest string, + info fs.FileInfo, + resultErr error, +) { + if err := requirePinnedRegularPath(root, collection.ManifestName); err != nil { + return manifest, "", nil, err + } + info, err := root.Lstat(collection.ManifestName) + if err != nil { + return manifest, "", nil, fmt.Errorf("inspect pinned collection manifest: %w", err) + } + file, err := root.Open(collection.ManifestName) + if err != nil { + return manifest, "", nil, fmt.Errorf("open pinned collection manifest: %w", err) + } + defer func() { + resultErr = errors.Join(resultErr, wrapTarCloseError("pinned collection manifest", file.Close())) + }() + openedInfo, err := file.Stat() + if err != nil { + return manifest, "", nil, fmt.Errorf("inspect open pinned collection manifest: %w", err) + } + if !openedInfo.Mode().IsRegular() || !os.SameFile(info, openedInfo) { + return manifest, "", nil, fmt.Errorf("pinned collection manifest changed while opening") + } + + hasher := sha256.New() + decoder := json.NewDecoder(io.TeeReader(file, hasher)) + decoder.DisallowUnknownFields() + if err := decoder.Decode(&manifest); err != nil { + return manifest, "", nil, fmt.Errorf("decode pinned collection manifest: %w", err) + } + var trailing any + if err := decoder.Decode(&trailing); !errors.Is(err, io.EOF) { + if err == nil { + return manifest, "", nil, fmt.Errorf("decode pinned collection manifest: trailing JSON value") + } + return manifest, "", nil, fmt.Errorf("decode pinned collection manifest trailing data: %w", err) + } + if err := manifest.Validate(); err != nil { + return manifest, "", nil, fmt.Errorf("validate pinned collection manifest: %w", err) + } + return manifest, hex.EncodeToString(hasher.Sum(nil)), info, nil +} + +func plannedTarFile(files []collectionTarFile, relative string) (collectionTarFile, bool) { + for _, file := range files { + if file.path == relative { + return file, true + } + } + return collectionTarFile{}, false +} + +func tarPlanPaths(files []collectionTarFile) []string { + paths := make([]string, len(files)) + for index, file := range files { + paths[index] = file.path + } + return paths +} + +func compareCollectionTarPlans(expected, actual collectionTarPlan) error { + if !os.SameFile(expected.rootInfo, actual.rootInfo) { + return fmt.Errorf("collection root changed during TAR creation") + } + if len(expected.files) != len(actual.files) { + return fmt.Errorf("collection file set changed during TAR creation") + } + for index := range expected.files { + if expected.files[index].path != actual.files[index].path || + !samePlannedFile(expected.files[index].info, actual.files[index].info) { + return fmt.Errorf("collection file %q changed during TAR creation", expected.files[index].path) + } + } + return nil +} + +func writeCollectionTar( + ctx context.Context, + destination io.Writer, + root string, + plan collectionTarPlan, + observer observe.Observer, +) (resultErr error) { + if destination == nil { + return fmt.Errorf("collection TAR destination is required") + } + if err := ctx.Err(); err != nil { + return fmt.Errorf("write collection TAR: %w", err) + } + + pinnedRoot, err := os.OpenRoot(root) + if err != nil { + return fmt.Errorf("open collection root: %w", err) + } + defer func() { + resultErr = errors.Join(resultErr, wrapTarCloseError("collection root", pinnedRoot.Close())) + }() + pinnedRootInfo, err := pinnedRoot.Stat(".") + if err != nil { + return fmt.Errorf("inspect open collection root: %w", err) + } + if !pinnedRootInfo.IsDir() || !os.SameFile(plan.rootInfo, pinnedRootInfo) { + return fmt.Errorf("collection root changed before TAR creation") + } + pinnedManifest, manifestDigest, manifestInfo, err := readPinnedCollectionManifest(pinnedRoot) + if err != nil { + return err + } + if !reflect.DeepEqual(plan.manifest, pinnedManifest) { + return fmt.Errorf("collection manifest changed after verification") + } + plannedManifest, ok := plannedTarFile(plan.files, collection.ManifestName) + if !ok || !samePlannedFile(plannedManifest.info, manifestInfo) { + return fmt.Errorf("collection manifest changed after inventory") + } + + writer := tar.NewWriter(destination) + defer func() { + resultErr = errors.Join(resultErr, wrapTarCloseError("TAR writer", writer.Close())) + }() + + for _, planned := range plan.files { + if err := ctx.Err(); err != nil { + return fmt.Errorf("write collection TAR: %w", err) + } + if err := requirePinnedRegularPath(pinnedRoot, planned.path); err != nil { + return err + } + info, err := pinnedRoot.Lstat(filepath.FromSlash(planned.path)) + if err != nil { + return fmt.Errorf("inspect TAR file %q: %w", planned.path, err) + } + if !samePlannedFile(planned.info, info) { + return fmt.Errorf("collection file %q changed before TAR creation", planned.path) + } + + file, err := pinnedRoot.Open(filepath.FromSlash(planned.path)) + if err != nil { + return fmt.Errorf("open TAR file %q: %w", planned.path, err) + } + expectedDigest := plan.artifactSHA256s[planned.path] + if planned.path == collection.ManifestName { + expectedDigest = manifestDigest + } + if err := writeCollectionTarFile(writer, file, planned, expectedDigest); err != nil { + return errors.Join(err, wrapTarCloseError("TAR file "+planned.path, file.Close())) + } + if err := file.Close(); err != nil { + return fmt.Errorf("close TAR file %q: %w", planned.path, err) + } + + observe.Emit(ctx, observer, observe.ArchiveEntryProcessed{ + Operation: "pack", + Path: planned.path, + Size: planned.info.Size(), + }) + if err := ctx.Err(); err != nil { + return fmt.Errorf("write collection TAR after %q: %w", planned.path, err) + } + } + finalPlan, err := inventoryPinnedCollectionTarFiles(pinnedRoot, tarPlanPaths(plan.files)) + if err != nil { + return fmt.Errorf("final collection TAR inventory: %w", err) + } + if err := compareCollectionTarPlans(plan, finalPlan); err != nil { + return err + } + return nil +} + +func requirePinnedRegularPath(root *os.Root, relative string) error { + current := "" + components := strings.Split(filepath.FromSlash(relative), string(filepath.Separator)) + for index, component := range components { + current = filepath.Join(current, component) + info, err := root.Lstat(current) + if err != nil { + return fmt.Errorf("inspect TAR path component %q: %w", component, err) + } + if info.Mode()&os.ModeSymlink != 0 { + return fmt.Errorf("TAR path %q contains a symlink component %q", relative, component) + } + if index < len(components)-1 { + if !info.IsDir() { + return fmt.Errorf("TAR path component %q is not a directory", component) + } + continue + } + if !info.Mode().IsRegular() { + return fmt.Errorf("TAR file %q is not a regular file", relative) + } + } + return nil +} + +func writeCollectionTarFile( + writer *tar.Writer, + file *os.File, + planned collectionTarFile, + expectedDigest string, +) error { + openedInfo, err := file.Stat() + if err != nil { + return fmt.Errorf("inspect open TAR file %q: %w", planned.path, err) + } + if !openedInfo.Mode().IsRegular() || !samePlannedFile(planned.info, openedInfo) { + return fmt.Errorf("collection file %q changed while opening", planned.path) + } + + header := &tar.Header{ + Name: planned.path, + Mode: collectionTarFileMode, + Uid: 0, + Gid: 0, + Size: planned.info.Size(), + ModTime: time.Unix(0, 0).UTC(), + AccessTime: time.Time{}, + ChangeTime: time.Time{}, + Typeflag: tar.TypeReg, + Uname: "", + Gname: "", + Format: tar.FormatPAX, + } + if err := writer.WriteHeader(header); err != nil { + return fmt.Errorf("write TAR header %q: %w", planned.path, err) + } + hasher := sha256.New() + copied, err := io.CopyN(io.MultiWriter(writer, hasher), file, planned.info.Size()) + if err != nil { + return fmt.Errorf("write TAR contents %q after %d bytes: %w", planned.path, copied, err) + } + var trailing [1]byte + n, err := file.Read(trailing[:]) + if err != nil && !errors.Is(err, io.EOF) { + return fmt.Errorf("check TAR file size %q: %w", planned.path, err) + } + if n != 0 { + return fmt.Errorf("collection file %q grew during TAR creation", planned.path) + } + actualDigest := hex.EncodeToString(hasher.Sum(nil)) + if expectedDigest == "" || actualDigest != expectedDigest { + return fmt.Errorf( + "collection file %q SHA-256 is %s, want %s", + planned.path, + actualDigest, + expectedDigest, + ) + } + return nil +} + +func samePlannedFile(expected, actual fs.FileInfo) bool { + return os.SameFile(expected, actual) && + expected.Mode() == actual.Mode() && + expected.Size() == actual.Size() && + expected.ModTime().Equal(actual.ModTime()) +} + +func wrapTarCloseError(name string, err error) error { + if err == nil { + return nil + } + return fmt.Errorf("close %s: %w", name, err) +} diff --git a/ret/archive/tar_test.go b/ret/archive/tar_test.go new file mode 100644 index 00000000..99d11d1a --- /dev/null +++ b/ret/archive/tar_test.go @@ -0,0 +1,355 @@ +package archive + +import ( + "archive/tar" + "bytes" + "context" + "io" + "os" + "path/filepath" + "sort" + "strings" + "testing" + "time" + + "github.com/specterops/dawgs/ret/collection" + "github.com/specterops/dawgs/ret/entity" + "github.com/specterops/dawgs/ret/jsonl" + "github.com/specterops/dawgs/ret/metrics" + "github.com/specterops/dawgs/ret/parquet" + "github.com/stretchr/testify/require" +) + +func TestCollectionTarIsDeterministicSortedAndHasFixedMetadata(t *testing.T) { + // Break caught: allowing host directory iteration order or filesystem + // metadata to alter the plaintext TAR bytes. + root := writeArchiveTestCollection(t, true, true) + first := createPlainTestTar(t, root) + + require.NoError(t, filepath.WalkDir(root, func(path string, entry os.DirEntry, err error) error { + require.NoError(t, err) + require.NoError(t, os.Chmod(path, 0o777)) + return os.Chtimes( + path, + time.Date(2038, time.January, 19, 3, 14, 7, 0, time.UTC), + time.Date(2040, time.February, 20, 4, 15, 8, 0, time.UTC), + ) + })) + + second := createPlainTestTar(t, root) + require.Equal(t, first, second) + + reader := tar.NewReader(bytes.NewReader(first)) + var names []string + for { + header, err := reader.Next() + if err == io.EOF { + break + } + require.NoError(t, err) + names = append(names, header.Name) + require.Equal(t, int64(0o600), header.Mode) + require.Zero(t, header.Uid) + require.Zero(t, header.Gid) + require.Empty(t, header.Uname) + require.Empty(t, header.Gname) + require.Equal(t, byte(tar.TypeReg), header.Typeflag) + require.True(t, header.ModTime.Equal(time.Unix(0, 0).UTC())) + require.True(t, header.AccessTime.IsZero()) + require.True(t, header.ChangeTime.IsZero()) + require.Empty(t, header.PAXRecords) + } + + require.Equal(t, []string{ + "graphs/example/nodes/000001.jsonl", + "graphs/example/nodes/000001.parquet", + "graphs/example/relationships/000001.jsonl", + "graphs/example/relationships/000001.parquet", + "manifest.json", + }, names) + require.True(t, sort.StringsAreSorted(names)) +} + +func TestCollectionTarRejectsAnythingOutsideTheDeclaredFileSet(t *testing.T) { + tests := []struct { + name string + mutate func(*testing.T, string) + match string + }{ + { + name: "undeclared regular file", + mutate: func(t *testing.T, root string) { + require.NoError(t, os.WriteFile(filepath.Join(root, "extra.txt"), []byte("extra"), 0o600)) + }, + match: "unexpected", + }, + { + name: "undeclared empty directory", + mutate: func(t *testing.T, root string) { + require.NoError(t, os.Mkdir(filepath.Join(root, "extra"), 0o700)) + }, + match: "unexpected", + }, + { + name: "undeclared symlink", + mutate: func(t *testing.T, root string) { + if err := os.Symlink(collection.ManifestName, filepath.Join(root, "extra")); err != nil { + t.Skipf("symlinks unavailable: %v", err) + } + }, + match: "symlink", + }, + { + name: "declared symlink", + mutate: func(t *testing.T, root string) { + manifest := readArchiveTestManifest(t, root) + path := manifest.Graphs[0].NodeShards[0].JSONL.Path + require.NoError(t, os.Remove(filepath.Join(root, filepath.FromSlash(path)))) + if err := os.Symlink( + filepath.Join(root, collection.ManifestName), + filepath.Join(root, filepath.FromSlash(path)), + ); err != nil { + t.Skipf("symlinks unavailable: %v", err) + } + }, + match: "symlink", + }, + { + name: "missing declared file", + mutate: func(t *testing.T, root string) { + manifest := readArchiveTestManifest(t, root) + path := manifest.Graphs[0].NodeShards[0].JSONL.Path + require.NoError(t, os.Remove(filepath.Join(root, filepath.FromSlash(path)))) + }, + match: "no such file", + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + root := writeArchiveTestCollection(t, true, false) + test.mutate(t, root) + + var destination bytes.Buffer + paths, err := collectionTarPaths(context.Background(), root, nil) + if err == nil { + err = writeCollectionTar(context.Background(), &destination, root, paths, nil) + } + + require.ErrorContains(t, err, test.match) + require.Empty(t, destination.Bytes()) + }) + } +} + +func TestCollectionTarFullyVerifiesBeforeWriting(t *testing.T) { + // Break caught: packaging a dual-format collection after checking only its + // JSONL artifacts or after beginning to emit TAR bytes. + root := writeArchiveTestCollection(t, true, true) + manifest := readArchiveTestManifest(t, root) + parquetPath := manifest.Graphs[0].NodeShards[0].Parquet.Path + require.NoError(t, os.WriteFile( + filepath.Join(root, filepath.FromSlash(parquetPath)), + []byte("corrupt parquet"), + 0o600, + )) + + var destination bytes.Buffer + paths, err := collectionTarPaths(context.Background(), root, nil) + if err == nil { + err = writeCollectionTar(context.Background(), &destination, root, paths, nil) + } + + require.Error(t, err) + require.Empty(t, destination.Bytes()) +} + +func TestCollectionTarRejectsInPlaceMutationAfterVerification(t *testing.T) { + // Break caught: publishing bytes changed in place after full verification + // while preserving the file identity, size, mode, and modification time. + root := writeArchiveTestCollection(t, true, true) + plan, err := collectionTarPaths(context.Background(), root, nil) + require.NoError(t, err) + require.NotEmpty(t, plan.files) + planned := plan.files[0] + require.NotEqual(t, collection.ManifestName, planned.path) + absolute := filepath.Join(root, filepath.FromSlash(planned.path)) + payload := mustReadArchiveTestFile(t, absolute) + require.NotEmpty(t, payload) + payload[0] ^= 0x01 + require.NoError(t, os.WriteFile(absolute, payload, planned.info.Mode().Perm())) + require.NoError(t, os.Chmod(absolute, planned.info.Mode().Perm())) + require.NoError(t, os.Chtimes(absolute, planned.info.ModTime(), planned.info.ModTime())) + + var destination bytes.Buffer + err = writeCollectionTar(context.Background(), &destination, root, plan, nil) + + require.ErrorContains(t, err, "SHA-256") +} + +func TestCollectionTarRoundTripsLongValidPAXPath(t *testing.T) { + // Break caught: forcing USTAR and rejecting a valid collection path longer + // than its name/prefix fields can represent. + graphName := strings.Repeat("g", 240) + root := writeArchiveTestCollectionNamed(t, graphName, true, false) + recipient, identity, err := GenerateKeyPair() + require.NoError(t, err) + archivePath := filepath.Join(t.TempDir(), "collection.ret") + output := filepath.Join(t.TempDir(), "collection") + + require.NoError(t, Create(context.Background(), CreateConfig{ + CollectionDirectory: root, + ArchivePath: archivePath, + Recipient: recipient, + })) + require.NoError(t, Extract(context.Background(), ExtractConfig{ + ArchivePath: archivePath, + OutputDirectory: output, + Identity: identity, + })) + _, err = collection.Verify(context.Background(), output, nil) + require.NoError(t, err) +} + +func createPlainTestTar(t *testing.T, root string) []byte { + t.Helper() + paths, err := collectionTarPaths(context.Background(), root, nil) + require.NoError(t, err) + var destination bytes.Buffer + require.NoError(t, writeCollectionTar(context.Background(), &destination, root, paths, nil)) + return append([]byte(nil), destination.Bytes()...) +} + +func writeArchiveTestCollection(t *testing.T, withJSONL, withParquet bool) string { + t.Helper() + return writeArchiveTestCollectionNamed(t, "example", withJSONL, withParquet) +} + +func writeArchiveTestCollectionNamed( + t *testing.T, + graphName string, + withJSONL, withParquet bool, +) string { + t.Helper() + root := t.TempDir() + nodes := []entity.Node{ + {SourceID: "node-1", Kinds: []string{"User", "Principal"}, Properties: map[string]any{"name": "Alice"}}, + {SourceID: "node-2", Kinds: []string{"Group"}, Properties: map[string]any{"name": "Admins"}}, + } + relationships := []entity.Relationship{{ + SourceID: "relationship-1", + StartID: "node-1", + EndID: "node-2", + Kind: "MEMBER_OF", + Properties: map[string]any{"active": true}, + }} + + builder := metrics.NewBuilder() + for _, node := range nodes { + require.NoError(t, builder.ObserveNode(node)) + } + for _, relationship := range relationships { + require.NoError(t, builder.ObserveRelationship(relationship)) + } + + graph := collection.Graph{ + Name: graphName, + NodeCount: int64(len(nodes)), + RelationshipCount: int64(len(relationships)), + KindCatalog: []string{"User", "Principal", "Group", "MEMBER_OF"}, + NodeShards: []collection.NodeShard{{Index: 1, Count: int64(len(nodes)), LastSourceID: 2}}, + RelationshipShards: []collection.RelationshipShard{{Index: 1, Count: int64(len(relationships)), LastSourceID: 1}}, + Metrics: builder.Finalize(), + } + outputs := collection.OutputConfig{} + + if withJSONL { + config := jsonl.Config{Codec: jsonl.CodecNone} + outputs.JSONL = &collection.JSONLOutput{ + SchemaVersion: jsonl.SchemaVersion, + Codec: string(config.Codec), + Level: config.Level, + } + nodePath := collection.NodeJSONLPath(graph.Name, 1, config.Codec) + nodeTemporary := filepath.Join(root, "nodes.jsonl.tmp") + nodeFile, err := os.Create(nodeTemporary) + require.NoError(t, err) + nodeWriter, err := jsonl.NewNodeWriter(nodeFile, config) + require.NoError(t, err) + require.NoError(t, nodeWriter.Push(nodes)) + require.NoError(t, nodeWriter.Close()) + nodeArtifact, err := nodeWriter.Result() + require.NoError(t, err) + require.NoError(t, nodeFile.Close()) + installArchiveTestArtifact(t, root, nodeTemporary, nodePath) + graph.NodeShards[0].JSONL = &collection.JSONLArtifact{Path: nodePath, Artifact: nodeArtifact} + + relationshipPath := collection.RelationshipJSONLPath(graph.Name, 1, config.Codec) + relationshipTemporary := filepath.Join(root, "relationships.jsonl.tmp") + relationshipFile, err := os.Create(relationshipTemporary) + require.NoError(t, err) + relationshipWriter, err := jsonl.NewRelationshipWriter(relationshipFile, config) + require.NoError(t, err) + require.NoError(t, relationshipWriter.Push(relationships)) + require.NoError(t, relationshipWriter.Close()) + relationshipArtifact, err := relationshipWriter.Result() + require.NoError(t, err) + require.NoError(t, relationshipFile.Close()) + installArchiveTestArtifact(t, root, relationshipTemporary, relationshipPath) + graph.RelationshipShards[0].JSONL = &collection.JSONLArtifact{Path: relationshipPath, Artifact: relationshipArtifact} + } + + if withParquet { + config := parquet.Config{} + outputs.Parquet = &collection.ParquetOutput{SchemaVersion: parquet.SchemaVersion} + nodePath := collection.NodeParquetPath(graph.Name, 1) + nodeTemporary := filepath.Join(root, "nodes.parquet.tmp") + nodeFile, err := os.Create(nodeTemporary) + require.NoError(t, err) + nodeWriter, err := parquet.NewNodeWriter(nodeFile, config) + require.NoError(t, err) + require.NoError(t, nodeWriter.Push(nodes)) + require.NoError(t, nodeWriter.Close()) + nodeArtifact, err := nodeWriter.Result() + require.NoError(t, err) + require.NoError(t, nodeFile.Close()) + installArchiveTestArtifact(t, root, nodeTemporary, nodePath) + graph.NodeShards[0].Parquet = &collection.ParquetArtifact{Path: nodePath, Artifact: nodeArtifact} + + relationshipPath := collection.RelationshipParquetPath(graph.Name, 1) + relationshipTemporary := filepath.Join(root, "relationships.parquet.tmp") + relationshipFile, err := os.Create(relationshipTemporary) + require.NoError(t, err) + relationshipWriter, err := parquet.NewRelationshipWriter(relationshipFile, config) + require.NoError(t, err) + require.NoError(t, relationshipWriter.Push(relationships)) + require.NoError(t, relationshipWriter.Close()) + relationshipArtifact, err := relationshipWriter.Result() + require.NoError(t, err) + require.NoError(t, relationshipFile.Close()) + installArchiveTestArtifact(t, root, relationshipTemporary, relationshipPath) + graph.RelationshipShards[0].Parquet = &collection.ParquetArtifact{Path: relationshipPath, Artifact: relationshipArtifact} + } + + require.NoError(t, collection.Write(root, collection.Manifest{ + Format: collection.Format, + CreatedAt: time.Date(2026, time.July, 28, 12, 0, 0, 0, time.UTC), + Outputs: outputs, + Graphs: []collection.Graph{graph}, + })) + return root +} + +func installArchiveTestArtifact(t *testing.T, root, temporary, relative string) { + t.Helper() + final := filepath.Join(root, filepath.FromSlash(relative)) + require.NoError(t, os.MkdirAll(filepath.Dir(final), 0o700)) + require.NoError(t, os.Rename(temporary, final)) +} + +func readArchiveTestManifest(t *testing.T, root string) collection.Manifest { + t.Helper() + manifest, err := collection.Read(root) + require.NoError(t, err) + return manifest +} diff --git a/ret/archive/unpack.go b/ret/archive/unpack.go new file mode 100644 index 00000000..67bb3ea3 --- /dev/null +++ b/ret/archive/unpack.go @@ -0,0 +1,817 @@ +package archive + +import ( + "archive/tar" + "context" + "crypto/rand" + "crypto/sha256" + "encoding/hex" + "errors" + "fmt" + "io" + "os" + "path" + "path/filepath" + "reflect" + "runtime" + "sort" + "strings" + + "github.com/specterops/dawgs/ret/collection" + "github.com/specterops/dawgs/ret/observe" +) + +type ExtractConfig struct { + ArchivePath string + OutputDirectory string + Identity PrivateKey + Observer observe.Observer +} + +func Extract(ctx context.Context, config ExtractConfig) (resultErr error) { + return extract(ctx, config, runtime.GOOS, archiveOperations{}) +} + +func extract( + ctx context.Context, + config ExtractConfig, + platform string, + operations archiveOperations, +) (resultErr error) { + if err := requirePlatformSupport(platform); err != nil { + return err + } + if err := validateExtractConfig(config); err != nil { + return err + } + if err := ctx.Err(); err != nil { + return fmt.Errorf("extract archive: %w", err) + } + + archiveFile, err := openArchiveForExtraction(config.ArchivePath) + if err != nil { + return err + } + defer func() { + if archiveFile != nil { + resultErr = errors.Join(resultErr, wrapExtractCloseError("archive", archiveFile.Close())) + } + }() + + outputParentPath, outputBase, err := prepareExtractDestination(config.OutputDirectory) + if err != nil { + return err + } + outputParent, err := os.OpenRoot(outputParentPath) + if err != nil { + return fmt.Errorf("open extraction destination parent: %w", err) + } + defer func() { + resultErr = errors.Join( + resultErr, + wrapExtractCloseError("extraction destination parent", outputParent.Close()), + ) + }() + outputParentDirectory, err := openDirectoryMatchingRoot(outputParentPath, outputParent) + if err != nil { + return err + } + defer func() { + resultErr = errors.Join( + resultErr, + wrapExtractCloseError("extraction destination parent directory", outputParentDirectory.Close()), + ) + }() + if _, err := outputParent.Lstat(outputBase); err == nil { + return fmt.Errorf("extraction destination exists: %s", config.OutputDirectory) + } else if !errors.Is(err, os.ErrNotExist) { + return fmt.Errorf("inspect extraction destination: %w", err) + } + + stageName, stageRoot, createdStageInfo, err := createExtractStage( + outputParent, + outputBase, + operations, + ) + if err != nil { + return err + } + published := false + defer func() { + if stageRoot == nil { + return + } + if published { + resultErr = errors.Join( + resultErr, + wrapExtractCloseError("extraction stage", stageRoot.Close()), + ) + } else { + resultErr = errors.Join( + resultErr, + removeOwnedDirectory( + outputParent, + stageName, + stageRoot, + createdStageInfo, + "extraction stage", + operations, + ), + ) + } + stageRoot = nil + }() + stagePath := filepath.Join(outputParentPath, stageName) + if err := requireRootPathIdentity(stagePath, stageRoot, "extraction stage"); err != nil { + return err + } + + decrypted, err := newDecryptReader(archiveFile, config.Identity) + if err != nil { + return fmt.Errorf("open encrypted archive: %w", err) + } + defer func() { + if decrypted != nil { + resultErr = errors.Join( + resultErr, + wrapExtractCloseError("archive decryption reader", decrypted.Close()), + ) + } + }() + + seen, err := extractCollectionTar( + ctx, + decrypted, + stageRoot, + config.Observer, + operations, + ) + if err != nil { + return err + } + if err := requireZeroTarTail(decrypted); err != nil { + return err + } + if err := decrypted.Close(); err != nil { + decrypted = nil + return fmt.Errorf("authenticate archive final frame: %w", err) + } + decrypted = nil + if err := archiveFile.Close(); err != nil { + archiveFile = nil + return fmt.Errorf("close archive: %w", err) + } + archiveFile = nil + + if err := requireRootPathIdentity(stagePath, stageRoot, "extraction stage before verification"); err != nil { + return err + } + manifest, manifestDigest, _, err := readPinnedCollectionManifest(stageRoot) + if err != nil { + return fmt.Errorf("read extracted collection manifest: %w", err) + } + declared, err := manifestTarPaths(manifest) + if err != nil { + return err + } + if err := compareExtractedFileSet(seen, declared); err != nil { + return err + } + initialPlan, err := inventoryPinnedCollectionTarFiles(stageRoot, declared) + if err != nil { + return fmt.Errorf("validate extracted collection file set: %w", err) + } + verification, err := collection.Verify(ctx, stagePath, config.Observer) + if err != nil { + return fmt.Errorf("verify extracted collection: %w", err) + } + if !reflect.DeepEqual(manifest, verification.Manifest) { + return fmt.Errorf("verified collection manifest differs from authenticated stage manifest") + } + if err := ctx.Err(); err != nil { + return fmt.Errorf("promote extracted collection: %w", err) + } + if err := syncExtractedCollection( + ctx, + stageRoot, + declared, + operations, + ); err != nil { + return err + } + if err := ctx.Err(); err != nil { + return fmt.Errorf("promote synced extracted collection: %w", err) + } + if err := verifyPinnedExtractedCollection( + stageRoot, + initialPlan, + manifest, + manifestDigest, + declared, + ); err != nil { + return err + } + if err := requireRootPathIdentity(stagePath, stageRoot, "synced extraction stage"); err != nil { + return err + } + if err := ctx.Err(); err != nil { + return fmt.Errorf("promote bound extracted collection: %w", err) + } + stageInfo, err := stageRoot.Stat(".") + if err != nil { + return fmt.Errorf("inspect synced extraction stage: %w", err) + } + if _, err := outputParent.Lstat(outputBase); err == nil { + return fmt.Errorf("extraction destination exists: %s", config.OutputDirectory) + } else if !errors.Is(err, os.ErrNotExist) { + return fmt.Errorf("inspect extraction destination before promotion: %w", err) + } + if err := ctx.Err(); err != nil { + return fmt.Errorf("promote extracted collection: %w", err) + } + if err := requireDirectoryPathIdentity(outputParentPath, outputParentDirectory); err != nil { + return fmt.Errorf("validate extraction destination parent before promotion: %w", err) + } + if err := requireRootAndPhysicalEntryIdentity( + outputParent, + outputParentPath, + stageName, + "verified extraction stage", + stageInfo, + ); err != nil { + return err + } + if err := renameNoReplace(outputParentDirectory, outputParentPath, stageName, outputBase); err != nil { + return fmt.Errorf("promote extracted collection: %w", err) + } + published = true + if err := requireRootAndPhysicalEntryIdentity( + outputParent, + outputParentPath, + outputBase, + "promoted collection", + stageInfo, + ); err != nil { + return err + } + if err := outputParentDirectory.Sync(); err != nil { + return fmt.Errorf("sync extraction destination parent: %w", err) + } + return nil +} + +func validateExtractConfig(config ExtractConfig) error { + if strings.TrimSpace(config.ArchivePath) == "" { + return fmt.Errorf("archive path is required") + } + if strings.TrimSpace(config.OutputDirectory) == "" { + return fmt.Errorf("output directory is required") + } + if !config.Identity.valid { + return fmt.Errorf("identity private key is required") + } + return nil +} + +func openArchiveForExtraction(archivePath string) (*os.File, error) { + absolute, err := filepath.Abs(archivePath) + if err != nil { + return nil, fmt.Errorf("resolve archive path: %w", err) + } + physicalParent, err := filepath.EvalSymlinks(filepath.Dir(absolute)) + if err != nil { + return nil, fmt.Errorf("resolve archive parent: %w", err) + } + parent, err := os.OpenRoot(physicalParent) + if err != nil { + return nil, fmt.Errorf("open archive parent: %w", err) + } + base := filepath.Base(absolute) + info, err := parent.Lstat(base) + if err != nil { + _ = parent.Close() + return nil, fmt.Errorf("inspect archive: %w", err) + } + if info.Mode()&os.ModeSymlink != 0 || !info.Mode().IsRegular() { + _ = parent.Close() + return nil, fmt.Errorf("archive is not a non-symlink regular file") + } + file, err := parent.Open(base) + if err != nil { + _ = parent.Close() + return nil, fmt.Errorf("open archive: %w", err) + } + openedInfo, statErr := file.Stat() + closeParentErr := parent.Close() + if statErr != nil { + return nil, errors.Join( + fmt.Errorf("inspect open archive: %w", statErr), + wrapExtractCloseError("archive", file.Close()), + wrapExtractCloseError("archive parent", closeParentErr), + ) + } + if !openedInfo.Mode().IsRegular() || !os.SameFile(info, openedInfo) { + return nil, errors.Join( + fmt.Errorf("archive changed while opening"), + wrapExtractCloseError("archive", file.Close()), + wrapExtractCloseError("archive parent", closeParentErr), + ) + } + if closeParentErr != nil { + return nil, errors.Join( + fmt.Errorf("close archive parent: %w", closeParentErr), + wrapExtractCloseError("archive", file.Close()), + ) + } + return file, nil +} + +func prepareExtractDestination(output string) (string, string, error) { + if _, err := os.Lstat(output); err == nil { + return "", "", fmt.Errorf("extraction destination exists: %s", output) + } else if !errors.Is(err, os.ErrNotExist) { + return "", "", fmt.Errorf("inspect extraction destination: %w", err) + } + absolute, err := filepath.Abs(output) + if err != nil { + return "", "", fmt.Errorf("resolve extraction destination: %w", err) + } + physicalParent, err := filepath.EvalSymlinks(filepath.Dir(absolute)) + if err != nil { + return "", "", fmt.Errorf("resolve extraction destination parent: %w", err) + } + physicalOutput := filepath.Join(physicalParent, filepath.Base(absolute)) + if _, err := os.Lstat(physicalOutput); err == nil { + return "", "", fmt.Errorf("extraction destination exists: %s", output) + } else if !errors.Is(err, os.ErrNotExist) { + return "", "", fmt.Errorf("inspect physical extraction destination: %w", err) + } + return physicalParent, filepath.Base(physicalOutput), nil +} + +func requireRootPathIdentity(path string, root *os.Root, description string) error { + pathInfo, err := os.Lstat(path) + if err != nil { + return fmt.Errorf("inspect %s path: %w", description, err) + } + if pathInfo.Mode()&os.ModeSymlink != 0 || !pathInfo.IsDir() { + return fmt.Errorf("%s path is not a non-symlink directory", description) + } + pinnedInfo, err := root.Stat(".") + if err != nil { + return fmt.Errorf("inspect pinned %s: %w", description, err) + } + if !pinnedInfo.IsDir() || !os.SameFile(pathInfo, pinnedInfo) { + return fmt.Errorf("%s changed while processing", description) + } + return nil +} + +func createExtractStage( + parent *os.Root, + outputBase string, + operations archiveOperations, +) (string, *os.Root, os.FileInfo, error) { + for range 100 { + var random [16]byte + if _, err := rand.Read(random[:]); err != nil { + return "", nil, nil, fmt.Errorf("generate extraction stage name: %w", err) + } + name := "." + outputBase + ".extract-" + hex.EncodeToString(random[:]) + ".tmp" + if err := parent.Mkdir(name, 0o700); err == nil { + info, statErr := parent.Lstat(name) + if statErr != nil { + return "", nil, nil, errors.Join( + fmt.Errorf("inspect extraction stage: %w", statErr), + fmt.Errorf( + "ownership cleanup for extraction stage: created identity is unavailable; preserving pathname", + ), + ) + } + root, openErr := operations.runOpenRoot(parent, name) + if openErr != nil { + return name, nil, info, errors.Join( + fmt.Errorf("open extraction stage: %w", openErr), + fmt.Errorf( + "ownership cleanup for extraction stage: retained stage root is unavailable; preserving pathname", + ), + ) + } + pinnedInfo, pinnedStatErr := operations.runStatRoot(root) + if pinnedStatErr != nil || + pinnedInfo == nil || + !pinnedInfo.IsDir() || + !os.SameFile(info, pinnedInfo) { + if pinnedStatErr == nil { + pinnedStatErr = fmt.Errorf("created stage identity changed while opening") + } + return name, nil, info, errors.Join( + fmt.Errorf("inspect pinned extraction stage: %w", pinnedStatErr), + wrapExtractCloseError("extraction stage", root.Close()), + fmt.Errorf( + "ownership cleanup for extraction stage: pinned stage identity is unproven; preserving pathname", + ), + ) + } + if err := root.Chmod(".", 0o700); err != nil { + return "", nil, nil, errors.Join( + fmt.Errorf("set extraction stage mode: %w", err), + removeOwnedDirectory( + parent, + name, + root, + info, + "extraction stage", + operations, + ), + ) + } + return name, root, info, nil + } else if !errors.Is(err, os.ErrExist) { + return "", nil, nil, fmt.Errorf("create extraction stage: %w", err) + } + } + return "", nil, nil, fmt.Errorf("create unique extraction stage: name attempts exhausted") +} + +func extractCollectionTar( + ctx context.Context, + source io.Reader, + stage *os.Root, + observer observe.Observer, + operations archiveOperations, +) (map[string]struct{}, error) { + reader := tar.NewReader(source) + seen := make(map[string]struct{}) + for { + if err := ctx.Err(); err != nil { + return nil, fmt.Errorf("extract collection TAR: %w", err) + } + header, err := reader.Next() + if errors.Is(err, io.EOF) { + return seen, nil + } + if err != nil { + return nil, fmt.Errorf("read collection TAR header: %w", err) + } + if hasSparseTarMetadata(header) { + return nil, fmt.Errorf("TAR entry %q uses unsupported sparse metadata", header.Name) + } + if header.Typeflag != tar.TypeReg { + return nil, fmt.Errorf("TAR entry %q is not a regular file", header.Name) + } + if _, found := seen[header.Name]; found { + return nil, fmt.Errorf("duplicate TAR entry %q", header.Name) + } + if _, err := collection.SafeJoin(".", header.Name); err != nil { + return nil, fmt.Errorf("unsafe TAR entry %q: %w", header.Name, err) + } + if err := extractExclusive( + stage, + header.Name, + header.Size, + reader, + operations, + ); err != nil { + return nil, err + } + seen[header.Name] = struct{}{} + observe.Emit(ctx, observer, observe.ArchiveEntryProcessed{ + Operation: "unpack", + Path: header.Name, + Size: header.Size, + }) + if err := ctx.Err(); err != nil { + return nil, fmt.Errorf("extract collection TAR after %q: %w", header.Name, err) + } + } +} + +func hasSparseTarMetadata(header *tar.Header) bool { + if header.Typeflag == tar.TypeGNUSparse { + return true + } + for key, value := range header.PAXRecords { + if strings.HasPrefix(key, "GNU.sparse.") || + (key == "SCHILY.filetype" && strings.EqualFold(value, "sparse")) { + return true + } + } + return false +} + +func requireZeroTarTail(source io.Reader) error { + var buffer [32 * 1024]byte + for { + count, err := source.Read(buffer[:]) + for _, value := range buffer[:count] { + if value != 0 { + return fmt.Errorf("archive contains nonzero data after the TAR end marker") + } + } + if errors.Is(err, io.EOF) { + return nil + } + if err != nil { + return fmt.Errorf("authenticate archive remainder: %w", err) + } + if count == 0 { + return fmt.Errorf("authenticate archive remainder: %w", io.ErrNoProgress) + } + } +} + +func extractExclusive( + root *os.Root, + relative string, + size int64, + source io.Reader, + operations archiveOperations, +) (resultErr error) { + if size < 0 { + return fmt.Errorf("TAR entry %q has negative size", relative) + } + parentPath := path.Dir(relative) + if err := ensureExtractDirectories(root, parentPath); err != nil { + return err + } + parent, err := root.OpenRoot(filepath.FromSlash(parentPath)) + if err != nil { + return fmt.Errorf("open TAR entry parent %q: %w", parentPath, err) + } + defer func() { + resultErr = errors.Join(resultErr, wrapExtractCloseError("TAR entry parent "+parentPath, parent.Close())) + }() + + base := path.Base(relative) + file, err := parent.OpenFile(base, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o600) + if err != nil { + return fmt.Errorf("create TAR entry %q: %w", relative, err) + } + createdInfo, err := file.Stat() + if err != nil { + return errors.Join( + fmt.Errorf("inspect created TAR entry %q: %w", relative, err), + wrapExtractCloseError("TAR entry "+relative, file.Close()), + fmt.Errorf( + "ownership cleanup for incomplete TAR entry %q: created identity is unavailable; preserving pathname", + relative, + ), + ) + } + complete := false + defer func() { + if !complete { + resultErr = errors.Join( + resultErr, + sanitizeAndRemoveOwnedEntry( + parent, + base, + createdInfo, + file, + "incomplete TAR entry "+relative, + operations, + ), + ) + } + if file != nil { + resultErr = errors.Join(resultErr, wrapExtractCloseError("TAR entry "+relative, file.Close())) + } + }() + if err := file.Chmod(0o600); err != nil { + return fmt.Errorf("set TAR entry mode %q: %w", relative, err) + } + copied, err := io.CopyN(file, source, size) + if err != nil { + return fmt.Errorf("extract TAR entry %q after %d bytes: %w", relative, copied, err) + } + if err := file.Sync(); err != nil { + return fmt.Errorf("sync TAR entry %q: %w", relative, err) + } + closeErr := file.Close() + file = nil + if closeErr != nil { + return fmt.Errorf("close TAR entry %q: %w", relative, closeErr) + } + complete = true + return nil +} + +func ensureExtractDirectories(root *os.Root, relative string) error { + if relative == "." { + return nil + } + current := "" + for _, component := range strings.Split(filepath.FromSlash(relative), string(filepath.Separator)) { + current = filepath.Join(current, component) + info, err := root.Lstat(current) + if errors.Is(err, os.ErrNotExist) { + if err := root.Mkdir(current, 0o700); err != nil { + return fmt.Errorf("create TAR directory %q: %w", filepath.ToSlash(current), err) + } + continue + } + if err != nil { + return fmt.Errorf("inspect TAR directory %q: %w", filepath.ToSlash(current), err) + } + if info.Mode()&os.ModeSymlink != 0 || !info.IsDir() { + return fmt.Errorf("TAR directory %q is not a non-symlink directory", filepath.ToSlash(current)) + } + } + return nil +} + +func compareExtractedFileSet(seen map[string]struct{}, declared []string) error { + declaredSet := make(map[string]struct{}, len(declared)) + for _, relative := range declared { + declaredSet[relative] = struct{}{} + if _, found := seen[relative]; !found { + return fmt.Errorf("declared collection file %q is missing from archive", relative) + } + } + var unexpected []string + for relative := range seen { + if _, declared := declaredSet[relative]; !declared { + unexpected = append(unexpected, relative) + } + } + if len(unexpected) != 0 { + sort.Strings(unexpected) + return fmt.Errorf("unexpected archive file %q", unexpected[0]) + } + return nil +} + +func verifyPinnedExtractedCollection( + root *os.Root, + initialPlan collectionTarPlan, + manifest collection.Manifest, + manifestDigest string, + declared []string, +) error { + currentManifest, currentManifestDigest, currentManifestInfo, err := readPinnedCollectionManifest(root) + if err != nil { + return fmt.Errorf("re-read authenticated stage manifest: %w", err) + } + if !reflect.DeepEqual(manifest, currentManifest) || manifestDigest != currentManifestDigest { + return fmt.Errorf("authenticated stage manifest changed during verification") + } + plannedManifest, ok := plannedTarFile(initialPlan.files, collection.ManifestName) + if !ok || !samePlannedFile(plannedManifest.info, currentManifestInfo) { + return fmt.Errorf("authenticated stage manifest identity changed during verification") + } + + expectedDigests := manifestArtifactSHA256s(manifest) + for relative, expectedDigest := range expectedDigests { + planned, ok := plannedTarFile(initialPlan.files, relative) + if !ok { + return fmt.Errorf("verified artifact %q was not inventoried", relative) + } + actualDigest, err := hashPinnedCollectionFile(root, planned) + if err != nil { + return err + } + if actualDigest != expectedDigest { + return fmt.Errorf( + "verified artifact %q SHA-256 is %s, want %s", + relative, + actualDigest, + expectedDigest, + ) + } + } + finalPlan, err := inventoryPinnedCollectionTarFiles(root, declared) + if err != nil { + return fmt.Errorf("final authenticated stage inventory: %w", err) + } + if err := compareCollectionTarPlans(initialPlan, finalPlan); err != nil { + return fmt.Errorf("authenticated stage changed during verification: %w", err) + } + return nil +} + +func hashPinnedCollectionFile(root *os.Root, planned collectionTarFile) ( + digest string, + resultErr error, +) { + if err := requirePinnedRegularPath(root, planned.path); err != nil { + return "", err + } + info, err := root.Lstat(filepath.FromSlash(planned.path)) + if err != nil { + return "", fmt.Errorf("inspect verified artifact %q: %w", planned.path, err) + } + if !samePlannedFile(planned.info, info) { + return "", fmt.Errorf("verified artifact %q changed before hashing", planned.path) + } + file, err := root.Open(filepath.FromSlash(planned.path)) + if err != nil { + return "", fmt.Errorf("open verified artifact %q: %w", planned.path, err) + } + defer func() { + resultErr = errors.Join( + resultErr, + wrapExtractCloseError("verified artifact "+planned.path, file.Close()), + ) + }() + openedInfo, err := file.Stat() + if err != nil { + return "", fmt.Errorf("inspect open verified artifact %q: %w", planned.path, err) + } + if !samePlannedFile(planned.info, openedInfo) { + return "", fmt.Errorf("verified artifact %q changed while opening", planned.path) + } + hasher := sha256.New() + if _, err := io.Copy(hasher, file); err != nil { + return "", fmt.Errorf("hash verified artifact %q: %w", planned.path, err) + } + return hex.EncodeToString(hasher.Sum(nil)), nil +} + +func syncExtractedCollection( + ctx context.Context, + root *os.Root, + declared []string, + operations archiveOperations, +) error { + for _, relative := range declared { + if err := ctx.Err(); err != nil { + return fmt.Errorf("sync extracted collection: %w", err) + } + file, err := root.Open(filepath.FromSlash(relative)) + if err != nil { + return fmt.Errorf("open extracted file %q for sync: %w", relative, err) + } + _, statErr := file.Stat() + syncErr := file.Sync() + closeErr := file.Close() + if err := errors.Join( + wrapExtractStatError("extracted file "+relative, statErr), + wrapExtractSyncError("extracted file "+relative, syncErr), + wrapExtractCloseError("extracted file "+relative, closeErr), + ); err != nil { + return err + } + if err := operations.runAfterExtractFileSync(root, relative); err != nil { + return fmt.Errorf("after extracted file sync %q: %w", relative, err) + } + if err := ctx.Err(); err != nil { + return fmt.Errorf("sync extracted collection after %q: %w", relative, err) + } + } + + directorySet := map[string]struct{}{".": {}} + for _, relative := range declared { + for directory := path.Dir(relative); directory != "."; directory = path.Dir(directory) { + directorySet[directory] = struct{}{} + } + } + directories := make([]string, 0, len(directorySet)) + for directory := range directorySet { + directories = append(directories, directory) + } + sort.Slice(directories, func(left, right int) bool { + leftDepth := strings.Count(directories[left], "/") + rightDepth := strings.Count(directories[right], "/") + if leftDepth != rightDepth { + return leftDepth > rightDepth + } + return directories[left] > directories[right] + }) + for _, directory := range directories { + file, err := root.Open(filepath.FromSlash(directory)) + if err != nil { + return fmt.Errorf("open extracted directory %q for sync: %w", directory, err) + } + syncErr := file.Sync() + closeErr := file.Close() + if err := errors.Join( + wrapExtractSyncError("extracted directory "+directory, syncErr), + wrapExtractCloseError("extracted directory "+directory, closeErr), + ); err != nil { + return err + } + } + return nil +} + +func wrapExtractCloseError(name string, err error) error { + if err == nil { + return nil + } + return fmt.Errorf("close %s: %w", name, err) +} + +func wrapExtractSyncError(name string, err error) error { + if err == nil { + return nil + } + return fmt.Errorf("sync %s: %w", name, err) +} + +func wrapExtractStatError(name string, err error) error { + if err == nil { + return nil + } + return fmt.Errorf("inspect %s: %w", name, err) +} diff --git a/ret/archive/unpack_test.go b/ret/archive/unpack_test.go new file mode 100644 index 00000000..1b19f910 --- /dev/null +++ b/ret/archive/unpack_test.go @@ -0,0 +1,904 @@ +package archive + +import ( + "archive/tar" + "bytes" + "context" + "errors" + "io" + "io/fs" + "os" + "path/filepath" + "runtime" + "sort" + "testing" + + "github.com/specterops/dawgs/ret/collection" + "github.com/specterops/dawgs/ret/observe" + "github.com/stretchr/testify/require" +) + +type archiveTestTarEntry struct { + header tar.Header + payload []byte +} + +type archiveTestBoundedZeroReader struct { + remaining int + maxRequest int +} + +type archiveTestReadFunc func([]byte) (int, error) + +func (s archiveTestReadFunc) Read(destination []byte) (int, error) { + return s(destination) +} + +func (s *archiveTestBoundedZeroReader) Read(destination []byte) (int, error) { + if len(destination) > s.maxRequest { + return 0, errors.New("read request exceeded fixed bound") + } + if s.remaining == 0 { + return 0, io.EOF + } + count := min(len(destination), s.remaining) + clear(destination[:count]) + s.remaining -= count + return count, nil +} + +func TestExtractPromotesOnlyACompleteVerifiedCollection(t *testing.T) { + for _, outputs := range []struct { + name string + jsonl, parquet bool + expectedFileCount int + }{ + {name: "JSONL only", jsonl: true, expectedFileCount: 3}, + {name: "Parquet only", parquet: true, expectedFileCount: 3}, + {name: "dual", jsonl: true, parquet: true, expectedFileCount: 5}, + } { + t.Run(outputs.name, func(t *testing.T) { + source := writeArchiveTestCollection(t, outputs.jsonl, outputs.parquet) + recipient, identity, err := GenerateKeyPair() + require.NoError(t, err) + archivePath := writeIndependentArchiveTestEnvelope( + t, + recipient, + archiveTestRegularEntries(t, source), + ) + output := filepath.Join(t.TempDir(), "collection") + + require.NoError(t, Extract(context.Background(), ExtractConfig{ + ArchivePath: archivePath, + OutputDirectory: output, + Identity: identity, + })) + + _, err = collection.Verify(context.Background(), output, nil) + require.NoError(t, err) + require.Len(t, archiveTestRegularEntries(t, output), outputs.expectedFileCount) + require.Empty(t, archiveExtractStagePaths(t, output)) + requireArchiveTestTreesEqual(t, source, output) + }) + } +} + +func TestExtractRejectsUnsafeTarEntriesWithoutFilesystemEscape(t *testing.T) { + recipient, identity, err := GenerateKeyPair() + require.NoError(t, err) + + tests := []struct { + name string + entry archiveTestTarEntry + match string + escape string + }{ + { + name: "absolute path", + entry: archiveTestEntry( + filepath.Join(string(filepath.Separator), "tmp", "ret-archive-absolute-escape"), + []byte("bad"), + ), + match: "absolute", + }, + { + name: "parent traversal", + entry: archiveTestEntry("../escape", []byte("bad")), + match: "traverses", + escape: "escape", + }, + { + name: "unclean path", + entry: archiveTestEntry("graphs/../escape", []byte("bad")), + match: "not clean", + }, + { + name: "backslash path", + entry: archiveTestEntry(`graphs\escape`, []byte("bad")), + match: "backslash", + }, + { + name: "symlink", + entry: archiveTestTarEntry{header: tar.Header{ + Name: collection.ManifestName, + Typeflag: tar.TypeSymlink, + Linkname: "target", + }}, + match: "regular", + }, + { + name: "hardlink", + entry: archiveTestTarEntry{header: tar.Header{ + Name: collection.ManifestName, + Typeflag: tar.TypeLink, + Linkname: "target", + }}, + match: "regular", + }, + { + name: "character device", + entry: archiveTestTarEntry{header: tar.Header{ + Name: collection.ManifestName, + Typeflag: tar.TypeChar, + Devmajor: 1, + Devminor: 3, + }}, + match: "regular", + }, + { + name: "directory", + entry: archiveTestTarEntry{header: tar.Header{ + Name: "graphs", + Typeflag: tar.TypeDir, + Mode: 0o700, + }}, + match: "regular", + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + parent := t.TempDir() + output := filepath.Join(parent, "collection") + archivePath := writeIndependentArchiveTestEnvelope(t, recipient, []archiveTestTarEntry{test.entry}) + + err := Extract(context.Background(), ExtractConfig{ + ArchivePath: archivePath, + OutputDirectory: output, + Identity: identity, + }) + + require.ErrorContains(t, err, test.match) + require.NoFileExists(t, output) + require.NoDirExists(t, output) + require.Empty(t, archiveExtractStagePaths(t, output)) + if test.escape != "" { + require.NoFileExists(t, filepath.Join(parent, test.escape)) + } + }) + } +} + +func TestExtractRejectsSparsePAXMetadata(t *testing.T) { + // Break caught: allowing archive/tar to expand a tiny authenticated GNU + // sparse payload into an attacker-selected logical file size on disk. + recipient, identity, err := GenerateKeyPair() + require.NoError(t, err) + entry := archiveTestEntry(collection.ManifestName, []byte("abc")) + entry.header.Format = tar.FormatPAX + entry.header.PAXRecords = map[string]string{ + "ABC.sparse.map": "0,3", + "ABC.sparse.numblocks": "1", + "ABC.sparse.size": "3", + } + plaintext := archiveTestTarPayload(t, []archiveTestTarEntry{entry}) + for _, field := range []string{"map", "numblocks", "size"} { + before := []byte("ABC.sparse." + field) + after := []byte("GNU.sparse." + field) + require.True(t, bytes.Contains(plaintext, before)) + plaintext = bytes.ReplaceAll(plaintext, before, after) + } + archivePath := writeIndependentArchiveTestEnvelopePayload(t, recipient, plaintext) + output := filepath.Join(t.TempDir(), "collection") + + err = Extract(context.Background(), ExtractConfig{ + ArchivePath: archivePath, + OutputDirectory: output, + Identity: identity, + }) + + require.ErrorContains(t, err, "sparse") + require.NoDirExists(t, output) + require.Empty(t, archiveExtractStagePaths(t, output)) +} + +func TestExtractRejectsDuplicateAndNonDeclaredFileSets(t *testing.T) { + source := writeArchiveTestCollection(t, true, true) + recipient, identity, err := GenerateKeyPair() + require.NoError(t, err) + valid := archiveTestRegularEntries(t, source) + require.Greater(t, len(valid), 2) + + tests := []struct { + name string + entries []archiveTestTarEntry + match string + }{ + { + name: "duplicate", + entries: append(append([]archiveTestTarEntry(nil), valid...), valid[0]), + match: "duplicate", + }, + { + name: "undeclared file", + entries: append(append([]archiveTestTarEntry(nil), valid...), archiveTestEntry("extra.txt", []byte("extra"))), + match: "unexpected", + }, + { + name: "missing manifest", + entries: func() []archiveTestTarEntry { + result := append([]archiveTestTarEntry(nil), valid...) + for index := range result { + if result[index].header.Name == collection.ManifestName { + return append(result[:index], result[index+1:]...) + } + } + t.Fatal("valid archive did not contain manifest") + return nil + }(), + match: "manifest", + }, + { + name: "missing artifact", + entries: append([]archiveTestTarEntry(nil), valid[1:]...), + match: "missing", + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + output := filepath.Join(t.TempDir(), "collection") + archivePath := writeIndependentArchiveTestEnvelope(t, recipient, test.entries) + + err := Extract(context.Background(), ExtractConfig{ + ArchivePath: archivePath, + OutputDirectory: output, + Identity: identity, + }) + + require.ErrorContains(t, err, test.match) + require.NoDirExists(t, output) + require.NoFileExists(t, output) + require.Empty(t, archiveExtractStagePaths(t, output)) + }) + } +} + +func TestExtractRejectsAuthenticatedCorruptParquetBeforePromotion(t *testing.T) { + source := writeArchiveTestCollection(t, true, true) + recipient, identity, err := GenerateKeyPair() + require.NoError(t, err) + entries := archiveTestRegularEntries(t, source) + for index := range entries { + if filepath.Ext(entries[index].header.Name) == ".parquet" { + entries[index].payload = []byte("authenticated but corrupt parquet") + entries[index].header.Size = int64(len(entries[index].payload)) + break + } + } + output := filepath.Join(t.TempDir(), "collection") + archivePath := writeIndependentArchiveTestEnvelope(t, recipient, entries) + + err = Extract(context.Background(), ExtractConfig{ + ArchivePath: archivePath, + OutputDirectory: output, + Identity: identity, + }) + + require.Error(t, err) + require.NoDirExists(t, output) + require.Empty(t, archiveExtractStagePaths(t, output)) +} + +func TestExtractAuthenticatesFinalFrameBeforePromotion(t *testing.T) { + source := writeArchiveTestCollection(t, true, false) + recipient, identity, err := GenerateKeyPair() + require.NoError(t, err) + archivePath := writeIndependentArchiveTestEnvelope( + t, + recipient, + archiveTestRegularEntries(t, source), + ) + wire := mustReadArchiveTestFile(t, archivePath) + _, _, frames := parseTestEnvelope(t, wire) + require.GreaterOrEqual(t, len(frames), 2) + require.NoError(t, os.WriteFile(archivePath, wire[:frames[len(frames)-1].start], 0o600)) + output := filepath.Join(t.TempDir(), "collection") + + err = Extract(context.Background(), ExtractConfig{ + ArchivePath: archivePath, + OutputDirectory: output, + Identity: identity, + }) + + require.ErrorContains(t, err, "final frame") + require.NoDirExists(t, output) + require.Empty(t, archiveExtractStagePaths(t, output)) +} + +func TestRequireZeroTarTailUsesBoundedStreamingReads(t *testing.T) { + // Break caught: buffering an attacker-controlled authenticated zero trailer + // after the TAR terminator instead of scanning it with fixed memory. + reader := &archiveTestBoundedZeroReader{ + remaining: 4 * 1024 * 1024, + maxRequest: 64 * 1024, + } + + require.NoError(t, requireZeroTarTail(reader)) + require.Zero(t, reader.remaining) + require.ErrorContains(t, requireZeroTarTail(bytes.NewReader([]byte{0, 0, 1, 0})), "nonzero") +} + +func TestExtractPreservesAnyExistingOutput(t *testing.T) { + source := writeArchiveTestCollection(t, true, false) + recipient, identity, err := GenerateKeyPair() + require.NoError(t, err) + archivePath := writeIndependentArchiveTestEnvelope( + t, + recipient, + archiveTestRegularEntries(t, source), + ) + + t.Run("directory", func(t *testing.T) { + output := filepath.Join(t.TempDir(), "collection") + require.NoError(t, os.Mkdir(output, 0o700)) + marker := filepath.Join(output, "preserve") + require.NoError(t, os.WriteFile(marker, []byte("preserve"), 0o600)) + + err := Extract(context.Background(), ExtractConfig{ + ArchivePath: archivePath, + OutputDirectory: output, + Identity: identity, + }) + + require.ErrorContains(t, err, "exists") + require.Equal(t, []byte("preserve"), mustReadArchiveTestFile(t, marker)) + require.Empty(t, archiveExtractStagePaths(t, output)) + }) + + t.Run("file", func(t *testing.T) { + output := filepath.Join(t.TempDir(), "collection") + require.NoError(t, os.WriteFile(output, []byte("preserve"), 0o600)) + + err := Extract(context.Background(), ExtractConfig{ + ArchivePath: archivePath, + OutputDirectory: output, + Identity: identity, + }) + + require.ErrorContains(t, err, "exists") + require.Equal(t, []byte("preserve"), mustReadArchiveTestFile(t, output)) + require.Empty(t, archiveExtractStagePaths(t, output)) + }) +} + +func TestExtractObserverCancellationSanitizesStageBeforeQuarantine(t *testing.T) { + source := writeArchiveTestCollection(t, true, true) + recipient, identity, err := GenerateKeyPair() + require.NoError(t, err) + archivePath := writeIndependentArchiveTestEnvelope( + t, + recipient, + archiveTestRegularEntries(t, source), + ) + outputParent := t.TempDir() + output := filepath.Join(outputParent, "collection") + ctx, cancel := context.WithCancel(context.Background()) + observer := observe.ObserverFunc(func(_ context.Context, event observe.Event) { + if entry, ok := event.(observe.ArchiveEntryProcessed); ok && entry.Operation == "unpack" { + cancel() + } + }) + + err = Extract(ctx, ExtractConfig{ + ArchivePath: archivePath, + OutputDirectory: output, + Identity: identity, + Observer: observer, + }) + + require.ErrorIs(t, err, context.Canceled) + require.ErrorContains(t, err, "ownership cleanup") + require.NoDirExists(t, output) + require.Empty(t, archiveExtractStagePaths(t, output)) + quarantines := archiveCleanupQuarantinePaths(t, outputParent) + require.Len(t, quarantines, 1) + info, statErr := os.Stat(quarantines[0]) + require.NoError(t, statErr) + require.True(t, info.IsDir()) + entries, readErr := os.ReadDir(quarantines[0]) + require.NoError(t, readErr) + require.Empty(t, entries) +} + +func TestExtractCleanupPreservesReplacedStageAndCleansOwnedContents(t *testing.T) { + // Break caught: recursively deleting a replacement installed at the stage + // pathname instead of cleaning only the pinned stage created by Extract. + source := writeArchiveTestCollection(t, true, false) + recipient, identity, err := GenerateKeyPair() + require.NoError(t, err) + archivePath := writeIndependentArchiveTestEnvelope( + t, + recipient, + archiveTestRegularEntries(t, source), + ) + parent := t.TempDir() + output := filepath.Join(parent, "collection") + movedStage := filepath.Join(parent, "owned-stage-moved") + ctx, cancel := context.WithCancel(context.Background()) + replaced := false + observer := observe.ObserverFunc(func(_ context.Context, event observe.Event) { + entry, ok := event.(observe.ArchiveEntryProcessed) + if !ok || entry.Operation != "unpack" || replaced { + return + } + stages := archiveExtractStagePaths(t, output) + require.Len(t, stages, 1) + require.NoError(t, os.Rename(stages[0], movedStage)) + require.NoError(t, os.Mkdir(stages[0], 0o700)) + require.NoError(t, os.WriteFile( + filepath.Join(stages[0], "replacement-marker"), + []byte("preserve replacement"), + 0o600, + )) + replaced = true + cancel() + }) + + err = Extract(ctx, ExtractConfig{ + ArchivePath: archivePath, + OutputDirectory: output, + Identity: identity, + Observer: observer, + }) + + require.ErrorIs(t, err, context.Canceled) + require.ErrorContains(t, err, "ownership cleanup") + require.True(t, replaced) + require.NoDirExists(t, output) + stages := archiveExtractStagePaths(t, output) + require.Len(t, stages, 1) + require.Equal( + t, + []byte("preserve replacement"), + mustReadArchiveTestFile(t, filepath.Join(stages[0], "replacement-marker")), + ) + ownedEntries, readErr := os.ReadDir(movedStage) + require.NoError(t, readErr) + require.Empty(t, ownedEntries) +} + +func TestCreateExtractStagePreservesNameWhenPinningFails(t *testing.T) { + // Break caught: removing a stage using stale FileInfo when no retained root + // was obtained, or after the retained root could not be proven identical. + tests := []struct { + name string + operations archiveOperations + match string + mismatch bool + nilInfo bool + nonDir bool + }{ + { + name: "open root", + operations: archiveOperations{ + openRoot: func(_ *os.Root, _ string) (*os.Root, error) { + return nil, errors.New("injected open root failure") + }, + }, + match: "injected open root failure", + }, + { + name: "stat pinned root", + operations: archiveOperations{ + statRoot: func(_ *os.Root) (fs.FileInfo, error) { + return nil, errors.New("injected pinned stat failure") + }, + }, + match: "injected pinned stat failure", + }, + { + name: "mismatched pinned root", + match: "identity changed", + mismatch: true, + }, + { + name: "nil pinned info", + match: "identity changed", + nilInfo: true, + }, + { + name: "non-directory pinned info", + match: "identity changed", + nonDir: true, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + parentPath := t.TempDir() + parent, err := os.OpenRoot(parentPath) + require.NoError(t, err) + t.Cleanup(func() { + require.NoError(t, parent.Close()) + }) + operations := test.operations + if test.mismatch { + differentPath := t.TempDir() + different, statErr := os.Stat(differentPath) + require.NoError(t, statErr) + operations.statRoot = func(_ *os.Root) (fs.FileInfo, error) { + return different, nil + } + } + if test.nilInfo { + operations.statRoot = func(_ *os.Root) (fs.FileInfo, error) { + return nil, nil + } + } + if test.nonDir { + differentPath := filepath.Join(t.TempDir(), "regular") + require.NoError(t, os.WriteFile(differentPath, []byte("regular"), 0o600)) + different, statErr := os.Stat(differentPath) + require.NoError(t, statErr) + operations.statRoot = func(_ *os.Root) (fs.FileInfo, error) { + return different, nil + } + } + + name, root, _, err := createExtractStage(parent, "collection", operations) + + require.ErrorContains(t, err, test.match) + require.ErrorContains(t, err, "ownership cleanup") + require.Nil(t, root) + require.NotEmpty(t, name) + info, statErr := parent.Lstat(name) + require.NoError(t, statErr) + require.True(t, info.IsDir()) + }) + } +} + +func TestExtractExclusiveCleanupPreservesReplacementEntry(t *testing.T) { + // Break caught: unlinking a replacement file when extraction of the owned + // file fails after its pathname has been substituted. + stagePath := t.TempDir() + stage, err := os.OpenRoot(stagePath) + require.NoError(t, err) + t.Cleanup(func() { + require.NoError(t, stage.Close()) + }) + entryPath := filepath.Join(stagePath, "artifact.bin") + replaced := false + source := archiveTestReadFunc(func(_ []byte) (int, error) { + require.NoError(t, os.Remove(entryPath)) + require.NoError(t, os.WriteFile(entryPath, []byte("preserve replacement"), 0o600)) + replaced = true + return 0, errors.New("source failed") + }) + + err = extractExclusive(stage, "artifact.bin", 32, source, archiveOperations{}) + + require.ErrorContains(t, err, "source failed") + require.ErrorContains(t, err, "ownership cleanup") + require.True(t, replaced) + require.Equal(t, []byte("preserve replacement"), mustReadArchiveTestFile(t, entryPath)) +} + +func TestExtractExclusiveSanitizesIncompleteEntryBeforeQuarantine(t *testing.T) { + stagePath := t.TempDir() + stage, err := os.OpenRoot(stagePath) + require.NoError(t, err) + t.Cleanup(func() { + require.NoError(t, stage.Close()) + }) + source := io.MultiReader( + bytes.NewReader([]byte("partial secret")), + archiveTestReadFunc(func(_ []byte) (int, error) { + return 0, errors.New("source failed") + }), + ) + + err = extractExclusive(stage, "artifact.bin", 32, source, archiveOperations{}) + + require.ErrorContains(t, err, "source failed") + require.ErrorContains(t, err, "ownership cleanup") + require.NoFileExists(t, filepath.Join(stagePath, "artifact.bin")) + quarantines := archiveCleanupQuarantinePaths(t, stagePath) + require.Len(t, quarantines, 1) + info, statErr := os.Stat(quarantines[0]) + require.NoError(t, statErr) + require.True(t, info.Mode().IsRegular()) + require.Zero(t, info.Size()) +} + +func TestExtractExclusivePreservesIncompleteEntryWhenSanitizationFails(t *testing.T) { + // Break caught: quarantining partial artifact bytes when retained-handle + // truncate or sync reports failure. + tests := []struct { + name string + operations archiveOperations + wantPayload bool + match string + }{ + { + name: "truncate", + operations: archiveOperations{ + truncateOwnedFile: func(_ *os.File, _ int64) error { + return errors.New("injected truncate failure") + }, + }, + wantPayload: true, + match: "injected truncate failure", + }, + { + name: "sync", + operations: archiveOperations{ + syncOwnedFile: func(_ *os.File) error { + return errors.New("injected sync failure") + }, + }, + match: "injected sync failure", + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + stagePath := t.TempDir() + stage, err := os.OpenRoot(stagePath) + require.NoError(t, err) + t.Cleanup(func() { + require.NoError(t, stage.Close()) + }) + source := io.MultiReader( + bytes.NewReader([]byte("partial secret")), + archiveTestReadFunc(func(_ []byte) (int, error) { + return 0, errors.New("source failed") + }), + ) + + err = extractExclusive(stage, "artifact.bin", 32, source, test.operations) + + require.ErrorContains(t, err, "source failed") + require.ErrorContains(t, err, test.match) + require.ErrorContains(t, err, "ownership cleanup") + payload := mustReadArchiveTestFile(t, filepath.Join(stagePath, "artifact.bin")) + if test.wantPayload { + require.Equal(t, []byte("partial secret"), payload) + } else { + require.Empty(t, payload) + } + require.Empty(t, archiveCleanupQuarantinePaths(t, stagePath)) + }) + } +} + +func TestExtractPreservesDestinationCreatedBeforePromotion(t *testing.T) { + // Break caught: replacing a destination created after extraction started + // but before the verified stage was promoted. + source := writeArchiveTestCollection(t, true, true) + recipient, identity, err := GenerateKeyPair() + require.NoError(t, err) + archivePath := writeIndependentArchiveTestEnvelope( + t, + recipient, + archiveTestRegularEntries(t, source), + ) + output := filepath.Join(t.TempDir(), "collection") + created := false + observer := observe.ObserverFunc(func(_ context.Context, event observe.Event) { + if _, ok := event.(observe.ArtifactVerified); ok && !created { + created = true + require.NoError(t, os.WriteFile(output, []byte("concurrent owner"), 0o600)) + } + }) + + err = Extract(context.Background(), ExtractConfig{ + ArchivePath: archivePath, + OutputDirectory: output, + Identity: identity, + Observer: observer, + }) + + require.Error(t, err) + require.True(t, created) + require.Equal(t, []byte("concurrent owner"), mustReadArchiveTestFile(t, output)) + require.Empty(t, archiveExtractStagePaths(t, output)) +} + +func TestExtractRejectsInPlaceMutationAfterPathVerification(t *testing.T) { + // Break caught: promoting stage bytes changed in place immediately after + // pathname-based collection verification consumed them. + source := writeArchiveTestCollection(t, true, false) + recipient, identity, err := GenerateKeyPair() + require.NoError(t, err) + archivePath := writeIndependentArchiveTestEnvelope( + t, + recipient, + archiveTestRegularEntries(t, source), + ) + parent := t.TempDir() + output := filepath.Join(parent, "collection") + mutated := false + observer := observe.ObserverFunc(func(_ context.Context, event observe.Event) { + verified, ok := event.(observe.ArtifactVerified) + if !ok || mutated { + return + } + stages := archiveExtractStagePaths(t, output) + require.Len(t, stages, 1) + artifactPath := filepath.Join(stages[0], filepath.FromSlash(verified.Path)) + info, err := os.Stat(artifactPath) + require.NoError(t, err) + payload := mustReadArchiveTestFile(t, artifactPath) + require.NotEmpty(t, payload) + payload[0] ^= 0x01 + require.NoError(t, os.WriteFile(artifactPath, payload, info.Mode().Perm())) + require.NoError(t, os.Chtimes(artifactPath, info.ModTime(), info.ModTime())) + mutated = true + }) + + err = Extract(context.Background(), ExtractConfig{ + ArchivePath: archivePath, + OutputDirectory: output, + Identity: identity, + Observer: observer, + }) + + require.ErrorContains(t, err, "SHA-256") + require.True(t, mutated) + require.NoDirExists(t, output) + require.Empty(t, archiveExtractStagePaths(t, output)) +} + +func TestExtractRejectsMutationAfterStageSyncBeforePromotion(t *testing.T) { + // Break caught: promoting bytes changed after their fsync because the last + // rooted digest and exact-inventory binding ran before the sync pass. + source := writeArchiveTestCollection(t, true, false) + manifest := readArchiveTestManifest(t, source) + artifact := manifest.Graphs[0].NodeShards[0].JSONL.Path + recipient, identity, err := GenerateKeyPair() + require.NoError(t, err) + archivePath := writeIndependentArchiveTestEnvelope( + t, + recipient, + archiveTestRegularEntries(t, source), + ) + output := filepath.Join(t.TempDir(), "collection") + mutated := false + operations := archiveOperations{ + afterExtractFileSync: func(_ *os.Root, relative string) error { + if relative != artifact || mutated { + return nil + } + stages := archiveExtractStagePaths(t, output) + require.Len(t, stages, 1) + artifactPath := filepath.Join(stages[0], filepath.FromSlash(artifact)) + info, err := os.Stat(artifactPath) + require.NoError(t, err) + payload := mustReadArchiveTestFile(t, artifactPath) + require.NotEmpty(t, payload) + payload[0] ^= 0x01 + require.NoError(t, os.WriteFile(artifactPath, payload, info.Mode().Perm())) + require.NoError(t, os.Chtimes(artifactPath, info.ModTime(), info.ModTime())) + mutated = true + return nil + }, + } + + err = extract(context.Background(), ExtractConfig{ + ArchivePath: archivePath, + OutputDirectory: output, + Identity: identity, + }, runtime.GOOS, operations) + + require.ErrorContains(t, err, "SHA-256") + require.True(t, mutated) + require.NoDirExists(t, output) + require.Empty(t, archiveExtractStagePaths(t, output)) +} + +func archiveTestRegularEntries(t *testing.T, root string) []archiveTestTarEntry { + t.Helper() + var entries []archiveTestTarEntry + require.NoError(t, filepath.WalkDir(root, func(candidate string, entry fs.DirEntry, err error) error { + require.NoError(t, err) + if entry.IsDir() { + return nil + } + relative, err := filepath.Rel(root, candidate) + require.NoError(t, err) + entries = append(entries, archiveTestEntry( + filepath.ToSlash(relative), + mustReadArchiveTestFile(t, candidate), + )) + return nil + })) + sort.Slice(entries, func(left, right int) bool { + return entries[left].header.Name < entries[right].header.Name + }) + return entries +} + +func archiveTestEntry(name string, payload []byte) archiveTestTarEntry { + return archiveTestTarEntry{ + header: tar.Header{ + Name: name, + Mode: 0o600, + Size: int64(len(payload)), + Typeflag: tar.TypeReg, + Format: tar.FormatUSTAR, + }, + payload: append([]byte(nil), payload...), + } +} + +func writeIndependentArchiveTestEnvelope( + t *testing.T, + recipient PublicKey, + entries []archiveTestTarEntry, +) string { + t.Helper() + return writeIndependentArchiveTestEnvelopePayload(t, recipient, archiveTestTarPayload(t, entries)) +} + +func archiveTestTarPayload(t *testing.T, entries []archiveTestTarEntry) []byte { + t.Helper() + var plaintext bytes.Buffer + tarWriter := tar.NewWriter(&plaintext) + for index := range entries { + header := entries[index].header + require.NoError(t, tarWriter.WriteHeader(&header)) + if len(entries[index].payload) != 0 { + _, err := tarWriter.Write(entries[index].payload) + require.NoError(t, err) + } + } + require.NoError(t, tarWriter.Close()) + return append([]byte(nil), plaintext.Bytes()...) +} + +func writeIndependentArchiveTestEnvelopePayload(t *testing.T, recipient PublicKey, plaintext []byte) string { + t.Helper() + var encrypted bytes.Buffer + writer, err := newEncryptWriter(&encrypted, recipient) + require.NoError(t, err) + _, err = writer.Write(plaintext) + require.NoError(t, err) + require.NoError(t, writer.Close()) + + path := filepath.Join(t.TempDir(), "archive.ret") + require.NoError(t, os.WriteFile(path, encrypted.Bytes(), 0o600)) + return path +} + +func archiveExtractStagePaths(t *testing.T, output string) []string { + t.Helper() + matches, err := filepath.Glob(filepath.Join( + filepath.Dir(output), + "."+filepath.Base(output)+".extract-*.tmp", + )) + require.NoError(t, err) + return matches +} + +func requireArchiveTestTreesEqual(t *testing.T, expected, actual string) { + t.Helper() + expectedEntries := archiveTestRegularEntries(t, expected) + actualEntries := archiveTestRegularEntries(t, actual) + require.Len(t, actualEntries, len(expectedEntries)) + for index := range expectedEntries { + require.Equal(t, expectedEntries[index].header.Name, actualEntries[index].header.Name) + require.Equal(t, expectedEntries[index].payload, actualEntries[index].payload) + } +} diff --git a/ret/archive_platform_unsupported_test.go b/ret/archive_platform_unsupported_test.go new file mode 100644 index 00000000..cbf25676 --- /dev/null +++ b/ret/archive_platform_unsupported_test.go @@ -0,0 +1,25 @@ +//go:build !linux && !darwin + +package ret + +import ( + "path/filepath" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestKeygenRemainsPortableWithoutArchivePublication(t *testing.T) { + root := t.TempDir() + privatePath := filepath.Join(root, "private.json") + publicPath := filepath.Join(root, "public.json") + + err := Keygen(KeygenConfig{ + PrivateKeyPath: privatePath, + PublicKeyPath: publicPath, + }) + + require.NoError(t, err) + require.FileExists(t, privatePath) + require.FileExists(t, publicPath) +} diff --git a/ret/archive_test.go b/ret/archive_test.go new file mode 100644 index 00000000..d948aa6b --- /dev/null +++ b/ret/archive_test.go @@ -0,0 +1,321 @@ +package ret + +import ( + "context" + "os" + "path/filepath" + "testing" + "time" + + retarchive "github.com/specterops/dawgs/ret/archive" + "github.com/specterops/dawgs/ret/collection" + "github.com/specterops/dawgs/ret/entity" + "github.com/specterops/dawgs/ret/jsonl" + "github.com/specterops/dawgs/ret/metrics" + "github.com/specterops/dawgs/ret/observe" + "github.com/specterops/dawgs/ret/parquet" + "github.com/stretchr/testify/require" +) + +func TestArchiveFacadeConfigsValidateEveryRequiredInput(t *testing.T) { + public, private, err := retarchive.GenerateKeyPair() + require.NoError(t, err) + + validPack := PackConfig{ + CollectionDirectory: "collection", + ArchivePath: "archive.ret", + Recipient: public, + } + validUnpack := UnpackConfig{ + ArchivePath: "archive.ret", + OutputDirectory: "collection", + Identity: private, + } + validKeygen := KeygenConfig{PrivateKeyPath: "private.json", PublicKeyPath: "public.json"} + + for _, test := range []struct { + name string + validate func() error + }{ + {name: "pack collection", validate: func() error { + config := validPack + config.CollectionDirectory = "" + return config.Validate() + }}, + {name: "pack archive", validate: func() error { + config := validPack + config.ArchivePath = "" + return config.Validate() + }}, + {name: "pack recipient", validate: func() error { + config := validPack + config.Recipient = retarchive.PublicKey{} + return config.Validate() + }}, + {name: "unpack archive", validate: func() error { + config := validUnpack + config.ArchivePath = "" + return config.Validate() + }}, + {name: "unpack output", validate: func() error { + config := validUnpack + config.OutputDirectory = "" + return config.Validate() + }}, + {name: "unpack identity", validate: func() error { + config := validUnpack + config.Identity = retarchive.PrivateKey{} + return config.Validate() + }}, + {name: "keygen private", validate: func() error { + config := validKeygen + config.PrivateKeyPath = "" + return config.Validate() + }}, + {name: "keygen public", validate: func() error { + config := validKeygen + config.PublicKeyPath = "" + return config.Validate() + }}, + {name: "keygen same destination", validate: func() error { + config := validKeygen + config.PublicKeyPath = config.PrivateKeyPath + return config.Validate() + }}, + } { + t.Run(test.name, func(t *testing.T) { + require.ErrorIs(t, test.validate(), ErrInvalidConfig) + }) + } + + require.NoError(t, validPack.Validate()) + require.NoError(t, validUnpack.Validate()) + require.NoError(t, validKeygen.Validate()) +} + +func TestPackAndUnpackDelegateAllFormatsAndObserver(t *testing.T) { + for _, outputs := range []struct { + name string + jsonl, parquet bool + }{ + {name: "JSONL only", jsonl: true}, + {name: "Parquet only", parquet: true}, + {name: "dual", jsonl: true, parquet: true}, + } { + t.Run(outputs.name, func(t *testing.T) { + source := writeRootArchiveTestCollection(t, outputs.jsonl, outputs.parquet) + public, private, err := retarchive.GenerateKeyPair() + require.NoError(t, err) + archivePath := filepath.Join(t.TempDir(), "collection.ret") + output := filepath.Join(t.TempDir(), "collection") + var events []observe.Event + observer := observe.ObserverFunc(func(_ context.Context, event observe.Event) { + events = append(events, event) + }) + + require.NoError(t, Pack(context.Background(), PackConfig{ + CollectionDirectory: source, + ArchivePath: archivePath, + Recipient: public, + Observer: observer, + })) + require.NoError(t, Unpack(context.Background(), UnpackConfig{ + ArchivePath: archivePath, + OutputDirectory: output, + Identity: private, + Observer: observer, + })) + + _, err = collection.Verify(context.Background(), output, nil) + require.NoError(t, err) + requireRootArchiveTestTreesEqual(t, source, output) + require.IsType(t, observe.OperationStarted{}, events[0]) + require.Equal(t, "pack", events[0].(observe.OperationStarted).Operation) + require.IsType(t, observe.OperationCompleted{}, events[len(events)-1]) + require.Equal(t, "unpack", events[len(events)-1].(observe.OperationCompleted).Operation) + require.NoError(t, events[len(events)-1].(observe.OperationCompleted).Err) + }) + } +} + +func TestPackFullyVerifiesParquetBeforePublishing(t *testing.T) { + source := writeRootArchiveTestCollection(t, true, true) + manifest, err := collection.Read(source) + require.NoError(t, err) + parquetPath := manifest.Graphs[0].NodeShards[0].Parquet.Path + require.NoError(t, os.WriteFile( + filepath.Join(source, filepath.FromSlash(parquetPath)), + []byte("corrupt"), + 0o600, + )) + public, _, err := retarchive.GenerateKeyPair() + require.NoError(t, err) + archivePath := filepath.Join(t.TempDir(), "collection.ret") + + err = Pack(context.Background(), PackConfig{ + CollectionDirectory: source, + ArchivePath: archivePath, + Recipient: public, + }) + + require.Error(t, err) + require.NoFileExists(t, archivePath) +} + +func TestKeygenPublishesPrivateThenPublicExclusively(t *testing.T) { + t.Run("success", func(t *testing.T) { + root := t.TempDir() + privatePath := filepath.Join(root, "private.json") + publicPath := filepath.Join(root, "public.json") + + require.NoError(t, Keygen(KeygenConfig{ + PrivateKeyPath: privatePath, + PublicKeyPath: publicPath, + })) + + public, err := retarchive.ReadPublicKey(publicPath) + require.NoError(t, err) + private, err := retarchive.ReadPrivateKey(privatePath) + require.NoError(t, err) + + roundTripArchive := filepath.Join(root, "probe.ret") + collectionRoot := writeRootArchiveTestCollection(t, true, false) + require.NoError(t, retarchive.Create(context.Background(), retarchive.CreateConfig{ + CollectionDirectory: collectionRoot, + ArchivePath: roundTripArchive, + Recipient: public, + })) + require.NoError(t, retarchive.Extract(context.Background(), retarchive.ExtractConfig{ + ArchivePath: roundTripArchive, + OutputDirectory: filepath.Join(root, "round-trip"), + Identity: private, + })) + }) + + t.Run("private exists", func(t *testing.T) { + root := t.TempDir() + privatePath := filepath.Join(root, "private.json") + publicPath := filepath.Join(root, "public.json") + require.NoError(t, os.WriteFile(privatePath, []byte("preserve-private"), 0o600)) + + err := Keygen(KeygenConfig{PrivateKeyPath: privatePath, PublicKeyPath: publicPath}) + + require.Error(t, err) + require.Equal(t, []byte("preserve-private"), mustReadRootArchiveTestFile(t, privatePath)) + require.NoFileExists(t, publicPath) + }) + + t.Run("public exists rolls back private", func(t *testing.T) { + root := t.TempDir() + privatePath := filepath.Join(root, "private.json") + publicPath := filepath.Join(root, "public.json") + require.NoError(t, os.WriteFile(publicPath, []byte("preserve-public"), 0o600)) + + err := Keygen(KeygenConfig{PrivateKeyPath: privatePath, PublicKeyPath: publicPath}) + + require.Error(t, err) + require.NoFileExists(t, privatePath) + require.Equal(t, []byte("preserve-public"), mustReadRootArchiveTestFile(t, publicPath)) + }) +} + +func TestKeygenIsIndependentOfArchivePublicationPlatform(t *testing.T) { + root := t.TempDir() + privatePath := filepath.Join(root, "private.json") + publicPath := filepath.Join(root, "public.json") + + err := Keygen(KeygenConfig{ + PrivateKeyPath: privatePath, + PublicKeyPath: publicPath, + }) + + require.NoError(t, err) + _, err = retarchive.ReadPrivateKey(privatePath) + require.NoError(t, err) + _, err = retarchive.ReadPublicKey(publicPath) + require.NoError(t, err) +} + +func writeRootArchiveTestCollection(t *testing.T, withJSONL, withParquet bool) string { + t.Helper() + root := t.TempDir() + nodes := []entity.Node{{ + SourceID: "node-1", + Kinds: []string{"User"}, + Properties: map[string]any{"name": "Alice"}, + }} + builder := metrics.NewBuilder() + require.NoError(t, builder.ObserveNode(nodes[0])) + graph := collection.Graph{ + Name: "example", + NodeCount: 1, + KindCatalog: []string{"User"}, + NodeShards: []collection.NodeShard{{Index: 1, Count: 1, LastSourceID: 1}}, + Metrics: builder.Finalize(), + } + outputs := collection.OutputConfig{} + + if withJSONL { + config := jsonl.Config{Codec: jsonl.CodecNone} + outputs.JSONL = &collection.JSONLOutput{ + SchemaVersion: jsonl.SchemaVersion, + Codec: string(config.Codec), + Level: config.Level, + } + relative := collection.NodeJSONLPath(graph.Name, 1, config.Codec) + temporary := filepath.Join(root, "nodes.jsonl.tmp") + artifact, err := writeJSONLNodeFile(temporary, relative, config, nodes) + require.NoError(t, err) + installRootArchiveTestArtifact(t, root, temporary, relative) + graph.NodeShards[0].JSONL = &artifact + } + if withParquet { + config := parquet.Config{} + outputs.Parquet = &collection.ParquetOutput{SchemaVersion: parquet.SchemaVersion} + relative := collection.NodeParquetPath(graph.Name, 1) + temporary := filepath.Join(root, "nodes.parquet.tmp") + artifact, err := writeParquetNodeFile(temporary, relative, config, nodes) + require.NoError(t, err) + installRootArchiveTestArtifact(t, root, temporary, relative) + graph.NodeShards[0].Parquet = &artifact + } + require.NoError(t, collection.Write(root, collection.Manifest{ + Format: collection.Format, + CreatedAt: time.Date(2026, time.July, 28, 12, 0, 0, 0, time.UTC), + Outputs: outputs, + Graphs: []collection.Graph{graph}, + })) + return root +} + +func installRootArchiveTestArtifact(t *testing.T, root, temporary, relative string) { + t.Helper() + final := filepath.Join(root, filepath.FromSlash(relative)) + require.NoError(t, os.MkdirAll(filepath.Dir(final), 0o700)) + require.NoError(t, os.Rename(temporary, final)) +} + +func requireRootArchiveTestTreesEqual(t *testing.T, expected, actual string) { + t.Helper() + require.NoError(t, filepath.WalkDir(expected, func(candidate string, entry os.DirEntry, err error) error { + require.NoError(t, err) + if entry.IsDir() { + return nil + } + relative, err := filepath.Rel(expected, candidate) + require.NoError(t, err) + require.Equal( + t, + mustReadRootArchiveTestFile(t, candidate), + mustReadRootArchiveTestFile(t, filepath.Join(actual, relative)), + ) + return nil + })) +} + +func mustReadRootArchiveTestFile(t *testing.T, path string) []byte { + t.Helper() + value, err := os.ReadFile(path) + require.NoError(t, err) + return value +} diff --git a/ret/artifact_io.go b/ret/artifact_io.go new file mode 100644 index 00000000..07f8a4db --- /dev/null +++ b/ret/artifact_io.go @@ -0,0 +1,91 @@ +package ret + +import ( + "errors" + "fmt" + "io" + "os" + + "github.com/specterops/dawgs/ret/collection" + "github.com/specterops/dawgs/ret/entity" + "github.com/specterops/dawgs/ret/jsonl" + "github.com/specterops/dawgs/ret/parquet" +) + +type artifactWriter[E, A any] interface { + Push([]E) error + Close() error + Result() (A, error) +} + +func writeArtifactFile[E, A any]( + temporary string, + values []E, + newWriter func(io.Writer) (artifactWriter[E, A], error), +) (A, error) { + var zero A + file, err := os.OpenFile(temporary, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o600) + if err != nil { + return zero, fmt.Errorf("open temporary artifact: %w", err) + } + writer, err := newWriter(file) + if err != nil { + return zero, errors.Join(err, file.Close()) + } + pushErr := writer.Push(values) + closeWriterErr := writer.Close() + var artifact A + var resultErr error + if pushErr == nil && closeWriterErr == nil { + artifact, resultErr = writer.Result() + } + closeFileErr := file.Close() + if err := errors.Join(pushErr, closeWriterErr, resultErr, closeFileErr); err != nil { + return zero, err + } + return artifact, nil +} + +func writeJSONLNodeFile(temporary, relative string, config jsonl.Config, values []entity.Node) (collection.JSONLArtifact, error) { + artifact, err := writeArtifactFile(temporary, values, func(output io.Writer) (artifactWriter[entity.Node, jsonl.Artifact], error) { + writer, err := jsonl.NewNodeWriter(output, config) + return &writer, err + }) + if err != nil { + return collection.JSONLArtifact{}, err + } + return collection.JSONLArtifact{Path: relative, Artifact: artifact}, nil +} + +func writeJSONLRelationshipFile(temporary, relative string, config jsonl.Config, values []entity.Relationship) (collection.JSONLArtifact, error) { + artifact, err := writeArtifactFile(temporary, values, func(output io.Writer) (artifactWriter[entity.Relationship, jsonl.Artifact], error) { + writer, err := jsonl.NewRelationshipWriter(output, config) + return &writer, err + }) + if err != nil { + return collection.JSONLArtifact{}, err + } + return collection.JSONLArtifact{Path: relative, Artifact: artifact}, nil +} + +func writeParquetNodeFile(temporary, relative string, config parquet.Config, values []entity.Node) (collection.ParquetArtifact, error) { + artifact, err := writeArtifactFile(temporary, values, func(output io.Writer) (artifactWriter[entity.Node, parquet.Artifact], error) { + writer, err := parquet.NewNodeWriter(output, config) + return &writer, err + }) + if err != nil { + return collection.ParquetArtifact{}, err + } + return collection.ParquetArtifact{Path: relative, Artifact: artifact}, nil +} + +func writeParquetRelationshipFile(temporary, relative string, config parquet.Config, values []entity.Relationship) (collection.ParquetArtifact, error) { + artifact, err := writeArtifactFile(temporary, values, func(output io.Writer) (artifactWriter[entity.Relationship, parquet.Artifact], error) { + writer, err := parquet.NewRelationshipWriter(output, config) + return &writer, err + }) + if err != nil { + return collection.ParquetArtifact{}, err + } + return collection.ParquetArtifact{Path: relative, Artifact: artifact}, nil +} diff --git a/ret/checkpoint/identity.go b/ret/checkpoint/identity.go new file mode 100644 index 00000000..712231e7 --- /dev/null +++ b/ret/checkpoint/identity.go @@ -0,0 +1,100 @@ +package checkpoint + +import ( + "errors" + "fmt" + "slices" + "strings" +) + +func ValidateIdentity(expected, actual Identity) error { + var differences []string + if !slices.Equal(expected.Graphs, actual.Graphs) { + differences = append(differences, fmt.Sprintf( + "ordered graph names differ: got %q want %q", + actual.Graphs, + expected.Graphs, + )) + } + if expected.EntityBatchSize != actual.EntityBatchSize { + differences = append(differences, fmt.Sprintf( + "entity batch size: got %d want %d", + actual.EntityBatchSize, + expected.EntityBatchSize, + )) + } + if expected.ShardSize != actual.ShardSize { + differences = append(differences, fmt.Sprintf( + "shard size: got %d want %d", + actual.ShardSize, + expected.ShardSize, + )) + } + if expected.JSONLEnabled != actual.JSONLEnabled { + differences = append(differences, fmt.Sprintf( + "JSONL enabled: got %t want %t", + actual.JSONLEnabled, + expected.JSONLEnabled, + )) + } + if expected.JSONLCodec != actual.JSONLCodec { + differences = append(differences, fmt.Sprintf( + "JSONL codec: got %q want %q", + actual.JSONLCodec, + expected.JSONLCodec, + )) + } + if expected.JSONLLevel != actual.JSONLLevel { + differences = append(differences, fmt.Sprintf( + "JSONL level: got %d want %d", + actual.JSONLLevel, + expected.JSONLLevel, + )) + } + if expected.ParquetEnabled != actual.ParquetEnabled { + differences = append(differences, fmt.Sprintf( + "Parquet enabled: got %t want %t", + actual.ParquetEnabled, + expected.ParquetEnabled, + )) + } + if expected.JSONLSchemaVersion != actual.JSONLSchemaVersion { + differences = append(differences, fmt.Sprintf( + "JSONL schema version: got %q want %q", + actual.JSONLSchemaVersion, + expected.JSONLSchemaVersion, + )) + } + if expected.ParquetSchemaVersion != actual.ParquetSchemaVersion { + differences = append(differences, fmt.Sprintf( + "Parquet schema version: got %q want %q", + actual.ParquetSchemaVersion, + expected.ParquetSchemaVersion, + )) + } + if expected.ScrubEnabled != actual.ScrubEnabled { + differences = append(differences, fmt.Sprintf( + "scrub enabled: got %t want %t", + actual.ScrubEnabled, + expected.ScrubEnabled, + )) + } + if expected.ScrubRulesFingerprint != actual.ScrubRulesFingerprint { + differences = append(differences, fmt.Sprintf( + "scrub rules fingerprint: got %q want %q", + actual.ScrubRulesFingerprint, + expected.ScrubRulesFingerprint, + )) + } + if expected.ScrubSaltFingerprint != actual.ScrubSaltFingerprint { + differences = append(differences, fmt.Sprintf( + "scrub salt fingerprint: got %q want %q", + actual.ScrubSaltFingerprint, + expected.ScrubSaltFingerprint, + )) + } + if len(differences) != 0 { + return errors.New(strings.Join(differences, "; ")) + } + return nil +} diff --git a/ret/checkpoint/identity_test.go b/ret/checkpoint/identity_test.go new file mode 100644 index 00000000..f4b75150 --- /dev/null +++ b/ret/checkpoint/identity_test.go @@ -0,0 +1,103 @@ +package checkpoint + +import ( + "strings" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestValidateIdentityAcceptsIdenticalIdentity(t *testing.T) { + identity := fixtureIdentity() + + require.NoError(t, ValidateIdentity(identity, identity)) +} + +func TestValidateIdentityReportsEveryFieldMismatchInStableOrder(t *testing.T) { + expected := fixtureIdentity() + actual := Identity{ + Graphs: []string{"beta", "alpha"}, + EntityBatchSize: 101, + ShardSize: 3, + JSONLEnabled: false, + JSONLCodec: "gzip", + JSONLLevel: 7, + ParquetEnabled: false, + JSONLSchemaVersion: "ret-jsonl-v0", + ParquetSchemaVersion: "ret-parquet-v0", + ScrubEnabled: false, + ScrubRulesFingerprint: strings.Repeat("c", 64), + ScrubSaltFingerprint: strings.Repeat("d", 64), + } + + err := ValidateIdentity(expected, actual) + require.EqualError(t, err, + `ordered graph names differ: got ["beta" "alpha"] want ["alpha" "beta"]; `+ + `entity batch size: got 101 want 100; `+ + `shard size: got 3 want 2; `+ + `JSONL enabled: got false want true; `+ + `JSONL codec: got "gzip" want "zstd"; `+ + `JSONL level: got 7 want 3; `+ + `Parquet enabled: got false want true; `+ + `JSONL schema version: got "ret-jsonl-v0" want "retriever-jsonl-v1"; `+ + `Parquet schema version: got "ret-parquet-v0" want "ret-parquet-v1"; `+ + `scrub enabled: got false want true; `+ + `scrub rules fingerprint: got "`+strings.Repeat("c", 64)+`" want "`+strings.Repeat("a", 64)+`"; `+ + `scrub salt fingerprint: got "`+strings.Repeat("d", 64)+`" want "`+strings.Repeat("b", 64)+`"`, + ) +} + +func TestValidateIdentityDetectsEveryFieldIndividually(t *testing.T) { + tests := []struct { + name string + label string + mutate func(*Identity) + }{ + {name: "ordered graphs", label: "ordered graph names", mutate: func(value *Identity) { + value.Graphs = []string{"beta", "alpha"} + }}, + {name: "entity batch size", label: "entity batch size", mutate: func(value *Identity) { + value.EntityBatchSize++ + }}, + {name: "shard size", label: "shard size", mutate: func(value *Identity) { + value.ShardSize++ + }}, + {name: "JSONL enabled", label: "JSONL enabled", mutate: func(value *Identity) { + value.JSONLEnabled = !value.JSONLEnabled + }}, + {name: "JSONL codec", label: "JSONL codec", mutate: func(value *Identity) { + value.JSONLCodec = "gzip" + }}, + {name: "JSONL level", label: "JSONL level", mutate: func(value *Identity) { + value.JSONLLevel++ + }}, + {name: "Parquet enabled", label: "Parquet enabled", mutate: func(value *Identity) { + value.ParquetEnabled = !value.ParquetEnabled + }}, + {name: "JSONL schema", label: "JSONL schema version", mutate: func(value *Identity) { + value.JSONLSchemaVersion = "changed" + }}, + {name: "Parquet schema", label: "Parquet schema version", mutate: func(value *Identity) { + value.ParquetSchemaVersion = "changed" + }}, + {name: "scrub enabled", label: "scrub enabled", mutate: func(value *Identity) { + value.ScrubEnabled = !value.ScrubEnabled + }}, + {name: "scrub rules", label: "scrub rules fingerprint", mutate: func(value *Identity) { + value.ScrubRulesFingerprint = strings.Repeat("c", 64) + }}, + {name: "scrub salt", label: "scrub salt fingerprint", mutate: func(value *Identity) { + value.ScrubSaltFingerprint = strings.Repeat("d", 64) + }}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + expected := fixtureIdentity() + actual := fixtureIdentity() + test.mutate(&actual) + + require.ErrorContains(t, ValidateIdentity(expected, actual), test.label) + }) + } +} diff --git a/ret/checkpoint/model.go b/ret/checkpoint/model.go new file mode 100644 index 00000000..5c2b6bbb --- /dev/null +++ b/ret/checkpoint/model.go @@ -0,0 +1,60 @@ +// Package checkpoint persists validated resumable dump progress. +package checkpoint + +import ( + "github.com/specterops/dawgs/ret/collection" + "github.com/specterops/dawgs/ret/dawgs" +) + +const ( + Format = "ret-checkpoint-v1" + FileName = ".ret-checkpoint.json" +) + +type Identity struct { + Graphs []string `json:"graphs"` + EntityBatchSize int `json:"entity_batch_size"` + ShardSize int `json:"shard_size"` + JSONLEnabled bool `json:"jsonl_enabled"` + JSONLCodec string `json:"jsonl_codec"` + JSONLLevel int `json:"jsonl_level"` + ParquetEnabled bool `json:"parquet_enabled"` + JSONLSchemaVersion string `json:"jsonl_schema_version"` + ParquetSchemaVersion string `json:"parquet_schema_version"` + ScrubEnabled bool `json:"scrub_enabled"` + ScrubRulesFingerprint string `json:"scrub_rules_fingerprint"` + ScrubSaltFingerprint string `json:"scrub_salt_fingerprint"` +} + +type Phase string + +const ( + PhaseNodes Phase = "nodes" + PhaseRelationships Phase = "relationships" + PhaseComplete Phase = "complete" +) + +type GraphState struct { + Name string `json:"name"` + Snapshot dawgs.Snapshot `json:"snapshot"` + Phase Phase `json:"phase"` + NodeCursor uint64 `json:"node_cursor"` + RelationshipCursor uint64 `json:"relationship_cursor"` + NodeShards []collection.NodeShard `json:"node_shards"` + RelationshipShards []collection.RelationshipShard `json:"relationship_shards"` +} + +type State struct { + Format string `json:"format"` + Identity Identity `json:"identity"` + Graphs []GraphState `json:"graphs"` +} + +// Store persists one dump's checkpoint beneath Root. +// +// A dump directory must have only one active writer. Store does not provide a +// locking protocol; CleanupOrphans may remove staging files left by a crashed +// writer after the published checkpoint has been loaded and validated. +type Store struct { + Root string +} diff --git a/ret/checkpoint/store.go b/ret/checkpoint/store.go new file mode 100644 index 00000000..63e149bc --- /dev/null +++ b/ret/checkpoint/store.go @@ -0,0 +1,886 @@ +package checkpoint + +import ( + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "io" + "io/fs" + "math" + "net/url" + "os" + "path" + "path/filepath" + "strings" + + "github.com/specterops/dawgs/ret/collection" + "github.com/specterops/dawgs/ret/jsonl" + "github.com/specterops/dawgs/ret/parquet" + "github.com/specterops/dawgs/ret/scrub" +) + +const checkpointStagingPrefix = FileName + ".tmp-" + +var ( + checkpointRename = os.Rename + checkpointRemove = os.Remove + checkpointOpenRoot = os.OpenRoot + checkpointPinnedRemove = func(root *os.Root, name string) error { + return root.Remove(name) + } +) + +type orphanRemoval struct { + relative string + basename string + parent *os.Root +} + +func (s Store) Load() (State, bool, error) { + if err := s.validateRoot(); err != nil { + return State{}, true, err + } + file, err := os.Open(filepath.Join(s.Root, FileName)) + if errors.Is(err, os.ErrNotExist) { + return State{}, false, nil + } + if err != nil { + return State{}, true, fmt.Errorf("open checkpoint: %w", err) + } + defer file.Close() + + decoder := json.NewDecoder(file) + decoder.DisallowUnknownFields() + var state State + if err := decoder.Decode(&state); err != nil { + return State{}, true, fmt.Errorf("decode checkpoint: %w", err) + } + var trailing any + if err := decoder.Decode(&trailing); !errors.Is(err, io.EOF) { + if err == nil { + return State{}, true, fmt.Errorf("decode checkpoint: trailing JSON value") + } + return State{}, true, fmt.Errorf("decode checkpoint trailing data: %w", err) + } + if err := validateState(state); err != nil { + return State{}, true, fmt.Errorf("validate checkpoint: %w", err) + } + return state, true, nil +} + +func (s Store) Save(state State) (resultErr error) { + if err := validateState(state); err != nil { + return fmt.Errorf("validate checkpoint: %w", err) + } + if err := s.validateRoot(); err != nil { + return err + } + + final := filepath.Join(s.Root, FileName) + file, err := os.CreateTemp(s.Root, checkpointStagingPrefix+"*") + if err != nil { + return fmt.Errorf("create temporary checkpoint: %w", err) + } + temporary := file.Name() + published := false + defer func() { + joined := []error{resultErr} + if file != nil { + if closeErr := file.Close(); closeErr != nil { + joined = append(joined, fmt.Errorf("cleanup close temporary checkpoint: %w", closeErr)) + } + } + if !published { + if removeErr := checkpointRemove(temporary); removeErr != nil && !errors.Is(removeErr, os.ErrNotExist) { + joined = append(joined, fmt.Errorf("cleanup remove temporary checkpoint: %w", removeErr)) + } + } + resultErr = errors.Join(joined...) + }() + + if err := json.NewEncoder(file).Encode(state); err != nil { + return fmt.Errorf("encode checkpoint: %w", err) + } + if err := file.Sync(); err != nil { + return fmt.Errorf("sync temporary checkpoint: %w", err) + } + if err := file.Close(); err != nil { + file = nil + return fmt.Errorf("close temporary checkpoint: %w", err) + } + file = nil + if err := checkpointRename(temporary, final); err != nil { + return fmt.Errorf("publish checkpoint: %w", err) + } + published = true + return nil +} + +func (s Store) Remove() error { + if err := s.validateRoot(); err != nil { + return err + } + if err := checkpointRemove(filepath.Join(s.Root, FileName)); err != nil && !errors.Is(err, os.ErrNotExist) { + return fmt.Errorf("remove checkpoint: %w", err) + } + return nil +} + +func (s Store) CleanupOrphans(state State) error { + if err := validateState(state); err != nil { + return fmt.Errorf("validate checkpoint before orphan cleanup: %w", err) + } + if err := s.validateRoot(); err != nil { + return err + } + + rootInfo, err := os.Lstat(s.Root) + if err != nil { + return fmt.Errorf("inspect checkpoint collection root: %w", err) + } + if rootInfo.Mode()&os.ModeSymlink != 0 { + return fmt.Errorf("checkpoint collection root is a symbolic link") + } + openRoot, err := checkpointOpenRoot(s.Root) + if err != nil { + return fmt.Errorf("open checkpoint collection root: %w", err) + } + pinnedRootInfo, err := openRoot.Stat(".") + if err != nil { + return errors.Join( + fmt.Errorf("inspect pinned checkpoint collection root: %w", err), + wrapRootCloseError("checkpoint collection root", openRoot.Close()), + ) + } + if !rootInfo.IsDir() || !pinnedRootInfo.IsDir() || !os.SameFile(rootInfo, pinnedRootInfo) { + return errors.Join( + fmt.Errorf("checkpoint collection root changed while opening"), + wrapRootCloseError("checkpoint collection root", openRoot.Close()), + ) + } + + committed, candidates := cleanupPaths(state) + allowedDirectories := cleanupDirectories(state.Identity) + directoryIdentities := make(map[string]fs.FileInfo, len(allowedDirectories)) + var removals []orphanRemoval + walkErr := fs.WalkDir(openRoot.FS(), ".", func(relative string, entry fs.DirEntry, walkErr error) error { + if walkErr != nil { + return walkErr + } + if entry.Type()&os.ModeSymlink != 0 { + return fmt.Errorf("resume entry %q is a symbolic link", relative) + } + if entry.IsDir() { + if _, ok := allowedDirectories[relative]; !ok { + return fmt.Errorf("unknown resume directory %q", relative) + } + info, err := entry.Info() + if err != nil { + return fmt.Errorf("inspect resume directory %q: %w", relative, err) + } + if !info.IsDir() { + return fmt.Errorf("resume directory %q changed during inventory", relative) + } + directoryIdentities[relative] = info + return nil + } + info, err := entry.Info() + if err != nil { + return fmt.Errorf("inspect resume entry %q: %w", relative, err) + } + if !info.Mode().IsRegular() { + return fmt.Errorf("resume entry %q is not a regular file", relative) + } + if relative == FileName { + return nil + } + if _, ok := committed[relative]; ok { + return nil + } + if isCheckpointStaging(relative) || isCleanupCandidate(relative, candidates) { + parentPath := path.Dir(relative) + parentIdentity, ok := directoryIdentities[parentPath] + if !ok { + return fmt.Errorf("resume candidate parent %q was not inventoried", parentPath) + } + parentRoot, err := openRoot.OpenRoot(parentPath) + if err != nil { + return fmt.Errorf("pin resume candidate parent %q: %w", parentPath, err) + } + pinnedParentInfo, err := parentRoot.Stat(".") + if err != nil { + return errors.Join( + fmt.Errorf("inspect pinned resume candidate parent %q: %w", parentPath, err), + wrapRootCloseError("resume candidate parent "+parentPath, parentRoot.Close()), + ) + } + if !pinnedParentInfo.IsDir() || !os.SameFile(parentIdentity, pinnedParentInfo) { + return errors.Join( + fmt.Errorf("resume candidate parent %q changed while pinning", parentPath), + wrapRootCloseError("resume candidate parent "+parentPath, parentRoot.Close()), + ) + } + removals = append(removals, orphanRemoval{ + relative: relative, + basename: path.Base(relative), + parent: parentRoot, + }) + return nil + } + return fmt.Errorf("unknown resume file %q", relative) + }) + if walkErr != nil { + return errors.Join( + fmt.Errorf("inventory checkpoint collection: %w", walkErr), + closeOrphanRemovalRoots(removals), + wrapRootCloseError("checkpoint collection root", openRoot.Close()), + ) + } + + var removeErrs []error + for _, removal := range removals { + if err := checkpointPinnedRemove(removal.parent, removal.basename); err != nil && !errors.Is(err, os.ErrNotExist) { + removeErrs = append(removeErrs, fmt.Errorf("remove uncommitted artifact %q: %w", removal.relative, err)) + } + } + removeErrs = append(removeErrs, closeOrphanRemovalRoots(removals)) + removeErrs = append(removeErrs, wrapRootCloseError("checkpoint collection root", openRoot.Close())) + return errors.Join(removeErrs...) +} + +func closeOrphanRemovalRoots(removals []orphanRemoval) error { + var closeErrs []error + for _, removal := range removals { + closeErrs = append( + closeErrs, + wrapRootCloseError("resume candidate parent "+path.Dir(removal.relative), removal.parent.Close()), + ) + } + return errors.Join(closeErrs...) +} + +func wrapRootCloseError(name string, err error) error { + if err == nil { + return nil + } + return fmt.Errorf("close %s: %w", name, err) +} + +func (s Store) validateRoot() error { + if s.Root == "" { + return fmt.Errorf("checkpoint root is empty") + } + return nil +} + +func validateState(state State) error { + if state.Format != Format { + return fmt.Errorf("checkpoint format %q does not match %q", state.Format, Format) + } + if err := validateIdentity(state.Identity); err != nil { + return fmt.Errorf("identity: %w", err) + } + if len(state.Graphs) == 0 { + return fmt.Errorf("checkpoint must contain at least one graph state") + } + if len(state.Graphs) > len(state.Identity.Graphs) { + return fmt.Errorf("checkpoint graph state count %d exceeds identity graph count %d", len(state.Graphs), len(state.Identity.Graphs)) + } + + seen := make(map[string]struct{}, len(state.Graphs)) + for index, graph := range state.Graphs { + if _, ok := seen[graph.Name]; ok { + return fmt.Errorf("duplicate graph state %q", graph.Name) + } + seen[graph.Name] = struct{}{} + if graph.Name != state.Identity.Graphs[index] { + return fmt.Errorf( + "graph state order at position %d: got %q want %q", + index+1, + graph.Name, + state.Identity.Graphs[index], + ) + } + } + + artifactPaths := make(map[string]struct{}) + for index, graph := range state.Graphs { + if index < len(state.Graphs)-1 && graph.Phase != PhaseComplete { + return fmt.Errorf("graph %q phase progression: only the last graph state may be incomplete", graph.Name) + } + if err := validateGraphState(graph, state.Identity, artifactPaths); err != nil { + return fmt.Errorf("graph %q: %w", graph.Name, err) + } + } + return nil +} + +func validateIdentity(identity Identity) error { + if len(identity.Graphs) == 0 { + return fmt.Errorf("at least one graph is required") + } + seen := make(map[string]struct{}, len(identity.Graphs)) + for index, graph := range identity.Graphs { + if err := validateGraphName(graph); err != nil { + return fmt.Errorf("graph %d must be a safe name: %w", index+1, err) + } + if _, ok := seen[graph]; ok { + return fmt.Errorf("duplicate graph name %q", graph) + } + seen[graph] = struct{}{} + } + if identity.EntityBatchSize <= 0 { + return fmt.Errorf("entity batch size must be positive") + } + if identity.ShardSize <= 0 { + return fmt.Errorf("shard size must be positive") + } + if !identity.JSONLEnabled && !identity.ParquetEnabled { + return fmt.Errorf("at least one output must be enabled") + } + if identity.JSONLSchemaVersion != jsonl.SchemaVersion { + return fmt.Errorf("JSONL schema %q does not match %q", identity.JSONLSchemaVersion, jsonl.SchemaVersion) + } + if identity.ParquetSchemaVersion != parquet.SchemaVersion { + return fmt.Errorf("Parquet schema %q does not match %q", identity.ParquetSchemaVersion, parquet.SchemaVersion) + } + if identity.JSONLEnabled { + if err := (jsonl.Config{ + Codec: jsonl.Codec(identity.JSONLCodec), + Level: identity.JSONLLevel, + }).Validate(); err != nil { + return fmt.Errorf("JSONL output: %w", err) + } + } + if identity.ScrubEnabled { + if !isLowerHexDigest(identity.ScrubRulesFingerprint) || !isLowerHexDigest(identity.ScrubSaltFingerprint) { + return fmt.Errorf("enabled scrub fingerprints must be 64 lowercase hexadecimal characters") + } + } else if identity.ScrubRulesFingerprint != "" || identity.ScrubSaltFingerprint != "" { + return fmt.Errorf("disabled scrub must not contain fingerprints") + } + return nil +} + +func validateGraphName(name string) error { + if strings.TrimSpace(name) == "" { + return fmt.Errorf("name is empty") + } + if name == "." || name == ".." || path.Clean(name) != name { + return fmt.Errorf("%q is not a clean path segment", name) + } + if strings.ContainsAny(name, `/\`) || strings.ContainsRune(name, '\x00') { + return fmt.Errorf("%q is not a single path segment", name) + } + return nil +} + +func validateGraphState(graph GraphState, identity Identity, artifactPaths map[string]struct{}) error { + if graph.Snapshot.NodeCount < 0 || graph.Snapshot.RelationshipCount < 0 { + return fmt.Errorf( + "snapshot counts must be nonnegative: nodes=%d relationships=%d", + graph.Snapshot.NodeCount, + graph.Snapshot.RelationshipCount, + ) + } + nodeTotal, nodeCursor, err := validateNodeShards(graph.Name, graph.NodeShards, identity, artifactPaths) + if err != nil { + return err + } + relationshipTotal, relationshipCursor, err := validateRelationshipShards( + graph.Name, + graph.RelationshipShards, + identity, + artifactPaths, + ) + if err != nil { + return err + } + if graph.NodeCursor != nodeCursor { + return fmt.Errorf("node cursor %d does not match last committed node shard cursor %d", graph.NodeCursor, nodeCursor) + } + if graph.RelationshipCursor != relationshipCursor { + return fmt.Errorf( + "relationship cursor %d does not match last committed relationship shard cursor %d", + graph.RelationshipCursor, + relationshipCursor, + ) + } + if nodeTotal > graph.Snapshot.NodeCount || relationshipTotal > graph.Snapshot.RelationshipCount { + return fmt.Errorf( + "committed shard totals exceed snapshot: nodes=%d/%d relationships=%d/%d", + nodeTotal, + graph.Snapshot.NodeCount, + relationshipTotal, + graph.Snapshot.RelationshipCount, + ) + } + if len(graph.NodeShards) != 0 && + nodeTotal < graph.Snapshot.NodeCount && + graph.NodeShards[len(graph.NodeShards)-1].Count != int64(identity.ShardSize) { + return fmt.Errorf( + "partial node shard is only legal at the snapshot boundary: got %d want %d", + graph.NodeShards[len(graph.NodeShards)-1].Count, + identity.ShardSize, + ) + } + if len(graph.RelationshipShards) != 0 && + relationshipTotal < graph.Snapshot.RelationshipCount && + graph.RelationshipShards[len(graph.RelationshipShards)-1].Count != int64(identity.ShardSize) { + return fmt.Errorf( + "partial relationship shard is only legal at the snapshot boundary: got %d want %d", + graph.RelationshipShards[len(graph.RelationshipShards)-1].Count, + identity.ShardSize, + ) + } + switch graph.Phase { + case PhaseNodes: + if relationshipTotal != 0 || graph.RelationshipCursor != 0 { + return fmt.Errorf("nodes phase must not contain committed relationship progress") + } + case PhaseRelationships: + if nodeTotal != graph.Snapshot.NodeCount { + return fmt.Errorf( + "relationships phase requires all snapshot nodes committed: got %d want %d", + nodeTotal, + graph.Snapshot.NodeCount, + ) + } + case PhaseComplete: + if nodeTotal != graph.Snapshot.NodeCount || relationshipTotal != graph.Snapshot.RelationshipCount { + return fmt.Errorf( + "complete phase requires snapshot totals: nodes=%d/%d relationships=%d/%d", + nodeTotal, + graph.Snapshot.NodeCount, + relationshipTotal, + graph.Snapshot.RelationshipCount, + ) + } + default: + return fmt.Errorf("unsupported phase %q", graph.Phase) + } + return nil +} + +func validateNodeShards( + graph string, + shards []collection.NodeShard, + identity Identity, + artifactPaths map[string]struct{}, +) (int64, uint64, error) { + var total int64 + var cursor uint64 + for offset, shard := range shards { + if err := validateLogicalShard( + "node", + offset, + shard.Index, + shard.Count, + shard.LastSourceID, + cursor, + shard.ScrubCounts, + shard.JSONL != nil, + shard.Parquet != nil, + identity, + ); err != nil { + return 0, 0, err + } + if offset < len(shards)-1 && shard.Count != int64(identity.ShardSize) { + return 0, 0, fmt.Errorf( + "node shard %d count %d does not match configured shard size %d", + shard.Index, + shard.Count, + identity.ShardSize, + ) + } + if shard.JSONL != nil { + expected := collection.NodeJSONLPath(graph, shard.Index, jsonl.Codec(identity.JSONLCodec)) + if err := validateJSONLArtifact( + "node", + shard.Index, + shard.Count, + shard.JSONL.SchemaVersion, + shard.JSONL.Path, + string(shard.JSONL.Codec), + shard.JSONL.SHA256, + shard.JSONL.Level, + shard.JSONL.Count, + shard.JSONL.UncompressedBytes, + shard.JSONL.StoredBytes, + expected, + identity, + artifactPaths, + ); err != nil { + return 0, 0, err + } + } + if shard.Parquet != nil { + expected := collection.NodeParquetPath(graph, shard.Index) + if err := validateParquetArtifact( + "node", + shard.Index, + shard.Count, + shard.Parquet.SchemaVersion, + shard.Parquet.Path, + shard.Parquet.SHA256, + shard.Parquet.Count, + shard.Parquet.StoredBytes, + expected, + identity, + artifactPaths, + ); err != nil { + return 0, 0, err + } + } + if total > math.MaxInt64-shard.Count { + return 0, 0, fmt.Errorf("node shard total overflows int64") + } + total += shard.Count + cursor = shard.LastSourceID + } + return total, cursor, nil +} + +func validateRelationshipShards( + graph string, + shards []collection.RelationshipShard, + identity Identity, + artifactPaths map[string]struct{}, +) (int64, uint64, error) { + var total int64 + var cursor uint64 + for offset, shard := range shards { + if err := validateLogicalShard( + "relationship", + offset, + shard.Index, + shard.Count, + shard.LastSourceID, + cursor, + shard.ScrubCounts, + shard.JSONL != nil, + shard.Parquet != nil, + identity, + ); err != nil { + return 0, 0, err + } + if offset < len(shards)-1 && shard.Count != int64(identity.ShardSize) { + return 0, 0, fmt.Errorf( + "relationship shard %d count %d does not match configured shard size %d", + shard.Index, + shard.Count, + identity.ShardSize, + ) + } + if shard.JSONL != nil { + expected := collection.RelationshipJSONLPath(graph, shard.Index, jsonl.Codec(identity.JSONLCodec)) + if err := validateJSONLArtifact( + "relationship", + shard.Index, + shard.Count, + shard.JSONL.SchemaVersion, + shard.JSONL.Path, + string(shard.JSONL.Codec), + shard.JSONL.SHA256, + shard.JSONL.Level, + shard.JSONL.Count, + shard.JSONL.UncompressedBytes, + shard.JSONL.StoredBytes, + expected, + identity, + artifactPaths, + ); err != nil { + return 0, 0, err + } + } + if shard.Parquet != nil { + expected := collection.RelationshipParquetPath(graph, shard.Index) + if err := validateParquetArtifact( + "relationship", + shard.Index, + shard.Count, + shard.Parquet.SchemaVersion, + shard.Parquet.Path, + shard.Parquet.SHA256, + shard.Parquet.Count, + shard.Parquet.StoredBytes, + expected, + identity, + artifactPaths, + ); err != nil { + return 0, 0, err + } + } + if total > math.MaxInt64-shard.Count { + return 0, 0, fmt.Errorf("relationship shard total overflows int64") + } + total += shard.Count + cursor = shard.LastSourceID + } + return total, cursor, nil +} + +func validateLogicalShard( + entityType string, + offset, index int, + count int64, + cursor, previousCursor uint64, + counts scrub.ActionCounts, + hasJSONL, hasParquet bool, + identity Identity, +) error { + if index != offset+1 { + return fmt.Errorf("%s shard index: got %d want %d", entityType, index, offset+1) + } + if count <= 0 { + return fmt.Errorf("%s shard %d count must be positive", entityType, index) + } + if count > int64(identity.ShardSize) { + return fmt.Errorf( + "%s shard %d count %d exceeds configured shard size %d", + entityType, + index, + count, + identity.ShardSize, + ) + } + if cursor == 0 { + return fmt.Errorf("%s shard %d last source ID must be nonzero", entityType, index) + } + if previousCursor != 0 && cursor <= previousCursor { + return fmt.Errorf( + "%s shard %d last source ID %d does not increase after %d", + entityType, + index, + cursor, + previousCursor, + ) + } + if hasJSONL != identity.JSONLEnabled || hasParquet != identity.ParquetEnabled { + return fmt.Errorf("%s shard %d output mismatch with checkpoint identity", entityType, index) + } + if !identity.ScrubEnabled && !counts.IsZero() { + return fmt.Errorf("%s shard %d has scrub counts while scrubbing is disabled", entityType, index) + } + if counts.Preserve < 0 { + return fmt.Errorf("%s shard %d has invalid scrub count %q=%d", entityType, index, "preserve", counts.Preserve) + } + if counts.Pseudonymize < 0 { + return fmt.Errorf("%s shard %d has invalid scrub count %q=%d", entityType, index, "pseudonymize", counts.Pseudonymize) + } + if counts.Redact < 0 { + return fmt.Errorf("%s shard %d has invalid scrub count %q=%d", entityType, index, "redact", counts.Redact) + } + if counts.ShiftTimestamp < 0 { + return fmt.Errorf("%s shard %d has invalid scrub count %q=%d", entityType, index, "shift_timestamp", counts.ShiftTimestamp) + } + return nil +} + +func validateJSONLArtifact( + entityType string, + index int, + shardCount int64, + schema, artifactPath, codec, digest string, + level int, + count, uncompressedBytes, storedBytes int64, + expectedPath string, + identity Identity, + artifactPaths map[string]struct{}, +) error { + prefix := fmt.Sprintf("%s shard %d JSONL", entityType, index) + if schema != identity.JSONLSchemaVersion { + return fmt.Errorf("%s schema %q does not match %q", prefix, schema, identity.JSONLSchemaVersion) + } + if codec != identity.JSONLCodec { + return fmt.Errorf("%s codec %q does not match %q", prefix, codec, identity.JSONLCodec) + } + if level != identity.JSONLLevel { + return fmt.Errorf("%s level %d does not match %d", prefix, level, identity.JSONLLevel) + } + if count != shardCount { + return fmt.Errorf("%s count %d does not match shard count %d", prefix, count, shardCount) + } + if !isLowerHexDigest(digest) { + return fmt.Errorf("%s SHA-256 must be 64 lowercase hexadecimal characters", prefix) + } + if uncompressedBytes <= 0 { + return fmt.Errorf("%s uncompressed bytes must be positive", prefix) + } + if storedBytes <= 0 { + return fmt.Errorf("%s stored bytes must be positive", prefix) + } + return validateArtifactPath(prefix, artifactPath, expectedPath, artifactPaths) +} + +func validateParquetArtifact( + entityType string, + index int, + shardCount int64, + schema, artifactPath, digest string, + count, storedBytes int64, + expectedPath string, + identity Identity, + artifactPaths map[string]struct{}, +) error { + prefix := fmt.Sprintf("%s shard %d Parquet", entityType, index) + if schema != identity.ParquetSchemaVersion { + return fmt.Errorf("%s schema %q does not match %q", prefix, schema, identity.ParquetSchemaVersion) + } + if count != shardCount { + return fmt.Errorf("%s count %d does not match shard count %d", prefix, count, shardCount) + } + if !isLowerHexDigest(digest) { + return fmt.Errorf("%s SHA-256 must be 64 lowercase hexadecimal characters", prefix) + } + if storedBytes <= 0 { + return fmt.Errorf("%s stored bytes must be positive", prefix) + } + return validateArtifactPath(prefix, artifactPath, expectedPath, artifactPaths) +} + +func validateArtifactPath(prefix, artifactPath, expectedPath string, paths map[string]struct{}) error { + if _, err := collection.SafeJoin(".", artifactPath); err != nil { + return fmt.Errorf("%s path: %w", prefix, err) + } + if artifactPath != expectedPath { + return fmt.Errorf("%s path %q does not match deterministic path %q", prefix, artifactPath, expectedPath) + } + if _, ok := paths[artifactPath]; ok { + return fmt.Errorf("%s path %q is duplicated", prefix, artifactPath) + } + paths[artifactPath] = struct{}{} + return nil +} + +func isLowerHexDigest(value string) bool { + if len(value) != 64 { + return false + } + if _, err := hex.DecodeString(value); err != nil { + return false + } + return strings.ToLower(value) == value +} + +func cleanupPaths(state State) (map[string]struct{}, map[string]struct{}) { + committed := make(map[string]struct{}) + candidates := make(map[string]struct{}) + for _, graph := range state.Graphs { + for _, shard := range graph.NodeShards { + if shard.JSONL != nil { + committed[shard.JSONL.Path] = struct{}{} + } + if shard.Parquet != nil { + committed[shard.Parquet.Path] = struct{}{} + } + } + for _, shard := range graph.RelationshipShards { + if shard.JSONL != nil { + committed[shard.JSONL.Path] = struct{}{} + } + if shard.Parquet != nil { + committed[shard.Parquet.Path] = struct{}{} + } + } + + switch graph.Phase { + case PhaseNodes: + if shardCountNodes(graph.NodeShards) < graph.Snapshot.NodeCount { + index := len(graph.NodeShards) + 1 + addNodeCandidates(candidates, state.Identity, graph.Name, index) + } + case PhaseRelationships: + if shardCountRelationships(graph.RelationshipShards) < graph.Snapshot.RelationshipCount { + index := len(graph.RelationshipShards) + 1 + addRelationshipCandidates(candidates, state.Identity, graph.Name, index) + } + } + } + return committed, candidates +} + +func addNodeCandidates(paths map[string]struct{}, identity Identity, graph string, index int) { + if identity.JSONLEnabled { + paths[collection.NodeJSONLPath(graph, index, jsonl.Codec(identity.JSONLCodec))] = struct{}{} + } + if identity.ParquetEnabled { + paths[collection.NodeParquetPath(graph, index)] = struct{}{} + } +} + +func addRelationshipCandidates(paths map[string]struct{}, identity Identity, graph string, index int) { + if identity.JSONLEnabled { + paths[collection.RelationshipJSONLPath(graph, index, jsonl.Codec(identity.JSONLCodec))] = struct{}{} + } + if identity.ParquetEnabled { + paths[collection.RelationshipParquetPath(graph, index)] = struct{}{} + } +} + +func shardCountNodes(shards []collection.NodeShard) int64 { + var count int64 + for _, shard := range shards { + count += shard.Count + } + return count +} + +func shardCountRelationships(shards []collection.RelationshipShard) int64 { + var count int64 + for _, shard := range shards { + count += shard.Count + } + return count +} + +func cleanupDirectories(identity Identity) map[string]struct{} { + directories := map[string]struct{}{ + ".": {}, + "graphs": {}, + } + for _, graph := range identity.Graphs { + graphDirectory := "graphs/" + url.PathEscape(graph) + directories[graphDirectory] = struct{}{} + directories[graphDirectory+"/nodes"] = struct{}{} + directories[graphDirectory+"/relationships"] = struct{}{} + } + return directories +} + +func isCleanupCandidate(relative string, candidates map[string]struct{}) bool { + if _, ok := candidates[relative]; ok { + return true + } + for candidate := range candidates { + prefix := candidate + ".tmp-" + if strings.HasPrefix(relative, prefix) && isURLSafeNonce(strings.TrimPrefix(relative, prefix)) { + return true + } + } + return false +} + +func isCheckpointStaging(relative string) bool { + if !strings.HasPrefix(relative, checkpointStagingPrefix) { + return false + } + return isURLSafeNonce(strings.TrimPrefix(relative, checkpointStagingPrefix)) +} + +func isURLSafeNonce(value string) bool { + if value == "" { + return false + } + for _, character := range value { + if character >= 'a' && character <= 'z' || + character >= 'A' && character <= 'Z' || + character >= '0' && character <= '9' || + character == '_' || character == '-' { + continue + } + return false + } + return true +} diff --git a/ret/checkpoint/store_nonregular_linux_test.go b/ret/checkpoint/store_nonregular_linux_test.go new file mode 100644 index 00000000..e44c645d --- /dev/null +++ b/ret/checkpoint/store_nonregular_linux_test.go @@ -0,0 +1,27 @@ +//go:build linux + +package checkpoint + +import ( + "path/filepath" + "syscall" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestCleanupOrphansRejectsCheckpointStagingNonRegularFile(t *testing.T) { + root := t.TempDir() + store := Store{Root: root} + require.NoError(t, store.Save(fixtureState())) + loaded, found, err := store.Load() + require.NoError(t, err) + require.True(t, found) + + staging := filepath.Join(root, FileName+".tmp-valid_nonce") + require.NoError(t, syscall.Mkfifo(staging, 0o600)) + + err = store.CleanupOrphans(loaded) + require.ErrorContains(t, err, "not a regular file") + requireArtifactExists(t, root, FileName+".tmp-valid_nonce") +} diff --git a/ret/checkpoint/store_test.go b/ret/checkpoint/store_test.go new file mode 100644 index 00000000..b97fa0ad --- /dev/null +++ b/ret/checkpoint/store_test.go @@ -0,0 +1,897 @@ +package checkpoint + +import ( + "crypto/sha256" + "encoding/json" + "errors" + "fmt" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/specterops/dawgs/ret/collection" + "github.com/specterops/dawgs/ret/dawgs" + "github.com/specterops/dawgs/ret/jsonl" + "github.com/specterops/dawgs/ret/parquet" + "github.com/specterops/dawgs/ret/scrub" + "github.com/stretchr/testify/require" +) + +func fixtureIdentity() Identity { + return Identity{ + Graphs: []string{"alpha", "beta"}, + EntityBatchSize: 100, + ShardSize: 2, + JSONLEnabled: true, + JSONLCodec: string(jsonl.CodecZstd), + JSONLLevel: 3, + ParquetEnabled: true, + JSONLSchemaVersion: jsonl.SchemaVersion, + ParquetSchemaVersion: parquet.SchemaVersion, + ScrubEnabled: true, + ScrubRulesFingerprint: strings.Repeat("a", 64), + ScrubSaltFingerprint: strings.Repeat("b", 64), + } +} + +func fixtureState() State { + identity := fixtureIdentity() + return State{ + Format: Format, + Identity: identity, + Graphs: []GraphState{{ + Name: "alpha", + Snapshot: dawgs.Snapshot{NodeCount: 5, RelationshipCount: 3}, + Phase: PhaseNodes, + NodeCursor: 20, + NodeShards: []collection.NodeShard{ + fixtureNodeShard(identity, "alpha", 1, 2, 20), + }, + }}, + } +} + +func fixtureNodeShard(identity Identity, graph string, index int, count int64, cursor uint64) collection.NodeShard { + shard := collection.NodeShard{ + Index: index, + Count: count, + LastSourceID: cursor, + ScrubCounts: scrub.ActionCounts{Redact: 1}, + } + if identity.JSONLEnabled { + shard.JSONL = &collection.JSONLArtifact{ + Path: collection.NodeJSONLPath(graph, index, jsonl.Codec(identity.JSONLCodec)), + Artifact: jsonl.Artifact{ + SchemaVersion: identity.JSONLSchemaVersion, + Codec: jsonl.Codec(identity.JSONLCodec), + SHA256: strings.Repeat("c", 64), + Level: identity.JSONLLevel, + Count: count, + UncompressedBytes: 100, + StoredBytes: 50, + }, + } + } + if identity.ParquetEnabled { + shard.Parquet = &collection.ParquetArtifact{ + Path: collection.NodeParquetPath(graph, index), + Artifact: parquet.Artifact{ + SchemaVersion: identity.ParquetSchemaVersion, + SHA256: strings.Repeat("d", 64), + Count: count, + StoredBytes: 75, + }, + } + } + return shard +} + +func fixtureRelationshipShard(identity Identity, graph string, index int, count int64, cursor uint64) collection.RelationshipShard { + shard := collection.RelationshipShard{ + Index: index, + Count: count, + LastSourceID: cursor, + ScrubCounts: scrub.ActionCounts{Redact: 1}, + } + if identity.JSONLEnabled { + shard.JSONL = &collection.JSONLArtifact{ + Path: collection.RelationshipJSONLPath(graph, index, jsonl.Codec(identity.JSONLCodec)), + Artifact: jsonl.Artifact{ + SchemaVersion: identity.JSONLSchemaVersion, + Codec: jsonl.Codec(identity.JSONLCodec), + SHA256: strings.Repeat("e", 64), + Level: identity.JSONLLevel, + Count: count, + UncompressedBytes: 120, + StoredBytes: 60, + }, + } + } + if identity.ParquetEnabled { + shard.Parquet = &collection.ParquetArtifact{ + Path: collection.RelationshipParquetPath(graph, index), + Artifact: parquet.Artifact{ + SchemaVersion: identity.ParquetSchemaVersion, + SHA256: strings.Repeat("f", 64), + Count: count, + StoredBytes: 80, + }, + } + } + return shard +} + +func TestStoreRoundTripDoesNotPersistSalt(t *testing.T) { + root := t.TempDir() + state := fixtureState() + privateSalt := "private-salt" + state.Identity.ScrubSaltFingerprint = fmt.Sprintf("%x", sha256.Sum256([]byte(privateSalt))) + store := Store{Root: root} + + require.NoError(t, store.Save(state)) + payload, err := os.ReadFile(filepath.Join(root, FileName)) + require.NoError(t, err) + require.NotContains(t, string(payload), privateSalt) + require.Contains(t, string(payload), state.Identity.ScrubSaltFingerprint) + + loaded, found, err := store.Load() + require.NoError(t, err) + require.True(t, found) + require.Equal(t, state, loaded) +} + +func TestStoreAcceptsEveryLegalPhaseBoundary(t *testing.T) { + tests := []struct { + name string + state func() State + }{ + {name: "nodes before first shard", state: func() State { + state := fixtureState() + state.Graphs[0].NodeCursor = 0 + state.Graphs[0].NodeShards = nil + return state + }}, + {name: "nodes fully committed before transition", state: func() State { + state := fixtureState() + state.Graphs[0].Snapshot.NodeCount = 2 + return state + }}, + {name: "relationships before first shard", state: func() State { + state := fixtureState() + state.Graphs[0].Snapshot.NodeCount = 2 + state.Graphs[0].Phase = PhaseRelationships + return state + }}, + {name: "relationships partially committed", state: func() State { + state := fixtureState() + state.Graphs[0].Snapshot.NodeCount = 2 + state.Graphs[0].Phase = PhaseRelationships + state.Graphs[0].RelationshipCursor = 30 + state.Graphs[0].RelationshipShards = []collection.RelationshipShard{ + fixtureRelationshipShard(state.Identity, "alpha", 1, 2, 30), + } + return state + }}, + {name: "relationships fully committed before transition", state: func() State { + state := fixtureState() + state.Graphs[0].Snapshot = dawgs.Snapshot{NodeCount: 2, RelationshipCount: 1} + state.Graphs[0].Phase = PhaseRelationships + state.Graphs[0].RelationshipCursor = 30 + state.Graphs[0].RelationshipShards = []collection.RelationshipShard{ + fixtureRelationshipShard(state.Identity, "alpha", 1, 1, 30), + } + return state + }}, + {name: "complete", state: func() State { + state := fixtureState() + state.Graphs[0].Snapshot = dawgs.Snapshot{NodeCount: 2, RelationshipCount: 1} + state.Graphs[0].Phase = PhaseComplete + state.Graphs[0].RelationshipCursor = 30 + state.Graphs[0].RelationshipShards = []collection.RelationshipShard{ + fixtureRelationshipShard(state.Identity, "alpha", 1, 1, 30), + } + return state + }}, + {name: "empty graph complete", state: func() State { + state := fixtureState() + state.Graphs[0] = GraphState{ + Name: "alpha", + Snapshot: dawgs.Snapshot{}, + Phase: PhaseComplete, + } + return state + }}, + {name: "completed prefix and active next graph", state: func() State { + state := fixtureState() + state.Graphs[0] = GraphState{ + Name: "alpha", + Snapshot: dawgs.Snapshot{}, + Phase: PhaseComplete, + } + state.Graphs = append(state.Graphs, GraphState{ + Name: "beta", + Snapshot: dawgs.Snapshot{}, + Phase: PhaseNodes, + }) + return state + }}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + store := Store{Root: t.TempDir()} + require.NoError(t, store.Save(test.state())) + }) + } +} + +func TestStoreLoadStrictlyRejectsMalformedCheckpoint(t *testing.T) { + tests := []struct { + name string + payload string + label string + }{ + {name: "invalid JSON", payload: "{", label: "decode checkpoint"}, + {name: "unknown field", payload: `{"format":"ret-checkpoint-v1","unknown":true}`, label: "unknown field"}, + {name: "trailing JSON", payload: `{}` + "\n" + `{}`, label: "trailing JSON value"}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + root := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(root, FileName), []byte(test.payload), 0o600)) + + _, found, err := (Store{Root: root}).Load() + require.True(t, found) + require.ErrorContains(t, err, test.label) + }) + } +} + +func TestStoreLoadMissingIsTheOnlyNotFoundResult(t *testing.T) { + store := Store{Root: t.TempDir()} + + _, found, err := store.Load() + require.NoError(t, err) + require.False(t, found) + + require.NoError(t, os.Mkdir(filepath.Join(store.Root, FileName), 0o700)) + _, found, err = store.Load() + require.True(t, found) + require.Error(t, err) +} + +func TestStoreSaveValidatesBeforeTouchingCurrentCheckpoint(t *testing.T) { + root := t.TempDir() + store := Store{Root: root} + require.NoError(t, store.Save(fixtureState())) + before, err := os.ReadFile(filepath.Join(root, FileName)) + require.NoError(t, err) + + invalid := fixtureState() + invalid.Format = "wrong" + err = store.Save(invalid) + require.ErrorContains(t, err, "format") + + after, readErr := os.ReadFile(filepath.Join(root, FileName)) + require.NoError(t, readErr) + require.Equal(t, before, after) + require.Empty(t, checkpointStagingNames(t, root)) +} + +func TestStoreSavePreservesCurrentCheckpointAndCleansTemporaryOnPublishFailure(t *testing.T) { + root := t.TempDir() + store := Store{Root: root} + require.NoError(t, store.Save(fixtureState())) + before, err := os.ReadFile(filepath.Join(root, FileName)) + require.NoError(t, err) + + originalRename := checkpointRename + checkpointRename = func(_, _ string) error { return errors.New("injected rename failure") } + t.Cleanup(func() { checkpointRename = originalRename }) + + updated := fixtureState() + updated.Graphs[0].Snapshot.NodeCount = 6 + err = store.Save(updated) + require.ErrorContains(t, err, "injected rename failure") + + after, readErr := os.ReadFile(filepath.Join(root, FileName)) + require.NoError(t, readErr) + require.Equal(t, before, after) + require.Empty(t, checkpointStagingNames(t, root)) +} + +func TestStoreSaveJoinsPublishAndTemporaryCleanupFailures(t *testing.T) { + root := t.TempDir() + store := Store{Root: root} + originalRename := checkpointRename + originalRemove := checkpointRemove + checkpointRename = func(_, _ string) error { return errors.New("primary publish failure") } + checkpointRemove = func(string) error { return errors.New("cleanup failure") } + t.Cleanup(func() { + checkpointRename = originalRename + checkpointRemove = originalRemove + }) + + err := store.Save(fixtureState()) + require.ErrorContains(t, err, "primary publish failure") + require.ErrorContains(t, err, "cleanup failure") +} + +func TestStoreUniqueStagingDoesNotBlockSaveAndCleanupRemovesCrashLeftover(t *testing.T) { + root := t.TempDir() + store := Store{Root: root} + initial := fixtureState() + require.NoError(t, store.Save(initial)) + + updated := fixtureState() + updated.Graphs[0].Snapshot.NodeCount = 6 + originalRename := checkpointRename + originalRemove := checkpointRemove + var crashedStage string + checkpointRename = func(oldPath, _ string) error { + crashedStage = oldPath + return errors.New("simulated crash before checkpoint publish") + } + checkpointRemove = func(name string) error { + if name == crashedStage { + return errors.New("simulated process loss before stage cleanup") + } + return originalRemove(name) + } + t.Cleanup(func() { + checkpointRename = originalRename + checkpointRemove = originalRemove + }) + + err := store.Save(updated) + require.ErrorContains(t, err, "simulated crash before checkpoint publish") + require.ErrorContains(t, err, "simulated process loss before stage cleanup") + checkpointRename = originalRename + checkpointRemove = originalRemove + + staging := checkpointStagingNames(t, root) + require.Len(t, staging, 1) + require.NotEqual(t, FileName+".tmp", staging[0]) + nonce := strings.TrimPrefix(staging[0], FileName+".tmp-") + require.NotEmpty(t, nonce) + require.Regexp(t, `^[A-Za-z0-9_-]+$`, nonce) + + require.NoError(t, store.Save(updated), "a unique stage must not be blocked by the crash leftover") + loaded, found, err := store.Load() + require.NoError(t, err) + require.True(t, found) + require.Equal(t, updated, loaded) + require.NoError(t, store.CleanupOrphans(loaded)) + require.Empty(t, checkpointStagingNames(t, root)) +} + +func TestStoreRemoveIsIdempotent(t *testing.T) { + root := t.TempDir() + store := Store{Root: root} + require.NoError(t, store.Remove()) + require.NoError(t, store.Save(fixtureState())) + require.NoError(t, store.Remove()) + require.NoError(t, store.Remove()) + _, err := os.Stat(filepath.Join(root, FileName)) + require.ErrorIs(t, err, os.ErrNotExist) +} + +func TestStoreAcceptsEveryConcreteOutputCombination(t *testing.T) { + tests := []struct { + name string + jsonl bool + parquet bool + }{ + {name: "JSONL only", jsonl: true}, + {name: "Parquet only", parquet: true}, + {name: "dual", jsonl: true, parquet: true}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + state := fixtureStateForOutputs(test.jsonl, test.parquet) + store := Store{Root: t.TempDir()} + require.NoError(t, store.Save(state)) + _, found, err := store.Load() + require.NoError(t, err) + require.True(t, found) + }) + } +} + +func TestStoreAcceptsDisabledScrubWithoutFingerprintsOrCounts(t *testing.T) { + state := fixtureState() + state.Identity.ScrubEnabled = false + state.Identity.ScrubRulesFingerprint = "" + state.Identity.ScrubSaltFingerprint = "" + state.Graphs[0].NodeShards[0].ScrubCounts = scrub.ActionCounts{} + + require.NoError(t, (Store{Root: t.TempDir()}).Save(state)) +} + +func TestStoreAcceptsParquetOnlyWithZeroDisabledJSONLConfig(t *testing.T) { + state := fixtureStateForOutputs(false, true) + state.Identity.JSONLCodec = "" + state.Identity.JSONLLevel = 0 + + require.NoError(t, (Store{Root: t.TempDir()}).Save(state)) +} + +func fixtureStateForOutputs(jsonlEnabled, parquetEnabled bool) State { + state := fixtureState() + state.Identity.JSONLEnabled = jsonlEnabled + state.Identity.ParquetEnabled = parquetEnabled + if !jsonlEnabled { + state.Graphs[0].NodeShards[0].JSONL = nil + } + if !parquetEnabled { + state.Graphs[0].NodeShards[0].Parquet = nil + } + return state +} + +func TestStoreRejectsMalformedState(t *testing.T) { + tests := []struct { + name string + label string + mutate func(*State) + }{ + {name: "format", label: "format", mutate: func(value *State) { value.Format = "wrong" }}, + {name: "empty graph identity", label: "graph", mutate: func(value *State) { value.Identity.Graphs = nil }}, + {name: "duplicate identity graph", label: "duplicate", mutate: func(value *State) { + value.Identity.Graphs = []string{"alpha", "alpha"} + }}, + {name: "unsafe identity graph", label: "safe", mutate: func(value *State) { + value.Identity.Graphs[0] = "../alpha" + value.Graphs[0].Name = "../alpha" + }}, + {name: "nonpositive batch", label: "batch", mutate: func(value *State) { value.Identity.EntityBatchSize = 0 }}, + {name: "nonpositive shard size", label: "shard size", mutate: func(value *State) { value.Identity.ShardSize = 0 }}, + {name: "no outputs", label: "output", mutate: func(value *State) { + value.Identity.JSONLEnabled = false + value.Identity.ParquetEnabled = false + }}, + {name: "invalid JSONL codec", label: "codec", mutate: func(value *State) { value.Identity.JSONLCodec = "zip" }}, + {name: "wrong JSONL schema", label: "JSONL schema", mutate: func(value *State) { + value.Identity.JSONLSchemaVersion = "wrong" + }}, + {name: "wrong Parquet schema", label: "Parquet schema", mutate: func(value *State) { + value.Identity.ParquetSchemaVersion = "wrong" + }}, + {name: "invalid scrub fingerprint", label: "fingerprint", mutate: func(value *State) { + value.Identity.ScrubSaltFingerprint = "private-salt" + }}, + {name: "disabled scrub fingerprints", label: "disabled scrub", mutate: func(value *State) { + value.Identity.ScrubEnabled = false + }}, + {name: "disabled scrub shard counts", label: "scrub counts", mutate: func(value *State) { + value.Identity.ScrubEnabled = false + value.Identity.ScrubRulesFingerprint = "" + value.Identity.ScrubSaltFingerprint = "" + }}, + {name: "negative preserve scrub count", label: "scrub count", mutate: func(value *State) { + value.Graphs[0].NodeShards[0].ScrubCounts.Preserve = -1 + }}, + {name: "negative pseudonymize scrub count", label: "scrub count", mutate: func(value *State) { + value.Graphs[0].NodeShards[0].ScrubCounts.Pseudonymize = -1 + }}, + {name: "negative redact scrub count", label: "scrub count", mutate: func(value *State) { + value.Graphs[0].NodeShards[0].ScrubCounts.Redact = -1 + }}, + {name: "negative shift timestamp scrub count", label: "scrub count", mutate: func(value *State) { + value.Graphs[0].NodeShards[0].ScrubCounts.ShiftTimestamp = -1 + }}, + {name: "graph order", label: "order", mutate: func(value *State) { value.Graphs[0].Name = "beta" }}, + {name: "duplicate graph state", label: "duplicate", mutate: func(value *State) { + value.Graphs = append(value.Graphs, value.Graphs[0]) + }}, + {name: "negative snapshot", label: "snapshot", mutate: func(value *State) { + value.Graphs[0].Snapshot.NodeCount = -1 + }}, + {name: "unknown phase", label: "phase", mutate: func(value *State) { value.Graphs[0].Phase = "unknown" }}, + {name: "cursor without matching shard", label: "cursor", mutate: func(value *State) { value.Graphs[0].NodeCursor++ }}, + {name: "relationship cursor without shard", label: "relationship cursor", mutate: func(value *State) { + value.Graphs[0].RelationshipCursor = 1 + }}, + {name: "relationship shard during nodes", label: "nodes phase", mutate: func(value *State) { + value.Graphs[0].RelationshipShards = []collection.RelationshipShard{ + fixtureRelationshipShard(value.Identity, "alpha", 1, 2, 30), + } + value.Graphs[0].RelationshipCursor = 30 + }}, + {name: "node shard gap", label: "node shard index", mutate: func(value *State) { + value.Graphs[0].NodeShards[0].Index = 2 + }}, + {name: "nonpositive shard count", label: "count", mutate: func(value *State) { + value.Graphs[0].NodeShards[0].Count = 0 + }}, + {name: "shard exceeds configured size", label: "shard size", mutate: func(value *State) { + value.Graphs[0].NodeShards[0].Count = 3 + value.Graphs[0].NodeShards[0].JSONL.Count = 3 + value.Graphs[0].NodeShards[0].Parquet.Count = 3 + }}, + {name: "partial shard before snapshot boundary", label: "partial", mutate: func(value *State) { + value.Graphs[0].NodeShards[0].Count = 1 + value.Graphs[0].NodeShards[0].JSONL.Count = 1 + value.Graphs[0].NodeShards[0].Parquet.Count = 1 + }}, + {name: "shards exceed snapshot", label: "snapshot", mutate: func(value *State) { + value.Graphs[0].Snapshot.NodeCount = 1 + }}, + {name: "nonincreasing shard cursor", label: "source ID", mutate: func(value *State) { + value.Graphs[0].NodeShards = append(value.Graphs[0].NodeShards, + fixtureNodeShard(value.Identity, "alpha", 2, 1, 19)) + }}, + {name: "output mismatch", label: "output mismatch", mutate: func(value *State) { + value.Graphs[0].NodeShards[0].JSONL = nil + }}, + {name: "traversing JSONL path", label: "traverses", mutate: func(value *State) { + value.Graphs[0].NodeShards[0].JSONL.Path = "../outside" + }}, + {name: "wrong JSONL codec metadata", label: "codec", mutate: func(value *State) { + value.Graphs[0].NodeShards[0].JSONL.Codec = "gzip" + }}, + {name: "wrong JSONL count metadata", label: "count", mutate: func(value *State) { + value.Graphs[0].NodeShards[0].JSONL.Count++ + }}, + {name: "bad JSONL digest", label: "SHA-256", mutate: func(value *State) { + value.Graphs[0].NodeShards[0].JSONL.SHA256 = "bad" + }}, + {name: "bad Parquet bytes", label: "stored bytes", mutate: func(value *State) { + value.Graphs[0].NodeShards[0].Parquet.StoredBytes = 0 + }}, + {name: "relationships phase before nodes complete", label: "relationships phase", mutate: func(value *State) { + value.Graphs[0].Phase = PhaseRelationships + }}, + {name: "complete before relationships complete", label: "complete phase", mutate: func(value *State) { + value.Graphs[0].Phase = PhaseComplete + }}, + {name: "incomplete graph before later state", label: "phase progression", mutate: func(value *State) { + second := GraphState{Name: "beta", Snapshot: dawgs.Snapshot{}, Phase: PhaseComplete} + value.Graphs = append(value.Graphs, second) + }}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + root := t.TempDir() + state := fixtureState() + test.mutate(&state) + + err := (Store{Root: root}).Save(state) + require.ErrorContains(t, err, test.label) + _, statErr := os.Stat(filepath.Join(root, FileName)) + require.ErrorIs(t, statErr, os.ErrNotExist) + }) + } +} + +func TestStoreLoadRejectsMalformedStateWrittenAsJSON(t *testing.T) { + root := t.TempDir() + state := fixtureState() + state.Graphs[0].NodeCursor++ + payload, err := json.Marshal(state) + require.NoError(t, err) + require.NoError(t, os.WriteFile(filepath.Join(root, FileName), payload, 0o600)) + + _, found, err := (Store{Root: root}).Load() + require.True(t, found) + require.ErrorContains(t, err, "cursor") +} + +func TestCleanupOrphansDeletesOnlyRecognizedNextShardArtifacts(t *testing.T) { + root := t.TempDir() + state := fixtureState() + committed := state.Graphs[0].NodeShards[0] + installArtifact(t, root, committed.JSONL.Path) + installArtifact(t, root, committed.Parquet.Path) + + nextJSONL := collection.NodeJSONLPath("alpha", 2, jsonl.CodecZstd) + nextParquet := collection.NodeParquetPath("alpha", 2) + installArtifact(t, root, nextJSONL) + installArtifact(t, root, nextParquet) + installArtifact(t, root, nextJSONL+".tmp-A_z09") + installArtifact(t, root, nextParquet+".tmp-nonce") + + require.NoError(t, (Store{Root: root}).CleanupOrphans(state)) + + requireArtifactExists(t, root, committed.JSONL.Path) + requireArtifactExists(t, root, committed.Parquet.Path) + requireArtifactMissing(t, root, nextJSONL) + requireArtifactMissing(t, root, nextParquet) + requireArtifactMissing(t, root, nextJSONL+".tmp-A_z09") + requireArtifactMissing(t, root, nextParquet+".tmp-nonce") +} + +func TestCleanupOrphansRejectsUnknownEntryWithoutDeletingAnything(t *testing.T) { + root := t.TempDir() + state := fixtureState() + next := collection.NodeJSONLPath("alpha", 2, jsonl.CodecZstd) + installArtifact(t, root, next) + installArtifact(t, root, "notes.txt") + + err := (Store{Root: root}).CleanupOrphans(state) + require.ErrorContains(t, err, "unknown resume") + requireArtifactExists(t, root, next) + requireArtifactExists(t, root, "notes.txt") +} + +func TestCleanupOrphansRejectsInvalidCheckpointStagingNameWithoutDeletingValidStage(t *testing.T) { + tests := []string{ + FileName + ".tmp", + FileName + ".tmp-", + FileName + ".tmp-a.b", + FileName + ".tmp-a=b", + FileName + ".tmp-z.bad", + FileName + ".tmpx-nonce", + "x" + FileName + ".tmp-nonce", + } + + for _, invalid := range tests { + t.Run(invalid, func(t *testing.T) { + root := t.TempDir() + store := Store{Root: root} + require.NoError(t, store.Save(fixtureState())) + loaded, found, err := store.Load() + require.NoError(t, err) + require.True(t, found) + + valid := FileName + ".tmp-valid_nonce-09" + installArtifact(t, root, valid) + installArtifact(t, root, invalid) + + err = store.CleanupOrphans(loaded) + require.ErrorContains(t, err, "unknown resume") + requireArtifactExists(t, root, valid) + requireArtifactExists(t, root, invalid) + }) + } +} + +func TestCleanupOrphansRejectsCheckpointStagingSymlink(t *testing.T) { + root := t.TempDir() + store := Store{Root: root} + require.NoError(t, store.Save(fixtureState())) + loaded, found, err := store.Load() + require.NoError(t, err) + require.True(t, found) + + outside := filepath.Join(t.TempDir(), "outside") + require.NoError(t, os.WriteFile(outside, []byte("keep"), 0o600)) + staging := filepath.Join(root, FileName+".tmp-valid_nonce") + require.NoError(t, os.Symlink(outside, staging)) + + err = store.CleanupOrphans(loaded) + require.ErrorContains(t, err, "symbolic link") + requireArtifactExists(t, root, FileName+".tmp-valid_nonce") + payload, readErr := os.ReadFile(outside) + require.NoError(t, readErr) + require.Equal(t, "keep", string(payload)) +} + +func TestCleanupOrphansRejectsCheckpointStagingDirectory(t *testing.T) { + root := t.TempDir() + store := Store{Root: root} + require.NoError(t, store.Save(fixtureState())) + loaded, found, err := store.Load() + require.NoError(t, err) + require.True(t, found) + + staging := filepath.Join(root, FileName+".tmp-valid_nonce") + require.NoError(t, os.Mkdir(staging, 0o700)) + + err = store.CleanupOrphans(loaded) + require.ErrorContains(t, err, "unknown resume directory") + requireArtifactExists(t, root, FileName+".tmp-valid_nonce") +} + +func TestCleanupOrphansRejectsUnsafeAndOverbroadLookalikesWithoutDeletion(t *testing.T) { + tests := []struct { + name string + path string + }{ + {name: "later final index", path: collection.NodeJSONLPath("alpha", 3, jsonl.CodecZstd)}, + {name: "wrong extension", path: "graphs/alpha/nodes/000002.json"}, + {name: "empty nonce", path: collection.NodeJSONLPath("alpha", 2, jsonl.CodecZstd) + ".tmp-"}, + {name: "non URL safe nonce", path: collection.NodeJSONLPath("alpha", 2, jsonl.CodecZstd) + ".tmp-a.b"}, + {name: "unrelated graph", path: collection.NodeJSONLPath("other", 2, jsonl.CodecZstd)}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + root := t.TempDir() + state := fixtureState() + recognized := collection.NodeParquetPath("alpha", 2) + installArtifact(t, root, recognized) + installArtifact(t, root, test.path) + + err := (Store{Root: root}).CleanupOrphans(state) + require.Error(t, err) + requireArtifactExists(t, root, recognized) + requireArtifactExists(t, root, test.path) + }) + } +} + +func TestCleanupOrphansRejectsSymlinkWithoutDeletingRecognizedArtifact(t *testing.T) { + root := t.TempDir() + outside := filepath.Join(t.TempDir(), "outside") + require.NoError(t, os.WriteFile(outside, []byte("keep"), 0o600)) + state := fixtureState() + recognized := collection.NodeJSONLPath("alpha", 2, jsonl.CodecZstd) + installArtifact(t, root, recognized) + link := filepath.Join(root, filepath.FromSlash("graphs/alpha/nodes/link")) + require.NoError(t, os.Symlink(outside, link)) + + err := (Store{Root: root}).CleanupOrphans(state) + require.ErrorContains(t, err, "symbolic link") + requireArtifactExists(t, root, recognized) + payload, readErr := os.ReadFile(outside) + require.NoError(t, readErr) + require.Equal(t, "keep", string(payload)) +} + +func TestCleanupOrphansRejectsRootReplacedByOutsideSymlinkBeforePinning(t *testing.T) { + parent := t.TempDir() + root := filepath.Join(parent, "collection") + require.NoError(t, os.Mkdir(root, 0o700)) + state := fixtureState() + recognized := collection.NodeJSONLPath("alpha", 2, jsonl.CodecZstd) + installArtifact(t, root, recognized) + + outside := filepath.Join(parent, "outside") + require.NoError(t, os.Mkdir(outside, 0o700)) + installArtifact(t, outside, recognized) + movedRoot := filepath.Join(parent, "collection-before-swap") + + originalOpenRoot := checkpointOpenRoot + checkpointOpenRoot = func(name string) (*os.Root, error) { + require.NoError(t, os.Rename(root, movedRoot)) + require.NoError(t, os.Symlink(outside, root)) + return originalOpenRoot(name) + } + t.Cleanup(func() { checkpointOpenRoot = originalOpenRoot }) + + err := (Store{Root: root}).CleanupOrphans(state) + require.ErrorContains(t, err, "changed while opening") + requireArtifactExists(t, movedRoot, recognized) + requireArtifactExists(t, outside, recognized) +} + +func TestCleanupOrphansCannotEscapeRootWhenDirectoryBecomesSymlinkBeforeRemoval(t *testing.T) { + root := t.TempDir() + state := fixtureState() + recognized := collection.NodeJSONLPath("alpha", 2, jsonl.CodecZstd) + installArtifact(t, root, recognized) + + outside := t.TempDir() + outsideArtifact := filepath.Join(outside, filepath.Base(recognized)) + require.NoError(t, os.WriteFile(outsideArtifact, []byte("outside"), 0o600)) + + originalRemove := checkpointPinnedRemove + checkpointPinnedRemove = func(parentRoot *os.Root, basename string) error { + nodes := filepath.Join(root, filepath.FromSlash("graphs/alpha/nodes")) + moved := filepath.Join(root, filepath.FromSlash("graphs/alpha/nodes-inventoried")) + require.NoError(t, os.Rename(nodes, moved)) + require.NoError(t, os.Symlink(outside, nodes)) + return originalRemove(parentRoot, basename) + } + t.Cleanup(func() { checkpointPinnedRemove = originalRemove }) + + require.NoError(t, (Store{Root: root}).CleanupOrphans(state)) + payload, readErr := os.ReadFile(outsideArtifact) + require.NoError(t, readErr) + require.Equal(t, "outside", string(payload)) + requireArtifactMissing(t, root, "graphs/alpha/nodes-inventoried/"+filepath.Base(recognized)) +} + +func TestCleanupOrphansPinsCandidateParentAgainstInRootCommittedRedirect(t *testing.T) { + root := t.TempDir() + state := fixtureState() + state.Graphs[0].Snapshot = dawgs.Snapshot{NodeCount: 2} + state.Graphs[0].Phase = PhaseComplete + state.Graphs = append(state.Graphs, GraphState{ + Name: "beta", + Snapshot: dawgs.Snapshot{NodeCount: 2}, + Phase: PhaseNodes, + }) + + committed := state.Graphs[0].NodeShards[0].JSONL.Path + candidate := collection.NodeJSONLPath("beta", 1, jsonl.CodecZstd) + installArtifact(t, root, committed) + installArtifact(t, root, candidate) + + originalRemove := checkpointPinnedRemove + checkpointPinnedRemove = func(parentRoot *os.Root, basename string) error { + betaNodes := filepath.Join(root, filepath.FromSlash("graphs/beta/nodes")) + moved := filepath.Join(root, filepath.FromSlash("graphs/beta/nodes-inventoried")) + require.NoError(t, os.Rename(betaNodes, moved)) + require.NoError(t, os.Symlink("../alpha/nodes", betaNodes)) + return originalRemove(parentRoot, basename) + } + t.Cleanup(func() { checkpointPinnedRemove = originalRemove }) + + require.NoError(t, (Store{Root: root}).CleanupOrphans(state)) + requireArtifactExists(t, root, committed) + requireArtifactMissing(t, root, "graphs/beta/nodes-inventoried/"+filepath.Base(candidate)) +} + +func TestCleanupOrphansRejectsInvalidTraversalStateBeforeDeleting(t *testing.T) { + root := t.TempDir() + state := fixtureState() + recognized := collection.NodeJSONLPath("alpha", 2, jsonl.CodecZstd) + installArtifact(t, root, recognized) + state.Graphs[0].NodeShards[0].JSONL.Path = "../outside" + + err := (Store{Root: root}).CleanupOrphans(state) + require.ErrorContains(t, err, "path") + requireArtifactExists(t, root, recognized) +} + +func TestCleanupOrphansHandlesEveryOutputCombination(t *testing.T) { + tests := []struct { + name string + jsonl bool + parquet bool + }{ + {name: "JSONL only", jsonl: true}, + {name: "Parquet only", parquet: true}, + {name: "dual", jsonl: true, parquet: true}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + root := t.TempDir() + state := fixtureStateForOutputs(test.jsonl, test.parquet) + var paths []string + if test.jsonl { + paths = append(paths, collection.NodeJSONLPath("alpha", 2, jsonl.CodecZstd)) + } + if test.parquet { + paths = append(paths, collection.NodeParquetPath("alpha", 2)) + } + for _, path := range paths { + installArtifact(t, root, path) + } + + require.NoError(t, (Store{Root: root}).CleanupOrphans(state)) + for _, path := range paths { + requireArtifactMissing(t, root, path) + } + }) + } +} + +func installArtifact(t *testing.T, root, relative string) { + t.Helper() + absolute := filepath.Join(root, filepath.FromSlash(relative)) + require.NoError(t, os.MkdirAll(filepath.Dir(absolute), 0o700)) + require.NoError(t, os.WriteFile(absolute, []byte("artifact"), 0o600)) +} + +func requireArtifactExists(t *testing.T, root, relative string) { + t.Helper() + _, err := os.Lstat(filepath.Join(root, filepath.FromSlash(relative))) + require.NoError(t, err) +} + +func requireArtifactMissing(t *testing.T, root, relative string) { + t.Helper() + _, err := os.Lstat(filepath.Join(root, filepath.FromSlash(relative))) + require.ErrorIs(t, err, os.ErrNotExist) +} + +func checkpointStagingNames(t *testing.T, root string) []string { + t.Helper() + entries, err := os.ReadDir(root) + require.NoError(t, err) + var names []string + for _, entry := range entries { + if strings.HasPrefix(entry.Name(), FileName+".tmp") { + names = append(names, entry.Name()) + } + } + return names +} diff --git a/ret/codec_helpers_test.go b/ret/codec_helpers_test.go new file mode 100644 index 00000000..137720c8 --- /dev/null +++ b/ret/codec_helpers_test.go @@ -0,0 +1,24 @@ +package ret + +import ( + "github.com/specterops/dawgs/ret/collection" + "github.com/specterops/dawgs/ret/entity" +) + +func readJSONLNodesForTest(root string, artifact collection.JSONLArtifact) ([]entity.Node, error) { + var values []entity.Node + err := collection.ReadJSONLNodes(root, artifact, func(value entity.Node) error { + values = append(values, value) + return nil + }) + return values, err +} + +func readJSONLRelationshipsForTest(root string, artifact collection.JSONLArtifact) ([]entity.Relationship, error) { + var values []entity.Relationship + err := collection.ReadJSONLRelationships(root, artifact, func(value entity.Relationship) error { + values = append(values, value) + return nil + }) + return values, err +} diff --git a/ret/collection/artifact_test.go b/ret/collection/artifact_test.go new file mode 100644 index 00000000..45b5eeef --- /dev/null +++ b/ret/collection/artifact_test.go @@ -0,0 +1,35 @@ +package collection + +import ( + "encoding/json" + "testing" + + "github.com/specterops/dawgs/ret/jsonl" + "github.com/specterops/dawgs/ret/parquet" + "github.com/stretchr/testify/require" +) + +func TestCodecArtifactWrappersMarshalFlat(t *testing.T) { + jsonlValue, err := json.Marshal(JSONLArtifact{ + Path: "nodes.jsonl", + Artifact: jsonl.Artifact{ + SchemaVersion: jsonl.SchemaVersion, + Codec: jsonl.CodecNone, + SHA256: "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + Count: 1, + }, + }) + require.NoError(t, err) + require.JSONEq(t, `{"path":"nodes.jsonl","SchemaVersion":"retriever-jsonl-v1","Codec":"none","SHA256":"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa","Level":0,"Count":1,"UncompressedBytes":0,"StoredBytes":0}`, string(jsonlValue)) + + parquetValue, err := json.Marshal(ParquetArtifact{ + Path: "nodes.parquet", + Artifact: parquet.Artifact{ + SchemaVersion: parquet.SchemaVersion, + SHA256: "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb", + Count: 1, + }, + }) + require.NoError(t, err) + require.JSONEq(t, `{"path":"nodes.parquet","SchemaVersion":"ret-parquet-v1","SHA256":"bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb","Count":1,"StoredBytes":0}`, string(parquetValue)) +} diff --git a/ret/collection/codec.go b/ret/collection/codec.go new file mode 100644 index 00000000..a48c5389 --- /dev/null +++ b/ret/collection/codec.go @@ -0,0 +1,145 @@ +package collection + +import ( + "errors" + "fmt" + "os" + + "github.com/specterops/dawgs/ret/entity" + "github.com/specterops/dawgs/ret/jsonl" + "github.com/specterops/dawgs/ret/parquet" +) + +const codecReadBatchSize = 256 + +func ReadJSONLNodes(root string, artifact JSONLArtifact, visit func(entity.Node) error) error { + file, err := openArtifact(root, artifact.Path) + if err != nil { + return err + } + reader, err := jsonl.NewNodeReader(file, artifact.Artifact) + if err != nil { + return errors.Join(err, file.Close()) + } + readErr := visitReader(&reader, visit) + resultErr := reader.Result() + closeReaderErr := reader.Close() + closeFileErr := file.Close() + return errors.Join(readErr, resultErr, closeReaderErr, closeFileErr) +} + +func ReadJSONLRelationships(root string, artifact JSONLArtifact, visit func(entity.Relationship) error) error { + file, err := openArtifact(root, artifact.Path) + if err != nil { + return err + } + reader, err := jsonl.NewRelationshipReader(file, artifact.Artifact) + if err != nil { + return errors.Join(err, file.Close()) + } + readErr := visitReader(&reader, visit) + resultErr := reader.Result() + closeReaderErr := reader.Close() + closeFileErr := file.Close() + return errors.Join(readErr, resultErr, closeReaderErr, closeFileErr) +} + +func ReadParquetNodes(root string, artifact ParquetArtifact, visit func(entity.Node) error) error { + file, size, err := openRandomAccessArtifact(root, artifact.Path) + if err != nil { + return err + } + reader, err := parquet.NewNodeReader(file, size, artifact.Artifact) + if err != nil { + return errors.Join(err, file.Close()) + } + readErr := visitReader(&reader, visit) + resultErr := reader.Result() + closeReaderErr := reader.Close() + closeFileErr := file.Close() + return errors.Join(readErr, resultErr, closeReaderErr, closeFileErr) +} + +func ReadParquetRelationships(root string, artifact ParquetArtifact, visit func(entity.Relationship) error) error { + file, size, err := openRandomAccessArtifact(root, artifact.Path) + if err != nil { + return err + } + reader, err := parquet.NewRelationshipReader(file, size, artifact.Artifact) + if err != nil { + return errors.Join(err, file.Close()) + } + readErr := visitReader(&reader, visit) + resultErr := reader.Result() + closeReaderErr := reader.Close() + closeFileErr := file.Close() + return errors.Join(readErr, resultErr, closeReaderErr, closeFileErr) +} + +type pullReader[E any] interface { + Pull(int) ([]E, error) + Done() bool +} + +func visitReader[E any](reader pullReader[E], visit func(E) error) error { + var index int64 + for !reader.Done() { + batch, err := reader.Pull(codecReadBatchSize) + if err != nil { + return err + } + for _, value := range batch { + index++ + if visit != nil { + if err := visit(value); err != nil { + return fmt.Errorf("visit record %d: %w", index, err) + } + } + } + } + return nil +} + +func openArtifact(root, relative string) (*os.File, error) { + if root == "" { + root = "." + } + if err := inspectNonSymlinkArtifact(root, relative); err != nil { + return nil, err + } + path, err := SafeJoin(root, relative) + if err != nil { + return nil, err + } + file, err := os.Open(path) + if err != nil { + return nil, fmt.Errorf("open artifact: %w", err) + } + pathInfo, err := os.Lstat(path) + if err != nil { + return nil, errors.Join(fmt.Errorf("reinspect artifact: %w", err), file.Close()) + } + openedInfo, err := file.Stat() + if err != nil { + return nil, errors.Join(fmt.Errorf("inspect open artifact: %w", err), file.Close()) + } + if pathInfo.Mode()&os.ModeSymlink != 0 || !openedInfo.Mode().IsRegular() || !os.SameFile(pathInfo, openedInfo) { + return nil, errors.Join(fmt.Errorf("artifact changed while opening: %q", relative), file.Close()) + } + return file, nil +} + +func openRandomAccessArtifact(root, relative string) (*os.File, int64, error) { + file, err := openArtifact(root, relative) + if err != nil { + return nil, 0, err + } + info, err := file.Stat() + if err != nil { + return nil, 0, errors.Join(fmt.Errorf("inspect open artifact: %w", err), file.Close()) + } + if !info.Mode().IsRegular() { + return nil, 0, errors.Join(fmt.Errorf("artifact is not a regular file: %q", relative), file.Close()) + } + return file, info.Size(), nil +} diff --git a/ret/collection/io.go b/ret/collection/io.go new file mode 100644 index 00000000..aab77720 --- /dev/null +++ b/ret/collection/io.go @@ -0,0 +1,100 @@ +package collection + +import ( + "encoding/json" + "errors" + "fmt" + "io" + "os" + "path/filepath" +) + +func Read(root string) (Manifest, error) { + manifestPath := filepath.Join(root, ManifestName) + file, err := os.Open(manifestPath) + if err != nil { + return Manifest{}, fmt.Errorf("open manifest: %w", err) + } + defer file.Close() + + decoder := json.NewDecoder(file) + decoder.DisallowUnknownFields() + + var manifest Manifest + if err := decoder.Decode(&manifest); err != nil { + return Manifest{}, fmt.Errorf("decode manifest: %w", err) + } + var trailing any + if err := decoder.Decode(&trailing); !errors.Is(err, io.EOF) { + if err == nil { + return Manifest{}, fmt.Errorf("decode manifest: trailing JSON value") + } + return Manifest{}, fmt.Errorf("decode manifest trailing data: %w", err) + } + if err := manifest.Validate(); err != nil { + return Manifest{}, fmt.Errorf("validate manifest: %w", err) + } + + return manifest, nil +} + +func Write(root string, manifest Manifest) error { + return writeWithEncoder(root, manifest, func(writer io.Writer, value Manifest) error { + if err := json.NewEncoder(writer).Encode(value); err != nil { + return fmt.Errorf("encode manifest: %w", err) + } + return nil + }, os.Remove) +} + +func writeWithEncoder( + root string, + manifest Manifest, + encode func(io.Writer, Manifest) error, + remove func(string) error, +) (resultErr error) { + if err := manifest.Validate(); err != nil { + return fmt.Errorf("validate manifest: %w", err) + } + + temporary := filepath.Join(root, ManifestName+".tmp") + final := filepath.Join(root, ManifestName) + file, err := os.OpenFile(temporary, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o600) + if err != nil { + return fmt.Errorf("create temporary manifest: %w", err) + } + published := false + defer func() { + errorsToJoin := []error{resultErr} + if file != nil { + if closeErr := file.Close(); closeErr != nil { + errorsToJoin = append(errorsToJoin, fmt.Errorf("cleanup close temporary manifest: %w", closeErr)) + } + } + if !published { + if removeErr := remove(temporary); removeErr != nil && !errors.Is(removeErr, os.ErrNotExist) { + errorsToJoin = append(errorsToJoin, fmt.Errorf("cleanup remove temporary manifest: %w", removeErr)) + } + } + resultErr = errors.Join(errorsToJoin...) + }() + + if err := encode(file, manifest); err != nil { + return err + } + if err := file.Sync(); err != nil { + return fmt.Errorf("sync temporary manifest: %w", err) + } + if err := file.Close(); err != nil { + file = nil + return fmt.Errorf("close temporary manifest: %w", err) + } + file = nil + + if err := os.Rename(temporary, final); err != nil { + return fmt.Errorf("publish manifest: %w", err) + } + published = true + + return nil +} diff --git a/ret/collection/io_test.go b/ret/collection/io_test.go new file mode 100644 index 00000000..d2aa76ff --- /dev/null +++ b/ret/collection/io_test.go @@ -0,0 +1,170 @@ +package collection + +import ( + "bytes" + "errors" + "io" + "os" + "path/filepath" + "testing" + + "github.com/specterops/dawgs/ret/jsonl" + "github.com/stretchr/testify/require" +) + +func TestSafeJoinKeepsCleanSlashSeparatedPathsBeneathRoot(t *testing.T) { + root := t.TempDir() + + got, err := SafeJoin(root, "graphs/example/nodes/000001.jsonl") + + require.NoError(t, err) + require.Equal(t, filepath.Join(root, "graphs", "example", "nodes", "000001.jsonl"), got) +} + +func TestSafeJoinRejectsUnsafeOrAmbiguousPaths(t *testing.T) { + for _, relative := range []string{ + "", + ".", + "/absolute", + "../escape", + "graphs/../../escape", + "graphs/../nodes", + "graphs/./nodes", + "graphs//nodes", + `graphs\example\nodes`, + } { + t.Run(relative, func(t *testing.T) { + _, err := SafeJoin(t.TempDir(), relative) + require.Error(t, err) + }) + } +} + +func TestPathHelpersUseEscapedGraphsOneBasedPaddedIndicesAndCodecSuffixes(t *testing.T) { + require.Equal(t, "graphs/graph%20name/nodes/000001.jsonl", NodeJSONLPath("graph name", 1, jsonl.CodecNone)) + require.Equal(t, "graphs/graph%20name/nodes/000012.jsonl.gz", NodeJSONLPath("graph name", 12, jsonl.CodecGzip)) + require.Equal(t, "graphs/graph%20name/relationships/000123.jsonl.zst", RelationshipJSONLPath("graph name", 123, jsonl.CodecZstd)) + require.Equal(t, "graphs/graph%20name/nodes/000001.parquet", NodeParquetPath("graph name", 1)) + require.Equal(t, "graphs/graph%20name/relationships/000001.parquet", RelationshipParquetPath("graph name", 1)) + require.Panics(t, func() { NodeParquetPath("graph", 0) }) + require.Panics(t, func() { NodeJSONLPath("graph", 1, jsonl.Codec("zip")) }) +} + +func TestWriteAndReadRoundTripValidatedManifest(t *testing.T) { + root := t.TempDir() + want := fixtureManifest() + + require.NoError(t, Write(root, want)) + got, err := Read(root) + + require.NoError(t, err) + require.Equal(t, want, got) + _, err = os.Stat(filepath.Join(root, ManifestName+".tmp")) + require.ErrorIs(t, err, os.ErrNotExist) +} + +func TestReadRejectsInvalidAndTrailingManifestJSON(t *testing.T) { + for name, contents := range map[string]string{ + "invalid manifest": `{"format":"wrong"}`, + "trailing JSON": `{"format":"wrong"} {}`, + "unknown field": `{"format":"wrong","unknown":true}`, + } { + t.Run(name, func(t *testing.T) { + root := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(root, ManifestName), []byte(contents), 0o600)) + + _, err := Read(root) + + require.Error(t, err) + }) + } +} + +func TestWritePreservesExistingManifestWhenEncodingFailsBeforePublication(t *testing.T) { + root := t.TempDir() + original := fixtureManifest() + require.NoError(t, Write(root, original)) + manifestPath := filepath.Join(root, ManifestName) + before, err := os.ReadFile(manifestPath) + require.NoError(t, err) + + injectedErr := errors.New("injected encoding failure") + err = writeWithEncoder(root, fixtureManifest(), func(writer io.Writer, _ Manifest) error { + _, writeErr := writer.Write([]byte(`{"partial":`)) + require.NoError(t, writeErr) + return injectedErr + }, os.Remove) + + require.ErrorIs(t, err, injectedErr) + after, err := os.ReadFile(manifestPath) + require.NoError(t, err) + require.True(t, bytes.Equal(before, after)) + _, err = os.Stat(filepath.Join(root, ManifestName+".tmp")) + require.ErrorIs(t, err, os.ErrNotExist) +} + +func TestWriteJoinsPrimaryAndCleanupErrorsWithoutReplacingManifest(t *testing.T) { + root := t.TempDir() + original := fixtureManifest() + require.NoError(t, Write(root, original)) + manifestPath := filepath.Join(root, ManifestName) + before, err := os.ReadFile(manifestPath) + require.NoError(t, err) + + primaryErr := errors.New("injected encoding failure") + err = writeWithEncoder(root, fixtureManifest(), func(writer io.Writer, _ Manifest) error { + file, ok := writer.(*os.File) + require.True(t, ok) + require.NoError(t, file.Close()) + return primaryErr + }, os.Remove) + + require.ErrorIs(t, err, primaryErr) + require.ErrorIs(t, err, os.ErrClosed) + require.ErrorContains(t, err, "close temporary manifest") + after, err := os.ReadFile(manifestPath) + require.NoError(t, err) + require.Equal(t, before, after) +} + +func TestWriteJoinsPrimaryAndRemovalErrorsWithoutReplacingManifest(t *testing.T) { + root := t.TempDir() + original := fixtureManifest() + require.NoError(t, Write(root, original)) + manifestPath := filepath.Join(root, ManifestName) + before, err := os.ReadFile(manifestPath) + require.NoError(t, err) + + primaryErr := errors.New("injected encoding failure") + cleanupErr := errors.New("injected removal failure") + err = writeWithEncoder(root, fixtureManifest(), func(writer io.Writer, _ Manifest) error { + _, writeErr := writer.Write([]byte(`{"partial":`)) + require.NoError(t, writeErr) + return primaryErr + }, func(string) error { + return cleanupErr + }) + + require.ErrorIs(t, err, primaryErr) + require.ErrorIs(t, err, cleanupErr) + require.ErrorContains(t, err, "remove temporary manifest") + after, err := os.ReadFile(manifestPath) + require.NoError(t, err) + require.Equal(t, before, after) +} + +func TestWriteValidatesBeforeReplacingExistingManifest(t *testing.T) { + root := t.TempDir() + require.NoError(t, Write(root, fixtureManifest())) + manifestPath := filepath.Join(root, ManifestName) + before, err := os.ReadFile(manifestPath) + require.NoError(t, err) + + invalid := fixtureManifest() + invalid.Format = "wrong" + require.Error(t, Write(root, invalid)) + + after, err := os.ReadFile(manifestPath) + require.NoError(t, err) + require.Equal(t, before, after) +} diff --git a/ret/collection/manifest.go b/ret/collection/manifest.go new file mode 100644 index 00000000..88f91b3c --- /dev/null +++ b/ret/collection/manifest.go @@ -0,0 +1,666 @@ +package collection + +import ( + "encoding/hex" + "fmt" + "math" + "path" + "strconv" + "strings" + "time" + + "github.com/specterops/dawgs/ret/jsonl" + "github.com/specterops/dawgs/ret/metrics" + "github.com/specterops/dawgs/ret/parquet" + "github.com/specterops/dawgs/ret/scrub" +) + +const ( + Format = "ret-collection-v1" + ManifestName = "manifest.json" +) + +type Manifest struct { + Format string `json:"format"` + CreatedAt time.Time `json:"created_at"` + Outputs OutputConfig `json:"outputs"` + Scrub ScrubMetadata `json:"scrub"` + Graphs []Graph `json:"graphs"` +} + +type OutputConfig struct { + JSONL *JSONLOutput `json:"jsonl,omitempty"` + Parquet *ParquetOutput `json:"parquet,omitempty"` +} + +type JSONLOutput struct { + SchemaVersion string `json:"schema_version"` + Codec string `json:"codec"` + Level int `json:"level"` +} + +type ParquetOutput struct { + SchemaVersion string `json:"schema_version"` +} + +type ScrubMetadata struct { + Enabled bool `json:"enabled"` + RulesFingerprint string `json:"rules_fingerprint,omitempty"` + SaltFingerprint string `json:"salt_fingerprint,omitempty"` +} + +type Graph struct { + Name string `json:"name"` + NodeCount int64 `json:"node_count"` + RelationshipCount int64 `json:"relationship_count"` + KindCatalog []string `json:"kind_catalog"` + NodeShards []NodeShard `json:"node_shards"` + RelationshipShards []RelationshipShard `json:"relationship_shards"` + Metrics metrics.GraphMetrics `json:"metrics"` +} + +func (s Manifest) Validate() error { + if s.Format != Format { + return fmt.Errorf("collection format %q does not match %q", s.Format, Format) + } + if s.CreatedAt.IsZero() { + return fmt.Errorf("collection created_at is required") + } + _, offset := s.CreatedAt.Zone() + if offset != 0 { + return fmt.Errorf("collection created_at must be UTC") + } + if err := s.Outputs.validate(); err != nil { + return err + } + if err := s.Scrub.validate(); err != nil { + return err + } + + graphNames := make(map[string]struct{}, len(s.Graphs)) + artifactPaths := make(map[string]struct{}) + for graphIndex := range s.Graphs { + graph := &s.Graphs[graphIndex] + if err := validateGraphName(graph.Name); err != nil { + return fmt.Errorf("graph %d graph name: %w", graphIndex+1, err) + } + if _, found := graphNames[graph.Name]; found { + return fmt.Errorf("duplicate graph name %q", graph.Name) + } + graphNames[graph.Name] = struct{}{} + + if err := validateGraph(*graph, s.Outputs, s.Scrub.Enabled, artifactPaths); err != nil { + return fmt.Errorf("graph %q: %w", graph.Name, err) + } + } + + return nil +} + +func (s OutputConfig) validate() error { + if s.JSONL == nil && s.Parquet == nil { + return fmt.Errorf("collection must enable at least one output") + } + if s.JSONL != nil { + if s.JSONL.SchemaVersion != jsonl.SchemaVersion { + return fmt.Errorf("collection JSONL schema %q does not match %q", s.JSONL.SchemaVersion, jsonl.SchemaVersion) + } + config := jsonl.Config{ + Codec: jsonl.Codec(s.JSONL.Codec), + Level: s.JSONL.Level, + } + if err := config.Validate(); err != nil { + return fmt.Errorf("collection JSONL output: %w", err) + } + } + if s.Parquet != nil && s.Parquet.SchemaVersion != parquet.SchemaVersion { + return fmt.Errorf("collection Parquet schema %q does not match %q", s.Parquet.SchemaVersion, parquet.SchemaVersion) + } + + return nil +} + +func (s ScrubMetadata) validate() error { + if s.Enabled { + if !isLowerHexDigest(s.RulesFingerprint) { + return fmt.Errorf("enabled scrub rules fingerprint must be 64 lowercase hexadecimal characters") + } + if !isLowerHexDigest(s.SaltFingerprint) { + return fmt.Errorf("enabled scrub salt fingerprint must be 64 lowercase hexadecimal characters") + } + return nil + } + if s.RulesFingerprint != "" || s.SaltFingerprint != "" { + return fmt.Errorf("disabled scrub metadata must not contain fingerprints") + } + + return nil +} + +func validateGraphName(name string) error { + if strings.TrimSpace(name) == "" { + return fmt.Errorf("is empty") + } + if name == "." || name == ".." || path.Clean(name) != name { + return fmt.Errorf("%q is not a clean path segment", name) + } + if strings.ContainsAny(name, `/\`) || strings.ContainsRune(name, '\x00') { + return fmt.Errorf("%q is not a single safe path segment", name) + } + + return nil +} + +func validateGraph(graph Graph, outputs OutputConfig, scrubEnabled bool, artifactPaths map[string]struct{}) error { + if graph.NodeCount < 0 { + return fmt.Errorf("node count must not be negative: %d", graph.NodeCount) + } + if graph.RelationshipCount < 0 { + return fmt.Errorf("relationship count must not be negative: %d", graph.RelationshipCount) + } + if err := validateKindCatalog(graph.KindCatalog); err != nil { + return err + } + + empty := graph.NodeCount == 0 && graph.RelationshipCount == 0 + if empty && (len(graph.NodeShards) != 0 || len(graph.RelationshipShards) != 0) { + return fmt.Errorf("empty graph must not contain shards") + } + + nodeTotal, err := validateNodeShards(graph.Name, graph.NodeShards, outputs, scrubEnabled, artifactPaths) + if err != nil { + return err + } + if nodeTotal != graph.NodeCount { + return fmt.Errorf("node shard total %d does not match graph node count %d", nodeTotal, graph.NodeCount) + } + + relationshipTotal, err := validateRelationshipShards(graph.Name, graph.RelationshipShards, outputs, scrubEnabled, artifactPaths) + if err != nil { + return err + } + if relationshipTotal != graph.RelationshipCount { + return fmt.Errorf("relationship shard total %d does not match graph relationship count %d", relationshipTotal, graph.RelationshipCount) + } + + if err := validateMetrics(graph.Metrics, graph.NodeCount, graph.RelationshipCount); err != nil { + return err + } + + return nil +} + +func validateKindCatalog(catalog []string) error { + seen := make(map[string]struct{}, len(catalog)) + for index, kind := range catalog { + if kind == "" { + return fmt.Errorf("kind catalog entry %d is empty", index+1) + } + if _, found := seen[kind]; found { + return fmt.Errorf("kind catalog entry %d duplicates %q", index+1, kind) + } + seen[kind] = struct{}{} + } + + return nil +} + +func validateNodeShards( + graph string, + shards []NodeShard, + outputs OutputConfig, + scrubEnabled bool, + artifactPaths map[string]struct{}, +) (int64, error) { + var total int64 + var lastSourceID uint64 + for offset, shard := range shards { + if err := validateLogicalShard( + "node", + offset, + shard.Index, + shard.Count, + shard.LastSourceID, + shard.ScrubCounts, + scrubEnabled, + shard.JSONL != nil, + shard.Parquet != nil, + outputs, + lastSourceID, + ); err != nil { + return 0, err + } + lastSourceID = shard.LastSourceID + + if shard.JSONL != nil { + expected := NodeJSONLPath(graph, shard.Index, jsonl.Codec(outputs.JSONL.Codec)) + if err := validateJSONLArtifact( + "node", + shard.Index, + shard.Count, + shard.JSONL.SchemaVersion, + shard.JSONL.Path, + string(shard.JSONL.Codec), + shard.JSONL.SHA256, + shard.JSONL.Level, + shard.JSONL.Count, + shard.JSONL.UncompressedBytes, + shard.JSONL.StoredBytes, + expected, + *outputs.JSONL, + artifactPaths, + ); err != nil { + return 0, err + } + } + if shard.Parquet != nil { + expected := NodeParquetPath(graph, shard.Index) + if err := validateParquetArtifact( + "node", + shard.Index, + shard.Count, + shard.Parquet.SchemaVersion, + shard.Parquet.Path, + shard.Parquet.SHA256, + shard.Parquet.Count, + shard.Parquet.StoredBytes, + expected, + artifactPaths, + ); err != nil { + return 0, err + } + } + + var ok bool + total, ok = addNonnegative(total, shard.Count) + if !ok { + return 0, fmt.Errorf("node shard total overflows int64") + } + } + + return total, nil +} + +func validateRelationshipShards( + graph string, + shards []RelationshipShard, + outputs OutputConfig, + scrubEnabled bool, + artifactPaths map[string]struct{}, +) (int64, error) { + var total int64 + var lastSourceID uint64 + for offset, shard := range shards { + if err := validateLogicalShard( + "relationship", + offset, + shard.Index, + shard.Count, + shard.LastSourceID, + shard.ScrubCounts, + scrubEnabled, + shard.JSONL != nil, + shard.Parquet != nil, + outputs, + lastSourceID, + ); err != nil { + return 0, err + } + lastSourceID = shard.LastSourceID + + if shard.JSONL != nil { + expected := RelationshipJSONLPath(graph, shard.Index, jsonl.Codec(outputs.JSONL.Codec)) + if err := validateJSONLArtifact( + "relationship", + shard.Index, + shard.Count, + shard.JSONL.SchemaVersion, + shard.JSONL.Path, + string(shard.JSONL.Codec), + shard.JSONL.SHA256, + shard.JSONL.Level, + shard.JSONL.Count, + shard.JSONL.UncompressedBytes, + shard.JSONL.StoredBytes, + expected, + *outputs.JSONL, + artifactPaths, + ); err != nil { + return 0, err + } + } + if shard.Parquet != nil { + expected := RelationshipParquetPath(graph, shard.Index) + if err := validateParquetArtifact( + "relationship", + shard.Index, + shard.Count, + shard.Parquet.SchemaVersion, + shard.Parquet.Path, + shard.Parquet.SHA256, + shard.Parquet.Count, + shard.Parquet.StoredBytes, + expected, + artifactPaths, + ); err != nil { + return 0, err + } + } + + var ok bool + total, ok = addNonnegative(total, shard.Count) + if !ok { + return 0, fmt.Errorf("relationship shard total overflows int64") + } + } + + return total, nil +} + +func validateLogicalShard( + entityType string, + offset, index int, + count int64, + lastSourceID uint64, + counts scrub.ActionCounts, + scrubEnabled, hasJSONL, hasParquet bool, + outputs OutputConfig, + previousSourceID uint64, +) error { + expectedIndex := offset + 1 + if index != expectedIndex { + return fmt.Errorf("%s shard index: got %d want %d", entityType, index, expectedIndex) + } + if hasJSONL != (outputs.JSONL != nil) || hasParquet != (outputs.Parquet != nil) { + return fmt.Errorf("%s shard %d output mismatch with globally enabled outputs", entityType, index) + } + if count <= 0 { + return fmt.Errorf("%s shard %d count must be positive", entityType, index) + } + if lastSourceID == 0 { + return fmt.Errorf("%s shard %d last source ID must be nonzero", entityType, index) + } + if previousSourceID != 0 && lastSourceID <= previousSourceID { + return fmt.Errorf("%s shard %d last source ID %d does not increase after %d", entityType, index, lastSourceID, previousSourceID) + } + if err := validateScrubCounts(counts, scrubEnabled); err != nil { + return fmt.Errorf("%s shard %d: %w", entityType, index, err) + } + + return nil +} + +func validateScrubCounts(counts scrub.ActionCounts, enabled bool) error { + if !enabled && !counts.IsZero() { + return fmt.Errorf("scrub counts are present while scrubbing is disabled") + } + if counts.Preserve < 0 { + return fmt.Errorf("scrub count for %q must not be negative", "preserve") + } + if counts.Pseudonymize < 0 { + return fmt.Errorf("scrub count for %q must not be negative", "pseudonymize") + } + if counts.Redact < 0 { + return fmt.Errorf("scrub count for %q must not be negative", "redact") + } + if counts.ShiftTimestamp < 0 { + return fmt.Errorf("scrub count for %q must not be negative", "shift_timestamp") + } + + return nil +} + +func validateJSONLArtifact( + entityType string, + shardIndex int, + shardCount int64, + schemaVersion, artifactPath, codec, sha256 string, + level int, + count, uncompressedBytes, storedBytes int64, + expectedPath string, + output JSONLOutput, + artifactPaths map[string]struct{}, +) error { + prefix := fmt.Sprintf("%s shard %d JSONL", entityType, shardIndex) + if schemaVersion != output.SchemaVersion { + return fmt.Errorf("%s schema %q does not match configured JSONL schema %q", prefix, schemaVersion, output.SchemaVersion) + } + if codec != output.Codec { + return fmt.Errorf("%s codec %q does not match configured JSONL codec %q", prefix, codec, output.Codec) + } + if level != output.Level { + return fmt.Errorf("%s level %d does not match configured JSONL level %d", prefix, level, output.Level) + } + if count != shardCount { + return fmt.Errorf("%s count %d does not match shard count %d", prefix, count, shardCount) + } + if !isLowerHexDigest(sha256) { + return fmt.Errorf("%s SHA-256 must be 64 lowercase hexadecimal characters", prefix) + } + if uncompressedBytes <= 0 { + return fmt.Errorf("%s uncompressed bytes must be positive", prefix) + } + if storedBytes <= 0 { + return fmt.Errorf("%s stored bytes must be positive", prefix) + } + if err := validateArtifactPath(artifactPath, expectedPath, artifactPaths); err != nil { + return fmt.Errorf("%s path: %w", prefix, err) + } + + return nil +} + +func validateParquetArtifact( + entityType string, + shardIndex int, + shardCount int64, + schemaVersion, artifactPath, sha256 string, + count, storedBytes int64, + expectedPath string, + artifactPaths map[string]struct{}, +) error { + prefix := fmt.Sprintf("%s shard %d Parquet", entityType, shardIndex) + if schemaVersion != parquet.SchemaVersion { + return fmt.Errorf("%s schema %q does not match %q", prefix, schemaVersion, parquet.SchemaVersion) + } + if count != shardCount { + return fmt.Errorf("%s count %d does not match shard count %d", prefix, count, shardCount) + } + if !isLowerHexDigest(sha256) { + return fmt.Errorf("%s SHA-256 must be 64 lowercase hexadecimal characters", prefix) + } + if storedBytes <= 0 { + return fmt.Errorf("%s stored bytes must be positive", prefix) + } + if err := validateArtifactPath(artifactPath, expectedPath, artifactPaths); err != nil { + return fmt.Errorf("%s path: %w", prefix, err) + } + + return nil +} + +func validateArtifactPath(artifactPath, expectedPath string, paths map[string]struct{}) error { + if _, err := SafeJoin(".", artifactPath); err != nil { + return err + } + if artifactPath != expectedPath { + return fmt.Errorf("%q does not match deterministic path %q", artifactPath, expectedPath) + } + if _, found := paths[artifactPath]; found { + return fmt.Errorf("%q is duplicated", artifactPath) + } + paths[artifactPath] = struct{}{} + + return nil +} + +func validateMetrics(value metrics.GraphMetrics, nodeCount, relationshipCount int64) error { + if value.NodeCount != nodeCount { + return fmt.Errorf("metrics node count %d does not match graph node count %d", value.NodeCount, nodeCount) + } + if value.RelationshipCount != relationshipCount { + return fmt.Errorf("metrics relationship count %d does not match graph relationship count %d", value.RelationshipCount, relationshipCount) + } + + nodeSequenceTotal, err := validateHistogram(value.NodeKindSequences, func(key string) error { + _, err := parseOrderedKindsKey(key) + return err + }) + if err != nil { + return fmt.Errorf("metrics node kind sequences: %w", err) + } + if nodeSequenceTotal != nodeCount { + return fmt.Errorf("metrics node kind sequences sum %d does not match node count %d", nodeSequenceTotal, nodeCount) + } + + relationshipKindTotal, err := validateHistogram(value.RelationshipKinds, func(key string) error { + if key == "" { + return fmt.Errorf("key is empty") + } + return nil + }) + if err != nil { + return fmt.Errorf("metrics relationship kinds: %w", err) + } + if relationshipKindTotal != relationshipCount { + return fmt.Errorf("metrics relationship kinds sum %d does not match relationship count %d", relationshipKindTotal, relationshipCount) + } + + if err := validateDegreeHistogram("inbound", value.InboundDegreeHistogram, nodeCount, relationshipCount); err != nil { + return err + } + if err := validateDegreeHistogram("outbound", value.OutboundDegreeHistogram, nodeCount, relationshipCount); err != nil { + return err + } + + endpointTotal, err := validateHistogram(value.EndpointShapeHistogram, validateEndpointShapeKey) + if err != nil { + return fmt.Errorf("metrics endpoint shape histogram: %w", err) + } + if endpointTotal != relationshipCount { + return fmt.Errorf("metrics endpoint shape histogram sum %d does not match relationship count %d", endpointTotal, relationshipCount) + } + if !strings.HasPrefix(value.Fingerprint, "sha256:") || !isLowerHexDigest(strings.TrimPrefix(value.Fingerprint, "sha256:")) { + return fmt.Errorf("metrics fingerprint must have sha256: followed by 64 lowercase hexadecimal characters") + } + + return nil +} + +func validateHistogram(histogram map[string]int64, validateKey func(string) error) (int64, error) { + var total int64 + for key, count := range histogram { + if err := validateKey(key); err != nil { + return 0, fmt.Errorf("invalid key %q: %w", key, err) + } + if count <= 0 { + return 0, fmt.Errorf("count for %q must be positive", key) + } + var ok bool + total, ok = addNonnegative(total, count) + if !ok { + return 0, fmt.Errorf("counts overflow int64") + } + } + + return total, nil +} + +func validateDegreeHistogram(name string, histogram map[string]int64, nodeCount, relationshipCount int64) error { + var weightedTotal int64 + histogramTotal, err := validateHistogram(histogram, func(key string) error { + degree, parseErr := strconv.ParseInt(key, 10, 64) + if parseErr != nil || degree < 0 || strconv.FormatInt(degree, 10) != key { + return fmt.Errorf("degree must be a canonical nonnegative integer") + } + count := histogram[key] + if degree != 0 && count > math.MaxInt64/degree { + return fmt.Errorf("weighted degree overflows int64") + } + if weightedTotal > math.MaxInt64-degree*count { + return fmt.Errorf("weighted degree total overflows int64") + } + weightedTotal += degree * count + return nil + }) + if err != nil { + return fmt.Errorf("metrics %s degree histogram: %w", name, err) + } + if histogramTotal != nodeCount { + return fmt.Errorf("metrics %s degree histogram sum %d does not match node count %d", name, histogramTotal, nodeCount) + } + if weightedTotal != relationshipCount { + return fmt.Errorf( + "metrics %s degree histogram: %s degree total %d does not match relationship count %d", + name, + name, + weightedTotal, + relationshipCount, + ) + } + + return nil +} + +func validateEndpointShapeKey(key string) error { + segments, err := parseOrderedKindsKey(key) + if err != nil { + return err + } + if len(segments) != 3 { + return fmt.Errorf("endpoint shape must contain exactly three segments") + } + if _, err := parseOrderedKindsKey(segments[0]); err != nil { + return fmt.Errorf("start kind sequence: %w", err) + } + if segments[1] == "" { + return fmt.Errorf("relationship kind is empty") + } + if _, err := parseOrderedKindsKey(segments[2]); err != nil { + return fmt.Errorf("end kind sequence: %w", err) + } + + return nil +} + +func parseOrderedKindsKey(key string) ([]string, error) { + segments := make([]string, 0) + for offset := 0; offset < len(key); { + colon := strings.IndexByte(key[offset:], ':') + if colon < 1 { + return nil, fmt.Errorf("missing length prefix") + } + colon += offset + lengthText := key[offset:colon] + length, err := strconv.Atoi(lengthText) + if err != nil || length < 0 || strconv.Itoa(length) != lengthText { + return nil, fmt.Errorf("invalid length prefix %q", lengthText) + } + start := colon + 1 + end := start + length + if end < start || end > len(key) { + return nil, fmt.Errorf("segment length %d exceeds remaining key", length) + } + segments = append(segments, key[start:end]) + offset = end + } + + return segments, nil +} + +func isLowerHexDigest(value string) bool { + if len(value) != 64 { + return false + } + if _, err := hex.DecodeString(value); err != nil { + return false + } + return strings.ToLower(value) == value +} + +func addNonnegative(left, right int64) (int64, bool) { + if left < 0 || right < 0 || left > math.MaxInt64-right { + return 0, false + } + return left + right, true +} diff --git a/ret/collection/manifest_test.go b/ret/collection/manifest_test.go new file mode 100644 index 00000000..ec53b826 --- /dev/null +++ b/ret/collection/manifest_test.go @@ -0,0 +1,480 @@ +package collection + +import ( + "strings" + "testing" + "time" + + "github.com/specterops/dawgs/ret/entity" + "github.com/specterops/dawgs/ret/jsonl" + "github.com/specterops/dawgs/ret/metrics" + "github.com/specterops/dawgs/ret/parquet" + "github.com/specterops/dawgs/ret/scrub" + "github.com/stretchr/testify/require" +) + +func TestManifestAllowsEverySupportedOutputCombination(t *testing.T) { + t.Run("JSONL only", func(t *testing.T) { + manifest := fixtureManifest() + manifest.Outputs.Parquet = nil + manifest.Graphs[0].NodeShards[0].Parquet = nil + manifest.Graphs[0].RelationshipShards[0].Parquet = nil + + require.NoError(t, manifest.Validate()) + }) + + t.Run("Parquet only", func(t *testing.T) { + manifest := fixtureManifest() + manifest.Outputs.JSONL = nil + manifest.Graphs[0].NodeShards[0].JSONL = nil + manifest.Graphs[0].RelationshipShards[0].JSONL = nil + + require.NoError(t, manifest.Validate()) + }) + + t.Run("both", func(t *testing.T) { + require.NoError(t, fixtureManifest().Validate()) + }) +} + +func TestManifestAllowsEmptyGraphWithoutShards(t *testing.T) { + manifest := fixtureManifest() + manifest.Graphs = append(manifest.Graphs, Graph{ + Name: "empty", + KindCatalog: []string{}, + Metrics: metrics.NewBuilder().Finalize(), + }) + + require.NoError(t, manifest.Validate()) +} + +func TestManifestRejectsInvalidTopLevelMetadata(t *testing.T) { + tests := map[string]struct { + mutate func(*Manifest) + message string + }{ + "format": { + mutate: func(value *Manifest) { value.Format = "ret-collection-v2" }, + message: "format", + }, + "zero creation time": { + mutate: func(value *Manifest) { value.CreatedAt = time.Time{} }, + message: "created_at", + }, + "non-UTC creation time": { + mutate: func(value *Manifest) { + value.CreatedAt = time.Date(2026, time.July, 28, 12, 0, 0, 0, time.FixedZone("east", 3600)) + }, + message: "UTC", + }, + "no output": { + mutate: func(value *Manifest) { value.Outputs = OutputConfig{} }, + message: "output", + }, + "JSONL schema": { + mutate: func(value *Manifest) { value.Outputs.JSONL.SchemaVersion = "wrong" }, + message: "JSONL schema", + }, + "JSONL codec": { + mutate: func(value *Manifest) { value.Outputs.JSONL.Codec = "zip" }, + message: "JSONL", + }, + "JSONL level": { + mutate: func(value *Manifest) { value.Outputs.JSONL.Level = 99 }, + message: "level", + }, + "Parquet schema": { + mutate: func(value *Manifest) { value.Outputs.Parquet.SchemaVersion = "wrong" }, + message: "Parquet schema", + }, + "enabled scrub rules fingerprint": { + mutate: func(value *Manifest) { value.Scrub.RulesFingerprint = "" }, + message: "rules fingerprint", + }, + "enabled scrub salt fingerprint": { + mutate: func(value *Manifest) { value.Scrub.SaltFingerprint = strings.Repeat("A", 64) }, + message: "salt fingerprint", + }, + "disabled scrub fingerprints": { + mutate: func(value *Manifest) { + value.Scrub.Enabled = false + }, + message: "disabled", + }, + } + + for name, testCase := range tests { + t.Run(name, func(t *testing.T) { + manifest := fixtureManifest() + testCase.mutate(&manifest) + + require.ErrorContains(t, manifest.Validate(), testCase.message) + }) + } +} + +func TestManifestRejectsInvalidGraphMetadata(t *testing.T) { + tests := map[string]struct { + mutate func(*Manifest) + message string + }{ + "empty graph name": { + mutate: func(value *Manifest) { value.Graphs[0].Name = "" }, + message: "graph name", + }, + "unsafe graph traversal": { + mutate: func(value *Manifest) { value.Graphs[0].Name = "../escape" }, + message: "graph name", + }, + "unsafe graph slash": { + mutate: func(value *Manifest) { value.Graphs[0].Name = "a/b" }, + message: "graph name", + }, + "unsafe graph backslash": { + mutate: func(value *Manifest) { value.Graphs[0].Name = `a\b` }, + message: "graph name", + }, + "unsafe graph dot": { + mutate: func(value *Manifest) { value.Graphs[0].Name = "." }, + message: "graph name", + }, + "duplicate graph": { + mutate: func(value *Manifest) { + value.Graphs = append(value.Graphs, Graph{Name: value.Graphs[0].Name}) + }, + message: "duplicate graph", + }, + "empty catalog entry": { + mutate: func(value *Manifest) { value.Graphs[0].KindCatalog[1] = "" }, + message: "kind catalog", + }, + "duplicate catalog entry": { + mutate: func(value *Manifest) { value.Graphs[0].KindCatalog[1] = value.Graphs[0].KindCatalog[0] }, + message: "kind catalog", + }, + "node total": { + mutate: func(value *Manifest) { value.Graphs[0].NodeCount++ }, + message: "node shard total", + }, + "relationship total": { + mutate: func(value *Manifest) { value.Graphs[0].RelationshipCount++ }, + message: "relationship shard total", + }, + "empty graph shards": { + mutate: func(value *Manifest) { + value.Graphs[0].NodeCount = 0 + value.Graphs[0].RelationshipCount = 0 + }, + message: "empty graph", + }, + } + + for name, testCase := range tests { + t.Run(name, func(t *testing.T) { + manifest := fixtureManifest() + testCase.mutate(&manifest) + + require.ErrorContains(t, manifest.Validate(), testCase.message) + }) + } +} + +func TestManifestRejectsInvalidLogicalShards(t *testing.T) { + tests := map[string]struct { + mutate func(*Manifest) + message string + }{ + "outputless node shard": { + mutate: func(value *Manifest) { + value.Graphs[0].NodeShards[0].JSONL = nil + value.Graphs[0].NodeShards[0].Parquet = nil + }, + message: "output mismatch", + }, + "missing global output": { + mutate: func(value *Manifest) { value.Graphs[0].NodeShards[0].JSONL = nil }, + message: "output mismatch", + }, + "noncontiguous index": { + mutate: func(value *Manifest) { value.Graphs[0].NodeShards[0].Index = 2 }, + message: "index", + }, + "zero count": { + mutate: func(value *Manifest) { value.Graphs[0].NodeShards[0].Count = 0 }, + message: "count", + }, + "zero cursor": { + mutate: func(value *Manifest) { value.Graphs[0].NodeShards[0].LastSourceID = 0 }, + message: "last source ID", + }, + "nonincreasing cursor": { + mutate: func(value *Manifest) { + first := value.Graphs[0].NodeShards[0] + first.Index = 2 + first.LastSourceID = value.Graphs[0].NodeShards[0].LastSourceID + value.Graphs[0].NodeShards = append(value.Graphs[0].NodeShards, first) + value.Graphs[0].NodeCount += first.Count + value.Graphs[0].Metrics.NodeCount += first.Count + }, + message: "last source ID", + }, + "negative preserve scrub count": { + mutate: func(value *Manifest) { value.Graphs[0].NodeShards[0].ScrubCounts.Preserve = -1 }, + message: "scrub count", + }, + "negative pseudonymize scrub count": { + mutate: func(value *Manifest) { value.Graphs[0].NodeShards[0].ScrubCounts.Pseudonymize = -1 }, + message: "scrub count", + }, + "negative redact scrub count": { + mutate: func(value *Manifest) { value.Graphs[0].NodeShards[0].ScrubCounts.Redact = -1 }, + message: "scrub count", + }, + "negative shift timestamp scrub count": { + mutate: func(value *Manifest) { value.Graphs[0].NodeShards[0].ScrubCounts.ShiftTimestamp = -1 }, + message: "scrub count", + }, + "disabled scrub count": { + mutate: func(value *Manifest) { + value.Scrub = ScrubMetadata{} + }, + message: "scrubbing is disabled", + }, + } + + for name, testCase := range tests { + t.Run(name, func(t *testing.T) { + manifest := fixtureManifest() + testCase.mutate(&manifest) + + require.ErrorContains(t, manifest.Validate(), testCase.message) + }) + } +} + +func TestManifestRejectsInvalidConcreteArtifactMetadata(t *testing.T) { + tests := map[string]struct { + mutate func(*Manifest) + message string + }{ + "JSONL schema": { + mutate: func(value *Manifest) { value.Graphs[0].NodeShards[0].JSONL.SchemaVersion = "wrong" }, + message: "JSONL schema", + }, + "JSONL codec": { + mutate: func(value *Manifest) { value.Graphs[0].NodeShards[0].JSONL.Codec = jsonl.CodecGzip }, + message: "JSONL codec", + }, + "JSONL level": { + mutate: func(value *Manifest) { value.Graphs[0].NodeShards[0].JSONL.Level = 1 }, + message: "JSONL level", + }, + "JSONL count": { + mutate: func(value *Manifest) { value.Graphs[0].NodeShards[0].JSONL.Count++ }, + message: "JSONL count", + }, + "JSONL checksum": { + mutate: func(value *Manifest) { value.Graphs[0].NodeShards[0].JSONL.SHA256 = "abc" }, + message: "JSONL SHA-256", + }, + "JSONL bytes": { + mutate: func(value *Manifest) { value.Graphs[0].NodeShards[0].JSONL.StoredBytes = 0 }, + message: "JSONL stored bytes", + }, + "Parquet schema": { + mutate: func(value *Manifest) { value.Graphs[0].RelationshipShards[0].Parquet.SchemaVersion = "wrong" }, + message: "Parquet schema", + }, + "Parquet count": { + mutate: func(value *Manifest) { value.Graphs[0].RelationshipShards[0].Parquet.Count++ }, + message: "Parquet count", + }, + "Parquet checksum": { + mutate: func(value *Manifest) { value.Graphs[0].RelationshipShards[0].Parquet.SHA256 = "" }, + message: "Parquet SHA-256", + }, + "Parquet bytes": { + mutate: func(value *Manifest) { value.Graphs[0].RelationshipShards[0].Parquet.StoredBytes = -1 }, + message: "Parquet stored bytes", + }, + "unsafe path": { + mutate: func(value *Manifest) { value.Graphs[0].NodeShards[0].JSONL.Path = "../nodes.jsonl" }, + message: "path", + }, + "wrong deterministic path": { + mutate: func(value *Manifest) { value.Graphs[0].NodeShards[0].JSONL.Path = "nodes.jsonl" }, + message: "path", + }, + "duplicate path": { + mutate: func(value *Manifest) { + value.Graphs[0].RelationshipShards[0].Parquet.Path = value.Graphs[0].NodeShards[0].Parquet.Path + }, + message: "path", + }, + } + + for name, testCase := range tests { + t.Run(name, func(t *testing.T) { + manifest := fixtureManifest() + testCase.mutate(&manifest) + + require.ErrorContains(t, manifest.Validate(), testCase.message) + }) + } +} + +func TestManifestRejectsStructurallyInvalidMetrics(t *testing.T) { + tests := map[string]struct { + mutate func(*metrics.GraphMetrics) + message string + }{ + "top-level count": { + mutate: func(value *metrics.GraphMetrics) { value.NodeCount++ }, + message: "metrics node count", + }, + "negative histogram count": { + mutate: func(value *metrics.GraphMetrics) { value.NodeKindSequences["4:User"] = -1 }, + message: "node kind sequences", + }, + "malformed node kind key": { + mutate: func(value *metrics.GraphMetrics) { value.NodeKindSequences = map[string]int64{"User": 2} }, + message: "node kind sequences", + }, + "empty relationship kind": { + mutate: func(value *metrics.GraphMetrics) { value.RelationshipKinds = map[string]int64{"": 1} }, + message: "relationship kinds", + }, + "degree key": { + mutate: func(value *metrics.GraphMetrics) { value.InboundDegreeHistogram = map[string]int64{"01": 2} }, + message: "inbound degree histogram", + }, + "degree count": { + mutate: func(value *metrics.GraphMetrics) { value.OutboundDegreeHistogram = map[string]int64{"0": 2} }, + message: "outbound degree histogram", + }, + "degree total": { + mutate: func(value *metrics.GraphMetrics) { value.InboundDegreeHistogram = map[string]int64{"1": 2} }, + message: "inbound degree total", + }, + "endpoint shape key": { + mutate: func(value *metrics.GraphMetrics) { value.EndpointShapeHistogram = map[string]int64{"broken": 1} }, + message: "endpoint shape histogram", + }, + "fingerprint": { + mutate: func(value *metrics.GraphMetrics) { value.Fingerprint = "sha256:ABC" }, + message: "fingerprint", + }, + } + + for name, testCase := range tests { + t.Run(name, func(t *testing.T) { + manifest := fixtureManifest() + testCase.mutate(&manifest.Graphs[0].Metrics) + + require.ErrorContains(t, manifest.Validate(), testCase.message) + }) + } +} + +func fixtureManifest() Manifest { + builder := metrics.NewBuilder() + mustObserveNode(builder, entity.Node{SourceID: "1", Kinds: []string{"User"}}) + mustObserveNode(builder, entity.Node{SourceID: "2", Kinds: []string{"Group"}}) + mustObserveRelationship(builder, entity.Relationship{ + SourceID: "10", + StartID: "1", + EndID: "2", + Kind: "MEMBER_OF", + }) + + return Manifest{ + Format: Format, + CreatedAt: time.Date(2026, time.July, 28, 12, 0, 0, 0, time.UTC), + Outputs: OutputConfig{ + JSONL: &JSONLOutput{ + SchemaVersion: jsonl.SchemaVersion, + Codec: string(jsonl.CodecZstd), + Level: 3, + }, + Parquet: &ParquetOutput{SchemaVersion: parquet.SchemaVersion}, + }, + Scrub: ScrubMetadata{ + Enabled: true, + RulesFingerprint: strings.Repeat("a", 64), + SaltFingerprint: strings.Repeat("b", 64), + }, + Graphs: []Graph{{ + Name: "bloodhound", + NodeCount: 2, + RelationshipCount: 1, + KindCatalog: []string{"User", "Group", "MEMBER_OF"}, + NodeShards: []NodeShard{{ + Index: 1, + Count: 2, + LastSourceID: 10, + ScrubCounts: scrub.ActionCounts{Pseudonymize: 2}, + JSONL: &JSONLArtifact{ + Path: NodeJSONLPath("bloodhound", 1, jsonl.CodecZstd), + Artifact: jsonl.Artifact{ + SchemaVersion: jsonl.SchemaVersion, + Codec: jsonl.CodecZstd, + SHA256: strings.Repeat("c", 64), + Level: 3, + Count: 2, + UncompressedBytes: 128, + StoredBytes: 80, + }, + }, + Parquet: &ParquetArtifact{ + Path: NodeParquetPath("bloodhound", 1), + Artifact: parquet.Artifact{ + SchemaVersion: parquet.SchemaVersion, + SHA256: strings.Repeat("d", 64), + Count: 2, + StoredBytes: 256, + }, + }, + }}, + RelationshipShards: []RelationshipShard{{ + Index: 1, + Count: 1, + LastSourceID: 20, + ScrubCounts: scrub.ActionCounts{Preserve: 1}, + JSONL: &JSONLArtifact{ + Path: RelationshipJSONLPath("bloodhound", 1, jsonl.CodecZstd), + Artifact: jsonl.Artifact{ + SchemaVersion: jsonl.SchemaVersion, + Codec: jsonl.CodecZstd, + SHA256: strings.Repeat("e", 64), + Level: 3, + Count: 1, + UncompressedBytes: 96, + StoredBytes: 64, + }, + }, + Parquet: &ParquetArtifact{ + Path: RelationshipParquetPath("bloodhound", 1), + Artifact: parquet.Artifact{ + SchemaVersion: parquet.SchemaVersion, + SHA256: strings.Repeat("f", 64), + Count: 1, + StoredBytes: 192, + }, + }, + }}, + Metrics: builder.Finalize(), + }}, + } +} + +func mustObserveNode(builder *metrics.Builder, node entity.Node) { + if err := builder.ObserveNode(node); err != nil { + panic(err) + } +} + +func mustObserveRelationship(builder *metrics.Builder, relationship entity.Relationship) { + if err := builder.ObserveRelationship(relationship); err != nil { + panic(err) + } +} diff --git a/ret/collection/paths.go b/ret/collection/paths.go new file mode 100644 index 00000000..3bea8fcf --- /dev/null +++ b/ret/collection/paths.go @@ -0,0 +1,94 @@ +package collection + +import ( + "fmt" + "net/url" + "path" + "path/filepath" + "strings" + + "github.com/specterops/dawgs/ret/jsonl" +) + +func SafeJoin(root, relative string) (string, error) { + if root == "" { + return "", fmt.Errorf("safe join root is empty") + } + if relative == "" || relative == "." { + return "", fmt.Errorf("safe join path is empty") + } + if strings.Contains(relative, `\`) { + return "", fmt.Errorf("safe join path contains a backslash: %q", relative) + } + if strings.ContainsRune(relative, '\x00') { + return "", fmt.Errorf("safe join path contains NUL: %q", relative) + } + if path.IsAbs(relative) || filepath.IsAbs(relative) { + return "", fmt.Errorf("safe join path is absolute: %q", relative) + } + if clean := path.Clean(relative); clean != relative { + return "", fmt.Errorf("safe join path is not clean: %q", relative) + } + if relative == ".." || strings.HasPrefix(relative, "../") { + return "", fmt.Errorf("safe join path traverses its root: %q", relative) + } + + absoluteRoot, err := filepath.Abs(root) + if err != nil { + return "", fmt.Errorf("resolve safe join root: %w", err) + } + joined := filepath.Join(absoluteRoot, filepath.FromSlash(relative)) + contained, err := filepath.Rel(absoluteRoot, joined) + if err != nil { + return "", fmt.Errorf("check safe join containment: %w", err) + } + if contained == "." || contained == ".." || strings.HasPrefix(contained, ".."+string(filepath.Separator)) { + return "", fmt.Errorf("safe join path escapes its root: %q", relative) + } + + return joined, nil +} + +func NodeJSONLPath(graph string, shard int, codec jsonl.Codec) string { + return entityJSONLPath(graph, "nodes", shard, codec) +} + +func RelationshipJSONLPath(graph string, shard int, codec jsonl.Codec) string { + return entityJSONLPath(graph, "relationships", shard, codec) +} + +func NodeParquetPath(graph string, shard int) string { + return entityParquetPath(graph, "nodes", shard) +} + +func RelationshipParquetPath(graph string, shard int) string { + return entityParquetPath(graph, "relationships", shard) +} + +func entityJSONLPath(graph, entityType string, shard int, codec jsonl.Codec) string { + var suffix string + switch codec { + case jsonl.CodecNone: + suffix = ".jsonl" + case jsonl.CodecGzip: + suffix = ".jsonl.gz" + case jsonl.CodecZstd: + suffix = ".jsonl.zst" + default: + panic(fmt.Sprintf("unsupported JSONL codec %q", codec)) + } + + return entityPath(graph, entityType, shard, suffix) +} + +func entityParquetPath(graph, entityType string, shard int) string { + return entityPath(graph, entityType, shard, ".parquet") +} + +func entityPath(graph, entityType string, shard int, suffix string) string { + if shard < 1 { + panic(fmt.Sprintf("shard index must be at least one: %d", shard)) + } + + return fmt.Sprintf("graphs/%s/%s/%06d%s", url.PathEscape(graph), entityType, shard, suffix) +} diff --git a/ret/collection/shard.go b/ret/collection/shard.go new file mode 100644 index 00000000..f7054710 --- /dev/null +++ b/ret/collection/shard.go @@ -0,0 +1,35 @@ +package collection + +import ( + "github.com/specterops/dawgs/ret/jsonl" + "github.com/specterops/dawgs/ret/parquet" + "github.com/specterops/dawgs/ret/scrub" +) + +type JSONLArtifact struct { + Path string `json:"path"` + jsonl.Artifact +} + +type ParquetArtifact struct { + Path string `json:"path"` + parquet.Artifact +} + +type NodeShard struct { + Index int `json:"index"` + Count int64 `json:"count"` + LastSourceID uint64 `json:"last_source_id"` + ScrubCounts scrub.ActionCounts `json:"scrub_counts"` + JSONL *JSONLArtifact `json:"jsonl,omitempty"` + Parquet *ParquetArtifact `json:"parquet,omitempty"` +} + +type RelationshipShard struct { + Index int `json:"index"` + Count int64 `json:"count"` + LastSourceID uint64 `json:"last_source_id"` + ScrubCounts scrub.ActionCounts `json:"scrub_counts"` + JSONL *JSONLArtifact `json:"jsonl,omitempty"` + Parquet *ParquetArtifact `json:"parquet,omitempty"` +} diff --git a/ret/collection/verify.go b/ret/collection/verify.go new file mode 100644 index 00000000..2326252a --- /dev/null +++ b/ret/collection/verify.go @@ -0,0 +1,733 @@ +package collection + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "math" + "math/big" + "os" + "path/filepath" + "reflect" + "slices" + "sort" + "strconv" + "strings" + + "github.com/specterops/dawgs/ret/entity" + "github.com/specterops/dawgs/ret/metrics" + "github.com/specterops/dawgs/ret/observe" +) + +type Verification struct { + Manifest Manifest + Graphs []GraphVerification +} + +type GraphVerification struct { + Name string + NodeCount int64 + RelationshipCount int64 +} + +// Verify validates every configured concrete artifact in a collection. +func Verify(ctx context.Context, root string, observer observe.Observer) (Verification, error) { + manifest, err := readManifestWithoutSymlinks(root) + if err != nil { + return Verification{}, err + } + if err := preflightArtifactPaths(root, manifest, false); err != nil { + return Verification{}, err + } + + result := Verification{Manifest: manifest, Graphs: make([]GraphVerification, 0, len(manifest.Graphs))} + for _, graph := range manifest.Graphs { + if err := ctx.Err(); err != nil { + return Verification{}, fmt.Errorf("verify collection: %w", err) + } + graphResult, err := verifyGraph(ctx, root, graph, observer) + if err != nil { + return Verification{}, err + } + result.Graphs = append(result.Graphs, graphResult) + } + return result, nil +} + +// VerifyJSONLForLoad validates only JSONL artifacts and requires one for every +// logical shard. Parquet paths and bytes are deliberately not inspected. +func VerifyJSONLForLoad(ctx context.Context, root string, observer observe.Observer) (Verification, error) { + manifest, err := readManifestWithoutSymlinks(root) + if err != nil { + return Verification{}, err + } + if err := preflightArtifactPaths(root, manifest, true); err != nil { + return Verification{}, err + } + + result := Verification{Manifest: manifest, Graphs: make([]GraphVerification, 0, len(manifest.Graphs))} + for _, graph := range manifest.Graphs { + if err := ctx.Err(); err != nil { + return Verification{}, fmt.Errorf("verify JSONL collection for load: %w", err) + } + graphResult, err := verifyJSONLGraph(ctx, root, graph, observer) + if err != nil { + return Verification{}, err + } + result.Graphs = append(result.Graphs, graphResult) + } + return result, nil +} + +// ReplayGraph visits JSONL nodes in shard order, followed by JSONL +// relationships in shard order. Callers first run VerifyJSONLForLoad. +func ReplayGraph( + ctx context.Context, + root string, + graph Graph, + visitNode func(entity.Node) error, + visitRelationship func(entity.Relationship) error, +) error { + for _, shard := range graph.NodeShards { + if shard.JSONL == nil { + return fmt.Errorf("replay graph %q node shard %d requires JSONL output", graph.Name, shard.Index) + } + if err := inspectNonSymlinkArtifact(root, shard.JSONL.Path); err != nil { + return fmt.Errorf("replay graph %q JSONL node shard %d path: %w", graph.Name, shard.Index, err) + } + } + for _, shard := range graph.RelationshipShards { + if shard.JSONL == nil { + return fmt.Errorf("replay graph %q relationship shard %d requires JSONL output", graph.Name, shard.Index) + } + if err := inspectNonSymlinkArtifact(root, shard.JSONL.Path); err != nil { + return fmt.Errorf("replay graph %q JSONL relationship shard %d path: %w", graph.Name, shard.Index, err) + } + } + + for _, shard := range graph.NodeShards { + if err := ctx.Err(); err != nil { + return fmt.Errorf("replay graph %q nodes: %w", graph.Name, err) + } + err := ReadJSONLNodes(root, *shard.JSONL, func(node entity.Node) error { + if err := ctx.Err(); err != nil { + return err + } + if visitNode != nil { + if err := visitNode(node); err != nil { + return err + } + } + return nil + }) + if err != nil { + return fmt.Errorf("replay graph %q JSONL node shard %d: %w", graph.Name, shard.Index, err) + } + } + for _, shard := range graph.RelationshipShards { + if err := ctx.Err(); err != nil { + return fmt.Errorf("replay graph %q relationships: %w", graph.Name, err) + } + err := ReadJSONLRelationships(root, *shard.JSONL, func(relationship entity.Relationship) error { + if err := ctx.Err(); err != nil { + return err + } + if visitRelationship != nil { + if err := visitRelationship(relationship); err != nil { + return err + } + } + return nil + }) + if err != nil { + return fmt.Errorf("replay graph %q JSONL relationship shard %d: %w", graph.Name, shard.Index, err) + } + } + return nil +} + +func verifyGraph(ctx context.Context, root string, graph Graph, observer observe.Observer) (GraphVerification, error) { + builder := metrics.NewBuilder() + catalog := newKindCatalog() + parquetRelationshipIDs := make(map[string]struct{}) + + for _, shard := range graph.NodeShards { + var jsonNodes []entity.Node + if shard.JSONL != nil { + err := ReadJSONLNodes(root, *shard.JSONL, func(node entity.Node) error { + if err := ctx.Err(); err != nil { + return err + } + if err := builder.ObserveNode(node); err != nil { + return err + } + catalog.observeNode(node) + if shard.Parquet != nil { + jsonNodes = append(jsonNodes, node) + } + return nil + }) + if err != nil { + return GraphVerification{}, fmt.Errorf("verify graph %q JSONL node shard %d: %w", graph.Name, shard.Index, err) + } + emitJSONLNodeArtifactVerified(ctx, observer, graph.Name, *shard.JSONL) + } + + if shard.Parquet != nil { + parquetRow := 0 + err := ReadParquetNodes(root, *shard.Parquet, func(node entity.Node) error { + if err := ctx.Err(); err != nil { + return err + } + parquetRow++ + if shard.JSONL != nil { + if parquetRow > len(jsonNodes) { + return fmt.Errorf("dual node row %d has no JSONL row", parquetRow) + } + if err := compareNodes(jsonNodes[parquetRow-1], node); err != nil { + return fmt.Errorf("dual node row %d differs: %w", parquetRow, err) + } + return nil + } + if err := builder.ObserveNode(node); err != nil { + return err + } + catalog.observeNode(node) + return nil + }) + if err != nil { + return GraphVerification{}, fmt.Errorf("verify graph %q Parquet node shard %d: %w", graph.Name, shard.Index, err) + } + if shard.JSONL != nil && parquetRow != len(jsonNodes) { + return GraphVerification{}, fmt.Errorf( + "verify graph %q Parquet node shard %d: dual row count differs: JSONL %d, Parquet %d", + graph.Name, shard.Index, len(jsonNodes), parquetRow, + ) + } + emitParquetNodeArtifactVerified(ctx, observer, graph.Name, *shard.Parquet) + } + } + + for _, shard := range graph.RelationshipShards { + var jsonRelationships []entity.Relationship + if shard.JSONL != nil { + err := ReadJSONLRelationships(root, *shard.JSONL, func(relationship entity.Relationship) error { + if err := ctx.Err(); err != nil { + return err + } + if err := builder.ObserveRelationship(relationship); err != nil { + return err + } + catalog.observeRelationship(relationship) + if shard.Parquet != nil { + jsonRelationships = append(jsonRelationships, relationship) + } + return nil + }) + if err != nil { + return GraphVerification{}, fmt.Errorf("verify graph %q JSONL relationship shard %d: %w", graph.Name, shard.Index, err) + } + emitJSONLRelationshipArtifactVerified(ctx, observer, graph.Name, *shard.JSONL) + } + + if shard.Parquet != nil { + parquetRow := 0 + err := ReadParquetRelationships(root, *shard.Parquet, func(relationship entity.Relationship) error { + if err := ctx.Err(); err != nil { + return err + } + parquetRow++ + if relationship.SourceID == "" { + return fmt.Errorf("Parquet relationship row %d source ID is empty", parquetRow) + } + if _, found := parquetRelationshipIDs[relationship.SourceID]; found { + return fmt.Errorf("Parquet relationship source ID %q is duplicate within graph", relationship.SourceID) + } + parquetRelationshipIDs[relationship.SourceID] = struct{}{} + + if shard.JSONL != nil { + if parquetRow > len(jsonRelationships) { + return fmt.Errorf("dual relationship row %d has no JSONL row", parquetRow) + } + if err := compareRelationships(jsonRelationships[parquetRow-1], relationship); err != nil { + return fmt.Errorf("dual relationship row %d differs: %w", parquetRow, err) + } + return nil + } + if err := builder.ObserveRelationship(relationship); err != nil { + return err + } + catalog.observeRelationship(relationship) + return nil + }) + if err != nil { + return GraphVerification{}, fmt.Errorf("verify graph %q Parquet relationship shard %d: %w", graph.Name, shard.Index, err) + } + if shard.JSONL != nil && parquetRow != len(jsonRelationships) { + return GraphVerification{}, fmt.Errorf( + "verify graph %q Parquet relationship shard %d: dual row count differs: JSONL %d, Parquet %d", + graph.Name, shard.Index, len(jsonRelationships), parquetRow, + ) + } + emitParquetRelationshipArtifactVerified(ctx, observer, graph.Name, *shard.Parquet) + } + } + + return finalizeGraphVerification(graph, builder, catalog.values) +} + +func verifyJSONLGraph(ctx context.Context, root string, graph Graph, observer observe.Observer) (GraphVerification, error) { + builder := metrics.NewBuilder() + catalog := newKindCatalog() + + for _, shard := range graph.NodeShards { + err := ReadJSONLNodes(root, *shard.JSONL, func(node entity.Node) error { + if err := ctx.Err(); err != nil { + return err + } + if err := builder.ObserveNode(node); err != nil { + return err + } + catalog.observeNode(node) + return nil + }) + if err != nil { + return GraphVerification{}, fmt.Errorf("verify graph %q JSONL node shard %d for load: %w", graph.Name, shard.Index, err) + } + emitJSONLNodeArtifactVerified(ctx, observer, graph.Name, *shard.JSONL) + } + + for _, shard := range graph.RelationshipShards { + err := ReadJSONLRelationships(root, *shard.JSONL, func(relationship entity.Relationship) error { + if err := ctx.Err(); err != nil { + return err + } + if err := builder.ObserveRelationship(relationship); err != nil { + return err + } + catalog.observeRelationship(relationship) + return nil + }) + if err != nil { + return GraphVerification{}, fmt.Errorf("verify graph %q JSONL relationship shard %d for load: %w", graph.Name, shard.Index, err) + } + emitJSONLRelationshipArtifactVerified(ctx, observer, graph.Name, *shard.JSONL) + } + + return finalizeGraphVerification(graph, builder, catalog.values) +} + +func finalizeGraphVerification(graph Graph, builder *metrics.Builder, catalog []string) (GraphVerification, error) { + if !slices.Equal(graph.KindCatalog, catalog) { + return GraphVerification{}, fmt.Errorf( + "verify graph %q kind catalog differs: expected %v, actual %v", + graph.Name, graph.KindCatalog, catalog, + ) + } + actualMetrics := builder.Finalize() + if err := metrics.Compare(graph.Metrics, actualMetrics); err != nil { + return GraphVerification{}, fmt.Errorf("verify graph %q metrics: %w", graph.Name, err) + } + if graph.NodeCount != actualMetrics.NodeCount { + return GraphVerification{}, fmt.Errorf( + "verify graph %q node count differs: expected %d, actual %d", + graph.Name, graph.NodeCount, actualMetrics.NodeCount, + ) + } + if graph.RelationshipCount != actualMetrics.RelationshipCount { + return GraphVerification{}, fmt.Errorf( + "verify graph %q relationship count differs: expected %d, actual %d", + graph.Name, graph.RelationshipCount, actualMetrics.RelationshipCount, + ) + } + return GraphVerification{ + Name: graph.Name, + NodeCount: actualMetrics.NodeCount, + RelationshipCount: actualMetrics.RelationshipCount, + }, nil +} + +type kindCatalog struct { + seen map[string]struct{} + values []string +} + +func newKindCatalog() *kindCatalog { + return &kindCatalog{seen: make(map[string]struct{})} +} + +func (s *kindCatalog) observeNode(node entity.Node) { + for _, kind := range node.Kinds { + s.add(kind) + } +} + +func (s *kindCatalog) observeRelationship(relationship entity.Relationship) { + s.add(relationship.Kind) +} + +func (s *kindCatalog) add(kind string) { + if _, found := s.seen[kind]; found { + return + } + s.seen[kind] = struct{}{} + s.values = append(s.values, kind) +} + +func compareNodes(jsonNode, parquetNode entity.Node) error { + if jsonNode.SourceID != parquetNode.SourceID { + return fmt.Errorf("source ID: JSONL %q, Parquet %q", jsonNode.SourceID, parquetNode.SourceID) + } + if !slices.Equal(jsonNode.Kinds, parquetNode.Kinds) { + return fmt.Errorf("kinds: JSONL %v, Parquet %v", jsonNode.Kinds, parquetNode.Kinds) + } + if err := compareProperties(jsonNode.Properties, parquetNode.Properties); err != nil { + return fmt.Errorf("properties: %w", err) + } + return nil +} + +func compareRelationships(jsonRelationship, parquetRelationship entity.Relationship) error { + if jsonRelationship.StartID != parquetRelationship.StartID { + return fmt.Errorf("start ID: JSONL %q, Parquet %q", jsonRelationship.StartID, parquetRelationship.StartID) + } + if jsonRelationship.EndID != parquetRelationship.EndID { + return fmt.Errorf("end ID: JSONL %q, Parquet %q", jsonRelationship.EndID, parquetRelationship.EndID) + } + if jsonRelationship.Kind != parquetRelationship.Kind { + return fmt.Errorf("kind: JSONL %q, Parquet %q", jsonRelationship.Kind, parquetRelationship.Kind) + } + if err := compareProperties(jsonRelationship.Properties, parquetRelationship.Properties); err != nil { + return fmt.Errorf("properties: %w", err) + } + return nil +} + +func compareProperties(jsonProperties, parquetProperties map[string]any) error { + jsonCanonical, err := canonicalJSONValue(jsonProperties) + if err != nil { + return fmt.Errorf("JSONL value is not JSON-compatible: %w", err) + } + parquetCanonical, err := canonicalJSONValue(parquetProperties) + if err != nil { + return fmt.Errorf("Parquet value is not JSON-compatible: %w", err) + } + if !bytes.Equal(jsonCanonical, parquetCanonical) { + return fmt.Errorf("JSONL and Parquet values differ") + } + return nil +} + +func canonicalJSONValue(value any) ([]byte, error) { + var output bytes.Buffer + if err := appendCanonicalJSONValue(&output, reflect.ValueOf(value), "$"); err != nil { + return nil, err + } + return output.Bytes(), nil +} + +func appendCanonicalJSONValue(output *bytes.Buffer, value reflect.Value, path string) error { + if !value.IsValid() { + output.WriteByte('z') + return nil + } + for value.Kind() == reflect.Interface { + if value.IsNil() { + output.WriteByte('z') + return nil + } + value = value.Elem() + } + + if value.CanInterface() { + if number, ok := value.Interface().(json.Number); ok { + canonical, err := canonicalDecimal(number.String()) + if err != nil { + return fmt.Errorf("%s has invalid JSON number %q: %w", path, number, err) + } + writeLengthPrefixed(output, 'n', canonical) + return nil + } + } + + switch value.Kind() { + case reflect.Bool: + if value.Bool() { + output.WriteString("b1") + } else { + output.WriteString("b0") + } + return nil + case reflect.String: + writeLengthPrefixed(output, 's', value.String()) + return nil + case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64: + canonical, err := canonicalDecimal(strconv.FormatInt(value.Int(), 10)) + if err != nil { + return fmt.Errorf("%s has invalid integer: %w", path, err) + } + writeLengthPrefixed(output, 'n', canonical) + return nil + case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64: + canonical, err := canonicalDecimal(strconv.FormatUint(value.Uint(), 10)) + if err != nil { + return fmt.Errorf("%s has invalid unsigned integer: %w", path, err) + } + writeLengthPrefixed(output, 'n', canonical) + return nil + case reflect.Float32, reflect.Float64: + number := value.Float() + if math.IsNaN(number) || math.IsInf(number, 0) { + return fmt.Errorf("%s has non-finite number", path) + } + canonical, err := canonicalDecimal(strconv.FormatFloat(number, 'g', -1, value.Type().Bits())) + if err != nil { + return fmt.Errorf("%s has invalid floating-point number: %w", path, err) + } + writeLengthPrefixed(output, 'n', canonical) + return nil + case reflect.Map: + if value.IsNil() { + output.WriteByte('z') + return nil + } + if value.Type().Key().Kind() != reflect.String { + return fmt.Errorf("%s has object key type %s, want string", path, value.Type().Key()) + } + keys := value.MapKeys() + sort.Slice(keys, func(left, right int) bool { + return keys[left].String() < keys[right].String() + }) + output.WriteByte('{') + for _, key := range keys { + writeLengthPrefixed(output, 'k', key.String()) + if err := appendCanonicalJSONValue(output, value.MapIndex(key), path+"."+key.String()); err != nil { + return err + } + } + output.WriteByte('}') + return nil + case reflect.Slice: + if value.IsNil() { + output.WriteByte('z') + return nil + } + if value.Type().Elem().Kind() == reflect.Uint8 { + return fmt.Errorf("%s has binary value, which JSON cannot represent", path) + } + fallthrough + case reflect.Array: + output.WriteByte('[') + for index := 0; index < value.Len(); index++ { + if err := appendCanonicalJSONValue(output, value.Index(index), fmt.Sprintf("%s[%d]", path, index)); err != nil { + return err + } + } + output.WriteByte(']') + return nil + default: + return fmt.Errorf("%s has unsupported type %s", path, value.Type()) + } +} + +func canonicalDecimal(value string) (string, error) { + if value == "" { + return "", fmt.Errorf("empty number") + } + sign := "" + if value[0] == '-' || value[0] == '+' { + if value[0] == '-' { + sign = "-" + } + value = value[1:] + if value == "" { + return "", fmt.Errorf("missing digits") + } + } + + exponent := new(big.Int) + if separator := strings.IndexAny(value, "eE"); separator >= 0 { + if strings.IndexAny(value[separator+1:], "eE") >= 0 { + return "", fmt.Errorf("multiple exponents") + } + exponentText := value[separator+1:] + if exponentText == "" { + return "", fmt.Errorf("missing exponent") + } + if _, ok := exponent.SetString(exponentText, 10); !ok { + return "", fmt.Errorf("invalid exponent") + } + value = value[:separator] + } + + whole, fraction, found := strings.Cut(value, ".") + if found && strings.Contains(fraction, ".") { + return "", fmt.Errorf("multiple decimal points") + } + if whole == "" && fraction == "" { + return "", fmt.Errorf("missing digits") + } + if whole == "" { + whole = "0" + } + if !decimalDigits(whole) || !decimalDigits(fraction) { + return "", fmt.Errorf("invalid decimal digits") + } + + digits := strings.TrimLeft(whole+fraction, "0") + if digits == "" { + return "0e0", nil + } + exponent.Sub(exponent, big.NewInt(int64(len(fraction)))) + trailing := len(digits) - len(strings.TrimRight(digits, "0")) + if trailing != 0 { + digits = digits[:len(digits)-trailing] + exponent.Add(exponent, big.NewInt(int64(trailing))) + } + return sign + digits + "e" + exponent.String(), nil +} + +func decimalDigits(value string) bool { + for _, digit := range value { + if digit < '0' || digit > '9' { + return false + } + } + return true +} + +func writeLengthPrefixed(output *bytes.Buffer, kind byte, value string) { + output.WriteByte(kind) + output.WriteString(strconv.Itoa(len(value))) + output.WriteByte(':') + output.WriteString(value) +} + +func readManifestWithoutSymlinks(root string) (Manifest, error) { + if err := inspectNonSymlinkArtifact(root, ManifestName); err != nil { + return Manifest{}, fmt.Errorf("verify collection manifest path: %w", err) + } + manifest, err := Read(root) + if err != nil { + return Manifest{}, fmt.Errorf("verify collection manifest: %w", err) + } + return manifest, nil +} + +func preflightArtifactPaths(root string, manifest Manifest, jsonlOnly bool) error { + for _, graph := range manifest.Graphs { + for _, shard := range graph.NodeShards { + if jsonlOnly && shard.JSONL == nil { + return fmt.Errorf("graph %q node shard %d requires JSONL output", graph.Name, shard.Index) + } + if shard.JSONL != nil { + if err := inspectNonSymlinkArtifact(root, shard.JSONL.Path); err != nil { + return fmt.Errorf("verify graph %q JSONL node shard %d path: %w", graph.Name, shard.Index, err) + } + } + if !jsonlOnly && shard.Parquet != nil { + if err := inspectNonSymlinkArtifact(root, shard.Parquet.Path); err != nil { + return fmt.Errorf("verify graph %q Parquet node shard %d path: %w", graph.Name, shard.Index, err) + } + } + } + for _, shard := range graph.RelationshipShards { + if jsonlOnly && shard.JSONL == nil { + return fmt.Errorf("graph %q relationship shard %d requires JSONL output", graph.Name, shard.Index) + } + if shard.JSONL != nil { + if err := inspectNonSymlinkArtifact(root, shard.JSONL.Path); err != nil { + return fmt.Errorf("verify graph %q JSONL relationship shard %d path: %w", graph.Name, shard.Index, err) + } + } + if !jsonlOnly && shard.Parquet != nil { + if err := inspectNonSymlinkArtifact(root, shard.Parquet.Path); err != nil { + return fmt.Errorf("verify graph %q Parquet relationship shard %d path: %w", graph.Name, shard.Index, err) + } + } + } + } + return nil +} + +// inspectNonSymlinkArtifact is a portable local-filesystem best effort. It +// checks every collection-relative path component before a reader opens the +// file, but pathname replacement cannot be made atomic on every supported OS. +func inspectNonSymlinkArtifact(root, relative string) error { + inspectionRoot := root + if inspectionRoot == "" { + inspectionRoot = "." + } + rootInfo, err := os.Lstat(inspectionRoot) + if err != nil { + return fmt.Errorf("inspect collection root: %w", err) + } + if rootInfo.Mode()&os.ModeSymlink != 0 || !rootInfo.IsDir() { + return fmt.Errorf("collection root is not a non-symlink directory: %q", root) + } + if _, err := SafeJoin(inspectionRoot, relative); err != nil { + return err + } + + components := strings.Split(filepath.FromSlash(relative), string(filepath.Separator)) + current := inspectionRoot + for index, component := range components { + current = filepath.Join(current, component) + info, err := os.Lstat(current) + if err != nil { + return fmt.Errorf("inspect artifact path component %q: %w", component, err) + } + if info.Mode()&os.ModeSymlink != 0 { + return fmt.Errorf("artifact path contains symlink component %q", component) + } + if index < len(components)-1 { + if !info.IsDir() { + return fmt.Errorf("artifact path component %q is not a directory", component) + } + continue + } + if !info.Mode().IsRegular() { + return fmt.Errorf("artifact path is not a regular file: %q", relative) + } + } + return nil +} + +func emitJSONLNodeArtifactVerified(ctx context.Context, observer observe.Observer, graph string, artifact JSONLArtifact) { + observe.Emit(ctx, observer, observe.ArtifactVerified{ + Graph: graph, EntityType: "node", Format: "JSONL", Path: artifact.Path, + Count: artifact.Count, Bytes: artifact.StoredBytes, + }) +} + +func emitJSONLRelationshipArtifactVerified( + ctx context.Context, + observer observe.Observer, + graph string, + artifact JSONLArtifact, +) { + observe.Emit(ctx, observer, observe.ArtifactVerified{ + Graph: graph, EntityType: "relationship", Format: "JSONL", Path: artifact.Path, + Count: artifact.Count, Bytes: artifact.StoredBytes, + }) +} + +func emitParquetNodeArtifactVerified(ctx context.Context, observer observe.Observer, graph string, artifact ParquetArtifact) { + observe.Emit(ctx, observer, observe.ArtifactVerified{ + Graph: graph, EntityType: "node", Format: "Parquet", Path: artifact.Path, + Count: artifact.Count, Bytes: artifact.StoredBytes, + }) +} + +func emitParquetRelationshipArtifactVerified( + ctx context.Context, + observer observe.Observer, + graph string, + artifact ParquetArtifact, +) { + observe.Emit(ctx, observer, observe.ArtifactVerified{ + Graph: graph, EntityType: "relationship", Format: "Parquet", Path: artifact.Path, + Count: artifact.Count, Bytes: artifact.StoredBytes, + }) +} diff --git a/ret/collection/verify_test.go b/ret/collection/verify_test.go new file mode 100644 index 00000000..150e9d30 --- /dev/null +++ b/ret/collection/verify_test.go @@ -0,0 +1,579 @@ +package collection_test + +import ( + "context" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/specterops/dawgs/graph" + "github.com/specterops/dawgs/ret/collection" + "github.com/specterops/dawgs/ret/entity" + "github.com/specterops/dawgs/ret/jsonl" + "github.com/specterops/dawgs/ret/metrics" + "github.com/specterops/dawgs/ret/observe" + "github.com/specterops/dawgs/ret/parquet" + "github.com/stretchr/testify/require" +) + +func TestVerifySupportsJSONLParquetDualAndEmptyGraphs(t *testing.T) { + nodes, relationships := verificationEntities() + + for _, testCase := range []struct { + name string + jsonl bool + parquet bool + }{ + {name: "JSONL only", jsonl: true}, + {name: "Parquet only", parquet: true}, + {name: "dual", jsonl: true, parquet: true}, + } { + t.Run(testCase.name, func(t *testing.T) { + root := writeVerificationCollection(t, nodes, nodes, relationships, relationships, testCase.jsonl, testCase.parquet) + + got, err := collection.Verify(context.Background(), root, nil) + + require.NoError(t, err) + require.Len(t, got.Graphs, 1) + require.Equal(t, collection.GraphVerification{Name: "example", NodeCount: 2, RelationshipCount: 1}, got.Graphs[0]) + }) + } + + t.Run("empty graph", func(t *testing.T) { + root := writeVerificationCollection(t, nil, nil, nil, nil, true, true) + + got, err := collection.Verify(context.Background(), root, nil) + + require.NoError(t, err) + require.Equal(t, []collection.GraphVerification{{Name: "example"}}, got.Graphs) + }) +} + +func TestVerifyJSONLForLoadIgnoresCorruptAndMissingParquet(t *testing.T) { + nodes, relationships := verificationEntities() + + t.Run("corrupt", func(t *testing.T) { + root := writeVerificationCollection(t, nodes, nodes, relationships, relationships, true, true) + manifest := readManifest(t, root) + require.NoError(t, os.WriteFile(filepath.Join(root, filepath.FromSlash(manifest.Graphs[0].NodeShards[0].Parquet.Path)), []byte("corrupt"), 0o600)) + + _, err := collection.VerifyJSONLForLoad(context.Background(), root, nil) + require.NoError(t, err) + + _, err = collection.Verify(context.Background(), root, nil) + require.Error(t, err) + }) + + t.Run("missing", func(t *testing.T) { + root := writeVerificationCollection(t, nodes, nodes, relationships, relationships, true, true) + manifest := readManifest(t, root) + for _, path := range parquetPaths(manifest) { + require.NoError(t, os.Remove(filepath.Join(root, filepath.FromSlash(path)))) + } + + _, err := collection.VerifyJSONLForLoad(context.Background(), root, nil) + require.NoError(t, err) + }) +} + +func TestVerifyJSONLForLoadRejectsParquetOnly(t *testing.T) { + nodes, relationships := verificationEntities() + root := writeVerificationCollection(t, nil, nodes, nil, relationships, false, true) + + _, err := collection.VerifyJSONLForLoad(context.Background(), root, nil) + + require.ErrorContains(t, err, "JSONL") +} + +func TestVerifyRejectsDuplicateNodesAndMissingRelationshipEndpoints(t *testing.T) { + nodes, relationships := verificationEntities() + + t.Run("duplicate node source ID", func(t *testing.T) { + root := writeVerificationCollection(t, nodes, nil, relationships, nil, true, false) + manifest := readManifest(t, root) + manifest.Graphs[0].NodeShards[0].JSONL = installJSONLNodes(t, root, "example", []entity.Node{ + nodes[0], + {SourceID: nodes[0].SourceID, Kinds: nodes[1].Kinds, Properties: nodes[1].Properties}, + }) + writeManifest(t, root, manifest) + + _, err := collection.Verify(context.Background(), root, nil) + + require.ErrorContains(t, err, "duplicate") + require.ErrorContains(t, err, nodes[0].SourceID) + }) + + t.Run("missing endpoint", func(t *testing.T) { + root := writeVerificationCollection(t, nodes, nil, relationships, nil, true, false) + manifest := readManifest(t, root) + broken := relationships + broken[0].EndID = "missing" + manifest.Graphs[0].RelationshipShards[0].JSONL = installJSONLRelationships(t, root, "example", broken) + writeManifest(t, root, manifest) + + _, err := collection.Verify(context.Background(), root, nil) + + require.ErrorContains(t, err, "missing endpoint") + }) +} + +func TestVerifyRequiresUniqueParquetRelationshipSourceIDs(t *testing.T) { + nodes, relationships := verificationEntities() + relationships = append(relationships, entity.Relationship{ + SourceID: "relationship-2", + StartID: "node-2", + EndID: "node-1", + Kind: "OWNS", + }) + root := writeVerificationCollection(t, nil, nodes, nil, relationships, false, true) + manifest := readManifest(t, root) + duplicate := append([]entity.Relationship(nil), relationships...) + duplicate[1].SourceID = duplicate[0].SourceID + manifest.Graphs[0].RelationshipShards[0].Parquet = installParquetRelationships(t, root, "example", duplicate) + writeManifest(t, root, manifest) + + _, err := collection.Verify(context.Background(), root, nil) + + require.ErrorContains(t, err, "duplicate") + require.ErrorContains(t, err, duplicate[0].SourceID) +} + +func TestVerifyRejectsCatalogAndMetricsMismatch(t *testing.T) { + nodes, relationships := verificationEntities() + + t.Run("first-seen catalog order", func(t *testing.T) { + root := writeVerificationCollection(t, nodes, nil, relationships, nil, true, false) + manifest := readManifest(t, root) + manifest.Graphs[0].KindCatalog[0], manifest.Graphs[0].KindCatalog[1] = + manifest.Graphs[0].KindCatalog[1], manifest.Graphs[0].KindCatalog[0] + writeManifest(t, root, manifest) + + _, err := collection.Verify(context.Background(), root, nil) + + require.ErrorContains(t, err, "kind catalog") + }) + + t.Run("metrics", func(t *testing.T) { + root := writeVerificationCollection(t, nodes, nil, relationships, nil, true, false) + manifest := readManifest(t, root) + delete(manifest.Graphs[0].Metrics.NodeKindSequences, metrics.OrderedKindsKey(nodes[0].Kinds)) + manifest.Graphs[0].Metrics.NodeKindSequences[metrics.OrderedKindsKey([]string{"Wrong"})] = 1 + writeManifest(t, root, manifest) + + _, err := collection.Verify(context.Background(), root, nil) + + require.ErrorContains(t, err, "metrics") + }) +} + +func TestVerifyComparesDualRowsAndCanonicalProperties(t *testing.T) { + nodes, relationships := verificationEntities() + + t.Run("equivalent numeric and container values", func(t *testing.T) { + nodes[0].Properties = map[string]any{ + "large": int64(9_007_199_254_740_993), + "whole": int64(42), + "nested": map[string]any{ + "array": []any{int64(1), 2.5, map[string]any{"enabled": true}}, + "null": nil, + }, + } + root := writeVerificationCollection(t, nodes, nodes, relationships, relationships, true, true) + + _, err := collection.Verify(context.Background(), root, nil) + + require.NoError(t, err) + }) + + t.Run("node value mismatch", func(t *testing.T) { + root := writeVerificationCollection(t, nodes, nodes, relationships, relationships, true, true) + manifest := readManifest(t, root) + different := append([]entity.Node(nil), nodes...) + different[0].Properties = map[string]any{"different": true} + manifest.Graphs[0].NodeShards[0].Parquet = installParquetNodes(t, root, "example", different) + writeManifest(t, root, manifest) + + _, err := collection.Verify(context.Background(), root, nil) + + require.ErrorContains(t, err, `graph "example"`) + require.ErrorContains(t, err, "node shard 1") + require.ErrorContains(t, err, "row 1") + }) + + t.Run("node order mismatch", func(t *testing.T) { + root := writeVerificationCollection(t, nodes, nodes, relationships, relationships, true, true) + manifest := readManifest(t, root) + reversed := []entity.Node{nodes[1], nodes[0]} + manifest.Graphs[0].NodeShards[0].Parquet = installParquetNodes(t, root, "example", reversed) + writeManifest(t, root, manifest) + + _, err := collection.Verify(context.Background(), root, nil) + + require.ErrorContains(t, err, "row 1") + }) + + t.Run("relationship value mismatch", func(t *testing.T) { + root := writeVerificationCollection(t, nodes, nodes, relationships, relationships, true, true) + manifest := readManifest(t, root) + different := append([]entity.Relationship(nil), relationships...) + different[0].Properties = map[string]any{"different": true} + manifest.Graphs[0].RelationshipShards[0].Parquet = installParquetRelationships(t, root, "example", different) + writeManifest(t, root, manifest) + + _, err := collection.Verify(context.Background(), root, nil) + + require.ErrorContains(t, err, "relationship shard 1") + require.ErrorContains(t, err, "row 1") + }) + + t.Run("non-JSON-compatible Parquet value", func(t *testing.T) { + nodes[0].Properties = map[string]any{"binary": []byte{0x01, 0x02}} + root := writeVerificationCollection(t, nodes, nodes, relationships, relationships, true, true) + + _, err := collection.Verify(context.Background(), root, nil) + + require.ErrorContains(t, err, "not JSON-compatible") + }) +} + +func TestVerifyRejectsSymlinksBeforeReadingAnyArtifact(t *testing.T) { + nodes, relationships := verificationEntities() + root := writeVerificationCollection(t, nodes, nodes, relationships, relationships, true, true) + manifest := readManifest(t, root) + target := filepath.Join(root, "outside.parquet") + require.NoError(t, os.WriteFile(target, []byte("outside"), 0o600)) + link := filepath.Join(root, filepath.FromSlash(manifest.Graphs[0].RelationshipShards[0].Parquet.Path)) + require.NoError(t, os.Remove(link)) + require.NoError(t, os.Symlink(target, link)) + var events []observe.Event + + _, err := collection.Verify(context.Background(), root, observe.ObserverFunc(func(_ context.Context, event observe.Event) { + events = append(events, event) + })) + + require.ErrorContains(t, err, "symlink") + require.Empty(t, events, "all artifact paths must be checked before any reader emits success") +} + +func TestArtifactVerifiedIsEmittedOnlyAfterConcreteReaderSucceeds(t *testing.T) { + nodes, relationships := verificationEntities() + root := writeVerificationCollection(t, nodes, nil, relationships, nil, true, false) + manifest := readManifest(t, root) + manifest.Graphs[0].NodeShards[0].JSONL = installJSONLNodes(t, root, "example", []entity.Node{ + nodes[0], + {SourceID: nodes[0].SourceID, Kinds: nodes[1].Kinds}, + }) + writeManifest(t, root, manifest) + var events []observe.Event + + _, err := collection.Verify(context.Background(), root, observe.ObserverFunc(func(_ context.Context, event observe.Event) { + events = append(events, event) + })) + + require.Error(t, err) + require.Empty(t, events) +} + +func TestVerifyEmitsArtifactsInConcreteReaderCompletionOrder(t *testing.T) { + nodes, relationships := verificationEntities() + root := writeVerificationCollection(t, nodes, nodes, relationships, relationships, true, true) + var got []string + + _, err := collection.Verify(context.Background(), root, observe.ObserverFunc(func(_ context.Context, event observe.Event) { + if verified, ok := event.(observe.ArtifactVerified); ok { + got = append(got, verified.EntityType+":"+verified.Format) + } + })) + + require.NoError(t, err) + require.Equal(t, []string{ + "node:JSONL", + "node:Parquet", + "relationship:JSONL", + "relationship:Parquet", + }, got) +} + +func TestReplayGraphVisitsAllNodesBeforeRelationshipsAndNeverTouchesParquet(t *testing.T) { + nodes, relationships := verificationEntities() + root := writeVerificationCollection(t, nodes, nodes, relationships, relationships, true, true) + manifest := readManifest(t, root) + for _, path := range parquetPaths(manifest) { + require.NoError(t, os.Remove(filepath.Join(root, filepath.FromSlash(path)))) + } + _, err := collection.VerifyJSONLForLoad(context.Background(), root, nil) + require.NoError(t, err) + var order []string + + err = collection.ReplayGraph( + context.Background(), + root, + manifest.Graphs[0], + func(node entity.Node) error { + order = append(order, "node:"+node.SourceID) + return nil + }, + func(relationship entity.Relationship) error { + order = append(order, "relationship:"+relationship.Kind) + return nil + }, + ) + + require.NoError(t, err) + require.Equal(t, []string{"node:node-1", "node:node-2", "relationship:MEMBER_OF"}, order) +} + +func TestReplayGraphReturnsBackendSupportedNativeNumbers(t *testing.T) { + nodes, relationships := verificationEntities() + nodes[0].Properties = map[string]any{ + "integer": int64(9_007_199_254_740_993), + "fraction": 1.25, + "nested": []any{int64(2)}, + } + root := writeVerificationCollection(t, nodes, nil, relationships, nil, true, false) + manifest := readManifest(t, root) + _, err := collection.VerifyJSONLForLoad(context.Background(), root, nil) + require.NoError(t, err) + + err = collection.ReplayGraph( + context.Background(), + root, + manifest.Graphs[0], + func(node entity.Node) error { + if node.SourceID != "node-1" { + return nil + } + require.IsType(t, int64(0), node.Properties["integer"]) + require.Equal(t, int64(9_007_199_254_740_993), node.Properties["integer"]) + require.IsType(t, float64(0), node.Properties["fraction"]) + nested := node.Properties["nested"].([]any) + require.IsType(t, int64(0), nested[0]) + mapped, mapErr := graph.AsProperties(node.Properties).Get("integer").Int64() + require.NoError(t, mapErr) + require.Equal(t, int64(9_007_199_254_740_993), mapped) + return nil + }, + nil, + ) + + require.NoError(t, err) +} + +func verificationEntities() ([]entity.Node, []entity.Relationship) { + return []entity.Node{ + {SourceID: "node-1", Kinds: []string{"User", "Principal"}, Properties: map[string]any{"name": "Alice"}}, + {SourceID: "node-2", Kinds: []string{"Group"}, Properties: map[string]any{"name": "Admins"}}, + }, []entity.Relationship{{ + SourceID: "relationship-1", + StartID: "node-1", + EndID: "node-2", + Kind: "MEMBER_OF", + Properties: map[string]any{"active": true}, + }} +} + +func writeVerificationCollection( + t *testing.T, + jsonNodes, parquetNodes []entity.Node, + jsonRelationships, parquetRelationships []entity.Relationship, + withJSONL, withParquet bool, +) string { + t.Helper() + root := t.TempDir() + graph := collection.Graph{Name: "example"} + outputs := collection.OutputConfig{} + if withJSONL { + outputs.JSONL = &collection.JSONLOutput{SchemaVersion: jsonl.SchemaVersion, Codec: string(jsonl.CodecNone)} + } + if withParquet { + outputs.Parquet = &collection.ParquetOutput{SchemaVersion: parquet.SchemaVersion} + } + + canonicalNodes := parquetNodes + canonicalRelationships := parquetRelationships + if withJSONL { + canonicalNodes = jsonNodes + canonicalRelationships = jsonRelationships + } + graph.NodeCount = int64(len(canonicalNodes)) + graph.RelationshipCount = int64(len(canonicalRelationships)) + graph.KindCatalog = firstSeenCatalog(canonicalNodes, canonicalRelationships) + graph.Metrics = buildMetrics(t, canonicalNodes, canonicalRelationships) + + if len(canonicalNodes) != 0 { + shard := collection.NodeShard{Index: 1, Count: int64(len(canonicalNodes)), LastSourceID: 100} + if withJSONL { + shard.JSONL = installJSONLNodes(t, root, graph.Name, jsonNodes) + } + if withParquet { + shard.Parquet = installParquetNodes(t, root, graph.Name, parquetNodes) + } + graph.NodeShards = []collection.NodeShard{shard} + } + if len(canonicalRelationships) != 0 { + shard := collection.RelationshipShard{Index: 1, Count: int64(len(canonicalRelationships)), LastSourceID: 200} + if withJSONL { + shard.JSONL = installJSONLRelationships(t, root, graph.Name, jsonRelationships) + } + if withParquet { + shard.Parquet = installParquetRelationships(t, root, graph.Name, parquetRelationships) + } + graph.RelationshipShards = []collection.RelationshipShard{shard} + } + + writeManifest(t, root, collection.Manifest{ + Format: collection.Format, + CreatedAt: time.Date(2026, time.July, 28, 12, 0, 0, 0, time.UTC), + Outputs: outputs, + Graphs: []collection.Graph{graph}, + }) + return root +} + +func installJSONLNodes(t *testing.T, root, graph string, nodes []entity.Node) *collection.JSONLArtifact { + t.Helper() + path := collection.NodeJSONLPath(graph, 1, jsonl.CodecNone) + temporary := filepath.Join(root, "nodes.jsonl.tmp") + file, err := os.Create(temporary) + require.NoError(t, err) + writer, err := jsonl.NewNodeWriter(file, jsonl.Config{Codec: jsonl.CodecNone}) + require.NoError(t, err) + require.NoError(t, writer.Push(nodes)) + require.NoError(t, writer.Close()) + artifact, err := writer.Result() + require.NoError(t, err) + require.NoError(t, file.Close()) + installArtifact(t, root, temporary, path) + return &collection.JSONLArtifact{Path: path, Artifact: artifact} +} + +func installJSONLRelationships(t *testing.T, root, graph string, relationships []entity.Relationship) *collection.JSONLArtifact { + t.Helper() + path := collection.RelationshipJSONLPath(graph, 1, jsonl.CodecNone) + temporary := filepath.Join(root, "relationships.jsonl.tmp") + file, err := os.Create(temporary) + require.NoError(t, err) + writer, err := jsonl.NewRelationshipWriter(file, jsonl.Config{Codec: jsonl.CodecNone}) + require.NoError(t, err) + require.NoError(t, writer.Push(relationships)) + require.NoError(t, writer.Close()) + artifact, err := writer.Result() + require.NoError(t, err) + require.NoError(t, file.Close()) + installArtifact(t, root, temporary, path) + return &collection.JSONLArtifact{Path: path, Artifact: artifact} +} + +func installParquetNodes(t *testing.T, root, graph string, nodes []entity.Node) *collection.ParquetArtifact { + t.Helper() + path := collection.NodeParquetPath(graph, 1) + temporary := filepath.Join(root, "nodes.parquet.tmp") + file, err := os.Create(temporary) + require.NoError(t, err) + writer, err := parquet.NewNodeWriter(file, parquet.Config{}) + require.NoError(t, err) + require.NoError(t, writer.Push(nodes)) + require.NoError(t, writer.Close()) + artifact, err := writer.Result() + require.NoError(t, err) + require.NoError(t, file.Close()) + installArtifact(t, root, temporary, path) + return &collection.ParquetArtifact{Path: path, Artifact: artifact} +} + +func installParquetRelationships(t *testing.T, root, graph string, relationships []entity.Relationship) *collection.ParquetArtifact { + t.Helper() + path := collection.RelationshipParquetPath(graph, 1) + temporary := filepath.Join(root, "relationships.parquet.tmp") + file, err := os.Create(temporary) + require.NoError(t, err) + writer, err := parquet.NewRelationshipWriter(file, parquet.Config{}) + require.NoError(t, err) + require.NoError(t, writer.Push(relationships)) + require.NoError(t, writer.Close()) + artifact, err := writer.Result() + require.NoError(t, err) + require.NoError(t, file.Close()) + installArtifact(t, root, temporary, path) + return &collection.ParquetArtifact{Path: path, Artifact: artifact} +} + +func installArtifact(t *testing.T, root, temporary, relative string) { + t.Helper() + final := filepath.Join(root, filepath.FromSlash(relative)) + require.NoError(t, os.MkdirAll(filepath.Dir(final), 0o700)) + require.NoError(t, os.Rename(temporary, final)) +} + +func writeManifest(t *testing.T, root string, manifest collection.Manifest) { + t.Helper() + require.NoError(t, collection.Write(root, manifest)) +} + +func readManifest(t *testing.T, root string) collection.Manifest { + t.Helper() + manifest, err := collection.Read(root) + require.NoError(t, err) + return manifest +} + +func buildMetrics(t *testing.T, nodes []entity.Node, relationships []entity.Relationship) metrics.GraphMetrics { + t.Helper() + builder := metrics.NewBuilder() + for _, node := range nodes { + require.NoError(t, builder.ObserveNode(node)) + } + for _, relationship := range relationships { + require.NoError(t, builder.ObserveRelationship(relationship)) + } + return builder.Finalize() +} + +func firstSeenCatalog(nodes []entity.Node, relationships []entity.Relationship) []string { + seen := map[string]struct{}{} + var result []string + add := func(kind string) { + if _, found := seen[kind]; !found { + seen[kind] = struct{}{} + result = append(result, kind) + } + } + for _, node := range nodes { + for _, kind := range node.Kinds { + add(kind) + } + } + for _, relationship := range relationships { + add(relationship.Kind) + } + return result +} + +func parquetPaths(manifest collection.Manifest) []string { + var result []string + for _, graph := range manifest.Graphs { + for _, shard := range graph.NodeShards { + if shard.Parquet != nil { + result = append(result, shard.Parquet.Path) + } + } + for _, shard := range graph.RelationshipShards { + if shard.Parquet != nil { + result = append(result, shard.Parquet.Path) + } + } + } + return result +} + +func TestVerifyJSONLForLoadErrorNamesMissingJSONLShard(t *testing.T) { + nodes, _ := verificationEntities() + root := writeVerificationCollection(t, nil, nodes, nil, nil, false, true) + + _, err := collection.VerifyJSONLForLoad(context.Background(), root, nil) + + require.Error(t, err) + require.True(t, strings.Contains(err.Error(), "node shard 1") && strings.Contains(err.Error(), "JSONL")) +} diff --git a/ret/config.go b/ret/config.go new file mode 100644 index 00000000..08628e3e --- /dev/null +++ b/ret/config.go @@ -0,0 +1,113 @@ +package ret + +import ( + "errors" + "fmt" + "path" + "strings" + + "github.com/specterops/dawgs/ret/jsonl" + "github.com/specterops/dawgs/ret/observe" + "github.com/specterops/dawgs/ret/parquet" + "github.com/specterops/dawgs/ret/scrub" +) + +type DumpConfig struct { + Directory string + Graphs []string + EntityBatchSize int + ShardSize int + Resume bool + JSONL *jsonl.Config + Parquet *parquet.Config + Scrub *scrub.Config // Nil disables scrubbing. + Observer observe.Observer +} + +type LoadConfig struct { + Directory string + BatchSize int + Observer observe.Observer +} + +type VerifyCollectionConfig struct { + Directory string + Observer observe.Observer +} + +type VerifyDatabaseConfig struct { + Directory string + BatchSize int + Observer observe.Observer +} + +func (s DumpConfig) Validate() error { + if strings.TrimSpace(s.Directory) == "" || len(s.Graphs) == 0 { + return fmt.Errorf("%w: dump directory and at least one graph are required", ErrInvalidConfig) + } + if err := validateGraphNames(s.Graphs); err != nil { + return fmt.Errorf("%w: dump graphs: %w", ErrInvalidConfig, err) + } + if s.EntityBatchSize <= 0 || s.ShardSize <= 0 { + return fmt.Errorf("%w: batch and shard sizes must be positive", ErrInvalidConfig) + } + if s.JSONL == nil && s.Parquet == nil { + return fmt.Errorf("%w: at least one output is required", ErrInvalidConfig) + } + var scrubErr error + if s.Scrub != nil { + scrubErr = s.Scrub.Validate() + } + var jsonlErr, parquetErr error + if s.JSONL != nil { + jsonlErr = s.JSONL.Validate() + } + if s.Parquet != nil { + parquetErr = s.Parquet.Validate() + } + if err := errors.Join(jsonlErr, parquetErr, scrubErr); err != nil { + return fmt.Errorf("%w: dump output and scrub configuration: %w", ErrInvalidConfig, err) + } + return nil +} + +func (s LoadConfig) Validate() error { + if strings.TrimSpace(s.Directory) == "" { + return fmt.Errorf("%w: load directory is required", ErrInvalidConfig) + } + if s.BatchSize <= 0 { + return fmt.Errorf("%w: load batch size must be positive", ErrInvalidConfig) + } + return nil +} + +func (s VerifyCollectionConfig) Validate() error { + if strings.TrimSpace(s.Directory) == "" { + return fmt.Errorf("%w: collection directory is required", ErrInvalidConfig) + } + return nil +} + +func (s VerifyDatabaseConfig) Validate() error { + if strings.TrimSpace(s.Directory) == "" { + return fmt.Errorf("%w: database verification directory is required", ErrInvalidConfig) + } + if s.BatchSize <= 0 { + return fmt.Errorf("%w: database verification batch size must be positive", ErrInvalidConfig) + } + return nil +} + +func validateGraphNames(graphs []string) error { + seen := make(map[string]struct{}, len(graphs)) + for _, graph := range graphs { + if strings.TrimSpace(graph) == "" || graph == "." || graph == ".." || path.Clean(graph) != graph || strings.ContainsAny(graph, "/\\") || strings.ContainsRune(graph, '\x00') { + return fmt.Errorf("graph name %q is not a safe path segment", graph) + } + if _, found := seen[graph]; found { + return fmt.Errorf("duplicate graph name %q", graph) + } + seen[graph] = struct{}{} + } + return nil +} diff --git a/ret/config_test.go b/ret/config_test.go new file mode 100644 index 00000000..78680869 --- /dev/null +++ b/ret/config_test.go @@ -0,0 +1,115 @@ +package ret + +import ( + "testing" + + "github.com/specterops/dawgs/ret/jsonl" + "github.com/specterops/dawgs/ret/parquet" + "github.com/specterops/dawgs/ret/scrub" + "github.com/stretchr/testify/require" +) + +func TestDumpConfigRequiresAnOutput(t *testing.T) { + // Break caught: accepting a dump that has no artifact writer enabled. + config := validDumpConfig(t) + config.JSONL = nil + config.Parquet = nil + + require.ErrorContains(t, config.Validate(), "output") +} + +func TestParquetOnlyDumpConfigIsValid(t *testing.T) { + // Break caught: rejecting a valid concrete Parquet-only dump. + config := validDumpConfig(t) + config.JSONL = nil + config.Parquet = pointerTo(parquet.Config{}) + + require.NoError(t, config.Validate()) +} + +func TestDumpConfigRejectsDuplicateAndUnsafeGraphNames(t *testing.T) { + // Break caught: accepting graph names that would collide or escape collection paths. + for _, graphs := range [][]string{{"asset", "asset"}, {"../asset"}, {"asset/node"}, {"."}} { + config := validDumpConfig(t) + config.Graphs = graphs + + require.ErrorIs(t, config.Validate(), ErrInvalidConfig) + } +} + +func TestDumpConfigRejectsNonPositiveBatchAndShardSizes(t *testing.T) { + // Break caught: accepting sizes that cannot make forward progress during a dump. + for _, mutate := range []func(*DumpConfig){ + func(config *DumpConfig) { config.EntityBatchSize = 0 }, + func(config *DumpConfig) { config.ShardSize = 0 }, + } { + config := validDumpConfig(t) + mutate(&config) + + require.ErrorIs(t, config.Validate(), ErrInvalidConfig) + } +} + +func TestDumpConfigReturnsFormatAndScrubValidationErrors(t *testing.T) { + // Break caught: allowing invalid delegated output or scrub configurations into a dump. + for _, mutate := range []func(*DumpConfig){ + func(config *DumpConfig) { config.JSONL.Codec = jsonl.Codec("zip") }, + func(config *DumpConfig) { config.JSONL.Level = 99 }, + func(config *DumpConfig) { + config.Scrub.Rules.Classifier.ValueShapePatterns = []scrub.ValueShapeConfig{{Name: "invalid", Pattern: "("}} + }, + } { + config := validDumpConfig(t) + mutate(&config) + + require.ErrorIs(t, config.Validate(), ErrInvalidConfig) + } +} + +func TestDumpConfigAllowsDisabledScrubbing(t *testing.T) { + // Break caught: requiring a scrub policy when nil is the library-level + // opt-out from scrubbing. + config := validDumpConfig(t) + config.Scrub = nil + + require.NoError(t, config.Validate()) +} + +func TestDumpConfigRetainsEveryDelegatedValidationCause(t *testing.T) { + // Break caught: classifying a root validation error while silently discarding one delegated failure. + config := validDumpConfig(t) + config.JSONL.Codec = jsonl.Codec("zip") + config.Scrub.Rules.Classifier.ValueShapePatterns = []scrub.ValueShapeConfig{{Name: "invalid", Pattern: "("}} + + err := config.Validate() + require.ErrorIs(t, err, ErrInvalidConfig) + require.ErrorContains(t, err, `unsupported JSONL codec "zip"`) + require.ErrorContains(t, err, `compile value shape "invalid"`) +} + +func TestOtherFacadeConfigsValidateRequiredInputs(t *testing.T) { + // Break caught: allowing a later load or verification operation to start without its required path or batch size. + require.ErrorIs(t, (LoadConfig{Directory: t.TempDir()}).Validate(), ErrInvalidConfig) + require.ErrorIs(t, (VerifyDatabaseConfig{Directory: t.TempDir()}).Validate(), ErrInvalidConfig) + require.ErrorIs(t, (VerifyCollectionConfig{}).Validate(), ErrInvalidConfig) + require.NoError(t, (LoadConfig{Directory: t.TempDir(), BatchSize: 1}).Validate()) + require.NoError(t, (VerifyCollectionConfig{Directory: t.TempDir()}).Validate()) + require.NoError(t, (VerifyDatabaseConfig{Directory: t.TempDir(), BatchSize: 1}).Validate()) +} + +func validDumpConfig(t *testing.T) DumpConfig { + t.Helper() + return DumpConfig{ + Directory: t.TempDir(), + Graphs: []string{"asset"}, + EntityBatchSize: 1, + ShardSize: 1, + JSONL: pointerTo(jsonl.Config{Codec: jsonl.CodecNone}), + Parquet: pointerTo(parquet.Config{}), + Scrub: pointerTo(scrub.DefaultConfig()), + } +} + +func pointerTo[T any](value T) *T { + return &value +} diff --git a/ret/dawgs/resolver.go b/ret/dawgs/resolver.go new file mode 100644 index 00000000..629edf5a --- /dev/null +++ b/ret/dawgs/resolver.go @@ -0,0 +1,66 @@ +package dawgs + +import ( + "strconv" + + "github.com/specterops/dawgs/graph" +) + +// Resolver maps collection source node IDs to IDs assigned by the target +// database. Canonical decimal IDs use a compact numeric map; all other IDs +// retain their exact string representation. +type Resolver struct { + numeric map[uint64]graph.ID + fallback map[string]graph.ID +} + +func NewResolver(expected int64) *Resolver { + capacity := int(expected) + if int64(capacity) != expected || capacity < 0 { + capacity = 0 + } + + return &Resolver{numeric: make(map[uint64]graph.ID, capacity)} +} + +func (s *Resolver) Put(sourceID string, destinationID graph.ID) bool { + if numericID, ok := canonicalNumericSourceID(sourceID); ok { + if _, exists := s.numeric[numericID]; exists { + return false + } + s.numeric[numericID] = destinationID + return true + } + + if s.fallback == nil { + s.fallback = make(map[string]graph.ID) + } + if _, exists := s.fallback[sourceID]; exists { + return false + } + s.fallback[sourceID] = destinationID + return true +} + +func (s *Resolver) Resolve(sourceID string) (graph.ID, bool) { + if numericID, ok := canonicalNumericSourceID(sourceID); ok { + resolved, found := s.numeric[numericID] + return resolved, found + } + + resolved, found := s.fallback[sourceID] + return resolved, found +} + +func canonicalNumericSourceID(sourceID string) (uint64, bool) { + if sourceID == "" { + return 0, false + } + + value, err := strconv.ParseUint(sourceID, 10, 64) + if err != nil || strconv.FormatUint(value, 10) != sourceID { + return 0, false + } + + return value, true +} diff --git a/ret/dawgs/resolver_benchmark_test.go b/ret/dawgs/resolver_benchmark_test.go new file mode 100644 index 00000000..d8983d9b --- /dev/null +++ b/ret/dawgs/resolver_benchmark_test.go @@ -0,0 +1,37 @@ +package dawgs_test + +import ( + "testing" + + "github.com/specterops/dawgs/graph" + "github.com/specterops/dawgs/ret/dawgs" +) + +var ( + resolverBenchmarkID graph.ID + resolverBenchmarkFound bool +) + +func BenchmarkResolverOperations(b *testing.B) { + b.Run("Numeric", func(b *testing.B) { + resolver := dawgs.NewResolver(1) + resolver.Put("184467", graph.ID(42)) + + b.ReportAllocs() + b.ResetTimer() + for range b.N { + resolverBenchmarkID, resolverBenchmarkFound = resolver.Resolve("184467") + } + }) + + b.Run("String", func(b *testing.B) { + resolver := dawgs.NewResolver(1) + resolver.Put("node-alpha", graph.ID(42)) + + b.ReportAllocs() + b.ResetTimer() + for range b.N { + resolverBenchmarkID, resolverBenchmarkFound = resolver.Resolve("node-alpha") + } + }) +} diff --git a/ret/dawgs/resolver_test.go b/ret/dawgs/resolver_test.go new file mode 100644 index 00000000..412545df --- /dev/null +++ b/ret/dawgs/resolver_test.go @@ -0,0 +1,35 @@ +package dawgs_test + +import ( + "testing" + + "github.com/specterops/dawgs/graph" + "github.com/specterops/dawgs/ret/dawgs" + "github.com/stretchr/testify/require" +) + +func TestResolverKeepsNonCanonicalNumericIDsDistinct(t *testing.T) { + resolver := dawgs.NewResolver(2) + require.True(t, resolver.Put("1", graph.ID(10))) + require.True(t, resolver.Put("01", graph.ID(11))) + require.Equal(t, graph.ID(10), mustResolve(t, resolver, "1")) + require.Equal(t, graph.ID(11), mustResolve(t, resolver, "01")) +} + +func TestResolverRejectsDuplicateSourceIDsAndLeavesMissingIDsUnresolved(t *testing.T) { + resolver := dawgs.NewResolver(3) + require.True(t, resolver.Put("42", graph.ID(100))) + require.True(t, resolver.Put("node-a", graph.ID(101))) + require.False(t, resolver.Put("42", graph.ID(999))) + require.False(t, resolver.Put("node-a", graph.ID(999))) + + _, found := resolver.Resolve("missing") + require.False(t, found) +} + +func mustResolve(t *testing.T, resolver *dawgs.Resolver, sourceID string) graph.ID { + t.Helper() + value, found := resolver.Resolve(sourceID) + require.Truef(t, found, "resolve %q", sourceID) + return value +} diff --git a/ret/dawgs/snapshot.go b/ret/dawgs/snapshot.go new file mode 100644 index 00000000..aa0c906f --- /dev/null +++ b/ret/dawgs/snapshot.go @@ -0,0 +1,18 @@ +package dawgs + +import "github.com/specterops/dawgs/ret/entity" + +type Snapshot struct { + NodeCount int64 + RelationshipCount int64 +} + +type NodeBatch struct { + Entities []entity.Node + LastID uint64 +} + +type RelationshipBatch struct { + Entities []entity.Relationship + LastID uint64 +} diff --git a/ret/dawgs/source.go b/ret/dawgs/source.go new file mode 100644 index 00000000..a4e55dc5 --- /dev/null +++ b/ret/dawgs/source.go @@ -0,0 +1,138 @@ +package dawgs + +import ( + "context" + "fmt" + "strconv" + "strings" + + "github.com/specterops/dawgs/graph" + "github.com/specterops/dawgs/query" + "github.com/specterops/dawgs/ret/entity" +) + +type Source struct { + database graph.Database + graph graph.Graph + batchSize int + nodeCursor uint64 + relationshipCursor uint64 +} + +func NewSource(database graph.Database, graphName string, batchSize int) (*Source, error) { + if strings.TrimSpace(graphName) == "" { + return nil, fmt.Errorf("graph name is required") + } + if batchSize <= 0 { + return nil, fmt.Errorf("batch size must be positive") + } + + return &Source{ + database: database, + graph: graph.Graph{Name: graphName}, + batchSize: batchSize, + }, nil +} + +func (s *Source) Snapshot(ctx context.Context) (Snapshot, error) { + var snapshot Snapshot + if err := s.database.ReadTransaction(ctx, func(tx graph.Transaction) error { + tx = tx.WithGraph(s.graph) + + var err error + if snapshot.NodeCount, err = tx.Nodes().Count(); err != nil { + return fmt.Errorf("snapshot nodes for graph %q: %w", s.graph.Name, err) + } + if snapshot.RelationshipCount, err = tx.Relationships().Count(); err != nil { + return fmt.Errorf("snapshot relationships for graph %q: %w", s.graph.Name, err) + } + return nil + }); err != nil { + return Snapshot{}, fmt.Errorf("snapshot graph %q: %w", s.graph.Name, err) + } + + return snapshot, nil +} + +func (s *Source) SetNodeCursor(lastID uint64) { + s.nodeCursor = lastID +} + +func (s *Source) SetRelationshipCursor(lastID uint64) { + s.relationshipCursor = lastID +} + +func (s *Source) NextNodes(ctx context.Context) (NodeBatch, error) { + batch := NodeBatch{} + if err := s.database.ReadTransaction(ctx, func(tx graph.Transaction) error { + tx = tx.WithGraph(s.graph) + return tx.Nodes(). + OrderBy(query.NodeID()). + Filter(query.GreaterThan(query.NodeID(), graph.ID(s.nodeCursor))). + Limit(s.batchSize). + Fetch(func(cursor graph.Cursor[*graph.Node]) error { + for node := range cursor.Chan() { + batch.Entities = append(batch.Entities, entity.Node{ + SourceID: strconv.FormatUint(node.ID.Uint64(), 10), + Kinds: copyKinds(node.Kinds), + Properties: copyProperties(node.Properties), + }) + batch.LastID = node.ID.Uint64() + } + return cursor.Error() + }) + }); err != nil { + return NodeBatch{}, fmt.Errorf("read nodes for graph %q: %w", s.graph.Name, err) + } + + if len(batch.Entities) > 0 { + s.nodeCursor = batch.LastID + } + return batch, nil +} + +func (s *Source) NextRelationships(ctx context.Context) (RelationshipBatch, error) { + batch := RelationshipBatch{} + if err := s.database.ReadTransaction(ctx, func(tx graph.Transaction) error { + tx = tx.WithGraph(s.graph) + return tx.Relationships(). + OrderBy(query.RelationshipID()). + Filter(query.GreaterThan(query.RelationshipID(), graph.ID(s.relationshipCursor))). + Limit(s.batchSize). + Fetch(func(cursor graph.Cursor[*graph.Relationship]) error { + for relationship := range cursor.Chan() { + batch.Entities = append(batch.Entities, entity.Relationship{ + SourceID: strconv.FormatUint(relationship.ID.Uint64(), 10), + StartID: strconv.FormatUint(relationship.StartID.Uint64(), 10), + EndID: strconv.FormatUint(relationship.EndID.Uint64(), 10), + Kind: relationship.Kind.String(), + Properties: copyProperties(relationship.Properties), + }) + batch.LastID = relationship.ID.Uint64() + } + return cursor.Error() + }) + }); err != nil { + return RelationshipBatch{}, fmt.Errorf("read relationships for graph %q: %w", s.graph.Name, err) + } + + if len(batch.Entities) > 0 { + s.relationshipCursor = batch.LastID + } + return batch, nil +} + +func copyKinds(kinds graph.Kinds) []string { + converted := make([]string, len(kinds)) + for index, kind := range kinds { + converted[index] = kind.String() + } + return converted +} + +func copyProperties(properties *graph.Properties) map[string]any { + if properties == nil || properties.Map == nil { + return nil + } + return entity.CloneProperties(properties.Map) +} diff --git a/ret/dawgs/source_test.go b/ret/dawgs/source_test.go new file mode 100644 index 00000000..06987926 --- /dev/null +++ b/ret/dawgs/source_test.go @@ -0,0 +1,441 @@ +package dawgs + +import ( + "context" + "errors" + "fmt" + "reflect" + "strings" + "testing" + + cypherModel "github.com/specterops/dawgs/cypher/models/cypher" + "github.com/specterops/dawgs/graph" +) + +type sourceTestCursor[T any] struct { + values chan T + err error +} + +func newSourceTestCursor[T any](values []T) *sourceTestCursor[T] { + channel := make(chan T, len(values)) + for _, value := range values { + channel <- value + } + close(channel) + + return &sourceTestCursor[T]{values: channel} +} + +func (s *sourceTestCursor[T]) Error() error { return s.err } +func (s *sourceTestCursor[T]) Close() {} +func (s *sourceTestCursor[T]) Chan() chan T { return s.values } + +type sourceTestDatabase struct { + graph.Database + + nodes []*graph.Node + relationships []*graph.Relationship + nodeCount int64 + edgeCount int64 + + nodeCountErr error + edgeCountErr error + nodeFetchErr error + edgeFetchErr error + + contexts []context.Context + graphs []graph.Graph + nodeOps []sourceTestQueryOperation + edgeOps []sourceTestQueryOperation +} + +type sourceTestQueryOperation struct { + orderBy graph.Criteria + filter graph.Criteria + limit int +} + +func (s *sourceTestDatabase) ReadTransaction(ctx context.Context, delegate graph.TransactionDelegate, _ ...graph.TransactionOption) error { + s.contexts = append(s.contexts, ctx) + return delegate(&sourceTestTransaction{database: s}) +} + +type sourceTestTransaction struct { + graph.Transaction + database *sourceTestDatabase +} + +func (s *sourceTestTransaction) WithGraph(target graph.Graph) graph.Transaction { + s.database.graphs = append(s.database.graphs, target) + return s +} + +func (s *sourceTestTransaction) Nodes() graph.NodeQuery { + return &sourceTestNodeQuery{database: s.database} +} + +func (s *sourceTestTransaction) Relationships() graph.RelationshipQuery { + return &sourceTestRelationshipQuery{database: s.database} +} + +type sourceTestNodeQuery struct { + graph.NodeQuery + database *sourceTestDatabase + operation sourceTestQueryOperation +} + +func (s *sourceTestNodeQuery) OrderBy(criteria ...graph.Criteria) graph.NodeQuery { + if len(criteria) != 1 { + panic(fmt.Sprintf("unexpected node order criteria: %d", len(criteria))) + } + s.operation.orderBy = criteria[0] + return s +} + +func (s *sourceTestNodeQuery) Filter(criteria graph.Criteria) graph.NodeQuery { + s.operation.filter = criteria + return s +} + +func (s *sourceTestNodeQuery) Limit(limit int) graph.NodeQuery { + s.operation.limit = limit + return s +} + +func (s *sourceTestNodeQuery) Count() (int64, error) { + if s.database.nodeCountErr != nil { + return 0, s.database.nodeCountErr + } + if s.database.nodeCount != 0 { + return s.database.nodeCount, nil + } + return int64(len(s.database.nodes)), nil +} + +func (s *sourceTestNodeQuery) Fetch(delegate func(graph.Cursor[*graph.Node]) error, _ ...graph.Criteria) error { + s.database.nodeOps = append(s.database.nodeOps, s.operation) + if s.database.nodeFetchErr != nil { + return s.database.nodeFetchErr + } + + return delegate(newSourceTestCursor(sourceTestNodesAfter(s.database.nodes, s.operation))) +} + +type sourceTestRelationshipQuery struct { + graph.RelationshipQuery + database *sourceTestDatabase + operation sourceTestQueryOperation +} + +func (s *sourceTestRelationshipQuery) OrderBy(criteria ...graph.Criteria) graph.RelationshipQuery { + if len(criteria) != 1 { + panic(fmt.Sprintf("unexpected relationship order criteria: %d", len(criteria))) + } + s.operation.orderBy = criteria[0] + return s +} + +func (s *sourceTestRelationshipQuery) Filter(criteria graph.Criteria) graph.RelationshipQuery { + s.operation.filter = criteria + return s +} + +func (s *sourceTestRelationshipQuery) Limit(limit int) graph.RelationshipQuery { + s.operation.limit = limit + return s +} + +func (s *sourceTestRelationshipQuery) Count() (int64, error) { + if s.database.edgeCountErr != nil { + return 0, s.database.edgeCountErr + } + if s.database.edgeCount != 0 { + return s.database.edgeCount, nil + } + return int64(len(s.database.relationships)), nil +} + +func (s *sourceTestRelationshipQuery) Fetch(delegate func(graph.Cursor[*graph.Relationship]) error) error { + s.database.edgeOps = append(s.database.edgeOps, s.operation) + if s.database.edgeFetchErr != nil { + return s.database.edgeFetchErr + } + + return delegate(newSourceTestCursor(sourceTestRelationshipsAfter(s.database.relationships, s.operation))) +} + +func sourceTestNodesAfter(nodes []*graph.Node, operation sourceTestQueryOperation) []*graph.Node { + afterID := sourceTestAfterID(operation.filter) + values := make([]*graph.Node, 0, operation.limit) + for _, node := range nodes { + if node.ID > afterID && len(values) < operation.limit { + values = append(values, node) + } + } + return values +} + +func sourceTestRelationshipsAfter(relationships []*graph.Relationship, operation sourceTestQueryOperation) []*graph.Relationship { + afterID := sourceTestAfterID(operation.filter) + values := make([]*graph.Relationship, 0, operation.limit) + for _, relationship := range relationships { + if relationship.ID > afterID && len(values) < operation.limit { + values = append(values, relationship) + } + } + return values +} + +func sourceTestAfterID(criteria graph.Criteria) graph.ID { + if criteria == nil { + return 0 + } + comparison := criteria.(*cypherModel.Comparison) + return comparison.Partials[0].Right.(*cypherModel.Parameter).Value.(graph.ID) +} + +func newSourceTestDatabase(nodes []*graph.Node, relationships []*graph.Relationship) *sourceTestDatabase { + return &sourceTestDatabase{nodes: nodes, relationships: relationships} +} + +func newSourceForTest(t *testing.T, database *sourceTestDatabase, batchSize int) *Source { + t.Helper() + source, err := NewSource(database, "example", batchSize) + if err != nil { + t.Fatalf("new source: %v", err) + } + return source +} + +func TestNewSourceRejectsMissingGraphAndNonPositiveBatchSize(t *testing.T) { + database := newSourceTestDatabase(nil, nil) + for name, input := range map[string]struct { + graphName string + batchSize int + }{ + "missing graph": {batchSize: 1}, + "blank graph": {graphName: " \t", batchSize: 1}, + "zero batch size": {graphName: "example"}, + "negative batch size": {graphName: "example", batchSize: -1}, + } { + t.Run(name, func(t *testing.T) { + if _, err := NewSource(database, input.graphName, input.batchSize); err == nil { + t.Fatal("expected validation error") + } + }) + } +} + +func TestSnapshotCountsEachEntityKindInSelectedGraph(t *testing.T) { + database := newSourceTestDatabase( + []*graph.Node{graph.NewNode(1, graph.NewProperties(), graph.StringKind("User")), graph.NewNode(2, graph.NewProperties(), graph.StringKind("Group"))}, + []*graph.Relationship{graph.NewRelationship(3, 1, 2, graph.NewProperties(), graph.StringKind("Member"))}, + ) + source := newSourceForTest(t, database, 2) + ctx := context.WithValue(context.Background(), sourceTestContextKey{}, "snapshot") + + snapshot, err := source.Snapshot(ctx) + if err != nil { + t.Fatalf("snapshot: %v", err) + } + if want := (Snapshot{NodeCount: 2, RelationshipCount: 1}); snapshot != want { + t.Fatalf("snapshot=%+v want=%+v", snapshot, want) + } + if len(database.graphs) != 1 || database.graphs[0].Name != "example" { + t.Fatalf("graphs=%+v", database.graphs) + } + if len(database.contexts) != 1 || database.contexts[0] != ctx { + t.Fatal("snapshot did not pass context to ReadTransaction") + } +} + +func TestNextNodesConvertsBatchesAndAdvancesCursorForOneEntity(t *testing.T) { + nested := map[string]any{"active": true} + properties := map[string]any{"name": "Ada", "metadata": nested} + database := newSourceTestDatabase([]*graph.Node{ + graph.NewNode(7, graph.AsProperties(properties), graph.StringKind("User"), graph.StringKind("Admin"), graph.StringKind("User")), + }, nil) + source := newSourceForTest(t, database, 2) + ctx := context.WithValue(context.Background(), sourceTestContextKey{}, "nodes") + + first, err := source.NextNodes(ctx) + if err != nil { + t.Fatalf("next nodes: %v", err) + } + if first.LastID != 7 || len(first.Entities) != 1 { + t.Fatalf("first batch=%+v", first) + } + if got, want := first.Entities[0].SourceID, "7"; got != want { + t.Fatalf("source ID=%q want=%q", got, want) + } + if got, want := first.Entities[0].Kinds, []string{"User", "Admin", "User"}; !reflect.DeepEqual(got, want) { + t.Fatalf("kinds=%v want=%v", got, want) + } + if got, want := first.Entities[0].Properties, properties; !reflect.DeepEqual(got, want) { + t.Fatalf("properties=%v want=%v", got, want) + } + + first.Entities[0].Properties["name"] = "Grace" + if properties["name"] != "Ada" { + t.Fatal("returned top-level properties map aliases the database map") + } + first.Entities[0].Properties["metadata"].(map[string]any)["active"] = false + if nested["active"] != false { + t.Fatal("returned nested property value was deep copied") + } + + second, err := source.NextNodes(ctx) + if err != nil { + t.Fatalf("next nodes after one-row batch: %v", err) + } + if len(second.Entities) != 0 || second.LastID != 0 { + t.Fatalf("second batch=%+v", second) + } + if len(database.nodeOps) != 2 { + t.Fatalf("node operations=%d", len(database.nodeOps)) + } + assertSourceTestKeysetQuery(t, database.nodeOps[0], "n", 0, 2) + assertSourceTestKeysetQuery(t, database.nodeOps[1], "n", 7, 2) + if len(database.contexts) != 2 || database.contexts[0] != ctx || database.contexts[1] != ctx { + t.Fatal("node reads did not pass context to ReadTransaction") + } +} + +func TestNextRelationshipsUsesIndependentCursorAndCanonicalEndpoints(t *testing.T) { + database := newSourceTestDatabase( + []*graph.Node{graph.NewNode(9, graph.NewProperties(), graph.StringKind("User"))}, + []*graph.Relationship{graph.NewRelationship(4, 10, 20, graph.AsProperties(map[string]any{"role": "owner"}), graph.StringKind("Owns"))}, + ) + source := newSourceForTest(t, database, 1) + source.SetNodeCursor(9) + source.SetRelationshipCursor(3) + + batch, err := source.NextRelationships(context.Background()) + if err != nil { + t.Fatalf("next relationships: %v", err) + } + if batch.LastID != 4 || len(batch.Entities) != 1 { + t.Fatalf("relationship batch=%+v", batch) + } + if want := (struct { + SourceID string + StartID string + EndID string + Kind string + }{"4", "10", "20", "Owns"}); struct { + SourceID string + StartID string + EndID string + Kind string + }{batch.Entities[0].SourceID, batch.Entities[0].StartID, batch.Entities[0].EndID, batch.Entities[0].Kind} != want { + t.Fatalf("relationship=%+v", batch.Entities[0]) + } + assertSourceTestKeysetQuery(t, database.edgeOps[0], "r", 3, 1) + + nodes, err := source.NextNodes(context.Background()) + if err != nil { + t.Fatalf("next nodes: %v", err) + } + if len(nodes.Entities) != 0 { + t.Fatalf("node cursor was not independent: %+v", nodes) + } + assertSourceTestKeysetQuery(t, database.nodeOps[0], "n", 9, 1) +} + +func TestNextNodesReturnsEmptyBatchWithoutChangingCursor(t *testing.T) { + database := newSourceTestDatabase(nil, nil) + source := newSourceForTest(t, database, 3) + source.SetNodeCursor(12) + + batch, err := source.NextNodes(context.Background()) + if err != nil { + t.Fatalf("next nodes: %v", err) + } + if len(batch.Entities) != 0 || batch.LastID != 0 { + t.Fatalf("batch=%+v", batch) + } + + _, err = source.NextNodes(context.Background()) + if err != nil { + t.Fatalf("second next nodes: %v", err) + } + assertSourceTestKeysetQuery(t, database.nodeOps[1], "n", 12, 3) +} + +func TestSourceErrorsIncludeGraphAndPhaseAndWrapCause(t *testing.T) { + nodeCountErr := errors.New("node count failed") + database := newSourceTestDatabase(nil, nil) + database.nodeCountErr = nodeCountErr + source := newSourceForTest(t, database, 1) + + _, err := source.Snapshot(context.Background()) + assertSourceTestError(t, err, nodeCountErr, "graph \"example\"", "snapshot", "nodes") + + nodeFetchErr := errors.New("node query failed") + database.nodeCountErr = nil + database.nodeFetchErr = nodeFetchErr + _, err = source.NextNodes(context.Background()) + assertSourceTestError(t, err, nodeFetchErr, "graph \"example\"", "nodes") + + edgeCountErr := errors.New("relationship count failed") + database.nodeFetchErr = nil + database.edgeCountErr = edgeCountErr + _, err = source.Snapshot(context.Background()) + assertSourceTestError(t, err, edgeCountErr, "graph \"example\"", "snapshot", "relationships") + + edgeFetchErr := errors.New("relationship query failed") + database.edgeCountErr = nil + database.edgeFetchErr = edgeFetchErr + _, err = source.NextRelationships(context.Background()) + assertSourceTestError(t, err, edgeFetchErr, "graph \"example\"", "relationships") +} + +type sourceTestContextKey struct{} + +func assertSourceTestKeysetQuery(t *testing.T, operation sourceTestQueryOperation, variable string, afterID graph.ID, limit int) { + t.Helper() + if operation.limit != limit { + t.Fatalf("limit=%d want=%d", operation.limit, limit) + } + assertSourceTestIDExpression(t, operation.orderBy, variable) + comparison, ok := operation.filter.(*cypherModel.Comparison) + if !ok { + t.Fatalf("filter=%T want keyset comparison", operation.filter) + } + assertSourceTestIDExpression(t, comparison.Left, variable) + if len(comparison.Partials) != 1 { + t.Fatalf("filter partials=%d", len(comparison.Partials)) + } + if comparison.Partials[0].Operator != cypherModel.OperatorGreaterThan { + t.Fatalf("filter operator=%q want=%q", comparison.Partials[0].Operator, cypherModel.OperatorGreaterThan) + } + parameter, ok := comparison.Partials[0].Right.(*cypherModel.Parameter) + if !ok || parameter.Value != afterID { + t.Fatalf("filter right=%#v want graph.ID(%d)", comparison.Partials[0].Right, afterID) + } +} + +func assertSourceTestIDExpression(t *testing.T, criteria graph.Criteria, variable string) { + t.Helper() + function, ok := criteria.(*cypherModel.FunctionInvocation) + if !ok || function.Name != "id" || len(function.Arguments) != 1 { + t.Fatalf("ID criteria=%#v", criteria) + } + argument, ok := function.Arguments[0].(*cypherModel.Variable) + if !ok || argument.Symbol != variable { + t.Fatalf("ID variable=%#v want %q", function.Arguments[0], variable) + } +} + +func assertSourceTestError(t *testing.T, err, cause error, fragments ...string) { + t.Helper() + if !errors.Is(err, cause) { + t.Fatalf("error %v does not wrap %v", err, cause) + } + for _, fragment := range fragments { + if !strings.Contains(err.Error(), fragment) { + t.Fatalf("error %q does not contain %q", err, fragment) + } + } +} diff --git a/ret/dawgs/target.go b/ret/dawgs/target.go new file mode 100644 index 00000000..02116717 --- /dev/null +++ b/ret/dawgs/target.go @@ -0,0 +1,187 @@ +package dawgs + +import ( + "context" + "errors" + "fmt" + "strings" + + "github.com/specterops/dawgs/graph" + "github.com/specterops/dawgs/ret/entity" +) + +// ErrTargetNotEmpty indicates that an emptiness snapshot completed and found +// at least one node or relationship. +var ErrTargetNotEmpty = errors.New("target graph is not empty") + +// Target writes canonical entities to one empty DAWGS graph. +type Target struct { + database graph.Database + graphName string + batchSize int + source *Source +} + +func NewTarget(database graph.Database, graphName string, batchSize int) (*Target, error) { + if strings.TrimSpace(graphName) == "" { + return nil, fmt.Errorf("graph name is required") + } + if batchSize <= 0 { + return nil, fmt.Errorf("batch size must be positive") + } + + source, err := NewSource(database, graphName, batchSize) + if err != nil { + return nil, err + } + + return &Target{ + database: database, + graphName: graphName, + batchSize: batchSize, + source: source, + }, nil +} + +func (s *Target) RequireEmpty(ctx context.Context) error { + snapshot, err := s.source.Snapshot(ctx) + if err != nil { + return err + } + if snapshot.NodeCount != 0 || snapshot.RelationshipCount != 0 { + return fmt.Errorf( + "%w: graph %q is not empty: nodes=%d relationships=%d", + ErrTargetNotEmpty, + s.graphName, + snapshot.NodeCount, + snapshot.RelationshipCount, + ) + } + return nil +} + +func (s *Target) AssertSchema(ctx context.Context, kindCatalog []string) error { + seen := make(map[string]struct{}, len(kindCatalog)) + kinds := make(graph.Kinds, 0, len(kindCatalog)) + for _, kind := range kindCatalog { + if _, found := seen[kind]; found { + continue + } + seen[kind] = struct{}{} + kinds = append(kinds, graph.StringKind(kind)) + } + + targetGraph := graph.Graph{ + Name: s.graphName, + Nodes: kinds, + Edges: kinds.Copy(), + } + if err := s.database.AssertSchema(ctx, graph.Schema{ + Graphs: []graph.Graph{targetGraph}, + DefaultGraph: targetGraph, + }); err != nil { + return fmt.Errorf("assert schema for graph %q: %w", s.graphName, err) + } + return nil +} + +func (s *Target) CreateNodes(ctx context.Context, nodes []entity.Node, resolver *Resolver) error { + if resolver == nil { + return errors.New("source ID resolver is required") + } + if len(nodes) == 0 { + return nil + } + if err := s.preflightNodeSourceIDs(nodes, resolver); err != nil { + return err + } + + staged := make([]resolvedNode, len(nodes)) + if err := s.database.BatchOperation(ctx, func(batch graph.Batch) error { + batch = batch.WithGraph(graph.Graph{Name: s.graphName}) + creator, ok := batch.(graph.NodeBatchCreator) + if !ok { + return errors.New("database batch does not support correlated node creation") + } + + graphNodes := make([]*graph.Node, len(nodes)) + for index, value := range nodes { + graphNodes[index] = graph.NewNode( + 0, + graph.AsProperties(entity.CloneProperties(value.Properties)), + graph.StringsToKinds(entity.CloneKinds(value.Kinds))..., + ) + } + + destinationIDs, err := creator.CreateNodes(graphNodes) + if err != nil { + return err + } + if len(destinationIDs) != len(nodes) { + return fmt.Errorf("created node ID count: got %d want %d", len(destinationIDs), len(nodes)) + } + for index, value := range nodes { + staged[index] = resolvedNode{sourceID: value.SourceID, destinationID: destinationIDs[index]} + } + return nil + }, graph.WithBatchSize(s.batchSize)); err != nil { + return err + } + + for _, value := range staged { + if !resolver.Put(value.sourceID, value.destinationID) { + return fmt.Errorf("duplicate source node ID %q in graph %q", value.sourceID, s.graphName) + } + } + return nil +} + +type resolvedNode struct { + sourceID string + destinationID graph.ID +} + +func (s *Target) preflightNodeSourceIDs(nodes []entity.Node, resolver *Resolver) error { + pending := NewResolver(int64(len(nodes))) + for _, value := range nodes { + if _, found := resolver.Resolve(value.SourceID); found || !pending.Put(value.SourceID, 0) { + return fmt.Errorf("duplicate source node ID %q in graph %q", value.SourceID, s.graphName) + } + } + return nil +} + +func (s *Target) CreateRelationships(ctx context.Context, relationships []entity.Relationship, resolver *Resolver) error { + if resolver == nil { + return errors.New("source ID resolver is required") + } + if len(relationships) == 0 { + return nil + } + + return s.database.BatchOperation(ctx, func(batch graph.Batch) error { + batch = batch.WithGraph(graph.Graph{Name: s.graphName}) + for index, value := range relationships { + startID, startOK := resolver.Resolve(value.StartID) + endID, endOK := resolver.Resolve(value.EndID) + if !startOK || !endOK { + return fmt.Errorf( + "graph %q relationship %d has unresolved endpoints %q -> %q", + s.graphName, + index, + value.StartID, + value.EndID, + ) + } + if err := batch.CreateRelationshipByIDs( + startID, + endID, + graph.StringKind(value.Kind), + graph.AsProperties(entity.CloneProperties(value.Properties)), + ); err != nil { + return err + } + } + return nil + }, graph.WithBatchSize(s.batchSize)) +} diff --git a/ret/dawgs/target_test.go b/ret/dawgs/target_test.go new file mode 100644 index 00000000..d14218d1 --- /dev/null +++ b/ret/dawgs/target_test.go @@ -0,0 +1,407 @@ +package dawgs_test + +import ( + "context" + "errors" + "testing" + + "github.com/specterops/dawgs/graph" + "github.com/specterops/dawgs/ret/dawgs" + "github.com/specterops/dawgs/ret/entity" + "github.com/stretchr/testify/require" +) + +func TestNewTargetRejectsMissingGraphAndNonPositiveBatchSize(t *testing.T) { + database := &targetTestDatabase{} + for name, input := range map[string]struct { + graphName string + batchSize int + }{ + "missing graph": {batchSize: 1}, + "blank graph": {graphName: " \t", batchSize: 1}, + "zero batch size": {graphName: "example"}, + "negative batch size": {graphName: "example", batchSize: -1}, + } { + t.Run(name, func(t *testing.T) { + _, err := dawgs.NewTarget(database, input.graphName, input.batchSize) + require.Error(t, err) + }) + } +} + +func TestTargetRequireEmptyRejectsExistingNodesAndRelationships(t *testing.T) { + for name, input := range map[string]struct { + database *targetTestDatabase + counts string + }{ + "nodes": {database: &targetTestDatabase{nodeCount: 1}, counts: "nodes=1 relationships=0"}, + "relationships": {database: &targetTestDatabase{relationshipCount: 2}, counts: "nodes=0 relationships=2"}, + } { + t.Run(name, func(t *testing.T) { + target := newTargetForTest(t, input.database, 2) + + err := target.RequireEmpty(context.Background()) + + require.ErrorIs(t, err, dawgs.ErrTargetNotEmpty) + require.ErrorContains(t, err, `graph "example" is not empty`) + require.ErrorContains(t, err, input.counts) + require.Len(t, input.database.readGraphs, 1) + require.Equal(t, "example", input.database.readGraphs[0].Name) + }) + } +} + +func TestTargetRequireEmptyPreservesSnapshotFailureClassification(t *testing.T) { + // Break caught: making every failed emptiness probe indistinguishable from + // a graph whose nonzero counts were actually observed. + injected := errors.New("injected snapshot failure") + target := newTargetForTest(t, &targetTestDatabase{readErr: injected}, 1) + + err := target.RequireEmpty(context.Background()) + + require.ErrorIs(t, err, injected) + require.NotErrorIs(t, err, dawgs.ErrTargetNotEmpty) +} + +func TestTargetRequireEmptyPreservesCancellationClassification(t *testing.T) { + // Break caught: converting context cancellation into a not-empty result. + ctx, cancel := context.WithCancel(context.Background()) + cancel() + target := newTargetForTest(t, &targetTestDatabase{}, 1) + + err := target.RequireEmpty(ctx) + + require.ErrorIs(t, err, context.Canceled) + require.NotErrorIs(t, err, dawgs.ErrTargetNotEmpty) +} + +func TestTargetAssertSchemaScopesFirstSeenKindsToSelectedGraph(t *testing.T) { + database := &targetTestDatabase{} + target := newTargetForTest(t, database, 3) + + require.NoError(t, target.AssertSchema(context.Background(), []string{"User", "MEMBER_OF", "User", "Group"})) + require.Len(t, database.schemas, 1) + + schema := database.schemas[0] + require.Len(t, schema.Graphs, 1) + require.Equal(t, "example", schema.Graphs[0].Name) + require.Equal(t, []string{"User", "MEMBER_OF", "Group"}, schema.Graphs[0].Nodes.Strings()) + require.Equal(t, []string{"User", "MEMBER_OF", "Group"}, schema.Graphs[0].Edges.Strings()) + require.Equal(t, "example", schema.DefaultGraph.Name) +} + +func TestTargetCreateNodesPreservesKindOrderPropertiesAndCorrelatedIDs(t *testing.T) { + database := &targetTestDatabase{createdIDs: [][]graph.ID{{20, 10}}} + target := newTargetForTest(t, database, 7) + resolver := dawgs.NewResolver(2) + nested := map[string]any{"value": "shared"} + nodes := []entity.Node{ + {SourceID: "first", Kinds: []string{"User", "Admin", "User"}, Properties: map[string]any{"name": "Ada", "nested": nested}}, + {SourceID: "second", Kinds: []string{"Group"}, Properties: map[string]any{"name": "Operators"}}, + } + + require.NoError(t, target.CreateNodes(context.Background(), nodes, resolver)) + + require.Len(t, database.nodes, 2) + require.Equal(t, []string{"User", "Admin", "User"}, database.nodes[0].Kinds.Strings()) + require.Equal(t, "Ada", database.nodes[0].Properties.Map["name"]) + nodes[0].Properties["name"] = "Grace" + nested["value"] = "still-shared" + require.Equal(t, "Ada", database.nodes[0].Properties.Map["name"]) + require.Equal(t, "still-shared", database.nodes[0].Properties.Map["nested"].(map[string]any)["value"]) + require.Equal(t, graph.ID(20), mustResolve(t, resolver, "first")) + require.Equal(t, graph.ID(10), mustResolve(t, resolver, "second")) + require.Equal(t, []targetTestBatchRecord{{graphName: "example", batchSize: 7}}, database.batches) +} + +func TestTargetCreateNodesRejectsDestinationIDCountMismatchWithoutMappings(t *testing.T) { + database := &targetTestDatabase{createdIDs: [][]graph.ID{{10}}} + target := newTargetForTest(t, database, 2) + resolver := dawgs.NewResolver(2) + require.True(t, resolver.Put("already-present", 99)) + + err := target.CreateNodes(context.Background(), []entity.Node{{SourceID: "one"}, {SourceID: "two"}}, resolver) + + require.ErrorContains(t, err, "created node ID count: got 1 want 2") + require.Empty(t, database.nodes) + require.Equal(t, graph.ID(99), mustResolve(t, resolver, "already-present")) + _, found := resolver.Resolve("one") + require.False(t, found) +} + +func TestTargetCreateNodesPreflightsDuplicateAndExistingSourceIDs(t *testing.T) { + for name, input := range map[string]struct { + present bool + nodes []entity.Node + destination []graph.ID + }{ + "duplicate input": { + nodes: []entity.Node{{SourceID: "same"}, {SourceID: "same"}}, + destination: []graph.ID{10, 20}, + }, + "existing canonical ID": { + present: true, + nodes: []entity.Node{{SourceID: "1"}}, + destination: []graph.ID{10}, + }, + } { + t.Run(name, func(t *testing.T) { + database := &targetTestDatabase{createdIDs: [][]graph.ID{input.destination}} + target := newTargetForTest(t, database, 2) + resolver := dawgs.NewResolver(2) + if input.present { + require.True(t, resolver.Put("1", 99)) + } + + err := target.CreateNodes(context.Background(), input.nodes, resolver) + + require.ErrorContains(t, err, "duplicate source node ID") + require.Empty(t, database.nodes) + require.Zero(t, database.createNodesCalls) + if input.present { + require.Equal(t, graph.ID(99), mustResolve(t, resolver, "1")) + } else { + _, found := resolver.Resolve("same") + require.False(t, found) + } + }) + } +} + +func TestTargetCreateNodesLeavesResolverUnchangedWhenBatchFailsAfterCreation(t *testing.T) { + batchErr := errors.New("injected batch failure") + database := &targetTestDatabase{createdIDs: [][]graph.ID{{10}}, batchErr: batchErr} + target := newTargetForTest(t, database, 1) + resolver := dawgs.NewResolver(2) + require.True(t, resolver.Put("already-present", 99)) + + err := target.CreateNodes(context.Background(), []entity.Node{{SourceID: "new"}}, resolver) + + require.ErrorIs(t, err, batchErr) + require.Empty(t, database.nodes) + require.Equal(t, graph.ID(99), mustResolve(t, resolver, "already-present")) + _, found := resolver.Resolve("new") + require.False(t, found) +} + +func TestTargetCreateNodesRequiresCorrelatedBatchCreator(t *testing.T) { + database := &targetTestDatabase{supportsNodeBatchCreator: false} + target := newTargetForTest(t, database, 1) + + err := target.CreateNodes(context.Background(), []entity.Node{{SourceID: "one"}}, dawgs.NewResolver(1)) + + require.ErrorContains(t, err, "does not support correlated node creation") + require.Empty(t, database.nodes) +} + +func TestTargetCreateRelationshipsResolvesEndpointsAndIgnoresSourceID(t *testing.T) { + database := &targetTestDatabase{} + target := newTargetForTest(t, database, 4) + resolver := dawgs.NewResolver(2) + require.True(t, resolver.Put("node-a", 30)) + require.True(t, resolver.Put("node-b", 40)) + nested := map[string]any{"value": "shared"} + relationships := []entity.Relationship{{ + StartID: "node-a", + EndID: "node-b", + Kind: "MEMBER_OF", + Properties: map[string]any{"nested": nested}, + }} + + require.NoError(t, target.CreateRelationships(context.Background(), relationships, resolver)) + + require.Equal(t, []targetTestRelationship{{startID: 30, endID: 40, kind: "MEMBER_OF", properties: map[string]any{"nested": nested}}}, database.relationships) + nested["value"] = "still-shared" + require.Equal(t, "still-shared", database.relationships[0].properties["nested"].(map[string]any)["value"]) + require.Equal(t, []targetTestBatchRecord{{graphName: "example", batchSize: 4}}, database.batches) +} + +func TestTargetCreateRelationshipsRejectsMissingEndpointBeforeCreatingRelationships(t *testing.T) { + database := &targetTestDatabase{} + target := newTargetForTest(t, database, 1) + resolver := dawgs.NewResolver(1) + require.True(t, resolver.Put("start", 10)) + + err := target.CreateRelationships(context.Background(), []entity.Relationship{{StartID: "start", EndID: "missing", Kind: "REL"}}, resolver) + + require.ErrorContains(t, err, `relationship 0 has unresolved endpoints "start" -> "missing"`) + require.Empty(t, database.relationships) +} + +func newTargetForTest(t *testing.T, database *targetTestDatabase, batchSize int) *dawgs.Target { + t.Helper() + target, err := dawgs.NewTarget(database, "example", batchSize) + require.NoError(t, err) + return target +} + +type targetTestDatabase struct { + graph.Database + + readErr error + nodeCount int64 + relationshipCount int64 + readGraphs []graph.Graph + schemas []graph.Schema + + createdIDs [][]graph.ID + supportsNodeBatchCreator bool + rollbackErr error + batchErr error + rollbacks int + createNodesCalls int + nodes []*graph.Node + relationships []targetTestRelationship + batches []targetTestBatchRecord +} + +func (s *targetTestDatabase) ReadTransaction(ctx context.Context, delegate graph.TransactionDelegate, _ ...graph.TransactionOption) error { + if s.readErr != nil { + return s.readErr + } + if err := ctx.Err(); err != nil { + return err + } + return delegate(&targetTestTransaction{database: s}) +} + +func (s *targetTestDatabase) AssertSchema(_ context.Context, schema graph.Schema) error { + s.schemas = append(s.schemas, schema) + return nil +} + +func (s *targetTestDatabase) BatchOperation(_ context.Context, delegate graph.BatchDelegate, options ...graph.BatchOption) error { + config := graph.BatchConfig{} + for _, option := range options { + option(&config) + } + + if !s.supportsNodeBatchCreator && s.createdIDs == nil { + batch := &targetTestUncorrelatedBatch{database: s, batchSize: config.BatchSize} + if err := delegate(batch); err != nil { + return s.rollback(err) + } + s.commit(batch.nodes, batch.relationships, batch.graphName, batch.batchSize) + return nil + } + + batch := &targetTestBatch{database: s, batchSize: config.BatchSize} + if err := delegate(batch); err != nil { + return s.rollback(err) + } + if s.batchErr != nil { + return s.batchErr + } + s.commit(batch.nodes, batch.relationships, batch.graphName, batch.batchSize) + return nil +} + +func (s *targetTestDatabase) rollback(delegateErr error) error { + s.rollbacks++ + if s.rollbackErr != nil { + return s.rollbackErr + } + return delegateErr +} + +func (s *targetTestDatabase) commit(nodes []*graph.Node, relationships []targetTestRelationship, graphName string, batchSize int) { + s.nodes = append(s.nodes, nodes...) + s.relationships = append(s.relationships, relationships...) + s.batches = append(s.batches, targetTestBatchRecord{graphName: graphName, batchSize: batchSize}) +} + +type targetTestTransaction struct { + graph.Transaction + database *targetTestDatabase +} + +func (s *targetTestTransaction) WithGraph(target graph.Graph) graph.Transaction { + s.database.readGraphs = append(s.database.readGraphs, target) + return s +} + +func (s *targetTestTransaction) Nodes() graph.NodeQuery { + return targetTestNodeQuery{count: s.database.nodeCount} +} + +func (s *targetTestTransaction) Relationships() graph.RelationshipQuery { + return targetTestRelationshipQuery{count: s.database.relationshipCount} +} + +type targetTestNodeQuery struct { + graph.NodeQuery + count int64 +} + +func (s targetTestNodeQuery) Count() (int64, error) { return s.count, nil } + +type targetTestRelationshipQuery struct { + graph.RelationshipQuery + count int64 +} + +func (s targetTestRelationshipQuery) Count() (int64, error) { return s.count, nil } + +type targetTestBatchRecord struct { + graphName string + batchSize int +} + +type targetTestRelationship struct { + startID graph.ID + endID graph.ID + kind string + properties map[string]any +} + +type targetTestBatch struct { + graph.Batch + database *targetTestDatabase + graphName string + batchSize int + nodes []*graph.Node + relationships []targetTestRelationship +} + +func (s *targetTestBatch) WithGraph(target graph.Graph) graph.Batch { + s.graphName = target.Name + return s +} + +func (s *targetTestBatch) CreateNodes(nodes []*graph.Node) ([]graph.ID, error) { + s.database.createNodesCalls++ + s.nodes = append(s.nodes, nodes...) + call := len(s.database.batches) + if call < len(s.database.createdIDs) { + return append([]graph.ID(nil), s.database.createdIDs[call]...), nil + } + ids := make([]graph.ID, len(nodes)) + for index := range ids { + ids[index] = graph.ID(len(s.database.nodes) + index + 1) + } + return ids, nil +} + +func (s *targetTestBatch) CreateRelationshipByIDs(startID, endID graph.ID, kind graph.Kind, properties *graph.Properties) error { + s.relationships = append(s.relationships, targetTestRelationship{startID: startID, endID: endID, kind: kind.String(), properties: properties.Map}) + return nil +} + +type targetTestUncorrelatedBatch struct { + graph.Batch + database *targetTestDatabase + graphName string + batchSize int + nodes []*graph.Node + relationships []targetTestRelationship +} + +func (s *targetTestUncorrelatedBatch) WithGraph(target graph.Graph) graph.Batch { + s.graphName = target.Name + return s +} + +func (s *targetTestUncorrelatedBatch) CreateRelationshipByIDs(startID, endID graph.ID, kind graph.Kind, properties *graph.Properties) error { + s.relationships = append(s.relationships, targetTestRelationship{startID: startID, endID: endID, kind: kind.String(), properties: properties.Map}) + return nil +} diff --git a/ret/doc.go b/ret/doc.go new file mode 100644 index 00000000..8e195bdf --- /dev/null +++ b/ret/doc.go @@ -0,0 +1,2 @@ +// Package ret contains components for exporting graph data. +package ret diff --git a/ret/dump.go b/ret/dump.go new file mode 100644 index 00000000..8f8143bf --- /dev/null +++ b/ret/dump.go @@ -0,0 +1,644 @@ +package ret + +import ( + "context" + "errors" + "fmt" + "os" + "path/filepath" + "strconv" + "time" + + "github.com/specterops/dawgs/graph" + "github.com/specterops/dawgs/ret/checkpoint" + "github.com/specterops/dawgs/ret/collection" + "github.com/specterops/dawgs/ret/dawgs" + "github.com/specterops/dawgs/ret/entity" + "github.com/specterops/dawgs/ret/jsonl" + "github.com/specterops/dawgs/ret/metrics" + "github.com/specterops/dawgs/ret/observe" + "github.com/specterops/dawgs/ret/parquet" + "github.com/specterops/dawgs/ret/scrub" +) + +const dumpOperationName = "dump" + +var writeCollection = collection.Write + +// Dump exports the configured graphs into a validated collection. +func Dump(ctx context.Context, database graph.Database, config DumpConfig) (result DumpResult, resultErr error) { + started := time.Now() + observe.Emit(ctx, config.Observer, observe.OperationStarted{Operation: dumpOperationName}) + defer func() { + observe.Emit(ctx, config.Observer, observe.OperationCompleted{ + Operation: dumpOperationName, + Duration: time.Since(started), + Err: resultErr, + }) + }() + + if err := ctx.Err(); err != nil { + return DumpResult{}, fmt.Errorf("dump: %w", err) + } + if err := config.Validate(); err != nil { + return DumpResult{}, err + } + var compiledScrubber *scrub.Scrubber + if config.Scrub != nil { + var err error + compiledScrubber, err = scrub.New(*config.Scrub) + if err != nil { + return DumpResult{}, fmt.Errorf("%w: compile scrub configuration: %w", ErrInvalidConfig, err) + } + } + + runner := dumpRunner{ + ctx: ctx, + database: database, + config: config, + scrubber: compiledScrubber, + store: checkpoint.Store{Root: config.Directory}, + identity: dumpCheckpointIdentity(config, compiledScrubber), + } + if config.Resume { + return runner.runResume() + } + return runner.runFresh() +} + +type dumpRunner struct { + ctx context.Context + database graph.Database + config DumpConfig + scrubber *scrub.Scrubber + store checkpoint.Store + identity checkpoint.Identity + state checkpoint.State + graphs []collection.Graph +} + +type dumpGraphRuntime struct { + source *dawgs.Source + builder *metrics.Builder + catalog *dumpKindCatalog + totalScrubCounts scrub.ActionCounts + nodeCount int64 + relationshipCount int64 +} + +func (s *dumpRunner) runFresh() (DumpResult, error) { + if err := s.prepareFreshDestination(); err != nil { + return DumpResult{}, err + } + s.state = checkpoint.State{Format: checkpoint.Format, Identity: s.identity} + + for index, graphName := range s.config.Graphs { + if err := s.ctx.Err(); err != nil { + return DumpResult{}, fmt.Errorf("dump graph %q: %w", graphName, err) + } + runtime, err := s.startFreshGraph(graphName) + if err != nil { + return DumpResult{}, err + } + if err := s.processGraph(index, runtime); err != nil { + return DumpResult{}, err + } + } + + return s.finalize() +} + +func (s *dumpRunner) prepareFreshDestination() error { + if _, err := os.Lstat(s.config.Directory); err == nil { + return fmt.Errorf("%w: %s", ErrDestinationExists, s.config.Directory) + } else if !errors.Is(err, os.ErrNotExist) { + return fmt.Errorf("inspect dump destination: %w", err) + } + if err := os.Mkdir(s.config.Directory, 0o755); err != nil { + if errors.Is(err, os.ErrExist) { + return fmt.Errorf("%w: %s", ErrDestinationExists, s.config.Directory) + } + return fmt.Errorf("create dump destination: %w", err) + } + return nil +} + +func (s *dumpRunner) startFreshGraph(graphName string) (*dumpGraphRuntime, error) { + source, err := dawgs.NewSource(s.database, graphName, s.config.EntityBatchSize) + if err != nil { + return nil, fmt.Errorf("prepare dump source for graph %q: %w", graphName, err) + } + snapshot, err := source.Snapshot(s.ctx) + if err != nil { + return nil, fmt.Errorf("snapshot dump graph %q: %w", graphName, err) + } + if err := s.ctx.Err(); err != nil { + return nil, fmt.Errorf("snapshot dump graph %q: %w", graphName, err) + } + s.state.Graphs = append(s.state.Graphs, checkpoint.GraphState{ + Name: graphName, + Snapshot: snapshot, + Phase: checkpoint.PhaseNodes, + }) + if err := s.store.Save(s.state); err != nil { + return nil, fmt.Errorf("save initial checkpoint for graph %q: %w", graphName, err) + } + return &dumpGraphRuntime{ + source: source, + builder: metrics.NewBuilder(), + catalog: newDumpKindCatalog(), + totalScrubCounts: scrub.ActionCounts{}, + }, nil +} + +func (s *dumpRunner) processGraph(index int, runtime *dumpGraphRuntime) error { + graphState := &s.state.Graphs[index] + graphStarted := time.Now() + observe.Emit(s.ctx, s.config.Observer, observe.GraphStarted{ + Operation: dumpOperationName, + Graph: graphState.Name, + }) + if err := s.ctx.Err(); err != nil { + return fmt.Errorf("dump graph %q: %w", graphState.Name, err) + } + + if graphState.Phase == checkpoint.PhaseNodes { + if err := s.processNodes(graphState, runtime); err != nil { + return err + } + if err := s.ctx.Err(); err != nil { + return fmt.Errorf("dump graph %q: %w", graphState.Name, err) + } + } + if graphState.Phase == checkpoint.PhaseRelationships { + if err := s.processRelationships(graphState, runtime); err != nil { + return err + } + } + if graphState.Phase != checkpoint.PhaseComplete { + return fmt.Errorf("dump graph %q has unsupported checkpoint phase %q", graphState.Name, graphState.Phase) + } + if err := s.ctx.Err(); err != nil { + return fmt.Errorf("dump graph %q: %w", graphState.Name, err) + } + + graph := collection.Graph{ + Name: graphState.Name, + NodeCount: runtime.nodeCount, + RelationshipCount: runtime.relationshipCount, + KindCatalog: append([]string(nil), runtime.catalog.values...), + NodeShards: append([]collection.NodeShard(nil), graphState.NodeShards...), + RelationshipShards: append([]collection.RelationshipShard(nil), graphState.RelationshipShards...), + Metrics: runtime.builder.Finalize(), + } + s.graphs = append(s.graphs, graph) + observe.Emit(s.ctx, s.config.Observer, observe.GraphCompleted{ + Operation: dumpOperationName, + Graph: graphState.Name, + Nodes: runtime.nodeCount, + Relationships: runtime.relationshipCount, + Duration: time.Since(graphStarted), + }) + return nil +} + +func (s *dumpRunner) finalize() (DumpResult, error) { + if err := s.recountEveryGraph(); err != nil { + return DumpResult{}, err + } + manifest := s.manifest() + if err := s.ctx.Err(); err != nil { + return DumpResult{}, fmt.Errorf("publish dump manifest: %w", err) + } + if err := writeCollection(s.config.Directory, manifest); err != nil { + return DumpResult{}, fmt.Errorf("publish dump manifest: %w", err) + } + if err := s.store.Remove(); err != nil { + return DumpResult{}, fmt.Errorf("remove completed dump checkpoint: %w", err) + } + + result := DumpResult{ + ManifestPath: filepath.Join(s.config.Directory, collection.ManifestName), + GraphCount: len(s.graphs), + } + for _, graph := range s.graphs { + result.NodeCount += graph.NodeCount + result.RelationshipCount += graph.RelationshipCount + } + return result, nil +} + +func (s *dumpRunner) processNodes(graphState *checkpoint.GraphState, runtime *dumpGraphRuntime) error { + phaseStarted := time.Now() + observe.Emit(s.ctx, s.config.Observer, observe.PhaseStarted{ + Operation: dumpOperationName, + Graph: graphState.Name, + Phase: string(checkpoint.PhaseNodes), + Completed: runtime.nodeCount, + Total: graphState.Snapshot.NodeCount, + }) + + var active []entity.Node + activeCounts := scrub.ActionCounts{} + var activeLastID uint64 + flush := func() error { + if len(active) == 0 { + return nil + } + if err := s.ctx.Err(); err != nil { + return fmt.Errorf("dump graph %q nodes: %w", graphState.Name, err) + } + shard, err := writeNodeShard( + s.config.Directory, + graphState.Name, + len(graphState.NodeShards)+1, + activeLastID, + activeCounts, + active, + s.config.JSONL, + s.config.Parquet, + ) + if err != nil { + return err + } + graphState.NodeShards = append(graphState.NodeShards, shard) + graphState.NodeCursor = shard.LastSourceID + if err := s.store.Save(s.state); err != nil { + return fmt.Errorf("checkpoint graph %q node shard %d: %w", graphState.Name, shard.Index, err) + } + emitNodeShardCommitted(s.ctx, s.config.Observer, graphState.Name, shard) + active = nil + activeCounts = scrub.ActionCounts{} + activeLastID = 0 + return nil + } + + for { + if err := s.ctx.Err(); err != nil { + return fmt.Errorf("dump graph %q nodes: %w", graphState.Name, err) + } + batch, err := runtime.source.NextNodes(s.ctx) + if err != nil { + return fmt.Errorf("dump graph %q nodes: %w", graphState.Name, err) + } + if err := s.ctx.Err(); err != nil { + return fmt.Errorf("dump graph %q nodes: %w", graphState.Name, err) + } + if len(batch.Entities) == 0 { + break + } + for _, node := range batch.Entities { + if err := s.ctx.Err(); err != nil { + return fmt.Errorf("dump graph %q nodes: %w", graphState.Name, err) + } + sourceID, err := parseDumpSourceID(node.SourceID) + if err != nil { + return fmt.Errorf("dump graph %q node source ID: %w", graphState.Name, err) + } + if s.scrubber != nil { + counts := s.scrubber.Scrub(node.Properties) + activeCounts.Add(counts) + runtime.totalScrubCounts.Add(counts) + } + if err := runtime.builder.ObserveNode(node); err != nil { + return fmt.Errorf("dump graph %q node metrics: %w", graphState.Name, err) + } + runtime.catalog.observeNode(node) + active = append(active, node) + activeLastID = sourceID + runtime.nodeCount++ + if runtime.nodeCount > graphState.Snapshot.NodeCount { + return fmt.Errorf( + "%w: graph %q got at least nodes=%d want nodes=%d", + ErrSourceCountChanged, + graphState.Name, + runtime.nodeCount, + graphState.Snapshot.NodeCount, + ) + } + observe.Emit(s.ctx, s.config.Observer, observe.PhaseProgress{ + Operation: dumpOperationName, + Graph: graphState.Name, + Phase: string(checkpoint.PhaseNodes), + Completed: runtime.nodeCount, + Total: graphState.Snapshot.NodeCount, + }) + if len(active) == s.config.ShardSize { + if err := flush(); err != nil { + return err + } + } + } + } + if runtime.nodeCount != graphState.Snapshot.NodeCount { + return fmt.Errorf( + "%w: graph %q got nodes=%d want nodes=%d", + ErrSourceCountChanged, + graphState.Name, + runtime.nodeCount, + graphState.Snapshot.NodeCount, + ) + } + if err := flush(); err != nil { + return err + } + if err := s.ctx.Err(); err != nil { + return fmt.Errorf("dump graph %q nodes: %w", graphState.Name, err) + } + + graphState.Phase = checkpoint.PhaseRelationships + if err := s.store.Save(s.state); err != nil { + return fmt.Errorf("checkpoint graph %q node phase completion: %w", graphState.Name, err) + } + observe.Emit(s.ctx, s.config.Observer, observe.PhaseCompleted{ + Operation: dumpOperationName, + Graph: graphState.Name, + Phase: string(checkpoint.PhaseNodes), + Completed: runtime.nodeCount, + Duration: time.Since(phaseStarted), + }) + return nil +} + +func (s *dumpRunner) processRelationships(graphState *checkpoint.GraphState, runtime *dumpGraphRuntime) error { + phaseStarted := time.Now() + observe.Emit(s.ctx, s.config.Observer, observe.PhaseStarted{ + Operation: dumpOperationName, + Graph: graphState.Name, + Phase: string(checkpoint.PhaseRelationships), + Completed: runtime.relationshipCount, + Total: graphState.Snapshot.RelationshipCount, + }) + + var active []entity.Relationship + activeCounts := scrub.ActionCounts{} + var activeLastID uint64 + flush := func() error { + if len(active) == 0 { + return nil + } + if err := s.ctx.Err(); err != nil { + return fmt.Errorf("dump graph %q relationships: %w", graphState.Name, err) + } + shard, err := writeRelationshipShard( + s.config.Directory, + graphState.Name, + len(graphState.RelationshipShards)+1, + activeLastID, + activeCounts, + active, + s.config.JSONL, + s.config.Parquet, + ) + if err != nil { + return err + } + graphState.RelationshipShards = append(graphState.RelationshipShards, shard) + graphState.RelationshipCursor = shard.LastSourceID + if err := s.store.Save(s.state); err != nil { + return fmt.Errorf("checkpoint graph %q relationship shard %d: %w", graphState.Name, shard.Index, err) + } + emitRelationshipShardCommitted(s.ctx, s.config.Observer, graphState.Name, shard) + active = nil + activeCounts = scrub.ActionCounts{} + activeLastID = 0 + return nil + } + + for { + if err := s.ctx.Err(); err != nil { + return fmt.Errorf("dump graph %q relationships: %w", graphState.Name, err) + } + batch, err := runtime.source.NextRelationships(s.ctx) + if err != nil { + return fmt.Errorf("dump graph %q relationships: %w", graphState.Name, err) + } + if err := s.ctx.Err(); err != nil { + return fmt.Errorf("dump graph %q relationships: %w", graphState.Name, err) + } + if len(batch.Entities) == 0 { + break + } + for _, relationship := range batch.Entities { + if err := s.ctx.Err(); err != nil { + return fmt.Errorf("dump graph %q relationships: %w", graphState.Name, err) + } + sourceID, err := parseDumpSourceID(relationship.SourceID) + if err != nil { + return fmt.Errorf("dump graph %q relationship source ID: %w", graphState.Name, err) + } + if s.scrubber != nil { + counts := s.scrubber.Scrub(relationship.Properties) + activeCounts.Add(counts) + runtime.totalScrubCounts.Add(counts) + } + if err := runtime.builder.ObserveRelationship(relationship); err != nil { + return fmt.Errorf("dump graph %q relationship metrics: %w", graphState.Name, err) + } + runtime.catalog.observeRelationship(relationship) + active = append(active, relationship) + activeLastID = sourceID + runtime.relationshipCount++ + if runtime.relationshipCount > graphState.Snapshot.RelationshipCount { + return fmt.Errorf( + "%w: graph %q got at least relationships=%d want relationships=%d", + ErrSourceCountChanged, + graphState.Name, + runtime.relationshipCount, + graphState.Snapshot.RelationshipCount, + ) + } + observe.Emit(s.ctx, s.config.Observer, observe.PhaseProgress{ + Operation: dumpOperationName, + Graph: graphState.Name, + Phase: string(checkpoint.PhaseRelationships), + Completed: runtime.relationshipCount, + Total: graphState.Snapshot.RelationshipCount, + }) + if len(active) == s.config.ShardSize { + if err := flush(); err != nil { + return err + } + } + } + } + if runtime.relationshipCount != graphState.Snapshot.RelationshipCount { + return fmt.Errorf( + "%w: graph %q got relationships=%d want relationships=%d", + ErrSourceCountChanged, + graphState.Name, + runtime.relationshipCount, + graphState.Snapshot.RelationshipCount, + ) + } + if err := flush(); err != nil { + return err + } + if err := s.ctx.Err(); err != nil { + return fmt.Errorf("dump graph %q relationships: %w", graphState.Name, err) + } + + graphState.Phase = checkpoint.PhaseComplete + if err := s.store.Save(s.state); err != nil { + return fmt.Errorf("checkpoint graph %q relationship phase completion: %w", graphState.Name, err) + } + observe.Emit(s.ctx, s.config.Observer, observe.PhaseCompleted{ + Operation: dumpOperationName, + Graph: graphState.Name, + Phase: string(checkpoint.PhaseRelationships), + Completed: runtime.relationshipCount, + Duration: time.Since(phaseStarted), + }) + return nil +} + +func (s *dumpRunner) recountEveryGraph() error { + for index, graphName := range s.config.Graphs { + if err := s.ctx.Err(); err != nil { + return fmt.Errorf("recount dump graph %q: %w", graphName, err) + } + source, err := dawgs.NewSource(s.database, graphName, s.config.EntityBatchSize) + if err != nil { + return fmt.Errorf("prepare recount source for graph %q: %w", graphName, err) + } + current, err := source.Snapshot(s.ctx) + if err != nil { + return fmt.Errorf("recount dump graph %q: %w", graphName, err) + } + if err := s.ctx.Err(); err != nil { + return fmt.Errorf("recount dump graph %q: %w", graphName, err) + } + snapshot := s.state.Graphs[index].Snapshot + if current != snapshot { + return fmt.Errorf( + "%w: graph %q got nodes=%d relationships=%d want nodes=%d relationships=%d", + ErrSourceCountChanged, + graphName, + current.NodeCount, + current.RelationshipCount, + snapshot.NodeCount, + snapshot.RelationshipCount, + ) + } + } + return nil +} + +func (s *dumpRunner) manifest() collection.Manifest { + outputs := collection.OutputConfig{} + if s.config.JSONL != nil { + outputs.JSONL = &collection.JSONLOutput{ + SchemaVersion: jsonl.SchemaVersion, + Codec: string(s.config.JSONL.Codec), + Level: s.config.JSONL.Level, + } + } + if s.config.Parquet != nil { + outputs.Parquet = &collection.ParquetOutput{SchemaVersion: parquet.SchemaVersion} + } + scrubMetadata := collection.ScrubMetadata{Enabled: s.scrubber != nil} + if s.scrubber != nil { + scrubMetadata.RulesFingerprint = s.scrubber.RulesFingerprint() + scrubMetadata.SaltFingerprint = s.scrubber.SaltFingerprint() + } + return collection.Manifest{ + Format: collection.Format, + CreatedAt: time.Now().UTC(), + Outputs: outputs, + Scrub: scrubMetadata, + Graphs: append([]collection.Graph(nil), s.graphs...), + } +} + +func dumpCheckpointIdentity(config DumpConfig, compiled *scrub.Scrubber) checkpoint.Identity { + identity := checkpoint.Identity{ + Graphs: append([]string(nil), config.Graphs...), + EntityBatchSize: config.EntityBatchSize, + ShardSize: config.ShardSize, + JSONLEnabled: config.JSONL != nil, + ParquetEnabled: config.Parquet != nil, + JSONLSchemaVersion: jsonl.SchemaVersion, + ParquetSchemaVersion: parquet.SchemaVersion, + ScrubEnabled: compiled != nil, + } + if config.JSONL != nil { + identity.JSONLCodec = string(config.JSONL.Codec) + identity.JSONLLevel = config.JSONL.Level + } + if compiled != nil { + identity.ScrubRulesFingerprint = compiled.RulesFingerprint() + identity.ScrubSaltFingerprint = compiled.SaltFingerprint() + } + return identity +} + +func parseDumpSourceID(value string) (uint64, error) { + parsed, err := strconv.ParseUint(value, 10, 64) + if err != nil || parsed == 0 || strconv.FormatUint(parsed, 10) != value { + return 0, fmt.Errorf("expected canonical positive Dawgs ID, got %q", value) + } + return parsed, nil +} + +type dumpKindCatalog struct { + seen map[string]struct{} + values []string +} + +func newDumpKindCatalog() *dumpKindCatalog { + return &dumpKindCatalog{seen: make(map[string]struct{})} +} + +func (s *dumpKindCatalog) observeNode(node entity.Node) { + for _, kind := range node.Kinds { + s.add(kind) + } +} + +func (s *dumpKindCatalog) observeRelationship(relationship entity.Relationship) { + s.add(relationship.Kind) +} + +func (s *dumpKindCatalog) add(kind string) { + if _, found := s.seen[kind]; found { + return + } + s.seen[kind] = struct{}{} + s.values = append(s.values, kind) +} + +func emitNodeShardCommitted(ctx context.Context, observer observe.Observer, graphName string, shard collection.NodeShard) { + event := observe.ShardCommitted{ + Graph: graphName, + EntityType: "node", + Index: shard.Index, + Count: shard.Count, + } + if shard.JSONL != nil { + event.JSONLPath = shard.JSONL.Path + event.JSONLBytes = shard.JSONL.StoredBytes + } + if shard.Parquet != nil { + event.ParquetPath = shard.Parquet.Path + event.ParquetBytes = shard.Parquet.StoredBytes + } + observe.Emit(ctx, observer, event) +} + +func emitRelationshipShardCommitted(ctx context.Context, observer observe.Observer, graphName string, shard collection.RelationshipShard) { + event := observe.ShardCommitted{ + Graph: graphName, + EntityType: "relationship", + Index: shard.Index, + Count: shard.Count, + } + if shard.JSONL != nil { + event.JSONLPath = shard.JSONL.Path + event.JSONLBytes = shard.JSONL.StoredBytes + } + if shard.Parquet != nil { + event.ParquetPath = shard.Parquet.Path + event.ParquetBytes = shard.Parquet.StoredBytes + } + observe.Emit(ctx, observer, event) +} diff --git a/ret/dump_resume.go b/ret/dump_resume.go new file mode 100644 index 00000000..284d5732 --- /dev/null +++ b/ret/dump_resume.go @@ -0,0 +1,213 @@ +package ret + +import ( + "errors" + "fmt" + "os" + "path/filepath" + + "github.com/specterops/dawgs/ret/checkpoint" + "github.com/specterops/dawgs/ret/collection" + "github.com/specterops/dawgs/ret/dawgs" + "github.com/specterops/dawgs/ret/entity" + "github.com/specterops/dawgs/ret/metrics" + "github.com/specterops/dawgs/ret/scrub" +) + +func (s *dumpRunner) runResume() (DumpResult, error) { + if err := s.prepareResume(); err != nil { + return DumpResult{}, err + } + + for index, graphName := range s.config.Graphs { + if err := s.ctx.Err(); err != nil { + return DumpResult{}, fmt.Errorf("resume dump graph %q: %w", graphName, err) + } + + var runtime *dumpGraphRuntime + var err error + if index < len(s.state.Graphs) { + runtime, err = s.reconstructGraph(index) + } else { + runtime, err = s.startFreshGraph(graphName) + } + if err != nil { + return DumpResult{}, err + } + if err := s.processGraph(index, runtime); err != nil { + return DumpResult{}, err + } + } + + return s.finalize() +} + +func (s *dumpRunner) prepareResume() error { + rootInfo, err := os.Lstat(s.config.Directory) + if errors.Is(err, os.ErrNotExist) { + return fmt.Errorf("%w: %s", ErrCheckpointMissing, s.config.Directory) + } + if err != nil { + return fmt.Errorf("inspect resume destination: %w", err) + } + if !rootInfo.IsDir() || rootInfo.Mode()&os.ModeSymlink != 0 { + return fmt.Errorf("%w: resume destination is not a non-symlink directory", ErrInvalidCollection) + } + + manifestPath := filepath.Join(s.config.Directory, collection.ManifestName) + if _, err := os.Lstat(manifestPath); err == nil { + return fmt.Errorf("%w: resume destination already contains %s", ErrDestinationExists, collection.ManifestName) + } else if !errors.Is(err, os.ErrNotExist) { + return fmt.Errorf("inspect resume manifest: %w", err) + } + + state, exists, err := s.store.Load() + if err != nil { + return fmt.Errorf("%w: load resume checkpoint: %w", ErrInvalidCollection, err) + } + if !exists { + return fmt.Errorf("%w: %s", ErrCheckpointMissing, filepath.Join(s.config.Directory, checkpoint.FileName)) + } + if err := checkpoint.ValidateIdentity(s.identity, state.Identity); err != nil { + return fmt.Errorf("%w: resume checkpoint identity: %w", ErrInvalidConfig, err) + } + if err := s.validateResumeCounts(state); err != nil { + return err + } + if err := s.store.CleanupOrphans(state); err != nil { + return fmt.Errorf("%w: clean resume crash artifacts: %w", ErrInvalidCollection, err) + } + s.state = state + return nil +} + +func (s *dumpRunner) validateResumeCounts(state checkpoint.State) error { + for _, graphState := range state.Graphs { + if err := s.ctx.Err(); err != nil { + return fmt.Errorf("validate resume graph %q counts: %w", graphState.Name, err) + } + source, err := dawgs.NewSource(s.database, graphState.Name, s.config.EntityBatchSize) + if err != nil { + return fmt.Errorf("prepare resume count source for graph %q: %w", graphState.Name, err) + } + current, err := source.Snapshot(s.ctx) + if err != nil { + return fmt.Errorf("validate resume graph %q counts: %w", graphState.Name, err) + } + if err := s.ctx.Err(); err != nil { + return fmt.Errorf("validate resume graph %q counts: %w", graphState.Name, err) + } + if current != graphState.Snapshot { + return fmt.Errorf( + "%w: graph %q got nodes=%d relationships=%d want nodes=%d relationships=%d", + ErrSourceCountChanged, + graphState.Name, + current.NodeCount, + current.RelationshipCount, + graphState.Snapshot.NodeCount, + graphState.Snapshot.RelationshipCount, + ) + } + } + return nil +} + +func (s *dumpRunner) reconstructGraph(index int) (*dumpGraphRuntime, error) { + graphState := &s.state.Graphs[index] + source, err := dawgs.NewSource(s.database, graphState.Name, s.config.EntityBatchSize) + if err != nil { + return nil, fmt.Errorf("prepare resumed dump source for graph %q: %w", graphState.Name, err) + } + source.SetNodeCursor(graphState.NodeCursor) + source.SetRelationshipCursor(graphState.RelationshipCursor) + + runtime := &dumpGraphRuntime{ + source: source, + builder: metrics.NewBuilder(), + catalog: newDumpKindCatalog(), + totalScrubCounts: scrub.ActionCounts{}, + } + for _, shard := range graphState.NodeShards { + if err := s.reconstructNodeShard(graphState.Name, shard, runtime); err != nil { + return nil, err + } + runtime.totalScrubCounts.Add(shard.ScrubCounts) + } + for _, shard := range graphState.RelationshipShards { + if err := s.reconstructRelationshipShard(graphState.Name, shard, runtime); err != nil { + return nil, err + } + runtime.totalScrubCounts.Add(shard.ScrubCounts) + } + return runtime, nil +} + +func (s *dumpRunner) reconstructNodeShard( + graphName string, + shard collection.NodeShard, + runtime *dumpGraphRuntime, +) error { + visit := func(node entity.Node) error { + if err := s.ctx.Err(); err != nil { + return err + } + if err := runtime.builder.ObserveNode(node); err != nil { + return err + } + runtime.catalog.observeNode(node) + runtime.nodeCount++ + return nil + } + + var err error + switch { + case shard.JSONL != nil: + err = collection.ReadJSONLNodes(s.config.Directory, *shard.JSONL, visit) + case shard.Parquet != nil: + err = collection.ReadParquetNodes(s.config.Directory, *shard.Parquet, visit) + default: + err = errors.New("committed node shard has no concrete artifact") + } + if err != nil { + if contextErr := s.ctx.Err(); contextErr != nil { + return fmt.Errorf("reconstruct graph %q node shard %d: %w", graphName, shard.Index, contextErr) + } + return fmt.Errorf("%w: reconstruct graph %q node shard %d: %w", ErrArtifactIntegrity, graphName, shard.Index, err) + } + return nil +} + +func (s *dumpRunner) reconstructRelationshipShard( + graphName string, + shard collection.RelationshipShard, + runtime *dumpGraphRuntime, +) error { + visit := func(relationship entity.Relationship) error { + if err := s.ctx.Err(); err != nil { + return err + } + if err := runtime.builder.ObserveRelationship(relationship); err != nil { + return err + } + runtime.catalog.observeRelationship(relationship) + runtime.relationshipCount++ + return nil + } + + var err error + switch { + case shard.JSONL != nil: + err = collection.ReadJSONLRelationships(s.config.Directory, *shard.JSONL, visit) + case shard.Parquet != nil: + err = collection.ReadParquetRelationships(s.config.Directory, *shard.Parquet, visit) + default: + err = errors.New("committed relationship shard has no concrete artifact") + } + if err != nil { + if contextErr := s.ctx.Err(); contextErr != nil { + return fmt.Errorf("reconstruct graph %q relationship shard %d: %w", graphName, shard.Index, contextErr) + } + return fmt.Errorf("%w: reconstruct graph %q relationship shard %d: %w", ErrArtifactIntegrity, graphName, shard.Index, err) + } + return nil +} diff --git a/ret/dump_resume_test.go b/ret/dump_resume_test.go new file mode 100644 index 00000000..73b6cb08 --- /dev/null +++ b/ret/dump_resume_test.go @@ -0,0 +1,711 @@ +package ret + +import ( + "context" + "errors" + "os" + "path/filepath" + "sync" + "testing" + + "github.com/specterops/dawgs/graph" + "github.com/specterops/dawgs/ret/checkpoint" + "github.com/specterops/dawgs/ret/collection" + "github.com/specterops/dawgs/ret/dawgs" + "github.com/specterops/dawgs/ret/entity" + "github.com/specterops/dawgs/ret/jsonl" + "github.com/specterops/dawgs/ret/observe" + "github.com/specterops/dawgs/ret/parquet" + "github.com/specterops/dawgs/ret/scrub" + "github.com/stretchr/testify/require" +) + +func TestDumpResumeReconstructsMetricsAndContinuesAfterCommittedShard(t *testing.T) { + // Break caught: resuming from the cursor without replaying committed nodes, + // which loses metrics endpoint state, kinds, or the first shard. + config, database := interruptedDumpAfterFirstNodeShard(t, true, true) + config.Resume = true + var starts []observe.PhaseStarted + config.Observer = observe.ObserverFunc(func(_ context.Context, event observe.Event) { + if value, ok := event.(observe.PhaseStarted); ok { + starts = append(starts, value) + } + }) + + result, err := Dump(context.Background(), database, config) + + require.NoError(t, err) + manifest := readDumpManifest(t, result.ManifestPath) + require.EqualValues(t, manifest.Graphs[0].NodeCount, manifest.Graphs[0].Metrics.NodeCount) + require.EqualValues(t, manifest.Graphs[0].RelationshipCount, manifest.Graphs[0].Metrics.RelationshipCount) + require.Equal(t, []int{1, 2, 3}, dumpNodeShardIndices(manifest.Graphs[0])) + require.NoFileExists(t, filepath.Join(config.Directory, checkpoint.FileName)) + require.Len(t, starts, 2) + require.EqualValues(t, 1, starts[0].Completed) + require.EqualValues(t, 0, starts[1].Completed) +} + +func TestDumpResumeReconstructsParquetOnlyCheckpoint(t *testing.T) { + // Break caught: making resume depend on loadable JSONL even though Parquet + // contains enough canonical entities to rebuild dump state. + config, database := interruptedDumpAfterFirstNodeShard(t, false, true) + config.Resume = true + + result, err := Dump(context.Background(), database, config) + + require.NoError(t, err) + manifest := readDumpManifest(t, result.ManifestPath) + require.Equal(t, []int{1, 2, 3}, dumpNodeShardIndices(manifest.Graphs[0])) + require.Nil(t, manifest.Graphs[0].NodeShards[0].JSONL) + require.NotNil(t, manifest.Graphs[0].NodeShards[0].Parquet) + require.EqualValues(t, 3, manifest.Graphs[0].Metrics.NodeCount) + require.EqualValues(t, 1, manifest.Graphs[0].Metrics.RelationshipCount) +} + +func TestDumpResumeContinuesRelationshipPhaseAfterCommittedShard(t *testing.T) { + // Break caught: treating every checkpoint as a nodes-phase checkpoint and + // either rescanning nodes or omitting committed relationship metrics. + database := newDumpTestDatabase(map[string]*dumpTestGraph{ + "asset": { + nodes: dumpTestNodes(1, 2, 3), + relationships: []*graph.Relationship{ + graph.NewRelationship(10, 1, 2, graph.NewProperties(), graph.StringKind("FIRST")), + graph.NewRelationship(11, 2, 3, graph.NewProperties(), graph.StringKind("SECOND")), + }, + }, + }) + config := validRootDumpConfig(t) + config.EntityBatchSize = 3 + config.ShardSize = 1 + originalWrite := writeJSONLRelationships + calls := 0 + injected := errors.New("injected second relationship shard failure") + writeJSONLRelationships = func( + tempPath, relativePath string, + output jsonl.Config, + relationships []entity.Relationship, + ) (collection.JSONLArtifact, error) { + calls++ + if calls == 2 { + return collection.JSONLArtifact{}, injected + } + return originalWrite(tempPath, relativePath, output, relationships) + } + t.Cleanup(func() { writeJSONLRelationships = originalWrite }) + + _, err := Dump(context.Background(), database, config) + writeJSONLRelationships = originalWrite + require.ErrorIs(t, err, injected) + state, exists, err := (checkpoint.Store{Root: config.Directory}).Load() + require.NoError(t, err) + require.True(t, exists) + require.Equal(t, checkpoint.PhaseRelationships, state.Graphs[0].Phase) + require.Len(t, state.Graphs[0].RelationshipShards, 1) + + config.Resume = true + result, err := Dump(context.Background(), database, config) + + require.NoError(t, err) + manifest := readDumpManifest(t, result.ManifestPath) + require.Equal(t, []int{1, 2}, dumpRelationshipShardIndices(manifest.Graphs[0])) + require.EqualValues(t, 3, manifest.Graphs[0].Metrics.NodeCount) + require.EqualValues(t, 2, manifest.Graphs[0].Metrics.RelationshipCount) + require.Equal(t, []string{"Entity", "FIRST", "SECOND"}, manifest.Graphs[0].KindCatalog) + require.NoFileExists(t, filepath.Join(config.Directory, checkpoint.FileName)) +} + +func TestDumpResumePublishesFromCompleteCheckpoint(t *testing.T) { + // Break caught: requiring an active scan phase on resume instead of + // reconstructing a complete checkpoint and retrying final publication. + database := newDumpTestDatabase(map[string]*dumpTestGraph{ + "asset": { + nodes: dumpTestNodes(1, 2), + relationships: dumpTestRelationships(10, 1, 2, "LINKED"), + }, + }) + config := validRootDumpConfig(t) + injected := errors.New("injected manifest publication failure") + originalWrite := writeCollection + writeCollection = func(string, collection.Manifest) error { return injected } + t.Cleanup(func() { writeCollection = originalWrite }) + + _, err := Dump(context.Background(), database, config) + writeCollection = originalWrite + require.ErrorIs(t, err, injected) + state, exists, err := (checkpoint.Store{Root: config.Directory}).Load() + require.NoError(t, err) + require.True(t, exists) + require.Equal(t, checkpoint.PhaseComplete, state.Graphs[0].Phase) + require.NoFileExists(t, filepath.Join(config.Directory, collection.ManifestName)) + + config.Resume = true + result, err := Dump(context.Background(), database, config) + + require.NoError(t, err) + manifest := readDumpManifest(t, result.ManifestPath) + require.EqualValues(t, 2, manifest.Graphs[0].Metrics.NodeCount) + require.EqualValues(t, 1, manifest.Graphs[0].Metrics.RelationshipCount) + require.NoFileExists(t, filepath.Join(config.Directory, checkpoint.FileName)) +} + +func TestDumpResumePreservesScrubbedArtifactsAndActionCounts(t *testing.T) { + // Break caught: losing committed scrub metadata on resume or applying + // scrubbing only to the pre-interruption portion of the source. + database := newDumpTestDatabase(map[string]*dumpTestGraph{ + "asset": { + nodes: []*graph.Node{ + graph.NewNode(1, graph.AsProperties(map[string]any{"password": "one"}), graph.StringKind("Entity")), + graph.NewNode(2, graph.AsProperties(map[string]any{"password": "two"}), graph.StringKind("Entity")), + graph.NewNode(3, graph.AsProperties(map[string]any{"password": "three"}), graph.StringKind("Entity")), + }, + }, + }) + config := validRootDumpConfig(t) + config.EntityBatchSize = 3 + config.ShardSize = 1 + config.Parquet = nil + originalWrite := writeJSONLNodes + calls := 0 + injected := errors.New("injected scrub resume writer failure") + writeJSONLNodes = func( + tempPath, relativePath string, + output jsonl.Config, + nodes []entity.Node, + ) (collection.JSONLArtifact, error) { + calls++ + if calls == 2 { + return collection.JSONLArtifact{}, injected + } + return originalWrite(tempPath, relativePath, output, nodes) + } + t.Cleanup(func() { writeJSONLNodes = originalWrite }) + + _, err := Dump(context.Background(), database, config) + writeJSONLNodes = originalWrite + require.ErrorIs(t, err, injected) + state, exists, err := (checkpoint.Store{Root: config.Directory}).Load() + require.NoError(t, err) + require.True(t, exists) + require.Equal(t, scrub.ActionCounts{Redact: 1}, state.Graphs[0].NodeShards[0].ScrubCounts) + + config.Resume = true + result, err := Dump(context.Background(), database, config) + + require.NoError(t, err) + manifest := readDumpManifest(t, result.ManifestPath) + require.Len(t, manifest.Graphs[0].NodeShards, 3) + for _, shard := range manifest.Graphs[0].NodeShards { + require.Equal(t, scrub.ActionCounts{Redact: 1}, shard.ScrubCounts) + nodes, err := readJSONLNodesForTest(config.Directory, *shard.JSONL) + require.NoError(t, err) + for _, node := range nodes { + require.Equal(t, "[REDACTED]", node.Properties["password"]) + } + } +} + +func TestDumpResumeReplayCancellationIsNotArtifactIntegrityFailure(t *testing.T) { + // Break caught: wrapping an active context cancellation from a replay visitor + // with ErrArtifactIntegrity and falsely classifying durable bytes as corrupt. + for _, test := range []struct { + name string + setup func(*testing.T) (DumpConfig, *dumpTestDatabase) + cancelOnCall int + errorContext string + }{ + { + name: "node replay", + setup: func(t *testing.T) (DumpConfig, *dumpTestDatabase) { + return interruptedDumpAfterFirstNodeShard(t, true, false) + }, + cancelOnCall: 3, + errorContext: `reconstruct graph "asset" node shard 1`, + }, + { + name: "relationship replay", + setup: func(t *testing.T) (DumpConfig, *dumpTestDatabase) { + database := newDumpTestDatabase(map[string]*dumpTestGraph{ + "asset": { + nodes: dumpTestNodes(1, 2), + relationships: dumpTestRelationships(10, 1, 2, "LINKED"), + }, + }) + config := validRootDumpConfig(t) + injected := errors.New("injected complete checkpoint manifest failure") + originalWrite := writeCollection + writeCollection = func(string, collection.Manifest) error { return injected } + t.Cleanup(func() { writeCollection = originalWrite }) + _, err := Dump(context.Background(), database, config) + writeCollection = originalWrite + require.ErrorIs(t, err, injected) + return config, database + }, + cancelOnCall: 5, + errorContext: `reconstruct graph "asset" relationship shard 1`, + }, + } { + t.Run(test.name, func(t *testing.T) { + config, database := test.setup(t) + checkpointBefore := mustReadFile(t, filepath.Join(config.Directory, checkpoint.FileName)) + ctx := newStagedCancelContext() + armed := false + database.onRelationshipCount = func(context.Context, string) { + if !armed { + armed = true + ctx.arm(test.cancelOnCall) + } + } + config.Resume = true + + _, err := Dump(ctx, database, config) + + require.ErrorIs(t, err, context.Canceled) + require.False(t, errors.Is(err, ErrArtifactIntegrity)) + require.ErrorContains(t, err, test.errorContext) + require.Equal(t, checkpointBefore, mustReadFile(t, filepath.Join(config.Directory, checkpoint.FileName))) + require.NoFileExists(t, filepath.Join(config.Directory, collection.ManifestName)) + }) + } +} + +func TestDumpResumeRejectsIdentityChanges(t *testing.T) { + // Break caught: continuing with configuration that changes cursor behavior, + // logical artifacts, concrete encoding, or deterministic scrubbing. + for _, test := range []struct { + name string + mutate func(*DumpConfig) + match string + }{ + { + name: "graph order", + mutate: func(config *DumpConfig) { + config.Graphs = []string{"other", "asset"} + }, + match: "ordered graph names", + }, + { + name: "batch size", + mutate: func(config *DumpConfig) { config.EntityBatchSize++ }, + match: "entity batch size", + }, + { + name: "shard size", + mutate: func(config *DumpConfig) { config.ShardSize++ }, + match: "shard size", + }, + { + name: "JSONL codec", + mutate: func(config *DumpConfig) { + config.JSONL.Codec = jsonl.CodecZstd + }, + match: "JSONL codec", + }, + { + name: "JSONL level", + mutate: func(config *DumpConfig) { config.JSONL.Level++ }, + match: "JSONL level", + }, + { + name: "enabled output", + mutate: func(config *DumpConfig) { config.Parquet = nil }, + match: "Parquet enabled", + }, + { + name: "scrub enabled", + mutate: func(config *DumpConfig) { config.Scrub = nil }, + match: "scrub enabled", + }, + { + name: "scrub rules", + mutate: func(config *DumpConfig) { + config.Scrub.Rules.FakeDomain = "changed.example" + }, + match: "scrub rules fingerprint", + }, + { + name: "scrub salt", + mutate: func(config *DumpConfig) { config.Scrub.Salt = "changed-salt" }, + match: "scrub salt fingerprint", + }, + } { + t.Run(test.name, func(t *testing.T) { + database := newDumpTestDatabase(map[string]*dumpTestGraph{ + "asset": { + nodes: dumpTestNodesWithKinds([]string{"User"}, 1, 2, 3), + relationships: dumpTestRelationships(10, 1, 3, "MEMBER_OF"), + }, + "other": {}, + }) + config := validRootDumpConfig(t) + config.Graphs = []string{"asset", "other"} + config.EntityBatchSize = 3 + config.ShardSize = 1 + config.JSONL.Codec = jsonl.CodecGzip + config.JSONL.Level = 1 + config, database = interruptedDumpFromConfig(t, config, database) + config.Resume = true + test.mutate(&config) + + _, err := Dump(context.Background(), database, config) + + require.ErrorIs(t, err, ErrInvalidConfig) + require.ErrorContains(t, err, test.match) + require.NoFileExists(t, filepath.Join(config.Directory, collection.ManifestName)) + require.FileExists(t, filepath.Join(config.Directory, checkpoint.FileName)) + }) + } +} + +func TestDumpResumeRemovesRecognizedUncheckpointedArtifact(t *testing.T) { + // Break caught: either retaining the known next-shard crash-window artifact + // or deleting it only after a writer collides with its deterministic path. + config, database := interruptedDumpAfterFirstNodeShard(t, true, false) + orphan, err := writeNodeShard( + config.Directory, + "asset", + 2, + 99, + scrub.ActionCounts{}, + []entity.Node{{SourceID: "99", Kinds: []string{"Orphan"}}}, + config.JSONL, + config.Parquet, + ) + require.NoError(t, err) + require.FileExists(t, filepath.Join(config.Directory, filepath.FromSlash(orphan.JSONL.Path))) + + originalWrite := writeJSONLNodes + writeJSONLNodes = func(string, string, jsonl.Config, []entity.Node) (collection.JSONLArtifact, error) { + return collection.JSONLArtifact{}, errors.New("stop after resume cleanup") + } + t.Cleanup(func() { writeJSONLNodes = originalWrite }) + config.Resume = true + + _, err = Dump(context.Background(), database, config) + + writeJSONLNodes = originalWrite + require.ErrorContains(t, err, "stop after resume cleanup") + require.NoFileExists(t, filepath.Join(config.Directory, filepath.FromSlash(orphan.JSONL.Path))) + require.FileExists(t, filepath.Join(config.Directory, checkpoint.FileName)) +} + +func TestDumpResumeRejectsUnknownFileWithoutDeletingIt(t *testing.T) { + // Break caught: broad cleanup that deletes caller data merely because it is + // present in a resumable collection directory. + config, database := interruptedDumpAfterFirstNodeShard(t, true, false) + unknown := filepath.Join(config.Directory, "caller-owned.txt") + require.NoError(t, os.WriteFile(unknown, []byte("keep"), 0o600)) + config.Resume = true + + _, err := Dump(context.Background(), database, config) + + require.ErrorIs(t, err, ErrInvalidCollection) + require.FileExists(t, unknown) + require.Equal(t, "keep", string(mustReadFile(t, unknown))) + require.FileExists(t, filepath.Join(config.Directory, checkpoint.FileName)) +} + +func TestDumpResumeRejectsSourceCountChangeBeforeScanning(t *testing.T) { + // Break caught: resuming a cursor against a source whose entity totals no + // longer match the checkpoint snapshot. + config, database := interruptedDumpAfterFirstNodeShard(t, true, true) + database.graphs["asset"].snapshots = append(database.graphs["asset"].snapshots, dawgs.Snapshot{ + NodeCount: 4, + RelationshipCount: 1, + }) + config.Resume = true + + _, err := Dump(context.Background(), database, config) + + require.ErrorIs(t, err, ErrSourceCountChanged) + require.NoFileExists(t, filepath.Join(config.Directory, collection.ManifestName)) + require.FileExists(t, filepath.Join(config.Directory, checkpoint.FileName)) +} + +func TestDumpResumeCancellationDuringCountSnapshotStopsBeforeOrphanCleanup(t *testing.T) { + // Break caught: continuing from a successful resume count result after the + // count callback canceled, including deleting the recognized crash artifact. + config, database := interruptedDumpAfterFirstNodeShard(t, true, false) + orphan, err := writeNodeShard( + config.Directory, + "asset", + 2, + 99, + scrub.ActionCounts{}, + []entity.Node{{SourceID: "99", Kinds: []string{"Orphan"}}}, + config.JSONL, + config.Parquet, + ) + require.NoError(t, err) + orphanPath := filepath.Join(config.Directory, filepath.FromSlash(orphan.JSONL.Path)) + require.FileExists(t, orphanPath) + checkpointBefore := mustReadFile(t, filepath.Join(config.Directory, checkpoint.FileName)) + + ctx, cancel := context.WithCancel(context.Background()) + database.onRelationshipCount = func(context.Context, string) { + cancel() + } + config.Resume = true + + _, err = Dump(ctx, database, config) + + require.ErrorIs(t, err, context.Canceled) + require.FileExists(t, orphanPath) + require.Equal(t, checkpointBefore, mustReadFile(t, filepath.Join(config.Directory, checkpoint.FileName))) + require.NoFileExists(t, filepath.Join(config.Directory, collection.ManifestName)) +} + +func TestDumpRecountsEveryGraphBeforePublishingManifest(t *testing.T) { + // Break caught: publishing after successful scans without detecting a later + // total change in one of the source graphs. + database := newDumpTestDatabase(map[string]*dumpTestGraph{ + "asset": { + nodes: dumpTestNodes(1, 2, 3), + snapshots: []dawgs.Snapshot{{NodeCount: 3}, {NodeCount: 4}}, + }, + }) + config := validRootDumpConfig(t) + + _, err := Dump(context.Background(), database, config) + + require.ErrorIs(t, err, ErrSourceCountChanged) + require.NoFileExists(t, filepath.Join(config.Directory, collection.ManifestName)) + require.FileExists(t, filepath.Join(config.Directory, checkpoint.FileName)) +} + +func TestDumpClassifiesPhaseScanTotalMismatchBeforePartialShardPublication(t *testing.T) { + // Break caught: attempting to checkpoint a partial shard against a larger + // snapshot and returning a checkpoint-validation error instead of the public + // source-count change classification. + for _, test := range []struct { + name string + value *dumpTestGraph + artifact string + }{ + { + name: "nodes", + value: &dumpTestGraph{ + nodes: dumpTestNodes(1), + snapshots: []dawgs.Snapshot{{NodeCount: 3}}, + }, + artifact: collection.NodeJSONLPath("asset", 1, jsonl.CodecNone), + }, + { + name: "relationships", + value: &dumpTestGraph{ + nodes: dumpTestNodes(1, 2), + relationships: dumpTestRelationships(10, 1, 2, "LINKED"), + snapshots: []dawgs.Snapshot{{NodeCount: 2, RelationshipCount: 2}}, + }, + artifact: collection.RelationshipJSONLPath("asset", 1, jsonl.CodecNone), + }, + } { + t.Run(test.name, func(t *testing.T) { + database := newDumpTestDatabase(map[string]*dumpTestGraph{"asset": test.value}) + config := validRootDumpConfig(t) + config.ShardSize = 2 + + _, err := Dump(context.Background(), database, config) + + require.ErrorIs(t, err, ErrSourceCountChanged) + require.NoFileExists(t, filepath.Join(config.Directory, filepath.FromSlash(test.artifact))) + require.NoFileExists(t, filepath.Join(config.Directory, collection.ManifestName)) + require.FileExists(t, filepath.Join(config.Directory, checkpoint.FileName)) + }) + } +} + +func TestDumpDoesNotClaimToDetectSameCountSourceMutation(t *testing.T) { + // Break caught: accidentally promising content-level source consistency when + // the publication contract intentionally checks totals only. + database := newDumpTestDatabase(map[string]*dumpTestGraph{ + "asset": {nodes: dumpTestNodes(1, 2)}, + }) + fetches := 0 + database.onNodeFetch = func(graphName string) { + fetches++ + if fetches == 2 { + database.graphs[graphName].nodes[0] = graph.NewNode( + 1, + graph.NewProperties(), + graph.StringKind("MutatedAfterScan"), + ) + } + } + config := validRootDumpConfig(t) + config.EntityBatchSize = 2 + + result, err := Dump(context.Background(), database, config) + + require.NoError(t, err) + manifest := readDumpManifest(t, result.ManifestPath) + require.Equal(t, []string{"Entity"}, manifest.Graphs[0].KindCatalog) + require.Equal(t, "MutatedAfterScan", database.graphs["asset"].nodes[0].Kinds[0].String()) +} + +func TestDumpWriterFailurePreservesCheckpointedShardsAndNoManifest(t *testing.T) { + // Break caught: deleting durable progress or publishing a manifest after the + // next concrete shard writer fails. + config, _ := interruptedDumpAfterFirstNodeShard(t, true, true) + state, exists, err := (checkpoint.Store{Root: config.Directory}).Load() + + require.NoError(t, err) + require.True(t, exists) + require.Len(t, state.Graphs[0].NodeShards, 1) + require.FileExists(t, filepath.Join(config.Directory, filepath.FromSlash(state.Graphs[0].NodeShards[0].JSONL.Path))) + require.FileExists(t, filepath.Join(config.Directory, filepath.FromSlash(state.Graphs[0].NodeShards[0].Parquet.Path))) + require.NoFileExists(t, filepath.Join(config.Directory, collection.ManifestName)) +} + +func TestDumpResumeRequiresCheckpointAndRejectsPublishedManifest(t *testing.T) { + // Break caught: treating resume as a fresh dump when no durable checkpoint + // exists, or overwriting a completed collection. + missing := validRootDumpConfig(t) + missing.Resume = true + _, err := Dump(context.Background(), newDumpTestDatabase(map[string]*dumpTestGraph{"asset": {}}), missing) + require.ErrorIs(t, err, ErrCheckpointMissing) + + database := newDumpTestDatabase(map[string]*dumpTestGraph{"asset": {}}) + published := validRootDumpConfig(t) + result, err := Dump(context.Background(), database, published) + require.NoError(t, err) + require.FileExists(t, result.ManifestPath) + published.Resume = true + + _, err = Dump(context.Background(), database, published) + + require.ErrorIs(t, err, ErrDestinationExists) + require.FileExists(t, result.ManifestPath) +} + +func TestDumpValidatesBeforeCreatingDestination(t *testing.T) { + // Break caught: leaving an empty destination behind for an invalid request. + config := validRootDumpConfig(t) + config.ShardSize = 0 + + _, err := Dump(context.Background(), newDumpTestDatabase(map[string]*dumpTestGraph{"asset": {}}), config) + + require.ErrorIs(t, err, ErrInvalidConfig) + require.NoDirExists(t, config.Directory) +} + +func interruptedDumpAfterFirstNodeShard(t *testing.T, jsonlEnabled, parquetEnabled bool) (DumpConfig, *dumpTestDatabase) { + t.Helper() + return interruptedDumpWithGraphs(t, jsonlEnabled, parquetEnabled, []string{"asset"}) +} + +func interruptedDumpWithGraphs(t *testing.T, jsonlEnabled, parquetEnabled bool, graphs []string) (DumpConfig, *dumpTestDatabase) { + t.Helper() + values := map[string]*dumpTestGraph{ + "asset": { + nodes: dumpTestNodesWithKinds([]string{"User"}, 1, 2, 3), + relationships: dumpTestRelationships(10, 1, 3, "MEMBER_OF"), + }, + "other": {}, + } + database := newDumpTestDatabase(values) + config := validRootDumpConfig(t) + config.Graphs = append([]string(nil), graphs...) + config.EntityBatchSize = 3 + config.ShardSize = 1 + if !jsonlEnabled { + config.JSONL = nil + } + if !parquetEnabled { + config.Parquet = nil + } + return interruptedDumpFromConfig(t, config, database) +} + +func interruptedDumpFromConfig(t *testing.T, config DumpConfig, database *dumpTestDatabase) (DumpConfig, *dumpTestDatabase) { + t.Helper() + config.Directory = filepath.Join(t.TempDir(), "collection") + injected := errors.New("injected second node shard failure") + if config.JSONL != nil { + originalWrite := writeJSONLNodes + calls := 0 + writeJSONLNodes = func(tempPath, relativePath string, output jsonl.Config, nodes []entity.Node) (collection.JSONLArtifact, error) { + calls++ + if calls == 2 { + return collection.JSONLArtifact{}, injected + } + return originalWrite(tempPath, relativePath, output, nodes) + } + t.Cleanup(func() { writeJSONLNodes = originalWrite }) + _, err := Dump(context.Background(), database, config) + writeJSONLNodes = originalWrite + require.ErrorIs(t, err, injected) + } else { + originalWrite := writeParquetNodes + calls := 0 + writeParquetNodes = func(tempPath, relativePath string, output parquet.Config, nodes []entity.Node) (collection.ParquetArtifact, error) { + calls++ + if calls == 2 { + return collection.ParquetArtifact{}, injected + } + return originalWrite(tempPath, relativePath, output, nodes) + } + t.Cleanup(func() { writeParquetNodes = originalWrite }) + _, err := Dump(context.Background(), database, config) + writeParquetNodes = originalWrite + require.ErrorIs(t, err, injected) + } + require.FileExists(t, filepath.Join(config.Directory, checkpoint.FileName)) + require.NoFileExists(t, filepath.Join(config.Directory, collection.ManifestName)) + return config, database +} + +func dumpNodeShardIndices(graph collection.Graph) []int { + indices := make([]int, len(graph.NodeShards)) + for index, shard := range graph.NodeShards { + indices[index] = shard.Index + } + return indices +} + +func dumpRelationshipShardIndices(graph collection.Graph) []int { + indices := make([]int, len(graph.RelationshipShards)) + for index, shard := range graph.RelationshipShards { + indices[index] = shard.Index + } + return indices +} + +func mustReadFile(t *testing.T, path string) []byte { + t.Helper() + contents, err := os.ReadFile(path) + require.NoError(t, err) + return contents +} + +type stagedCancelContext struct { + context.Context + + cancel context.CancelFunc + mu sync.Mutex + remaining int + armed bool +} + +func newStagedCancelContext() *stagedCancelContext { + ctx, cancel := context.WithCancel(context.Background()) + return &stagedCancelContext{Context: ctx, cancel: cancel} +} + +func (s *stagedCancelContext) arm(calls int) { + s.mu.Lock() + defer s.mu.Unlock() + s.armed = true + s.remaining = calls +} + +func (s *stagedCancelContext) Err() error { + s.mu.Lock() + if s.armed && s.Context.Err() == nil { + s.remaining-- + if s.remaining == 0 { + s.cancel() + } + } + s.mu.Unlock() + return s.Context.Err() +} diff --git a/ret/dump_shard.go b/ret/dump_shard.go new file mode 100644 index 00000000..0dd23262 --- /dev/null +++ b/ret/dump_shard.go @@ -0,0 +1,207 @@ +package ret + +import ( + "errors" + "fmt" + "os" + "path/filepath" + "sync/atomic" + + "github.com/specterops/dawgs/ret/collection" + "github.com/specterops/dawgs/ret/entity" + "github.com/specterops/dawgs/ret/jsonl" + "github.com/specterops/dawgs/ret/parquet" + "github.com/specterops/dawgs/ret/scrub" +) + +var ( + writeJSONLNodes = writeJSONLNodeFile + writeParquetNodes = writeParquetNodeFile + writeJSONLRelationships = writeJSONLRelationshipFile + writeParquetRelationships = writeParquetRelationshipFile + shardRename = os.Rename + shardRemove = os.Remove + shardNonce atomic.Uint64 +) + +func writeNodeShard( + root, graphName string, + index int, + lastSourceID uint64, + counts scrub.ActionCounts, + nodes []entity.Node, + jsonlConfig *jsonl.Config, + parquetConfig *parquet.Config, +) (collection.NodeShard, error) { + jsonlPath := "" + if jsonlConfig != nil { + jsonlPath = collection.NodeJSONLPath(graphName, index, jsonlConfig.Codec) + } + parquetPath := "" + if parquetConfig != nil { + parquetPath = collection.NodeParquetPath(graphName, index) + } + + temporary, finals, err := shardPaths(root, jsonlPath, parquetPath) + if err != nil { + return collection.NodeShard{}, fmt.Errorf("%w: prepare node shard %d for graph %q: %w", ErrArtifactIntegrity, index, graphName, err) + } + cleanup := newShardCleanup(temporary) + if err := cleanup.ensureDestinationsUnused(finals); err != nil { + return collection.NodeShard{}, fmt.Errorf("%w: node shard %d for graph %q: %w", ErrDestinationExists, index, graphName, err) + } + + var jsonlArtifact *collection.JSONLArtifact + if jsonlConfig != nil { + artifact, err := writeJSONLNodes(temporary[0], jsonlPath, *jsonlConfig, nodes) + if err != nil { + return collection.NodeShard{}, cleanup.fail(fmt.Errorf("%w: write JSONL node shard %d for graph %q: %w", ErrArtifactIntegrity, index, graphName, err)) + } + if artifact.Count != int64(len(nodes)) { + return collection.NodeShard{}, cleanup.fail(fmt.Errorf("%w: JSONL node shard %d for graph %q count %d does not match %d", ErrArtifactIntegrity, index, graphName, artifact.Count, len(nodes))) + } + jsonlArtifact = &artifact + } + + var parquetArtifact *collection.ParquetArtifact + if parquetConfig != nil { + temporaryIndex := len(temporary) - 1 + artifact, err := writeParquetNodes(temporary[temporaryIndex], parquetPath, *parquetConfig, nodes) + if err != nil { + return collection.NodeShard{}, cleanup.fail(fmt.Errorf("%w: write Parquet node shard %d for graph %q: %w", ErrArtifactIntegrity, index, graphName, err)) + } + if artifact.Count != int64(len(nodes)) { + return collection.NodeShard{}, cleanup.fail(fmt.Errorf("%w: Parquet node shard %d for graph %q count %d does not match %d", ErrArtifactIntegrity, index, graphName, artifact.Count, len(nodes))) + } + parquetArtifact = &artifact + } + + if err := cleanup.publish(finals); err != nil { + return collection.NodeShard{}, cleanup.fail(fmt.Errorf("%w: publish node shard %d for graph %q: %w", ErrArtifactIntegrity, index, graphName, err)) + } + + return collection.NodeShard{Index: index, Count: int64(len(nodes)), LastSourceID: lastSourceID, ScrubCounts: counts, JSONL: jsonlArtifact, Parquet: parquetArtifact}, nil +} + +func writeRelationshipShard( + root, graphName string, + index int, + lastSourceID uint64, + counts scrub.ActionCounts, + relationships []entity.Relationship, + jsonlConfig *jsonl.Config, + parquetConfig *parquet.Config, +) (collection.RelationshipShard, error) { + jsonlPath := "" + if jsonlConfig != nil { + jsonlPath = collection.RelationshipJSONLPath(graphName, index, jsonlConfig.Codec) + } + parquetPath := "" + if parquetConfig != nil { + parquetPath = collection.RelationshipParquetPath(graphName, index) + } + + temporary, finals, err := shardPaths(root, jsonlPath, parquetPath) + if err != nil { + return collection.RelationshipShard{}, fmt.Errorf("%w: prepare relationship shard %d for graph %q: %w", ErrArtifactIntegrity, index, graphName, err) + } + cleanup := newShardCleanup(temporary) + if err := cleanup.ensureDestinationsUnused(finals); err != nil { + return collection.RelationshipShard{}, fmt.Errorf("%w: relationship shard %d for graph %q: %w", ErrDestinationExists, index, graphName, err) + } + + var jsonlArtifact *collection.JSONLArtifact + if jsonlConfig != nil { + artifact, err := writeJSONLRelationships(temporary[0], jsonlPath, *jsonlConfig, relationships) + if err != nil { + return collection.RelationshipShard{}, cleanup.fail(fmt.Errorf("%w: write JSONL relationship shard %d for graph %q: %w", ErrArtifactIntegrity, index, graphName, err)) + } + if artifact.Count != int64(len(relationships)) { + return collection.RelationshipShard{}, cleanup.fail(fmt.Errorf("%w: JSONL relationship shard %d for graph %q count %d does not match %d", ErrArtifactIntegrity, index, graphName, artifact.Count, len(relationships))) + } + jsonlArtifact = &artifact + } + + var parquetArtifact *collection.ParquetArtifact + if parquetConfig != nil { + temporaryIndex := len(temporary) - 1 + artifact, err := writeParquetRelationships(temporary[temporaryIndex], parquetPath, *parquetConfig, relationships) + if err != nil { + return collection.RelationshipShard{}, cleanup.fail(fmt.Errorf("%w: write Parquet relationship shard %d for graph %q: %w", ErrArtifactIntegrity, index, graphName, err)) + } + if artifact.Count != int64(len(relationships)) { + return collection.RelationshipShard{}, cleanup.fail(fmt.Errorf("%w: Parquet relationship shard %d for graph %q count %d does not match %d", ErrArtifactIntegrity, index, graphName, artifact.Count, len(relationships))) + } + parquetArtifact = &artifact + } + + if err := cleanup.publish(finals); err != nil { + return collection.RelationshipShard{}, cleanup.fail(fmt.Errorf("%w: publish relationship shard %d for graph %q: %w", ErrArtifactIntegrity, index, graphName, err)) + } + + return collection.RelationshipShard{Index: index, Count: int64(len(relationships)), LastSourceID: lastSourceID, ScrubCounts: counts, JSONL: jsonlArtifact, Parquet: parquetArtifact}, nil +} + +type shardCleanup struct { + temporary []string + published []string +} + +func newShardCleanup(temporary []string) *shardCleanup { + return &shardCleanup{temporary: temporary} +} + +func (s *shardCleanup) ensureDestinationsUnused(finals []string) error { + for _, final := range finals { + if _, err := os.Lstat(final); err == nil { + return fmt.Errorf("artifact already exists: %s", final) + } else if !errors.Is(err, os.ErrNotExist) { + return fmt.Errorf("inspect destination %s: %w", final, err) + } + } + return nil +} + +func (s *shardCleanup) publish(finals []string) error { + for index, temporary := range s.temporary { + if err := shardRename(temporary, finals[index]); err != nil { + return err + } + s.published = append(s.published, finals[index]) + } + return nil +} + +func (s *shardCleanup) fail(primary error) error { + errorsToJoin := []error{primary} + for _, filename := range append(append([]string(nil), s.temporary...), s.published...) { + if err := shardRemove(filename); err != nil && !errors.Is(err, os.ErrNotExist) { + errorsToJoin = append(errorsToJoin, fmt.Errorf("cleanup shard artifact %s: %w", filename, err)) + } + } + return errors.Join(errorsToJoin...) +} + +func shardPaths(root string, relativePaths ...string) ([]string, []string, error) { + finals := make([]string, 0, len(relativePaths)) + temporary := make([]string, 0, len(relativePaths)) + nonce := shardNonce.Add(1) + for _, relativePath := range relativePaths { + if relativePath == "" { + continue + } + final, err := collection.SafeJoin(root, relativePath) + if err != nil { + return nil, nil, err + } + if err := os.MkdirAll(filepath.Dir(final), 0o755); err != nil { + return nil, nil, fmt.Errorf("create artifact directory: %w", err) + } + finals = append(finals, final) + temporary = append(temporary, fmt.Sprintf("%s.tmp-%d", final, nonce)) + } + if len(finals) == 0 { + return nil, nil, fmt.Errorf("no shard output is enabled") + } + return temporary, finals, nil +} diff --git a/ret/dump_shard_test.go b/ret/dump_shard_test.go new file mode 100644 index 00000000..0534412a --- /dev/null +++ b/ret/dump_shard_test.go @@ -0,0 +1,198 @@ +package ret + +import ( + "errors" + "io/fs" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/specterops/dawgs/ret/checkpoint" + "github.com/specterops/dawgs/ret/collection" + "github.com/specterops/dawgs/ret/dawgs" + "github.com/specterops/dawgs/ret/entity" + "github.com/specterops/dawgs/ret/jsonl" + "github.com/specterops/dawgs/ret/parquet" + "github.com/specterops/dawgs/ret/scrub" + "github.com/stretchr/testify/require" +) + +func TestWriteNodeShardPublishesBothConcreteArtifacts(t *testing.T) { + // Break caught: publishing only one enabled artifact or returning metadata before both files exist. + root := t.TempDir() + shard, err := writeNodeShard( + root, "asset", 1, 42, scrub.ActionCounts{Redact: 1}, + []entity.Node{{SourceID: "42", Kinds: []string{"User"}, Properties: map[string]any{"name": "Ada"}}}, + pointerTo(jsonl.Config{Codec: jsonl.CodecZstd, Level: 3}), + pointerTo(parquet.Config{}), + ) + require.NoError(t, err) + require.NotNil(t, shard.JSONL) + require.NotNil(t, shard.Parquet) + require.EqualValues(t, 1, shard.Count) + require.Equal(t, 42, int(shard.LastSourceID)) + require.FileExists(t, filepath.Join(root, "graphs", "asset", "nodes", "000001.jsonl.zst")) + require.FileExists(t, filepath.Join(root, "graphs", "asset", "nodes", "000001.parquet")) +} + +func TestWriteNodeShardPublishesExactlyTheEnabledConcreteOutput(t *testing.T) { + // Break caught: emitting a disabled format or omitting the one enabled format. + for _, test := range []struct { + name string + jsonl *jsonl.Config + parquet *parquet.Config + wantJSONL bool + wantParquet bool + }{ + {name: "jsonl only", jsonl: pointerTo(jsonl.Config{Codec: jsonl.CodecNone}), wantJSONL: true}, + {name: "parquet only", parquet: pointerTo(parquet.Config{}), wantParquet: true}, + } { + t.Run(test.name, func(t *testing.T) { + root := t.TempDir() + shard, err := writeNodeShard(root, "asset", 1, 42, scrub.ActionCounts{}, []entity.Node{{SourceID: "42"}}, test.jsonl, test.parquet) + + require.NoError(t, err) + require.Equal(t, test.wantJSONL, shard.JSONL != nil) + require.Equal(t, test.wantParquet, shard.Parquet != nil) + require.EqualValues(t, 1, shard.Count) + require.Len(t, regularFiles(t, root), 1) + }) + } +} + +func TestWriteRelationshipShardPublishesBothConcreteArtifacts(t *testing.T) { + // Break caught: applying the logical dual-output commit only to node shards. + root := t.TempDir() + shard, err := writeRelationshipShard( + root, "asset", 1, 99, scrub.ActionCounts{Redact: 1}, + []entity.Relationship{{SourceID: "99", StartID: "1", EndID: "2", Kind: "MemberOf"}}, + pointerTo(jsonl.Config{Codec: jsonl.CodecNone}), + pointerTo(parquet.Config{}), + ) + require.NoError(t, err) + require.NotNil(t, shard.JSONL) + require.NotNil(t, shard.Parquet) + require.EqualValues(t, 1, shard.Count) + require.FileExists(t, filepath.Join(root, "graphs", "asset", "relationships", "000001.jsonl")) + require.FileExists(t, filepath.Join(root, "graphs", "asset", "relationships", "000001.parquet")) +} + +func TestWriteNodeShardCleansAllArtifactsWhenSecondWriterFails(t *testing.T) { + // Break caught: leaving the first concrete artifact behind when the second writer fails. + originalWriteParquetNodes := writeParquetNodes + writeParquetNodes = func(string, string, parquet.Config, []entity.Node) (collection.ParquetArtifact, error) { + return collection.ParquetArtifact{}, errors.New("injected Parquet failure") + } + t.Cleanup(func() { writeParquetNodes = originalWriteParquetNodes }) + + root := t.TempDir() + _, err := writeNodeShard(root, "asset", 1, 42, scrub.ActionCounts{}, []entity.Node{{SourceID: "42"}}, pointerTo(jsonl.Config{Codec: jsonl.CodecNone}), pointerTo(parquet.Config{})) + require.ErrorContains(t, err, "injected Parquet failure") + require.Empty(t, regularFiles(t, root)) +} + +func TestArtifactWriterDoesNotOverwriteExistingTemporaryPath(t *testing.T) { + temporary := filepath.Join(t.TempDir(), "artifact.tmp") + sentinel := []byte("existing temporary contents") + require.NoError(t, os.WriteFile(temporary, sentinel, 0o600)) + + _, err := writeJSONLNodeFile( + temporary, + "graphs/asset/nodes/000001.jsonl", + jsonl.Config{Codec: jsonl.CodecNone}, + []entity.Node{{SourceID: "1"}}, + ) + require.ErrorIs(t, err, fs.ErrExist) + contents, readErr := os.ReadFile(temporary) + require.NoError(t, readErr) + require.Equal(t, sentinel, contents) +} + +func TestWriteNodeShardRejectsMismatchedWriterCount(t *testing.T) { + // Break caught: publishing metadata whose artifact count disagrees with the logical shard count. + originalWriteJSONLNodes := writeJSONLNodes + writeJSONLNodes = func(tempPath, finalRelativePath string, config jsonl.Config, nodes []entity.Node) (collection.JSONLArtifact, error) { + artifact, err := originalWriteJSONLNodes(tempPath, finalRelativePath, config, nodes) + artifact.Count++ + return artifact, err + } + t.Cleanup(func() { writeJSONLNodes = originalWriteJSONLNodes }) + + root := t.TempDir() + _, err := writeNodeShard(root, "asset", 1, 42, scrub.ActionCounts{}, []entity.Node{{SourceID: "42"}}, pointerTo(jsonl.Config{Codec: jsonl.CodecNone}), pointerTo(parquet.Config{})) + require.ErrorIs(t, err, ErrArtifactIntegrity) + require.Empty(t, regularFiles(t, root)) +} + +func TestWriteNodeShardCleansPublishedArtifactWhenSecondRenameFails(t *testing.T) { + // Break caught: leaving the first published artifact behind when the second publication fails. + originalRename := shardRename + renames := 0 + shardRename = func(oldPath, newPath string) error { + renames++ + if renames == 2 { + return errors.New("injected second rename failure") + } + return originalRename(oldPath, newPath) + } + t.Cleanup(func() { shardRename = originalRename }) + + root := t.TempDir() + _, err := writeNodeShard(root, "asset", 1, 42, scrub.ActionCounts{}, []entity.Node{{SourceID: "42"}}, pointerTo(jsonl.Config{Codec: jsonl.CodecNone}), pointerTo(parquet.Config{})) + require.ErrorContains(t, err, "injected second rename failure") + require.Empty(t, regularFiles(t, root)) +} + +func TestShardStageIsRemovedByCheckpointOrphanCleanup(t *testing.T) { + // Break caught: a crash-leftover shard stage uses a name the checkpoint's strict cleanup inventory rejects. + root := t.TempDir() + state := checkpoint.State{ + Format: checkpoint.Format, + Identity: checkpoint.Identity{ + Graphs: []string{"asset"}, + EntityBatchSize: 1, + ShardSize: 1, + JSONLEnabled: true, + JSONLCodec: string(jsonl.CodecNone), + JSONLSchemaVersion: jsonl.SchemaVersion, + ParquetSchemaVersion: parquet.SchemaVersion, + }, + Graphs: []checkpoint.GraphState{{ + Name: "asset", + Snapshot: dawgs.Snapshot{NodeCount: 1}, + Phase: checkpoint.PhaseNodes, + }}, + } + store := checkpoint.Store{Root: root} + require.NoError(t, store.Save(state)) + + finalRelativePath := "graphs/asset/nodes/000001.jsonl" + temporary, _, err := shardPaths(root, finalRelativePath) + require.NoError(t, err) + require.Len(t, temporary, 1) + _, err = writeJSONLNodes(temporary[0], finalRelativePath, jsonl.Config{Codec: jsonl.CodecNone}, []entity.Node{{SourceID: "1"}}) + require.NoError(t, err) + + stageRelativePath, err := filepath.Rel(root, temporary[0]) + require.NoError(t, err) + require.True(t, strings.HasPrefix(filepath.ToSlash(stageRelativePath), finalRelativePath+".tmp-")) + + require.NoError(t, store.CleanupOrphans(state)) + require.NoFileExists(t, temporary[0]) +} + +func regularFiles(t *testing.T, root string) []string { + t.Helper() + var paths []string + require.NoError(t, filepath.WalkDir(root, func(path string, entry fs.DirEntry, err error) error { + if err != nil { + return err + } + if entry.Type().IsRegular() { + paths = append(paths, path) + } + return nil + })) + return paths +} diff --git a/ret/dump_test.go b/ret/dump_test.go new file mode 100644 index 00000000..21f32ef0 --- /dev/null +++ b/ret/dump_test.go @@ -0,0 +1,816 @@ +package ret + +import ( + "context" + "errors" + "fmt" + "os" + "path/filepath" + "testing" + + cypherModel "github.com/specterops/dawgs/cypher/models/cypher" + "github.com/specterops/dawgs/graph" + "github.com/specterops/dawgs/ret/checkpoint" + "github.com/specterops/dawgs/ret/collection" + "github.com/specterops/dawgs/ret/dawgs" + "github.com/specterops/dawgs/ret/entity" + "github.com/specterops/dawgs/ret/jsonl" + "github.com/specterops/dawgs/ret/observe" + "github.com/specterops/dawgs/ret/parquet" + "github.com/specterops/dawgs/ret/scrub" + "github.com/stretchr/testify/require" +) + +func TestDumpSplitsDatabaseBatchAcrossLogicalShards(t *testing.T) { + // Break caught: treating each database batch as one shard instead of enforcing + // exact logical shard boundaries within that batch. + database := newDumpTestDatabase(map[string]*dumpTestGraph{ + "asset": {nodes: dumpTestNodes(1, 2, 3, 4, 5)}, + }) + config := validRootDumpConfig(t) + config.EntityBatchSize = 5 + config.ShardSize = 2 + + result, err := Dump(context.Background(), database, config) + + require.NoError(t, err) + manifest := readDumpManifest(t, result.ManifestPath) + require.Equal(t, []int64{2, 2, 1}, dumpNodeShardCounts(manifest.Graphs[0])) + require.Equal(t, []uint64{2, 4, 5}, dumpNodeShardCursors(manifest.Graphs[0])) + + var sourceIDs []string + for _, shard := range manifest.Graphs[0].NodeShards { + nodes, err := readJSONLNodesForTest(config.Directory, *shard.JSONL) + require.NoError(t, err) + for _, node := range nodes { + sourceIDs = append(sourceIDs, node.SourceID) + } + } + require.Equal(t, []string{"1", "2", "3", "4", "5"}, sourceIDs) +} + +func TestDumpWritesEmptyGraphWithoutShards(t *testing.T) { + // Break caught: inventing an empty artifact or omitting a valid empty graph + // from the final collection. + database := newDumpTestDatabase(map[string]*dumpTestGraph{"asset": {}}) + config := validRootDumpConfig(t) + + result, err := Dump(context.Background(), database, config) + + require.NoError(t, err) + manifest := readDumpManifest(t, result.ManifestPath) + require.Len(t, manifest.Graphs, 1) + require.Empty(t, manifest.Graphs[0].NodeShards) + require.Empty(t, manifest.Graphs[0].RelationshipShards) + require.Zero(t, manifest.Graphs[0].NodeCount) + require.Zero(t, manifest.Graphs[0].RelationshipCount) + require.Equal(t, 1, result.GraphCount) + require.Zero(t, result.NodeCount) + require.Zero(t, result.RelationshipCount) +} + +func TestDumpWritesOnePartialShardForEveryEnabledOutputMode(t *testing.T) { + // Break caught: coupling JSONL and Parquet publication, or dropping a final + // partial logical shard when the phase ends below ShardSize. + for _, test := range []struct { + name string + jsonl bool + parquet bool + wantJSONL bool + wantParquet bool + }{ + {name: "JSONL only", jsonl: true, wantJSONL: true}, + {name: "Parquet only", parquet: true, wantParquet: true}, + {name: "dual output", jsonl: true, parquet: true, wantJSONL: true, wantParquet: true}, + } { + t.Run(test.name, func(t *testing.T) { + database := newDumpTestDatabase(map[string]*dumpTestGraph{ + "asset": {nodes: dumpTestNodes(7)}, + }) + config := validRootDumpConfig(t) + config.ShardSize = 3 + if !test.jsonl { + config.JSONL = nil + } + if !test.parquet { + config.Parquet = nil + } + + result, err := Dump(context.Background(), database, config) + + require.NoError(t, err) + manifest := readDumpManifest(t, result.ManifestPath) + require.Len(t, manifest.Graphs[0].NodeShards, 1) + shard := manifest.Graphs[0].NodeShards[0] + require.EqualValues(t, 1, shard.Count) + require.Equal(t, test.wantJSONL, shard.JSONL != nil) + require.Equal(t, test.wantParquet, shard.Parquet != nil) + }) + } +} + +func TestDumpPreservesCallerGraphAndEntityPhaseOrder(t *testing.T) { + // Break caught: sorting graphs or interleaving relationship reads before all + // nodes for a graph have been processed. + database := newDumpTestDatabase(map[string]*dumpTestGraph{ + "second": { + nodes: dumpTestNodesWithKinds([]string{"User", "Admin", "User"}, 1, 2), + relationships: dumpTestRelationships(10, 1, 2, "MEMBER_OF"), + }, + "first": {nodes: dumpTestNodes(20)}, + }) + config := validRootDumpConfig(t) + config.Graphs = []string{"second", "first"} + config.EntityBatchSize = 10 + config.ShardSize = 10 + + result, err := Dump(context.Background(), database, config) + + require.NoError(t, err) + manifest := readDumpManifest(t, result.ManifestPath) + require.Equal(t, []string{"second", "first"}, []string{manifest.Graphs[0].Name, manifest.Graphs[1].Name}) + require.Equal(t, []string{"second:nodes", "second:nodes", "second:relationships", "second:relationships", "first:nodes", "first:nodes", "first:relationships"}, database.fetches) + require.Equal(t, []string{"User", "Admin", "MEMBER_OF"}, manifest.Graphs[0].KindCatalog) + require.EqualValues(t, 2, manifest.Graphs[0].Metrics.NodeCount) + require.EqualValues(t, 1, manifest.Graphs[0].Metrics.RelationshipCount) + require.Equal(t, 2, result.GraphCount) + require.EqualValues(t, 3, result.NodeCount) + require.EqualValues(t, 1, result.RelationshipCount) +} + +func TestDumpRejectsExistingDestinationWithoutTouchingIt(t *testing.T) { + // Break caught: accepting a pre-existing destination and mixing a fresh dump + // with files the caller already owns. + root := t.TempDir() + sentinel := filepath.Join(root, "sentinel") + require.NoError(t, os.WriteFile(sentinel, []byte("keep"), 0o600)) + config := validRootDumpConfig(t) + config.Directory = root + + _, err := Dump(context.Background(), newDumpTestDatabase(map[string]*dumpTestGraph{"asset": {}}), config) + + require.ErrorIs(t, err, ErrDestinationExists) + require.FileExists(t, sentinel) + require.NoFileExists(t, filepath.Join(root, checkpointFileNameForTest)) +} + +func TestDumpObserverEventsAreOrderedAndShardEventsFollowCheckpoint(t *testing.T) { + // Break caught: emitting lifecycle events out of order, using a different + // entity vocabulary from artifact events, or announcing a shard before its + // checkpoint is durable. + database := newDumpTestDatabase(map[string]*dumpTestGraph{ + "asset": { + nodes: dumpTestNodes(1, 2), + relationships: dumpTestRelationships(10, 1, 2, "LINKED"), + }, + }) + config := validRootDumpConfig(t) + config.EntityBatchSize = 2 + config.ShardSize = 2 + var names []string + config.Observer = observe.ObserverFunc(func(_ context.Context, event observe.Event) { + names = append(names, dumpTestEventName(event)) + shard, ok := event.(observe.ShardCommitted) + if !ok { + return + } + state, exists, err := (checkpoint.Store{Root: config.Directory}).Load() + require.NoError(t, err) + require.True(t, exists) + if shard.EntityType == "node" { + require.Len(t, state.Graphs[0].NodeShards, 1) + } else { + require.Len(t, state.Graphs[0].RelationshipShards, 1) + } + }) + + _, err := Dump(context.Background(), database, config) + + require.NoError(t, err) + require.Equal(t, []string{ + "operation_started", + "graph_started:asset", + "phase_started:nodes:0", + "phase_progress:nodes:1", + "phase_progress:nodes:2", + "shard:node:1", + "phase_completed:nodes", + "phase_started:relationships:0", + "phase_progress:relationships:1", + "shard:relationship:1", + "phase_completed:relationships", + "graph_completed:asset", + "operation_completed", + }, names) +} + +func TestDumpObserverTerminalEventContainsFailureAndIsLast(t *testing.T) { + // Break caught: dropping the terminal cause from observation or emitting + // success/graph events after a writer has failed. + injected := errors.New("injected observer-order writer failure") + originalWrite := writeJSONLNodes + writeJSONLNodes = func(string, string, jsonl.Config, []entity.Node) (collection.JSONLArtifact, error) { + return collection.JSONLArtifact{}, injected + } + t.Cleanup(func() { writeJSONLNodes = originalWrite }) + + database := newDumpTestDatabase(map[string]*dumpTestGraph{"asset": {nodes: dumpTestNodes(1)}}) + config := validRootDumpConfig(t) + var events []observe.Event + config.Observer = observe.ObserverFunc(func(_ context.Context, event observe.Event) { + events = append(events, event) + }) + + _, err := Dump(context.Background(), database, config) + + require.ErrorIs(t, err, injected) + require.NotEmpty(t, events) + completed, ok := events[len(events)-1].(observe.OperationCompleted) + require.True(t, ok) + require.ErrorIs(t, completed.Err, injected) + for _, event := range events[:len(events)-1] { + _, graphCompleted := event.(observe.GraphCompleted) + require.False(t, graphCompleted) + } +} + +func TestDumpObserverCancellationStopsLaterNonTerminalLifecycleEvents(t *testing.T) { + // Break caught: synchronously observing cancellation but emitting the next + // phase or graph lifecycle event before checking the context again. + for _, test := range []struct { + name string + cancelOn func(observe.Event) bool + wantPhase checkpoint.Phase + }{ + { + name: "after graph started", + cancelOn: func(event observe.Event) bool { + _, ok := event.(observe.GraphStarted) + return ok + }, + wantPhase: checkpoint.PhaseNodes, + }, + { + name: "between node and relationship phases", + cancelOn: func(event observe.Event) bool { + value, ok := event.(observe.PhaseCompleted) + return ok && value.Phase == string(checkpoint.PhaseNodes) + }, + wantPhase: checkpoint.PhaseRelationships, + }, + { + name: "before graph completed", + cancelOn: func(event observe.Event) bool { + value, ok := event.(observe.PhaseCompleted) + return ok && value.Phase == string(checkpoint.PhaseRelationships) + }, + wantPhase: checkpoint.PhaseComplete, + }, + } { + t.Run(test.name, func(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + database := newDumpTestDatabase(map[string]*dumpTestGraph{"asset": {}}) + config := validRootDumpConfig(t) + var events []observe.Event + cancelEvent := -1 + config.Observer = observe.ObserverFunc(func(_ context.Context, event observe.Event) { + events = append(events, event) + if cancelEvent < 0 && test.cancelOn(event) { + cancelEvent = len(events) - 1 + cancel() + } + }) + + _, err := Dump(ctx, database, config) + + require.ErrorIs(t, err, context.Canceled) + require.GreaterOrEqual(t, cancelEvent, 0) + require.Len(t, events[cancelEvent+1:], 1) + completed, ok := events[len(events)-1].(observe.OperationCompleted) + require.True(t, ok) + require.ErrorIs(t, completed.Err, context.Canceled) + state, exists, loadErr := (checkpoint.Store{Root: config.Directory}).Load() + require.NoError(t, loadErr) + require.True(t, exists) + require.Equal(t, test.wantPhase, state.Graphs[0].Phase) + require.NoFileExists(t, filepath.Join(config.Directory, collection.ManifestName)) + }) + } +} + +func TestDumpCancellationDuringScanningDoesNotPublishOrCheckpointProgress(t *testing.T) { + // Break caught: processing a batch returned concurrently with cancellation + // and publishing artifacts or a later cursor before returning context.Canceled. + ctx, cancel := context.WithCancel(context.Background()) + database := newDumpTestDatabase(map[string]*dumpTestGraph{"asset": {nodes: dumpTestNodes(1, 2)}}) + config := validRootDumpConfig(t) + var initialCheckpoint []byte + database.onNodeFetch = func(string) { + initialCheckpoint = mustReadFile(t, filepath.Join(config.Directory, checkpoint.FileName)) + cancel() + } + var events []observe.Event + config.Observer = observe.ObserverFunc(func(_ context.Context, event observe.Event) { + events = append(events, event) + }) + + _, err := Dump(ctx, database, config) + + require.ErrorIs(t, err, context.Canceled) + require.Equal(t, initialCheckpoint, mustReadFile(t, filepath.Join(config.Directory, checkpoint.FileName))) + require.Len(t, regularFiles(t, config.Directory), 1) + require.NoFileExists(t, filepath.Join(config.Directory, collection.ManifestName)) + completed, ok := events[len(events)-1].(observe.OperationCompleted) + require.True(t, ok) + require.ErrorIs(t, completed.Err, context.Canceled) + for _, event := range events { + _, committed := event.(observe.ShardCommitted) + require.False(t, committed) + } +} + +func TestDumpCancellationWithEmptyTerminalBatchDoesNotAdvancePhaseCheckpoint(t *testing.T) { + // Break caught: overlooking cancellation returned alongside the terminal + // empty scan batch and checkpointing a phase transition afterward. + ctx, cancel := context.WithCancel(context.Background()) + database := newDumpTestDatabase(map[string]*dumpTestGraph{"asset": {nodes: dumpTestNodes(1, 2)}}) + config := validRootDumpConfig(t) + config.EntityBatchSize = 2 + config.ShardSize = 2 + fetches := 0 + var checkpointAtCancellation []byte + database.onNodeFetch = func(string) { + fetches++ + if fetches == 2 { + checkpointAtCancellation = mustReadFile(t, filepath.Join(config.Directory, checkpoint.FileName)) + cancel() + } + } + + _, err := Dump(ctx, database, config) + + require.ErrorIs(t, err, context.Canceled) + require.Equal(t, checkpointAtCancellation, mustReadFile(t, filepath.Join(config.Directory, checkpoint.FileName))) + state, exists, loadErr := (checkpoint.Store{Root: config.Directory}).Load() + require.NoError(t, loadErr) + require.True(t, exists) + require.Equal(t, checkpoint.PhaseNodes, state.Graphs[0].Phase) + require.Len(t, state.Graphs[0].NodeShards, 1) + require.NoFileExists(t, filepath.Join(config.Directory, collection.ManifestName)) +} + +func TestDumpCancellationDuringInitialSnapshotDoesNotPublishCheckpoint(t *testing.T) { + // Break caught: treating a successful snapshot as proof the context remains + // active and publishing the initial checkpoint after cancellation. + ctx, cancel := context.WithCancel(context.Background()) + database := newDumpTestDatabase(map[string]*dumpTestGraph{"asset": {}}) + database.onRelationshipCount = func(context.Context, string) { + cancel() + } + config := validRootDumpConfig(t) + + _, err := Dump(ctx, database, config) + + require.ErrorIs(t, err, context.Canceled) + require.NoFileExists(t, filepath.Join(config.Directory, checkpoint.FileName)) + require.NoFileExists(t, filepath.Join(config.Directory, collection.ManifestName)) +} + +func TestDumpCancellationDuringFinalSnapshotDoesNotPublishManifest(t *testing.T) { + // Break caught: accepting the final recount result after cancellation and + // replacing the resumable checkpoint with a manifest. + ctx, cancel := context.WithCancel(context.Background()) + database := newDumpTestDatabase(map[string]*dumpTestGraph{"asset": {nodes: dumpTestNodes(1)}}) + counts := 0 + database.onRelationshipCount = func(context.Context, string) { + counts++ + if counts == 2 { + cancel() + } + } + config := validRootDumpConfig(t) + + _, err := Dump(ctx, database, config) + + require.ErrorIs(t, err, context.Canceled) + require.FileExists(t, filepath.Join(config.Directory, checkpoint.FileName)) + require.NoFileExists(t, filepath.Join(config.Directory, collection.ManifestName)) +} + +func TestDumpCancellationBeforeManifestPublicationDoesNotPublishManifest(t *testing.T) { + // Break caught: removing the context gate immediately before manifest + // publication after the final recount's post-snapshot gate has succeeded. + ctx := newStagedCancelContext() + database := newDumpTestDatabase(map[string]*dumpTestGraph{"asset": {}}) + snapshots := 0 + database.onRelationshipCount = func(context.Context, string) { + snapshots++ + if snapshots == 2 { + // The recount post-snapshot gate observes an active context; the + // immediately following pre-publication gate observes cancellation. + ctx.arm(2) + } + } + config := validRootDumpConfig(t) + + _, err := Dump(ctx, database, config) + + require.ErrorIs(t, err, context.Canceled) + require.FileExists(t, filepath.Join(config.Directory, checkpoint.FileName)) + require.NoFileExists(t, filepath.Join(config.Directory, collection.ManifestName)) +} + +func TestDumpCancellationFromFinalShardObserverDoesNotAdvancePhase(t *testing.T) { + // Break caught: allowing synchronous observer cancellation after the final + // partial shard commit to publish a later phase checkpoint or phase event. + for _, test := range []struct { + name string + cancelOn string + value *dumpTestGraph + wantPhase checkpoint.Phase + }{ + { + name: "node shard", + cancelOn: "node", + value: &dumpTestGraph{nodes: dumpTestNodes(1)}, + wantPhase: checkpoint.PhaseNodes, + }, + { + name: "relationship shard", + cancelOn: "relationship", + value: &dumpTestGraph{ + nodes: dumpTestNodes(1, 2), + relationships: dumpTestRelationships(10, 1, 2, "LINKED"), + }, + wantPhase: checkpoint.PhaseRelationships, + }, + } { + t.Run(test.name, func(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + database := newDumpTestDatabase(map[string]*dumpTestGraph{"asset": test.value}) + config := validRootDumpConfig(t) + config.EntityBatchSize = 2 + config.ShardSize = 3 + var events []observe.Event + cancelEvent := -1 + config.Observer = observe.ObserverFunc(func(_ context.Context, event observe.Event) { + events = append(events, event) + if shard, ok := event.(observe.ShardCommitted); ok && shard.EntityType == test.cancelOn { + cancelEvent = len(events) - 1 + cancel() + } + }) + + _, err := Dump(ctx, database, config) + + require.ErrorIs(t, err, context.Canceled) + require.GreaterOrEqual(t, cancelEvent, 0) + state, exists, loadErr := (checkpoint.Store{Root: config.Directory}).Load() + require.NoError(t, loadErr) + require.True(t, exists) + require.Equal(t, test.wantPhase, state.Graphs[0].Phase) + require.IsType(t, observe.OperationCompleted{}, events[len(events)-1]) + require.Len(t, events[cancelEvent+1:], 1) + require.NoFileExists(t, filepath.Join(config.Directory, collection.ManifestName)) + }) + } +} + +func TestDumpScrubsConcreteArtifactsAndRecordsPerShardActions(t *testing.T) { + // Break caught: bypassing root scrubbing, writing pre-scrub properties, or + // dropping node/relationship action counts from logical shard metadata. + database := newDumpTestDatabase(map[string]*dumpTestGraph{ + "asset": { + nodes: []*graph.Node{ + graph.NewNode(1, graph.AsProperties(map[string]any{"password": "one"}), graph.StringKind("Entity")), + graph.NewNode(2, graph.AsProperties(map[string]any{"password": "two"}), graph.StringKind("Entity")), + }, + relationships: []*graph.Relationship{ + graph.NewRelationship( + 10, + 1, + 2, + graph.AsProperties(map[string]any{"password": "edge-secret"}), + graph.StringKind("LINKED"), + ), + }, + }, + }) + config := validRootDumpConfig(t) + config.Parquet = nil + config.EntityBatchSize = 2 + config.ShardSize = 2 + + result, err := Dump(context.Background(), database, config) + + require.NoError(t, err) + manifest := readDumpManifest(t, result.ManifestPath) + require.Equal(t, scrub.ActionCounts{Redact: 2}, manifest.Graphs[0].NodeShards[0].ScrubCounts) + require.Equal(t, scrub.ActionCounts{Redact: 1}, manifest.Graphs[0].RelationshipShards[0].ScrubCounts) + nodes, err := readJSONLNodesForTest( + config.Directory, + *manifest.Graphs[0].NodeShards[0].JSONL, + ) + require.NoError(t, err) + nodePasswords := make([]any, 0, len(nodes)) + for _, node := range nodes { + nodePasswords = append(nodePasswords, node.Properties["password"]) + } + require.Equal(t, []any{"[REDACTED]", "[REDACTED]"}, nodePasswords) + relationships, err := readJSONLRelationshipsForTest( + config.Directory, + *manifest.Graphs[0].RelationshipShards[0].JSONL, + ) + require.NoError(t, err) + require.Len(t, relationships, 1) + require.Equal(t, "[REDACTED]", relationships[0].Properties["password"]) +} + +func TestDumpWithoutScrubConfigPreservesPropertiesAndDisablesScrubMetadata(t *testing.T) { + // Break caught: compiling or applying a zero-value scrub policy when the + // library caller disables scrubbing with a nil configuration. + database := newDumpTestDatabase(map[string]*dumpTestGraph{ + "asset": { + nodes: []*graph.Node{ + graph.NewNode(1, graph.AsProperties(map[string]any{"password": "secret"}), graph.StringKind("Entity")), + }, + }, + }) + config := validRootDumpConfig(t) + config.Parquet = nil + config.Scrub = nil + + result, err := Dump(context.Background(), database, config) + + require.NoError(t, err) + manifest := readDumpManifest(t, result.ManifestPath) + require.False(t, manifest.Scrub.Enabled) + require.Empty(t, manifest.Scrub.RulesFingerprint) + require.Empty(t, manifest.Scrub.SaltFingerprint) + require.True(t, manifest.Graphs[0].NodeShards[0].ScrubCounts.IsZero()) + nodes, err := readJSONLNodesForTest( + config.Directory, + *manifest.Graphs[0].NodeShards[0].JSONL, + ) + require.NoError(t, err) + require.Len(t, nodes, 1) + require.Equal(t, "secret", nodes[0].Properties["password"]) +} + +const checkpointFileNameForTest = ".ret-checkpoint.json" + +func validRootDumpConfig(t *testing.T) DumpConfig { + t.Helper() + return DumpConfig{ + Directory: filepath.Join(t.TempDir(), "collection"), + Graphs: []string{"asset"}, + EntityBatchSize: 2, + ShardSize: 2, + JSONL: pointerTo(jsonl.Config{Codec: jsonl.CodecNone}), + Parquet: pointerTo(parquet.Config{}), + Scrub: pointerTo(scrub.DefaultConfig()), + } +} + +func readDumpManifest(t *testing.T, manifestPath string) collection.Manifest { + t.Helper() + require.Equal(t, collection.ManifestName, filepath.Base(manifestPath)) + manifest, err := collection.Read(filepath.Dir(manifestPath)) + require.NoError(t, err) + return manifest +} + +func dumpNodeShardCounts(graph collection.Graph) []int64 { + counts := make([]int64, len(graph.NodeShards)) + for index, shard := range graph.NodeShards { + counts[index] = shard.Count + } + return counts +} + +func dumpNodeShardCursors(graph collection.Graph) []uint64 { + cursors := make([]uint64, len(graph.NodeShards)) + for index, shard := range graph.NodeShards { + cursors[index] = shard.LastSourceID + } + return cursors +} + +func dumpTestEventName(event observe.Event) string { + switch value := event.(type) { + case observe.OperationStarted: + return "operation_started" + case observe.OperationCompleted: + return "operation_completed" + case observe.GraphStarted: + return "graph_started:" + value.Graph + case observe.GraphCompleted: + return "graph_completed:" + value.Graph + case observe.PhaseStarted: + return fmt.Sprintf("phase_started:%s:%d", value.Phase, value.Completed) + case observe.PhaseProgress: + return fmt.Sprintf("phase_progress:%s:%d", value.Phase, value.Completed) + case observe.PhaseCompleted: + return "phase_completed:" + value.Phase + case observe.ShardCommitted: + return fmt.Sprintf("shard:%s:%d", value.EntityType, value.Index) + default: + return fmt.Sprintf("unexpected:%T", event) + } +} + +type dumpTestGraph struct { + nodes []*graph.Node + relationships []*graph.Relationship + snapshots []dawgs.Snapshot + countRound int +} + +type dumpTestDatabase struct { + graph.Database + + graphs map[string]*dumpTestGraph + fetches []string + onNodeCount func(context.Context, string) + onRelationshipCount func(context.Context, string) + onNodeFetch func(string) + onRelationshipFetch func(string) +} + +func newDumpTestDatabase(graphs map[string]*dumpTestGraph) *dumpTestDatabase { + for _, value := range graphs { + if len(value.snapshots) == 0 { + value.snapshots = []dawgs.Snapshot{{ + NodeCount: int64(len(value.nodes)), + RelationshipCount: int64(len(value.relationships)), + }} + } + } + return &dumpTestDatabase{graphs: graphs} +} + +func (s *dumpTestDatabase) ReadTransaction(ctx context.Context, delegate graph.TransactionDelegate, _ ...graph.TransactionOption) error { + if err := ctx.Err(); err != nil { + return err + } + return delegate(&dumpTestTransaction{database: s, ctx: ctx}) +} + +type dumpTestTransaction struct { + graph.Transaction + database *dumpTestDatabase + ctx context.Context + graphName string +} + +func (s *dumpTestTransaction) WithGraph(target graph.Graph) graph.Transaction { + s.graphName = target.Name + return s +} + +func (s *dumpTestTransaction) Nodes() graph.NodeQuery { + return &dumpTestNodeQuery{database: s.database, ctx: s.ctx, graphName: s.graphName} +} + +func (s *dumpTestTransaction) Relationships() graph.RelationshipQuery { + return &dumpTestRelationshipQuery{database: s.database, ctx: s.ctx, graphName: s.graphName} +} + +type dumpTestNodeQuery struct { + graph.NodeQuery + database *dumpTestDatabase + ctx context.Context + graphName string + afterID graph.ID + limit int +} + +func (s *dumpTestNodeQuery) OrderBy(...graph.Criteria) graph.NodeQuery { return s } +func (s *dumpTestNodeQuery) Filter(criteria graph.Criteria) graph.NodeQuery { + s.afterID = dumpTestAfterID(criteria) + return s +} +func (s *dumpTestNodeQuery) Limit(limit int) graph.NodeQuery { + s.limit = limit + return s +} +func (s *dumpTestNodeQuery) Count() (int64, error) { + if s.database.onNodeCount != nil { + s.database.onNodeCount(s.ctx, s.graphName) + } + value := s.database.graphs[s.graphName] + return value.snapshots[min(value.countRound, len(value.snapshots)-1)].NodeCount, nil +} +func (s *dumpTestNodeQuery) Fetch(delegate func(graph.Cursor[*graph.Node]) error, _ ...graph.Criteria) error { + s.database.fetches = append(s.database.fetches, s.graphName+":nodes") + if s.database.onNodeFetch != nil { + s.database.onNodeFetch(s.graphName) + } + return delegate(newDumpTestCursor(dumpTestNodesAfter(s.database.graphs[s.graphName].nodes, s.afterID, s.limit))) +} + +type dumpTestRelationshipQuery struct { + graph.RelationshipQuery + database *dumpTestDatabase + ctx context.Context + graphName string + afterID graph.ID + limit int +} + +func (s *dumpTestRelationshipQuery) OrderBy(...graph.Criteria) graph.RelationshipQuery { return s } +func (s *dumpTestRelationshipQuery) Filter(criteria graph.Criteria) graph.RelationshipQuery { + s.afterID = dumpTestAfterID(criteria) + return s +} +func (s *dumpTestRelationshipQuery) Limit(limit int) graph.RelationshipQuery { + s.limit = limit + return s +} +func (s *dumpTestRelationshipQuery) Count() (int64, error) { + if s.database.onRelationshipCount != nil { + s.database.onRelationshipCount(s.ctx, s.graphName) + } + value := s.database.graphs[s.graphName] + round := min(value.countRound, len(value.snapshots)-1) + value.countRound++ + return value.snapshots[round].RelationshipCount, nil +} +func (s *dumpTestRelationshipQuery) Fetch(delegate func(graph.Cursor[*graph.Relationship]) error) error { + s.database.fetches = append(s.database.fetches, s.graphName+":relationships") + if s.database.onRelationshipFetch != nil { + s.database.onRelationshipFetch(s.graphName) + } + return delegate(newDumpTestCursor(dumpTestRelationshipsAfter(s.database.graphs[s.graphName].relationships, s.afterID, s.limit))) +} + +type dumpTestCursor[T any] struct { + values chan T +} + +func newDumpTestCursor[T any](values []T) *dumpTestCursor[T] { + channel := make(chan T, len(values)) + for _, value := range values { + channel <- value + } + close(channel) + return &dumpTestCursor[T]{values: channel} +} + +func (s *dumpTestCursor[T]) Error() error { return nil } +func (s *dumpTestCursor[T]) Close() {} +func (s *dumpTestCursor[T]) Chan() chan T { return s.values } + +func dumpTestAfterID(criteria graph.Criteria) graph.ID { + comparison, ok := criteria.(*cypherModel.Comparison) + if !ok || len(comparison.Partials) != 1 { + panic(fmt.Sprintf("unexpected cursor criteria %T", criteria)) + } + parameter, ok := comparison.Partials[0].Right.(*cypherModel.Parameter) + if !ok { + panic(fmt.Sprintf("unexpected cursor parameter %T", comparison.Partials[0].Right)) + } + return parameter.Value.(graph.ID) +} + +func dumpTestNodesAfter(values []*graph.Node, after graph.ID, limit int) []*graph.Node { + result := make([]*graph.Node, 0, limit) + for _, value := range values { + if value.ID > after && len(result) < limit { + result = append(result, value) + } + } + return result +} + +func dumpTestRelationshipsAfter(values []*graph.Relationship, after graph.ID, limit int) []*graph.Relationship { + result := make([]*graph.Relationship, 0, limit) + for _, value := range values { + if value.ID > after && len(result) < limit { + result = append(result, value) + } + } + return result +} + +func dumpTestNodes(ids ...uint64) []*graph.Node { + return dumpTestNodesWithKinds([]string{"Entity"}, ids...) +} + +func dumpTestNodesWithKinds(kinds []string, ids ...uint64) []*graph.Node { + graphKinds := make([]graph.Kind, len(kinds)) + for index, kind := range kinds { + graphKinds[index] = graph.StringKind(kind) + } + nodes := make([]*graph.Node, len(ids)) + for index, id := range ids { + nodes[index] = graph.NewNode(graph.ID(id), graph.AsProperties(map[string]any{"id": id}), graphKinds...) + } + return nodes +} + +func dumpTestRelationships(id, startID, endID uint64, kind string) []*graph.Relationship { + return []*graph.Relationship{ + graph.NewRelationship(graph.ID(id), graph.ID(startID), graph.ID(endID), graph.NewProperties(), graph.StringKind(kind)), + } +} diff --git a/ret/entity/entity.go b/ret/entity/entity.go new file mode 100644 index 00000000..183aed92 --- /dev/null +++ b/ret/entity/entity.go @@ -0,0 +1,62 @@ +// Package entity defines canonical graph-export entity values. +package entity + +import "errors" + +type Entity interface { + Node | Relationship + Validate() error +} + +type Node struct { + SourceID string + Kinds []string + Properties map[string]any +} + +type Relationship struct { + SourceID string + StartID string + EndID string + Kind string + Properties map[string]any +} + +func CloneKinds(kinds []string) []string { + return append([]string(nil), kinds...) +} + +func CloneProperties(properties map[string]any) map[string]any { + if properties == nil { + return nil + } + + cloned := make(map[string]any, len(properties)) + for key, value := range properties { + cloned[key] = value + } + + return cloned +} + +func (s Node) Validate() error { + if s.SourceID == "" { + return errors.New("node source ID is required") + } + + return nil +} + +func (s Relationship) Validate() error { + if s.StartID == "" { + return errors.New("relationship start ID is required") + } + if s.EndID == "" { + return errors.New("relationship end ID is required") + } + if s.Kind == "" { + return errors.New("relationship kind is required") + } + + return nil +} diff --git a/ret/entity/entity_test.go b/ret/entity/entity_test.go new file mode 100644 index 00000000..5fcb8216 --- /dev/null +++ b/ret/entity/entity_test.go @@ -0,0 +1,44 @@ +package entity_test + +import ( + "testing" + + "github.com/specterops/dawgs/ret/entity" + "github.com/stretchr/testify/require" +) + +func TestCloneKindsPreservesOrderAndDuplicates(t *testing.T) { + input := []string{"User", "Admin", "User"} + + got := entity.CloneKinds(input) + + require.Equal(t, input, got) + got[0] = "Changed" + require.Equal(t, "User", input[0]) +} + +func TestClonePropertiesIsShallow(t *testing.T) { + nested := map[string]any{"secret": "shared"} + input := map[string]any{"name": "Ada", "nested": nested} + + got := entity.CloneProperties(input) + + got["name"] = "Grace" + got["nested"].(map[string]any)["secret"] = "mutated" + require.Equal(t, "Ada", input["name"]) + require.Equal(t, "mutated", nested["secret"]) +} + +func TestNodeValidateRequiresSourceID(t *testing.T) { + require.Error(t, (entity.Node{}).Validate()) + require.NoError(t, (entity.Node{SourceID: "node-1"}).Validate()) +} + +func TestRelationshipValidateRequiresEndpointsAndKindButNotSourceID(t *testing.T) { + valid := entity.Relationship{StartID: "node-1", EndID: "node-2", Kind: "MEMBER_OF"} + + require.NoError(t, valid.Validate()) + require.Error(t, (entity.Relationship{EndID: valid.EndID, Kind: valid.Kind}).Validate()) + require.Error(t, (entity.Relationship{StartID: valid.StartID, Kind: valid.Kind}).Validate()) + require.Error(t, (entity.Relationship{StartID: valid.StartID, EndID: valid.EndID}).Validate()) +} diff --git a/ret/errors.go b/ret/errors.go new file mode 100644 index 00000000..2d225f69 --- /dev/null +++ b/ret/errors.go @@ -0,0 +1,16 @@ +package ret + +import "errors" + +var ( + ErrInvalidConfig = errors.New("ret invalid configuration") + ErrInvalidCollection = errors.New("ret invalid collection") + ErrArtifactIntegrity = errors.New("ret artifact integrity failure") + ErrDestinationExists = errors.New("ret destination exists") + ErrResumeRequired = errors.New("ret checkpoint requires resume") + ErrCheckpointMissing = errors.New("ret resume checkpoint missing") + ErrCollectionNotLoadable = errors.New("ret collection is not loadable") + ErrNonEmptyTarget = errors.New("ret target graph is not empty") + ErrSourceCountChanged = errors.New("ret source count changed") + ErrMetricsMismatch = errors.New("ret database metrics mismatch") +) diff --git a/ret/jsonl/artifact.go b/ret/jsonl/artifact.go new file mode 100644 index 00000000..7eed50c6 --- /dev/null +++ b/ret/jsonl/artifact.go @@ -0,0 +1,33 @@ +package jsonl + +import ( + "crypto/sha256" + "encoding/hex" + "fmt" +) + +type Artifact struct { + SchemaVersion string + Codec Codec + SHA256 string + Level int + Count int64 + UncompressedBytes int64 + StoredBytes int64 +} + +func (s Artifact) validate() error { + if s.SchemaVersion != SchemaVersion { + return fmt.Errorf("unsupported JSONL artifact schema %q", s.SchemaVersion) + } else if err := validateCodecLevel(s.Codec, s.Level); err != nil { + return fmt.Errorf("validate JSONL artifact codec: %w", err) + } else if s.Count < 0 || s.UncompressedBytes < 0 || s.StoredBytes < 0 { + return fmt.Errorf("JSONL artifact sizes and count must be non-negative") + } + decoded, err := hex.DecodeString(s.SHA256) + if err != nil || len(decoded) != sha256.Size { + return fmt.Errorf("JSONL artifact SHA-256 is invalid: %q", s.SHA256) + } + + return nil +} diff --git a/ret/jsonl/config.go b/ret/jsonl/config.go new file mode 100644 index 00000000..4389d54e --- /dev/null +++ b/ret/jsonl/config.go @@ -0,0 +1,75 @@ +// Package jsonl writes and verifies concrete JSON Lines graph artifacts. +package jsonl + +import ( + "compress/gzip" + "fmt" + + "github.com/klauspost/compress/zstd" +) + +const ( + // SchemaVersion identifies the JSONL artifact metadata and record layout. + SchemaVersion = "retriever-jsonl-v1" + + minZstdLevel = -5 + maxZstdLevel = 22 +) + +// Codec specifies the encoding used to store an artifact. +type Codec string + +const ( + CodecNone Codec = "none" + CodecGzip Codec = "gzip" + CodecZstd Codec = "zstd" +) + +type Config struct { + Codec Codec + Level int +} + +func (s Config) validate() error { + return validateCodecLevel(s.Codec, s.Level) +} + +func (s Config) Validate() error { + return s.validate() +} + +func validateCodecLevel(codec Codec, level int) error { + switch codec { + case CodecNone: + if level != 0 { + return fmt.Errorf("compression level %d is invalid for codec %q", level, codec) + } + case CodecGzip: + if level < gzip.HuffmanOnly || level > gzip.BestCompression { + return fmt.Errorf("gzip compression level %d is outside %d..%d", level, gzip.HuffmanOnly, gzip.BestCompression) + } + case CodecZstd: + if level < minZstdLevel || level > maxZstdLevel { + return fmt.Errorf("zstd compression level %d is outside %d..%d", level, minZstdLevel, maxZstdLevel) + } + default: + return fmt.Errorf("unsupported JSONL codec %q", codec) + } + + return nil +} + +func (s Config) gzipLevel() int { + if s.Level == 0 { + return gzip.DefaultCompression + } + return s.Level +} + +func (s Config) zstdLevel() zstd.EncoderLevel { + level := s.Level + if level == 0 { + level = 3 + } + return zstd.EncoderLevelFromZstd(level) +} diff --git a/ret/jsonl/jsonl.go b/ret/jsonl/jsonl.go new file mode 100644 index 00000000..77fad60c --- /dev/null +++ b/ret/jsonl/jsonl.go @@ -0,0 +1,9 @@ +package jsonl + +type State int + +const ( + Open State = iota + Failed + Closed +) diff --git a/ret/jsonl/reader.go b/ret/jsonl/reader.go new file mode 100644 index 00000000..8f738668 --- /dev/null +++ b/ret/jsonl/reader.go @@ -0,0 +1,315 @@ +package jsonl + +import ( + "bufio" + "bytes" + "compress/gzip" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "hash" + "io" + "math" + "strconv" + "strings" + + "github.com/klauspost/compress/zstd" + "github.com/specterops/dawgs/ret/entity" +) + +const ( + initialLineBuffer = 64 * 1024 + maxPhysicalLine = 10 * 1024 * 1024 +) + +var ( + ErrReaderNotOpen = fmt.Errorf("reader is not open") + ErrReaderNotDone = fmt.Errorf("reader is not done reading") + ErrReaderFailed = fmt.Errorf("reader failed") + ErrInvalidLimit = fmt.Errorf("pull limit must be positive") +) + +func NewNodeReader(reader io.Reader, artifact Artifact) (Reader[entity.Node, NodeRecord], error) { + if err := artifact.validate(); err != nil { + return Reader[entity.Node, NodeRecord]{}, err + } + + hasher := sha256.New() + fileReader := newCountingReader(io.TeeReader(reader, hasher)) + + decompressor, err := newDecompressionReader(fileReader, artifact.Codec) + if err != nil { + return Reader[entity.Node, NodeRecord]{}, err + } + + decomCounter := newCountingReader(decompressor) + scanner := bufio.NewScanner(decomCounter) + scanner.Buffer(make([]byte, initialLineBuffer), maxPhysicalLine+1) + + return Reader[entity.Node, NodeRecord]{ + artifact: artifact, + + hasher: hasher, + fileReader: fileReader, + decomReadCloser: decompressor, + decomCountingReader: decomCounter, + scanner: scanner, + + recordToEntity: nodeEntity, + }, nil +} + +func NewRelationshipReader(reader io.Reader, artifact Artifact) (Reader[entity.Relationship, RelationshipRecord], error) { + if err := artifact.validate(); err != nil { + return Reader[entity.Relationship, RelationshipRecord]{}, err + } + + hasher := sha256.New() + fileReader := newCountingReader(io.TeeReader(reader, hasher)) + + decompressor, err := newDecompressionReader(fileReader, artifact.Codec) + if err != nil { + return Reader[entity.Relationship, RelationshipRecord]{}, err + } + + decomCounter := newCountingReader(decompressor) + scanner := bufio.NewScanner(decomCounter) + scanner.Buffer(make([]byte, initialLineBuffer), maxPhysicalLine+1) + + return Reader[entity.Relationship, RelationshipRecord]{ + artifact: artifact, + + hasher: hasher, + fileReader: fileReader, + decomReadCloser: decompressor, + decomCountingReader: decomCounter, + scanner: scanner, + + recordToEntity: relationshipEntity, + }, nil +} + +type Reader[E entity.Entity, R record] struct { + artifact Artifact + + hasher hash.Hash + fileReader *countingReader + decomReadCloser io.ReadCloser + decomCountingReader *countingReader + scanner *bufio.Scanner + + recordToEntity func(R) E + recordCount int64 + + state State + resourceClosed bool + failure error +} + +func (s *Reader[E, R]) Pull(limit int) ([]E, error) { + if limit <= 0 { + return nil, ErrInvalidLimit + } + if s.resourceClosed { + return nil, ErrReaderNotOpen + } + switch s.state { + case Closed: + return nil, ErrReaderNotOpen + case Failed: + return nil, errors.Join(ErrReaderFailed, s.failure) + } + + entities := make([]E, 0, limit) + + for len(entities) < limit { + if !s.scanner.Scan() { + if err := s.scanner.Err(); err != nil { + return nil, s.fail(fmt.Errorf("read JSONL record %d: %w", s.recordCount+1, err)) + } + + s.state = Closed + return entities, nil + } + + var ( + line = s.scanner.Bytes() + record R + decoder = json.NewDecoder(bytes.NewReader(line)) + ) + decoder.DisallowUnknownFields() + decoder.UseNumber() + + if err := decoder.Decode(&record); err != nil { + return nil, s.fail(fmt.Errorf("decode JSONL record %d: %w", s.recordCount+1, err)) + } + if err := decoder.Decode(&struct{}{}); err != io.EOF { + if err == nil { + return nil, s.fail(fmt.Errorf("decode JSONL record %d: multiple JSON values", s.recordCount+1)) + } + return nil, s.fail(fmt.Errorf("decode JSONL record %d: %w", s.recordCount+1, err)) + } + + entity := s.recordToEntity(record) + if err := normalizeProperties(entityProperties(entity)); err != nil { + return nil, s.fail(fmt.Errorf("normalize JSONL record %d properties: %w", s.recordCount+1, err)) + } + if err := entity.Validate(); err != nil { + return nil, s.fail(fmt.Errorf("validate JSONL record %d: %w", s.recordCount+1, err)) + } + + entities = append(entities, entity) + s.recordCount++ + } + + return entities, nil +} + +func entityProperties[E entity.Entity](value E) map[string]any { + switch typed := any(value).(type) { + case entity.Node: + return typed.Properties + case entity.Relationship: + return typed.Properties + default: + return nil + } +} + +func normalizeProperties(properties map[string]any) error { + for key, value := range properties { + normalized, err := normalizeJSONValue(value, "properties."+key) + if err != nil { + return err + } + properties[key] = normalized + } + return nil +} + +func normalizeJSONValue(value any, path string) (any, error) { + switch typed := value.(type) { + case nil, bool, string: + return typed, nil + case json.Number: + literal := typed.String() + if !strings.ContainsAny(literal, ".eE") { + integer, err := strconv.ParseInt(literal, 10, 64) + if err != nil { + return nil, fmt.Errorf("JSON integer at %s is outside the int64 domain: %q", path, literal) + } + return integer, nil + } + fractional, err := strconv.ParseFloat(literal, 64) + if err != nil || math.IsNaN(fractional) || math.IsInf(fractional, 0) { + return nil, fmt.Errorf("JSON number at %s is not a finite float64: %q", path, literal) + } + return fractional, nil + case []any: + for index, element := range typed { + normalized, err := normalizeJSONValue(element, fmt.Sprintf("%s[%d]", path, index)) + if err != nil { + return nil, err + } + typed[index] = normalized + } + return typed, nil + case map[string]any: + for key, element := range typed { + normalized, err := normalizeJSONValue(element, path+"."+key) + if err != nil { + return nil, err + } + typed[key] = normalized + } + return typed, nil + default: + return nil, fmt.Errorf("JSON value at %s has unsupported decoded type %T", path, value) + } +} + +func (s *Reader[E, R]) Result() error { + if s.state == Open { + return ErrReaderNotDone + } else if s.state == Failed { + return errors.Join(ErrReaderFailed, s.failure) + } else if actual := hex.EncodeToString(s.hasher.Sum(nil)); s.artifact.SHA256 != actual { + return s.fail(fmt.Errorf("JSONL stored SHA-256 mismatch: got %s, want %s", actual, s.artifact.SHA256)) + } else if s.artifact.Count != s.recordCount { + return s.fail(fmt.Errorf("JSONL record count mismatch: got %d, want %d", s.recordCount, s.artifact.Count)) + } else if s.artifact.UncompressedBytes != s.decomCountingReader.count { + return s.fail(fmt.Errorf("JSONL uncompressed size mismatch: got %d, want %d", s.decomCountingReader.count, s.artifact.UncompressedBytes)) + } else if s.artifact.StoredBytes != s.fileReader.count { + return s.fail(fmt.Errorf("JSONL stored size mismatch: got %d, want %d", s.fileReader.count, s.artifact.StoredBytes)) + } + + return nil +} + +func (s *Reader[E, R]) Done() bool { + return s.state != Open +} + +func (s *Reader[E, R]) Close() error { + if s.resourceClosed { + if s.state == Failed { + return errors.Join(ErrReaderFailed, s.failure) + } + return nil + } + s.resourceClosed = true + if err := s.decomReadCloser.Close(); err != nil { + s.fail(fmt.Errorf("close JSONL reader: %w", err)) + } + if s.state == Failed { + return errors.Join(ErrReaderFailed, s.failure) + } + + return nil +} + +func (s *Reader[E, R]) fail(err error) error { + if s.failure == nil { + s.failure = err + } else { + s.failure = errors.Join(s.failure, err) + } + s.state = Failed + return err +} + +func newCountingReader(reader io.Reader) *countingReader { + return &countingReader{ + reader: reader, + } +} + +type countingReader struct { + reader io.Reader + count int64 +} + +func (s *countingReader) Read(value []byte) (int, error) { + read, err := s.reader.Read(value) + s.count += int64(read) + return read, err +} + +func newDecompressionReader(reader io.Reader, codec Codec) (io.ReadCloser, error) { + switch codec { + case CodecNone: + return io.NopCloser(reader), nil + case CodecGzip: + return gzip.NewReader(reader) + case CodecZstd: + decoder, err := zstd.NewReader(reader) + if err != nil { + return nil, err + } + return decoder.IOReadCloser(), nil + default: + return nil, fmt.Errorf("unsupported JSONL codec %q", codec) + } +} diff --git a/ret/jsonl/reader_test.go b/ret/jsonl/reader_test.go new file mode 100644 index 00000000..f71bed2a --- /dev/null +++ b/ret/jsonl/reader_test.go @@ -0,0 +1,169 @@ +package jsonl + +import ( + "bytes" + "crypto/sha256" + "encoding/hex" + "errors" + "testing" + + "github.com/specterops/dawgs/ret/entity" + "github.com/stretchr/testify/require" +) + +func TestReaderPullsBatchesAndVerifiesArtifact(t *testing.T) { + stored := []byte("{\"source_id\":\"1\",\"kinds\":null,\"properties\":null}\n{\"source_id\":\"2\",\"kinds\":null,\"properties\":null}\n") + reader, err := NewNodeReader(bytes.NewReader(stored), testArtifact(stored, 2)) + require.NoError(t, err) + + first, err := reader.Pull(1) + require.NoError(t, err) + require.Equal(t, []entity.Node{{SourceID: "1"}}, first) + require.False(t, reader.Done()) + + second, err := reader.Pull(1) + require.NoError(t, err) + require.Equal(t, []entity.Node{{SourceID: "2"}}, second) + + last, err := reader.Pull(1) + require.NoError(t, err) + require.Empty(t, last) + require.True(t, reader.Done()) + require.NoError(t, reader.Result()) + require.NoError(t, reader.Close()) +} + +func TestReaderRejectsNonPositiveLimit(t *testing.T) { + stored := []byte("{\"source_id\":\"1\"}\n") + reader, err := NewNodeReader(bytes.NewReader(stored), testArtifact(stored, 1)) + require.NoError(t, err) + + _, err = reader.Pull(0) + require.Error(t, err) + require.False(t, reader.Done()) + + _, err = reader.Pull(-1) + require.Error(t, err) + require.False(t, reader.Done()) +} + +func TestReaderRejectsMultipleJSONValuesAndLatchesFailure(t *testing.T) { + stored := []byte("{\"source_id\":\"1\"} {}\n") + reader, err := NewNodeReader(bytes.NewReader(stored), testArtifact(stored, 1)) + require.NoError(t, err) + + _, err = reader.Pull(1) + require.ErrorContains(t, err, "multiple JSON values") + require.True(t, reader.Done()) + require.ErrorIs(t, reader.Result(), ErrReaderFailed) + + _, err = reader.Pull(1) + require.ErrorIs(t, err, ErrReaderFailed) + require.ErrorIs(t, reader.Close(), ErrReaderFailed) + require.ErrorIs(t, reader.Close(), ErrReaderFailed) +} + +func TestReaderEnforcesOneKnownJSONRecordPerPhysicalLine(t *testing.T) { + tests := []struct { + name string + stored []byte + message string + }{ + { + name: "unknown field", + stored: []byte("{\"source_id\":\"1\",\"unexpected\":true}\n"), + message: "unknown field", + }, + { + name: "blank line", + stored: []byte("\n"), + message: "decode JSONL record 1", + }, + { + name: "oversized line", + stored: append(bytes.Repeat([]byte{' '}, maxPhysicalLine+1), '\n'), + message: "read JSONL record 1", + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + reader, err := NewNodeReader(bytes.NewReader(test.stored), testArtifact(test.stored, 1)) + require.NoError(t, err) + + _, err = reader.Pull(1) + require.ErrorContains(t, err, test.message) + require.True(t, reader.Done()) + require.ErrorIs(t, reader.Result(), ErrReaderFailed) + }) + } +} + +func TestReaderRejectsInvalidEntityAndLatchesFailure(t *testing.T) { + stored := []byte("{\"source_id\":\"\"}\n") + reader, err := NewNodeReader(bytes.NewReader(stored), testArtifact(stored, 1)) + require.NoError(t, err) + + _, err = reader.Pull(1) + require.ErrorContains(t, err, "source ID") + require.True(t, reader.Done()) + require.ErrorIs(t, reader.Result(), ErrReaderFailed) +} + +func TestReaderValidatesArtifactHashBeforeReading(t *testing.T) { + artifact := testArtifact(nil, 0) + artifact.SHA256 = "not-a-sha256" + + _, err := NewNodeReader(bytes.NewReader(nil), artifact) + require.ErrorContains(t, err, "SHA-256") +} + +func TestReaderResultReportsIntegrityMismatch(t *testing.T) { + stored := []byte("{\"source_id\":\"1\"}\n") + tests := []struct { + name string + mutate func(*Artifact) + message string + }{ + {name: "hash", mutate: func(value *Artifact) { value.SHA256 = string(bytes.Repeat([]byte{'0'}, 64)) }, message: "SHA-256 mismatch"}, + {name: "count", mutate: func(value *Artifact) { value.Count++ }, message: "record count mismatch"}, + {name: "uncompressed bytes", mutate: func(value *Artifact) { value.UncompressedBytes++ }, message: "uncompressed size mismatch"}, + {name: "stored bytes", mutate: func(value *Artifact) { value.StoredBytes++ }, message: "stored size mismatch"}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + artifact := testArtifact(stored, 1) + test.mutate(&artifact) + reader, err := NewNodeReader(bytes.NewReader(stored), artifact) + require.NoError(t, err) + + _, err = reader.Pull(2) + require.NoError(t, err) + require.ErrorContains(t, reader.Result(), test.message) + }) + } +} + +func TestReaderCloseIsIdempotent(t *testing.T) { + stored := []byte("{\"source_id\":\"1\"}\n") + reader, err := NewNodeReader(bytes.NewReader(stored), testArtifact(stored, 1)) + require.NoError(t, err) + + require.NoError(t, reader.Close()) + require.NoError(t, reader.Close()) + _, err = reader.Pull(1) + require.True(t, errors.Is(err, ErrReaderNotOpen) || errors.Is(err, ErrReaderFailed)) +} + +func testArtifact(stored []byte, count int64) Artifact { + hash := sha256.Sum256(stored) + return Artifact{ + SchemaVersion: SchemaVersion, + Codec: CodecNone, + SHA256: hex.EncodeToString(hash[:]), + Count: count, + UncompressedBytes: int64(len(stored)), + StoredBytes: int64(len(stored)), + } +} diff --git a/ret/jsonl/record.go b/ret/jsonl/record.go new file mode 100644 index 00000000..619e3f03 --- /dev/null +++ b/ret/jsonl/record.go @@ -0,0 +1,40 @@ +package jsonl + +import "github.com/specterops/dawgs/ret/entity" + +type record interface { + NodeRecord | RelationshipRecord +} + +// NodeRecord is the JSON representation of an entity.Node. +type NodeRecord struct { + SourceID string `json:"source_id"` + Kinds []string `json:"kinds"` + Properties map[string]any `json:"properties"` +} + +// RelationshipRecord is the JSON representation of an entity.Relationship. +// SourceID is intentionally omitted because JSONL relationships are identified +// by their endpoints and kind. +type RelationshipRecord struct { + StartID string `json:"start_id"` + EndID string `json:"end_id"` + Kind string `json:"kind"` + Properties map[string]any `json:"properties"` +} + +func nodeRecord(value entity.Node) NodeRecord { + return NodeRecord{SourceID: value.SourceID, Kinds: value.Kinds, Properties: value.Properties} +} + +func relationshipRecord(value entity.Relationship) RelationshipRecord { + return RelationshipRecord{StartID: value.StartID, EndID: value.EndID, Kind: value.Kind, Properties: value.Properties} +} + +func nodeEntity(value NodeRecord) entity.Node { + return entity.Node{SourceID: value.SourceID, Kinds: value.Kinds, Properties: value.Properties} +} + +func relationshipEntity(value RelationshipRecord) entity.Relationship { + return entity.Relationship{StartID: value.StartID, EndID: value.EndID, Kind: value.Kind, Properties: value.Properties} +} diff --git a/ret/jsonl/writer.go b/ret/jsonl/writer.go new file mode 100644 index 00000000..afd6c138 --- /dev/null +++ b/ret/jsonl/writer.go @@ -0,0 +1,207 @@ +package jsonl + +import ( + "compress/gzip" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "hash" + "io" + + "github.com/klauspost/compress/zstd" + "github.com/specterops/dawgs/ret/entity" +) + +var ( + ErrWriterNotOpen = fmt.Errorf("writer is not open") + ErrWriterNotClosed = fmt.Errorf("writer is not closed") + ErrWriterFailed = fmt.Errorf("writer failed") +) + +func NewNodeWriter(writer io.Writer, config Config) (EntityWriter[entity.Node, NodeRecord], error) { + if err := config.validate(); err != nil { + return EntityWriter[entity.Node, NodeRecord]{}, err + } + + hasher := sha256.New() + output := newCountingWriter(io.MultiWriter(writer, hasher)) + + compressionWriter, err := newCompressionWriter(output, config) + if err != nil { + return EntityWriter[entity.Node, NodeRecord]{}, err + } + + inputWriter := newCountingWriter(compressionWriter) + + return EntityWriter[entity.Node, NodeRecord]{ + outputWriter: output, + hasher: hasher, + + compressionWriter: compressionWriter, + inputWriter: inputWriter, + + entityToRecord: nodeRecord, + config: config, + + state: Open, + }, nil +} + +func NewRelationshipWriter(writer io.Writer, config Config) (EntityWriter[entity.Relationship, RelationshipRecord], error) { + if err := config.validate(); err != nil { + return EntityWriter[entity.Relationship, RelationshipRecord]{}, err + } + + hasher := sha256.New() + output := newCountingWriter(io.MultiWriter(writer, hasher)) + + compressionWriter, err := newCompressionWriter(output, config) + if err != nil { + return EntityWriter[entity.Relationship, RelationshipRecord]{}, err + } + + inputWriter := newCountingWriter(compressionWriter) + + return EntityWriter[entity.Relationship, RelationshipRecord]{ + outputWriter: output, + hasher: hasher, + + compressionWriter: compressionWriter, + inputWriter: inputWriter, + + entityToRecord: relationshipRecord, + config: config, + + state: Open, + }, nil +} + +type EntityWriter[E entity.Entity, R record] struct { + hasher hash.Hash + + outputWriter *countingWriter + compressionWriter io.WriteCloser + inputWriter *countingWriter + + recordCount int64 + + entityToRecord func(E) R + config Config + + state State + resourceClosed bool + failure error +} + +func (s *EntityWriter[E, R]) Push(entities []E) error { + switch s.state { + case Closed: + return ErrWriterNotOpen + case Failed: + return errors.Join(ErrWriterFailed, s.failure) + } + + for _, entity := range entities { + if err := entity.Validate(); err != nil { + return s.fail(fmt.Errorf("validate JSONL record %d: %w", s.recordCount+1, err)) + } + + record := s.entityToRecord(entity) + + if encoded, err := json.Marshal(record); err != nil { + return s.fail(fmt.Errorf("encode JSONL record %d: %w", s.recordCount+1, err)) + } else if _, err := s.inputWriter.Write(append(encoded, '\n')); err != nil { + return s.fail(fmt.Errorf("write JSONL record %d: %w", s.recordCount+1, err)) + } + + s.recordCount++ + } + + return nil +} + +func (s *EntityWriter[E, R]) Close() error { + if s.resourceClosed { + if s.state == Failed { + return errors.Join(ErrWriterFailed, s.failure) + } + return nil + } + s.resourceClosed = true + if err := s.compressionWriter.Close(); err != nil { + s.fail(fmt.Errorf("finish JSONL stream: %w", err)) + } + if s.state == Failed { + return errors.Join(ErrWriterFailed, s.failure) + } + s.state = Closed + + return nil +} + +func (s *EntityWriter[E, R]) Result() (Artifact, error) { + switch s.state { + case Open: + return Artifact{}, ErrWriterNotClosed + case Failed: + return Artifact{}, errors.Join(ErrWriterFailed, s.failure) + } + + return Artifact{ + SchemaVersion: SchemaVersion, + Codec: s.config.Codec, + SHA256: hex.EncodeToString(s.hasher.Sum(nil)), + Level: s.config.Level, + Count: s.recordCount, + UncompressedBytes: s.inputWriter.count, + StoredBytes: s.outputWriter.count, + }, nil +} + +func (s *EntityWriter[E, R]) fail(err error) error { + if s.failure == nil { + s.failure = err + } else { + s.failure = errors.Join(s.failure, err) + } + s.state = Failed + return err +} + +func newCountingWriter(writer io.Writer) *countingWriter { + return &countingWriter{writer: writer} +} + +type countingWriter struct { + writer io.Writer + count int64 +} + +func (s *countingWriter) Write(value []byte) (int, error) { + written, err := s.writer.Write(value) + s.count += int64(written) + return written, err +} + +type nopWriteCloser struct { + io.Writer +} + +func (nopWriteCloser) Close() error { + return nil +} + +func newCompressionWriter(writer io.Writer, config Config) (io.WriteCloser, error) { + switch config.Codec { + case CodecNone: + return nopWriteCloser{Writer: writer}, nil + case CodecGzip: + return gzip.NewWriterLevel(writer, config.gzipLevel()) + case CodecZstd: + return zstd.NewWriter(writer, zstd.WithEncoderLevel(config.zstdLevel())) + default: + return nil, fmt.Errorf("unsupported JSONL codec %q", config.Codec) + } +} diff --git a/ret/jsonl/writer_test.go b/ret/jsonl/writer_test.go new file mode 100644 index 00000000..f115b903 --- /dev/null +++ b/ret/jsonl/writer_test.go @@ -0,0 +1,124 @@ +package jsonl + +import ( + "bytes" + "errors" + "testing" + + "github.com/specterops/dawgs/ret/entity" + "github.com/stretchr/testify/require" +) + +func TestWriterProducesVerifiableArtifactsAcrossCodecs(t *testing.T) { + tests := []Config{ + {Codec: CodecNone}, + {Codec: CodecGzip}, + {Codec: CodecZstd, Level: 3}, + } + for _, config := range tests { + t.Run(string(config.Codec), func(t *testing.T) { + var stored bytes.Buffer + writer, err := NewNodeWriter(&stored, config) + require.NoError(t, err) + require.NoError(t, writer.Push([]entity.Node{{SourceID: "1"}, {SourceID: "2"}})) + _, err = writer.Result() + require.ErrorIs(t, err, ErrWriterNotClosed) + require.NoError(t, writer.Close()) + + artifact, err := writer.Result() + require.NoError(t, err) + require.EqualValues(t, 2, artifact.Count) + require.EqualValues(t, stored.Len(), artifact.StoredBytes) + + reader, err := NewNodeReader(bytes.NewReader(stored.Bytes()), artifact) + require.NoError(t, err) + values, err := reader.Pull(3) + require.NoError(t, err) + require.Equal(t, []entity.Node{{SourceID: "1"}, {SourceID: "2"}}, values) + require.NoError(t, reader.Result()) + require.NoError(t, reader.Close()) + }) + } +} + +func TestWriterRejectsInvalidEntityAndNeverReturnsArtifact(t *testing.T) { + var stored bytes.Buffer + writer, err := NewNodeWriter(&stored, Config{Codec: CodecNone}) + require.NoError(t, err) + + err = writer.Push([]entity.Node{{}}) + require.ErrorContains(t, err, "source ID") + require.ErrorIs(t, writer.Push([]entity.Node{{SourceID: "1"}}), ErrWriterFailed) + closeErr := writer.Close() + require.ErrorIs(t, closeErr, ErrWriterFailed) + require.ErrorContains(t, closeErr, "source ID") + closeErr = writer.Close() + require.ErrorIs(t, closeErr, ErrWriterFailed) + require.ErrorContains(t, closeErr, "source ID") + _, err = writer.Result() + require.ErrorIs(t, err, ErrWriterFailed) +} + +func TestRelationshipWriterRoundTripPreservesStoredFields(t *testing.T) { + want := entity.Relationship{ + SourceID: "relationship-1", + StartID: "node-1", + EndID: "node-2", + Kind: "MemberOf", + Properties: map[string]any{ + "weight": int64(3), + }, + } + var stored bytes.Buffer + writer, err := NewRelationshipWriter(&stored, Config{Codec: CodecNone}) + require.NoError(t, err) + require.NoError(t, writer.Push([]entity.Relationship{want})) + require.NoError(t, writer.Close()) + require.NoError(t, writer.Close()) + artifact, err := writer.Result() + require.NoError(t, err) + + reader, err := NewRelationshipReader(bytes.NewReader(stored.Bytes()), artifact) + require.NoError(t, err) + values, err := reader.Pull(2) + require.NoError(t, err) + // Relationship source IDs are pagination metadata and intentionally are not + // part of the current JSONL record shape. + want.SourceID = "" + require.Equal(t, []entity.Relationship{want}, values) + require.NoError(t, reader.Result()) + require.NoError(t, reader.Close()) +} + +func TestWriterLatchesUnderlyingWriteFailure(t *testing.T) { + writer, err := NewNodeWriter(failingWriter{}, Config{Codec: CodecNone}) + require.NoError(t, err) + + require.ErrorIs(t, writer.Push([]entity.Node{{SourceID: "1"}}), errInjectedWrite) + closeErr := writer.Close() + require.ErrorIs(t, closeErr, ErrWriterFailed) + require.ErrorIs(t, closeErr, errInjectedWrite) + closeErr = writer.Close() + require.ErrorIs(t, closeErr, ErrWriterFailed) + require.ErrorIs(t, closeErr, errInjectedWrite) + _, err = writer.Result() + require.ErrorIs(t, err, ErrWriterFailed) +} + +func TestWriterCloseIsIdempotent(t *testing.T) { + var stored bytes.Buffer + writer, err := NewNodeWriter(&stored, Config{Codec: CodecGzip}) + require.NoError(t, err) + require.NoError(t, writer.Close()) + require.NoError(t, writer.Close()) + _, err = writer.Result() + require.NoError(t, err) +} + +var errInjectedWrite = errors.New("injected write failure") + +type failingWriter struct{} + +func (failingWriter) Write([]byte) (int, error) { + return 0, errInjectedWrite +} diff --git a/ret/load.go b/ret/load.go new file mode 100644 index 00000000..f282eaa5 --- /dev/null +++ b/ret/load.go @@ -0,0 +1,302 @@ +package ret + +import ( + "context" + "errors" + "fmt" + "time" + + "github.com/specterops/dawgs/graph" + "github.com/specterops/dawgs/ret/collection" + "github.com/specterops/dawgs/ret/dawgs" + "github.com/specterops/dawgs/ret/entity" + "github.com/specterops/dawgs/ret/observe" +) + +const ( + loadOperationName = "load" + loadNodesPhase = "nodes" + loadRelationshipsPhase = "relationships" +) + +type replayGraphFunc func( + context.Context, + string, + collection.Graph, + func(entity.Node) error, + func(entity.Relationship) error, +) error + +// Load validates and replays every JSONL graph in a collection into empty +// target graphs. +func Load(ctx context.Context, database graph.Database, config LoadConfig) (LoadResult, error) { + return loadWithReplay(ctx, database, config, collection.ReplayGraph) +} + +func loadWithReplay( + ctx context.Context, + database graph.Database, + config LoadConfig, + replay replayGraphFunc, +) (result LoadResult, resultErr error) { + started := time.Now() + observe.Emit(ctx, config.Observer, observe.OperationStarted{Operation: loadOperationName}) + defer func() { + observe.Emit(ctx, config.Observer, observe.OperationCompleted{ + Operation: loadOperationName, + Duration: time.Since(started), + Err: resultErr, + }) + }() + + if err := ctx.Err(); err != nil { + return LoadResult{}, fmt.Errorf("load: %w", err) + } + if err := config.Validate(); err != nil { + return LoadResult{}, err + } + + verification, err := collection.VerifyJSONLForLoad(ctx, config.Directory, config.Observer) + if err != nil { + return LoadResult{}, fmt.Errorf("%w: %w", ErrCollectionNotLoadable, err) + } + if err := ctx.Err(); err != nil { + return LoadResult{}, fmt.Errorf("load collection preflight: %w", err) + } + + targets := make([]*dawgs.Target, len(verification.Manifest.Graphs)) + for index, graphEntry := range verification.Manifest.Graphs { + target, err := dawgs.NewTarget(database, graphEntry.Name, config.BatchSize) + if err != nil { + return LoadResult{}, fmt.Errorf("prepare load target for graph %q: %w", graphEntry.Name, err) + } + targets[index] = target + } + for index, target := range targets { + graphName := verification.Manifest.Graphs[index].Name + if err := target.RequireEmpty(ctx); err != nil { + if errors.Is(err, dawgs.ErrTargetNotEmpty) { + return LoadResult{}, fmt.Errorf("%w: graph %q: %w", ErrNonEmptyTarget, graphName, err) + } + return LoadResult{}, fmt.Errorf("load graph %q emptiness check: %w", graphName, err) + } + if err := ctx.Err(); err != nil { + return LoadResult{}, fmt.Errorf("load graph %q emptiness check: %w", graphName, err) + } + } + if err := ctx.Err(); err != nil { + return LoadResult{}, fmt.Errorf("load target emptiness checks: %w", err) + } + for index, target := range targets { + graphEntry := verification.Manifest.Graphs[index] + if err := ctx.Err(); err != nil { + return LoadResult{}, fmt.Errorf("load graph %q schema phase: %w", graphEntry.Name, err) + } + if err := target.AssertSchema(ctx, graphEntry.KindCatalog); err != nil { + return LoadResult{}, fmt.Errorf("load graph %q schema phase: %w", graphEntry.Name, err) + } + if err := ctx.Err(); err != nil { + return LoadResult{}, fmt.Errorf("load graph %q schema phase: %w", graphEntry.Name, err) + } + } + + result.GraphCount = len(verification.Manifest.Graphs) + for index, graphEntry := range verification.Manifest.Graphs { + nodes, relationships, err := loadGraph(ctx, config, graphEntry, targets[index], replay) + if err != nil { + return LoadResult{}, err + } + result.NodeCount += nodes + result.RelationshipCount += relationships + } + if err := ctx.Err(); err != nil { + return LoadResult{}, fmt.Errorf("load completion: %w", err) + } + return result, nil +} + +func loadGraph( + ctx context.Context, + config LoadConfig, + graphEntry collection.Graph, + target *dawgs.Target, + replay replayGraphFunc, +) (nodeCount int64, relationshipCount int64, resultErr error) { + graphStarted := time.Now() + observe.Emit(ctx, config.Observer, observe.GraphStarted{ + Operation: loadOperationName, + Graph: graphEntry.Name, + }) + if err := ctx.Err(); err != nil { + return 0, 0, partialLoadError(graphEntry.Name, loadNodesPhase, err) + } + + resolver := dawgs.NewResolver(graphEntry.NodeCount) + pendingNodes := make([]entity.Node, 0, config.BatchSize) + pendingRelationships := make([]entity.Relationship, 0, config.BatchSize) + nodesCompleted := false + relationshipsStarted := false + + observe.Emit(ctx, config.Observer, observe.PhaseStarted{ + Operation: loadOperationName, + Graph: graphEntry.Name, + Phase: loadNodesPhase, + Completed: 0, + Total: graphEntry.NodeCount, + }) + if err := ctx.Err(); err != nil { + return 0, 0, partialLoadError(graphEntry.Name, loadNodesPhase, err) + } + + flushNodes := func() error { + if len(pendingNodes) == 0 { + return nil + } + if err := target.CreateNodes(ctx, pendingNodes, resolver); err != nil { + return err + } + nodeCount += int64(len(pendingNodes)) + pendingNodes = pendingNodes[:0] + observe.Emit(ctx, config.Observer, observe.PhaseProgress{ + Operation: loadOperationName, + Graph: graphEntry.Name, + Phase: loadNodesPhase, + Completed: nodeCount, + Total: graphEntry.NodeCount, + }) + return ctx.Err() + } + finishNodes := func() error { + if nodesCompleted { + return nil + } + if err := flushNodes(); err != nil { + return err + } + if nodeCount != graphEntry.NodeCount { + return fmt.Errorf("count mismatch: got %d want %d", nodeCount, graphEntry.NodeCount) + } + nodesCompleted = true + observe.Emit(ctx, config.Observer, observe.PhaseCompleted{ + Operation: loadOperationName, + Graph: graphEntry.Name, + Phase: loadNodesPhase, + Completed: nodeCount, + Duration: time.Since(graphStarted), + }) + if err := ctx.Err(); err != nil { + return err + } + relationshipsStarted = true + observe.Emit(ctx, config.Observer, observe.PhaseStarted{ + Operation: loadOperationName, + Graph: graphEntry.Name, + Phase: loadRelationshipsPhase, + Completed: 0, + Total: graphEntry.RelationshipCount, + }) + return ctx.Err() + } + flushRelationships := func() error { + if len(pendingRelationships) == 0 { + return nil + } + if err := target.CreateRelationships(ctx, pendingRelationships, resolver); err != nil { + return err + } + relationshipCount += int64(len(pendingRelationships)) + pendingRelationships = pendingRelationships[:0] + observe.Emit(ctx, config.Observer, observe.PhaseProgress{ + Operation: loadOperationName, + Graph: graphEntry.Name, + Phase: loadRelationshipsPhase, + Completed: relationshipCount, + Total: graphEntry.RelationshipCount, + }) + return ctx.Err() + } + + err := replay( + ctx, + config.Directory, + graphEntry, + func(node entity.Node) error { + pendingNodes = append(pendingNodes, node) + if len(pendingNodes) == config.BatchSize { + return flushNodes() + } + return nil + }, + func(relationship entity.Relationship) error { + if err := finishNodes(); err != nil { + return err + } + pendingRelationships = append(pendingRelationships, relationship) + if len(pendingRelationships) == config.BatchSize { + return flushRelationships() + } + return nil + }, + ) + if err != nil { + phase := loadNodesPhase + if relationshipsStarted { + phase = loadRelationshipsPhase + } + return nodeCount, relationshipCount, partialLoadError(graphEntry.Name, phase, err) + } + if err := ctx.Err(); err != nil { + phase := loadNodesPhase + if relationshipsStarted { + phase = loadRelationshipsPhase + } + return nodeCount, relationshipCount, partialLoadError(graphEntry.Name, phase, err) + } + if err := finishNodes(); err != nil { + return nodeCount, relationshipCount, partialLoadError(graphEntry.Name, loadNodesPhase, err) + } + if err := flushRelationships(); err != nil { + return nodeCount, relationshipCount, partialLoadError(graphEntry.Name, loadRelationshipsPhase, err) + } + if relationshipCount != graphEntry.RelationshipCount { + return nodeCount, relationshipCount, partialLoadError( + graphEntry.Name, + loadRelationshipsPhase, + fmt.Errorf("count mismatch: got %d want %d", relationshipCount, graphEntry.RelationshipCount), + ) + } + if err := ctx.Err(); err != nil { + return nodeCount, relationshipCount, partialLoadError(graphEntry.Name, loadRelationshipsPhase, err) + } + + observe.Emit(ctx, config.Observer, observe.PhaseCompleted{ + Operation: loadOperationName, + Graph: graphEntry.Name, + Phase: loadRelationshipsPhase, + Completed: relationshipCount, + Duration: time.Since(graphStarted), + }) + if err := ctx.Err(); err != nil { + return nodeCount, relationshipCount, partialLoadError(graphEntry.Name, loadRelationshipsPhase, err) + } + observe.Emit(ctx, config.Observer, observe.GraphCompleted{ + Operation: loadOperationName, + Graph: graphEntry.Name, + Nodes: nodeCount, + Relationships: relationshipCount, + Duration: time.Since(graphStarted), + }) + if err := ctx.Err(); err != nil { + return nodeCount, relationshipCount, partialLoadError(graphEntry.Name, loadRelationshipsPhase, err) + } + return nodeCount, relationshipCount, nil +} + +func partialLoadError(graphName, phase string, err error) error { + return fmt.Errorf( + "load graph %q %s phase failed; partial graph must be cleared before retry: %w", + graphName, + phase, + err, + ) +} diff --git a/ret/load_test.go b/ret/load_test.go new file mode 100644 index 00000000..1afa5b96 --- /dev/null +++ b/ret/load_test.go @@ -0,0 +1,722 @@ +package ret + +import ( + "context" + "errors" + "fmt" + "os" + "path/filepath" + "testing" + + "github.com/specterops/dawgs/graph" + "github.com/specterops/dawgs/ret/collection" + "github.com/specterops/dawgs/ret/entity" + "github.com/specterops/dawgs/ret/jsonl" + "github.com/specterops/dawgs/ret/observe" + "github.com/specterops/dawgs/ret/parquet" + "github.com/stretchr/testify/require" +) + +func TestLoadCorruptJSONLFailsBeforeDatabaseMutation(t *testing.T) { + // Break caught: validating and writing one graph at a time, which can leave + // earlier graphs mutated before corruption in a later JSONL artifact is found. + root := writeLoadCollection(t, []string{"first", "second"}, map[string]*dumpTestGraph{ + "first": {nodes: dumpTestNodes(1)}, + "second": {nodes: dumpTestNodes(2)}, + }, true, false) + manifest, err := collection.Read(root) + require.NoError(t, err) + lastArtifact := manifest.Graphs[1].NodeShards[0].JSONL + require.NoError(t, os.WriteFile(filepath.Join(root, filepath.FromSlash(lastArtifact.Path)), []byte("corrupt"), 0o600)) + database := newLoadTestDatabase() + + _, err = Load(context.Background(), database, LoadConfig{Directory: root, BatchSize: 2}) + + require.ErrorIs(t, err, ErrCollectionNotLoadable) + require.Empty(t, database.schemas) + require.Empty(t, database.durableMutations()) +} + +func TestLoadIgnoresCorruptParquet(t *testing.T) { + // Break caught: using full collection verification for load and therefore + // opening an optional Parquet twin that JSONL replay does not consume. + root := writeLoadCollection(t, []string{"asset"}, map[string]*dumpTestGraph{ + "asset": { + nodes: dumpTestNodes(1, 2), + relationships: dumpTestRelationships(10, 1, 2, "LINKED"), + }, + }, true, true) + corruptLoadParquetArtifacts(t, root) + database := newLoadTestDatabase() + + result, err := Load(context.Background(), database, LoadConfig{Directory: root, BatchSize: 2}) + + require.NoError(t, err) + require.Equal(t, LoadResult{GraphCount: 1, NodeCount: 2, RelationshipCount: 1}, result) + require.Len(t, database.graphs["asset"].nodes, 2) + require.Len(t, database.graphs["asset"].relationships, 1) +} + +func TestLoadRejectsParquetOnlyBeforeDatabaseMutation(t *testing.T) { + // Break caught: silently treating Parquet as a supported load source or + // preparing the target before discovering that JSONL is absent. + root := writeLoadCollection(t, []string{"asset"}, map[string]*dumpTestGraph{ + "asset": {nodes: dumpTestNodes(1)}, + }, false, true) + database := newLoadTestDatabase() + + _, err := Load(context.Background(), database, LoadConfig{Directory: root, BatchSize: 1}) + + require.ErrorIs(t, err, ErrCollectionNotLoadable) + require.Empty(t, database.schemas) + require.Empty(t, database.durableMutations()) +} + +func TestLoadRequiresEveryTargetEmptyBeforeAssertingAnySchema(t *testing.T) { + // Break caught: asserting schema immediately after checking each graph, + // which mutates schema before a later nonempty graph rejects the load. + root := writeLoadCollection(t, []string{"first", "second"}, map[string]*dumpTestGraph{ + "first": {nodes: dumpTestNodes(1)}, + "second": {nodes: dumpTestNodes(2)}, + }, true, false) + database := newLoadTestDatabase() + database.graphs["second"] = &loadTestGraphState{ + nodes: []*graph.Node{graph.NewNode(900, graph.NewProperties(), graph.StringKind("Existing"))}, + } + + _, err := Load(context.Background(), database, LoadConfig{Directory: root, BatchSize: 2}) + + require.ErrorIs(t, err, ErrNonEmptyTarget) + require.Equal(t, []string{"empty:first", "empty:second"}, database.operations) + require.Empty(t, database.schemas) + require.Len(t, database.graphs["second"].nodes, 1) +} + +func TestLoadSnapshotFailureIsNotClassifiedAsNonEmpty(t *testing.T) { + // Break caught: wrapping an unavailable snapshot with ErrNonEmptyTarget even + // though no nonzero target count was observed. + root := writeLoadCollection(t, []string{"asset"}, map[string]*dumpTestGraph{ + "asset": {nodes: dumpTestNodes(1)}, + }, true, false) + injected := errors.New("injected load snapshot failure") + database := newLoadTestDatabase() + database.readErr = injected + + _, err := Load(context.Background(), database, LoadConfig{Directory: root, BatchSize: 1}) + + require.ErrorIs(t, err, injected) + require.NotErrorIs(t, err, ErrNonEmptyTarget) + require.Empty(t, database.schemas) +} + +func TestLoadEmptinessCancellationIsNotClassifiedAsNonEmpty(t *testing.T) { + // Break caught: telling callers cleanup is required when the emptiness query + // was merely canceled and never proved that the graph had data. + root := writeLoadCollection(t, []string{"asset"}, map[string]*dumpTestGraph{ + "asset": {nodes: dumpTestNodes(1)}, + }, true, false) + database := newLoadTestDatabase() + database.readErr = context.Canceled + + _, err := Load(context.Background(), database, LoadConfig{Directory: root, BatchSize: 1}) + + require.ErrorIs(t, err, context.Canceled) + require.NotErrorIs(t, err, ErrNonEmptyTarget) + require.Empty(t, database.schemas) +} + +func TestLoadPreservesManifestGraphPhaseBatchAndEntityOrder(t *testing.T) { + // Break caught: sorting manifest graphs, interleaving entity phases, sending + // oversized slices, or normalizing away kind order and shallow properties. + nested := map[string]any{"shared": "node"} + root := writeLoadCollection(t, []string{"second", "first"}, map[string]*dumpTestGraph{ + "second": { + nodes: []*graph.Node{ + graph.NewNode(1, graph.AsProperties(map[string]any{"name": "Ada", "nested": nested}), graph.StringKind("User"), graph.StringKind("Admin"), graph.StringKind("User")), + graph.NewNode(2, graph.AsProperties(map[string]any{"name": "Ops"}), graph.StringKind("Group")), + graph.NewNode(3, graph.AsProperties(map[string]any{"name": "Grace"}), graph.StringKind("User")), + }, + relationships: []*graph.Relationship{ + graph.NewRelationship(10, 1, 2, graph.AsProperties(map[string]any{"role": "owner"}), graph.StringKind("MEMBER_OF")), + graph.NewRelationship(11, 3, 2, graph.AsProperties(map[string]any{"role": "reader"}), graph.StringKind("MEMBER_OF")), + }, + }, + "first": {nodes: dumpTestNodes(20)}, + }, true, false) + database := newLoadTestDatabase() + var events []observe.Event + observer := observe.ObserverFunc(func(_ context.Context, event observe.Event) { + switch event.(type) { + case observe.OperationStarted, observe.OperationCompleted, observe.GraphStarted, observe.GraphCompleted, + observe.PhaseStarted, observe.PhaseProgress, observe.PhaseCompleted: + events = append(events, event) + } + }) + + result, err := Load(context.Background(), database, LoadConfig{ + Directory: root, + BatchSize: 2, + Observer: observer, + }) + + require.NoError(t, err) + require.Equal(t, LoadResult{GraphCount: 2, NodeCount: 4, RelationshipCount: 2}, result) + require.Equal(t, []string{ + "empty:second", "empty:first", + "schema:second", "schema:first", + "nodes:second:2", "nodes:second:1", "relationships:second:2", + "nodes:first:1", + }, database.operations) + require.Equal(t, []int{2, 1, 2, 1}, database.committedBatchSizes) + require.Equal(t, []int{2, 2, 2, 2}, database.requestedBatchSizes) + + second := database.graphs["second"] + require.Len(t, second.nodes, 3) + require.Equal(t, []string{"User", "Admin", "User"}, second.nodes[0].Kinds.Strings()) + require.Equal(t, "Ada", second.nodes[0].Properties.Map["name"]) + require.Equal(t, map[string]any{"shared": "node"}, second.nodes[0].Properties.Map["nested"]) + require.Len(t, second.relationships, 2) + require.Equal(t, "MEMBER_OF", second.relationships[0].Kind.String()) + require.Equal(t, "owner", second.relationships[0].Properties.Map["role"]) + + require.Equal(t, []string{ + "operation_started:load", + "graph_started:second", + "phase_started:second:nodes:0:3", + "phase_progress:second:nodes:2:3", + "phase_progress:second:nodes:3:3", + "phase_completed:second:nodes:3", + "phase_started:second:relationships:0:2", + "phase_progress:second:relationships:2:2", + "phase_completed:second:relationships:2", + "graph_completed:second:3:2", + "graph_started:first", + "phase_started:first:nodes:0:1", + "phase_progress:first:nodes:1:1", + "phase_completed:first:nodes:1", + "phase_started:first:relationships:0:0", + "phase_completed:first:relationships:0", + "graph_completed:first:1:0", + "operation_completed:load:ok", + }, loadEventNames(events)) + require.NoFileExists(t, filepath.Join(root, checkpointFileNameForTest)) +} + +func TestLoadDatabaseFailureReportsDurablePartialGraphAndRequiredRetryCleanup(t *testing.T) { + // Break caught: reporting only the driver error, or implying an automatic + // rollback across earlier committed batches when the graph is partial. + root := writeLoadCollection(t, []string{"asset"}, map[string]*dumpTestGraph{ + "asset": { + nodes: dumpTestNodes(1, 2, 3), + relationships: dumpTestRelationships(10, 1, 2, "LINKED"), + }, + }, true, false) + injected := errors.New("injected database batch failure") + database := newLoadTestDatabase() + database.failBatchAt = 2 + database.failBatchErr = injected + var events []observe.Event + + _, err := Load(context.Background(), database, LoadConfig{ + Directory: root, + BatchSize: 2, + Observer: observe.ObserverFunc(func(_ context.Context, event observe.Event) { + events = append(events, event) + }), + }) + + require.ErrorIs(t, err, injected) + require.ErrorContains(t, err, `graph "asset"`) + require.ErrorContains(t, err, "nodes phase") + require.ErrorContains(t, err, "partial graph must be cleared before retry") + require.Len(t, database.graphs["asset"].nodes, 2) + require.Empty(t, database.graphs["asset"].relationships) + completed, ok := events[len(events)-1].(observe.OperationCompleted) + require.True(t, ok) + require.ErrorIs(t, completed.Err, injected) +} + +func TestLoadMissingEndpointKeepsPriorRelationshipAndRollsBackFailingBatch(t *testing.T) { + // Break caught: losing a valid preceding relationship in the failing batch + // test, or committing the unresolved relationship despite delegate rollback. + root := writeLoadCollection(t, []string{"asset"}, map[string]*dumpTestGraph{ + "asset": { + nodes: []*graph.Node{ + graph.NewNode(1, graph.NewProperties(), graph.StringKind("Entity")), + graph.NewNode(2, graph.NewProperties(), graph.StringKind("Entity")), + }, + relationships: []*graph.Relationship{ + graph.NewRelationship(10, 1, 2, graph.NewProperties(), graph.StringKind("VALID")), + graph.NewRelationship(11, 2, 1, graph.NewProperties(), graph.StringKind("ALSO_VALID")), + }, + }, + }, true, false) + replay := func( + _ context.Context, + _ string, + _ collection.Graph, + visitNode func(entity.Node) error, + visitRelationship func(entity.Relationship) error, + ) error { + for _, node := range []entity.Node{ + {SourceID: "1", Kinds: []string{"Entity"}}, + {SourceID: "2", Kinds: []string{"Entity"}}, + } { + if err := visitNode(node); err != nil { + return err + } + } + for _, relationship := range []entity.Relationship{ + {SourceID: "10", StartID: "1", EndID: "2", Kind: "VALID"}, + {SourceID: "11", StartID: "2", EndID: "missing", Kind: "BROKEN"}, + } { + if err := visitRelationship(relationship); err != nil { + return err + } + } + return nil + } + database := newLoadTestDatabase() + + _, err := loadWithReplay( + context.Background(), + database, + LoadConfig{Directory: root, BatchSize: 1}, + replay, + ) + + require.ErrorContains(t, err, `unresolved endpoints "2" -> "missing"`) + require.ErrorContains(t, err, `graph "asset"`) + require.ErrorContains(t, err, "relationships phase") + require.ErrorContains(t, err, "partial graph must be cleared before retry") + require.Len(t, database.graphs["asset"].nodes, 2) + require.Len(t, database.graphs["asset"].relationships, 1) + require.Equal(t, "VALID", database.graphs["asset"].relationships[0].Kind.String()) +} + +func TestLoadDetectsReplayCountMismatchAfterSuccessfulBoundedWrites(t *testing.T) { + // Break caught: trusting the preflight manifest totals without comparing + // what the replay callback actually delivered at the load boundary. + root := writeLoadCollection(t, []string{"asset"}, map[string]*dumpTestGraph{ + "asset": {nodes: dumpTestNodes(1, 2)}, + }, true, false) + replay := func( + _ context.Context, + _ string, + _ collection.Graph, + visitNode func(entity.Node) error, + _ func(entity.Relationship) error, + ) error { + return visitNode(entity.Node{SourceID: "1", Kinds: []string{"Entity"}}) + } + database := newLoadTestDatabase() + + _, err := loadWithReplay( + context.Background(), + database, + LoadConfig{Directory: root, BatchSize: 2}, + replay, + ) + + require.ErrorContains(t, err, `graph "asset"`) + require.ErrorContains(t, err, "nodes phase") + require.ErrorContains(t, err, "count mismatch: got 1 want 2") + require.ErrorContains(t, err, "partial graph must be cleared before retry") + require.Len(t, database.graphs["asset"].nodes, 1) +} + +func TestLoadCancellationOnFinalPreflightArtifactStopsBeforeTargetAccess(t *testing.T) { + // Break caught: continuing from the synchronous final preflight artifact + // event into target reads after its observer cancels the operation. + root := writeLoadCollection(t, []string{"asset"}, map[string]*dumpTestGraph{ + "asset": {nodes: dumpTestNodes(1)}, + }, true, false) + database := newLoadTestDatabase() + ctx, cancel := context.WithCancel(context.Background()) + var events []observe.Event + + _, err := Load(ctx, database, LoadConfig{ + Directory: root, + BatchSize: 1, + Observer: observe.ObserverFunc(func(_ context.Context, event observe.Event) { + events = append(events, event) + if _, ok := event.(observe.ArtifactVerified); ok { + cancel() + } + }), + }) + + require.ErrorIs(t, err, context.Canceled) + require.NotErrorIs(t, err, ErrNonEmptyTarget) + require.Zero(t, database.readCalls) + require.Empty(t, database.schemas) + requireSingleCanceledTerminalEvent(t, events) +} + +func TestLoadCancellationAfterAllEmptinessChecksStopsBeforeSchema(t *testing.T) { + // Break caught: crossing the global empty-target barrier into schema mutation + // after the final synchronous emptiness read cancels the context. + root := writeLoadCollection(t, []string{"asset"}, map[string]*dumpTestGraph{ + "asset": {nodes: dumpTestNodes(1)}, + }, true, false) + database := newLoadTestDatabase() + ctx, cancel := context.WithCancel(context.Background()) + database.afterRead = cancel + var events []observe.Event + + _, err := Load(ctx, database, LoadConfig{ + Directory: root, + BatchSize: 1, + Observer: observe.ObserverFunc(func(_ context.Context, event observe.Event) { + events = append(events, event) + }), + }) + + require.ErrorIs(t, err, context.Canceled) + require.NotErrorIs(t, err, ErrNonEmptyTarget) + require.Empty(t, database.schemas) + require.Empty(t, database.durableMutations()) + requireSingleCanceledTerminalEvent(t, events) +} + +func TestLoadCancellationOnRelationshipPhaseStartStopsBeforeRelationshipMutation(t *testing.T) { + // Break caught: relying on the database to notice cancellation instead of + // gating the mutation immediately after synchronous PhaseStarted delivery. + root := writeLoadCollection(t, []string{"asset"}, map[string]*dumpTestGraph{ + "asset": { + nodes: dumpTestNodes(1, 2), + relationships: dumpTestRelationships(10, 1, 2, "LINKED"), + }, + }, true, false) + database := newLoadTestDatabase() + database.ignoreContext = true + ctx, cancel := context.WithCancel(context.Background()) + var events []observe.Event + + _, err := Load(ctx, database, LoadConfig{ + Directory: root, + BatchSize: 1, + Observer: observe.ObserverFunc(func(_ context.Context, event observe.Event) { + events = append(events, event) + if value, ok := event.(observe.PhaseStarted); ok && value.Phase == loadRelationshipsPhase { + cancel() + } + }), + }) + + require.ErrorIs(t, err, context.Canceled) + require.Len(t, database.graphs["asset"].nodes, 2) + require.Empty(t, database.graphs["asset"].relationships) + require.Equal(t, []string{ + "operation_started:load", + "graph_started:asset", + "phase_started:asset:nodes:0:2", + "phase_progress:asset:nodes:1:2", + "phase_progress:asset:nodes:2:2", + "phase_completed:asset:nodes:2", + "phase_started:asset:relationships:0:1", + "operation_completed:load:error", + }, loadLifecycleEventNames(events)) + requireSingleCanceledTerminalEvent(t, events) +} + +func TestLoadCancellationOnFinalPhaseProgressSuppressesSuccessEvents(t *testing.T) { + // Break caught: returning success and emitting phase/graph completion after + // the observer cancels on the final committed progress event. + root := writeLoadCollection(t, []string{"asset"}, map[string]*dumpTestGraph{ + "asset": { + nodes: dumpTestNodes(1, 2), + relationships: dumpTestRelationships(10, 1, 2, "LINKED"), + }, + }, true, false) + database := newLoadTestDatabase() + ctx, cancel := context.WithCancel(context.Background()) + var events []observe.Event + + _, err := Load(ctx, database, LoadConfig{ + Directory: root, + BatchSize: 2, + Observer: observe.ObserverFunc(func(_ context.Context, event observe.Event) { + events = append(events, event) + if value, ok := event.(observe.PhaseProgress); ok && + value.Phase == loadRelationshipsPhase && + value.Completed == value.Total { + cancel() + } + }), + }) + + require.ErrorIs(t, err, context.Canceled) + require.Len(t, database.graphs["asset"].relationships, 1) + names := loadLifecycleEventNames(events) + require.Equal(t, "phase_progress:asset:relationships:1:1", names[len(names)-2]) + require.Equal(t, "operation_completed:load:error", names[len(names)-1]) + requireSingleCanceledTerminalEvent(t, events) +} + +func writeLoadCollection( + t *testing.T, + graphOrder []string, + graphs map[string]*dumpTestGraph, + jsonlEnabled bool, + parquetEnabled bool, +) string { + t.Helper() + config := validRootDumpConfig(t) + config.Graphs = append([]string(nil), graphOrder...) + config.EntityBatchSize = 10 + config.ShardSize = 10 + config.JSONL = nil + if jsonlEnabled { + config.JSONL = pointerTo(jsonl.Config{Codec: jsonl.CodecNone}) + } + config.Parquet = nil + if parquetEnabled { + config.Parquet = pointerTo(parquet.Config{}) + } + config.Scrub = nil + _, err := Dump(context.Background(), newDumpTestDatabase(graphs), config) + require.NoError(t, err) + return config.Directory +} + +func corruptLoadParquetArtifacts(t *testing.T, root string) { + t.Helper() + manifest, err := collection.Read(root) + require.NoError(t, err) + for _, graphEntry := range manifest.Graphs { + for _, shard := range graphEntry.NodeShards { + if shard.Parquet != nil { + require.NoError(t, os.WriteFile(filepath.Join(root, filepath.FromSlash(shard.Parquet.Path)), []byte("corrupt"), 0o600)) + } + } + for _, shard := range graphEntry.RelationshipShards { + if shard.Parquet != nil { + require.NoError(t, os.WriteFile(filepath.Join(root, filepath.FromSlash(shard.Parquet.Path)), []byte("corrupt"), 0o600)) + } + } + } +} + +func loadEventNames(events []observe.Event) []string { + names := make([]string, len(events)) + for index, event := range events { + switch value := event.(type) { + case observe.OperationStarted: + names[index] = "operation_started:" + value.Operation + case observe.OperationCompleted: + status := "ok" + if value.Err != nil { + status = "error" + } + names[index] = "operation_completed:" + value.Operation + ":" + status + case observe.GraphStarted: + names[index] = "graph_started:" + value.Graph + case observe.GraphCompleted: + names[index] = fmt.Sprintf("graph_completed:%s:%d:%d", value.Graph, value.Nodes, value.Relationships) + case observe.PhaseStarted: + names[index] = fmt.Sprintf("phase_started:%s:%s:%d:%d", value.Graph, value.Phase, value.Completed, value.Total) + case observe.PhaseProgress: + names[index] = fmt.Sprintf("phase_progress:%s:%s:%d:%d", value.Graph, value.Phase, value.Completed, value.Total) + case observe.PhaseCompleted: + names[index] = fmt.Sprintf("phase_completed:%s:%s:%d", value.Graph, value.Phase, value.Completed) + default: + names[index] = fmt.Sprintf("unexpected:%T", event) + } + } + return names +} + +func loadLifecycleEventNames(events []observe.Event) []string { + var lifecycle []observe.Event + for _, event := range events { + switch event.(type) { + case observe.OperationStarted, observe.OperationCompleted, observe.GraphStarted, observe.GraphCompleted, + observe.PhaseStarted, observe.PhaseProgress, observe.PhaseCompleted: + lifecycle = append(lifecycle, event) + } + } + return loadEventNames(lifecycle) +} + +func requireSingleCanceledTerminalEvent(t *testing.T, events []observe.Event) { + t.Helper() + var completions []observe.OperationCompleted + for _, event := range events { + if value, ok := event.(observe.OperationCompleted); ok { + completions = append(completions, value) + } + } + require.Len(t, completions, 1) + require.ErrorIs(t, completions[0].Err, context.Canceled) + require.IsType(t, observe.OperationCompleted{}, events[len(events)-1]) +} + +type loadTestDatabase struct { + graph.Database + + graphs map[string]*loadTestGraphState + schemas []graph.Schema + operations []string + requestedBatchSizes []int + committedBatchSizes []int + nextID graph.ID + batchCalls int + failBatchAt int + failBatchErr error + failSchemaGraph string + failSchemaErr error + readErr error + readCalls int + afterRead func() + ignoreContext bool +} + +type loadTestGraphState struct { + nodes []*graph.Node + relationships []*graph.Relationship +} + +func newLoadTestDatabase() *loadTestDatabase { + return &loadTestDatabase{ + graphs: make(map[string]*loadTestGraphState), + nextID: 100, + } +} + +func (s *loadTestDatabase) graphState(name string) *loadTestGraphState { + state := s.graphs[name] + if state == nil { + state = &loadTestGraphState{} + s.graphs[name] = state + } + return state +} + +func (s *loadTestDatabase) durableMutations() []string { + var mutations []string + for name, state := range s.graphs { + if len(state.nodes) > 0 || len(state.relationships) > 0 { + mutations = append(mutations, name) + } + } + return mutations +} + +func (s *loadTestDatabase) ReadTransaction(ctx context.Context, delegate graph.TransactionDelegate, _ ...graph.TransactionOption) error { + s.readCalls++ + if s.readErr != nil { + return s.readErr + } + if err := ctx.Err(); err != nil && !s.ignoreContext { + return err + } + err := delegate(&loadTestTransaction{database: s}) + if s.afterRead != nil { + s.afterRead() + } + return err +} + +func (s *loadTestDatabase) AssertSchema(_ context.Context, schema graph.Schema) error { + name := schema.DefaultGraph.Name + s.operations = append(s.operations, "schema:"+name) + if name == s.failSchemaGraph { + return s.failSchemaErr + } + s.schemas = append(s.schemas, schema) + return nil +} + +func (s *loadTestDatabase) BatchOperation(ctx context.Context, delegate graph.BatchDelegate, options ...graph.BatchOption) error { + if err := ctx.Err(); err != nil && !s.ignoreContext { + return err + } + config := graph.BatchConfig{} + for _, option := range options { + option(&config) + } + s.requestedBatchSizes = append(s.requestedBatchSizes, config.BatchSize) + s.batchCalls++ + batch := &loadTestBatch{database: s} + if err := delegate(batch); err != nil { + return err + } + if s.failBatchAt == s.batchCalls { + return s.failBatchErr + } + + state := s.graphState(batch.graphName) + state.nodes = append(state.nodes, batch.nodes...) + state.relationships = append(state.relationships, batch.relationships...) + size := len(batch.nodes) + len(batch.relationships) + s.committedBatchSizes = append(s.committedBatchSizes, size) + entityType := "nodes" + if len(batch.relationships) > 0 { + entityType = "relationships" + } + s.operations = append(s.operations, fmt.Sprintf("%s:%s:%d", entityType, batch.graphName, size)) + return nil +} + +type loadTestTransaction struct { + graph.Transaction + database *loadTestDatabase + graphName string +} + +func (s *loadTestTransaction) WithGraph(target graph.Graph) graph.Transaction { + s.graphName = target.Name + s.database.operations = append(s.database.operations, "empty:"+target.Name) + return s +} + +func (s *loadTestTransaction) Nodes() graph.NodeQuery { + return loadTestNodeQuery{count: int64(len(s.database.graphState(s.graphName).nodes))} +} + +func (s *loadTestTransaction) Relationships() graph.RelationshipQuery { + return loadTestRelationshipQuery{count: int64(len(s.database.graphState(s.graphName).relationships))} +} + +type loadTestNodeQuery struct { + graph.NodeQuery + count int64 +} + +func (s loadTestNodeQuery) Count() (int64, error) { return s.count, nil } + +type loadTestRelationshipQuery struct { + graph.RelationshipQuery + count int64 +} + +func (s loadTestRelationshipQuery) Count() (int64, error) { return s.count, nil } + +type loadTestBatch struct { + graph.Batch + database *loadTestDatabase + graphName string + nodes []*graph.Node + relationships []*graph.Relationship +} + +func (s *loadTestBatch) WithGraph(target graph.Graph) graph.Batch { + s.graphName = target.Name + return s +} + +func (s *loadTestBatch) CreateNodes(nodes []*graph.Node) ([]graph.ID, error) { + ids := make([]graph.ID, len(nodes)) + for index, node := range nodes { + ids[index] = s.database.nextID + s.database.nextID++ + s.nodes = append(s.nodes, graph.NewNode(ids[index], node.Properties, node.Kinds...)) + } + return ids, nil +} + +func (s *loadTestBatch) CreateRelationshipByIDs(startID, endID graph.ID, kind graph.Kind, properties *graph.Properties) error { + s.relationships = append(s.relationships, graph.NewRelationship(0, startID, endID, properties, kind)) + return nil +} diff --git a/ret/metrics/builder.go b/ret/metrics/builder.go new file mode 100644 index 00000000..c3134aee --- /dev/null +++ b/ret/metrics/builder.go @@ -0,0 +1,121 @@ +package metrics + +import ( + "fmt" + "strconv" + "strings" + + "github.com/specterops/dawgs/ret/entity" +) + +// Builder incrementally aggregates graph metrics. Nodes must be observed before +// any relationship that references them. +type Builder struct { + nodeCount int64 + relationshipCount int64 + nodeKinds map[string]string + inbound map[string]int64 + outbound map[string]int64 + nodeSequences map[string]int64 + relationshipKinds map[string]int64 + endpointShapes map[string]int64 +} + +// NewBuilder creates an empty graph metrics builder. +func NewBuilder() *Builder { + return &Builder{ + nodeKinds: make(map[string]string), + inbound: make(map[string]int64), + outbound: make(map[string]int64), + nodeSequences: make(map[string]int64), + relationshipKinds: make(map[string]int64), + endpointShapes: make(map[string]int64), + } +} + +// ObserveNode adds one normalized graph node to the aggregate. +func (s *Builder) ObserveNode(node entity.Node) error { + if err := node.Validate(); err != nil { + return fmt.Errorf("metrics node observation: %w", err) + } + if _, found := s.nodeKinds[node.SourceID]; found { + return fmt.Errorf("metrics node observation has duplicate source ID %q", node.SourceID) + } + + kinds := OrderedKindsKey(node.Kinds) + s.nodeKinds[node.SourceID] = kinds + s.inbound[node.SourceID] = 0 + s.outbound[node.SourceID] = 0 + s.nodeSequences[kinds]++ + s.nodeCount++ + + return nil +} + +// ObserveRelationship adds one normalized graph relationship to the aggregate. +func (s *Builder) ObserveRelationship(relationship entity.Relationship) error { + if err := relationship.Validate(); err != nil { + return fmt.Errorf("metrics relationship observation: %w", err) + } + + startKinds, startOK := s.nodeKinds[relationship.StartID] + endKinds, endOK := s.nodeKinds[relationship.EndID] + if !startOK || !endOK { + return fmt.Errorf("relationship %q references missing endpoint", relationship.SourceID) + } + + s.relationshipKinds[relationship.Kind]++ + s.outbound[relationship.StartID]++ + s.inbound[relationship.EndID]++ + s.endpointShapes[endpointShapeKey(startKinds, relationship.Kind, endKinds)]++ + s.relationshipCount++ + + return nil +} + +// Finalize returns an independent graph metrics snapshot. +func (s *Builder) Finalize() GraphMetrics { + value := GraphMetrics{ + NodeCount: s.nodeCount, + RelationshipCount: s.relationshipCount, + NodeKindSequences: cloneHistogram(s.nodeSequences), + RelationshipKinds: cloneHistogram(s.relationshipKinds), + InboundDegreeHistogram: make(map[string]int64), + OutboundDegreeHistogram: make(map[string]int64), + EndpointShapeHistogram: cloneHistogram(s.endpointShapes), + } + + for sourceID := range s.nodeKinds { + value.InboundDegreeHistogram[strconv.FormatInt(s.inbound[sourceID], 10)]++ + value.OutboundDegreeHistogram[strconv.FormatInt(s.outbound[sourceID], 10)]++ + } + + value.Fingerprint = fingerprint(value) + + return value +} + +// OrderedKindsKey encodes every kind segment in order, preserving duplicates. +func OrderedKindsKey(kinds []string) string { + var key strings.Builder + for _, kind := range kinds { + key.WriteString(strconv.Itoa(len(kind))) + key.WriteByte(':') + key.WriteString(kind) + } + + return key.String() +} + +func endpointShapeKey(startKinds, relationshipKind, endKinds string) string { + return OrderedKindsKey([]string{startKinds, relationshipKind, endKinds}) +} + +func cloneHistogram(source map[string]int64) map[string]int64 { + clone := make(map[string]int64, len(source)) + for key, count := range source { + clone[key] = count + } + + return clone +} diff --git a/ret/metrics/builder_test.go b/ret/metrics/builder_test.go new file mode 100644 index 00000000..6428c8aa --- /dev/null +++ b/ret/metrics/builder_test.go @@ -0,0 +1,137 @@ +package metrics_test + +import ( + "encoding/json" + "testing" + + "github.com/specterops/dawgs/ret/entity" + "github.com/specterops/dawgs/ret/metrics" + "github.com/stretchr/testify/require" +) + +func TestOrderedKindsKeyPreservesOrderMultiplicityAndBoundaries(t *testing.T) { + first := metrics.OrderedKindsKey([]string{"A", "B", "A"}) + second := metrics.OrderedKindsKey([]string{"A", "A", "B"}) + + require.Equal(t, "1:A1:B1:A", first) + require.Equal(t, "1:A1:A1:B", second) + require.NotEqual(t, first, second) + require.NotEqual(t, metrics.OrderedKindsKey([]string{"12", "3"}), metrics.OrderedKindsKey([]string{"1", "23"})) +} + +func TestOrderedKindsRemainDistinct(t *testing.T) { + builder := metrics.NewBuilder() + require.NoError(t, builder.ObserveNode(entity.Node{SourceID: "1", Kinds: []string{"A", "B", "A"}})) + require.NoError(t, builder.ObserveNode(entity.Node{SourceID: "2", Kinds: []string{"A", "A", "B"}})) + + got := builder.Finalize() + + require.EqualValues(t, 1, got.NodeKindSequences[metrics.OrderedKindsKey([]string{"A", "B", "A"})]) + require.EqualValues(t, 1, got.NodeKindSequences[metrics.OrderedKindsKey([]string{"A", "A", "B"})]) +} + +func TestBuilderAggregatesGraphShapeIncludingIsolatedNodeDegrees(t *testing.T) { + builder := metrics.NewBuilder() + for _, node := range []entity.Node{ + {SourceID: "user", Kinds: []string{"User", "Person"}}, + {SourceID: "group", Kinds: []string{"Group"}}, + {SourceID: "isolated"}, + } { + require.NoError(t, builder.ObserveNode(node)) + } + for _, relationship := range []entity.Relationship{ + {SourceID: "member", StartID: "user", EndID: "group", Kind: "MEMBER_OF"}, + {SourceID: "admin", StartID: "group", EndID: "user", Kind: "ADMIN_TO"}, + } { + require.NoError(t, builder.ObserveRelationship(relationship)) + } + + got := builder.Finalize() + + user := metrics.OrderedKindsKey([]string{"User", "Person"}) + group := metrics.OrderedKindsKey([]string{"Group"}) + empty := metrics.OrderedKindsKey(nil) + require.EqualValues(t, 3, got.NodeCount) + require.EqualValues(t, 2, got.RelationshipCount) + require.Equal(t, map[string]int64{user: 1, group: 1, empty: 1}, got.NodeKindSequences) + require.Equal(t, map[string]int64{"MEMBER_OF": 1, "ADMIN_TO": 1}, got.RelationshipKinds) + require.Equal(t, map[string]int64{"0": 1, "1": 2}, got.InboundDegreeHistogram) + require.Equal(t, map[string]int64{"0": 1, "1": 2}, got.OutboundDegreeHistogram) + require.Equal(t, map[string]int64{ + metrics.OrderedKindsKey([]string{user, "MEMBER_OF", group}): 1, + metrics.OrderedKindsKey([]string{group, "ADMIN_TO", user}): 1, + }, got.EndpointShapeHistogram) + require.Regexp(t, `^sha256:[0-9a-f]{64}$`, got.Fingerprint) +} + +func TestBuilderRejectsInvalidNodesDuplicateIDsAndMissingRelationshipEndpoints(t *testing.T) { + builder := metrics.NewBuilder() + require.Error(t, builder.ObserveNode(entity.Node{})) + require.NoError(t, builder.ObserveNode(entity.Node{SourceID: "present"})) + require.ErrorContains(t, builder.ObserveNode(entity.Node{SourceID: "present"}), "duplicate") + require.ErrorContains(t, builder.ObserveRelationship(entity.Relationship{StartID: "present", EndID: "missing", Kind: "KNOWS"}), "missing endpoint") + require.Error(t, builder.ObserveRelationship(entity.Relationship{StartID: "present", EndID: "present"})) +} + +func TestFingerprintIsStableForEquivalentGraphsRegardlessOfObservationOrder(t *testing.T) { + first := buildGraph(t, []entity.Node{ + {SourceID: "one", Kinds: []string{"User", "Person"}}, + {SourceID: "two", Kinds: []string{"Group"}}, + }, []entity.Relationship{ + {StartID: "one", EndID: "two", Kind: "MEMBER_OF"}, + {StartID: "two", EndID: "one", Kind: "ADMIN_TO"}, + }) + second := buildGraph(t, []entity.Node{ + {SourceID: "two", Kinds: []string{"Group"}}, + {SourceID: "one", Kinds: []string{"User", "Person"}}, + }, []entity.Relationship{ + {StartID: "two", EndID: "one", Kind: "ADMIN_TO"}, + {StartID: "one", EndID: "two", Kind: "MEMBER_OF"}, + }) + + require.Equal(t, first, second) +} + +func TestMetricsExcludeSourceIDsAndProperties(t *testing.T) { + first := buildGraph(t, []entity.Node{{ + SourceID: "node-secret-one", + Kinds: []string{"User"}, + Properties: map[string]any{"email": "ada@example.test"}, + }}, nil) + second := buildGraph(t, []entity.Node{{ + SourceID: "node-secret-two", + Kinds: []string{"User"}, + Properties: map[string]any{"email": "grace@example.test"}, + }}, nil) + + payload, err := json.Marshal(first) + require.NoError(t, err) + require.Equal(t, first.Fingerprint, second.Fingerprint) + require.NotContains(t, string(payload), "node-secret-one") + require.NotContains(t, string(payload), "ada@example.test") +} + +func TestFinalizeReturnsIndependentMetrics(t *testing.T) { + builder := metrics.NewBuilder() + require.NoError(t, builder.ObserveNode(entity.Node{SourceID: "one", Kinds: []string{"User"}})) + + first := builder.Finalize() + originalFingerprint := first.Fingerprint + first.NodeKindSequences[metrics.OrderedKindsKey([]string{"User"})] = 99 + second := builder.Finalize() + + require.EqualValues(t, 1, second.NodeKindSequences[metrics.OrderedKindsKey([]string{"User"})]) + require.Equal(t, originalFingerprint, second.Fingerprint) +} + +func buildGraph(t *testing.T, nodes []entity.Node, relationships []entity.Relationship) metrics.GraphMetrics { + t.Helper() + builder := metrics.NewBuilder() + for _, node := range nodes { + require.NoError(t, builder.ObserveNode(node)) + } + for _, relationship := range relationships { + require.NoError(t, builder.ObserveRelationship(relationship)) + } + return builder.Finalize() +} diff --git a/ret/metrics/compare.go b/ret/metrics/compare.go new file mode 100644 index 00000000..ccfb3bb1 --- /dev/null +++ b/ret/metrics/compare.go @@ -0,0 +1,58 @@ +package metrics + +import ( + "fmt" + "sort" + "strings" +) + +// Compare returns all deterministic differences between two graph metrics. +func Compare(expected, actual GraphMetrics) error { + differences := make([]string, 0) + if expected.NodeCount != actual.NodeCount { + differences = append(differences, fmt.Sprintf("node count: expected %d, actual %d", expected.NodeCount, actual.NodeCount)) + } + if expected.RelationshipCount != actual.RelationshipCount { + differences = append(differences, fmt.Sprintf("relationship count: expected %d, actual %d", expected.RelationshipCount, actual.RelationshipCount)) + } + + differences = append(differences, compareHistogram("node kind sequences", expected.NodeKindSequences, actual.NodeKindSequences)...) + differences = append(differences, compareHistogram("relationship kinds", expected.RelationshipKinds, actual.RelationshipKinds)...) + differences = append(differences, compareHistogram("inbound degree histogram", expected.InboundDegreeHistogram, actual.InboundDegreeHistogram)...) + differences = append(differences, compareHistogram("outbound degree histogram", expected.OutboundDegreeHistogram, actual.OutboundDegreeHistogram)...) + differences = append(differences, compareHistogram("endpoint shape histogram", expected.EndpointShapeHistogram, actual.EndpointShapeHistogram)...) + + if expected.Fingerprint != actual.Fingerprint { + differences = append(differences, fmt.Sprintf("fingerprint: expected %q, actual %q", expected.Fingerprint, actual.Fingerprint)) + } + if len(differences) == 0 { + return nil + } + + return fmt.Errorf("graph metrics differ:\n%s", strings.Join(differences, "\n")) +} + +func compareHistogram(name string, expected, actual map[string]int64) []string { + keys := make(map[string]struct{}, len(expected)+len(actual)) + for key := range expected { + keys[key] = struct{}{} + } + for key := range actual { + keys[key] = struct{}{} + } + + sortedKeys := make([]string, 0, len(keys)) + for key := range keys { + sortedKeys = append(sortedKeys, key) + } + sort.Strings(sortedKeys) + + differences := make([]string, 0) + for _, key := range sortedKeys { + if expected[key] != actual[key] { + differences = append(differences, fmt.Sprintf("%s[%q]: expected %d, actual %d", name, key, expected[key], actual[key])) + } + } + + return differences +} diff --git a/ret/metrics/compare_test.go b/ret/metrics/compare_test.go new file mode 100644 index 00000000..903b56e2 --- /dev/null +++ b/ret/metrics/compare_test.go @@ -0,0 +1,54 @@ +package metrics_test + +import ( + "testing" + + "github.com/specterops/dawgs/ret/metrics" + "github.com/stretchr/testify/require" +) + +func TestCompareReportsEveryChangedCategoryDeterministically(t *testing.T) { + expected := metrics.GraphMetrics{ + NodeCount: 1, + RelationshipCount: 2, + NodeKindSequences: map[string]int64{"a": 1}, + RelationshipKinds: map[string]int64{"A": 2}, + InboundDegreeHistogram: map[string]int64{"0": 1}, + OutboundDegreeHistogram: map[string]int64{"0": 1}, + EndpointShapeHistogram: map[string]int64{"shape-a": 2}, + Fingerprint: "a", + } + actual := metrics.GraphMetrics{ + NodeCount: 3, + RelationshipCount: 4, + NodeKindSequences: map[string]int64{"z": 3}, + RelationshipKinds: map[string]int64{"Z": 4}, + InboundDegreeHistogram: map[string]int64{"1": 3}, + OutboundDegreeHistogram: map[string]int64{"1": 3}, + EndpointShapeHistogram: map[string]int64{"shape-z": 4}, + Fingerprint: "b", + } + + err := metrics.Compare(expected, actual) + + require.EqualError(t, err, "graph metrics differ:\n"+ + "node count: expected 1, actual 3\n"+ + "relationship count: expected 2, actual 4\n"+ + "node kind sequences[\"a\"]: expected 1, actual 0\n"+ + "node kind sequences[\"z\"]: expected 0, actual 3\n"+ + "relationship kinds[\"A\"]: expected 2, actual 0\n"+ + "relationship kinds[\"Z\"]: expected 0, actual 4\n"+ + "inbound degree histogram[\"0\"]: expected 1, actual 0\n"+ + "inbound degree histogram[\"1\"]: expected 0, actual 3\n"+ + "outbound degree histogram[\"0\"]: expected 1, actual 0\n"+ + "outbound degree histogram[\"1\"]: expected 0, actual 3\n"+ + "endpoint shape histogram[\"shape-a\"]: expected 2, actual 0\n"+ + "endpoint shape histogram[\"shape-z\"]: expected 0, actual 4\n"+ + "fingerprint: expected \"a\", actual \"b\"") +} + +func TestCompareReturnsNilForEqualMetrics(t *testing.T) { + value := metrics.GraphMetrics{NodeCount: 1, Fingerprint: "sha256:abc"} + + require.NoError(t, metrics.Compare(value, value)) +} diff --git a/ret/metrics/fingerprint.go b/ret/metrics/fingerprint.go new file mode 100644 index 00000000..58294253 --- /dev/null +++ b/ret/metrics/fingerprint.go @@ -0,0 +1,46 @@ +package metrics + +import ( + "crypto/sha256" + "encoding/hex" + "encoding/json" + "fmt" + "sort" +) + +func fingerprint(value GraphMetrics) string { + payload, err := json.Marshal(canonicalize(value)) + if err != nil { + panic(fmt.Sprintf("canonical graph metrics cannot be marshaled: %v", err)) + } + + sum := sha256.Sum256(payload) + return "sha256:" + hex.EncodeToString(sum[:]) +} + +func canonicalize(value GraphMetrics) canonicalGraphMetrics { + return canonicalGraphMetrics{ + NodeCount: value.NodeCount, + RelationshipCount: value.RelationshipCount, + NodeKindSequences: canonicalHistogram(value.NodeKindSequences), + RelationshipKinds: canonicalHistogram(value.RelationshipKinds), + InboundDegreeHistogram: canonicalHistogram(value.InboundDegreeHistogram), + OutboundDegreeHistogram: canonicalHistogram(value.OutboundDegreeHistogram), + EndpointShapeHistogram: canonicalHistogram(value.EndpointShapeHistogram), + } +} + +func canonicalHistogram(histogram map[string]int64) []metricHistogramEntry { + keys := make([]string, 0, len(histogram)) + for key := range histogram { + keys = append(keys, key) + } + sort.Strings(keys) + + entries := make([]metricHistogramEntry, 0, len(keys)) + for _, key := range keys { + entries = append(entries, metricHistogramEntry{Key: key, Count: histogram[key]}) + } + + return entries +} diff --git a/ret/metrics/model.go b/ret/metrics/model.go new file mode 100644 index 00000000..e9c600c1 --- /dev/null +++ b/ret/metrics/model.go @@ -0,0 +1,30 @@ +// Package metrics aggregates format-neutral graph shape metrics. +package metrics + +// GraphMetrics describes the observable graph shape without source identifiers +// or entity properties. +type GraphMetrics struct { + NodeCount int64 `json:"node_count"` + RelationshipCount int64 `json:"relationship_count"` + NodeKindSequences map[string]int64 `json:"node_kind_sequences"` + RelationshipKinds map[string]int64 `json:"relationship_kinds"` + InboundDegreeHistogram map[string]int64 `json:"inbound_degree_histogram"` + OutboundDegreeHistogram map[string]int64 `json:"outbound_degree_histogram"` + EndpointShapeHistogram map[string]int64 `json:"endpoint_shape_histogram"` + Fingerprint string `json:"fingerprint"` +} + +type metricHistogramEntry struct { + Key string `json:"key"` + Count int64 `json:"count"` +} + +type canonicalGraphMetrics struct { + NodeCount int64 `json:"node_count"` + RelationshipCount int64 `json:"relationship_count"` + NodeKindSequences []metricHistogramEntry `json:"node_kind_sequences"` + RelationshipKinds []metricHistogramEntry `json:"relationship_kinds"` + InboundDegreeHistogram []metricHistogramEntry `json:"inbound_degree_histogram"` + OutboundDegreeHistogram []metricHistogramEntry `json:"outbound_degree_histogram"` + EndpointShapeHistogram []metricHistogramEntry `json:"endpoint_shape_histogram"` +} diff --git a/ret/observe/event.go b/ret/observe/event.go new file mode 100644 index 00000000..cba03955 --- /dev/null +++ b/ret/observe/event.go @@ -0,0 +1,86 @@ +package observe + +import "time" + +// Event is an observation emitted during a graph export operation. +type Event interface { + isEvent() +} + +// OperationStarted indicates an operation has begun. +type OperationStarted struct{ Operation string } + +func (OperationStarted) isEvent() {} + +// OperationCompleted indicates an operation has completed. +type OperationCompleted struct { + Operation string + Duration time.Duration + Err error +} + +func (OperationCompleted) isEvent() {} + +// GraphStarted indicates processing has begun for a graph. +type GraphStarted struct{ Operation, Graph string } + +func (GraphStarted) isEvent() {} + +// GraphCompleted indicates processing has completed for a graph. +type GraphCompleted struct { + Operation, Graph string + Nodes, Relationships int64 + Duration time.Duration +} + +func (GraphCompleted) isEvent() {} + +// PhaseStarted indicates a graph-processing phase has begun. +type PhaseStarted struct { + Operation, Graph, Phase string + Completed, Total int64 +} + +func (PhaseStarted) isEvent() {} + +// PhaseProgress reports progress through a graph-processing phase. +type PhaseProgress struct { + Operation, Graph, Phase string + Completed, Total int64 +} + +func (PhaseProgress) isEvent() {} + +// PhaseCompleted indicates a graph-processing phase has completed. +type PhaseCompleted struct { + Operation, Graph, Phase string + Completed int64 + Duration time.Duration +} + +func (PhaseCompleted) isEvent() {} + +// ShardCommitted indicates output files for a shard have been committed. +type ShardCommitted struct { + Graph, EntityType, JSONLPath, ParquetPath string + Index int + Count, JSONLBytes, ParquetBytes int64 +} + +func (ShardCommitted) isEvent() {} + +// ArtifactVerified indicates an output artifact has been verified. +type ArtifactVerified struct { + Graph, EntityType, Format, Path string + Count, Bytes int64 +} + +func (ArtifactVerified) isEvent() {} + +// ArchiveEntryProcessed indicates an archive entry has been processed. +type ArchiveEntryProcessed struct { + Operation, Path string + Size int64 +} + +func (ArchiveEntryProcessed) isEvent() {} diff --git a/ret/observe/observer.go b/ret/observe/observer.go new file mode 100644 index 00000000..d054aa91 --- /dev/null +++ b/ret/observe/observer.go @@ -0,0 +1,23 @@ +package observe + +import "context" + +// Observer receives operation observation events. +type Observer interface { + Observe(context.Context, Event) +} + +// ObserverFunc adapts a function to the Observer interface. +type ObserverFunc func(context.Context, Event) + +// Observe delivers an event to s. +func (s ObserverFunc) Observe(ctx context.Context, event Event) { + s(ctx, event) +} + +// Emit delivers an event when an observer is configured. +func Emit(ctx context.Context, observer Observer, event Event) { + if observer != nil { + observer.Observe(ctx, event) + } +} diff --git a/ret/observe/observer_test.go b/ret/observe/observer_test.go new file mode 100644 index 00000000..94deb7af --- /dev/null +++ b/ret/observe/observer_test.go @@ -0,0 +1,24 @@ +package observe_test + +import ( + "context" + "testing" + + "github.com/specterops/dawgs/ret/observe" + "github.com/stretchr/testify/require" +) + +func TestEmitAllowsNilObserver(t *testing.T) { + require.NotPanics(t, func() { + observe.Emit(context.Background(), nil, observe.OperationStarted{Operation: "dump"}) + }) +} + +func TestObserverFuncReceivesTypedEvent(t *testing.T) { + var got observe.Event + observer := observe.ObserverFunc(func(_ context.Context, event observe.Event) { got = event }) + + observe.Emit(context.Background(), observer, observe.GraphStarted{Operation: "dump", Graph: "asset"}) + + require.Equal(t, observe.GraphStarted{Operation: "dump", Graph: "asset"}, got) +} diff --git a/ret/parquet/artifact.go b/ret/parquet/artifact.go new file mode 100644 index 00000000..87804ffb --- /dev/null +++ b/ret/parquet/artifact.go @@ -0,0 +1,28 @@ +package parquet + +import ( + "crypto/sha256" + "encoding/hex" + "fmt" +) + +type Artifact struct { + SchemaVersion string + SHA256 string + Count int64 + StoredBytes int64 +} + +func (s Artifact) validate() error { + if s.SchemaVersion != SchemaVersion { + return fmt.Errorf("unsupported Parquet artifact schema %q", s.SchemaVersion) + } + if s.Count < 0 || s.StoredBytes < 0 { + return fmt.Errorf("Parquet artifact size and count must be non-negative") + } + decoded, err := hex.DecodeString(s.SHA256) + if err != nil || len(decoded) != sha256.Size { + return fmt.Errorf("Parquet artifact SHA-256 is invalid: %q", s.SHA256) + } + return nil +} diff --git a/ret/parquet/config.go b/ret/parquet/config.go new file mode 100644 index 00000000..91f58030 --- /dev/null +++ b/ret/parquet/config.go @@ -0,0 +1,15 @@ +// Package parquet writes and verifies concrete Parquet graph artifacts. +package parquet + +const ( + // SchemaVersion identifies the Parquet artifact metadata and row layout. + SchemaVersion = "ret-parquet-v1" +) + +// Config controls creation of Parquet artifacts. +type Config struct{} + +// Validate verifies that the configuration can be used to write an artifact. +func (Config) Validate() error { + return nil +} diff --git a/ret/parquet/fixture_test.go b/ret/parquet/fixture_test.go new file mode 100644 index 00000000..d948f867 --- /dev/null +++ b/ret/parquet/fixture_test.go @@ -0,0 +1,181 @@ +package parquet + +import ( + "errors" + "fmt" + "os" + "path/filepath" + "strings" + + "github.com/specterops/dawgs/ret/entity" +) + +type nodeFixtureArtifact struct { + SchemaVersion, Path, SHA256 string + Count, StoredBytes int64 +} + +type relationshipFixtureArtifact struct { + SchemaVersion, Path, SHA256 string + Count, StoredBytes int64 +} + +var afterVerifiedSnapshotForTest func() + +func writeNodesFixture(temporary, relative string, config Config, values []entity.Node) (nodeFixtureArtifact, error) { + path, err := fixturePath(temporary, relative) + if err != nil { + return nodeFixtureArtifact{}, err + } + file, err := os.OpenFile(temporary, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0o600) + if err != nil { + return nodeFixtureArtifact{}, err + } + writer, err := NewNodeWriter(file, config) + if err != nil { + return nodeFixtureArtifact{}, errors.Join(err, file.Close()) + } + pushErr := writer.Push(values) + closeWriterErr := writer.Close() + artifact, resultErr := writer.Result() + closeFileErr := file.Close() + if err := errors.Join(pushErr, closeWriterErr, resultErr, closeFileErr); err != nil { + return nodeFixtureArtifact{}, err + } + return nodeFixtureArtifact{artifact.SchemaVersion, path, artifact.SHA256, artifact.Count, artifact.StoredBytes}, nil +} + +func writeRelationshipsFixture(temporary, relative string, config Config, values []entity.Relationship) (relationshipFixtureArtifact, error) { + path, err := fixturePath(temporary, relative) + if err != nil { + return relationshipFixtureArtifact{}, err + } + file, err := os.OpenFile(temporary, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0o600) + if err != nil { + return relationshipFixtureArtifact{}, err + } + writer, err := NewRelationshipWriter(file, config) + if err != nil { + return relationshipFixtureArtifact{}, errors.Join(err, file.Close()) + } + pushErr := writer.Push(values) + closeWriterErr := writer.Close() + artifact, resultErr := writer.Result() + closeFileErr := file.Close() + if err := errors.Join(pushErr, closeWriterErr, resultErr, closeFileErr); err != nil { + return relationshipFixtureArtifact{}, err + } + return relationshipFixtureArtifact{artifact.SchemaVersion, path, artifact.SHA256, artifact.Count, artifact.StoredBytes}, nil +} + +func readNodesFixture(root string, fixture nodeFixtureArtifact, visit func(entity.Node) error) error { + file, size, err := openFixture(root, fixture.Path) + if err != nil { + return err + } + reader, err := NewNodeReader(file, size, Artifact{fixture.SchemaVersion, fixture.SHA256, fixture.Count, fixture.StoredBytes}) + if err != nil { + return errors.Join(err, file.Close()) + } + if afterVerifiedSnapshotForTest != nil { + afterVerifiedSnapshotForTest() + } + values, readErr := drainFixtureReader(&reader) + resultErr := reader.Result() + closeReaderErr := reader.Close() + closeFileErr := file.Close() + if err := errors.Join(readErr, resultErr, closeReaderErr, closeFileErr); err != nil { + return err + } + for _, value := range values { + if visit != nil { + if err := visit(value); err != nil { + return err + } + } + } + return nil +} + +func readRelationshipsFixture(root string, fixture relationshipFixtureArtifact, visit func(entity.Relationship) error) error { + file, size, err := openFixture(root, fixture.Path) + if err != nil { + return err + } + reader, err := NewRelationshipReader(file, size, Artifact{fixture.SchemaVersion, fixture.SHA256, fixture.Count, fixture.StoredBytes}) + if err != nil { + return errors.Join(err, file.Close()) + } + values, readErr := drainFixtureReader(&reader) + resultErr := reader.Result() + closeReaderErr := reader.Close() + closeFileErr := file.Close() + if err := errors.Join(readErr, resultErr, closeReaderErr, closeFileErr); err != nil { + return err + } + for _, value := range values { + if visit != nil { + if err := visit(value); err != nil { + return err + } + } + } + return nil +} + +type fixturePullReader[E any] interface { + Pull(int) ([]E, error) + Done() bool +} + +func drainFixtureReader[E any](reader fixturePullReader[E]) ([]E, error) { + var values []E + for !reader.Done() { + batch, err := reader.Pull(256) + if err != nil { + return nil, err + } + values = append(values, batch...) + } + return values, nil +} + +func fixturePath(temporary, relative string) (string, error) { + if !filepath.IsAbs(temporary) { + return "", fmt.Errorf("Parquet temporary path must be absolute: %q", temporary) + } + if relative == "" || strings.Contains(relative, "\\") { + return "", fmt.Errorf("Parquet artifact path must be a slash-separated relative file: %q", relative) + } + clean := filepath.Clean(filepath.FromSlash(relative)) + if filepath.IsAbs(clean) || clean == "." || clean == ".." || strings.HasPrefix(clean, ".."+string(filepath.Separator)) || filepath.ToSlash(clean) != relative { + return "", fmt.Errorf("Parquet artifact path escapes collection: %q", relative) + } + return relative, nil +} + +func openFixture(root, relative string) (*os.File, int64, error) { + if _, err := fixturePath(filepath.Join(root, "fixture.tmp"), relative); err != nil { + return nil, 0, err + } + path := root + for _, component := range strings.Split(filepath.FromSlash(relative), string(filepath.Separator)) { + path = filepath.Join(path, component) + info, err := os.Lstat(path) + if err != nil { + return nil, 0, err + } + if info.Mode()&os.ModeSymlink != 0 { + return nil, 0, fmt.Errorf("Parquet artifact path contains symlink component %q", component) + } + } + file, err := os.Open(path) + if err != nil { + return nil, 0, err + } + info, err := file.Stat() + if err != nil { + return nil, 0, errors.Join(err, file.Close()) + } + return file, info.Size(), nil +} diff --git a/ret/parquet/lifecycle.go b/ret/parquet/lifecycle.go new file mode 100644 index 00000000..72a86cbc --- /dev/null +++ b/ret/parquet/lifecycle.go @@ -0,0 +1,303 @@ +package parquet + +import ( + "crypto/sha256" + "encoding/hex" + "errors" + "fmt" + "hash" + "io" + + parquetgo "github.com/parquet-go/parquet-go" + "github.com/parquet-go/parquet-go/compress/zstd" + "github.com/specterops/dawgs/ret/entity" +) + +var ( + ErrWriterNotOpen = errors.New("writer is not open") + ErrWriterNotClosed = errors.New("writer is not closed") + ErrWriterFailed = errors.New("writer failed") + ErrReaderNotOpen = errors.New("reader is not open") + ErrReaderNotDone = errors.New("reader is not done reading") + ErrReaderFailed = errors.New("reader failed") + ErrInvalidLimit = errors.New("pull limit must be positive") +) + +type lifecycleState uint8 + +const ( + stateOpen lifecycleState = iota + stateFailed + stateClosed +) + +type row interface { + NodeRow | RelationshipRow +} + +type countingWriter struct { + writer io.Writer + count int64 +} + +func (s *countingWriter) Write(value []byte) (int, error) { + written, err := s.writer.Write(value) + s.count += int64(written) + return written, err +} + +func NewNodeWriter(output io.Writer, config Config) (EntityWriter[entity.Node, NodeRow], error) { + return newEntityWriter(output, config, func(value entity.Node) (NodeRow, error) { + if err := value.Validate(); err != nil { + return NodeRow{}, err + } + if err := validateProperties(value.Properties); err != nil { + return NodeRow{}, fmt.Errorf("validate properties: %w", err) + } + return nodeRow(value), nil + }) +} + +func NewRelationshipWriter(output io.Writer, config Config) (EntityWriter[entity.Relationship, RelationshipRow], error) { + return newEntityWriter(output, config, func(value entity.Relationship) (RelationshipRow, error) { + if value.SourceID == "" { + return RelationshipRow{}, fmt.Errorf("relationship source ID is required") + } + if err := value.Validate(); err != nil { + return RelationshipRow{}, err + } + if err := validateProperties(value.Properties); err != nil { + return RelationshipRow{}, fmt.Errorf("validate properties: %w", err) + } + return relationshipRow(value), nil + }) +} + +type EntityWriter[E entity.Entity, R row] struct { + hasher hash.Hash + stored *countingWriter + writer *parquetgo.GenericWriter[R] + entityToRow func(E) (R, error) + recordCount int64 + state lifecycleState + resourceDone bool + failure error +} + +func newEntityWriter[E entity.Entity, R row](output io.Writer, config Config, convert func(E) (R, error)) (EntityWriter[E, R], error) { + if err := config.Validate(); err != nil { + return EntityWriter[E, R]{}, err + } + hasher := sha256.New() + stored := &countingWriter{writer: io.MultiWriter(output, hasher)} + return EntityWriter[E, R]{ + hasher: hasher, + stored: stored, + writer: parquetgo.NewGenericWriter[R](stored, parquetgo.Compression(&zstd.Codec{})), + entityToRow: convert, + state: stateOpen, + }, nil +} + +func (s *EntityWriter[E, R]) Push(entities []E) error { + switch s.state { + case stateClosed: + return ErrWriterNotOpen + case stateFailed: + return errors.Join(ErrWriterFailed, s.failure) + } + rows := make([]R, len(entities)) + for index, value := range entities { + row, err := s.entityToRow(value) + if err != nil { + return s.fail(fmt.Errorf("validate Parquet record %d: %w", s.recordCount+int64(index)+1, err)) + } + rows[index] = row + } + written, err := s.writer.Write(rows) + if err != nil { + return s.fail(fmt.Errorf("write Parquet rows: %w", err)) + } + if written != len(rows) { + return s.fail(fmt.Errorf("write Parquet rows: wrote %d, want %d", written, len(rows))) + } + s.recordCount += int64(written) + return nil +} + +func (s *EntityWriter[E, R]) Close() error { + if s.resourceDone { + if s.state == stateFailed { + return errors.Join(ErrWriterFailed, s.failure) + } + return nil + } + s.resourceDone = true + if err := s.writer.Close(); err != nil { + s.fail(fmt.Errorf("finish Parquet file: %w", err)) + } + if s.state == stateFailed { + return errors.Join(ErrWriterFailed, s.failure) + } + s.state = stateClosed + return nil +} + +func (s *EntityWriter[E, R]) Result() (Artifact, error) { + switch s.state { + case stateOpen: + return Artifact{}, ErrWriterNotClosed + case stateFailed: + return Artifact{}, errors.Join(ErrWriterFailed, s.failure) + } + return Artifact{ + SchemaVersion: SchemaVersion, + SHA256: hex.EncodeToString(s.hasher.Sum(nil)), + Count: s.recordCount, + StoredBytes: s.stored.count, + }, nil +} + +func (s *EntityWriter[E, R]) fail(err error) error { + if s.failure == nil { + s.failure = err + } else { + s.failure = errors.Join(s.failure, err) + } + s.state = stateFailed + return err +} + +func NewNodeReader(input io.ReaderAt, size int64, artifact Artifact) (Reader[entity.Node, NodeRow], error) { + return newEntityReader(input, size, artifact, func(value NodeRow) (entity.Node, error) { + return value.entity() + }) +} + +func NewRelationshipReader(input io.ReaderAt, size int64, artifact Artifact) (Reader[entity.Relationship, RelationshipRow], error) { + return newEntityReader(input, size, artifact, func(value RelationshipRow) (entity.Relationship, error) { + return value.entity() + }) +} + +type Reader[E entity.Entity, R row] struct { + artifact Artifact + reader *parquetgo.GenericReader[R] + rowToEntity func(R) (E, error) + recordCount int64 + state lifecycleState + resourceDone bool + failure error +} + +func newEntityReader[E entity.Entity, R row](input io.ReaderAt, size int64, artifact Artifact, convert func(R) (E, error)) (Reader[E, R], error) { + if err := artifact.validate(); err != nil { + return Reader[E, R]{}, err + } + if size != artifact.StoredBytes { + return Reader[E, R]{}, fmt.Errorf("Parquet stored size mismatch: got %d, want %d", size, artifact.StoredBytes) + } + hasher := sha256.New() + if _, err := io.Copy(hasher, io.NewSectionReader(input, 0, size)); err != nil { + return Reader[E, R]{}, fmt.Errorf("hash Parquet artifact: %w", err) + } + actual := hex.EncodeToString(hasher.Sum(nil)) + if actual != artifact.SHA256 { + return Reader[E, R]{}, fmt.Errorf("Parquet stored SHA-256 mismatch: got %s, want %s", actual, artifact.SHA256) + } + file, err := parquetgo.OpenFile(input, size) + if err != nil { + return Reader[E, R]{}, fmt.Errorf("open Parquet artifact: %w", err) + } + wantSchema := parquetgo.SchemaOf(new(R)) + if !parquetgo.EqualNodes(file.Schema(), wantSchema) { + return Reader[E, R]{}, fmt.Errorf("Parquet row schema does not match %s", wantSchema.Name()) + } + if file.NumRows() != artifact.Count { + return Reader[E, R]{}, fmt.Errorf("Parquet row count mismatch: got %d, want %d", file.NumRows(), artifact.Count) + } + return Reader[E, R]{ + artifact: artifact, + reader: parquetgo.NewGenericReader[R](file), + rowToEntity: convert, + state: stateOpen, + }, nil +} + +func (s *Reader[E, R]) Pull(limit int) ([]E, error) { + if limit <= 0 { + return nil, ErrInvalidLimit + } + if s.resourceDone { + return nil, ErrReaderNotOpen + } + switch s.state { + case stateClosed: + return nil, ErrReaderNotOpen + case stateFailed: + return nil, errors.Join(ErrReaderFailed, s.failure) + } + rows := make([]R, limit) + read, readErr := s.reader.Read(rows) + values := make([]E, 0, read) + for index := range read { + value, err := s.rowToEntity(rows[index]) + if err != nil { + return nil, s.fail(fmt.Errorf("validate Parquet row %d: %w", s.recordCount+int64(index)+1, err)) + } + values = append(values, value) + } + s.recordCount += int64(read) + if errors.Is(readErr, io.EOF) { + s.state = stateClosed + return values, nil + } + if readErr != nil { + return nil, s.fail(fmt.Errorf("read Parquet row %d: %w", s.recordCount+1, readErr)) + } + return values, nil +} + +func (s *Reader[E, R]) Done() bool { + return s.state != stateOpen +} + +func (s *Reader[E, R]) Close() error { + if s.resourceDone { + if s.state == stateFailed { + return errors.Join(ErrReaderFailed, s.failure) + } + return nil + } + s.resourceDone = true + if err := s.reader.Close(); err != nil { + s.fail(fmt.Errorf("close Parquet reader: %w", err)) + } + if s.state == stateFailed { + return errors.Join(ErrReaderFailed, s.failure) + } + return nil +} + +func (s *Reader[E, R]) Result() error { + switch s.state { + case stateOpen: + return ErrReaderNotDone + case stateFailed: + return errors.Join(ErrReaderFailed, s.failure) + } + if s.recordCount != s.artifact.Count { + return s.fail(fmt.Errorf("Parquet row count mismatch: got %d, want %d", s.recordCount, s.artifact.Count)) + } + return nil +} + +func (s *Reader[E, R]) fail(err error) error { + if s.failure == nil { + s.failure = err + } else { + s.failure = errors.Join(s.failure, err) + } + s.state = stateFailed + return err +} diff --git a/ret/parquet/lifecycle_test.go b/ret/parquet/lifecycle_test.go new file mode 100644 index 00000000..daec511b --- /dev/null +++ b/ret/parquet/lifecycle_test.go @@ -0,0 +1,150 @@ +package parquet + +import ( + "bytes" + "errors" + "testing" + + "github.com/specterops/dawgs/ret/entity" + "github.com/stretchr/testify/require" +) + +func TestStatefulNodeCodecRoundTrip(t *testing.T) { + var stored bytes.Buffer + writer, err := NewNodeWriter(&stored, Config{}) + require.NoError(t, err) + require.NoError(t, writer.Push([]entity.Node{ + {SourceID: "1", Kinds: []string{"User"}, Properties: map[string]any{"enabled": true}}, + {SourceID: "2", Kinds: []string{"Group"}}, + })) + _, err = writer.Result() + require.ErrorIs(t, err, ErrWriterNotClosed) + require.NoError(t, writer.Close()) + artifact, err := writer.Result() + require.NoError(t, err) + require.EqualValues(t, 2, artifact.Count) + require.EqualValues(t, stored.Len(), artifact.StoredBytes) + + reader, err := NewNodeReader(bytes.NewReader(stored.Bytes()), int64(stored.Len()), artifact) + require.NoError(t, err) + first, err := reader.Pull(1) + require.NoError(t, err) + require.Equal(t, "1", first[0].SourceID) + require.False(t, reader.Done()) + second, err := reader.Pull(4) + require.NoError(t, err) + require.Equal(t, []entity.Node{{SourceID: "2", Kinds: []string{"Group"}}}, second) + require.True(t, reader.Done()) + require.NoError(t, reader.Result()) + require.NoError(t, reader.Close()) +} + +func TestStatefulParquetReaderRejectsNonPositiveLimit(t *testing.T) { + stored, artifact := writeNodeArtifact(t, []entity.Node{{SourceID: "1"}}) + reader, err := NewNodeReader(bytes.NewReader(stored), int64(len(stored)), artifact) + require.NoError(t, err) + + _, err = reader.Pull(0) + require.ErrorIs(t, err, ErrInvalidLimit) + _, err = reader.Pull(-1) + require.ErrorIs(t, err, ErrInvalidLimit) +} + +func TestStatefulParquetReaderVerifiesStoredMetadataBeforeRows(t *testing.T) { + stored, artifact := writeNodeArtifact(t, []entity.Node{{SourceID: "1"}}) + + wrongSize := artifact + wrongSize.StoredBytes++ + _, err := NewNodeReader(bytes.NewReader(stored), int64(len(stored)), wrongSize) + require.ErrorContains(t, err, "stored size mismatch") + + wrongHash := artifact + wrongHash.SHA256 = "0000000000000000000000000000000000000000000000000000000000000000" + _, err = NewNodeReader(bytes.NewReader(stored), int64(len(stored)), wrongHash) + require.ErrorContains(t, err, "SHA-256 mismatch") +} + +func TestStatefulParquetWriterRejectsInvalidEntity(t *testing.T) { + var stored bytes.Buffer + writer, err := NewNodeWriter(&stored, Config{}) + require.NoError(t, err) + require.ErrorContains(t, writer.Push([]entity.Node{{}}), "source ID") + repeatedCloseErr := writer.Close() + require.ErrorIs(t, repeatedCloseErr, ErrWriterFailed) + require.ErrorContains(t, repeatedCloseErr, "source ID") + require.ErrorIs(t, writer.Close(), ErrWriterFailed) + _, err = writer.Result() + require.ErrorIs(t, err, ErrWriterFailed) +} + +func TestStatefulRelationshipCodecRoundTrip(t *testing.T) { + want := entity.Relationship{ + SourceID: "relationship-1", + StartID: "node-1", + EndID: "node-2", + Kind: "MemberOf", + Properties: map[string]any{"weight": int64(3)}, + } + var stored bytes.Buffer + writer, err := NewRelationshipWriter(&stored, Config{}) + require.NoError(t, err) + require.NoError(t, writer.Push([]entity.Relationship{want})) + require.NoError(t, writer.Close()) + artifact, err := writer.Result() + require.NoError(t, err) + + reader, err := NewRelationshipReader(bytes.NewReader(stored.Bytes()), int64(stored.Len()), artifact) + require.NoError(t, err) + values, err := reader.Pull(2) + require.NoError(t, err) + require.Equal(t, []entity.Relationship{want}, values) + require.True(t, reader.Done()) + require.NoError(t, reader.Result()) + require.NoError(t, reader.Close()) +} + +func TestStatefulParquetReaderCloseIsIdempotent(t *testing.T) { + stored, artifact := writeNodeArtifact(t, []entity.Node{{SourceID: "1"}}) + reader, err := NewNodeReader(bytes.NewReader(stored), int64(len(stored)), artifact) + require.NoError(t, err) + require.ErrorIs(t, reader.Result(), ErrReaderNotDone) + + require.NoError(t, reader.Close()) + require.NoError(t, reader.Close()) + _, err = reader.Pull(1) + require.ErrorIs(t, err, ErrReaderNotOpen) + require.ErrorIs(t, reader.Result(), ErrReaderNotDone) +} + +func TestStatefulParquetWriterLatchesUnderlyingFailure(t *testing.T) { + writer, err := NewNodeWriter(parquetFailingWriter{}, Config{}) + require.NoError(t, err) + pushErr := writer.Push([]entity.Node{{SourceID: "1"}}) + closeErr := writer.Close() + require.True(t, errors.Is(pushErr, errParquetWrite) || errors.Is(closeErr, errParquetWrite)) + repeatedCloseErr := writer.Close() + require.ErrorIs(t, repeatedCloseErr, ErrWriterFailed) + require.ErrorIs(t, repeatedCloseErr, errParquetWrite) + _, err = writer.Result() + require.ErrorIs(t, err, ErrWriterFailed) +} + +var errParquetWrite = errors.New("injected Parquet write failure") + +type parquetFailingWriter struct{} + +func (parquetFailingWriter) Write([]byte) (int, error) { + return 0, errParquetWrite +} + +func writeNodeArtifact(t *testing.T, nodes []entity.Node) ([]byte, Artifact) { + t.Helper() + var stored bytes.Buffer + writer, err := NewNodeWriter(&stored, Config{}) + require.NoError(t, err) + require.NoError(t, writer.Push(nodes)) + require.NoError(t, writer.Close()) + artifact, err := writer.Result() + require.NoError(t, err) + return stored.Bytes(), artifact +} diff --git a/ret/parquet/reader_test.go b/ret/parquet/reader_test.go new file mode 100644 index 00000000..8ffb7023 --- /dev/null +++ b/ret/parquet/reader_test.go @@ -0,0 +1,383 @@ +package parquet + +import ( + "bytes" + "crypto/sha256" + "encoding/hex" + "os" + "path/filepath" + "reflect" + "strings" + "testing" + + parquetgo "github.com/parquet-go/parquet-go" + "github.com/parquet-go/parquet-go/compress/zstd" + "github.com/specterops/dawgs/ret/entity" +) + +func TestVariantValuesAndRelationshipSourceIDSurviveRoundTrip(t *testing.T) { + want := entity.Relationship{ + SourceID: "relationship-7", + StartID: "node-1", + EndID: "node-2", + Kind: "MemberOf", + Properties: map[string]any{ + "null": nil, + "bool": true, + "integer": int64(42), + "float": 3.5, + "string": "value", + "list": []any{"first", int64(2), false}, + "object": map[string]any{ + "nested": "yes", + }, + }, + } + root := t.TempDir() + artifact, err := writeRelationshipsFixture(filepath.Join(root, "relationships.tmp"), "relationships.parquet", Config{}, []entity.Relationship{want}) + if err != nil { + t.Fatalf("write relationships: %v", err) + } + if err := os.Rename(filepath.Join(root, "relationships.tmp"), filepath.Join(root, artifact.Path)); err != nil { + t.Fatalf("install artifact: %v", err) + } + + var got []entity.Relationship + if err := readRelationshipsFixture(root, artifact, func(relationship entity.Relationship) error { + got = append(got, relationship) + return nil + }); err != nil { + t.Fatalf("read relationships: %v", err) + } + if len(got) != 1 || !reflect.DeepEqual(got[0], want) { + t.Fatalf("round trip relationships = %#v, want %#v", got, want) + } +} + +func TestReadNodesRejectsInvalidArtifactMetadataBeforeVisiting(t *testing.T) { + root, artifact := installedNodeArtifact(t, []entity.Node{{SourceID: "1"}}) + tests := []struct { + name string + mutate func(nodeFixtureArtifact) nodeFixtureArtifact + }{ + { + name: "schema version", + mutate: func(value nodeFixtureArtifact) nodeFixtureArtifact { + value.SchemaVersion = "wrong" + return value + }, + }, + { + name: "stored size", + mutate: func(value nodeFixtureArtifact) nodeFixtureArtifact { + value.StoredBytes++ + return value + }, + }, + { + name: "stored SHA", + mutate: func(value nodeFixtureArtifact) nodeFixtureArtifact { + value.SHA256 = strings.Repeat("0", sha256.Size*2) + return value + }, + }, + { + name: "row count", + mutate: func(value nodeFixtureArtifact) nodeFixtureArtifact { + value.Count++ + return value + }, + }, + { + name: "unsafe path", + mutate: func(value nodeFixtureArtifact) nodeFixtureArtifact { + value.Path = "../nodes.parquet" + return value + }, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + visited := 0 + err := readNodesFixture(root, test.mutate(artifact), func(entity.Node) error { + visited++ + return nil + }) + if err == nil { + t.Fatal("readNodesFixture unexpectedly succeeded") + } + if visited != 0 { + t.Fatalf("visited %d nodes before rejecting artifact", visited) + } + }) + } +} + +func TestReadNodesRejectsCorruptAndTruncatedFiles(t *testing.T) { + tests := []struct { + name string + mutate func([]byte) []byte + }{ + { + name: "corrupt", + mutate: func(contents []byte) []byte { + contents[len(contents)/2] ^= 0xff + return contents + }, + }, + { + name: "truncated", + mutate: func(contents []byte) []byte { + return contents[:len(contents)-1] + }, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + root, artifact := installedNodeArtifact(t, []entity.Node{{SourceID: "1"}}) + path := filepath.Join(root, artifact.Path) + contents, err := os.ReadFile(path) + if err != nil { + t.Fatalf("read artifact: %v", err) + } + if err := os.WriteFile(path, test.mutate(contents), 0o600); err != nil { + t.Fatalf("mutate artifact: %v", err) + } + + visited := 0 + if err := readNodesFixture(root, artifact, func(entity.Node) error { + visited++ + return nil + }); err == nil { + t.Fatal("readNodesFixture unexpectedly accepted damaged artifact") + } + if visited != 0 { + t.Fatalf("visited %d nodes from damaged artifact", visited) + } + }) + } +} + +func TestReadRelationshipsRequiresSourceIDBeforeAnyVisit(t *testing.T) { + root := t.TempDir() + contents := rawParquet(t, []RelationshipRow{ + {SourceID: "valid", StartID: "1", EndID: "2", Kind: "MemberOf", Properties: map[string]any{}}, + {StartID: "2", EndID: "3", Kind: "MemberOf", Properties: map[string]any{}}, + }) + path := "relationships.parquet" + writeContents(t, filepath.Join(root, path), contents) + hash := sha256.Sum256(contents) + artifact := relationshipFixtureArtifact{ + SchemaVersion: SchemaVersion, + Path: path, + SHA256: hex.EncodeToString(hash[:]), + Count: 2, + StoredBytes: int64(len(contents)), + } + + visited := 0 + err := readRelationshipsFixture(root, artifact, func(entity.Relationship) error { + visited++ + return nil + }) + if err == nil { + t.Fatal("readRelationshipsFixture unexpectedly accepted empty source ID") + } + if visited != 0 { + t.Fatalf("visited %d relationships before validating all source IDs", visited) + } +} + +func TestReadersValidateEveryEntityBeforeAnyVisit(t *testing.T) { + tests := []struct { + name string + read func(func() error) error + }{ + { + name: "node", + read: func(visit func() error) error { + root := t.TempDir() + contents := rawParquet(t, []NodeRow{ + {SourceID: "valid", Properties: map[string]any{}}, + {Properties: map[string]any{}}, + }) + path := "nodes.parquet" + writeContents(t, filepath.Join(root, path), contents) + hash := sha256.Sum256(contents) + return readNodesFixture(root, nodeFixtureArtifact{ + SchemaVersion: SchemaVersion, + Path: path, + SHA256: hex.EncodeToString(hash[:]), + Count: 2, + StoredBytes: int64(len(contents)), + }, func(entity.Node) error { + return visit() + }) + }, + }, + { + name: "relationship", + read: func(visit func() error) error { + root := t.TempDir() + contents := rawParquet(t, []RelationshipRow{ + {SourceID: "valid-1", StartID: "1", EndID: "2", Kind: "MemberOf", Properties: map[string]any{}}, + {SourceID: "valid-2", EndID: "3", Kind: "MemberOf", Properties: map[string]any{}}, + }) + path := "relationships.parquet" + writeContents(t, filepath.Join(root, path), contents) + hash := sha256.Sum256(contents) + return readRelationshipsFixture(root, relationshipFixtureArtifact{ + SchemaVersion: SchemaVersion, + Path: path, + SHA256: hex.EncodeToString(hash[:]), + Count: 2, + StoredBytes: int64(len(contents)), + }, func(entity.Relationship) error { + return visit() + }) + }, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + visited := 0 + err := test.read(func() error { + visited++ + return nil + }) + if err == nil { + t.Fatal("read unexpectedly accepted invalid entity") + } + if visited != 0 { + t.Fatalf("visited %d entities before validating all rows", visited) + } + }) + } +} + +func TestReadEmptyNodeArtifact(t *testing.T) { + root, artifact := installedNodeArtifact(t, nil) + visited := 0 + if err := readNodesFixture(root, artifact, func(entity.Node) error { + visited++ + return nil + }); err != nil { + t.Fatalf("read empty node artifact: %v", err) + } + if visited != 0 { + t.Fatalf("visited %d nodes, want none", visited) + } +} + +func TestReadNodesVisitsVerifiedSnapshotAfterPathReplacement(t *testing.T) { + root, artifact := installedNodeArtifact(t, []entity.Node{{SourceID: "old-1"}, {SourceID: "old-2"}}) + replacementTemporary := filepath.Join(root, "replacement.tmp") + replacementArtifact, err := writeNodesFixture(replacementTemporary, "replacement.parquet", Config{}, []entity.Node{{SourceID: "new-1"}, {SourceID: "new-2"}}) + if err != nil { + t.Fatalf("write replacement: %v", err) + } + replacement := filepath.Join(root, replacementArtifact.Path) + if err := os.Rename(replacementTemporary, replacement); err != nil { + t.Fatalf("install replacement: %v", err) + } + + afterVerifiedSnapshotForTest = func() { + if err := os.Rename(replacement, filepath.Join(root, artifact.Path)); err != nil { + t.Fatalf("replace artifact path: %v", err) + } + } + defer func() { afterVerifiedSnapshotForTest = nil }() + + var visited []string + if err := readNodesFixture(root, artifact, func(node entity.Node) error { + visited = append(visited, node.SourceID) + return nil + }); err != nil { + t.Fatalf("read nodes: %v", err) + } + if got, want := visited, []string{"old-1", "old-2"}; !reflect.DeepEqual(got, want) { + t.Fatalf("visited nodes = %v, want verified snapshot %v", got, want) + } +} + +func TestReadNodesRejectsSymlinksBeneathRootBeforeVisiting(t *testing.T) { + outsideRoot, outsideArtifact := installedNodeArtifact(t, []entity.Node{{SourceID: "outside"}}) + outsidePath := filepath.Join(outsideRoot, outsideArtifact.Path) + + tests := []struct { + name string + link func(string) (nodeFixtureArtifact, error) + }{ + { + name: "final file", + link: func(root string) (nodeFixtureArtifact, error) { + artifact := outsideArtifact + artifact.Path = "nodes.parquet" + return artifact, os.Symlink(outsidePath, filepath.Join(root, artifact.Path)) + }, + }, + { + name: "intermediate directory", + link: func(root string) (nodeFixtureArtifact, error) { + artifact := outsideArtifact + artifact.Path = "linked/nodes.parquet" + return artifact, os.Symlink(outsideRoot, filepath.Join(root, "linked")) + }, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + root := t.TempDir() + artifact, err := test.link(root) + if err != nil { + t.Fatalf("create in-root symlink to outside artifact: %v", err) + } + + visited := 0 + err = readNodesFixture(root, artifact, func(entity.Node) error { + visited++ + return nil + }) + if err == nil { + t.Fatal("readNodesFixture unexpectedly followed in-root symlink") + } + if visited != 0 { + t.Fatalf("visited %d nodes through in-root symlink", visited) + } + }) + } +} + +func installedNodeArtifact(t *testing.T, nodes []entity.Node) (string, nodeFixtureArtifact) { + t.Helper() + root := t.TempDir() + temporary := filepath.Join(root, "nodes.tmp") + artifact, err := writeNodesFixture(temporary, "nodes.parquet", Config{}, nodes) + if err != nil { + t.Fatalf("write nodes: %v", err) + } + if err := os.Rename(temporary, filepath.Join(root, artifact.Path)); err != nil { + t.Fatalf("install nodes: %v", err) + } + return root, artifact +} + +func rawParquet[T any](t *testing.T, rows []T) []byte { + t.Helper() + var output bytes.Buffer + writer := parquetgo.NewGenericWriter[T](&output, parquetgo.Compression(&zstd.Codec{})) + if _, err := writer.Write(rows); err != nil { + t.Fatalf("write raw Parquet rows: %v", err) + } + if err := writer.Close(); err != nil { + t.Fatalf("close raw Parquet writer: %v", err) + } + return output.Bytes() +} + +func writeContents(t *testing.T, path string, contents []byte) { + t.Helper() + if err := os.WriteFile(path, contents, 0o600); err != nil { + t.Fatalf("write artifact: %v", err) + } +} diff --git a/ret/parquet/row.go b/ret/parquet/row.go new file mode 100644 index 00000000..8f118436 --- /dev/null +++ b/ret/parquet/row.go @@ -0,0 +1,87 @@ +package parquet + +import ( + "fmt" + + "github.com/specterops/dawgs/ret/entity" +) + +type NodeRow struct { + SourceID string `parquet:"source_id"` + Kinds []string `parquet:"kinds,list"` + Properties any `parquet:"properties,variant"` +} + +type RelationshipRow struct { + SourceID string `parquet:"source_id"` + StartID string `parquet:"start_id"` + EndID string `parquet:"end_id"` + Kind string `parquet:"kind"` + Properties any `parquet:"properties,variant"` +} + +func nodeRow(value entity.Node) NodeRow { + return NodeRow{ + SourceID: value.SourceID, + Kinds: entity.CloneKinds(value.Kinds), + Properties: entity.CloneProperties(value.Properties), + } +} + +func relationshipRow(value entity.Relationship) RelationshipRow { + return RelationshipRow{ + SourceID: value.SourceID, + StartID: value.StartID, + EndID: value.EndID, + Kind: value.Kind, + Properties: entity.CloneProperties(value.Properties), + } +} + +func (s NodeRow) entity() (entity.Node, error) { + properties, err := propertyMap(s.Properties) + if err != nil { + return entity.Node{}, err + } + value := entity.Node{ + SourceID: s.SourceID, + Kinds: entity.CloneKinds(s.Kinds), + Properties: properties, + } + if err := value.Validate(); err != nil { + return entity.Node{}, err + } + return value, nil +} + +func (s RelationshipRow) entity() (entity.Relationship, error) { + if s.SourceID == "" { + return entity.Relationship{}, fmt.Errorf("relationship source ID is required") + } + properties, err := propertyMap(s.Properties) + if err != nil { + return entity.Relationship{}, err + } + value := entity.Relationship{ + SourceID: s.SourceID, + StartID: s.StartID, + EndID: s.EndID, + Kind: s.Kind, + Properties: properties, + } + if err := value.Validate(); err != nil { + return entity.Relationship{}, err + } + return value, nil +} + +func propertyMap(value any) (map[string]any, error) { + if value == nil { + return nil, nil + } + properties, ok := value.(map[string]any) + if !ok { + return nil, fmt.Errorf("properties VARIANT has type %T, want map[string]any or null", value) + } + return entity.CloneProperties(properties), nil +} diff --git a/ret/parquet/variant.go b/ret/parquet/variant.go new file mode 100644 index 00000000..f2ba3090 --- /dev/null +++ b/ret/parquet/variant.go @@ -0,0 +1,155 @@ +package parquet + +import ( + "fmt" + "math" + "reflect" + "strings" + + "github.com/parquet-go/parquet-go/variant" +) + +type variantReference struct { + kind reflect.Kind + pointer uintptr + length int + capacity int +} + +func validateProperties(properties map[string]any) error { + if err := validateVariantValue(reflect.ValueOf(properties), map[variantReference]string{}, "properties"); err != nil { + return err + } + _, _, err := variant.Marshal(properties) + return err +} + +func validateVariantValue(value reflect.Value, active map[variantReference]string, path string) error { + if !value.IsValid() { + return nil + } + if value.Kind() == reflect.Interface { + if value.IsNil() { + return nil + } + return validateVariantValue(value.Elem(), active, path) + } + + switch value.Kind() { + case reflect.Bool, + reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64, + reflect.Float32, reflect.Float64, + reflect.String: + return nil + + case reflect.Uint8, reflect.Uint16, reflect.Uint32: + return nil + + case reflect.Uint, reflect.Uint64: + if value.Uint() > math.MaxInt64 { + return fmt.Errorf("VARIANT value at %s overflows int64", path) + } + return nil + + case reflect.Pointer: + if value.IsNil() { + return nil + } + release, err := trackVariantReference(value, active, path) + if err != nil { + return err + } + defer release() + return validateVariantValue(value.Elem(), active, path) + + case reflect.Map: + if value.IsNil() { + return nil + } + if value.Type().Key().Kind() != reflect.String { + return fmt.Errorf("VARIANT map at %s has key type %s, want string", path, value.Type().Key()) + } + release, err := trackVariantReference(value, active, path) + if err != nil { + return err + } + defer release() + + iterator := value.MapRange() + for iterator.Next() { + fieldPath := path + "." + iterator.Key().String() + if err := validateVariantValue(iterator.Value(), active, fieldPath); err != nil { + return err + } + } + return nil + + case reflect.Slice: + if value.IsNil() || value.Type().Elem().Kind() == reflect.Uint8 { + return nil + } + release, err := trackVariantReference(value, active, path) + if err != nil { + return err + } + defer release() + return validateVariantElements(value, active, path) + + case reflect.Array: + return validateVariantElements(value, active, path) + + case reflect.Struct: + valueType := value.Type() + for index := range value.NumField() { + field := valueType.Field(index) + if !field.IsExported() || variantFieldSkipped(field) { + continue + } + if err := validateVariantValue(value.Field(index), active, path+"."+field.Name); err != nil { + return err + } + } + return nil + + default: + return fmt.Errorf("VARIANT value at %s has unsupported type %s", path, value.Type()) + } +} + +func validateVariantElements(value reflect.Value, active map[variantReference]string, path string) error { + for index := range value.Len() { + if err := validateVariantValue(value.Index(index), active, fmt.Sprintf("%s[%d]", path, index)); err != nil { + return err + } + } + return nil +} + +func trackVariantReference(value reflect.Value, active map[variantReference]string, path string) (func(), error) { + reference := variantReference{kind: value.Kind(), pointer: value.Pointer()} + if value.Kind() == reflect.Slice { + reference.length = value.Len() + reference.capacity = value.Cap() + } + if firstPath, found := active[reference]; found { + return nil, fmt.Errorf("VARIANT value cycle at %s references active value at %s", path, firstPath) + } + active[reference] = path + return func() { + delete(active, reference) + }, nil +} + +func variantFieldSkipped(field reflect.StructField) bool { + if tag, found := field.Tag.Lookup("variant"); found { + name, _, _ := strings.Cut(tag, ",") + if name != "" { + return name == "-" + } + } + if tag, found := field.Tag.Lookup("json"); found { + name, _, _ := strings.Cut(tag, ",") + return name == "-" + } + return false +} diff --git a/ret/parquet/writer_benchmark_test.go b/ret/parquet/writer_benchmark_test.go new file mode 100644 index 00000000..7930b9a1 --- /dev/null +++ b/ret/parquet/writer_benchmark_test.go @@ -0,0 +1,69 @@ +package parquet + +import ( + "fmt" + "path/filepath" + "testing" + + "github.com/specterops/dawgs/ret/entity" +) + +var ( + parquetBenchmarkArtifact nodeFixtureArtifact + parquetBenchmarkRow NodeRow +) + +func BenchmarkParquetVARIANTConversion(b *testing.B) { + node := entity.Node{ + SourceID: "1", + Kinds: []string{"User", "Principal"}, + Properties: map[string]any{ + "name": "alice", + "enabled": true, + "score": int64(42), + "nested": map[string]any{ + "labels": []any{"one", "two", "three"}, + }, + }, + } + + b.ReportAllocs() + b.ResetTimer() + for range b.N { + if err := validateProperties(node.Properties); err != nil { + b.Fatal(err) + } + parquetBenchmarkRow = nodeRow(node) + } +} + +func BenchmarkParquetVARIANTWriting(b *testing.B) { + nodes := make([]entity.Node, 256) + for index := range nodes { + nodes[index] = entity.Node{ + SourceID: fmt.Sprintf("%d", index+1), + Kinds: []string{"User", "Principal"}, + Properties: map[string]any{ + "name": fmt.Sprintf("user-%d", index), + "score": int64(index), + "nested": map[string]any{ + "enabled": index%2 == 0, + "groups": []any{"one", "two"}, + }, + }, + } + } + config := Config{} + directory := b.TempDir() + + b.ReportAllocs() + b.ResetTimer() + for index := 0; index < b.N; index++ { + path := filepath.Join(directory, fmt.Sprintf("nodes-%d.parquet", index)) + artifact, err := writeNodesFixture(path, filepath.Base(path), config, nodes) + if err != nil { + b.Fatal(err) + } + parquetBenchmarkArtifact = artifact + } +} diff --git a/ret/parquet/writer_test.go b/ret/parquet/writer_test.go new file mode 100644 index 00000000..793ba8a5 --- /dev/null +++ b/ret/parquet/writer_test.go @@ -0,0 +1,316 @@ +package parquet + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/hex" + "os" + "os/exec" + "path/filepath" + "reflect" + "runtime/debug" + "strings" + "testing" + "time" + + parquetgo "github.com/parquet-go/parquet-go" + "github.com/parquet-go/parquet-go/format" + "github.com/specterops/dawgs/ret/entity" +) + +func TestRelationshipSchemaContainsSourceIDAndUnshreddedVariantProperties(t *testing.T) { + root := t.TempDir() + temporary := filepath.Join(root, "relationships.tmp") + artifact, err := writeRelationshipsFixture(temporary, "relationships.parquet", Config{}, []entity.Relationship{{ + SourceID: "99", + StartID: "1", + EndID: "2", + Kind: "MemberOf", + Properties: map[string]any{ + "score": int64(7), + }, + }}) + if err != nil { + t.Fatalf("write relationships: %v", err) + } + + contents, err := os.ReadFile(temporary) + if err != nil { + t.Fatalf("read Parquet artifact: %v", err) + } + file, err := parquetgo.OpenFile(bytes.NewReader(contents), int64(len(contents))) + if err != nil { + t.Fatalf("open Parquet artifact: %v", err) + } + schema := strings.ToLower(file.Schema().String()) + if !strings.Contains(schema, "source_id") { + t.Fatalf("schema does not contain relationship source_id:\n%s", schema) + } + if !strings.Contains(schema, "required binary source_id (string)") { + t.Fatalf("schema does not require relationship source_id:\n%s", schema) + } + if !strings.Contains(schema, "group properties (variant)") { + t.Fatalf("schema does not identify properties as VARIANT:\n%s", schema) + } + if strings.Contains(schema, "typed_value") { + t.Fatalf("schema shreds properties VARIANT unexpectedly:\n%s", schema) + } + if got, want := artifact.SchemaVersion, SchemaVersion; got != want { + t.Fatalf("schema version = %q, want %q", got, want) + } + + var propertiesVariant bool + for _, element := range file.Metadata().Schema { + if element.Name == "properties" && element.LogicalType.Valid && element.LogicalType.V.Variant != nil { + propertiesVariant = true + } + } + if !propertiesVariant { + t.Fatal("properties field does not carry VARIANT logical type metadata") + } + for rowGroupIndex, rowGroup := range file.Metadata().RowGroups { + for columnIndex, column := range rowGroup.Columns { + if column.MetaData.Codec != format.Zstd { + t.Fatalf("row group %d column %d codec = %s, want ZSTD", rowGroupIndex, columnIndex, column.MetaData.Codec) + } + } + } +} + +func TestNodeRoundTripPreservesKindOrderAndDuplicates(t *testing.T) { + want := entity.Node{ + SourceID: "1", + Kinds: []string{"User", "Admin", "User"}, + Properties: map[string]any{ + "active": true, + }, + } + root := t.TempDir() + artifact, err := writeNodesFixture(filepath.Join(root, "nodes.tmp"), "nodes.parquet", Config{}, []entity.Node{want}) + if err != nil { + t.Fatalf("write nodes: %v", err) + } + if err := os.Rename(filepath.Join(root, "nodes.tmp"), filepath.Join(root, artifact.Path)); err != nil { + t.Fatalf("install artifact: %v", err) + } + + var got []entity.Node + if err := readNodesFixture(root, artifact, func(node entity.Node) error { + got = append(got, node) + return nil + }); err != nil { + t.Fatalf("read nodes: %v", err) + } + if len(got) != 1 || !reflect.DeepEqual(got[0], want) { + t.Fatalf("round trip nodes = %#v, want %#v", got, want) + } +} + +func TestWriteNodesRecordsStoredIntegrityMetadata(t *testing.T) { + root := t.TempDir() + temporary := filepath.Join(root, "nodes.tmp") + artifact, err := writeNodesFixture(temporary, "nested/nodes.parquet", Config{}, []entity.Node{{SourceID: "1"}}) + if err != nil { + t.Fatalf("write nodes: %v", err) + } + contents, err := os.ReadFile(temporary) + if err != nil { + t.Fatalf("read Parquet artifact: %v", err) + } + if got, want := artifact.Path, "nested/nodes.parquet"; got != want { + t.Fatalf("path = %q, want %q", got, want) + } + if got, want := artifact.Count, int64(1); got != want { + t.Fatalf("count = %d, want %d", got, want) + } + if got, want := artifact.StoredBytes, int64(len(contents)); got != want { + t.Fatalf("stored bytes = %d, want %d", got, want) + } + hash := sha256.Sum256(contents) + if got, want := artifact.SHA256, hex.EncodeToString(hash[:]); got != want { + t.Fatalf("SHA-256 = %q, want %q", got, want) + } +} + +func TestWritersEnforceFormatBoundaryAndSafePaths(t *testing.T) { + root := t.TempDir() + tests := []struct { + name string + write func() error + }{ + { + name: "relative temporary path", + write: func() error { + _, err := writeNodesFixture("nodes.tmp", "nodes.parquet", Config{}, []entity.Node{{SourceID: "1"}}) + return err + }, + }, + { + name: "escaping final path", + write: func() error { + _, err := writeNodesFixture(filepath.Join(root, "nodes.tmp"), "../nodes.parquet", Config{}, []entity.Node{{SourceID: "1"}}) + return err + }, + }, + { + name: "empty node source ID", + write: func() error { + _, err := writeNodesFixture(filepath.Join(root, "nodes.tmp"), "nodes.parquet", Config{}, []entity.Node{{}}) + return err + }, + }, + { + name: "empty relationship source ID", + write: func() error { + _, err := writeRelationshipsFixture(filepath.Join(root, "relationships.tmp"), "relationships.parquet", Config{}, []entity.Relationship{{ + StartID: "1", + EndID: "2", + Kind: "MemberOf", + }}) + return err + }, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + if err := test.write(); err == nil { + t.Fatal("write unexpectedly succeeded") + } + }) + } +} + +func TestWriteNodesReturnsErrorForUnsupportedVariantValue(t *testing.T) { + defer func() { + if recovered := recover(); recovered != nil { + t.Fatalf("writeNodesFixture panicked instead of returning an error: %v", recovered) + } + }() + + _, err := writeNodesFixture(filepath.Join(t.TempDir(), "nodes.tmp"), "nodes.parquet", Config{}, []entity.Node{{ + SourceID: "1", + Properties: map[string]any{ + "unsupported": make(chan int), + }, + }}) + if err == nil { + t.Fatal("writeNodesFixture unexpectedly accepted unsupported VARIANT value") + } +} + +func TestWriteNodesRejectsCyclicPropertiesBeforeOpeningTemporaryFile(t *testing.T) { + runCyclicPropertiesWriterTest(t, "node") +} + +func TestWriteRelationshipsRejectsCyclicPropertiesBeforeOpeningTemporaryFile(t *testing.T) { + runCyclicPropertiesWriterTest(t, "relationship") +} + +func TestWriteNodesAllowsRepeatedAcyclicPropertyReference(t *testing.T) { + shared := map[string]any{"value": "shared"} + _, err := writeNodesFixture(filepath.Join(t.TempDir(), "nodes.tmp"), "nodes.parquet", Config{}, []entity.Node{{ + SourceID: "node-1", + Properties: map[string]any{ + "first": shared, + "second": shared, + }, + }}) + if err != nil { + t.Fatalf("write node with repeated acyclic property reference: %v", err) + } +} + +func TestWriteNodesAllowsAcyclicOverlappingSlices(t *testing.T) { + values := make([]any, 1) + values[0] = values[:0] + _, err := writeNodesFixture(filepath.Join(t.TempDir(), "nodes.tmp"), "nodes.parquet", Config{}, []entity.Node{{ + SourceID: "node-1", + Properties: map[string]any{"values": values}, + }}) + if err != nil { + t.Fatalf("write node with acyclic overlapping slices: %v", err) + } +} + +func TestConfigHasNoStorageOrEnablementFields(t *testing.T) { + configType := reflect.TypeOf(Config{}) + if got, want := configType.NumField(), 0; got != want { + t.Fatalf("Config field count = %d, want %d", got, want) + } +} + +const cyclicPropertiesWriter = "DAWGS_PARQUET_CYCLIC_PROPERTIES_WRITER" + +func runCyclicPropertiesWriterTest(t *testing.T, writerName string) { + t.Helper() + if os.Getenv(cyclicPropertiesWriter) == writerName { + debug.SetMaxStack(1 << 20) + assertCyclicPropertiesWriter(t, writerName) + return + } + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + command := exec.CommandContext(ctx, os.Args[0], "-test.run=^"+t.Name()+"$") + command.Env = append(os.Environ(), cyclicPropertiesWriter+"="+writerName) + output, err := command.CombinedOutput() + if ctx.Err() != nil { + t.Fatalf("%s writer subprocess did not terminate: %v\n%s", writerName, ctx.Err(), output) + } + if err != nil { + t.Fatalf("%s writer subprocess crashed instead of returning an error: %v\n%s", writerName, err, output) + } +} + +func assertCyclicPropertiesWriter(t *testing.T, writerName string) { + t.Helper() + mapCycle := map[string]any{} + mapCycle["self"] = mapCycle + sliceCycle := make([]any, 1) + sliceCycle[0] = sliceCycle + type pointerCycle struct { + Self *pointerCycle + } + pointerValue := new(pointerCycle) + pointerValue.Self = pointerValue + + cycles := []struct { + name string + value any + }{ + {name: "map", value: mapCycle}, + {name: "slice", value: sliceCycle}, + {name: "pointer", value: pointerValue}, + } + for _, cycle := range cycles { + root := t.TempDir() + temporary := filepath.Join(root, "artifact.tmp") + properties := map[string]any{"cycle": cycle.value} + + var err error + switch writerName { + case "node": + _, err = writeNodesFixture(temporary, "nodes.parquet", Config{}, []entity.Node{{ + SourceID: "node-1", + Properties: properties, + }}) + case "relationship": + _, err = writeRelationshipsFixture(temporary, "relationships.parquet", Config{}, []entity.Relationship{{ + SourceID: "relationship-1", + StartID: "node-1", + EndID: "node-2", + Kind: "MemberOf", + Properties: properties, + }}) + default: + t.Fatalf("unknown writer %q", writerName) + } + if err == nil { + t.Fatalf("%s writer unexpectedly accepted cyclic %s properties", writerName, cycle.name) + } + if !strings.Contains(strings.ToLower(err.Error()), "cycle") { + t.Fatalf("%s writer error = %q, want descriptive cycle error", writerName, err) + } + } +} diff --git a/ret/result.go b/ret/result.go new file mode 100644 index 00000000..30b925ef --- /dev/null +++ b/ret/result.go @@ -0,0 +1,26 @@ +package ret + +type DumpResult struct { + ManifestPath string + GraphCount int + NodeCount int64 + RelationshipCount int64 +} + +type LoadResult struct { + GraphCount int + NodeCount int64 + RelationshipCount int64 +} + +type VerifyCollectionResult struct { + GraphCount int + NodeCount int64 + RelationshipCount int64 +} + +type VerifyDatabaseResult struct { + GraphCount int + NodeCount int64 + RelationshipCount int64 +} diff --git a/ret/scrub/action_counts_test.go b/ret/scrub/action_counts_test.go new file mode 100644 index 00000000..2c34c666 --- /dev/null +++ b/ret/scrub/action_counts_test.go @@ -0,0 +1,44 @@ +package scrub + +import ( + "encoding/json" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestActionCountsTotalAndIsZero(t *testing.T) { + require.True(t, (ActionCounts{}).IsZero()) + require.Zero(t, (ActionCounts{}).Total()) + + counts := ActionCounts{ + Preserve: 1, + Pseudonymize: 2, + Redact: 3, + ShiftTimestamp: 4, + } + require.False(t, counts.IsZero()) + require.EqualValues(t, 10, counts.Total()) +} + +func TestActionCountsJSONUsesNamedFieldsAndOmitsZeros(t *testing.T) { + payload, err := json.Marshal(ActionCounts{ + Pseudonymize: 2, + ShiftTimestamp: 1, + }) + + require.NoError(t, err) + require.JSONEq(t, `{"pseudonymize":2,"shift_timestamp":1}`, string(payload)) + + var decoded ActionCounts + require.NoError(t, json.Unmarshal( + []byte(`{"preserve":3,"pseudonymize":4,"redact":5,"shift_timestamp":6}`), + &decoded, + )) + require.Equal(t, ActionCounts{ + Preserve: 3, + Pseudonymize: 4, + Redact: 5, + ShiftTimestamp: 6, + }, decoded) +} diff --git a/ret/scrub/actions.go b/ret/scrub/actions.go new file mode 100644 index 00000000..770f767a --- /dev/null +++ b/ret/scrub/actions.go @@ -0,0 +1,334 @@ +package scrub + +import ( + "encoding/hex" + "fmt" + "regexp" + "strings" + "time" +) + +type propertyAction string + +const ( + actionPreserve propertyAction = "preserve" + actionPseudonymize propertyAction = "pseudonymize" + actionRedact propertyAction = "redact" + actionShiftTimestamp propertyAction = "shift_timestamp" +) + +type propertyPlan struct { + normalized string + reference bool + preserve bool + timestamp bool + freeText bool + path bool + script bool + sensitive bool + semantic bool +} + +var ( + emailPattern = regexp.MustCompile(`(?i)^[a-z0-9._%+\-]+@[a-z0-9.\-]+\.[a-z]{2,}$`) + uuidPattern = regexp.MustCompile(`(?i)^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$`) + domainSIDPattern = regexp.MustCompile(`^S-1-5-21-\d+-\d+-\d+$`) + objectSIDPattern = regexp.MustCompile(`^(S-1-5-21-\d+-\d+-\d+)-(\d+)$`) + ipv4Pattern = regexp.MustCompile(`^(\d{1,3}\.){3}\d{1,3}$`) + hostLikePattern = regexp.MustCompile(`(?i)^[a-z0-9][a-z0-9-]*(\.[a-z0-9][a-z0-9-]*)+$`) + secretValuePattern = regexp.MustCompile(`(?i)(password|secret|token|private[_-]?key|credential|apikey|api[_-]?key)`) +) + +func (s *Scrubber) planProperty(key string, value any) propertyAction { + plan := s.planKey(key) + + if plan.reference && isStringLike(value) { + return actionPseudonymize + } + if plan.preserve { + return actionPreserve + } + if plan.timestamp { + return actionShiftTimestamp + } + if plan.freeText { + return actionRedact + } + if plan.path || plan.script { + return actionPseudonymize + } + if s.shouldRedact(plan.normalized, value) { + return actionRedact + } + if s.classifyValue(value) != "" { + return actionPseudonymize + } + if plan.sensitive || plan.semantic || isStringLike(value) { + return actionPseudonymize + } + + return actionPreserve +} + +func isStringLike(value any) bool { + switch typed := value.(type) { + case string: + return strings.TrimSpace(typed) != "" + case []any: + for _, item := range typed { + if isStringLike(item) { + return true + } + } + case []string: + return len(typed) > 0 + } + + return false +} + +func (s *Scrubber) scrubWithAction(key string, value any, action propertyAction) any { + switch action { + case actionRedact: + return s.redact(value) + case actionShiftTimestamp: + return s.shiftTimestamp(value) + case actionPseudonymize: + return s.pseudonymizeValue(key, value, s.classifyValue(value)) + default: + return value + } +} + +func (s *Scrubber) shouldRedact(normalizedKey string, value any) bool { + if secretValuePattern.MatchString(normalizedKey) { + return true + } + switch typed := value.(type) { + case string: + return len(typed) > s.rules.Classifier.LongTextThreshold + case []any: + for _, item := range typed { + if s.shouldRedact(normalizedKey, item) { + return true + } + } + case []string: + for _, item := range typed { + if s.shouldRedact(normalizedKey, item) { + return true + } + } + } + + return false +} + +func (s *Scrubber) redact(value any) any { + switch typed := value.(type) { + case []any: + values := make([]any, len(typed)) + for index := range values { + values[index] = s.rules.RedactionMarker + } + return values + case []string: + values := make([]string, len(typed)) + for index := range values { + values[index] = s.rules.RedactionMarker + } + return values + case map[string]any: + values := make(map[string]any, len(typed)) + for key := range typed { + values[key] = s.rules.RedactionMarker + } + return values + default: + return s.rules.RedactionMarker + } +} + +func (s *Scrubber) shiftTimestamp(value any) any { + shift := time.Duration(s.rules.TimestampShiftDays) * 24 * time.Hour + switch typed := value.(type) { + case time.Time: + return typed.Add(shift).UTC().Format(time.RFC3339Nano) + case string: + if parsed, err := time.Parse(time.RFC3339Nano, typed); err == nil { + return parsed.Add(shift).UTC().Format(time.RFC3339Nano) + } + return s.pseudonymizeString(typed, "") + case int: + return typed + int(shift.Seconds()) + case int64: + return typed + int64(shift.Seconds()) + case float64: + return typed + shift.Seconds() + case []any: + values := make([]any, 0, len(typed)) + for _, item := range typed { + values = append(values, s.shiftTimestamp(item)) + } + return values + case []string: + values := make([]string, 0, len(typed)) + for _, item := range typed { + if shifted, ok := s.shiftTimestamp(item).(string); ok { + values = append(values, shifted) + } + } + return values + default: + return value + } +} + +func (s *Scrubber) pseudonymizeValue(key string, value any, shape string) any { + normalizedKey := normalizeKey(key) + switch typed := value.(type) { + case string: + if _, ok := s.referenceKeys[normalizedKey]; ok { + return s.pseudonymizeString(typed, s.classifyString(typed)) + } + return s.pseudonymizeString(typed, shape) + case []any: + values := make([]any, 0, len(typed)) + for _, item := range typed { + values = append(values, s.pseudonymizeValue(key, item, s.classifyValue(item))) + } + return values + case []string: + values := make([]string, 0, len(typed)) + for _, item := range typed { + if replacement, ok := s.pseudonymizeValue(key, item, s.classifyString(item)).(string); ok { + values = append(values, replacement) + } + } + return values + case map[string]any: + values := make(map[string]any, len(typed)) + for nestedKey, nestedValue := range typed { + values[nestedKey] = s.pseudonymizeValue(nestedKey, nestedValue, s.classifyValue(nestedValue)) + } + return values + default: + return value + } +} + +func (s *Scrubber) classifyValue(value any) string { + switch typed := value.(type) { + case string: + return s.classifyString(typed) + case []any: + for _, item := range typed { + if shape := s.classifyValue(item); shape != "" { + return shape + } + } + case []string: + for _, item := range typed { + if shape := s.classifyString(item); shape != "" { + return shape + } + } + } + + return "" +} + +func (s *Scrubber) classifyString(value string) string { + trimmed := strings.TrimSpace(value) + if trimmed == "" { + return "" + } + for _, rule := range s.shapeRules { + if rule.pattern.MatchString(trimmed) { + return rule.name + } + } + return "" +} + +func (s *Scrubber) pseudonymizeString(value string, shape string) string { + trimmed := strings.TrimSpace(value) + if trimmed == "" { + return value + } + digest := s.digest(trimmed) + switch { + case shape == "email" || emailPattern.MatchString(trimmed): + return "user-" + digest[:12] + "@" + s.rules.FakeDomain + case shape == "uuid" || uuidPattern.MatchString(trimmed): + return fmt.Sprintf("%s-%s-%s-%s-%s", digest[:8], digest[8:12], digest[12:16], digest[16:20], digest[20:32]) + case shape == "domain_sid" || domainSIDPattern.MatchString(trimmed): + return s.fakeDomainSID(digest) + case shape == "object_sid" || objectSIDPattern.MatchString(trimmed): + matches := objectSIDPattern.FindStringSubmatch(trimmed) + if len(matches) == 3 { + return s.pseudonymizeString(matches[1], "domain_sid") + "-" + matches[2] + } + return "value-" + digest[:16] + case shape == "ipv4" || ipv4Pattern.MatchString(trimmed): + return fmt.Sprintf("10.%d.%d.%d", intFromHex(digest[0:2]), intFromHex(digest[2:4]), intFromHex(digest[4:6])) + case shape == "host" || hostLikePattern.MatchString(trimmed): + return "host-" + digest[:12] + "." + s.rules.FakeDomain + default: + return "value-" + digest[:16] + } +} + +func (s *Scrubber) fakeDomainSID(digest string) string { + return fmt.Sprintf("S-1-5-21-%09d-%09d-%09d", intFromHex(digest[0:8])%1_000_000_000, intFromHex(digest[8:16])%1_000_000_000, intFromHex(digest[16:24])%1_000_000_000) +} + +func intFromHex(value string) int { + decoded, err := hex.DecodeString(value) + if err != nil { + return 0 + } + result := 0 + for _, next := range decoded { + result = result*256 + int(next) + } + return result +} + +func normalizeKey(value string) string { + value = strings.ToLower(strings.TrimSpace(value)) + value = strings.ReplaceAll(value, "-", "") + value = strings.ReplaceAll(value, "_", "") + return strings.ReplaceAll(value, " ", "") +} + +func isTimestampKey(normalizedKey string) bool { + return strings.Contains(normalizedKey, "time") || + strings.Contains(normalizedKey, "date") || + strings.Contains(normalizedKey, "created") || + strings.Contains(normalizedKey, "updated") || + strings.Contains(normalizedKey, "deleted") || + strings.Contains(normalizedKey, "modified") || + strings.HasSuffix(normalizedKey, "seenat") +} + +func isFreeTextKey(normalizedKey string) bool { + return strings.Contains(normalizedKey, "description") || strings.Contains(normalizedKey, "comment") || strings.Contains(normalizedKey, "note") || normalizedKey == "info" +} + +func isPathKey(normalizedKey string) bool { + return strings.Contains(normalizedKey, "path") || strings.Contains(normalizedKey, "directory") || strings.Contains(normalizedKey, "homedir") || strings.Contains(normalizedKey, "folder") +} + +func isScriptKey(normalizedKey string) bool { + return strings.Contains(normalizedKey, "script") +} + +func isSemanticOrgKey(normalizedKey string) bool { + switch normalizedKey { + case "title", "department", "division", "company", "organization", "office", "location": + return true + default: + return false + } +} diff --git a/ret/scrub/config.go b/ret/scrub/config.go new file mode 100644 index 00000000..532866dd --- /dev/null +++ b/ret/scrub/config.go @@ -0,0 +1,190 @@ +// Package scrub applies the retriever's deterministic property-scrubbing policy. +package scrub + +import ( + "fmt" + "os" + "regexp" + "strings" + + "github.com/pelletier/go-toml/v2" +) + +// Config configures a Scrubber. +type Config struct { + // Salt is runtime-only and is not decoded from or encoded into config files. + Salt string `toml:"-" json:"-"` + Rules +} + +// Rules contains the policy rules applied while scrubbing properties. +type Rules struct { + FakeDomain string `toml:"fake_domain"` + TimestampShiftDays int `toml:"timestamp_shift_days"` + RedactionMarker string `toml:"redaction_marker"` + GraphRules GraphRulesConfig `toml:"graph_rules"` + Classifier ClassifierConfig `toml:"classifier"` +} + +// GraphRulesConfig controls the reference-key portion of the policy. +type GraphRulesConfig struct { + DomainKind string `toml:"domain_kind"` + ObjectIDKey string `toml:"objectid_key"` + DomainNameKey string `toml:"domain_name_key"` + DomainSIDReferenceKeys []string `toml:"domain_sid_reference_keys"` + ObjectIDReferenceKeys []string `toml:"objectid_reference_keys"` + SelfObjectIDAliasKeys []string `toml:"self_objectid_alias_keys"` + DomainNameReferenceKeys []string `toml:"domain_name_reference_keys"` + CaseInsensitiveDomainNames bool `toml:"case_insensitive_domain_names"` + PreserveADSIDDomainPrefixes bool `toml:"preserve_ad_sid_domain_prefixes"` +} + +// ClassifierConfig controls property and value classification. +type ClassifierConfig struct { + LongTextThreshold int `toml:"long_text_threshold"` + PreserveKeys []string `toml:"preserve_keys"` + SensitiveKeyMarks []string `toml:"sensitive_key_markers"` + ValueShapePatterns []ValueShapeConfig `toml:"value_shapes"` +} + +// ValueShapeConfig identifies a named regular-expression value shape. +type ValueShapeConfig struct { + Name string `toml:"name"` + Pattern string `toml:"pattern"` +} + +// DefaultConfig returns the legacy scrub policy. +func DefaultConfig() Config { + return Config{ + Rules: Rules{ + FakeDomain: "example.invalid", + TimestampShiftDays: 17, + RedactionMarker: "[REDACTED]", + GraphRules: GraphRulesConfig{ + DomainKind: "Domain", + ObjectIDKey: "objectid", + DomainNameKey: "domain", + DomainSIDReferenceKeys: []string{"domainsid", "domain_sid"}, + ObjectIDReferenceKeys: []string{"objectid", "object_id", "sid", "owner_sid", "primarygroupid"}, + SelfObjectIDAliasKeys: []string{"objectsid"}, + DomainNameReferenceKeys: []string{"domain", "domain_name"}, + CaseInsensitiveDomainNames: true, + PreserveADSIDDomainPrefixes: true, + }, + Classifier: ClassifierConfig{ + LongTextThreshold: 512, + PreserveKeys: []string{"objectid", "domainsid", "kind"}, + SensitiveKeyMarks: []string{ + "password", + "secret", + "token", + "credential", + "privatekey", + "private_key", + "apikey", + "api_key", + "email", + "mail", + "phone", + "address", + "name", + "displayname", + "samaccountname", + "userprincipalname", + "dns", + "hostname", + }, + ValueShapePatterns: []ValueShapeConfig{ + {Name: "email", Pattern: `(?i)^[a-z0-9._%+\-]+@[a-z0-9.\-]+\.[a-z]{2,}$`}, + {Name: "uuid", Pattern: `(?i)^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$`}, + {Name: "domain_sid", Pattern: `^S-1-5-21-\d+-\d+-\d+$`}, + {Name: "object_sid", Pattern: `^(S-1-5-21-\d+-\d+-\d+)-(\d+)$`}, + {Name: "ipv4", Pattern: `^(\d{1,3}\.){3}\d{1,3}$`}, + {Name: "host", Pattern: `(?i)^[a-z0-9][a-z0-9-]*(\.[a-z0-9][a-z0-9-]*)+$`}, + }, + }, + }, + } +} + +// ReadConfig reads a TOML scrub policy. Missing values retain the defaults. +func ReadConfig(path string) (Config, error) { + contents, err := os.ReadFile(path) + if err != nil { + return Config{}, fmt.Errorf("read scrub config: %w", err) + } + + config := DefaultConfig() + if err := toml.Unmarshal(contents, &config); err != nil { + return Config{}, fmt.Errorf("parse scrub config: %w", err) + } + + return config, config.Validate() +} + +// Validate verifies that the scrub policy is complete, canonical, and compilable. +func (s Config) Validate() error { + if strings.TrimSpace(s.Rules.FakeDomain) == "" { + return fmt.Errorf("fake domain must be non-empty") + } + if strings.TrimSpace(s.Rules.FakeDomain) != s.Rules.FakeDomain { + return fmt.Errorf("fake domain must be trimmed") + } + if strings.ToLower(s.Rules.FakeDomain) != s.Rules.FakeDomain { + return fmt.Errorf("fake domain must be lowercase") + } + if strings.HasPrefix(s.Rules.FakeDomain, ".") || strings.HasSuffix(s.Rules.FakeDomain, ".") { + return fmt.Errorf("fake domain must not have a leading or trailing dot") + } + if strings.TrimSpace(s.Rules.RedactionMarker) == "" { + return fmt.Errorf("redaction marker must be non-empty") + } + if strings.TrimSpace(s.Rules.RedactionMarker) != s.Rules.RedactionMarker { + return fmt.Errorf("redaction marker must be trimmed") + } + if s.Rules.Classifier.LongTextThreshold <= 0 { + return fmt.Errorf("long text threshold must be greater than zero") + } + if s.Rules.TimestampShiftDays == 0 { + return fmt.Errorf("timestamp shift days must be non-zero") + } + + for index, shape := range s.Rules.Classifier.ValueShapePatterns { + if strings.TrimSpace(shape.Name) == "" { + return fmt.Errorf("value shape %d name must be non-empty", index) + } + if strings.TrimSpace(shape.Pattern) == "" { + return fmt.Errorf("value shape %q pattern must be non-empty", shape.Name) + } + if _, err := regexp.Compile(shape.Pattern); err != nil { + return fmt.Errorf("compile value shape %q: %w", shape.Name, err) + } + } + + return nil +} + +func cloneRules(rules Rules) Rules { + return Rules{ + FakeDomain: rules.FakeDomain, + TimestampShiftDays: rules.TimestampShiftDays, + RedactionMarker: rules.RedactionMarker, + GraphRules: cloneGraphRules(rules.GraphRules), + Classifier: cloneClassifier(rules.Classifier), + } +} + +func cloneGraphRules(rules GraphRulesConfig) GraphRulesConfig { + rules.DomainSIDReferenceKeys = append([]string(nil), rules.DomainSIDReferenceKeys...) + rules.ObjectIDReferenceKeys = append([]string(nil), rules.ObjectIDReferenceKeys...) + rules.SelfObjectIDAliasKeys = append([]string(nil), rules.SelfObjectIDAliasKeys...) + rules.DomainNameReferenceKeys = append([]string(nil), rules.DomainNameReferenceKeys...) + return rules +} + +func cloneClassifier(classifier ClassifierConfig) ClassifierConfig { + classifier.PreserveKeys = append([]string(nil), classifier.PreserveKeys...) + classifier.SensitiveKeyMarks = append([]string(nil), classifier.SensitiveKeyMarks...) + classifier.ValueShapePatterns = append([]ValueShapeConfig(nil), classifier.ValueShapePatterns...) + return classifier +} diff --git a/ret/scrub/config_test.go b/ret/scrub/config_test.go new file mode 100644 index 00000000..9847c841 --- /dev/null +++ b/ret/scrub/config_test.go @@ -0,0 +1,254 @@ +package scrub + +import ( + "os" + "path/filepath" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestDefaultConfigMatchesLegacyPolicy(t *testing.T) { + require.Equal(t, legacyDefaultConfig(), DefaultConfig()) + require.NoError(t, DefaultConfig().Validate()) +} + +func TestDefaultConfigReturnsIndependentMutableValues(t *testing.T) { + first := DefaultConfig() + second := DefaultConfig() + + first.Rules.GraphRules.DomainSIDReferenceKeys[0] = "changed" + first.Rules.GraphRules.ObjectIDReferenceKeys[0] = "changed" + first.Rules.GraphRules.SelfObjectIDAliasKeys[0] = "changed" + first.Rules.GraphRules.DomainNameReferenceKeys[0] = "changed" + first.Rules.Classifier.PreserveKeys[0] = "changed" + first.Rules.Classifier.SensitiveKeyMarks[0] = "changed" + first.Rules.Classifier.ValueShapePatterns[0].Name = "changed" + + require.Equal(t, legacyDefaultConfig(), second) +} + +func TestExampleConfigMatchesDefaultPolicy(t *testing.T) { + config, err := ReadConfig("example.toml") + + require.NoError(t, err) + config.Salt = "" + require.Equal(t, DefaultConfig(), config) +} + +func TestReadConfigDecodesDirectPolicyShapeAndKeepsSaltRuntimeOnly(t *testing.T) { + path := filepath.Join(t.TempDir(), "retriever.toml") + require.NoError(t, os.WriteFile(path, []byte(` +salt = "file-salt" +fake_domain = "scrub.example" +redaction_marker = "[X]" + +[graph_rules] +domain_kind = "CustomDomain" + +[classifier] +long_text_threshold = 8 +`), 0o600)) + + config, err := ReadConfig(path) + require.NoError(t, err) + require.Empty(t, config.Salt) + require.Equal(t, "scrub.example", config.Rules.FakeDomain) + require.Equal(t, "[X]", config.Rules.RedactionMarker) + require.Equal(t, "CustomDomain", config.Rules.GraphRules.DomainKind) + require.Equal(t, 8, config.Rules.Classifier.LongTextThreshold) + require.Equal(t, legacyDefaultConfig().Rules.GraphRules.ObjectIDReferenceKeys, config.Rules.GraphRules.ObjectIDReferenceKeys) +} + +func TestConfigValidateRejectsInvalidValueShapePattern(t *testing.T) { + config := DefaultConfig() + config.Rules.Classifier.ValueShapePatterns = []ValueShapeConfig{{ + Name: "invalid", + Pattern: "[", + }} + + require.EqualError(t, config.Validate(), "compile value shape \"invalid\": error parsing regexp: missing closing ]: `[`") +} + +func TestConfigValidationRejectsIncompleteOrNoncanonicalRules(t *testing.T) { + tests := []struct { + name string + change func(*Config) + want string + }{ + { + name: "empty fake domain", + change: func(config *Config) { + config.Rules.FakeDomain = "" + }, + want: "fake domain must be non-empty", + }, + { + name: "blank fake domain", + change: func(config *Config) { + config.Rules.FakeDomain = " \t" + }, + want: "fake domain must be non-empty", + }, + { + name: "untrimmed fake domain", + change: func(config *Config) { + config.Rules.FakeDomain = " example.invalid " + }, + want: "fake domain must be trimmed", + }, + { + name: "uppercase fake domain", + change: func(config *Config) { + config.Rules.FakeDomain = "Example.Invalid" + }, + want: "fake domain must be lowercase", + }, + { + name: "leading dot in fake domain", + change: func(config *Config) { + config.Rules.FakeDomain = ".example.invalid" + }, + want: "fake domain must not have a leading or trailing dot", + }, + { + name: "trailing dot in fake domain", + change: func(config *Config) { + config.Rules.FakeDomain = "example.invalid." + }, + want: "fake domain must not have a leading or trailing dot", + }, + { + name: "empty redaction marker", + change: func(config *Config) { + config.Rules.RedactionMarker = "" + }, + want: "redaction marker must be non-empty", + }, + { + name: "blank redaction marker", + change: func(config *Config) { + config.Rules.RedactionMarker = " \t" + }, + want: "redaction marker must be non-empty", + }, + { + name: "untrimmed redaction marker", + change: func(config *Config) { + config.Rules.RedactionMarker = " [REDACTED] " + }, + want: "redaction marker must be trimmed", + }, + { + name: "zero long text threshold", + change: func(config *Config) { + config.Rules.Classifier.LongTextThreshold = 0 + }, + want: "long text threshold must be greater than zero", + }, + { + name: "negative long text threshold", + change: func(config *Config) { + config.Rules.Classifier.LongTextThreshold = -1 + }, + want: "long text threshold must be greater than zero", + }, + { + name: "zero timestamp shift", + change: func(config *Config) { + config.Rules.TimestampShiftDays = 0 + }, + want: "timestamp shift days must be non-zero", + }, + { + name: "missing value shape name", + change: func(config *Config) { + config.Rules.Classifier.ValueShapePatterns = []ValueShapeConfig{{ + Pattern: "^value$", + }} + }, + want: "value shape 0 name must be non-empty", + }, + { + name: "missing value shape pattern", + change: func(config *Config) { + config.Rules.Classifier.ValueShapePatterns = []ValueShapeConfig{{ + Name: "custom", + }} + }, + want: "value shape \"custom\" pattern must be non-empty", + }, + } + + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + config := DefaultConfig() + testCase.change(&config) + + require.EqualError(t, config.Validate(), testCase.want) + scrubber, err := New(config) + require.Nil(t, scrubber) + require.EqualError(t, err, testCase.want) + }) + } +} + +func TestConfigValidateAllowsNoValueShapePatterns(t *testing.T) { + config := DefaultConfig() + config.Rules.Classifier.ValueShapePatterns = nil + + require.NoError(t, config.Validate()) +} + +func legacyDefaultConfig() Config { + return Config{ + Rules: Rules{ + FakeDomain: "example.invalid", + TimestampShiftDays: 17, + RedactionMarker: "[REDACTED]", + GraphRules: GraphRulesConfig{ + DomainKind: "Domain", + ObjectIDKey: "objectid", + DomainNameKey: "domain", + DomainSIDReferenceKeys: []string{"domainsid", "domain_sid"}, + ObjectIDReferenceKeys: []string{"objectid", "object_id", "sid", "owner_sid", "primarygroupid"}, + SelfObjectIDAliasKeys: []string{"objectsid"}, + DomainNameReferenceKeys: []string{"domain", "domain_name"}, + CaseInsensitiveDomainNames: true, + PreserveADSIDDomainPrefixes: true, + }, + Classifier: ClassifierConfig{ + LongTextThreshold: 512, + PreserveKeys: []string{"objectid", "domainsid", "kind"}, + SensitiveKeyMarks: []string{ + "password", + "secret", + "token", + "credential", + "privatekey", + "private_key", + "apikey", + "api_key", + "email", + "mail", + "phone", + "address", + "name", + "displayname", + "samaccountname", + "userprincipalname", + "dns", + "hostname", + }, + ValueShapePatterns: []ValueShapeConfig{ + {Name: "email", Pattern: `(?i)^[a-z0-9._%+\-]+@[a-z0-9.\-]+\.[a-z]{2,}$`}, + {Name: "uuid", Pattern: `(?i)^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$`}, + {Name: "domain_sid", Pattern: `^S-1-5-21-\d+-\d+-\d+$`}, + {Name: "object_sid", Pattern: `^(S-1-5-21-\d+-\d+-\d+)-(\d+)$`}, + {Name: "ipv4", Pattern: `^(\d{1,3}\.){3}\d{1,3}$`}, + {Name: "host", Pattern: `(?i)^[a-z0-9][a-z0-9-]*(\.[a-z0-9][a-z0-9-]*)+$`}, + }, + }, + }, + } +} diff --git a/ret/scrub/example.toml b/ret/scrub/example.toml new file mode 100644 index 00000000..57e4bf7d --- /dev/null +++ b/ret/scrub/example.toml @@ -0,0 +1,62 @@ +fake_domain = "example.invalid" +timestamp_shift_days = 17 +redaction_marker = "[REDACTED]" + +[graph_rules] +domain_kind = "Domain" +objectid_key = "objectid" +domain_name_key = "domain" +domain_sid_reference_keys = ["domainsid", "domain_sid"] +objectid_reference_keys = ["objectid", "object_id", "sid", "owner_sid", "primarygroupid"] +self_objectid_alias_keys = ["objectsid"] +domain_name_reference_keys = ["domain", "domain_name"] +case_insensitive_domain_names = true +preserve_ad_sid_domain_prefixes = true + +[classifier] +long_text_threshold = 512 +preserve_keys = ["objectid", "domainsid", "kind"] +sensitive_key_markers = [ + "password", + "secret", + "token", + "credential", + "privatekey", + "private_key", + "apikey", + "api_key", + "email", + "mail", + "phone", + "address", + "name", + "displayname", + "samaccountname", + "userprincipalname", + "dns", + "hostname", +] + +[[classifier.value_shapes]] +name = "email" +pattern = "(?i)^[a-z0-9._%+\\-]+@[a-z0-9.\\-]+\\.[a-z]{2,}$" + +[[classifier.value_shapes]] +name = "uuid" +pattern = "(?i)^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$" + +[[classifier.value_shapes]] +name = "domain_sid" +pattern = "^S-1-5-21-\\d+-\\d+-\\d+$" + +[[classifier.value_shapes]] +name = "object_sid" +pattern = "^(S-1-5-21-\\d+-\\d+-\\d+)-(\\d+)$" + +[[classifier.value_shapes]] +name = "ipv4" +pattern = "^(\\d{1,3}\\.){3}\\d{1,3}$" + +[[classifier.value_shapes]] +name = "host" +pattern = "(?i)^[a-z0-9][a-z0-9-]*(\\.[a-z0-9][a-z0-9-]*)+$" diff --git a/ret/scrub/scrubber.go b/ret/scrub/scrubber.go new file mode 100644 index 00000000..56ce17fa --- /dev/null +++ b/ret/scrub/scrubber.go @@ -0,0 +1,226 @@ +package scrub + +import ( + "crypto/hmac" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "regexp" + "strings" + "sync" +) + +const maxPropertyPlans = 4096 + +// ActionCounts reports the number of successful scrub actions. +type ActionCounts struct { + Preserve int64 `json:"preserve,omitempty"` + Pseudonymize int64 `json:"pseudonymize,omitempty"` + Redact int64 `json:"redact,omitempty"` + ShiftTimestamp int64 `json:"shift_timestamp,omitempty"` +} + +// Add accumulates other into counts. +func (s *ActionCounts) Add(other ActionCounts) { + s.Preserve += other.Preserve + s.Pseudonymize += other.Pseudonymize + s.Redact += other.Redact + s.ShiftTimestamp += other.ShiftTimestamp +} + +// Total returns the number of all successful scrub actions. +func (s ActionCounts) Total() int64 { + return s.Preserve + s.Pseudonymize + s.Redact + s.ShiftTimestamp +} + +// IsZero reports whether counts contains no scrub actions. +func (s ActionCounts) IsZero() bool { + return s == ActionCounts{} +} + +func (s *ActionCounts) increment(action propertyAction) { + switch action { + case actionPreserve: + s.Preserve++ + case actionPseudonymize: + s.Pseudonymize++ + case actionRedact: + s.Redact++ + case actionShiftTimestamp: + s.ShiftTimestamp++ + } +} + +type compiledShape struct { + name string + pattern *regexp.Regexp +} + +// Scrubber applies compiled, immutable rules to caller-owned property maps. +type Scrubber struct { + rules Rules + salt []byte + preserveKeys map[string]struct{} + referenceKeys map[string]struct{} + sensitiveKeyMarkers []string + shapeRules []compiledShape + propertyPlans map[string]propertyPlan + propertyPlansMu sync.RWMutex + rulesFingerprint string + saltFingerprint string +} + +// New compiles a scrub policy. +func New(config Config) (*Scrubber, error) { + config.Rules = cloneRules(config.Rules) + if err := config.Validate(); err != nil { + return nil, err + } + + preserveKeys := make(map[string]struct{}, len(config.Rules.Classifier.PreserveKeys)) + for _, key := range config.Rules.Classifier.PreserveKeys { + preserveKeys[normalizeKey(key)] = struct{}{} + } + + referenceKeys := map[string]struct{}{} + for _, keys := range [][]string{ + config.Rules.GraphRules.DomainSIDReferenceKeys, + config.Rules.GraphRules.ObjectIDReferenceKeys, + config.Rules.GraphRules.SelfObjectIDAliasKeys, + config.Rules.GraphRules.DomainNameReferenceKeys, + {config.Rules.GraphRules.ObjectIDKey, config.Rules.GraphRules.DomainNameKey}, + } { + for _, key := range keys { + if normalized := normalizeKey(key); normalized != "" { + referenceKeys[normalized] = struct{}{} + } + } + } + + shapeRules := make([]compiledShape, 0, len(config.Rules.Classifier.ValueShapePatterns)) + for _, shape := range config.Rules.Classifier.ValueShapePatterns { + pattern, err := regexp.Compile(shape.Pattern) + if err != nil { + return nil, err + } + shapeRules = append(shapeRules, compiledShape{name: shape.Name, pattern: pattern}) + } + + sensitiveKeyMarkers := make([]string, 0, len(config.Rules.Classifier.SensitiveKeyMarks)) + for _, marker := range config.Rules.Classifier.SensitiveKeyMarks { + if normalized := normalizeKey(marker); normalized != "" { + sensitiveKeyMarkers = append(sensitiveKeyMarkers, normalized) + } + } + + rulesBytes, err := json.Marshal(config.Rules) + if err != nil { + return nil, err + } + rulesDigest := sha256.Sum256(rulesBytes) + saltDigest := sha256.Sum256([]byte(config.Salt)) + + return &Scrubber{ + rules: config.Rules, + salt: append([]byte(nil), config.Salt...), + preserveKeys: preserveKeys, + referenceKeys: referenceKeys, + sensitiveKeyMarkers: sensitiveKeyMarkers, + shapeRules: shapeRules, + propertyPlans: map[string]propertyPlan{}, + rulesFingerprint: hex.EncodeToString(rulesDigest[:]), + saltFingerprint: hex.EncodeToString(saltDigest[:]), + }, nil +} + +// Scrub mutates properties in place and returns action counts. +func (s *Scrubber) Scrub(properties map[string]any) ActionCounts { + counts := ActionCounts{} + if s == nil { + return counts + } + s.scrubMap(properties, &counts) + return counts +} + +func (s *Scrubber) scrubMap(properties map[string]any, counts *ActionCounts) { + for key, value := range properties { + action := s.planProperty(key, value) + properties[key] = s.scrubWithAction(key, value, action) + counts.increment(action) + if action == actionPreserve { + s.scrubNested(value, counts) + } + } +} + +func (s *Scrubber) scrubNested(value any, counts *ActionCounts) { + switch typed := value.(type) { + case map[string]any: + s.scrubMap(typed, counts) + case []any: + for _, item := range typed { + s.scrubNested(item, counts) + } + } +} + +func (s *Scrubber) planKey(key string) propertyPlan { + normalized := normalizeKey(key) + s.propertyPlansMu.RLock() + plan, found := s.propertyPlans[normalized] + s.propertyPlansMu.RUnlock() + if found { + return plan + } + + _, plan.reference = s.referenceKeys[normalized] + _, plan.preserve = s.preserveKeys[normalized] + plan.normalized = normalized + plan.timestamp = isTimestampKey(normalized) + plan.freeText = isFreeTextKey(normalized) + plan.path = isPathKey(normalized) + plan.script = isScriptKey(normalized) + plan.sensitive = s.isSensitiveKey(normalized) + plan.semantic = isSemanticOrgKey(normalized) + + s.propertyPlansMu.Lock() + if existing, ok := s.propertyPlans[normalized]; ok { + plan = existing + } else if len(s.propertyPlans) < maxPropertyPlans { + s.propertyPlans[normalized] = plan + } + s.propertyPlansMu.Unlock() + return plan +} + +func (s *Scrubber) isSensitiveKey(normalizedKey string) bool { + for _, marker := range s.sensitiveKeyMarkers { + if strings.Contains(normalizedKey, marker) { + return true + } + } + return false +} + +func (s *Scrubber) digest(value string) string { + mac := hmac.New(sha256.New, s.salt) + _, _ = mac.Write([]byte(value)) + return hex.EncodeToString(mac.Sum(nil)) +} + +// RulesFingerprint returns the lowercase SHA-256 fingerprint of canonical rules. +func (s *Scrubber) RulesFingerprint() string { + if s == nil { + return "" + } + return s.rulesFingerprint +} + +// SaltFingerprint returns the lowercase SHA-256 fingerprint of the configured salt. +func (s *Scrubber) SaltFingerprint() string { + if s == nil { + return "" + } + return s.saltFingerprint +} diff --git a/ret/scrub/scrubber_benchmark_test.go b/ret/scrub/scrubber_benchmark_test.go new file mode 100644 index 00000000..f07c8b3d --- /dev/null +++ b/ret/scrub/scrubber_benchmark_test.go @@ -0,0 +1,31 @@ +package scrub_test + +import ( + "testing" + + "github.com/specterops/dawgs/ret/scrub" +) + +var scrubBenchmarkCounts scrub.ActionCounts + +func BenchmarkScrubPlanReuse(b *testing.B) { + config := scrub.DefaultConfig() + config.Salt = "benchmark-salt" + scrubber, err := scrub.New(config) + if err != nil { + b.Fatal(err) + } + properties := map[string]any{ + "name": "Alice Example", + "objectid": "S-1-5-21-111-222-333-1001", + "description": "fixed benchmark description", + "enabled": true, + } + scrubber.Scrub(properties) + + b.ReportAllocs() + + for b.Loop() { + scrubBenchmarkCounts = scrubber.Scrub(properties) + } +} diff --git a/ret/scrub/scrubber_test.go b/ret/scrub/scrubber_test.go new file mode 100644 index 00000000..bbfe0c9b --- /dev/null +++ b/ret/scrub/scrubber_test.go @@ -0,0 +1,400 @@ +package scrub + +import ( + "regexp" + "strconv" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func newTestScrubber(t *testing.T) *Scrubber { + t.Helper() + config := DefaultConfig() + config.Salt = "test-salt" + scrubber, err := New(config) + require.NoError(t, err) + return scrubber +} + +func TestScrubPseudonymizesSensitiveValues(t *testing.T) { + scrubber := newTestScrubber(t) + properties := map[string]any{ + "name": "Alice", + "email": "alice@example.com", + "password": "super-secret", + } + + counts := scrubber.Scrub(properties) + + for key, raw := range map[string]string{ + "name": "Alice", + "email": "alice@example.com", + "password": "super-secret", + } { + require.NotEqual(t, raw, properties[key], key) + } + require.EqualValues(t, 2, counts.Pseudonymize) + require.EqualValues(t, 1, counts.Redact) +} + +func TestScrubIsDeterministic(t *testing.T) { + left := newTestScrubber(t) + right := newTestScrubber(t) + leftProperties := map[string]any{"email": "alice@example.com"} + rightProperties := map[string]any{"email": "alice@example.com"} + + left.Scrub(leftProperties) + right.Scrub(rightProperties) + + require.Equal(t, leftProperties, rightProperties) +} + +func TestScrubShapeSpecificPseudonyms(t *testing.T) { + scrubber := newTestScrubber(t) + cases := []struct { + name string + key string + value string + pattern *regexp.Regexp + }{ + {"email", "email", "alice@example.com", regexp.MustCompile(`^user-[0-9a-f]{12}@example\.invalid$`)}, + {"uuid", "value", "00112233-4455-6677-8899-aabbccddeeff", regexp.MustCompile(`^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$`)}, + {"domain sid", "value", "S-1-5-21-1-2-3", regexp.MustCompile(`^S-1-5-21-\d{9}-\d{9}-\d{9}$`)}, + {"object sid", "value", "S-1-5-21-1-2-3-500", regexp.MustCompile(`^S-1-5-21-\d{9}-\d{9}-\d{9}-500$`)}, + {"ipv4", "value", "192.0.2.10", regexp.MustCompile(`^10\.\d+\.\d+\.\d+$`)}, + {"host", "value", "server.example.com", regexp.MustCompile(`^host-[0-9a-f]{12}\.example\.invalid$`)}, + {"generic", "value", "Alice", regexp.MustCompile(`^value-[0-9a-f]{16}$`)}, + } + + for _, testCase := range cases { + t.Run(testCase.name, func(t *testing.T) { + properties := map[string]any{testCase.key: testCase.value} + scrubber.Scrub(properties) + got, ok := properties[testCase.key].(string) + require.True(t, ok) + require.NotEqual(t, testCase.value, got) + require.Regexp(t, testCase.pattern, got) + }) + } +} + +func TestScrubObjectSIDShapeFallbackUsesGenericPseudonym(t *testing.T) { + config := DefaultConfig() + config.Salt = "test-salt" + config.Rules.Classifier.ValueShapePatterns = []ValueShapeConfig{{ + Name: "object_sid", + Pattern: "^not-a-sid$", + }} + scrubber, err := New(config) + require.NoError(t, err) + properties := map[string]any{"value": "not-a-sid"} + + scrubber.Scrub(properties) + + require.Regexp(t, `^value-[0-9a-f]{16}$`, properties["value"]) +} + +func TestScrubTimestampAndRedactionBranches(t *testing.T) { + scrubber := newTestScrubber(t) + properties := map[string]any{ + "description": strings.Repeat("x", DefaultConfig().Rules.Classifier.LongTextThreshold+1), + "email_map": map[string]any{"primary": "alice@example.com"}, + "seen_at": "2026-01-01T00:00:00Z", + "created_unix": 1767225600, + "updated_time": time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC), + "deleted_at": []string{"2026-01-01T00:00:00Z"}, + } + + counts := scrubber.Scrub(properties) + + require.Equal(t, "[REDACTED]", properties["description"]) + require.Equal(t, "2026-01-18T00:00:00Z", properties["seen_at"]) + require.Equal(t, 1768694400, properties["created_unix"]) + require.Equal(t, "2026-01-18T00:00:00Z", properties["updated_time"]) + require.Equal(t, []string{"2026-01-18T00:00:00Z"}, properties["deleted_at"]) + require.EqualValues(t, 1, counts.Redact) + require.EqualValues(t, 4, counts.ShiftTimestamp) + require.EqualValues(t, 1, counts.Pseudonymize) +} + +func TestScrubTimestampKeyHeuristics(t *testing.T) { + scrubber := newTestScrubber(t) + timestampKeys := []string{ + "timestamp", + "created_at", + "updated_at", + "deleted_at", + "modified_at", + "seen_at", + } + properties := make(map[string]any, len(timestampKeys)) + for _, key := range timestampKeys { + require.True(t, isTimestampKey(normalizeKey(key)), key) + require.Equal(t, actionShiftTimestamp, scrubber.planProperty(key, "2026-01-01T00:00:00Z"), key) + properties[key] = "2026-01-01T00:00:00Z" + } + + counts := scrubber.Scrub(properties) + + require.EqualValues(t, len(timestampKeys), counts.ShiftTimestamp) + for key, value := range properties { + require.Equal(t, "2026-01-18T00:00:00Z", value, key) + } + + nonTimestampKeys := []string{ + "format", + "seat", + "heat", + "coat", + "float", + } + nonTimestampProperties := make(map[string]any, len(nonTimestampKeys)) + for _, key := range nonTimestampKeys { + require.False(t, isTimestampKey(normalizeKey(key)), key) + require.NotEqual(t, actionShiftTimestamp, scrubber.planProperty(key, "json"), key) + nonTimestampProperties[key] = "json" + } + + counts = scrubber.Scrub(nonTimestampProperties) + + require.Zero(t, counts.ShiftTimestamp) +} + +func TestScrubRedactsFreeTextFields(t *testing.T) { + scrubber := newTestScrubber(t) + properties := map[string]any{ + "description": "Work item ABC123 service owner Example Person for Example Division", + "comments": "Read only access to placeholder application resource", + "info": "Example location operations notes", + } + + counts := scrubber.Scrub(properties) + + for key, value := range properties { + require.Equal(t, "[REDACTED]", value, key) + } + require.EqualValues(t, 3, counts.Redact) +} + +func TestScrubPseudonymizesPathAndScriptFields(t *testing.T) { + scrubber := newTestScrubber(t) + properties := map[string]any{ + "homedirectory": `\\fileserver01\share\account123`, + "profilepath": `\\profilehost01\profiles\group\account456`, + "logonscript": `startup\login.bat`, + } + + counts := scrubber.Scrub(properties) + + require.EqualValues(t, 3, counts.Pseudonymize) + for key, forbidden := range map[string][]string{ + "homedirectory": {"fileserver01", "share", "account123"}, + "profilepath": {"profilehost01", "profiles", "account456"}, + "logonscript": {"startup", "login"}, + } { + got, ok := properties[key].(string) + require.True(t, ok) + require.True(t, strings.HasPrefix(got, "value-")) + for _, value := range forbidden { + require.NotContains(t, got, value) + } + } +} + +func TestScrubPseudonymizesUnknownStringsAndPreservesSafeScalars(t *testing.T) { + scrubber := newTestScrubber(t) + properties := map[string]any{ + "business_justification": "Read only access to placeholder application for example organization", + "enabled": true, + "risk_score": 42, + } + + counts := scrubber.Scrub(properties) + + require.EqualValues(t, 1, counts.Pseudonymize) + require.EqualValues(t, 2, counts.Preserve) + require.Regexp(t, `^value-[0-9a-f]{16}$`, properties["business_justification"]) + require.Equal(t, true, properties["enabled"]) + require.Equal(t, 42, properties["risk_score"]) +} + +func TestScrubPseudonymizesTicketLikeValues(t *testing.T) { + properties := map[string]any{"request_id": "WORKITEM-12345"} + counts := newTestScrubber(t).Scrub(properties) + + require.EqualValues(t, 1, counts.Pseudonymize) + require.Regexp(t, `^value-[0-9a-f]{16}$`, properties["request_id"]) +} + +func TestScrubRecordFixturePreservesReferenceConsistencyAndActionCounts(t *testing.T) { + scrubber := newTestScrubber(t) + nested := map[string]any{ + "email": "alice@example.com", + "values": []any{"alpha", "beta"}, + } + records := []struct { + name string + properties map[string]any + want ActionCounts + }{ + { + name: "first node", + properties: map[string]any{ + "objectid": "S-1-5-21-1-2-3-500", + "owner_sid": "S-1-5-21-1-2-3-501", + "unrelated_repeat": "S-1-5-21-1-2-3-500", + "domain_sid_history": []any{"S-1-5-21-1-2-3", "", 42}, + "nested": nested, + "description": "free text", + "created_at": "2026-01-01T00:00:00Z", + "enabled": true, + }, + want: ActionCounts{Pseudonymize: 6, Preserve: 2, Redact: 1, ShiftTimestamp: 1}, + }, + { + name: "second node", + properties: map[string]any{ + "objectid": "S-1-5-21-1-2-3-501", + "owner_sid": " S-1-5-21-1-2-3-500 ", + "name": "ALICE", + "empty": "", + }, + want: ActionCounts{Pseudonymize: 3, Preserve: 1}, + }, + { + name: "first edge", + properties: map[string]any{ + "owner_sid": "S-1-5-21-1-2-3-500", + "path": `\\server\share\alice`, + "password": "secret", + }, + want: ActionCounts{Pseudonymize: 2, Redact: 1}, + }, + { + name: "third node", + properties: map[string]any{ + "objectid": "S-1-5-21-9-8-7-1000", + "owner_sid": "S-1-5-21-1-2-3-500", + }, + want: ActionCounts{Pseudonymize: 2}, + }, + { + name: "second edge", + properties: map[string]any{ + "owner_sid": "S-1-5-21-9-8-7-1000", + "office": "North", + }, + want: ActionCounts{Pseudonymize: 2}, + }, + } + + for _, record := range records { + t.Run(record.name, func(t *testing.T) { + require.Equal(t, record.want, scrubber.Scrub(record.properties)) + }) + } + + first := records[0].properties + second := records[1].properties + third := records[3].properties + + require.Equal(t, first["objectid"], first["unrelated_repeat"]) + require.Equal(t, first["objectid"], second["owner_sid"]) + require.NotEqual(t, first["objectid"], first["owner_sid"]) + require.NotEqual(t, first["objectid"], third["objectid"]) + + history, ok := first["domain_sid_history"].([]any) + require.True(t, ok) + require.Len(t, history, 3) + require.Regexp(t, `^S-1-5-21-\d{9}-\d{9}-\d{9}$`, history[0]) + require.Equal(t, "", history[1]) + require.Equal(t, 42, history[2]) + + // A preserved top-level map remains caller-owned and is recursively scrubbed. + require.Regexp(t, `^user-[0-9a-f]{12}@example\.invalid$`, nested["email"]) + values, ok := nested["values"].([]any) + require.True(t, ok) + require.Len(t, values, 2) + for _, value := range values { + require.Regexp(t, `^value-[0-9a-f]{16}$`, value) + } +} + +func TestScrubMutatesNestedValueBecauseInputCopyIsShallow(t *testing.T) { + nested := map[string]any{"password": "secret"} + properties := map[string]any{"nested": nested} + + newTestScrubber(t).Scrub(properties) + + require.Equal(t, "[REDACTED]", nested["password"]) +} + +func TestConstructedScrubberAlwaysScrubs(t *testing.T) { + config := DefaultConfig() + config.Salt = "test-salt" + scrubber, err := New(config) + require.NoError(t, err) + properties := map[string]any{"password": "secret"} + + counts := scrubber.Scrub(properties) + + require.Equal(t, ActionCounts{Redact: 1}, counts) + require.Equal(t, "[REDACTED]", properties["password"]) +} + +func TestNilScrubberDoesNotMutate(t *testing.T) { + var scrubber *Scrubber + properties := map[string]any{"password": "secret"} + + counts := scrubber.Scrub(properties) + + require.True(t, counts.IsZero()) + require.Equal(t, map[string]any{"password": "secret"}, properties) +} + +func TestScrubPropertyPlanCacheIsBounded(t *testing.T) { + scrubber := newTestScrubber(t) + for index := 0; index < maxPropertyPlans+100; index++ { + scrubber.planKey("unique-property-" + strconv.Itoa(index)) + } + + require.Len(t, scrubber.propertyPlans, maxPropertyPlans) +} + +func TestFingerprintsAreStableAndSaltIsNotExposed(t *testing.T) { + config := DefaultConfig() + config.Salt = "private" + first, err := New(config) + require.NoError(t, err) + second, err := New(config) + require.NoError(t, err) + + require.Equal(t, first.RulesFingerprint(), second.RulesFingerprint()) + require.Equal(t, first.SaltFingerprint(), second.SaltFingerprint()) + require.NotContains(t, first.SaltFingerprint(), "private") + require.Regexp(t, `^[0-9a-f]{64}$`, first.RulesFingerprint()) + require.Regexp(t, `^[0-9a-f]{64}$`, first.SaltFingerprint()) +} + +func TestNewPreservesExactSalt(t *testing.T) { + trimmedConfig := DefaultConfig() + trimmedConfig.Salt = "test-salt" + trimmed, err := New(trimmedConfig) + require.NoError(t, err) + spacedConfig := DefaultConfig() + spacedConfig.Salt = " test-salt " + spaced, err := New(spacedConfig) + require.NoError(t, err) + trimmedProperties := map[string]any{"email": "alice@example.com"} + spacedProperties := map[string]any{"email": "alice@example.com"} + + trimmed.Scrub(trimmedProperties) + spaced.Scrub(spacedProperties) + + require.NotEqual(t, trimmedProperties, spacedProperties) + require.NotEqual(t, trimmed.SaltFingerprint(), spaced.SaltFingerprint()) + require.Equal(t, "01320aaeeb335238814ed5c78b7f8f5fb73340eb9f4392e68f7861cc2f3a4bf8", spaced.SaltFingerprint()) +} diff --git a/ret/verify_collection.go b/ret/verify_collection.go new file mode 100644 index 00000000..1bba5cae --- /dev/null +++ b/ret/verify_collection.go @@ -0,0 +1,52 @@ +package ret + +import ( + "context" + "fmt" + "time" + + "github.com/specterops/dawgs/ret/collection" + "github.com/specterops/dawgs/ret/observe" +) + +const verifyCollectionOperationName = "verify_collection" + +// VerifyCollection validates every configured artifact in a collection and +// returns aggregate entity counts. +func VerifyCollection( + ctx context.Context, + config VerifyCollectionConfig, +) (result VerifyCollectionResult, resultErr error) { + started := time.Now() + observe.Emit(ctx, config.Observer, observe.OperationStarted{Operation: verifyCollectionOperationName}) + defer func() { + observe.Emit(ctx, config.Observer, observe.OperationCompleted{ + Operation: verifyCollectionOperationName, + Duration: time.Since(started), + Err: resultErr, + }) + }() + + if err := ctx.Err(); err != nil { + return VerifyCollectionResult{}, fmt.Errorf("verify collection: %w", err) + } + if err := config.Validate(); err != nil { + return VerifyCollectionResult{}, err + } + verification, err := collection.Verify(ctx, config.Directory, config.Observer) + if err != nil { + return VerifyCollectionResult{}, fmt.Errorf("%w: %w", ErrInvalidCollection, err) + } + if err := ctx.Err(); err != nil { + return VerifyCollectionResult{}, fmt.Errorf("verify collection completion: %w", err) + } + result.GraphCount = len(verification.Manifest.Graphs) + for _, graphEntry := range verification.Manifest.Graphs { + result.NodeCount += graphEntry.NodeCount + result.RelationshipCount += graphEntry.RelationshipCount + } + if err := ctx.Err(); err != nil { + return VerifyCollectionResult{}, fmt.Errorf("verify collection completion: %w", err) + } + return result, nil +} diff --git a/ret/verify_collection_test.go b/ret/verify_collection_test.go new file mode 100644 index 00000000..2575c4f2 --- /dev/null +++ b/ret/verify_collection_test.go @@ -0,0 +1,118 @@ +package ret + +import ( + "context" + "os" + "path/filepath" + "testing" + + "github.com/specterops/dawgs/ret/collection" + "github.com/specterops/dawgs/ret/observe" + "github.com/stretchr/testify/require" +) + +func TestVerifyCollectionFullyVerifiesArtifactsAndReturnsOnlyAggregateCounts(t *testing.T) { + // Break caught: forwarding the internal manifest instead of the public + // aggregate result, or skipping concrete artifact verification. + root := writeLoadCollection(t, []string{"second", "first"}, map[string]*dumpTestGraph{ + "second": { + nodes: dumpTestNodes(1, 2), + relationships: dumpTestRelationships(10, 1, 2, "LINKED"), + }, + "first": {nodes: dumpTestNodes(20)}, + }, true, true) + var events []observe.Event + + result, err := VerifyCollection(context.Background(), VerifyCollectionConfig{ + Directory: root, + Observer: observe.ObserverFunc(func(_ context.Context, event observe.Event) { + events = append(events, event) + }), + }) + + require.NoError(t, err) + require.Equal(t, VerifyCollectionResult{GraphCount: 2, NodeCount: 3, RelationshipCount: 1}, result) + require.IsType(t, observe.OperationStarted{}, events[0]) + require.IsType(t, observe.OperationCompleted{}, events[len(events)-1]) + started := events[0].(observe.OperationStarted) + completed := events[len(events)-1].(observe.OperationCompleted) + require.Equal(t, "verify_collection", started.Operation) + require.Equal(t, "verify_collection", completed.Operation) + require.NoError(t, completed.Err) + var verified int + for _, event := range events { + if _, ok := event.(observe.ArtifactVerified); ok { + verified++ + } + } + require.Equal(t, 6, verified) +} + +func TestVerifyCollectionCorruptParquetReturnsInvalidCollection(t *testing.T) { + // Break caught: accidentally switching the full verification facade to the + // JSONL-only load preflight and ignoring corrupt configured Parquet. + root := writeLoadCollection(t, []string{"asset"}, map[string]*dumpTestGraph{ + "asset": {nodes: dumpTestNodes(1)}, + }, true, true) + manifest, err := collection.Read(root) + require.NoError(t, err) + parquetArtifact := manifest.Graphs[0].NodeShards[0].Parquet + require.NoError(t, os.WriteFile(filepath.Join(root, filepath.FromSlash(parquetArtifact.Path)), []byte("corrupt"), 0o600)) + var events []observe.Event + + result, err := VerifyCollection(context.Background(), VerifyCollectionConfig{ + Directory: root, + Observer: observe.ObserverFunc(func(_ context.Context, event observe.Event) { + events = append(events, event) + }), + }) + + require.ErrorIs(t, err, ErrInvalidCollection) + require.Zero(t, result) + completed, ok := events[len(events)-1].(observe.OperationCompleted) + require.True(t, ok) + require.ErrorIs(t, completed.Err, ErrInvalidCollection) +} + +func TestVerifyCollectionRejectsInvalidConfigWithinObserverLifecycle(t *testing.T) { + // Break caught: calling collection verification with an empty root or + // omitting the terminal typed event on configuration failure. + var events []observe.Event + + _, err := VerifyCollection(context.Background(), VerifyCollectionConfig{ + Observer: observe.ObserverFunc(func(_ context.Context, event observe.Event) { + events = append(events, event) + }), + }) + + require.ErrorIs(t, err, ErrInvalidConfig) + require.Len(t, events, 2) + require.IsType(t, observe.OperationStarted{}, events[0]) + completed, ok := events[1].(observe.OperationCompleted) + require.True(t, ok) + require.ErrorIs(t, completed.Err, ErrInvalidConfig) +} + +func TestVerifyCollectionCancellationOnFinalArtifactSuppressesSuccess(t *testing.T) { + // Break caught: aggregating and returning success after the synchronous + // observer cancels while receiving the final verified artifact. + root := writeLoadCollection(t, []string{"asset"}, map[string]*dumpTestGraph{ + "asset": {nodes: dumpTestNodes(1)}, + }, true, true) + ctx, cancel := context.WithCancel(context.Background()) + var events []observe.Event + + result, err := VerifyCollection(ctx, VerifyCollectionConfig{ + Directory: root, + Observer: observe.ObserverFunc(func(_ context.Context, event observe.Event) { + events = append(events, event) + if value, ok := event.(observe.ArtifactVerified); ok && value.Format == "Parquet" { + cancel() + } + }), + }) + + require.ErrorIs(t, err, context.Canceled) + require.Zero(t, result) + requireSingleCanceledTerminalEvent(t, events) +} diff --git a/ret/verify_database.go b/ret/verify_database.go new file mode 100644 index 00000000..ad27d636 --- /dev/null +++ b/ret/verify_database.go @@ -0,0 +1,334 @@ +package ret + +import ( + "context" + "fmt" + "slices" + "strings" + "time" + + "github.com/specterops/dawgs/graph" + "github.com/specterops/dawgs/ret/collection" + "github.com/specterops/dawgs/ret/dawgs" + "github.com/specterops/dawgs/ret/entity" + "github.com/specterops/dawgs/ret/metrics" + "github.com/specterops/dawgs/ret/observe" +) + +const ( + verifyDatabaseOperationName = "verify_database" + verifyDatabaseNodesPhase = "nodes" + verifyDatabaseRelationshipsPhase = "relationships" +) + +// VerifyDatabase compares every graph in a collection manifest with its +// current database contents. It reads and validates only manifest.json; the +// collection's JSONL and Parquet artifacts are not opened. +func VerifyDatabase( + ctx context.Context, + database graph.Database, + config VerifyDatabaseConfig, +) (result VerifyDatabaseResult, resultErr error) { + started := time.Now() + observe.Emit(ctx, config.Observer, observe.OperationStarted{Operation: verifyDatabaseOperationName}) + defer func() { + observe.Emit(ctx, config.Observer, observe.OperationCompleted{ + Operation: verifyDatabaseOperationName, + Duration: time.Since(started), + Err: resultErr, + }) + }() + + if err := ctx.Err(); err != nil { + return VerifyDatabaseResult{}, fmt.Errorf("verify database: %w", err) + } + if err := config.Validate(); err != nil { + return VerifyDatabaseResult{}, err + } + manifest, err := collection.Read(config.Directory) + if err != nil { + return VerifyDatabaseResult{}, fmt.Errorf("%w: %w", ErrInvalidCollection, err) + } + if err := ctx.Err(); err != nil { + return VerifyDatabaseResult{}, fmt.Errorf("verify database collection preflight: %w", err) + } + + differences := make([]string, 0) + for _, graphEntry := range manifest.Graphs { + graphResult, graphDifferences, err := verifyDatabaseGraph(ctx, database, config, graphEntry) + if err != nil { + return VerifyDatabaseResult{}, err + } + result.GraphCount++ + result.NodeCount += graphResult.nodes + result.RelationshipCount += graphResult.relationships + differences = append(differences, graphDifferences...) + } + if err := ctx.Err(); err != nil { + return VerifyDatabaseResult{}, fmt.Errorf("verify database completion: %w", err) + } + if len(differences) != 0 { + return VerifyDatabaseResult{}, fmt.Errorf("%w: %s", ErrMetricsMismatch, strings.Join(differences, "; ")) + } + + return result, nil +} + +type verifyDatabaseGraphResult struct { + nodes int64 + relationships int64 +} + +func verifyDatabaseGraph( + ctx context.Context, + database graph.Database, + config VerifyDatabaseConfig, + graphEntry collection.Graph, +) (result verifyDatabaseGraphResult, differences []string, resultErr error) { + graphStarted := time.Now() + observe.Emit(ctx, config.Observer, observe.GraphStarted{ + Operation: verifyDatabaseOperationName, + Graph: graphEntry.Name, + }) + if err := ctx.Err(); err != nil { + return result, nil, verifyDatabaseGraphError(graphEntry.Name, err) + } + + source, err := dawgs.NewSource(database, graphEntry.Name, config.BatchSize) + if err != nil { + return result, nil, fmt.Errorf("prepare database verification source for graph %q: %w", graphEntry.Name, err) + } + builder := metrics.NewBuilder() + catalog := newVerifyDatabaseKindCatalog() + + if result.nodes, err = verifyDatabaseNodes(ctx, source, config.Observer, graphEntry, builder, catalog); err != nil { + return result, nil, err + } + if result.relationships, err = verifyDatabaseRelationships(ctx, source, config.Observer, graphEntry, builder, catalog); err != nil { + return result, nil, err + } + if err := ctx.Err(); err != nil { + return result, nil, verifyDatabaseGraphError(graphEntry.Name, err) + } + + if result.nodes != graphEntry.NodeCount { + differences = append(differences, fmt.Sprintf( + "graph %q node count differs: expected %d, actual %d", + graphEntry.Name, graphEntry.NodeCount, result.nodes, + )) + } + if result.relationships != graphEntry.RelationshipCount { + differences = append(differences, fmt.Sprintf( + "graph %q relationship count differs: expected %d, actual %d", + graphEntry.Name, graphEntry.RelationshipCount, result.relationships, + )) + } + if !slices.Equal(graphEntry.KindCatalog, catalog.values) { + differences = append(differences, fmt.Sprintf("graph %q kind catalog differs", graphEntry.Name)) + } + if err := metrics.Compare(graphEntry.Metrics, builder.Finalize()); err != nil { + differences = append(differences, fmt.Sprintf("graph %q: %v", graphEntry.Name, err)) + } + + observe.Emit(ctx, config.Observer, observe.GraphCompleted{ + Operation: verifyDatabaseOperationName, + Graph: graphEntry.Name, + Nodes: result.nodes, + Relationships: result.relationships, + Duration: time.Since(graphStarted), + }) + if err := ctx.Err(); err != nil { + return result, nil, verifyDatabaseGraphError(graphEntry.Name, err) + } + return result, differences, nil +} + +func verifyDatabaseNodes( + ctx context.Context, + source *dawgs.Source, + observer observe.Observer, + graphEntry collection.Graph, + builder *metrics.Builder, + catalog *verifyDatabaseKindCatalog, +) (int64, error) { + phaseStarted := time.Now() + observe.Emit(ctx, observer, observe.PhaseStarted{ + Operation: verifyDatabaseOperationName, + Graph: graphEntry.Name, + Phase: verifyDatabaseNodesPhase, + Completed: 0, + Total: graphEntry.NodeCount, + }) + if err := ctx.Err(); err != nil { + return 0, verifyDatabasePhaseError(graphEntry.Name, verifyDatabaseNodesPhase, err) + } + + var count int64 + for { + if err := ctx.Err(); err != nil { + return count, verifyDatabasePhaseError(graphEntry.Name, verifyDatabaseNodesPhase, err) + } + batch, err := source.NextNodes(ctx) + if err != nil { + return count, verifyDatabasePhaseError(graphEntry.Name, verifyDatabaseNodesPhase, err) + } + if err := ctx.Err(); err != nil { + return count, verifyDatabasePhaseError(graphEntry.Name, verifyDatabaseNodesPhase, err) + } + if len(batch.Entities) == 0 { + break + } + for _, node := range batch.Entities { + if err := observeVerifyDatabaseNode(ctx, observer, graphEntry, builder, catalog, node, &count); err != nil { + return count, err + } + } + } + + observe.Emit(ctx, observer, observe.PhaseCompleted{ + Operation: verifyDatabaseOperationName, + Graph: graphEntry.Name, + Phase: verifyDatabaseNodesPhase, + Completed: count, + Duration: time.Since(phaseStarted), + }) + if err := ctx.Err(); err != nil { + return count, verifyDatabasePhaseError(graphEntry.Name, verifyDatabaseNodesPhase, err) + } + return count, nil +} + +func observeVerifyDatabaseNode( + ctx context.Context, + observer observe.Observer, + graphEntry collection.Graph, + builder *metrics.Builder, + catalog *verifyDatabaseKindCatalog, + node entity.Node, + count *int64, +) error { + if err := ctx.Err(); err != nil { + return verifyDatabasePhaseError(graphEntry.Name, verifyDatabaseNodesPhase, err) + } + if err := builder.ObserveNode(node); err != nil { + return fmt.Errorf("verify database graph %q node metrics: %w", graphEntry.Name, err) + } + catalog.observeNode(node) + *count += 1 + observe.Emit(ctx, observer, observe.PhaseProgress{ + Operation: verifyDatabaseOperationName, + Graph: graphEntry.Name, + Phase: verifyDatabaseNodesPhase, + Completed: *count, + Total: graphEntry.NodeCount, + }) + if err := ctx.Err(); err != nil { + return verifyDatabasePhaseError(graphEntry.Name, verifyDatabaseNodesPhase, err) + } + return nil +} + +func verifyDatabaseRelationships( + ctx context.Context, + source *dawgs.Source, + observer observe.Observer, + graphEntry collection.Graph, + builder *metrics.Builder, + catalog *verifyDatabaseKindCatalog, +) (int64, error) { + phaseStarted := time.Now() + observe.Emit(ctx, observer, observe.PhaseStarted{ + Operation: verifyDatabaseOperationName, + Graph: graphEntry.Name, + Phase: verifyDatabaseRelationshipsPhase, + Completed: 0, + Total: graphEntry.RelationshipCount, + }) + if err := ctx.Err(); err != nil { + return 0, verifyDatabasePhaseError(graphEntry.Name, verifyDatabaseRelationshipsPhase, err) + } + + var count int64 + for { + if err := ctx.Err(); err != nil { + return count, verifyDatabasePhaseError(graphEntry.Name, verifyDatabaseRelationshipsPhase, err) + } + batch, err := source.NextRelationships(ctx) + if err != nil { + return count, verifyDatabasePhaseError(graphEntry.Name, verifyDatabaseRelationshipsPhase, err) + } + if err := ctx.Err(); err != nil { + return count, verifyDatabasePhaseError(graphEntry.Name, verifyDatabaseRelationshipsPhase, err) + } + if len(batch.Entities) == 0 { + break + } + for _, relationship := range batch.Entities { + if err := ctx.Err(); err != nil { + return count, verifyDatabasePhaseError(graphEntry.Name, verifyDatabaseRelationshipsPhase, err) + } + if err := builder.ObserveRelationship(relationship); err != nil { + return count, fmt.Errorf("verify database graph %q relationship metrics: %w", graphEntry.Name, err) + } + catalog.observeRelationship(relationship) + count++ + observe.Emit(ctx, observer, observe.PhaseProgress{ + Operation: verifyDatabaseOperationName, + Graph: graphEntry.Name, + Phase: verifyDatabaseRelationshipsPhase, + Completed: count, + Total: graphEntry.RelationshipCount, + }) + if err := ctx.Err(); err != nil { + return count, verifyDatabasePhaseError(graphEntry.Name, verifyDatabaseRelationshipsPhase, err) + } + } + } + + observe.Emit(ctx, observer, observe.PhaseCompleted{ + Operation: verifyDatabaseOperationName, + Graph: graphEntry.Name, + Phase: verifyDatabaseRelationshipsPhase, + Completed: count, + Duration: time.Since(phaseStarted), + }) + if err := ctx.Err(); err != nil { + return count, verifyDatabasePhaseError(graphEntry.Name, verifyDatabaseRelationshipsPhase, err) + } + return count, nil +} + +func verifyDatabaseGraphError(graphName string, err error) error { + return fmt.Errorf("verify database graph %q: %w", graphName, err) +} + +func verifyDatabasePhaseError(graphName, phase string, err error) error { + return fmt.Errorf("verify database graph %q %s phase: %w", graphName, phase, err) +} + +type verifyDatabaseKindCatalog struct { + seen map[string]struct{} + values []string +} + +func newVerifyDatabaseKindCatalog() *verifyDatabaseKindCatalog { + return &verifyDatabaseKindCatalog{seen: make(map[string]struct{})} +} + +func (s *verifyDatabaseKindCatalog) observeNode(node entity.Node) { + for _, kind := range node.Kinds { + s.add(kind) + } +} + +func (s *verifyDatabaseKindCatalog) observeRelationship(relationship entity.Relationship) { + s.add(relationship.Kind) +} + +func (s *verifyDatabaseKindCatalog) add(kind string) { + if _, found := s.seen[kind]; found { + return + } + s.seen[kind] = struct{}{} + s.values = append(s.values, kind) +} diff --git a/ret/verify_database_test.go b/ret/verify_database_test.go new file mode 100644 index 00000000..c38c0896 --- /dev/null +++ b/ret/verify_database_test.go @@ -0,0 +1,229 @@ +package ret + +import ( + "context" + "fmt" + "os" + "path/filepath" + "testing" + + "github.com/specterops/dawgs/graph" + "github.com/specterops/dawgs/ret/collection" + "github.com/specterops/dawgs/ret/observe" + "github.com/stretchr/testify/require" +) + +func TestVerifyDatabaseDoesNotOpenArtifacts(t *testing.T) { + // Break caught: using full collection verification and opening JSONL or + // Parquet even though database verification consumes only manifest metadata. + root, database := collectionAndMatchingDatabase(t, []string{"asset"}, map[string]*dumpTestGraph{ + "asset": { + nodes: dumpTestNodes(1, 2), + relationships: dumpTestRelationships(10, 1, 2, "LINKED"), + }, + }) + removeEveryArtifact(t, root) + + result, err := VerifyDatabase(context.Background(), database, VerifyDatabaseConfig{ + Directory: root, + BatchSize: 2, + }) + + require.NoError(t, err) + require.Equal(t, VerifyDatabaseResult{GraphCount: 1, NodeCount: 2, RelationshipCount: 1}, result) +} + +func TestVerifyDatabaseReportsEveryGraphMismatch(t *testing.T) { + // Break caught: returning after the first graph difference or comparing only + // totals and missing catalog/metric mismatches. + root, database := collectionAndMatchingDatabase(t, []string{"first", "second"}, map[string]*dumpTestGraph{ + "first": { + nodes: dumpTestNodes(1, 2), + relationships: dumpTestRelationships(10, 1, 2, "LINKED"), + }, + "second": {nodes: dumpTestNodesWithKinds([]string{"Original"}, 3)}, + }) + database.graphs["first"].relationships[0].Kind = graph.StringKind("CHANGED") + database.graphs["second"].nodes[0].Kinds = graph.Kinds{graph.StringKind("Changed")} + + _, err := VerifyDatabase(context.Background(), database, validVerifyDatabaseConfig(root)) + + require.ErrorIs(t, err, ErrMetricsMismatch) + require.ErrorContains(t, err, `graph "first" kind catalog differs`) + require.ErrorContains(t, err, "relationship kinds") + require.ErrorContains(t, err, `graph "second" kind catalog differs`) + require.ErrorContains(t, err, "node kind sequences") +} + +func TestVerifyDatabaseRejectsReorderedKindsWithSameCatalog(t *testing.T) { + // Break caught: treating the catalog as a set and omitting ordered node-kind + // sequence metrics. Both variants first see A then B, but their nodes differ. + root, database := collectionAndMatchingDatabase(t, []string{"asset"}, map[string]*dumpTestGraph{ + "asset": { + nodes: []*graph.Node{ + graph.NewNode(1, graph.NewProperties(), graph.StringKind("A"), graph.StringKind("B"), graph.StringKind("A")), + }, + }, + }) + database.graphs["asset"].nodes[0].Kinds = graph.Kinds{graph.StringKind("A"), graph.StringKind("A"), graph.StringKind("B")} + + _, err := VerifyDatabase(context.Background(), database, validVerifyDatabaseConfig(root)) + + require.ErrorIs(t, err, ErrMetricsMismatch) + require.NotContains(t, err.Error(), "kind catalog differs") + require.ErrorContains(t, err, "node kind sequences") +} + +func TestVerifyDatabasePreservesManifestGraphSourceAndLifecycleOrder(t *testing.T) { + // Break caught: sorting graphs, interleaving phases, exceeding source batch + // size, or emitting a lifecycle different from the dump-style scan. + root, database := collectionAndMatchingDatabase(t, []string{"second", "first"}, map[string]*dumpTestGraph{ + "second": { + nodes: dumpTestNodes(1, 2, 3), + relationships: dumpTestRelationships(10, 1, 2, "LINKED"), + }, + "first": {nodes: dumpTestNodes(20)}, + }) + database.fetches = nil + var events []observe.Event + + result, err := VerifyDatabase(context.Background(), database, VerifyDatabaseConfig{ + Directory: root, + BatchSize: 2, + Observer: observe.ObserverFunc(func(_ context.Context, event observe.Event) { + events = append(events, event) + }), + }) + + require.NoError(t, err) + require.Equal(t, VerifyDatabaseResult{GraphCount: 2, NodeCount: 4, RelationshipCount: 1}, result) + require.Equal(t, []string{ + "second:nodes", "second:nodes", "second:nodes", "second:relationships", "second:relationships", + "first:nodes", "first:nodes", "first:relationships", + }, database.fetches) + require.Equal(t, []string{ + "operation_started:verify_database", + "graph_started:second", + "phase_started:second:nodes:0:3", + "phase_progress:second:nodes:1:3", + "phase_progress:second:nodes:2:3", + "phase_progress:second:nodes:3:3", + "phase_completed:second:nodes:3", + "phase_started:second:relationships:0:1", + "phase_progress:second:relationships:1:1", + "phase_completed:second:relationships:1", + "graph_completed:second:3:1", + "graph_started:first", + "phase_started:first:nodes:0:1", + "phase_progress:first:nodes:1:1", + "phase_completed:first:nodes:1", + "phase_started:first:relationships:0:0", + "phase_completed:first:relationships:0", + "graph_completed:first:1:0", + "operation_completed:verify_database:ok", + }, verifyDatabaseEventNames(events)) +} + +func TestVerifyDatabaseObserverCancellationStopsBeforeLaterSourceQueriesOrEvents(t *testing.T) { + // Break caught: continuing into another source query or completion event after + // a synchronous progress observer cancels the operation. + root, database := collectionAndMatchingDatabase(t, []string{"asset"}, map[string]*dumpTestGraph{ + "asset": { + nodes: dumpTestNodes(1, 2), + relationships: dumpTestRelationships(10, 1, 2, "LINKED"), + }, + }) + database.fetches = nil + ctx, cancel := context.WithCancel(context.Background()) + var events []observe.Event + + result, err := VerifyDatabase(ctx, database, VerifyDatabaseConfig{ + Directory: root, + BatchSize: 1, + Observer: observe.ObserverFunc(func(_ context.Context, event observe.Event) { + events = append(events, event) + if value, ok := event.(observe.PhaseProgress); ok && value.Phase == "nodes" { + cancel() + } + }), + }) + + require.Zero(t, result) + require.ErrorIs(t, err, context.Canceled) + require.Equal(t, []string{"asset:nodes"}, database.fetches) + require.Equal(t, []string{ + "operation_started:verify_database", + "graph_started:asset", + "phase_started:asset:nodes:0:2", + "phase_progress:asset:nodes:1:2", + "operation_completed:verify_database:error", + }, verifyDatabaseEventNames(events)) + requireSingleCanceledTerminalEvent(t, events) +} + +func collectionAndMatchingDatabase( + t *testing.T, + graphOrder []string, + graphs map[string]*dumpTestGraph, +) (string, *dumpTestDatabase) { + t.Helper() + root := writeLoadCollection(t, graphOrder, graphs, true, true) + return root, newDumpTestDatabase(graphs) +} + +func validVerifyDatabaseConfig(root string) VerifyDatabaseConfig { + return VerifyDatabaseConfig{Directory: root, BatchSize: 2} +} + +func removeEveryArtifact(t *testing.T, root string) { + t.Helper() + manifest, err := collection.Read(root) + require.NoError(t, err) + for _, graphEntry := range manifest.Graphs { + for _, shard := range graphEntry.NodeShards { + if shard.JSONL != nil { + require.NoError(t, os.Remove(filepath.Join(root, filepath.FromSlash(shard.JSONL.Path)))) + } + if shard.Parquet != nil { + require.NoError(t, os.Remove(filepath.Join(root, filepath.FromSlash(shard.Parquet.Path)))) + } + } + for _, shard := range graphEntry.RelationshipShards { + if shard.JSONL != nil { + require.NoError(t, os.Remove(filepath.Join(root, filepath.FromSlash(shard.JSONL.Path)))) + } + if shard.Parquet != nil { + require.NoError(t, os.Remove(filepath.Join(root, filepath.FromSlash(shard.Parquet.Path)))) + } + } + } +} + +func verifyDatabaseEventNames(events []observe.Event) []string { + names := make([]string, len(events)) + for index, event := range events { + switch value := event.(type) { + case observe.OperationStarted: + names[index] = "operation_started:" + value.Operation + case observe.OperationCompleted: + status := "ok" + if value.Err != nil { + status = "error" + } + names[index] = "operation_completed:" + value.Operation + ":" + status + case observe.GraphStarted: + names[index] = "graph_started:" + value.Graph + case observe.GraphCompleted: + names[index] = fmt.Sprintf("graph_completed:%s:%d:%d", value.Graph, value.Nodes, value.Relationships) + case observe.PhaseStarted: + names[index] = fmt.Sprintf("phase_started:%s:%s:%d:%d", value.Graph, value.Phase, value.Completed, value.Total) + case observe.PhaseProgress: + names[index] = fmt.Sprintf("phase_progress:%s:%s:%d:%d", value.Graph, value.Phase, value.Completed, value.Total) + case observe.PhaseCompleted: + names[index] = fmt.Sprintf("phase_completed:%s:%s:%d", value.Graph, value.Phase, value.Completed) + default: + names[index] = fmt.Sprintf("unexpected:%T", event) + } + } + return names +}