From 7e9c91f5dd55384142527f16e68e34629b9702f4 Mon Sep 17 00:00:00 2001 From: John Hopper Date: Thu, 18 Jun 2026 10:46:30 -0700 Subject: [PATCH] feat(cmd): add retriever graph dump tooling --- GOLANG_CODING_STANDARD.md | 332 ++++++++ README.md | 19 +- cmd/retriever/README.md | 216 +++++ cmd/retriever/bench.go | 593 ++++++++++++++ cmd/retriever/bench_support.go | 231 ++++++ cmd/retriever/bench_test.go | 186 +++++ cmd/retriever/config.go | 151 ++++ cmd/retriever/config_test.go | 112 +++ cmd/retriever/database.go | 214 +++++ cmd/retriever/database_test.go | 89 +++ cmd/retriever/main.go | 417 ++++++++++ cmd/retriever/main_test.go | 93 +++ cmd/retriever/packaging_plan.md | 192 +++++ cmd/retriever/retriever_integration_test.go | 195 +++++ go.mod | 58 +- go.sum | 130 +-- query/model.go | 12 +- query/neo4j/neo4j_test.go | 14 + retriever/archive_envelope.go | 843 ++++++++++++++++++++ retriever/archive_envelope_test.go | 549 +++++++++++++ retriever/archive_keys.go | 393 +++++++++ retriever/archive_keys_test.go | 151 ++++ retriever/archive_tar.go | 403 ++++++++++ retriever/archive_tar_test.go | 295 +++++++ retriever/compression.go | 237 ++++++ retriever/compression_test.go | 175 ++++ retriever/defaults.toml | 63 ++ retriever/doc.go | 18 + retriever/dump.go | 766 ++++++++++++++++++ retriever/dump_test.go | 101 +++ retriever/errors.go | 85 ++ retriever/graph.go | 12 + retriever/load.go | 655 +++++++++++++++ retriever/load_test.go | 296 +++++++ retriever/manifest.go | 77 ++ retriever/manifest_test.go | 225 ++++++ retriever/metrics.go | 530 ++++++++++++ retriever/metrics_test.go | 367 +++++++++ retriever/options.go | 293 +++++++ retriever/options_test.go | 187 +++++ retriever/progress.go | 97 +++ retriever/progress_test.go | 64 ++ retriever/scan.go | 170 ++++ retriever/scan_test.go | 221 +++++ retriever/scrubber.go | 675 ++++++++++++++++ retriever/scrubber_test.go | 403 ++++++++++ retriever/types.go | 290 +++++++ retriever/verify.go | 357 +++++++++ 48 files changed, 12160 insertions(+), 92 deletions(-) create mode 100644 GOLANG_CODING_STANDARD.md create mode 100644 cmd/retriever/README.md create mode 100644 cmd/retriever/bench.go create mode 100644 cmd/retriever/bench_support.go create mode 100644 cmd/retriever/bench_test.go create mode 100644 cmd/retriever/config.go create mode 100644 cmd/retriever/config_test.go create mode 100644 cmd/retriever/database.go create mode 100644 cmd/retriever/database_test.go create mode 100644 cmd/retriever/main.go create mode 100644 cmd/retriever/main_test.go create mode 100644 cmd/retriever/packaging_plan.md create mode 100644 cmd/retriever/retriever_integration_test.go create mode 100644 retriever/archive_envelope.go create mode 100644 retriever/archive_envelope_test.go create mode 100644 retriever/archive_keys.go create mode 100644 retriever/archive_keys_test.go create mode 100644 retriever/archive_tar.go create mode 100644 retriever/archive_tar_test.go create mode 100644 retriever/compression.go create mode 100644 retriever/compression_test.go create mode 100644 retriever/defaults.toml create mode 100644 retriever/doc.go create mode 100644 retriever/dump.go create mode 100644 retriever/dump_test.go create mode 100644 retriever/errors.go create mode 100644 retriever/graph.go create mode 100644 retriever/load.go create mode 100644 retriever/load_test.go create mode 100644 retriever/manifest.go create mode 100644 retriever/manifest_test.go create mode 100644 retriever/metrics.go create mode 100644 retriever/metrics_test.go create mode 100644 retriever/options.go create mode 100644 retriever/options_test.go create mode 100644 retriever/progress.go create mode 100644 retriever/progress_test.go create mode 100644 retriever/scan.go create mode 100644 retriever/scan_test.go create mode 100644 retriever/scrubber.go create mode 100644 retriever/scrubber_test.go create mode 100644 retriever/types.go create mode 100644 retriever/verify.go diff --git a/GOLANG_CODING_STANDARD.md b/GOLANG_CODING_STANDARD.md new file mode 100644 index 00000000..c7958ca4 --- /dev/null +++ b/GOLANG_CODING_STANDARD.md @@ -0,0 +1,332 @@ +# Go Coding Standard + +This standard captures the preferred Go style used in `dawgs`. Follow `gofmt` +first, then apply the conventions below when choosing code shape. + +## Receiver Names + +Use `s` as the receiver name for methods, regardless of the concrete receiver +type. The goal is to remove per-type receiver-name churn and reduce cognitive +load while reading method bodies. + +```go +func (s *Server) Start() error { + go s.loop() + return nil +} + +func (s Config) Validate() error { + if s.Firewall.Backend != "nftables" { + return fmt.Errorf("unsupported firewall backend %q", s.Firewall.Backend) + } + + return nil +} +``` + +Avoid type-derived receiver names such as `srv`, `cfg`, `db`, or `wm` unless a +local collision makes `s` unusable. + +## Vertical Spacing + +Use single blank lines inside functions and closures to separate logical +paragraphs. The driver packages are the reference style: setup, validation, +resource acquisition, side-effecting work, mutation, and final return each get +their own visual space when they are independent steps. + +Keep validation blocks at the top of a function compact when they are a single +gate into the rest of the operation. + +```go +if len(kindIDs) == 0 { + return graph.Kinds{}, true +} + +if mappedKinds, err := kindMapper.MapKindIDs(ctx, kindIDs); err == nil { + return mappedKinds, true +} + +return nil, false +``` + +Separate a multi-line call from the next independent operation. This includes +logging, progress emission, database calls, and other side-effecting calls. + +```go +slog.Info("retriever load node phase started", + slog.String("graph", graphEntry.Name), + slog.Int64("node_count", graphEntry.NodeCount), +) +progress.emit(ProgressEvent{ + Operation: OperationLoad, + Message: "retriever load node phase started", + Graph: graphEntry.Name, +}) + +nodeMap, nodeCount, err := loadGraphNodes(ctx, db, graphEntry) +if err != nil { + return 0, 0, err +} +``` + +Within loops and callbacks, give each filter, transformation, mutation, and +checkpoint its own paragraph when the block does more than a couple of trivial +statements. + +```go +for _, node := range nodes { + kinds := node.Kinds.Strings() + sort.Strings(kinds) + addKindsToSet(nodeKinds, kinds) + + properties := node.Properties.MapOrEmpty() + if activeScrubber != nil { + properties = activeScrubber.scrubProperties(properties) + } + + item := FragmentNode{ + ID: node.ID.String(), + Kinds: kinds, + Properties: properties, + } + + items = append(items, item) + if len(items) >= shardSize { + if err := flush(); err != nil { + return err + } + } +} +``` + +Do not add blank lines between statements that are one tight operation, such as +creating a value and immediately returning it, assigning a field and returning +the receiver, or a short `if`/`else if` chain that represents one decision. + +## Error Handling and Scope + +Prefer initializer-backed `if` statements for operations that can fail. Handle +the error at the point it is raised, and keep successful values scoped to the +branch that uses them. + +```go +if cfg, err := server.ReadConfiguration(cfgPath); err != nil { + log.Fatalf("Error reading config: %v", err) +} else if dbInst, err := db.NewDatabase(cfg.DBPath); err != nil { + log.Fatalf("Error opening database: %v", err) +} else { + // cfg and dbInst are only available where they are valid. +} +``` + +Use `else if` chains when each later step depends on the successful output of +the previous step. This keeps the happy path close to the failure path and +prevents partially initialized variables from leaking into wider scopes. + +```go +if content, err := os.ReadFile(path); err != nil { + return cfg, err +} else if err := toml.Unmarshal(content, &cfg); err != nil { + return cfg, err +} else { + cfg = cfg.withDefaults() + return cfg, nil +} +``` + +When an error requires a special-case branch, keep that branch nested at the +error site. + +```go +if record, err := s.db.GetHostRecord(match.IPAddress); err != nil { + if !errors.Is(err, db.ErrNotFound) { + return err + } + + // Create the missing record here. +} else { + // Update the existing record here. +} +``` + +Use plain early returns when there is no useful success value to scope or when +the initializer chain would make the code harder to follow. + +## Variable Grouping + +Group related local variables aggressively with `var` blocks. This is preferred +when multiple locals establish the state for the same operation, especially when +some values are initialized and others are intentionally zero-valued. + +```go +var ( + records HostRecords + + txn = s.db.NewTransaction(false) + iter = txn.NewIterator(badger.DefaultIteratorOptions) + prefix = []byte(recordKeyPrefix) +) +``` + +Use blank lines inside the group to separate conceptual clusters. Prefer a +single grouped declaration over several adjacent `var` statements. + +For short-lived values used immediately, `:=` is still appropriate. + +```go +line := scanner.Text() +``` + +## Logical Spacing + +Use blank lines inside functions to separate logical phases. This is guidance, +not a hard formatting rule: prefer readability over mechanically inserting +empty lines. + +Common phase boundaries include setup before control flow, guard checks before +work, resource acquisition before deferred cleanup and use, and mutation before +the final return. + +```go +transaction, err := s.renderFirewallTransaction(state) +if err != nil { + return err +} + +log.Infof("Applying nftables transaction. Banned %d hosts.", state.Count()) +if output, err := s.runNftCommand(s.nftPath, transaction.Payload, "-j", "-f", "-"); err != nil { + return fmt.Errorf("backend apply failed running %s: %s: %w", s.nftPath, output, err) +} + +s.firewallState = state +return nil +``` + +Within loops, use blank lines to make each filtering or transformation step +stand on its own. + +```go +for _, record := range hostRecords { + if !now.Before(record.NextRefresh) { + continue + } + + ipAddress := net.ParseIP(record.IPAddress) + if ipAddress == nil || ipAddress.To4() == nil { + continue + } + + if staticCIDRAllows(cfg.StaticAllowCIDRs, ipAddress) { + continue + } + + state.BannedIPv4[record.IPAddress] = record +} +``` + +In `select` statements, separate cases with blank lines when each case performs +distinct work. + +```go +select { +case <-refreshTicker.C: + s.update() + +case <-s.updateC: + s.update() + +case <-s.joiner.StopC: + return +} +``` + +Avoid splitting tightly coupled statements when the second line is the immediate +effect of the first. + +## Function Ordering + +Prefer ordering functions in the same file so dependencies appear before the +functions that call them, when that makes the file read naturally. This is +guidance, not a hard formatting rule: public entry points, framework +conventions, and established local ordering may take precedence. + +## Struct Definitions + +Write struct type definitions across multiple lines, with one field per line. +Align naturally with `gofmt`; do not compress structs onto one line. + +```go +type FirewallConfig struct { + Backend string `toml:"backend"` + Table string `toml:"table"` + BanSet string `toml:"ban_set"` + Family string `toml:"family"` + DryRunSummaryOnly bool `toml:"dry_run_summary_only"` +} +``` + +Use field names in struct literals, especially for exported types, config +types, tests, and any literal with more than one field. + +```go +return Server{ + cfg: cfg, + db: dbInst, + firewallState: NewFirewallState(), + updateC: updateC, + joiner: NewJoinChannelPair(), + nftPath: nftPath, + nftBackend: NewNftablesBackend(cfg.Firewall), + runNftCommand: defaultNftCommandRunner, +}, nil +``` + +Prefer multi-line keyed literals even when a literal is currently small, because +they remain stable as fields are added and make diffs easier to read. + +```go +freshWatch := db.File{ + Path: event.Name, + Offset: 0, +} +``` + +## Practical Defaults + +Keep code direct and local. Favor readable control flow over abstractions that +only hide one or two calls. + +Return zero values explicitly on error for multi-return functions. + +```go +if watcher, err := fsnotify.NewWatcher(); err != nil { + return FileWatcher{}, err +} else { + return FileWatcher{ + cfg: cfg, + watcher: watcher, + }, nil +} +``` + +Defer cleanup immediately after acquiring a resource. + +```go +fin, err := os.Open(fileWatch.Path) +if err != nil { + return uniqueBanHosts, uniqueAllowHosts +} +defer fin.Close() +``` + +Use package-level grouped `const` and `var` declarations for related values. + +```go +const ( + ErrNotFound = errors.New("not found") + + fileWatchKey KeyFormat = "file_watch.%s" + hostRecordKey KeyFormat = "hosts.%s" + hostRecordKeyPrefix KeyFormat = "hosts." +) +``` diff --git a/README.md b/README.md index 697b3385..14e67c15 100644 --- a/README.md +++ b/README.md @@ -95,7 +95,9 @@ make quality FUZZ_REPORT=.coverage/fuzz.json MUTATION_REPORT=.coverage/mutation. `make quality_backend` captures PostgreSQL and Neo4j integration results for backend equivalence comparison. It requires `PG_CONNECTION_STRING` and `NEO4J_CONNECTION_STRING`. `make quality_bench` writes benchmark markdown and JSON captures -for later baseline comparison. +for later baseline comparison. Benchmark drift comparison is performed by +`make quality` through `tools/metrics`; there is no separate benchmark diff +command package. `make plan_corpus` captures plan diagnostics for the shared Cypher integration corpus. It accepts either `CONNECTION_STRING` for one backend or `PG_CONNECTION_STRING` and `NEO4J_CONNECTION_STRING` for both backends, then @@ -106,6 +108,21 @@ 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 compact OpenGraph-derived fragments. It supports +PostgreSQL and Neo4j, gzip and zstd compression, 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. + 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 anchors without matching JSON booleans or numbers. Simple relationship count fast paths depend on the schema's diff --git a/cmd/retriever/README.md b/cmd/retriever/README.md new file mode 100644 index 00000000..51e98aa2 --- /dev/null +++ b/cmd/retriever/README.md @@ -0,0 +1,216 @@ +# retriever + +`retriever` dumps and loads live Dawgs graph databases as manifest-based +collections of compact OpenGraph-derived JSON 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. + +## Dump + +```bash +retriever dump \ + -connection "$CONNECTION_STRING" \ + -out ./dumpdir \ + -graph default \ + -scrub none \ + -compression zstd \ + -zstd-level 11 \ + -shard-size 100000 +``` + +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. + +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, scrub pre-pass work, node and relationship +phase boundaries, periodic entity progress, manifest writing, and completion. + +## Encrypted Archives + +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: + +```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 +``` + +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. + +Unpack an encrypted archive with the private key: + +```bash +retriever unpack \ + -archive ./dump.tar.pq \ + -identity ./retriever-private.key \ + -out ./dumpdir +``` + +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`. + +## Scrubbed Dumps + +```bash +retriever dump \ + -connection "$CONNECTION_STRING" \ + -out ./scrubbed-dump \ + -graph default \ + -scrub full \ + -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. + +Classifier and graph-identifier settings can be overridden with `-config`; see +`../../retriever/defaults.toml` for the supported TOML shape. + +## Load + +```bash +retriever load \ + -connection "$CONNECTION_STRING" \ + -in ./dumpdir +``` + +Encrypted archives produced by `dump -archive-out` can be loaded directly: + +```bash +retriever load \ + -connection "$CONNECTION_STRING" \ + -archive ./dump.tar.pq \ + -identity ./retriever-private.key +``` + +Load reads and validates `manifest.json`, verifies every fragment checksum in a +separate pass, asserts destination graph schemas from the manifest metadata, and +then loads all nodes before relationships. Graph names from the dump are +preserved; load does not support overriding graph names. Archive loads decrypt +and validate into a temporary collection directory before opening the database. + +Load refuses to write into target graphs that already contain nodes or +relationships; clear the destination graph before restoring a collection. + +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. + +Pass `-verify-metrics` to scan the loaded destination graph after load and +compare it against the metrics stored in the manifest: + +```bash +retriever load \ + -connection "$CONNECTION_STRING" \ + -in ./dumpdir \ + -verify-metrics +``` + +Metrics verification adds a full post-load node and relationship scan. + +## Verify + +```bash +retriever verify \ + -connection "$CONNECTION_STRING" \ + -in ./dumpdir +``` + +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: + +```text +retriever: graph metrics mismatch: + graph "default" node_count: expected 884868, actual 884867 +``` + +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. + +## Bench + +```bash +retriever bench \ + -connection "$CONNECTION_STRING" \ + -graph default \ + -workers 1 \ + -batch-size 10000 \ + -sample-size 1000000 +``` + +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: + +```bash +retriever bench -connection "$CONNECTION_STRING" -graph default -compression zstd -json +``` + +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. + +## Testing Policy + +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 rather than heavy mocks of Dawgs database interfaces. This +keeps tests focused on durable behavior instead of line coverage for +orchestration code. diff --git a/cmd/retriever/bench.go b/cmd/retriever/bench.go new file mode 100644 index 00000000..6aec0730 --- /dev/null +++ b/cmd/retriever/bench.go @@ -0,0 +1,593 @@ +package main + +import ( + "context" + "fmt" + "io" + "log/slog" + "sort" + "sync" + "time" + + "github.com/specterops/dawgs/graph" + "github.com/specterops/dawgs/retriever" +) + +type benchReport struct { + Driver string `json:"driver"` + GeneratedAt time.Time `json:"generated_at"` + Graphs []benchGraphReport `json:"graphs"` +} + +type benchGraphReport struct { + Name string `json:"name"` + Results []benchResult `json:"results"` +} + +type benchResult struct { + Workers int `json:"workers"` + BatchSize int `json:"batch_size"` + SampleSize int `json:"sample_size,omitempty"` + NodeCount int64 `json:"node_count"` + EdgeCount int64 `json:"edge_count"` + NodeProcessed int64 `json:"node_processed"` + EdgeProcessed int64 `json:"edge_processed"` + NodeWallMillis int64 `json:"node_wall_millis"` + EdgeWallMillis int64 `json:"edge_wall_millis"` + NodeDBReadMillis int64 `json:"node_db_read_millis"` + EdgeDBReadMillis int64 `json:"edge_db_read_millis"` + NodeEncodeCompressMillis int64 `json:"node_encode_compress_millis,omitempty"` + EdgeEncodeCompressMillis int64 `json:"edge_encode_compress_millis,omitempty"` + TotalWallMillis int64 `json:"total_wall_millis"` + NodesPerSecond float64 `json:"nodes_per_second"` + EdgesPerSecond float64 `json:"edges_per_second"` + EntitiesPerSecond float64 `json:"entities_per_second"` + UncompressedBytes int64 `json:"uncompressed_bytes,omitempty"` + CompressedBytes int64 `json:"compressed_bytes,omitempty"` +} + +type benchPhaseResult struct { + Count int64 + WallElapsed time.Duration + DBReadElapsed time.Duration + EncodeCompressTime time.Duration + UncompressedByteSize int64 + CompressedByteSize int64 +} + +func Bench(ctx context.Context, db graph.Database, driverName string, targets []retriever.GraphTarget, options benchOptions) (benchReport, error) { + if err := options.validate(); err != nil { + return benchReport{}, err + } + + startedAt := time.Now() + slog.Info("retriever bench started", + slog.String("driver", driverName), + slog.Int("graph_count", len(targets)), + slog.Int("batch_size", options.BatchSize), + slog.Int("sample_size", options.SampleSize), + slog.Any("workers", options.Workers), + slog.String("compression", string(options.Compression)), + ) + + report := benchReport{ + Driver: driverName, + GeneratedAt: time.Now().UTC(), + Graphs: make([]benchGraphReport, 0, len(targets)), + } + + for targetIndex, target := range targets { + graphStartedAt := time.Now() + + slog.Info("retriever bench graph started", + slog.String("graph", target.Name), + slog.Int("graph_index", targetIndex+1), + slog.Int("graph_count", len(targets)), + ) + + targetGraph := graph.Graph{ + Name: target.Name, + } + + slog.Info("retriever bench counting graph entities", + slog.String("graph", target.Name), + ) + + nodeCount, edgeCount, err := countGraphEntities(ctx, db, targetGraph) + 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), + ) + + graphReport := benchGraphReport{ + Name: target.Name, + } + + 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 + } + + 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 + } + + 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.Duration("wall_elapsed", time.Since(graphStartedAt)), + ) + } + + slog.Info("retriever bench completed", + slog.String("driver", driverName), + slog.Int("graph_count", len(targets)), + slog.Duration("wall_elapsed", time.Since(startedAt)), + ) + + return report, nil +} + +type benchBatchProcessor[T any] struct { + ctx context.Context + cancel context.CancelFunc + process func([]T) (benchPhaseResult, error) + inline bool + jobs chan []T + wg sync.WaitGroup + mu sync.Mutex + result benchPhaseResult + err error +} + +func newBenchBatchProcessor[T any](ctx context.Context, workers int, process func([]T) (benchPhaseResult, error)) (*benchBatchProcessor[T], context.Context, error) { + if workers <= 0 { + return nil, nil, fmt.Errorf("bench workers must be > 0") + } + if process == nil { + return nil, nil, fmt.Errorf("bench batch processor is required") + } + + var ( + scanCtx, cancel = context.WithCancel(ctx) + processor = &benchBatchProcessor[T]{ + ctx: scanCtx, + cancel: cancel, + process: process, + inline: workers == 1, + } + ) + if processor.inline { + return processor, scanCtx, nil + } + + processor.jobs = make(chan []T, workers) + + for range workers { + processor.wg.Add(1) + go processor.run() + } + + return processor, scanCtx, nil +} + +func (s *benchBatchProcessor[T]) run() { + defer s.wg.Done() + for { + select { + case <-s.ctx.Done(): + return + case batch, ok := <-s.jobs: + if !ok { + return + } + + result, err := s.process(batch) + if err != nil { + s.setError(err) + return + } + + s.addResult(result) + } + } +} + +func (s *benchBatchProcessor[T]) handle(batch []T) error { + if err := s.currentError(); err != nil { + return err + } + + if s.inline { + result, err := s.process(batch) + if err != nil { + s.setError(err) + return err + } + + s.addResult(result) + + return nil + } + + select { + case <-s.ctx.Done(): + if err := s.currentError(); err != nil { + return err + } + + return s.ctx.Err() + + case s.jobs <- batch: + return nil + } +} + +func (s *benchBatchProcessor[T]) closeAndWait() (benchPhaseResult, error) { + if !s.inline { + close(s.jobs) + s.wg.Wait() + } + + s.cancel() + + if err := s.currentError(); err != nil { + return benchPhaseResult{}, err + } + + return s.snapshot(), nil +} + +func (s *benchBatchProcessor[T]) addDBReadElapsed(value time.Duration) { + s.mu.Lock() + defer s.mu.Unlock() + s.result.DBReadElapsed += value +} + +func (s *benchBatchProcessor[T]) addResult(value benchPhaseResult) { + s.mu.Lock() + defer s.mu.Unlock() + s.result.Count += value.Count + s.result.WallElapsed += value.WallElapsed + s.result.DBReadElapsed += value.DBReadElapsed + s.result.EncodeCompressTime += value.EncodeCompressTime + s.result.UncompressedByteSize += value.UncompressedByteSize + s.result.CompressedByteSize += value.CompressedByteSize +} + +func (s *benchBatchProcessor[T]) setError(err error) { + s.mu.Lock() + defer s.mu.Unlock() + if s.err == nil { + s.err = err + s.cancel() + } +} + +func (s *benchBatchProcessor[T]) currentError() error { + s.mu.Lock() + defer s.mu.Unlock() + return s.err +} + +func (s *benchBatchProcessor[T]) snapshot() benchPhaseResult { + s.mu.Lock() + defer s.mu.Unlock() + 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) + + processor, scanCtx, err := newBenchBatchProcessor(ctx, workers, func(nodes []*graph.Node) (benchPhaseResult, error) { + return benchNodeBatch(nodes, options) + }) + if err != nil { + return benchPhaseResult{}, err + } + + _, scanErr := scanEntityBatches(entityScanOptions[*graph.Node]{ + Total: planned, + BatchSize: options.BatchSize, + EntityName: "node", + Read: func(afterID graph.ID, hasAfterID bool, limit int) ([]*graph.Node, error) { + readStarted := time.Now() + nodes, err := readDatabaseNodes(scanCtx, db, targetGraph, afterID, hasAfterID, limit) + processor.addDBReadElapsed(time.Since(readStarted)) + + return nodes, err + }, + ID: func(node *graph.Node) graph.ID { + return node.ID + }, + Handle: func(nodes []*graph.Node) error { + return processor.handle(nodes) + }, + LogProgress: func(processed int64, startedAt time.Time, nextProgressAt int64) int64 { + progressResult := processor.snapshot() + progressResult.Count = processed + + return logBenchPhaseProgress(targetGraph.Name, retriever.PhaseNodes, workers, progressResult, planned, startedAt, nextProgressAt) + }, + }) + + result, processErr := processor.closeAndWait() + result.WallElapsed = time.Since(startedAt) + + if processErr != nil { + return benchPhaseResult{}, processErr + } + + if scanErr != nil { + return benchPhaseResult{}, scanErr + } + + return result, nil +} + +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) + + processor, scanCtx, err := newBenchBatchProcessor(ctx, workers, func(relationships []*graph.Relationship) (benchPhaseResult, error) { + return benchRelationshipBatch(relationships, options) + }) + if err != nil { + return benchPhaseResult{}, err + } + + _, scanErr := scanEntityBatches(entityScanOptions[*graph.Relationship]{ + Total: planned, + BatchSize: options.BatchSize, + EntityName: "relationship", + Read: func(afterID graph.ID, hasAfterID bool, limit int) ([]*graph.Relationship, error) { + readStarted := time.Now() + relationships, err := readDatabaseRelationships(scanCtx, db, targetGraph, afterID, hasAfterID, limit) + processor.addDBReadElapsed(time.Since(readStarted)) + + return relationships, err + }, + ID: func(relationship *graph.Relationship) graph.ID { + return relationship.ID + }, + Handle: func(relationships []*graph.Relationship) error { + return processor.handle(relationships) + }, + LogProgress: func(processed int64, startedAt time.Time, nextProgressAt int64) int64 { + progressResult := processor.snapshot() + progressResult.Count = processed + + return logBenchPhaseProgress(targetGraph.Name, retriever.PhaseEdges, workers, progressResult, planned, startedAt, nextProgressAt) + }, + }) + + result, processErr := processor.closeAndWait() + result.WallElapsed = time.Since(startedAt) + + if processErr != nil { + return benchPhaseResult{}, processErr + } + + if scanErr != nil { + return benchPhaseResult{}, scanErr + } + + return result, nil +} + +func benchNodeBatch(nodes []*graph.Node, options benchOptions) (benchPhaseResult, error) { + return benchCompressedBatch(len(nodes), options, func() retriever.NodeFragment { + 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 retriever.NodeFragment{ + Phase: retriever.PhaseNodes, + Items: items, + } + }) +} + +func benchRelationshipBatch(relationships []*graph.Relationship, options benchOptions) (benchPhaseResult, error) { + return benchCompressedBatch(len(relationships), options, func() retriever.EdgeFragment { + 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 retriever.EdgeFragment{ + Phase: retriever.PhaseEdges, + Items: items, + } + }) +} + +func benchCompressedBatch[T any](count int, options benchOptions, buildPayload func() T) (benchPhaseResult, error) { + result := benchPhaseResult{ + Count: int64(count), + } + if options.Compression == retriever.CompressionDisabled || count == 0 { + return result, nil + } + + encodeStarted := time.Now() + uncompressedBytes, compressedBytes, err := retriever.CompressedJSONSize(options.Compression, options.ZstdLevel, buildPayload()) + result.EncodeCompressTime = time.Since(encodeStarted) + + if err != nil { + return benchPhaseResult{}, err + } + + result.UncompressedByteSize = uncompressedBytes + result.CompressedByteSize = compressedBytes + + return result, nil +} + +func benchPlannedCount(total int64, sampleSize int) int64 { + if total <= 0 { + return 0 + } + + if sampleSize <= 0 { + return total + } + + sampleCount := int64(sampleSize) + if sampleCount > total { + return total + } + + return sampleCount +} + +func logBenchPhaseProgress(graphName string, phaseName retriever.Phase, 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.Int("worker_count", workers), + slog.Int64("processed", result.Count), + slog.Int64("planned_count", planned), + slog.Duration("wall_elapsed", time.Since(startedAt)), + slog.Duration("db_read_elapsed", result.DBReadElapsed), + slog.Duration("encode_compress_elapsed", result.EncodeCompressTime), + slog.Float64("entities_per_second", perSecond(result.Count, time.Since(startedAt))), + ) + + return retrieverNextProgressAt(result.Count, planned, nextProgressAt) +} + +func writeBenchReport(writer io.Writer, report benchReport) { + for _, graphReport := range report.Graphs { + fmt.Fprintf(writer, "graph: %s\n", graphReport.Name) + + 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", + result.Workers, + result.BatchSize, + result.SampleSize, + result.NodeProcessed, + result.NodeCount, + result.EdgeProcessed, + result.EdgeCount, + result.TotalWallMillis, + result.EntitiesPerSecond, + result.NodeDBReadMillis+result.EdgeDBReadMillis, + result.NodeEncodeCompressMillis+result.EdgeEncodeCompressMillis, + ) + } + } +} diff --git a/cmd/retriever/bench_support.go b/cmd/retriever/bench_support.go new file mode 100644 index 00000000..44493a1c --- /dev/null +++ b/cmd/retriever/bench_support.go @@ -0,0 +1,231 @@ +package main + +import ( + "context" + "fmt" + "time" + + "github.com/specterops/dawgs/graph" + "github.com/specterops/dawgs/query" +) + +const retrieverProgressEntityInterval int64 = 250_000 + +type graphEntitySnapshot struct { + NodeCount int64 + EdgeCount int64 +} + +type entityBatchReader[T any] func(afterID graph.ID, hasAfterID bool, batchSize int) ([]T, error) +type entityBatchHandler[T any] func([]T) error +type entityIDFunc[T any] func(T) graph.ID +type entityProgressLogger func(processed int64, startedAt time.Time, nextProgressAt int64) int64 + +type entityScanOptions[T any] struct { + Total int64 + BatchSize int + EntityName string + Read entityBatchReader[T] + ID entityIDFunc[T] + Handle entityBatchHandler[T] + LogProgress entityProgressLogger +} + +func scanEntityBatches[T any](options entityScanOptions[T]) (int64, error) { + entityName := options.EntityName + if entityName == "" { + entityName = "entity" + } + + if options.Total <= 0 { + return 0, nil + } + + if options.BatchSize <= 0 { + return 0, fmt.Errorf("%s batch size must be > 0", entityName) + } + + if options.Read == nil { + return 0, fmt.Errorf("%s batch reader is required", entityName) + } + + if options.ID == nil { + return 0, fmt.Errorf("%s ID accessor is required", entityName) + } + + var ( + lastID graph.ID + hasLastID bool + processed int64 + + startedAt = time.Now() + nextProgressAt = retrieverInitialProgressAt(options.Total) + ) + + for processed < options.Total { + remaining := options.Total - processed + batch, err := options.Read(lastID, hasLastID, retrieverBatchLimit(remaining, options.BatchSize)) + if err != nil { + return processed, err + } + + if int64(len(batch)) > remaining { + batch = batch[:int(remaining)] + } + + if len(batch) == 0 { + break + } + + nextID := options.ID(batch[len(batch)-1]) + if hasLastID && nextID <= lastID { + return processed, fmt.Errorf("%s keyset scan did not advance after ID %d", entityName, lastID.Uint64()) + } + + if options.Handle != nil { + if err := options.Handle(batch); err != nil { + return processed, err + } + } + + processed += int64(len(batch)) + lastID = nextID + hasLastID = true + + if options.LogProgress != nil { + nextProgressAt = options.LogProgress(processed, startedAt, nextProgressAt) + } + } + + return processed, nil +} + +func countGraphEntities(ctx context.Context, db graph.Database, targetGraph graph.Graph) (int64, int64, error) { + entitySnapshot, err := countGraphEntitySnapshot(ctx, db, targetGraph) + if err != nil { + return 0, 0, err + } + + return entitySnapshot.NodeCount, entitySnapshot.EdgeCount, nil +} + +func countGraphEntitySnapshot(ctx context.Context, db graph.Database, targetGraph graph.Graph) (graphEntitySnapshot, error) { + var entitySnapshot graphEntitySnapshot + if err := db.ReadTransaction(ctx, func(tx graph.Transaction) error { + tx = tx.WithGraph(targetGraph) + + var err error + if entitySnapshot.NodeCount, err = tx.Nodes().Count(); err != nil { + return fmt.Errorf("count nodes: %w", err) + } + + if entitySnapshot.EdgeCount, err = tx.Relationships().Count(); err != nil { + return fmt.Errorf("count relationships: %w", err) + } + + return nil + }); err != nil { + return graphEntitySnapshot{}, err + } + + return entitySnapshot, nil +} + +func readDatabaseNodes(ctx context.Context, db graph.Database, targetGraph graph.Graph, afterID graph.ID, hasAfterID bool, batchSize int) ([]*graph.Node, error) { + var nodes []*graph.Node + if err := db.ReadTransaction(ctx, func(tx graph.Transaction) error { + tx = tx.WithGraph(targetGraph) + nodeQuery := tx.Nodes(). + OrderBy(query.NodeID()). + Limit(batchSize) + + if hasAfterID { + nodeQuery = nodeQuery.Filter(query.GreaterThan(query.NodeID(), afterID)) + } + + return nodeQuery.Fetch(func(cursor graph.Cursor[*graph.Node]) error { + for node := range cursor.Chan() { + nodes = append(nodes, node) + } + + return cursor.Error() + }) + }); err != nil { + if hasAfterID { + return nil, fmt.Errorf("read node batch after ID %d: %w", afterID.Uint64(), err) + } + + return nil, fmt.Errorf("read initial node batch: %w", err) + } + + return nodes, nil +} + +func readDatabaseRelationships(ctx context.Context, db graph.Database, targetGraph graph.Graph, afterID graph.ID, hasAfterID bool, batchSize int) ([]*graph.Relationship, error) { + var relationships []*graph.Relationship + if err := db.ReadTransaction(ctx, func(tx graph.Transaction) error { + tx = tx.WithGraph(targetGraph) + relationshipQuery := tx.Relationships(). + OrderBy(query.RelationshipID()). + Limit(batchSize) + + if hasAfterID { + relationshipQuery = relationshipQuery.Filter(query.GreaterThan(query.RelationshipID(), afterID)) + } + + return relationshipQuery.Fetch(func(cursor graph.Cursor[*graph.Relationship]) error { + for relationship := range cursor.Chan() { + relationships = append(relationships, relationship) + } + + return cursor.Error() + }) + }); err != nil { + if hasAfterID { + return nil, fmt.Errorf("read relationship batch after ID %d: %w", afterID.Uint64(), err) + } + + return nil, fmt.Errorf("read initial relationship batch: %w", err) + } + + return relationships, nil +} + +func retrieverInitialProgressAt(planned int64) int64 { + if planned <= retrieverProgressEntityInterval { + return 0 + } + + return retrieverProgressEntityInterval +} + +func retrieverBatchLimit(remaining int64, batchSize int) int { + if remaining <= 0 { + return 0 + } + + if int64(batchSize) > remaining { + return int(remaining) + } + + return batchSize +} + +func retrieverNextProgressAt(processed int64, planned int64, nextProgressAt int64) int64 { + if nextProgressAt <= processed { + nextProgressAt += ((processed-nextProgressAt)/retrieverProgressEntityInterval + 1) * retrieverProgressEntityInterval + } + if nextProgressAt >= planned { + return 0 + } + + return nextProgressAt +} + +func perSecond(count int64, elapsed time.Duration) float64 { + if count == 0 || elapsed <= 0 { + return 0 + } + + return float64(count) / elapsed.Seconds() +} diff --git a/cmd/retriever/bench_test.go b/cmd/retriever/bench_test.go new file mode 100644 index 00000000..e902d3ba --- /dev/null +++ b/cmd/retriever/bench_test.go @@ -0,0 +1,186 @@ +package main + +import ( + "bytes" + "context" + "strings" + "testing" + "time" + + "github.com/specterops/dawgs/graph" + "github.com/specterops/dawgs/retriever" +) + +func TestBenchSamplingHelpers(t *testing.T) { + if got := benchPlannedCount(10, 3); got != 3 { + t.Fatalf("planned sampled count = %d", got) + } + if got := benchPlannedCount(10, 0); got != 10 { + t.Fatalf("planned full count = %d", got) + } + if got := benchPlannedCount(3, 10); got != 3 { + t.Fatalf("planned capped count = %d", got) + } + if got := benchPlannedCount(0, 10); got != 0 { + t.Fatalf("planned empty count = %d", got) + } + if got := retrieverBatchLimit(3, 10); got != 3 { + t.Fatalf("batch limit for remainder = %d", got) + } + if got := retrieverBatchLimit(20, 10); got != 10 { + t.Fatalf("batch limit for full batch = %d", got) + } + if got := retrieverInitialProgressAt(retrieverProgressEntityInterval); got != 0 { + t.Fatalf("unexpected progress threshold for exact interval: %d", got) + } + if got := retrieverInitialProgressAt(retrieverProgressEntityInterval + 1); got != retrieverProgressEntityInterval { + t.Fatalf("unexpected progress threshold: %d", got) + } +} + +func TestBenchFormattingHelpers(t *testing.T) { + if got := perSecond(10, 2*time.Second); got != 5 { + t.Fatalf("perSecond = %f", got) + } + if got := perSecond(10, 0); got != 0 { + t.Fatalf("perSecond with zero duration = %f", got) + } + + var buffer bytes.Buffer + writeBenchReport(&buffer, benchReport{ + Graphs: []benchGraphReport{{ + Name: "default", + Results: []benchResult{{ + Workers: 2, + BatchSize: 100, + SampleSize: 2, + NodeCount: 3, + EdgeCount: 4, + NodeProcessed: 2, + EdgeProcessed: 2, + TotalWallMillis: 50, + EntitiesPerSecond: 140, + NodeDBReadMillis: 10, + EdgeDBReadMillis: 20, + }}, + }}, + }) + 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"} { + if !strings.Contains(output, expected) { + t.Fatalf("bench report missing %q in %q", expected, output) + } + } +} + +func TestLogBenchPhaseProgressThresholds(t *testing.T) { + planned := retrieverProgressEntityInterval * 3 + nextProgressAt := retrieverProgressEntityInterval + startedAt := time.Now().Add(-time.Second) + + if got := logBenchPhaseProgress("default", retriever.PhaseNodes, 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{ + 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{ + 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{ + Count: planned, + }, planned, startedAt, nextProgressAt); got != nextProgressAt { + t.Fatalf("completed progress advanced to %d", got) + } + if got := logBenchPhaseProgress("default", retriever.PhaseNodes, 1, benchPhaseResult{ + Count: nextProgressAt, + }, planned, startedAt, 0); got != 0 { + t.Fatalf("disabled progress advanced to %d", got) + } +} + +func TestBenchBatchProcessorAggregatesConcurrentResults(t *testing.T) { + processor, _, err := newBenchBatchProcessor(context.Background(), 2, func(values []int) (benchPhaseResult, error) { + return benchPhaseResult{ + Count: int64(len(values)), + EncodeCompressTime: time.Duration(len(values)) * time.Millisecond, + UncompressedByteSize: int64(len(values) * 10), + CompressedByteSize: int64(len(values) * 5), + }, nil + }) + if err != nil { + t.Fatalf("create bench batch processor: %v", err) + } + + processor.addDBReadElapsed(3 * time.Millisecond) + if err := processor.handle([]int{1}); err != nil { + t.Fatalf("handle first batch: %v", err) + } + if err := processor.handle([]int{2, 3}); err != nil { + t.Fatalf("handle second batch: %v", err) + } + result, err := processor.closeAndWait() + if err != nil { + t.Fatalf("wait for bench batch processor: %v", err) + } + if result.Count != 3 { + t.Fatalf("processed count = %d", result.Count) + } + if result.DBReadElapsed != 3*time.Millisecond { + t.Fatalf("db read elapsed = %s", result.DBReadElapsed) + } + if result.EncodeCompressTime != 3*time.Millisecond { + t.Fatalf("encode/compress elapsed = %s", result.EncodeCompressTime) + } + if result.UncompressedByteSize != 30 || result.CompressedByteSize != 15 { + t.Fatalf("unexpected byte sizes: %+v", result) + } +} + +func TestBenchBatchCompression(t *testing.T) { + options := benchOptions{ + Compression: retriever.CompressionGzip, + ZstdLevel: retriever.DefaultZstdLevel, + } + 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")), + } + + nodeResult, err := benchNodeBatch(nodes, options) + if err != nil { + t.Fatalf("bench node batch: %v", err) + } + if nodeResult.Count != 2 || nodeResult.UncompressedByteSize <= 0 || nodeResult.CompressedByteSize <= 0 { + t.Fatalf("unexpected node batch result: %+v", nodeResult) + } + + noCompressionResult, err := benchNodeBatch(nodes, benchOptions{ + Compression: retriever.CompressionDisabled, + ZstdLevel: retriever.DefaultZstdLevel, + }) + if err != nil { + t.Fatalf("bench uncompressed node batch: %v", err) + } + if noCompressionResult.Count != 2 || noCompressionResult.UncompressedByteSize != 0 || noCompressionResult.CompressedByteSize != 0 { + t.Fatalf("unexpected no-compression node result: %+v", noCompressionResult) + } + + 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), + } + relationshipResult, err := benchRelationshipBatch(relationships, options) + if err != nil { + t.Fatalf("bench relationship batch: %v", err) + } + if relationshipResult.Count != 2 || relationshipResult.UncompressedByteSize <= 0 || relationshipResult.CompressedByteSize <= 0 { + t.Fatalf("unexpected relationship batch result: %+v", relationshipResult) + } +} diff --git a/cmd/retriever/config.go b/cmd/retriever/config.go new file mode 100644 index 00000000..a9aac300 --- /dev/null +++ b/cmd/retriever/config.go @@ -0,0 +1,151 @@ +package main + +import ( + "flag" + "fmt" + "os" + "strconv" + "strings" + + "github.com/specterops/dawgs/retriever" +) + +const defaultBenchSampleSize = 1_000_000 + +type stringList []string + +func (s *stringList) Set(value string) error { + trimmed := strings.TrimSpace(value) + if trimmed == "" { + return fmt.Errorf("value cannot be empty") + } + + *s = append(*s, trimmed) + + return nil +} + +func (s *stringList) String() string { + return strings.Join(*s, ",") +} + +type workerList []int + +func (s *workerList) Set(value string) error { + values, err := parseWorkerList(value) + if err != nil { + return err + } + + seen := make(map[int]struct{}, len(*s)+len(values)) + for _, value := range *s { + seen[value] = struct{}{} + } + + for _, value := range values { + if _, ok := seen[value]; ok { + continue + } + + seen[value] = struct{}{} + *s = append(*s, value) + } + + return nil +} + +func (s *workerList) String() string { + values := make([]string, 0, len(*s)) + for _, value := range *s { + values = append(values, strconv.Itoa(value)) + } + + return strings.Join(values, ",") +} + +func parseWorkerList(value string) ([]int, error) { + parts := strings.Split(value, ",") + workers := make([]int, 0, len(parts)) + seen := map[int]struct{}{} + + for _, part := range parts { + trimmed := strings.TrimSpace(part) + if trimmed == "" { + continue + } + + count, err := strconv.Atoi(trimmed) + if err != nil { + return nil, fmt.Errorf("parse worker count %q: %w", trimmed, err) + } + + if count <= 0 { + return nil, fmt.Errorf("worker counts must be > 0") + } + + if _, ok := seen[count]; ok { + continue + } + + seen[count] = struct{}{} + workers = append(workers, count) + } + + if len(workers) == 0 { + return nil, fmt.Errorf("at least one worker count is required") + } + + return workers, nil +} + +type benchOptions struct { + Workers []int + BatchSize int + SampleSize int + Compression retriever.CompressionCodec + ZstdLevel int + JSONOutput bool +} + +func (s benchOptions) validate() error { + if len(s.Workers) == 0 { + return fmt.Errorf("workers are required") + } + + for _, workerCount := range s.Workers { + if workerCount <= 0 { + return fmt.Errorf("worker counts must be > 0") + } + } + + if s.BatchSize <= 0 { + return fmt.Errorf("batch-size must be > 0") + } + + if s.SampleSize < 0 { + return fmt.Errorf("sample-size must be >= 0") + } + + if s.ZstdLevel <= 0 { + return fmt.Errorf("zstd-level must be > 0") + } + + if s.Compression != retriever.CompressionDisabled { + if err := retriever.ValidateCompression(s.Compression); err != nil { + return err + } + } + + return 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.") +} + +func fillConnectionFromEnv(cfg *databaseConfig) { + if strings.TrimSpace(cfg.Connection) == "" { + cfg.Connection = strings.TrimSpace(os.Getenv("CONNECTION_STRING")) + } +} diff --git a/cmd/retriever/config_test.go b/cmd/retriever/config_test.go new file mode 100644 index 00000000..041d0c07 --- /dev/null +++ b/cmd/retriever/config_test.go @@ -0,0 +1,112 @@ +package main + +import ( + "testing" + + "github.com/specterops/dawgs/retriever" +) + +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") + } +} + +func TestFlagListTypes(t *testing.T) { + var graphs stringList + 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") + } + + 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()) + } +} + +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") + } + + if workers.String() != "1,2,4" { + t.Fatalf("invalid worker update changed list to %q", workers.String()) + } +} + +func TestBenchOptionsValidate(t *testing.T) { + bench := benchOptions{ + Workers: []int{1}, + BatchSize: 1, + SampleSize: 1, + ZstdLevel: retriever.DefaultZstdLevel, + } + 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") + } + + 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") + } + + bench.Workers = []int{1} + bench.SampleSize = -1 + + if err := bench.validate(); err == nil { + t.Fatalf("expected invalid sample size") + } +} diff --git a/cmd/retriever/database.go b/cmd/retriever/database.go new file mode 100644 index 00000000..5c85760d --- /dev/null +++ b/cmd/retriever/database.go @@ -0,0 +1,214 @@ +package main + +import ( + "context" + "fmt" + "net/url" + "strings" + + "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" +) + +type databaseConfig struct { + Driver string + Connection string + Graph string +} + +func openDatabase(ctx context.Context, cfg databaseConfig) (graph.Database, string, error) { + connection := strings.TrimSpace(cfg.Connection) + if connection == "" { + return nil, "", fmt.Errorf("database connection is required; pass -connection or set CONNECTION_STRING") + } + + driverName := strings.TrimSpace(cfg.Driver) + if driverName == "" { + if inferredDriverName, err := driverFromConnectionString(connection); err != nil { + return nil, "", err + } else { + driverName = inferredDriverName + } + } + + openConfig := dawgs.Config{ + ConnectionString: connection, + GraphQueryMemoryLimit: size.Gibibyte, + } + + poolOwnedByDriver := false + + switch driverName { + case pg.DriverName: + if poolCfg, err := pgxpool.ParseConfig(connection); err != nil { + return nil, "", fmt.Errorf("parse PostgreSQL pool configuration: %w", err) + } else if pool, err := pg.NewPool(poolCfg); err != nil { + return nil, "", fmt.Errorf("open PostgreSQL pool: %w", err) + } else { + defer func() { + if !poolOwnedByDriver { + pool.Close() + } + }() + openConfig.Pool = pool + } + + case neo4j.DriverName: + // No driver-specific setup is required for Neo4j. + + default: + return nil, "", fmt.Errorf("unsupported driver %q; expected %s or %s", driverName, pg.DriverName, neo4j.DriverName) + } + + db, err := dawgs.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 graphName := strings.TrimSpace(cfg.Graph); graphName != "" { + if err := db.SetDefaultGraph(ctx, graph.Graph{ + Name: graphName, + }); err != nil { + return nil, "", fmt.Errorf("set graph target %q: %w", graphName, err) + } + } + + openSuccess = true + + return db, driverName, nil +} + +func driverFromConnectionString(connection string) (string, error) { + parsedURL, err := url.Parse(connection) + if err != nil { + return "", fmt.Errorf("parse connection string: %w", err) + } + + switch strings.ToLower(parsedURL.Scheme) { + case "postgres", "postgresql": + return pg.DriverName, nil + case neo4j.DriverName, "neo4j+s", "neo4j+ssc": + return neo4j.DriverName, nil + default: + return "", fmt.Errorf("unknown connection string scheme %q; expected postgres/postgresql or neo4j", parsedURL.Scheme) + } +} + +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") + } + + if allGraphs { + switch driverName { + case pg.DriverName: + return discoverPostgresGraphs(ctx, db) + case neo4j.DriverName: + return []retriever.GraphTarget{{ + Name: retriever.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 + } + + 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{}{} + + for _, name := range requested { + trimmed := strings.TrimSpace(name) + if trimmed == "" { + return nil, fmt.Errorf("graph name cannot be empty") + } + + if _, ok := seen[trimmed]; ok { + return nil, fmt.Errorf("duplicate graph target %q", trimmed) + } + + seen[trimmed] = struct{}{} + targets = append(targets, retriever.GraphTarget{ + Name: trimmed, + }) + } + + return targets, nil +} + +func discoverPostgresGraphs(ctx context.Context, db graph.Database) ([]retriever.GraphTarget, error) { + const graphQuery = ` +select + g.name, + to_regclass('node_' || g.id::text) is not null as has_node_partition, + to_regclass('edge_' || g.id::text) is not null as has_edge_partition +from graph g +order by g.name` + + var targets []retriever.GraphTarget + if err := db.ReadTransaction(ctx, func(tx graph.Transaction) error { + result := tx.Raw(graphQuery, nil) + defer result.Close() + + for result.Next() { + var ( + name string + hasNodePartition bool + hasEdgePartition bool + ) + + if err := result.Scan(&name, &hasNodePartition, &hasEdgePartition); err != nil { + return err + } + + if !hasNodePartition || !hasEdgePartition { + return fmt.Errorf("PostgreSQL graph %q is missing expected node/edge partitions", name) + } + + targets = append(targets, retriever.GraphTarget{ + Name: name, + }) + } + + return result.Error() + }); err != nil { + return nil, fmt.Errorf("discover PostgreSQL graphs: %w", err) + } + + if len(targets) == 0 { + return nil, fmt.Errorf("no PostgreSQL graphs were discovered") + } + + 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 new file mode 100644 index 00000000..2b0358fe --- /dev/null +++ b/cmd/retriever/database_test.go @@ -0,0 +1,89 @@ +package main + +import ( + "context" + "testing" + + "github.com/specterops/dawgs/drivers/neo4j" + "github.com/specterops/dawgs/drivers/pg" + "github.com/specterops/dawgs/retriever" +) + +func TestDriverFromConnectionString(t *testing.T) { + cases := map[string]string{ + "postgres://user:pass@example/db": pg.DriverName, + "postgresql://user:pass@example/db": pg.DriverName, + "neo4j://user:pass@example": neo4j.DriverName, + "neo4j+s://user:pass@example": neo4j.DriverName, + "neo4j+ssc://user:pass@example": neo4j.DriverName, + } + for connection, expected := range cases { + actual, err := driverFromConnectionString(connection) + 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") + } +} + +func TestResolveGraphTargets(t *testing.T) { + targets, err := resolveGraphTargets(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) + } + + targets, err = resolveGraphTargets(context.Background(), nil, pg.DriverName, []string{"a", "b"}, 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 _, err := resolveGraphTargets(context.Background(), nil, pg.DriverName, []string{"a", "a"}, false); err == nil { + t.Fatalf("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 := resolveGraphTargets(context.Background(), nil, neo4j.DriverName, []string{"a", "b"}, false); err == nil { + t.Fatalf("expected neo4j multi-graph error") + } + + targets, err = resolveGraphTargets(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) + } +} + +func TestGraphDirectoryName(t *testing.T) { + if got := graphDirectoryName("default"); got != "default" { + t.Fatalf("graphDirectoryName(default) = %q", got) + } + + if got := graphDirectoryName("graph/name"); got != "graph%2Fname" { + t.Fatalf("graphDirectoryName(graph/name) = %q", got) + } + + if got := graphDirectoryName(""); got != retriever.DefaultGraphName { + t.Fatalf("graphDirectoryName(empty) = %q", got) + } +} diff --git a/cmd/retriever/main.go b/cmd/retriever/main.go new file mode 100644 index 00000000..124b2ae5 --- /dev/null +++ b/cmd/retriever/main.go @@ -0,0 +1,417 @@ +package main + +import ( + "context" + "crypto/hpke" + "encoding/json" + "flag" + "fmt" + "io" + "os" + "strings" + + "github.com/specterops/dawgs/retriever" +) + +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. +` + +type commandRuntime struct { + stdout io.Writer + stderr io.Writer +} + +func main() { + runtime := commandRuntime{ + 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) + } +} + +func (s commandRuntime) run(ctx context.Context, args []string) error { + if len(args) == 0 { + fmt.Fprint(s.stderr, usage) + return fmt.Errorf("command is required") + } + + switch args[0] { + 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 "bench": + return s.runBench(ctx, args[1:]) + default: + fmt.Fprint(s.stderr, usage) + return fmt.Errorf("unknown command %q", args[0]) + } +} + +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 + ) + + 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.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 or gzip.") + 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.") + if err := flags.Parse(args); 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 strings.TrimSpace(scrubConfig) != "" { + file, err := os.Open(scrubConfig) + if err != nil { + return fmt.Errorf("open scrub config: %w", err) + } + defer file.Close() + + cfg.ScrubConfig = file + } + + if err := cfg.Validate(); err != nil { + return err + } + + var archiveRecipient hpke.PublicKey + if strings.TrimSpace(archiveOut) != "" { + if strings.TrimSpace(recipientPath) == "" { + return fmt.Errorf("-archive-out requires -recipient") + } + + var err error + archiveRecipient, err = retriever.LoadArchivePublicKey(recipientPath) + if err != nil { + return err + } + + if err := retriever.PreflightArchiveOutputPath(archiveOut); err != nil { + return err + } + } else if strings.TrimSpace(recipientPath) != "" { + return fmt.Errorf("-recipient requires -archive-out") + } + + db, driverName, err := openDatabase(ctx, dbCfg) + if err != nil { + return err + } + defer db.Close(ctx) + + targets, err := resolveGraphTargets(ctx, db, driverName, []string(graphs), *allGraphs) + if err != nil { + return err + } + + result, err := retriever.Dump(ctx, db, driverName, targets, cfg) + if err != nil { + return err + } + + var archiveLine string + if strings.TrimSpace(archiveOut) != "" { + if err := retriever.WriteEncryptedCollectionArchiveFile(cfg.OutputDir, archiveOut, archiveRecipient); err != nil { + return err + } + + archiveLine = fmt.Sprintf("archive: %s\n", archiveOut) + } + + 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) + return nil +} + +func (s commandRuntime) runKeygen(args []string) error { + var cfg retriever.KeygenOptions + + flags := flag.NewFlagSet("retriever keygen", flag.ContinueOnError) + flags.SetOutput(s.stderr) + flags.StringVar(&cfg.PrivatePath, "private", "", "Private key output path.") + flags.StringVar(&cfg.PublicPath, "public", "", "Public key output path.") + if err := flags.Parse(args); err != nil { + return err + } + + if err := cfg.Validate(); err != nil { + return err + } + + if err := retriever.Keygen(cfg); err != nil { + return err + } + + fmt.Fprintf(s.stdout, "private key: %s\npublic key: %s\n", cfg.PrivatePath, cfg.PublicPath) + return nil +} + +func (s commandRuntime) runUnpack(args []string) error { + var ( + archivePath string + identityPath string + outputDir string + force bool + ) + + flags := flag.NewFlagSet("retriever unpack", 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.") + 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") + } + + identity, err := retriever.LoadArchivePrivateKey(identityPath) + if err != nil { + return err + } + + if err := retriever.UnpackEncryptedCollectionArchiveFile(archivePath, outputDir, force, identity); err != nil { + return err + } + + fmt.Fprintf(s.stdout, "unpacked archive: %s\noutput: %s\n", archivePath, outputDir) + return nil +} + +func (s commandRuntime) runLoad(ctx context.Context, args []string) error { + var ( + dbCfg databaseConfig + cfg = retriever.DefaultLoadOptions("") + inputDir string + archivePath string + identityPath string + ) + + flags := flag.NewFlagSet("retriever load", 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.") + 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 err := cfg.Validate(); err != nil { + return err + } + + db, driverName, err := openDatabase(ctx, dbCfg) + if err != nil { + return err + } + defer db.Close(ctx) + + result, err := retriever.Load(ctx, db, driverName, cfg) + if err != nil { + return err + } + + fmt.Fprintf(s.stdout, "loaded %d graph(s)\nnodes: %d\nrelationships: %d\n", result.GraphCount, result.NodeCount, result.EdgeCount) + return nil +} + +func (s commandRuntime) runVerify(ctx context.Context, args []string) error { + var ( + dbCfg databaseConfig + cfg = retriever.DefaultVerifyOptions("") + ) + + flags := flag.NewFlagSet("retriever verify", 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.") + if err := flags.Parse(args); err != nil { + return err + } + + fillConnectionFromEnv(&dbCfg) + + if err := cfg.Validate(); err != nil { + return err + } + + db, driverName, err := openDatabase(ctx, dbCfg) + if err != nil { + return err + } + defer db.Close(ctx) + + result, err := retriever.Verify(ctx, db, driverName, cfg) + if err != nil { + return err + } + + fmt.Fprintf(s.stdout, "verified %d graph(s)\nnodes: %d\nrelationships: %d\n", result.GraphCount, result.NodeCount, result.EdgeCount) + return nil +} + +func (s commandRuntime) runBench(ctx context.Context, args []string) error { + var ( + dbCfg databaseConfig + cfg benchOptions + + graphs stringList + workers workerList + compressionVal string + ) + + cfg.BatchSize = retriever.DefaultBatchSize + cfg.SampleSize = defaultBenchSampleSize + cfg.ZstdLevel = retriever.DefaultZstdLevel + + flags := flag.NewFlagSet("retriever bench", 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 or gzip.") + flags.IntVar(&cfg.ZstdLevel, "zstd-level", cfg.ZstdLevel, "zstd compression level.") + flags.BoolVar(&cfg.JSONOutput, "json", false, "Emit machine-readable JSON.") + 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 { + return err + } + + db, driverName, err := openDatabase(ctx, dbCfg) + if err != nil { + return err + } + defer db.Close(ctx) + + targets, err := resolveGraphTargets(ctx, db, driverName, []string(graphs), *allGraphs) + if err != nil { + return err + } + + report, err := Bench(ctx, db, driverName, targets, cfg) + if err != nil { + return err + } + + if cfg.JSONOutput { + encoder := json.NewEncoder(s.stdout) + encoder.SetIndent("", " ") + + return encoder.Encode(report) + } + + writeBenchReport(s.stdout, report) + return nil +} diff --git a/cmd/retriever/main_test.go b/cmd/retriever/main_test.go new file mode 100644 index 00000000..7e265278 --- /dev/null +++ b/cmd/retriever/main_test.go @@ -0,0 +1,93 @@ +package main + +import ( + "bytes" + "context" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/specterops/dawgs/retriever" +) + +func TestCommandRuntimeHelpAndValidation(t *testing.T) { + runtime := commandRuntime{ + stdout: &bytes.Buffer{}, + stderr: &bytes.Buffer{}, + } + if err := runtime.run(context.Background(), []string{"help"}); err != nil { + 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) { + t.Fatalf("help output missing %s command", command) + } + } + + 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) + } + + 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) + } + + 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) + } + + 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) + } + + 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) + } + + 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) + } + + 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) + } +} + +func TestDumpArchiveOutputPreflightBeforeDatabase(t *testing.T) { + runtime := commandRuntime{ + stdout: &bytes.Buffer{}, + stderr: &bytes.Buffer{}, + } + 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) + } + + archivePath := filepath.Join(dir, "dump.tar.pq") + if err := os.WriteFile(archivePath, []byte("exists"), 0o600); err != nil { + t.Fatalf("write existing archive: %v", err) + } + + err := runtime.run(context.Background(), []string{ + "dump", + "-out", filepath.Join(dir, "dump"), + "-archive-out", archivePath, + "-recipient", publicPath, + }) + if err == nil || !strings.Contains(err.Error(), "already exists") { + t.Fatalf("expected archive preflight error before database open, got %v", err) + } +} diff --git a/cmd/retriever/packaging_plan.md b/cmd/retriever/packaging_plan.md new file mode 100644 index 00000000..58dbc18d --- /dev/null +++ b/cmd/retriever/packaging_plan.md @@ -0,0 +1,192 @@ +# Retriever Packaging And Encryption Plan + +## Goal + +Add optional packaging and encryption for retriever dumps. + +The package should preserve the existing manifest-based dump collection while +allowing operators to produce a single encrypted archive suitable for storage or +transfer. + +Assumption: "QC resistant" means quantum-computer-resistant cryptography. + +## Recommended Design + +Use a raw TAR packaging layer plus an encrypted streaming envelope: + +```text +retriever dump -> dump directory -> raw tar stream -> HPKE/ML-KEM encrypted archive +``` + +Default crypto suite: + +```text +HPKE: + KEM: ML-KEM-1024 + KDF: HKDF-SHA512 + AEAD: AES-256-GCM +``` + +NIST FIPS 203 standardizes ML-KEM for post-quantum key establishment. The +current Go toolchain exposes `crypto/mlkem` and `crypto/hpke`, so this should +not require a third-party post-quantum implementation. + +## CLI Shape + +Generate recipient keys: + +```bash +retriever keygen \ + -private retriever-private.key \ + -public retriever-public.key +``` + +Create an encrypted archive during dump: + +```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 +``` + +Unpack an encrypted archive: + +```bash +retriever unpack \ + -archive ./dump.tar.pq \ + -identity ./retriever-private.key \ + -out ./dumpdir +``` + +Optional later convenience: + +```bash +retriever load \ + -archive ./dump.tar.pq \ + -identity ./retriever-private.key \ + -connection "$CONNECTION_STRING" +``` + +## Archive Format + +Use uncompressed TAR only: + +- No `.tar.gz`. +- No `.tgz`. +- No zstd archive compression. +- Archive packaging uses Go's `archive/tar`. + +The existing fragment compression setting remains independent. If the full +plaintext dump also needs uncompressed fragments, add that separately as +`-compression none`. + +TAR writer rules: + +- Deterministic sorted paths. +- Normalized slash paths. +- No absolute paths. +- No `..` path traversal. +- No symlinks. +- Regular files only. +- Stable TAR metadata: + - uid/gid `0`. + - uname/gname empty. + - constrained file modes. + - mtime fixed to zero. +- Include `manifest.json` and fragment files. +- Exclude the archive output path if it lives under the dump directory. + +## Encrypted Envelope + +Create a retriever-specific archive envelope: + +```text +magic: RTRV-PQ-ARCHIVE-v1 +header_len: uint32 +header: JSON +frames: encrypted chunk frames +final: authenticated final frame +``` + +Header fields: + +```json +{ + "format": "retriever-encrypted-tar-v1", + "archive": { + "format": "tar", + "compression": "none" + }, + "crypto": { + "scheme": "hpke", + "kem": "ML-KEM-1024", + "kdf": "HKDF-SHA512", + "aead": "AES-256-GCM", + "encapsulated_key": "base64..." + }, + "chunk_size": 1048576 +} +``` + +Encrypt the TAR stream in chunks. Each chunk should authenticate additional data +containing: + +- Envelope version. +- Header hash. +- Frame index. +- Final/non-final marker. + +Require an authenticated final frame so truncation at a frame boundary is +detected. + +## Implementation Sequence + +1. Add key file types and `retriever keygen`. +2. Add key load and parse validation. +3. Add deterministic raw TAR packer. +4. Add TAR unpacker with path traversal protections. +5. Add encrypted streaming writer using an HPKE sender context. +6. Add encrypted streaming reader using an HPKE recipient context. +7. Wire `dump -archive-out -recipient`. +8. Add `retriever unpack`. +9. Optionally wire `load -archive -identity`. +10. Add `slog` progress for archive walking, TAR streaming, encryption, + finalization, decryption, and unpacking. +11. Document key handling and archive workflow. + +## Testing Plan + +Pure tests should cover: + +- TAR path sorting is deterministic. +- TAR metadata does not leak host user/group. +- TAR rejects symlinks, absolute paths, and `..`. +- Encrypted archive header contains no graph names or manifest contents. +- `decrypt(encrypt(tar))` round-trips byte-for-byte. +- Wrong key fails. +- Tampered header fails. +- Tampered frame fails. +- Truncated archive fails. +- Missing final frame fails. + +Integration-style tests without a database should cover: + +- Create a synthetic dump directory. +- Package and encrypt it. +- Unpack it. +- Compare file tree and checksums. +- Verify `readManifest` still works after unpack. + +Database integration can remain separate and gated behind `CONNECTION_STRING`. + +## References + +- NIST FIPS 203, Module-Lattice-Based Key-Encapsulation Mechanism Standard: + https://csrc.nist.gov/pubs/fips/203/final +- NIST announcement of finalized post-quantum encryption standards: + https://www.nist.gov/news-events/news/2024/08/nist-releases-first-3-finalized-post-quantum-encryption-standards diff --git a/cmd/retriever/retriever_integration_test.go b/cmd/retriever/retriever_integration_test.go new file mode 100644 index 00000000..bb004e80 --- /dev/null +++ b/cmd/retriever/retriever_integration_test.go @@ -0,0 +1,195 @@ +// Copyright 2026 Specter Ops, Inc. +// +// Licensed under the Apache License, Version 2.0 +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// SPDX-License-Identifier: Apache-2.0 + +//go:build manual_integration + +package main + +import ( + "context" + "errors" + "fmt" + "os" + "strings" + "testing" + "time" + + "github.com/specterops/dawgs/drivers/pg" + "github.com/specterops/dawgs/graph" + "github.com/specterops/dawgs/retriever" +) + +func TestPostgreSQLDumpLoadRoundTrip(t *testing.T) { + connection := os.Getenv("CONNECTION_STRING") + if connection == "" { + t.Skip("CONNECTION_STRING not set") + } + driverName, err := driverFromConnectionString(connection) + if err != nil { + t.Fatalf("infer driver: %v", err) + } + if driverName != pg.DriverName { + t.Skipf("CONNECTION_STRING selects %s, not PostgreSQL", driverName) + } + + ctx := context.Background() + db, _, err := openDatabase(ctx, databaseConfig{ + Connection: connection, + }) + if err != nil { + t.Fatalf("open database: %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}, + } + if err := db.AssertSchema(ctx, graph.Schema{ + Graphs: []graph.Graph{graphSchema}, + DefaultGraph: graphSchema, + }); err != nil { + t.Fatalf("assert schema: %v", err) + } + t.Cleanup(func() { + _ = db.WriteTransaction(ctx, func(tx graph.Transaction) error { + return tx.WithGraph(graph.Graph{ + Name: graphName, + }).Nodes().Delete() + }) + }) + + if err := db.WriteTransaction(ctx, func(tx graph.Transaction) error { + tx = tx.WithGraph(graph.Graph{ + Name: graphName, + }) + alice, err := tx.CreateNode(graph.AsProperties(map[string]any{"name": "alice"}), userKind) + if err != nil { + return err + } + system, err := tx.CreateNode(graph.AsProperties(map[string]any{"name": "server"}), systemKind) + if err != nil { + return err + } + _, err = tx.CreateRelationshipByIDs(alice.ID, system.ID, adminKind, graph.AsProperties(map[string]any{"source": "test"})) + return err + }); err != nil { + t.Fatalf("seed graph: %v", err) + } + + entitySnapshot, err := countGraphEntitySnapshot(ctx, db, graph.Graph{ + Name: graphName, + }) + if err != nil { + t.Fatalf("count graph entities: %v", err) + } + if entitySnapshot.NodeCount != 2 || entitySnapshot.EdgeCount != 1 { + t.Skipf("graph target %q is not isolated by the current PostgreSQL query path: nodes=%d edges=%d", graphName, entitySnapshot.NodeCount, entitySnapshot.EdgeCount) + } + + 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: 1, + BatchSize: 1, + }) + if err != nil { + t.Fatalf("dump: %v", err) + } + if dumpResult.NodeCount != 2 || dumpResult.EdgeCount != 1 { + t.Fatalf("unexpected dump counts: nodes=%d edges=%d", dumpResult.NodeCount, dumpResult.EdgeCount) + } + 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) + } + if dumpResult.Manifest.Metrics.Graphs[0].NodeCount != 2 || dumpResult.Manifest.Metrics.Graphs[0].EdgeCount != 1 { + t.Fatalf("unexpected dump metrics counts: %+v", dumpResult.Manifest.Metrics.Graphs[0]) + } + + if _, err := retriever.Load(ctx, db, driverName, retriever.LoadOptions{ + InputDir: dumpDir, + BatchSize: 1, + VerifyMetrics: true, + }); err == nil || !strings.Contains(err.Error(), "is not empty") { + t.Fatalf("expected non-empty target load error, got %v", err) + } + + if err := db.WriteTransaction(ctx, func(tx graph.Transaction) error { + return tx.WithGraph(graph.Graph{ + Name: graphName, + }).Nodes().Delete() + }); err != nil { + t.Fatalf("clear graph before load: %v", err) + } + + loadResult, err := retriever.Load(ctx, db, driverName, retriever.LoadOptions{ + InputDir: dumpDir, + BatchSize: 1, + VerifyMetrics: true, + }) + if err != nil { + t.Fatalf("load: %v", err) + } + if loadResult.NodeCount != 2 || loadResult.EdgeCount != 1 { + t.Fatalf("unexpected load counts: nodes=%d edges=%d", loadResult.NodeCount, loadResult.EdgeCount) + } + + verifyResult, err := retriever.Verify(ctx, db, driverName, retriever.VerifyOptions{ + InputDir: dumpDir, + BatchSize: 1, + }) + if err != nil { + t.Fatalf("verify: %v", err) + } + if verifyResult.NodeCount != 2 || verifyResult.EdgeCount != 1 { + t.Fatalf("unexpected verify counts: nodes=%d edges=%d", verifyResult.NodeCount, verifyResult.EdgeCount) + } + + 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) + return err + }); err != nil { + t.Fatalf("mutate loaded graph: %v", err) + } + + _, err = retriever.Verify(ctx, db, driverName, retriever.VerifyOptions{ + InputDir: dumpDir, + BatchSize: 1, + }) + 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) + } +} diff --git a/go.mod b/go.mod index e55abafa..1f380c05 100644 --- a/go.mod +++ b/go.mod @@ -1,35 +1,37 @@ module github.com/specterops/dawgs -go 1.25.7 +go 1.26.4 require ( - cuelang.org/go v0.16.0 - github.com/RoaringBitmap/roaring/v2 v2.16.0 + cuelang.org/go v0.17.0 + github.com/RoaringBitmap/roaring/v2 v2.19.0 github.com/antlr4-go/antlr/v4 v4.13.1 github.com/axiomhq/hyperloglog v0.2.6 - github.com/bits-and-blooms/bitset v1.24.4 + github.com/bits-and-blooms/bitset v1.24.5 github.com/cespare/xxhash/v2 v2.3.0 github.com/fzipp/gocyclo v0.6.0 github.com/gammazero/deque v1.2.1 github.com/jackc/pgtype v1.14.4 - github.com/jackc/pgx/v5 v5.9.2 + 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/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/tools v0.44.0 + golang.org/x/tools v0.47.0 ) // Dawgrun requirements require ( - github.com/alecthomas/chroma/v2 v2.23.1 + github.com/alecthomas/chroma/v2 v2.27.0 github.com/charmbracelet/lipgloss v1.1.0 github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510 - github.com/jedib0t/go-pretty/v6 v6.7.8 + github.com/jedib0t/go-pretty/v6 v6.8.2 github.com/kanmu/go-sqlfmt v0.0.2-0.20200215095417-d1e63e2ee5eb github.com/mitchellh/go-wordwrap v1.0.1 github.com/specterops/go-repl v1.0.1 - golang.org/x/term v0.43.0 + golang.org/x/term v0.44.0 ) require ( @@ -72,19 +74,21 @@ require ( github.com/catenacyber/perfsprint v0.10.1 // indirect github.com/ccojocar/zxcvbn-go v1.0.4 // indirect github.com/charithe/durationcheck v0.0.11 // indirect - github.com/charmbracelet/colorprofile v0.2.3-0.20250311203215-f60798e515dc // indirect - github.com/charmbracelet/x/ansi v0.10.1 // indirect - github.com/charmbracelet/x/cellbuf v0.0.13-0.20250311204145-2c3ea96c31dd // indirect - github.com/charmbracelet/x/term v0.2.1 // indirect + github.com/charmbracelet/colorprofile v0.4.3 // indirect + github.com/charmbracelet/x/ansi v0.11.7 // indirect + github.com/charmbracelet/x/cellbuf v0.0.15 // indirect + github.com/charmbracelet/x/term v0.2.2 // indirect github.com/ckaznocha/intrange v0.3.1 // indirect - github.com/clipperhouse/uax29/v2 v2.2.0 // indirect - github.com/cockroachdb/apd/v3 v3.2.2 // indirect + github.com/clipperhouse/displaywidth v0.11.0 // indirect + github.com/clipperhouse/uax29/v2 v2.7.0 // indirect + github.com/cockroachdb/apd/v3 v3.2.3 // indirect github.com/curioswitch/go-reassign v0.3.0 // indirect github.com/daixiang0/gci v0.13.7 // indirect github.com/dave/dst v0.27.3 // indirect github.com/denis-tingaikin/go-header v0.5.0 // indirect github.com/dgryski/go-metro v0.0.0-20250106013310-edb8663e5e33 // indirect - github.com/dlclark/regexp2 v1.11.5 // indirect + github.com/dlclark/regexp2 v1.12.0 // indirect + github.com/dlclark/regexp2/v2 v2.2.2 // indirect github.com/ettle/strcase v0.2.0 // indirect github.com/fatih/color v1.19.0 // indirect github.com/fatih/structtag v1.2.0 // indirect @@ -148,7 +152,7 @@ require ( github.com/ldez/tagliatelle v0.7.2 // indirect github.com/ldez/usetesting v0.5.0 // indirect github.com/leonklingele/grouper v1.1.2 // indirect - github.com/lucasb-eyer/go-colorful v1.2.0 // indirect + github.com/lucasb-eyer/go-colorful v1.4.0 // indirect github.com/macabu/inamedparam v0.2.0 // indirect github.com/manuelarte/embeddedstructfieldcheck v0.4.0 // indirect github.com/manuelarte/funcorder v0.5.0 // indirect @@ -156,8 +160,8 @@ require ( github.com/maratori/testpackage v1.1.2 // indirect github.com/matoous/godox v1.1.0 // indirect github.com/mattn/go-colorable v0.1.14 // indirect - github.com/mattn/go-isatty v0.0.21 // indirect - github.com/mattn/go-runewidth v0.0.19 // indirect + github.com/mattn/go-isatty v0.0.22 // indirect + github.com/mattn/go-runewidth v0.0.24 // indirect github.com/mgechev/revive v1.15.0 // indirect github.com/mitchellh/go-homedir v1.1.0 // indirect github.com/moricho/tparallel v0.3.2 // indirect @@ -168,7 +172,6 @@ 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/pelletier/go-toml/v2 v2.2.4 // 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 @@ -182,7 +185,7 @@ require ( github.com/quasilyte/stdinfo v0.0.0-20220114132959-f7386bf02567 // indirect github.com/raeperd/recvcheck v0.2.0 // indirect github.com/rivo/uniseg v0.4.7 // indirect - github.com/rogpeppe/go-internal v1.14.1 // indirect + github.com/rogpeppe/go-internal v1.15.0 // indirect github.com/ryancurrah/gomodguard v1.4.1 // indirect github.com/ryanrolds/sqlclosecheck v0.6.0 // indirect github.com/sagikazarmark/locafero v0.12.0 // indirect @@ -226,14 +229,13 @@ require ( go.uber.org/multierr v1.11.0 // indirect go.uber.org/zap v1.27.1 // indirect go.yaml.in/yaml/v3 v3.0.4 // indirect - golang.org/x/crypto v0.52.0 // indirect - golang.org/x/exp v0.0.0-20260410095643-746e56fc9e2f // indirect + golang.org/x/crypto v0.53.0 // indirect + golang.org/x/exp v0.0.0-20260611194520-c48552f49976 // indirect golang.org/x/exp/typeparams v0.0.0-20260209203927-2842357ff358 // indirect - golang.org/x/mod v0.35.0 // indirect - golang.org/x/net v0.54.0 // indirect - golang.org/x/sync v0.20.0 // indirect - golang.org/x/sys v0.45.0 // indirect - golang.org/x/text v0.37.0 // 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 honnef.co/go/tools v0.7.0 // indirect diff --git a/go.sum b/go.sum index 9799273b..bfc4bb5c 100644 --- a/go.sum +++ b/go.sum @@ -6,10 +6,10 @@ codeberg.org/chavacava/garif v0.2.0 h1:F0tVjhYbuOCnvNcU3YSpO6b3Waw6Bimy4K0mM8y6M codeberg.org/chavacava/garif v0.2.0/go.mod h1:P2BPbVbT4QcvLZrORc2T29szK3xEOlnl0GiPTJmEqBQ= codeberg.org/polyfloyd/go-errorlint v1.9.0 h1:VkdEEmA1VBpH6ecQoMR4LdphVI3fA4RrCh2an7YmodI= codeberg.org/polyfloyd/go-errorlint v1.9.0/go.mod h1:GPRRu2LzVijNn4YkrZYJfatQIdS+TrcK8rL5Xs24qw8= -cuelabs.dev/go/oci/ociregistry v0.0.0-20251212221603-3adeb8663819 h1:Zh+Ur3OsoWpvALHPLT45nOekHkgOt+IOfutBbPqM17I= -cuelabs.dev/go/oci/ociregistry v0.0.0-20251212221603-3adeb8663819/go.mod h1:WjmQxb+W6nVNCgj8nXrF24lIz95AHwnSl36tpjDZSU8= -cuelang.org/go v0.16.0 h1:mmt9SL/IzfSIiBKuP5wxdO4xLjvIHr3urpbjCDdMV5U= -cuelang.org/go v0.16.0/go.mod h1:4veMX+GpsK0B91b1seGXoozG80LJCczvG1M1Re/knxo= +cuelabs.dev/go/oci/ociregistry v0.0.0-20260601085548-328ff8e2c943 h1:XUtzi/yWlmuy8V6kkmVbbmirmUqcFe9Ce3gmEaHXf1Q= +cuelabs.dev/go/oci/ociregistry v0.0.0-20260601085548-328ff8e2c943/go.mod h1:WjmQxb+W6nVNCgj8nXrF24lIz95AHwnSl36tpjDZSU8= +cuelang.org/go v0.17.0 h1:PrijS5ofUD01yiG11w74I04laXKLaBiMhEYvdt8Gb/A= +cuelang.org/go v0.17.0/go.mod h1:xlly/o1wSLvxOsi5vkQGieU0rLOt7TvUIizOFtnxHRU= dev.gaijin.team/go/exhaustruct/v4 v4.0.0 h1:873r7aNneqoBB3IaFIzhvt2RFYTuHgmMjoKfwODoI1Y= dev.gaijin.team/go/exhaustruct/v4 v4.0.0/go.mod h1:aZ/k2o4Y05aMJtiux15x8iXaumE88YdiB0Ai4fXOzPI= dev.gaijin.team/go/golib v0.6.0 h1:v6nnznFTs4bppib/NyU1PQxobwDHwCXXl15P7DV5Zgo= @@ -40,12 +40,12 @@ github.com/MirrexOne/unqueryvet v1.5.4 h1:38QOxShO7JmMWT+eCdDMbcUgGCOeJphVkzzRgy github.com/MirrexOne/unqueryvet v1.5.4/go.mod h1:fs9Zq6eh1LRIhsDIsxf9PONVUjYdFHdtkHIgZdJnyPU= github.com/OpenPeeDeeP/depguard/v2 v2.2.1 h1:vckeWVESWp6Qog7UZSARNqfu/cZqvki8zsuj3piCMx4= github.com/OpenPeeDeeP/depguard/v2 v2.2.1/go.mod h1:q4DKzC4UcVaAvcfd41CZh0PWpGgzrVxUYBlgKNGquUo= -github.com/RoaringBitmap/roaring/v2 v2.16.0 h1:Kys1UNf49d5W8Tq3bpuAhIr/Z8/yPB+59CO8A6c/BbE= -github.com/RoaringBitmap/roaring/v2 v2.16.0/go.mod h1:eq4wdNXxtJIS/oikeCzdX1rBzek7ANzbth041hrU8Q4= +github.com/RoaringBitmap/roaring/v2 v2.19.0 h1:zsWtVE+biht4eVl0YDLvykUqGkftR3Qc5UCbeKB4UQ4= +github.com/RoaringBitmap/roaring/v2 v2.19.0/go.mod h1:SfT3of9nYh3vis1dIbCj4Yw6KQGujTN+f345nrN/0JA= github.com/alecthomas/assert/v2 v2.11.0 h1:2Q9r3ki8+JYXvGsDyBXwH3LcJ+WK5D0gc5E8vS6K3D0= github.com/alecthomas/assert/v2 v2.11.0/go.mod h1:Bze95FyfUr7x34QZrjL+XP+0qgp/zg8yS+TtBj1WA3k= -github.com/alecthomas/chroma/v2 v2.23.1 h1:nv2AVZdTyClGbVQkIzlDm/rnhk1E9bU9nXwmZ/Vk/iY= -github.com/alecthomas/chroma/v2 v2.23.1/go.mod h1:NqVhfBR0lte5Ouh3DcthuUCTUpDC9cxBOfyMbMQPs3o= +github.com/alecthomas/chroma/v2 v2.27.0 h1:FodwmyOBgJULFYmDqibcp9pvfDLWdtPRh9v/r5BXYZs= +github.com/alecthomas/chroma/v2 v2.27.0/go.mod h1:NjJ3ciIgrqBNeIkWZ4e46nseoLDslxU1LmfCoL+wcY8= github.com/alecthomas/go-check-sumtype v0.3.1 h1:u9aUvbGINJxLVXiFvHUlPEaD7VDULsrxJb4Aq31NLkU= github.com/alecthomas/go-check-sumtype v0.3.1/go.mod h1:A8TSiN3UPRw3laIgWEUOHHLPa6/r9MtoigdlP5h3K/E= github.com/alecthomas/repr v0.5.2 h1:SU73FTI9D1P5UNtvseffFSGmdNci/O6RsqzeXJtP0Qs= @@ -72,8 +72,8 @@ github.com/aymanbagabas/go-osc52/v2 v2.0.1 h1:HwpRHbFMcZLEVr42D4p7XBqjyuxQH5SMiE github.com/aymanbagabas/go-osc52/v2 v2.0.1/go.mod h1:uYgXzlJ7ZpABp8OJ+exZzJJhRNQ2ASbcXHWsFqH8hp8= github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM= github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw= -github.com/bits-and-blooms/bitset v1.24.4 h1:95H15Og1clikBrKr/DuzMXkQzECs1M6hhoGXLwLQOZE= -github.com/bits-and-blooms/bitset v1.24.4/go.mod h1:7hO7Gc7Pp1vODcmWvKMRA9BNmbv6a/7QIWpPxHddWR8= +github.com/bits-and-blooms/bitset v1.24.5 h1:654xBVHc23gJMAgOTkPNoCVfiRxuIOAUnAZFtopqJ4w= +github.com/bits-and-blooms/bitset v1.24.5/go.mod h1:7hO7Gc7Pp1vODcmWvKMRA9BNmbv6a/7QIWpPxHddWR8= github.com/bkielbasa/cyclop v1.2.3 h1:faIVMIGDIANuGPWH031CZJTi2ymOQBULs9H21HSMa5w= github.com/bkielbasa/cyclop v1.2.3/go.mod h1:kHTwA9Q0uZqOADdupvcFJQtp/ksSnytRMe8ztxG8Fuo= github.com/blizzy78/varnamelen v0.8.0 h1:oqSblyuQvFsW1hbBHh1zfwrKe3kcSj0rnXkKzsQ089M= @@ -98,23 +98,25 @@ github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UF github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= github.com/charithe/durationcheck v0.0.11 h1:g1/EX1eIiKS57NTWsYtHDZ/APfeXKhye1DidBcABctk= github.com/charithe/durationcheck v0.0.11/go.mod h1:x5iZaixRNl8ctbM+3B2RrPG5t856TxRyVQEnbIEM2X4= -github.com/charmbracelet/colorprofile v0.2.3-0.20250311203215-f60798e515dc h1:4pZI35227imm7yK2bGPcfpFEmuY1gc2YSTShr4iJBfs= -github.com/charmbracelet/colorprofile v0.2.3-0.20250311203215-f60798e515dc/go.mod h1:X4/0JoqgTIPSFcRA/P6INZzIuyqdFY5rm8tb41s9okk= +github.com/charmbracelet/colorprofile v0.4.3 h1:QPa1IWkYI+AOB+fE+mg/5/4HRMZcaXex9t5KX76i20Q= +github.com/charmbracelet/colorprofile v0.4.3/go.mod h1:/zT4BhpD5aGFpqQQqw7a+VtHCzu+zrQtt1zhMt9mR4Q= github.com/charmbracelet/lipgloss v1.1.0 h1:vYXsiLHVkK7fp74RkV7b2kq9+zDLoEU4MZoFqR/noCY= github.com/charmbracelet/lipgloss v1.1.0/go.mod h1:/6Q8FR2o+kj8rz4Dq0zQc3vYf7X+B0binUUBwA0aL30= -github.com/charmbracelet/x/ansi v0.10.1 h1:rL3Koar5XvX0pHGfovN03f5cxLbCF2YvLeyz7D2jVDQ= -github.com/charmbracelet/x/ansi v0.10.1/go.mod h1:3RQDQ6lDnROptfpWuUVIUG64bD2g2BgntdxH0Ya5TeE= -github.com/charmbracelet/x/cellbuf v0.0.13-0.20250311204145-2c3ea96c31dd h1:vy0GVL4jeHEwG5YOXDmi86oYw2yuYUGqz6a8sLwg0X8= -github.com/charmbracelet/x/cellbuf v0.0.13-0.20250311204145-2c3ea96c31dd/go.mod h1:xe0nKWGd3eJgtqZRaN9RjMtK7xUYchjzPr7q6kcvCCs= -github.com/charmbracelet/x/term v0.2.1 h1:AQeHeLZ1OqSXhrAWpYUtZyX1T3zVxfpZuEQMIQaGIAQ= -github.com/charmbracelet/x/term v0.2.1/go.mod h1:oQ4enTYFV7QN4m0i9mzHrViD7TQKvNEEkHUMCmsxdUg= +github.com/charmbracelet/x/ansi v0.11.7 h1:kzv1kJvjg2S3r9KHo8hDdHFQLEqn4RBCb39dAYC84jI= +github.com/charmbracelet/x/ansi v0.11.7/go.mod h1:9qGpnAVYz+8ACONkZBUWPtL7lulP9No6p1epAihUZwQ= +github.com/charmbracelet/x/cellbuf v0.0.15 h1:ur3pZy0o6z/R7EylET877CBxaiE1Sp1GMxoFPAIztPI= +github.com/charmbracelet/x/cellbuf v0.0.15/go.mod h1:J1YVbR7MUuEGIFPCaaZ96KDl5NoS0DAWkskup+mOY+Q= +github.com/charmbracelet/x/term v0.2.2 h1:xVRT/S2ZcKdhhOuSP4t5cLi5o+JxklsoEObBSgfgZRk= +github.com/charmbracelet/x/term v0.2.2/go.mod h1:kF8CY5RddLWrsgVwpw4kAa6TESp6EB5y3uxGLeCqzAI= github.com/ckaznocha/intrange v0.3.1 h1:j1onQyXvHUsPWujDH6WIjhyH26gkRt/txNlV7LspvJs= github.com/ckaznocha/intrange v0.3.1/go.mod h1:QVepyz1AkUoFQkpEqksSYpNpUo3c5W7nWh/s6SHIJJk= -github.com/clipperhouse/uax29/v2 v2.2.0 h1:ChwIKnQN3kcZteTXMgb1wztSgaU+ZemkgWdohwgs8tY= -github.com/clipperhouse/uax29/v2 v2.2.0/go.mod h1:EFJ2TJMRUaplDxHKj1qAEhCtQPW2tJSwu5BF98AuoVM= +github.com/clipperhouse/displaywidth v0.11.0 h1:lBc6kY44VFw+TDx4I8opi/EtL9m20WSEFgwIwO+UVM8= +github.com/clipperhouse/displaywidth v0.11.0/go.mod h1:bkrFNkf81G8HyVqmKGxsPufD3JhNl3dSqnGhOoSD/o0= +github.com/clipperhouse/uax29/v2 v2.7.0 h1:+gs4oBZ2gPfVrKPthwbMzWZDaAFPGYK72F0NJv2v7Vk= +github.com/clipperhouse/uax29/v2 v2.7.0/go.mod h1:EFJ2TJMRUaplDxHKj1qAEhCtQPW2tJSwu5BF98AuoVM= github.com/cockroachdb/apd v1.1.0/go.mod h1:8Sl8LxpKi29FqWXR16WEFZRNSz3SoPzUzeMeY4+DwBQ= -github.com/cockroachdb/apd/v3 v3.2.2 h1:R1VaDQkMR321HBM6+6b2eYZfxi0ybPJgUh0Ztr7twzU= -github.com/cockroachdb/apd/v3 v3.2.2/go.mod h1:klXJcjp+FffLTHlhIG69tezTDvdP065naDsHzKhYSqc= +github.com/cockroachdb/apd/v3 v3.2.3 h1:4Zx+I3R35bFXMnltzmjP79i2cravE4jTRL6ps9Aux80= +github.com/cockroachdb/apd/v3 v3.2.3/go.mod h1:klXJcjp+FffLTHlhIG69tezTDvdP065naDsHzKhYSqc= github.com/coreos/go-systemd v0.0.0-20190321100706-95778dfbb74e/go.mod h1:F5haX7vjVVG0kc13fIWeqUViNPyEJxv/OmvnBo0Yme4= github.com/coreos/go-systemd v0.0.0-20190719114852-fd7a80b32e1f/go.mod h1:F5haX7vjVVG0kc13fIWeqUViNPyEJxv/OmvnBo0Yme4= github.com/cpuguy83/go-md2man/v2 v2.0.6/go.mod h1:oOW0eioCTA6cOiMLiUPZOpcVxMig6NIQQ7OS05n1F4g= @@ -135,8 +137,10 @@ github.com/denis-tingaikin/go-header v0.5.0 h1:SRdnP5ZKvcO9KKRP1KJrhFR3RrlGuD+42 github.com/denis-tingaikin/go-header v0.5.0/go.mod h1:mMenU5bWrok6Wl2UsZjy+1okegmwQ3UgWl4V1D8gjlY= github.com/dgryski/go-metro v0.0.0-20250106013310-edb8663e5e33 h1:ucRHb6/lvW/+mTEIGbvhcYU3S8+uSNkuMjx/qZFfhtM= github.com/dgryski/go-metro v0.0.0-20250106013310-edb8663e5e33/go.mod h1:c9O8+fpSOX1DM8cPNSkX/qsBWdkD4yd2dpciOWQjpBw= -github.com/dlclark/regexp2 v1.11.5 h1:Q/sSnsKerHeCkc/jSTNq1oCm7KiVgUMZRDUoRu0JQZQ= -github.com/dlclark/regexp2 v1.11.5/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8= +github.com/dlclark/regexp2 v1.12.0 h1:0j4c5qQmnC6XOWNjP3PIXURXN2gWx76rd3KvgdPkCz8= +github.com/dlclark/regexp2 v1.12.0/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8= +github.com/dlclark/regexp2/v2 v2.2.2 h1:MYWvNYw8okuqNhwTYO587EZMiDruVa2vhV6fsGpfya0= +github.com/dlclark/regexp2/v2 v2.2.2/go.mod h1:avUrQvPaLz2DrFNHJF0taWAFFX2C1GMSSoeiqFjcBmU= github.com/emicklei/proto v1.14.3 h1:zEhlzNkpP8kN6utonKMzlPfIvy82t5Kb9mufaJxSe1Q= github.com/emicklei/proto v1.14.3/go.mod h1:rn1FgRS/FANiZdD2djyH7TMA9jdRDcYQ9IEN9yvjX0A= github.com/ettle/strcase v0.2.0 h1:fGNiVF21fHXpX1niBgk0aROov1LagYsOwV/xqKDKR/Q= @@ -163,8 +167,8 @@ github.com/go-kit/log v0.1.0/go.mod h1:zbhenjAZHb184qTLMA9ZjW7ThYL0H2mk7Q6pNt4vb github.com/go-logfmt/logfmt v0.5.0/go.mod h1:wCYkCAKZfumFQihp8CzCvQ3paCTfi41vtzG1KdI/P7A= github.com/go-logr/logr v1.4.3 h1:CjnDlHq8ikf6E492q6eKboGOC0T8CDaOvkHCIg8idEI= github.com/go-logr/logr v1.4.3/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY= -github.com/go-quicktest/qt v1.101.0 h1:O1K29Txy5P2OK0dGo59b7b0LR6wKfIhttaAhHUyn7eI= -github.com/go-quicktest/qt v1.101.0/go.mod h1:14Bz/f7NwaXPtdYEgzsx46kqSxVwTbzVZsDC26tQJow= +github.com/go-quicktest/qt v1.102.0 h1:HSQxCeh5YZH3EL3W39ixjtyaEhcWSXQHtHnMBzSs474= +github.com/go-quicktest/qt v1.102.0/go.mod h1:p4lGIVX+8Wa6ZPNDvqcxq36XpUDLh42FLetFU7odllI= github.com/go-stack/stack v1.8.0/go.mod h1:v0f6uXyyMGvRgIKkXu+yp6POWl0qKG85gN/melR3HDY= github.com/go-task/slim-sprig/v3 v3.0.0 h1:sUs3vkvUymDpBKi3qH1YSqBQk9+9D/8M2mN1vB6EwHI= github.com/go-task/slim-sprig/v3 v3.0.0/go.mod h1:W848ghGpv3Qj3dhTPRyJypKRiqCdHZiAzKg9hl15HA8= @@ -306,16 +310,16 @@ github.com/jackc/pgx/v4 v4.0.0-pre1.0.20190824185557-6972a5742186/go.mod h1:X+GQ github.com/jackc/pgx/v4 v4.12.1-0.20210724153913-640aa07df17c/go.mod h1:1QD0+tgSXP7iUjYm9C1NxKhny7lq6ee99u/z+IHFcgs= github.com/jackc/pgx/v4 v4.18.2 h1:xVpYkNR5pk5bMCZGfClbO962UIqVABcAGt7ha1s/FeU= github.com/jackc/pgx/v4 v4.18.2/go.mod h1:Ey4Oru5tH5sB6tV7hDmfWFahwF15Eb7DNXlRKx2CkVw= -github.com/jackc/pgx/v5 v5.9.2 h1:3ZhOzMWnR4yJ+RW1XImIPsD1aNSz4T4fyP7zlQb56hw= -github.com/jackc/pgx/v5 v5.9.2/go.mod h1:mal1tBGAFfLHvZzaYh77YS/eC6IX9OWbRV1QIIM0Jn4= +github.com/jackc/pgx/v5 v5.10.0 h1:VhSvgU2jSli8o3AqIEOTJr7rZwAEUVo4E4XhR94Zfr0= +github.com/jackc/pgx/v5 v5.10.0/go.mod h1:mal1tBGAFfLHvZzaYh77YS/eC6IX9OWbRV1QIIM0Jn4= github.com/jackc/puddle v0.0.0-20190413234325-e4ced69a3a2b/go.mod h1:m4B5Dj62Y0fbyuIc15OsIqK0+JU8nkqQjsgx7dvjSWk= github.com/jackc/puddle v0.0.0-20190608224051-11cab39313c9/go.mod h1:m4B5Dj62Y0fbyuIc15OsIqK0+JU8nkqQjsgx7dvjSWk= github.com/jackc/puddle v1.1.3/go.mod h1:m4B5Dj62Y0fbyuIc15OsIqK0+JU8nkqQjsgx7dvjSWk= github.com/jackc/puddle v1.3.0/go.mod h1:m4B5Dj62Y0fbyuIc15OsIqK0+JU8nkqQjsgx7dvjSWk= github.com/jackc/puddle/v2 v2.2.2 h1:PR8nw+E/1w0GLuRFSmiioY6UooMp6KJv0/61nB7icHo= github.com/jackc/puddle/v2 v2.2.2/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4= -github.com/jedib0t/go-pretty/v6 v6.7.8 h1:BVYrDy5DPBA3Qn9ICT+PokP9cvCv1KaHv2i+Hc8sr5o= -github.com/jedib0t/go-pretty/v6 v6.7.8/go.mod h1:YwC5CE4fJ1HFUDeivSV1r//AmANFHyqczZk+U6BDALU= +github.com/jedib0t/go-pretty/v6 v6.8.2 h1:FmKNr1GOyot/zqNQplE8HLhFguJaeHJTCArntnI4uxE= +github.com/jedib0t/go-pretty/v6 v6.8.2/go.mod h1:YwC5CE4fJ1HFUDeivSV1r//AmANFHyqczZk+U6BDALU= github.com/jgautheron/goconst v1.8.2 h1:y0XF7X8CikZ93fSNT6WBTb/NElBu9IjaY7CCYQrCMX4= github.com/jgautheron/goconst v1.8.2/go.mod h1:A0oxgBCHy55NQn6sYpO7UdnA9p+h7cPtoOZUmvNIako= github.com/jingyugao/rowserrcheck v1.1.1 h1:zibz55j/MJtLsjP1OF4bSdgXxwL1b+Vn7Tjzq7gFzUs= @@ -335,6 +339,8 @@ github.com/kisielk/errcheck v1.10.0/go.mod h1:kQxWMMVZgIkDq7U8xtG/n2juOjbLgZtedi github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+oQHNcck= github.com/kkHAIKE/contextcheck v1.1.6 h1:7HIyRcnyzxL9Lz06NGhiKvenXq7Zw6Q0UQu/ttjfJCE= github.com/kkHAIKE/contextcheck v1.1.6/go.mod h1:3dDbMRNBFaq8HFXWC1JyvDSPm43CmE6IuHam8Wr0rkg= +github.com/klauspost/compress v1.19.0 h1:sXLILfc9jV2QYWkzFOPWStmcUVH2RHEB1JCdY2oVvCQ= +github.com/klauspost/compress v1.19.0/go.mod h1:cwPg85FWrGar70rWktvGQj8/hthj3wpl0PGDogxkrSQ= github.com/konsorten/go-windows-terminal-sequences v1.0.1/go.mod h1:T0+1ngSBFLxvqU3pZ+m/2kptfBszLMUkC4ZK/EgS/cQ= github.com/konsorten/go-windows-terminal-sequences v1.0.2/go.mod h1:T0+1ngSBFLxvqU3pZ+m/2kptfBszLMUkC4ZK/EgS/cQ= github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo= @@ -371,8 +377,8 @@ github.com/lib/pq v1.2.0/go.mod h1:5WUZQaWbwv1U+lTReE5YruASi9Al49XbQIvNi/34Woo= github.com/lib/pq v1.10.2/go.mod h1:AlVN5x4E4T544tWzH6hKfbfQvm3HdbOxrmggDNAPY9o= github.com/lib/pq v1.11.2 h1:x6gxUeu39V0BHZiugWe8LXZYZ+Utk7hSJGThs8sdzfs= github.com/lib/pq v1.11.2/go.mod h1:/p+8NSbOcwzAEI7wiMXFlgydTwcgTr3OSKMsD2BitpA= -github.com/lucasb-eyer/go-colorful v1.2.0 h1:1nnpGOrhyZZuNyfu1QjKiUICQ74+3FNCN69Aj6K7nkY= -github.com/lucasb-eyer/go-colorful v1.2.0/go.mod h1:R4dSotOR9KMtayYi1e77YzuveK+i7ruzyGqttikkLy0= +github.com/lucasb-eyer/go-colorful v1.4.0 h1:UtrWVfLdarDgc44HcS7pYloGHJUjHV/4FwW4TvVgFr4= +github.com/lucasb-eyer/go-colorful v1.4.0/go.mod h1:R4dSotOR9KMtayYi1e77YzuveK+i7ruzyGqttikkLy0= github.com/macabu/inamedparam v0.2.0 h1:VyPYpOc10nkhI2qeNUdh3Zket4fcZjEWe35poddBCpE= github.com/macabu/inamedparam v0.2.0/go.mod h1:+Pee9/YfGe5LJ62pYXqB89lJ+0k5bsR8Wgz/C0Zlq3U= github.com/manuelarte/embeddedstructfieldcheck v0.4.0 h1:3mAIyaGRtjK6EO9E73JlXLtiy7ha80b2ZVGyacxgfww= @@ -394,10 +400,10 @@ github.com/mattn/go-colorable v0.1.14/go.mod h1:6LmQG8QLFO4G5z1gPvYEzlUgJ2wF+stg github.com/mattn/go-isatty v0.0.5/go.mod h1:Iq45c/XA43vh69/j3iqttzPXn0bhXyGjM0Hdxcsrc5s= github.com/mattn/go-isatty v0.0.7/go.mod h1:Iq45c/XA43vh69/j3iqttzPXn0bhXyGjM0Hdxcsrc5s= github.com/mattn/go-isatty v0.0.12/go.mod h1:cbi8OIDigv2wuxKPP5vlRcQ1OAZbq2CE4Kysco4FUpU= -github.com/mattn/go-isatty v0.0.21 h1:xYae+lCNBP7QuW4PUnNG61ffM4hVIfm+zUzDuSzYLGs= -github.com/mattn/go-isatty v0.0.21/go.mod h1:ZXfXG4SQHsB/w3ZeOYbR0PrPwLy+n6xiMrJlRFqopa4= -github.com/mattn/go-runewidth v0.0.19 h1:v++JhqYnZuu5jSKrk9RbgF5v4CGUjqRfBm05byFGLdw= -github.com/mattn/go-runewidth v0.0.19/go.mod h1:XBkDxAl56ILZc9knddidhrOlY5R/pDhgLpndooCuJAs= +github.com/mattn/go-isatty v0.0.22 h1:j8l17JJ9i6VGPUFUYoTUKPSgKe/83EYU2zBC7YNKMw4= +github.com/mattn/go-isatty v0.0.22/go.mod h1:ZXfXG4SQHsB/w3ZeOYbR0PrPwLy+n6xiMrJlRFqopa4= +github.com/mattn/go-runewidth v0.0.24 h1:cpokDiIn0MGnhdHwuWnJBITySJ20QyNGnY2kR/ay2DU= +github.com/mattn/go-runewidth v0.0.24/go.mod h1:XBkDxAl56ILZc9knddidhrOlY5R/pDhgLpndooCuJAs= github.com/mgechev/revive v1.15.0 h1:vJ0HzSBzfNyPbHKolgiFjHxLek9KUijhqh42yGoqZ8Q= github.com/mgechev/revive v1.15.0/go.mod h1:LlAKO3QQe9OJ0pVZzI2GPa8CbXGZ/9lNpCGvK4T/a8A= github.com/mitchellh/go-homedir v1.1.0 h1:lukF9ziXFxDFPkA1vsr5zpc1XuPDn/wFntq5mG+4E0Y= @@ -439,8 +445,8 @@ github.com/otiai10/mint v1.3.0/go.mod h1:F5AjcsTsWUqX+Na9fpHb52P8pcRX2CI6A3ctIT9 github.com/otiai10/mint v1.3.1/go.mod h1:/yxELlJQ0ufhjUwhshSj+wFjZ78CnZ48/1wtmBH1OTc= 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.2.4 h1:mye9XuhQ6gvn5h28+VilKrrPoQVanw5PMw/TB0t5Ec4= -github.com/pelletier/go-toml/v2 v2.2.4/go.mod h1:2gIqNv+qfxSVS7cM2xJQKtLSTLUE9V8t9Stt+h56mCY= +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/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= @@ -455,8 +461,8 @@ github.com/prometheus/common v0.63.0 h1:YR/EIY1o3mEFP/kZCD7iDMnLPlGyuU2Gb3HIcXnA github.com/prometheus/common v0.63.0/go.mod h1:VVFF/fBIoToEnWRVkYoXEkq3R3paCoxG9PXP74SnV18= github.com/prometheus/procfs v0.20.1 h1:XwbrGOIplXW/AU3YhIhLODXMJYyC1isLFfYCsTEycfc= github.com/prometheus/procfs v0.20.1/go.mod h1:o9EMBZGRyvDrSPH1RqdxhojkuXstoe4UlK79eF5TGGo= -github.com/protocolbuffers/txtpbfmt v0.0.0-20260217160748-a481f6a22f94 h1:2PC6Ql3jipz1KvBlqUHjjk6v4aMwE86mfDu1XMH0LR8= -github.com/protocolbuffers/txtpbfmt v0.0.0-20260217160748-a481f6a22f94/go.mod h1:JSbkp0BviKovYYt9XunS95M3mLPibE9bGg+Y95DsEEY= +github.com/protocolbuffers/txtpbfmt v0.0.0-20260420112717-c39628bde8b5 h1:Mckui8l+Wqz2Ve7XQvsE8SbHNmDWu8NA7Xce5NFJ/kM= +github.com/protocolbuffers/txtpbfmt v0.0.0-20260420112717-c39628bde8b5/go.mod h1:JSbkp0BviKovYYt9XunS95M3mLPibE9bGg+Y95DsEEY= github.com/quasilyte/go-ruleguard v0.4.5 h1:AGY0tiOT5hJX9BTdx/xBdoCubQUAE2grkqY2lSwvZcA= github.com/quasilyte/go-ruleguard v0.4.5/go.mod h1:Vl05zJ538vcEEwu16V/Hdu7IYZWyKSwIy4c88Ro1kRE= github.com/quasilyte/go-ruleguard/dsl v0.3.23 h1:lxjt5B6ZCiBeeNO8/oQsegE6fLeCzuMRoVWSkXC4uvY= @@ -472,8 +478,8 @@ github.com/raeperd/recvcheck v0.2.0/go.mod h1:n04eYkwIR0JbgD73wT8wL4JjPC3wm0nFtz github.com/rivo/uniseg v0.4.7 h1:WUdvkW8uEhrYfLC4ZzdpI2ztxP1I582+49Oc5Mq64VQ= github.com/rivo/uniseg v0.4.7/go.mod h1:FN3SvrM+Zdj16jyLfmOkMNblXMcoc8DfTHruCPUcx88= github.com/rogpeppe/go-internal v1.3.0/go.mod h1:M8bDsm7K2OlrFYOpmOWEs/qY81heoFRclV5y23lUDJ4= -github.com/rogpeppe/go-internal v1.14.1 h1:UQB4HGPB6osV0SQTLymcB4TgvyWu6ZyliaW0tI/otEQ= -github.com/rogpeppe/go-internal v1.14.1/go.mod h1:MaRKkUm5W0goXpeCfT7UZI6fk/L7L7so1lCWt35ZSgc= +github.com/rogpeppe/go-internal v1.15.0 h1:D0RCU5rMAp+SpgkiNdrjfJ+LX4J1M32V2NeCY7EJ6hc= +github.com/rogpeppe/go-internal v1.15.0/go.mod h1:DrUVZyrJU+txYW5/1kwtXQSMFio52ZOxX7yM1VHvnxs= github.com/rs/xid v1.2.1/go.mod h1:+uKXf+4Djp6Md1KODXJxgGQPKngRmWyn10oCKFzNHOQ= github.com/rs/zerolog v1.13.0/go.mod h1:YbFCdg8HfsridGWAh22vktObvhZbQsZXe4/zB0OKkWU= github.com/rs/zerolog v1.15.0/go.mod h1:xYTKnLHcpfU2225ny5qZjxnj9NvkumZYjJHlAThCjNc= @@ -632,10 +638,10 @@ golang.org/x/crypto v0.13.0/go.mod h1:y6Z2r+Rw4iayiXXAIxJIDAJ1zMW4yaTpebo8fPOliY golang.org/x/crypto v0.14.0/go.mod h1:MVFd36DqK4CsrnJYDkBA3VC4m2GkXAM0PvzMCn4JQf4= golang.org/x/crypto v0.19.0/go.mod h1:Iy9bg/ha4yyC70EfRS8jz+B6ybOBKMaSxLj6P6oBDfU= golang.org/x/crypto v0.20.0/go.mod h1:Xwo95rrVNIoSMx9wa1JroENMToLWn3RNVrTBpLHgZPQ= -golang.org/x/crypto v0.52.0 h1:RMs7fP2rXdep0CftQlK8Uf+kibLm7qkCcradZWYz988= -golang.org/x/crypto v0.52.0/go.mod h1:1QgfPxDqh0T2M/elOJtp9RvuR95kVjir0e6/BvEmGbc= -golang.org/x/exp v0.0.0-20260410095643-746e56fc9e2f h1:W3F4c+6OLc6H2lb//N1q4WpJkhzJCK5J6kUi1NTVXfM= -golang.org/x/exp v0.0.0-20260410095643-746e56fc9e2f/go.mod h1:J1xhfL/vlindoeF/aINzNzt2Bket5bjo9sdOYzOsU80= +golang.org/x/crypto v0.53.0 h1:QZ4Muo8THX6CizN2vPPd5fBGHyogrdK9fG4wLPFUsto= +golang.org/x/crypto v0.53.0/go.mod h1:DNLU434OwVakk9PzuwV8w62mAJpRJL3vsgcfp4Qnsio= +golang.org/x/exp v0.0.0-20260611194520-c48552f49976 h1:X8Hz2ImujgbmetVuW+w2YkyZChE3cBpZi2P158rTG9M= +golang.org/x/exp v0.0.0-20260611194520-c48552f49976/go.mod h1:vnf4pv9iKZXY58sQE1L86zmNWJ4159e1RkcWiLCkeEY= golang.org/x/exp/typeparams v0.0.0-20220428152302-39d4317da171/go.mod h1:AbB0pIl9nAr9wVwH+Z2ZpaocVmF5I4GyWCDIsVjR0bk= golang.org/x/exp/typeparams v0.0.0-20230203172020-98cc5a0785f9/go.mod h1:AbB0pIl9nAr9wVwH+Z2ZpaocVmF5I4GyWCDIsVjR0bk= golang.org/x/exp/typeparams v0.0.0-20260209203927-2842357ff358 h1:qWFG1Dj7TBjOjOvhEOkmyGPVoquqUKnIU0lEVLp8xyk= @@ -652,8 +658,8 @@ golang.org/x/mod v0.6.0-dev.0.20220419223038-86c51ed26bb4/go.mod h1:jJ57K6gSWd91 golang.org/x/mod v0.8.0/go.mod h1:iBbtSCu2XBx23ZKBPSOrRkjjQPZFPuis4dIYUhu/chs= golang.org/x/mod v0.12.0/go.mod h1:iBbtSCu2XBx23ZKBPSOrRkjjQPZFPuis4dIYUhu/chs= golang.org/x/mod v0.13.0/go.mod h1:hTbmBsO62+eylJbnUtE2MGJUyE7QWk4xUqPFrRgJ+7c= -golang.org/x/mod v0.35.0 h1:Ww1D637e6Pg+Zb2KrWfHQUnH2dQRLBQyAtpr/haaJeM= -golang.org/x/mod v0.35.0/go.mod h1:+GwiRhIInF8wPm+4AoT6L0FA1QWAad3OMdTRx4tFYlU= +golang.org/x/mod v0.37.0 h1:vF1DjpVEshcIqoEaauuHebaLk1O1forxjxBaVn884JQ= +golang.org/x/mod v0.37.0/go.mod h1:m8S8VeM9r4dzDwjrKO0a1sZP3YjeMamRRlD+fmR2Q/0= golang.org/x/net v0.0.0-20190311183353-d8887717615a/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg= golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg= golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s= @@ -670,10 +676,10 @@ golang.org/x/net v0.10.0/go.mod h1:0qNGK6F8kojg2nk9dLZ2mShWaEBan6FAoqfSigmmuDg= golang.org/x/net v0.15.0/go.mod h1:idbUs1IY1+zTqbi8yxTbhexhEEk5ur9LInksu6HrEpk= golang.org/x/net v0.16.0/go.mod h1:NxSsAGuq816PNPmqtQdLE42eU2Fs7NoRIZrHJAlaCOE= golang.org/x/net v0.21.0/go.mod h1:bIjVDfnllIU7BJ2DNgfnXvpSvtn8VRwhlsaeUTyUS44= -golang.org/x/net v0.54.0 h1:2zJIZAxAHV/OHCDTCOHAYehQzLfSXuf/5SoL/Dv6w/w= -golang.org/x/net v0.54.0/go.mod h1:Sj4oj8jK6XmHpBZU/zWHw3BV3abl4Kvi+Ut7cQcY+cQ= -golang.org/x/oauth2 v0.35.0 h1:Mv2mzuHuZuY2+bkyWXIHMfhNdJAdwW3FuWeCPYN5GVQ= -golang.org/x/oauth2 v0.35.0/go.mod h1:lzm5WQJQwKZ3nwavOZ3IS5Aulzxi68dUSgRHujetwEA= +golang.org/x/net v0.56.0 h1:Rw8j/hFzGvJUZwNBXnAtf5sVDVt+65SK2C7IxCxZt5o= +golang.org/x/net v0.56.0/go.mod h1:D3Ku6r+V6JROoZK144D2XfMHFcMq/0zSfLelVTCFKec= +golang.org/x/oauth2 v0.36.0 h1:peZ/1z27fi9hUOFCAZaHyrpWG5lwe0RJEEEeH0ThlIs= +golang.org/x/oauth2 v0.36.0/go.mod h1:YDBUJMTkDnJS+A4BP4eZBjCqtokkg1hODuPjwiGPO7Q= golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.0.0-20190911185100-cd5d95a43a6e/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.0.0-20200625203802-6e8e738ad208/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= @@ -683,8 +689,8 @@ golang.org/x/sync v0.0.0-20220722155255-886fb9371eb4/go.mod h1:RxMgew5VJxzue5/jJ golang.org/x/sync v0.1.0/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.3.0/go.mod h1:FU7BRWz2tNW+3quACPkgCx/L+uEAv1htQ0V83Z9Rj+Y= golang.org/x/sync v0.4.0/go.mod h1:FU7BRWz2tNW+3quACPkgCx/L+uEAv1htQ0V83Z9Rj+Y= -golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4= -golang.org/x/sync v0.20.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= +golang.org/x/sync v0.21.0 h1:HLII4xRRTtCRkxYp4HNFF0Js/Og6q2i++KXbg0gHCwM= +golang.org/x/sync v0.21.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= golang.org/x/sys v0.0.0-20180905080454-ebe1bf3edb33/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.0.0-20190222072716-a9d3bda3a223/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= @@ -712,8 +718,8 @@ golang.org/x/sys v0.8.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.12.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.13.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.17.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= -golang.org/x/sys v0.45.0 h1:dO4czNzziLiiXplLQgBCEpCvXQ3dnkn0SdaZSYdQ+FY= -golang.org/x/sys v0.45.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= +golang.org/x/sys v0.46.0 h1:noSf2Fq6F8DBgS+LysIkx7rIExoNHJsxOAtPp4rthXw= +golang.org/x/sys v0.46.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= golang.org/x/term v0.0.0-20201117132131-f5c789dd3221/go.mod h1:Nr5EML6q2oocZ2LXRh80K7BxOlk5/8JxuGnuhpl+muw= golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo= golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8= @@ -722,8 +728,8 @@ golang.org/x/term v0.8.0/go.mod h1:xPskH00ivmX89bAKVGSKKtLOWNx2+17Eiy94tnKShWo= golang.org/x/term v0.12.0/go.mod h1:owVbMEjm3cBLCHdkQu9b1opXd4ETQWc3BhuQGKgXgvU= golang.org/x/term v0.13.0/go.mod h1:LTmsnFJwVN6bCy1rVCoS+qHT1HhALEFxKncY3WNNh4U= golang.org/x/term v0.17.0/go.mod h1:lLRBjIVuehSbZlaOtGMbcMncT+aqLLLmKrsjNrUguwk= -golang.org/x/term v0.43.0 h1:S4RLU2sB31O/NCl+zFN9Aru9A/Cq2aqKpTZJ6B+DwT4= -golang.org/x/term v0.43.0/go.mod h1:lrhlHNdQJHO+1qVYiHfFKVuVioJIheAc3fBSMFYEIsk= +golang.org/x/term v0.44.0 h1:0rLvDRCtNj0gZkyIXhCyOb2OAzEhLVqc4B+hrsBhrmc= +golang.org/x/term v0.44.0/go.mod h1:7ze4MdzUzLXpSAoFP1H0bOI9aXDqveSvatT5vKcFh2Y= golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= golang.org/x/text v0.3.2/go.mod h1:bEr9sfX3Q8Zfm5fL9x+3itogRgK3+ptLWKqgva+5dAk= golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= @@ -734,8 +740,8 @@ golang.org/x/text v0.7.0/go.mod h1:mrYo+phRRbMaCq/xk9113O4dZlRixOauAjOtrjsXDZ8= golang.org/x/text v0.9.0/go.mod h1:e1OnstbJyHTd6l/uOt8jFFHp6TRDWZR/bV3emEE/zU8= golang.org/x/text v0.13.0/go.mod h1:TvPlkZtksWOMsz7fbANvkp4WM8x/WCo/om8BMLbz+aE= golang.org/x/text v0.14.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU= -golang.org/x/text v0.37.0 h1:Cqjiwd9eSg8e0QAkyCaQTNHFIIzWtidPahFWR83rTrc= -golang.org/x/text v0.37.0/go.mod h1:a5sjxXGs9hsn/AJVwuElvCAo9v8QYLzvavO5z2PiM38= +golang.org/x/text v0.39.0 h1:UbZz4pLOvn600D6Oh6GGEI6VAmndrEBLv8/6BEXzyus= +golang.org/x/text v0.39.0/go.mod h1:3UwRclnC2g0TU9x8PZiyfOajCd1zaUNHF9cvqcQZ+ZM= golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= golang.org/x/tools v0.0.0-20190311212946-11955173bddd/go.mod h1:LCzVGOaR6xXOjkQ3onu1FJEFr0SW1gC7cKk1uF8kGRs= golang.org/x/tools v0.0.0-20190425163242-31fd60d6bfdc/go.mod h1:RgjU9mgBXZiqYHBnxXauZ1Gv1EHHAz9KjViQ78xBX0Q= @@ -755,8 +761,8 @@ golang.org/x/tools v0.1.12/go.mod h1:hNGJHUnrk76NpqgfD5Aqm5Crs+Hm0VOH/i9J2+nxYbc golang.org/x/tools v0.6.0/go.mod h1:Xwgl3UAJ/d3gWutnCtw505GrjyAbvKui8lOU390QaIU= golang.org/x/tools v0.13.0/go.mod h1:HvlwmtVNQAhOuCjW7xxvovg8wbNq7LwfXh/k7wXUl58= golang.org/x/tools v0.14.0/go.mod h1:uYBEerGOWcJyEORxN+Ek8+TT266gXkNlHdJBwexUsBg= -golang.org/x/tools v0.44.0 h1:UP4ajHPIcuMjT1GqzDWRlalUEoY+uzoZKnhOjbIPD2c= -golang.org/x/tools v0.44.0/go.mod h1:KA0AfVErSdxRZIsOVipbv3rQhVXTnlU6UhKxHd1seDI= +golang.org/x/tools v0.47.0 h1:7Kn5x/d1svx/PzryTsqeoZN4TZwqeH5pGWjefhLi/1Q= +golang.org/x/tools v0.47.0/go.mod h1:dFHnyTvFWY212G+h7ZY4Vsp/K3U4/7W9TyVaAul8uCA= golang.org/x/tools/go/expect v0.1.1-deprecated h1:jpBZDwmgPhXsKZC6WhL20P4b/wmnpsEAGHaNy0n/rJM= golang.org/x/tools/go/expect v0.1.1-deprecated/go.mod h1:eihoPOH+FgIqa3FpoTwguz/bVUSGBlGQU67vpBeOrBY= golang.org/x/tools/go/packages/packagestest v0.1.1-deprecated h1:1h2MnaIAIXISqTFKdENegdpAgUXz6NrPEsbIeWaBRvM= diff --git a/query/model.go b/query/model.go index c5b85887..76159756 100644 --- a/query/model.go +++ b/query/model.go @@ -369,8 +369,18 @@ func Where(expression graph.Criteria) *cypherModel.Where { } func OrderBy(leaves ...graph.Criteria) *cypherModel.Order { + items := make([]*cypherModel.SortItem, 0, len(leaves)) + + for _, leaf := range leaves { + if sortItem, ok := leaf.(*cypherModel.SortItem); ok { + items = append(items, sortItem) + } else { + items = append(items, Order(leaf, Ascending())) + } + } + return &cypherModel.Order{ - Items: convertCriteria[*cypherModel.SortItem](leaves...), + Items: items, } } diff --git a/query/neo4j/neo4j_test.go b/query/neo4j/neo4j_test.go index 2efc3435..b6adea1e 100644 --- a/query/neo4j/neo4j_test.go +++ b/query/neo4j/neo4j_test.go @@ -438,6 +438,20 @@ func TestQueryBuilder_Render(t *testing.T) { ), ), "match (n) where (n)<-[]->() return n order by n.value asc")) + t.Run("Node has Relationships Order by Direct Node ID", assertQueryResult(query.SinglePartQuery( + query.Where( + query.HasRelationships(query.Node()), + ), + + query.Returning( + query.Node(), + ), + + query.OrderBy( + query.NodeID(), + ), + ), "match (n) where (n)<-[]->() return n order by id(n) asc")) + t.Run("Node has Relationships Order by Node Item", assertQueryResult(query.SinglePartQuery( query.Where( query.HasRelationships(query.Node()), diff --git a/retriever/archive_envelope.go b/retriever/archive_envelope.go new file mode 100644 index 00000000..45bbe657 --- /dev/null +++ b/retriever/archive_envelope.go @@ -0,0 +1,843 @@ +package retriever + +import ( + "bytes" + "crypto/hpke" + "crypto/sha256" + "encoding/base64" + "encoding/binary" + "encoding/json" + "errors" + "fmt" + "io" + "log/slog" + "math" + "os" + "path/filepath" + "strings" + "time" +) + +var removeUnpackBackupDirectory = os.RemoveAll + +const ( + encryptedArchiveMagic = "RTRV-PQ-ARCHIVE-v1" + encryptedArchiveFormat = "retriever-encrypted-tar-v1" + encryptedArchiveHPKEInfo = "retriever/archive/hpke/v1" + encryptedArchiveChunkSize = 1024 * 1024 + maxEncryptedArchiveFrameSize = encryptedArchiveChunkSize + 4096 + maxEncryptedArchiveHeaderLen = 64 * 1024 + + encryptedArchiveFrameData byte = 0 + encryptedArchiveFrameFinal byte = 1 +) + +type encryptedArchiveHeader struct { + Format string `json:"format"` + Archive encryptedArchiveFormatHeader `json:"archive"` + Crypto encryptedArchiveCryptoHeader `json:"crypto"` + ChunkSize int `json:"chunk_size"` +} + +type encryptedArchiveFormatHeader struct { + Format string `json:"format"` + Compression string `json:"compression"` +} + +type encryptedArchiveCryptoHeader struct { + Scheme string `json:"scheme"` + KEM string `json:"kem"` + KDF string `json:"kdf"` + AEAD string `json:"aead"` + EncapsulatedKey string `json:"encapsulated_key"` +} + +type encryptedArchiveWriter struct { + writer io.Writer + sender *hpke.Sender + headerHash [sha256.Size]byte + frameIndex uint64 + closed bool +} + +type encryptedArchiveReader struct { + reader io.Reader + recipient *hpke.Recipient + headerHash [sha256.Size]byte + frameIndex uint64 + plaintext []byte + final bool +} + +func WriteEncryptedCollectionArchive(writer io.Writer, dumpDir string, recipient hpke.PublicKey, excludePaths ...string) error { + return WriteEncryptedCollectionArchiveWithOptions(writer, dumpDir, recipient, ArchiveOptions{}, excludePaths...) +} + +func WriteEncryptedCollectionArchiveWithOptions(writer io.Writer, dumpDir string, recipient hpke.PublicKey, options ArchiveOptions, excludePaths ...string) error { + startedAt := time.Now() + options.Progress.emit(ProgressEvent{ + Operation: OperationArchive, + Message: "retriever encrypted archive streaming started", + InputDir: dumpDir, + }) + + archiveWriter, err := newEncryptedArchiveWriter(writer, recipient) + if err != nil { + return err + } + + if err := WriteCollectionTarWithOptions(archiveWriter, dumpDir, options, excludePaths...); err != nil { + _ = archiveWriter.Close() + return err + } + + if err := archiveWriter.Close(); err != nil { + return err + } + + options.Progress.emit(ProgressEvent{ + Operation: OperationArchive, + Message: "retriever encrypted archive streaming completed", + InputDir: dumpDir, + Elapsed: time.Since(startedAt), + }) + return nil +} + +func WriteEncryptedCollectionArchiveFile(dumpDir, archivePath string, recipient hpke.PublicKey) error { + return WriteEncryptedCollectionArchiveFileWithOptions(dumpDir, archivePath, recipient, ArchiveOptions{}) +} + +func writeEncryptedCollectionArchive(dumpDir, archivePath string, recipient hpke.PublicKey) error { + return WriteEncryptedCollectionArchiveFile(dumpDir, archivePath, recipient) +} + +func WriteEncryptedCollectionArchiveFileWithOptions(dumpDir, archivePath string, recipient hpke.PublicKey, options ArchiveOptions) error { + archivePath = strings.TrimSpace(archivePath) + if err := preflightArchiveOutputPath(archivePath); err != nil { + return err + } + + tempFile, err := os.CreateTemp(filepath.Dir(archivePath), filepath.Base(archivePath)+".*.tmp") + if err != nil { + return fmt.Errorf("create archive temp file: %w", err) + } + tempPath := tempFile.Name() + cleanupTemp := true + defer func() { + if cleanupTemp { + _ = os.Remove(tempPath) + } + }() + + slog.Info("retriever archive encryption started", + slog.String("input_dir", dumpDir), + slog.String("archive", archivePath), + ) + options.Progress.emit(ProgressEvent{ + Operation: OperationArchive, + Message: "retriever archive encryption started", + InputDir: dumpDir, + ArchivePath: archivePath, + }) + + if err := WriteEncryptedCollectionArchiveWithOptions(tempFile, dumpDir, recipient, options, archivePath, tempPath); err != nil { + _ = tempFile.Close() + return err + } + + if err := tempFile.Close(); err != nil { + return fmt.Errorf("close archive temp file: %w", err) + } + + if err := requirePathDoesNotExist(archivePath); err != nil { + return err + } + + if err := os.Rename(tempPath, archivePath); err != nil { + return fmt.Errorf("rename archive: %w", err) + } + + cleanupTemp = false + + slog.Info("retriever archive encryption completed", + slog.String("archive", archivePath), + ) + options.Progress.emit(ProgressEvent{ + Operation: OperationArchive, + Message: "retriever archive encryption completed", + InputDir: dumpDir, + ArchivePath: archivePath, + }) + return nil +} + +func preflightArchiveOutputPath(archivePath string) error { + archivePath = strings.TrimSpace(archivePath) + if archivePath == "" { + return fmt.Errorf("archive output path is required") + } + + if err := requirePathDoesNotExist(archivePath); err != nil { + return err + } + + if err := os.MkdirAll(filepath.Dir(archivePath), 0o755); err != nil { + return fmt.Errorf("create archive parent directory: %w", err) + } + + tempFile, err := os.CreateTemp(filepath.Dir(archivePath), filepath.Base(archivePath)+".preflight-*.tmp") + if err != nil { + return fmt.Errorf("create archive preflight temp file: %w", err) + } + tempPath := tempFile.Name() + closeErr := tempFile.Close() + removeErr := os.Remove(tempPath) + + if closeErr != nil { + _ = os.Remove(tempPath) + return fmt.Errorf("close archive preflight temp file: %w", closeErr) + } + + if removeErr != nil { + return fmt.Errorf("remove archive preflight temp file: %w", removeErr) + } + + return requirePathDoesNotExist(archivePath) +} + +func PreflightArchiveOutputPath(archivePath string) error { + return preflightArchiveOutputPath(archivePath) +} + +func UnpackEncryptedCollectionArchiveFile(archivePath, outputDir string, force bool, identity hpke.PrivateKey) error { + return UnpackEncryptedCollectionArchiveFileWithOptions(archivePath, outputDir, force, identity, ArchiveOptions{}) +} + +func unpackEncryptedCollectionArchive(archivePath, outputDir string, force bool, identity hpke.PrivateKey) error { + return UnpackEncryptedCollectionArchiveFile(archivePath, outputDir, force, identity) +} + +func UnpackEncryptedCollectionArchiveFileWithOptions(archivePath, outputDir string, force bool, identity hpke.PrivateKey, options ArchiveOptions) error { + file, err := os.Open(archivePath) + if err != nil { + return fmt.Errorf("open archive: %w", err) + } + defer file.Close() + + slog.Info("retriever archive decryption started", + slog.String("archive", archivePath), + slog.String("output_dir", outputDir), + ) + options.Progress.emit(ProgressEvent{ + Operation: OperationUnpack, + Message: "retriever archive decryption started", + ArchivePath: archivePath, + OutputDir: outputDir, + }) + + if err := Unpack(UnpackOptions{ + ArchiveReader: file, + ArchiveIdentity: identity, + OutputDir: outputDir, + Force: force, + Progress: options.Progress, + }); err != nil { + return err + } + + slog.Info("retriever archive decryption completed", + slog.String("archive", archivePath), + slog.String("output_dir", outputDir), + ) + options.Progress.emit(ProgressEvent{ + Operation: OperationUnpack, + Message: "retriever archive decryption completed", + ArchivePath: archivePath, + OutputDir: outputDir, + }) + return nil +} + +func Unpack(options UnpackOptions) error { + if err := options.validate(); err != nil { + return err + } + + outputDir := strings.TrimSpace(options.OutputDir) + stagingDir, err := createUnpackStagingDirectory(outputDir, options.Force) + if err != nil { + return err + } + cleanupStaging := true + defer func() { + if cleanupStaging { + _ = os.RemoveAll(stagingDir) + } + }() + + slog.Info("retriever archive stream unpack started", + slog.String("output_dir", outputDir), + slog.Bool("force", options.Force), + ) + options.Progress.emit(ProgressEvent{ + Operation: OperationUnpack, + Message: "retriever archive stream unpack started", + OutputDir: outputDir, + }) + + if err := UnpackEncryptedCollectionArchiveWithOptions(options.ArchiveReader, stagingDir, options.ArchiveIdentity, ArchiveOptions{Progress: options.Progress}); err != nil { + return err + } + + if err := promoteUnpackStagingDirectory(stagingDir, outputDir, options.Force); err != nil { + return err + } + + cleanupStaging = false + + slog.Info("retriever archive stream unpack completed", + slog.String("output_dir", outputDir), + ) + options.Progress.emit(ProgressEvent{ + Operation: OperationUnpack, + Message: "retriever archive stream unpack completed", + OutputDir: outputDir, + }) + return nil +} + +func UnpackEncryptedCollectionArchive(reader io.Reader, outputDir string, identity hpke.PrivateKey) error { + return UnpackEncryptedCollectionArchiveWithOptions(reader, outputDir, identity, ArchiveOptions{}) +} + +func UnpackEncryptedCollectionArchiveWithOptions(reader io.Reader, outputDir string, identity hpke.PrivateKey, options ArchiveOptions) error { + startedAt := time.Now() + options.Progress.emit(ProgressEvent{ + Operation: OperationUnpack, + Message: "retriever encrypted archive unpack started", + OutputDir: outputDir, + }) + + archiveReader, err := newEncryptedArchiveReader(reader, identity) + if err != nil { + return err + } + + if err := UnpackTarWithOptions(archiveReader, outputDir, false, options); err != nil { + return err + } + + if _, err := io.Copy(io.Discard, archiveReader); err != nil { + return fmt.Errorf("finish encrypted archive stream: %w", err) + } + + if err := validateUnpackedCollection(outputDir); err != nil { + return err + } + + options.Progress.emit(ProgressEvent{ + Operation: OperationUnpack, + Message: "retriever encrypted archive unpack completed", + OutputDir: outputDir, + Elapsed: time.Since(startedAt), + }) + return nil +} + +func unpackEncryptedCollectionArchiveToDirectory(reader io.Reader, outputDir string, identity hpke.PrivateKey) error { + return UnpackEncryptedCollectionArchive(reader, outputDir, identity) +} + +func createUnpackStagingDirectory(outputDir string, force bool) (string, error) { + outputDir = strings.TrimSpace(outputDir) + if outputDir == "" { + return "", fmt.Errorf("output directory is required; pass -out") + } + + if err := preflightUnpackOutputDirectory(outputDir, force); err != nil { + return "", err + } + + parentDir := filepath.Dir(outputDir) + if err := os.MkdirAll(parentDir, 0o755); err != nil { + return "", fmt.Errorf("create unpack parent directory: %w", err) + } + + stagingDir, err := os.MkdirTemp(parentDir, "."+filepath.Base(outputDir)+".unpack-*.tmp") + if err != nil { + return "", fmt.Errorf("create unpack staging directory: %w", err) + } + + return stagingDir, nil +} + +func preflightUnpackOutputDirectory(outputDir string, force bool) error { + info, err := os.Stat(outputDir) + if err == nil { + if !info.IsDir() { + return fmt.Errorf("output path %q exists and is not a directory", outputDir) + } + + entries, err := os.ReadDir(outputDir) + if err != nil { + return fmt.Errorf("read output directory: %w", err) + } + + if len(entries) > 0 && !force { + return fmt.Errorf("output directory %q is not empty; pass -force to replace it", outputDir) + } + + return nil + } + + if !os.IsNotExist(err) { + return fmt.Errorf("inspect output directory: %w", err) + } + + return nil +} + +func validateUnpackedCollection(outputDir string) error { + nextManifest, err := readManifest(outputDir) + if err != nil { + return err + } + if err := verifyManifestFiles(outputDir, nextManifest); err != nil { + return err + } + + expectedPaths, err := archivePathsFromManifest(nextManifest) + if err != nil { + return err + } + expected := make(map[string]struct{}, len(expectedPaths)) + for _, relativePath := range expectedPaths { + expected[relativePath] = struct{}{} + } + + if err := filepath.WalkDir(outputDir, func(filePath string, entry os.DirEntry, err error) error { + if err != nil { + return err + } + + if entry.IsDir() { + return nil + } + + relativePath, err := filepath.Rel(outputDir, filePath) + if err != nil { + return err + } + + archivePath := filepath.ToSlash(relativePath) + if _, ok := expected[archivePath]; !ok { + return fmt.Errorf("unpacked archive contains unexpected file %q", archivePath) + } + + return nil + }); err != nil { + return err + } + + return nil +} + +func promoteUnpackStagingDirectory(stagingDir, outputDir string, force bool) error { + if err := preflightUnpackOutputDirectory(outputDir, force); err != nil { + return err + } + + backupDir := "" + if _, err := os.Stat(outputDir); err == nil { + parentDir := filepath.Dir(outputDir) + reservedBackupDir, err := os.MkdirTemp(parentDir, "."+filepath.Base(outputDir)+".backup-*.tmp") + if err != nil { + return fmt.Errorf("create output directory backup path: %w", err) + } + + if err := os.Remove(reservedBackupDir); err != nil { + return fmt.Errorf("reserve output directory backup path: %w", err) + } + + if err := os.Rename(outputDir, reservedBackupDir); err != nil { + return fmt.Errorf("backup output directory: %w", err) + } + + backupDir = reservedBackupDir + } else if !os.IsNotExist(err) { + return fmt.Errorf("inspect output directory: %w", err) + } + + if err := os.Rename(stagingDir, outputDir); err != nil { + if backupDir != "" { + if restoreErr := os.Rename(backupDir, outputDir); restoreErr != nil { + return fmt.Errorf("promote unpacked archive: %w; restore output directory: %v", err, restoreErr) + } + } + + return fmt.Errorf("promote unpacked archive: %w", err) + } + + if backupDir != "" { + if err := removeUnpackBackupDirectory(backupDir); err != nil { + slog.Warn("retriever unpack output directory backup cleanup failed", + slog.String("backup_dir", backupDir), + slog.Any("error", err), + ) + } + } + + return nil +} + +func NewEncryptedArchiveWriter(writer io.Writer, recipient hpke.PublicKey) (io.WriteCloser, error) { + return newEncryptedArchiveWriter(writer, recipient) +} + +func newEncryptedArchiveWriter(writer io.Writer, recipient hpke.PublicKey) (*encryptedArchiveWriter, error) { + if recipient == nil { + return nil, fmt.Errorf("recipient public key is required") + } + + if recipient.KEM().ID() != defaultArchiveKEM().ID() { + return nil, fmt.Errorf("recipient public key uses unsupported KEM") + } + + encapsulatedKey, sender, err := hpke.NewSender(recipient, defaultArchiveKDF(), defaultArchiveAEAD(), []byte(encryptedArchiveHPKEInfo)) + if err != nil { + return nil, fmt.Errorf("create archive sender: %w", err) + } + + header := encryptedArchiveHeader{ + Format: encryptedArchiveFormat, + Archive: encryptedArchiveFormatHeader{ + Format: "tar", + Compression: "none", + }, + Crypto: encryptedArchiveCryptoHeader{ + Scheme: archiveCryptoScheme, + KEM: archiveKEMName, + KDF: archiveKDFName, + AEAD: archiveAEADName, + EncapsulatedKey: base64.StdEncoding.EncodeToString(encapsulatedKey), + }, + ChunkSize: encryptedArchiveChunkSize, + } + + headerBytes, err := json.Marshal(header) + if err != nil { + return nil, fmt.Errorf("encode archive header: %w", err) + } + + if len(headerBytes) > maxEncryptedArchiveHeaderLen { + return nil, fmt.Errorf("archive header is too large") + } + + if _, err := writer.Write([]byte(encryptedArchiveMagic)); err != nil { + return nil, fmt.Errorf("write archive magic: %w", err) + } + + var headerLen [4]byte + binary.BigEndian.PutUint32(headerLen[:], uint32(len(headerBytes))) + + if _, err := writer.Write(headerLen[:]); err != nil { + return nil, fmt.Errorf("write archive header length: %w", err) + } + + if _, err := writer.Write(headerBytes); err != nil { + return nil, fmt.Errorf("write archive header: %w", err) + } + + return &encryptedArchiveWriter{ + writer: writer, + sender: sender, + headerHash: sha256.Sum256(headerBytes), + }, nil +} + +func (s *encryptedArchiveWriter) Write(p []byte) (int, error) { + if s.closed { + return 0, fmt.Errorf("encrypted archive writer is closed") + } + + written := 0 + for len(p) > 0 { + chunkSize := len(p) + if chunkSize > encryptedArchiveChunkSize { + chunkSize = encryptedArchiveChunkSize + } + + if err := s.writeFrame(encryptedArchiveFrameData, p[:chunkSize]); err != nil { + return written, err + } + + p = p[chunkSize:] + written += chunkSize + } + + return written, nil +} + +func (s *encryptedArchiveWriter) Close() error { + if s.closed { + return fmt.Errorf("encrypted archive writer is already closed") + } + if err := s.writeFrame(encryptedArchiveFrameFinal, nil); err != nil { + return err + } + + s.closed = true + + slog.Info("retriever archive encryption finalized", + slog.Uint64("frame_count", s.frameIndex), + ) + + return nil +} + +func (s *encryptedArchiveWriter) writeFrame(frameType byte, plaintext []byte) error { + ciphertext, err := s.sender.Seal(archiveFrameAAD(s.headerHash, s.frameIndex, frameType), plaintext) + if err != nil { + return fmt.Errorf("encrypt archive frame %d: %w", s.frameIndex, err) + } + + if err := writeEncryptedArchiveFrame(s.writer, frameType, ciphertext); err != nil { + return err + } + + s.frameIndex++ + return nil +} + +func writeEncryptedArchiveFrame(writer io.Writer, frameType byte, ciphertext []byte) error { + if uint64(len(ciphertext)) > math.MaxUint32 { + return fmt.Errorf("encrypted archive frame is too large") + } + + var frameHeader [5]byte + frameHeader[0] = frameType + binary.BigEndian.PutUint32(frameHeader[1:], uint32(len(ciphertext))) + + if _, err := writer.Write(frameHeader[:]); err != nil { + return fmt.Errorf("write archive frame header: %w", err) + } + + if _, err := writer.Write(ciphertext); err != nil { + return fmt.Errorf("write archive frame: %w", err) + } + + return nil +} + +func NewEncryptedArchiveReader(reader io.Reader, identity hpke.PrivateKey) (io.Reader, error) { + return newEncryptedArchiveReader(reader, identity) +} + +func newEncryptedArchiveReader(reader io.Reader, identity hpke.PrivateKey) (*encryptedArchiveReader, error) { + if identity == nil { + return nil, fmt.Errorf("identity private key is required") + } + + if identity.KEM().ID() != defaultArchiveKEM().ID() { + return nil, fmt.Errorf("identity private key uses unsupported KEM") + } + + headerBytes, header, err := readEncryptedArchiveHeader(reader) + if err != nil { + return nil, err + } + encapsulatedKey, err := base64.StdEncoding.DecodeString(header.Crypto.EncapsulatedKey) + if err != nil { + return nil, fmt.Errorf("decode encapsulated key: %w", err) + } + + recipient, err := hpke.NewRecipient(encapsulatedKey, identity, defaultArchiveKDF(), defaultArchiveAEAD(), []byte(encryptedArchiveHPKEInfo)) + if err != nil { + return nil, fmt.Errorf("create archive recipient: %w", err) + } + + return &encryptedArchiveReader{ + reader: reader, + recipient: recipient, + headerHash: sha256.Sum256(headerBytes), + }, nil +} + +func readEncryptedArchiveHeader(reader io.Reader) ([]byte, encryptedArchiveHeader, error) { + var header encryptedArchiveHeader + magic := make([]byte, len(encryptedArchiveMagic)) + if _, err := io.ReadFull(reader, magic); err != nil { + return nil, header, fmt.Errorf("read archive magic: %w", err) + } + + if string(magic) != encryptedArchiveMagic { + return nil, header, fmt.Errorf("unsupported archive magic %q", string(magic)) + } + + var headerLenBytes [4]byte + if _, err := io.ReadFull(reader, headerLenBytes[:]); err != nil { + return nil, header, fmt.Errorf("read archive header length: %w", err) + } + + headerLen := binary.BigEndian.Uint32(headerLenBytes[:]) + if headerLen == 0 || headerLen > maxEncryptedArchiveHeaderLen { + return nil, header, fmt.Errorf("archive header length %d is invalid", headerLen) + } + + headerBytes := make([]byte, headerLen) + if _, err := io.ReadFull(reader, headerBytes); err != nil { + return nil, header, fmt.Errorf("read archive header: %w", err) + } + + if err := json.Unmarshal(headerBytes, &header); err != nil { + return nil, header, fmt.Errorf("decode archive header: %w", err) + } + + if err := validateEncryptedArchiveHeader(header); err != nil { + return nil, header, err + } + + return headerBytes, header, nil +} + +func validateEncryptedArchiveHeader(header encryptedArchiveHeader) error { + if header.Format != encryptedArchiveFormat { + return fmt.Errorf("unsupported encrypted archive format %q", header.Format) + } + + if header.Archive.Format != "tar" { + return fmt.Errorf("unsupported archive payload format %q", header.Archive.Format) + } + + if header.Archive.Compression != "none" { + return fmt.Errorf("unsupported archive payload compression %q", header.Archive.Compression) + } + + if err := validateArchiveCryptoMetadata(ArchiveCryptoMetadata{ + Scheme: header.Crypto.Scheme, + KEM: header.Crypto.KEM, + KDF: header.Crypto.KDF, + AEAD: header.Crypto.AEAD, + }); err != nil { + return err + } + + if strings.TrimSpace(header.Crypto.EncapsulatedKey) == "" { + return fmt.Errorf("archive header is missing encapsulated key") + } + + if header.ChunkSize != encryptedArchiveChunkSize { + return fmt.Errorf("unsupported archive chunk size %d", header.ChunkSize) + } + + return nil +} + +func (s *encryptedArchiveReader) Read(p []byte) (int, error) { + if len(p) == 0 { + return 0, nil + } + + for len(s.plaintext) == 0 { + if s.final { + return 0, io.EOF + } + + if err := s.readNextFrame(); err != nil { + return 0, err + } + } + + n := copy(p, s.plaintext) + s.plaintext = s.plaintext[n:] + + return n, nil +} + +func (s *encryptedArchiveReader) readNextFrame() error { + var frameHeader [5]byte + if _, err := io.ReadFull(s.reader, frameHeader[:]); err != nil { + if errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) { + return fmt.Errorf("encrypted archive missing final frame") + } + + return fmt.Errorf("read archive frame header: %w", err) + } + + frameType := frameHeader[0] + if frameType != encryptedArchiveFrameData && frameType != encryptedArchiveFrameFinal { + return fmt.Errorf("archive frame %d has unsupported type %d", s.frameIndex, frameType) + } + + ciphertextLen := binary.BigEndian.Uint32(frameHeader[1:]) + if ciphertextLen > maxEncryptedArchiveFrameSize { + return fmt.Errorf("archive frame %d is too large", s.frameIndex) + } + + ciphertext := make([]byte, ciphertextLen) + if _, err := io.ReadFull(s.reader, ciphertext); err != nil { + return fmt.Errorf("read archive frame %d: %w", s.frameIndex, err) + } + + plaintext, err := s.recipient.Open(archiveFrameAAD(s.headerHash, s.frameIndex, frameType), ciphertext) + if err != nil { + return fmt.Errorf("decrypt archive frame %d: %w", s.frameIndex, err) + } + + s.frameIndex++ + + if frameType == encryptedArchiveFrameFinal { + if len(plaintext) != 0 { + return fmt.Errorf("encrypted archive final frame contained plaintext") + } + + if err := requireEncryptedArchiveEOF(s.reader); err != nil { + return err + } + + s.final = true + + slog.Info("retriever archive decryption finalized", + slog.Uint64("frame_count", s.frameIndex), + ) + + return nil + } + + s.plaintext = plaintext + return nil +} + +func requireEncryptedArchiveEOF(reader io.Reader) error { + var extra [1]byte + n, err := reader.Read(extra[:]) + if n > 0 { + return fmt.Errorf("encrypted archive has trailing data after final frame") + } + + if err == nil { + return fmt.Errorf("encrypted archive stream did not end after final frame") + } + + if !errors.Is(err, io.EOF) { + return fmt.Errorf("read archive trailer: %w", err) + } + + return nil +} + +func archiveFrameAAD(headerHash [sha256.Size]byte, frameIndex uint64, frameType byte) []byte { + var aad bytes.Buffer + aad.WriteString(encryptedArchiveMagic) + aad.WriteByte(0) + aad.Write(headerHash[:]) + + var indexBytes [8]byte + binary.BigEndian.PutUint64(indexBytes[:], frameIndex) + + aad.Write(indexBytes[:]) + aad.WriteByte(frameType) + + return aad.Bytes() +} diff --git a/retriever/archive_envelope_test.go b/retriever/archive_envelope_test.go new file mode 100644 index 00000000..3183486b --- /dev/null +++ b/retriever/archive_envelope_test.go @@ -0,0 +1,549 @@ +package retriever + +import ( + "archive/tar" + "bytes" + "crypto/hpke" + "encoding/binary" + "errors" + "io" + "os" + "path/filepath" + "strings" + "testing" +) + +func TestEncryptedArchiveRoundTrip(t *testing.T) { + privateKey, publicKey := generateTestArchiveKeyPair(t) + payload := []byte("plain tar bytes") + + ciphertext := encryptArchiveBytes(t, payload, publicKey) + plaintext, err := decryptArchiveBytes(ciphertext, privateKey) + if err != nil { + t.Fatalf("decrypt archive: %v", err) + } + + if !bytes.Equal(plaintext, payload) { + t.Fatalf("plaintext = %q, want %q", plaintext, payload) + } +} + +func TestEncryptedArchiveHeaderDoesNotContainDumpContents(t *testing.T) { + _, publicKey := generateTestArchiveKeyPair(t) + dumpDir := writeArchiveFixture(t) + + var tarPayload bytes.Buffer + if err := writeCollectionTar(&tarPayload, dumpDir); err != nil { + t.Fatalf("write tar: %v", err) + } + + ciphertext := encryptArchiveBytes(t, tarPayload.Bytes(), publicKey) + headerBytes := encryptedArchiveHeaderBytes(t, ciphertext) + + for _, secret := range []string{"secret-graph", "alice", "MemberOf"} { + if bytes.Contains(headerBytes, []byte(secret)) { + t.Fatalf("archive header contains dump content %q: %s", secret, string(headerBytes)) + } + } +} + +func TestEncryptedArchiveWrongKeyAndTamperingFailures(t *testing.T) { + privateKey, publicKey := generateTestArchiveKeyPair(t) + wrongPrivateKey, _ := generateTestArchiveKeyPair(t) + ciphertext := encryptArchiveBytes(t, []byte("payload"), publicKey) + + if _, err := decryptArchiveBytes(ciphertext, wrongPrivateKey); err == nil { + t.Fatalf("expected wrong private key to fail") + } + + tamperedHeader := append([]byte(nil), ciphertext...) + oldChunk := []byte(`"chunk_size":1048576`) + newChunk := []byte(`"chunk_size":1048577`) + headerStart := len(encryptedArchiveMagic) + 4 + headerEnd := encryptedArchiveFrameStart(t, ciphertext) + replaced := bytes.Replace(tamperedHeader[headerStart:headerEnd], oldChunk, newChunk, 1) + if bytes.Equal(replaced, tamperedHeader[headerStart:headerEnd]) { + t.Fatalf("test did not tamper header chunk size") + } + + copy(tamperedHeader[headerStart:headerEnd], replaced) + + if _, err := decryptArchiveBytes(tamperedHeader, privateKey); err == nil { + t.Fatalf("expected tampered header to fail") + } + + tamperedFrame := append([]byte(nil), ciphertext...) + frameStart := encryptedArchiveFrameStart(t, tamperedFrame) + if len(tamperedFrame) <= frameStart+5 { + t.Fatalf("archive frame is too short") + } + + tamperedFrame[frameStart+5] ^= 0x01 + + if _, err := decryptArchiveBytes(tamperedFrame, privateKey); err == nil { + t.Fatalf("expected tampered frame to fail") + } + + oversizedFrame := append([]byte(nil), ciphertext...) + binary.BigEndian.PutUint32(oversizedFrame[frameStart+1:], maxEncryptedArchiveFrameSize+1) + if _, err := decryptArchiveBytes(oversizedFrame, privateKey); err == nil || !strings.Contains(err.Error(), "too large") { + t.Fatalf("expected oversized frame to fail before allocation, got %v", err) + } + + if _, err := decryptArchiveBytes(ciphertext[:len(ciphertext)-1], privateKey); err == nil { + t.Fatalf("expected truncated archive to fail") + } + + withoutFinal := stripFinalArchiveFrame(t, ciphertext) + if _, err := decryptArchiveBytes(withoutFinal, privateKey); err == nil || !strings.Contains(err.Error(), "final frame") { + t.Fatalf("expected missing final frame error, got %v", err) + } +} + +func TestEncryptedCollectionArchiveSyntheticRoundTrip(t *testing.T) { + dir := t.TempDir() + privatePath := filepath.Join(dir, "private.key") + publicPath := filepath.Join(dir, "public.key") + if err := generateArchiveKeyFiles(privatePath, publicPath); err != nil { + t.Fatalf("generate keys: %v", err) + } + + publicKey, err := loadArchivePublicKey(publicPath) + if err != nil { + t.Fatalf("load public key: %v", err) + } + + privateKey, err := loadArchivePrivateKey(privatePath) + if err != nil { + t.Fatalf("load private key: %v", err) + } + + dumpDir := writeArchiveFixture(t) + archivePath := filepath.Join(dir, "dump.tar.pq") + if err := writeEncryptedCollectionArchive(dumpDir, archivePath, publicKey); err != nil { + t.Fatalf("write encrypted collection archive: %v", err) + } + + outputDir := filepath.Join(dir, "unpacked") + if err := unpackEncryptedCollectionArchive(archivePath, outputDir, false, privateKey); err != nil { + t.Fatalf("unpack encrypted collection archive: %v", err) + } + + originalTree := readFileTree(t, dumpDir) + unpackedTree := readFileTree(t, outputDir) + if !equalFileTrees(originalTree, unpackedTree) { + t.Fatalf("unpacked file tree does not match original") + } + + if _, err := readManifest(outputDir); err != nil { + t.Fatalf("read unpacked Manifest: %v", err) + } + + if _, err := os.Stat(archivePath + ".tmp"); !os.IsNotExist(err) { + t.Fatalf("expected archive temp path cleanup, got %v", err) + } +} + +func TestEncryptedCollectionArchiveProgressCallbacks(t *testing.T) { + privateKey, publicKey := generateTestArchiveKeyPair(t) + dumpDir := writeArchiveFixture(t) + var archive bytes.Buffer + var events []ProgressEvent + options := ArchiveOptions{ + Progress: func(event ProgressEvent) { + events = append(events, event) + }, + } + + if err := WriteEncryptedCollectionArchiveWithOptions(&archive, dumpDir, publicKey, options); err != nil { + t.Fatalf("write encrypted collection archive: %v", err) + } + + outputDir := filepath.Join(t.TempDir(), "out") + if err := UnpackEncryptedCollectionArchiveWithOptions(bytes.NewReader(archive.Bytes()), outputDir, privateKey, options); err != nil { + t.Fatalf("unpack encrypted collection archive: %v", err) + } + + if !progressEventsContain(events, OperationArchive, "retriever archive tar streaming completed") { + t.Fatalf("missing archive completion progress event: %+v", events) + } + + if !progressEventsContain(events, OperationUnpack, "retriever archive unpacking completed") { + t.Fatalf("missing unpack completion progress event: %+v", events) + } +} + +func TestUnpackOptionsStreamRoundTripAndForceSafety(t *testing.T) { + privateKey, publicKey := generateTestArchiveKeyPair(t) + dumpDir := writeArchiveFixture(t) + var archive bytes.Buffer + if err := WriteEncryptedCollectionArchive(&archive, dumpDir, publicKey); err != nil { + t.Fatalf("write encrypted collection archive: %v", err) + } + + dir := t.TempDir() + outputDir := filepath.Join(dir, "out") + + if err := Unpack(UnpackOptions{ + ArchiveReader: bytes.NewReader(archive.Bytes()), + ArchiveIdentity: privateKey, + OutputDir: outputDir, + }); err != nil { + t.Fatalf("unpack with options: %v", err) + } + + if originalTree, unpackedTree := readFileTree(t, dumpDir), readFileTree(t, outputDir); !equalFileTrees(originalTree, unpackedTree) { + t.Fatalf("unpacked file tree does not match original") + } + + oldPath := filepath.Join(outputDir, "old") + if err := os.WriteFile(oldPath, []byte("keep me"), 0o600); err != nil { + t.Fatalf("write old output: %v", err) + } + + badArchive := stripFinalArchiveFrame(t, archive.Bytes()) + err := Unpack(UnpackOptions{ + ArchiveReader: bytes.NewReader(badArchive), + ArchiveIdentity: privateKey, + OutputDir: outputDir, + Force: true, + }) + if err == nil || !strings.Contains(err.Error(), "final frame") { + t.Fatalf("expected missing final frame error, got %v", err) + } + + if contents, err := os.ReadFile(oldPath); err != nil { + t.Fatalf("read old output after failed option unpack: %v", err) + } else if string(contents) != "keep me" { + t.Fatalf("old output was replaced after failed option unpack: %q", contents) + } +} + +func TestEncryptedCollectionArchiveFailureDoesNotReplaceOutput(t *testing.T) { + privateKey, publicKey := generateTestArchiveKeyPair(t) + dir := t.TempDir() + dumpDir := writeArchiveFixture(t) + goodArchivePath := filepath.Join(dir, "good.tar.pq") + if err := writeEncryptedCollectionArchive(dumpDir, goodArchivePath, publicKey); err != nil { + t.Fatalf("write encrypted collection archive: %v", err) + } + + goodPayload, err := os.ReadFile(goodArchivePath) + if err != nil { + t.Fatalf("read good archive: %v", err) + } + + badArchivePath := filepath.Join(dir, "bad.tar.pq") + if err := os.WriteFile(badArchivePath, stripFinalArchiveFrame(t, goodPayload), 0o600); err != nil { + t.Fatalf("write bad archive: %v", err) + } + + outputDir := filepath.Join(dir, "out") + if err := os.MkdirAll(outputDir, 0o755); err != nil { + t.Fatalf("create output dir: %v", err) + } + + oldPath := filepath.Join(outputDir, "old") + if err := os.WriteFile(oldPath, []byte("keep me"), 0o600); err != nil { + t.Fatalf("write old output: %v", err) + } + + err = unpackEncryptedCollectionArchive(badArchivePath, outputDir, true, privateKey) + if err == nil || !strings.Contains(err.Error(), "final frame") { + t.Fatalf("expected missing final frame error, got %v", err) + } + + contents, err := os.ReadFile(oldPath) + if err != nil { + t.Fatalf("read old output after failed unpack: %v", err) + } + + if string(contents) != "keep me" { + t.Fatalf("old output was replaced after failed unpack: %q", contents) + } +} + +func progressEventsContain(events []ProgressEvent, operation string, message string) bool { + for _, event := range events { + if event.Operation == operation && event.Message == message { + return true + } + } + + return false +} + +func TestPromoteUnpackStagingDirectoryForceReplacesOutput(t *testing.T) { + dir := t.TempDir() + outputDir := filepath.Join(dir, "out") + stagingDir := filepath.Join(dir, "staging") + if err := os.MkdirAll(outputDir, 0o755); err != nil { + t.Fatalf("create output dir: %v", err) + } + + if err := os.WriteFile(filepath.Join(outputDir, "old.txt"), []byte("old"), 0o600); err != nil { + t.Fatalf("write old output: %v", err) + } + + if err := os.MkdirAll(stagingDir, 0o755); err != nil { + t.Fatalf("create staging dir: %v", err) + } + + if err := os.WriteFile(filepath.Join(stagingDir, "new.txt"), []byte("new"), 0o600); err != nil { + t.Fatalf("write staging output: %v", err) + } + + if err := promoteUnpackStagingDirectory(stagingDir, outputDir, true); err != nil { + t.Fatalf("promote staging dir: %v", err) + } + + if contents, err := os.ReadFile(filepath.Join(outputDir, "new.txt")); err != nil { + t.Fatalf("read promoted output: %v", err) + } else if string(contents) != "new" { + t.Fatalf("promoted output = %q", contents) + } + + if _, err := os.Stat(filepath.Join(outputDir, "old.txt")); !os.IsNotExist(err) { + t.Fatalf("expected old output to be replaced, got %v", err) + } + + if _, err := os.Stat(stagingDir); !os.IsNotExist(err) { + t.Fatalf("expected staging dir to be moved, got %v", err) + } +} + +func TestPromoteUnpackStagingDirectoryIgnoresBackupCleanupFailure(t *testing.T) { + dir := t.TempDir() + outputDir := filepath.Join(dir, "out") + stagingDir := filepath.Join(dir, "staging") + + if err := os.MkdirAll(outputDir, 0o755); err != nil { + t.Fatalf("create output dir: %v", err) + } + + if err := os.MkdirAll(stagingDir, 0o755); err != nil { + t.Fatalf("create staging dir: %v", err) + } + + if err := os.WriteFile(filepath.Join(stagingDir, "new.txt"), []byte("new"), 0o600); err != nil { + t.Fatalf("write staging output: %v", err) + } + + cleanupErr := errors.New("cleanup failed") + originalRemoveUnpackBackupDirectory := removeUnpackBackupDirectory + removeUnpackBackupDirectory = func(string) error { + return cleanupErr + } + t.Cleanup(func() { + removeUnpackBackupDirectory = originalRemoveUnpackBackupDirectory + }) + + if err := promoteUnpackStagingDirectory(stagingDir, outputDir, true); err != nil { + t.Fatalf("promote staging dir: %v", err) + } + + if contents, err := os.ReadFile(filepath.Join(outputDir, "new.txt")); err != nil { + t.Fatalf("read promoted output: %v", err) + } else if string(contents) != "new" { + t.Fatalf("promoted output = %q", contents) + } +} + +func TestPromoteUnpackStagingDirectoryRestoresOutputOnPromotionFailure(t *testing.T) { + dir := t.TempDir() + outputDir := filepath.Join(dir, "out") + missingStagingDir := filepath.Join(dir, "missing-staging") + if err := os.MkdirAll(outputDir, 0o755); err != nil { + t.Fatalf("create output dir: %v", err) + } + + oldPath := filepath.Join(outputDir, "old.txt") + if err := os.WriteFile(oldPath, []byte("old"), 0o600); err != nil { + t.Fatalf("write old output: %v", err) + } + + err := promoteUnpackStagingDirectory(missingStagingDir, outputDir, true) + if err == nil || !strings.Contains(err.Error(), "promote unpacked archive") { + t.Fatalf("expected promotion failure, got %v", err) + } + + if contents, err := os.ReadFile(oldPath); err != nil { + t.Fatalf("read restored output: %v", err) + } else if string(contents) != "old" { + t.Fatalf("restored output = %q", contents) + } +} + +func TestEncryptedCollectionArchiveRejectsInvalidCollection(t *testing.T) { + privateKey, publicKey := generateTestArchiveKeyPair(t) + dir := t.TempDir() + archivePayload := encryptArchiveBytes(t, buildTarPayload(t, &tar.Header{ + Typeflag: tar.TypeReg, + Name: "extra.txt", + Mode: 0o600, + Size: int64(len("not a collection")), + }, []byte("not a collection")), publicKey) + archivePath := filepath.Join(dir, "invalid.tar.pq") + if err := os.WriteFile(archivePath, archivePayload, 0o600); err != nil { + t.Fatalf("write invalid archive: %v", err) + } + + outputDir := filepath.Join(dir, "out") + err := unpackEncryptedCollectionArchive(archivePath, outputDir, false, privateKey) + if err == nil || !strings.Contains(err.Error(), "read manifest") { + t.Fatalf("expected manifest validation error, got %v", err) + } + + if _, err := os.Stat(outputDir); !os.IsNotExist(err) { + t.Fatalf("invalid collection was promoted to output dir: %v", err) + } +} + +func TestEncryptedCollectionArchiveOutputCollisionFails(t *testing.T) { + _, publicKey := generateTestArchiveKeyPair(t) + dumpDir := writeArchiveFixture(t) + + if err := writeEncryptedCollectionArchive(dumpDir, filepath.Join(dumpDir, manifestFileName), publicKey); err == nil { + t.Fatalf("expected archive path collision to fail") + } +} + +func generateTestArchiveKeyPair(t *testing.T) (hpke.PrivateKey, hpke.PublicKey) { + t.Helper() + privateKey, err := defaultArchiveKEM().GenerateKey() + if err != nil { + t.Fatalf("generate archive key: %v", err) + } + + return privateKey, privateKey.PublicKey() +} + +func encryptArchiveBytes(t *testing.T, payload []byte, publicKey hpke.PublicKey) []byte { + t.Helper() + var buffer bytes.Buffer + writer, err := newEncryptedArchiveWriter(&buffer, publicKey) + if err != nil { + t.Fatalf("create encrypted archive writer: %v", err) + } + + if _, err := writer.Write(payload); err != nil { + t.Fatalf("write encrypted archive payload: %v", err) + } + + if err := writer.Close(); err != nil { + t.Fatalf("close encrypted archive writer: %v", err) + } + + return buffer.Bytes() +} + +func decryptArchiveBytes(payload []byte, privateKey hpke.PrivateKey) ([]byte, error) { + reader, err := newEncryptedArchiveReader(bytes.NewReader(payload), privateKey) + if err != nil { + return nil, err + } + + return io.ReadAll(reader) +} + +func encryptedArchiveHeaderBytes(t *testing.T, payload []byte) []byte { + t.Helper() + frameStart := encryptedArchiveFrameStart(t, payload) + headerStart := len(encryptedArchiveMagic) + 4 + return append([]byte(nil), payload[headerStart:frameStart]...) +} + +func encryptedArchiveFrameStart(t *testing.T, payload []byte) int { + t.Helper() + if len(payload) < len(encryptedArchiveMagic)+4 { + t.Fatalf("archive too short") + } + + if string(payload[:len(encryptedArchiveMagic)]) != encryptedArchiveMagic { + t.Fatalf("bad magic") + } + + headerLen := binary.BigEndian.Uint32(payload[len(encryptedArchiveMagic):]) + frameStart := len(encryptedArchiveMagic) + 4 + int(headerLen) + + if frameStart > len(payload) { + t.Fatalf("header length exceeds archive size") + } + + return frameStart +} + +func stripFinalArchiveFrame(t *testing.T, payload []byte) []byte { + t.Helper() + offset := encryptedArchiveFrameStart(t, payload) + for offset < len(payload) { + frameStart := offset + if offset+5 > len(payload) { + t.Fatalf("truncated frame header") + } + + frameType := payload[offset] + frameLen := int(binary.BigEndian.Uint32(payload[offset+1:])) + offset += 5 + + if offset+frameLen > len(payload) { + t.Fatalf("truncated frame payload") + } + + offset += frameLen + + if frameType == encryptedArchiveFrameFinal { + return append([]byte(nil), payload[:frameStart]...) + } + } + + t.Fatalf("missing final frame in test archive") + + return nil +} + +func readFileTree(t *testing.T, root string) map[string][]byte { + t.Helper() + tree := map[string][]byte{} + if err := filepath.WalkDir(root, func(path string, entry os.DirEntry, err error) error { + if err != nil { + return err + } + + if entry.IsDir() { + return nil + } + + relativePath, err := filepath.Rel(root, path) + if err != nil { + return err + } + + contents, err := os.ReadFile(path) + if err != nil { + return err + } + + tree[filepath.ToSlash(relativePath)] = contents + + return nil + }); err != nil { + t.Fatalf("walk file tree: %v", err) + } + + return tree +} + +func equalFileTrees(left, right map[string][]byte) bool { + if len(left) != len(right) { + return false + } + for path, leftContents := range left { + rightContents, ok := right[path] + if !ok || !bytes.Equal(leftContents, rightContents) { + return false + } + } + + return true +} diff --git a/retriever/archive_keys.go b/retriever/archive_keys.go new file mode 100644 index 00000000..ca72daf6 --- /dev/null +++ b/retriever/archive_keys.go @@ -0,0 +1,393 @@ +package retriever + +import ( + "crypto/hpke" + "encoding/base64" + "encoding/json" + "fmt" + "io" + "os" + "path/filepath" + "strings" +) + +const ( + archiveKeyEnvelopeFormat = "retriever-hpke-key-v1" + archiveKeyTypePrivate = "private" + archiveKeyTypePublic = "public" + + archiveCryptoScheme = "hpke" + archiveKEMName = "ML-KEM-1024" + archiveKDFName = "HKDF-SHA512" + archiveAEADName = "AES-256-GCM" +) + +type ArchiveCryptoMetadata struct { + Scheme string `json:"scheme"` + KEM string `json:"kem"` + KDF string `json:"kdf"` + AEAD string `json:"aead"` +} + +type ArchiveKeyEnvelope struct { + Format string `json:"format"` + Type string `json:"type"` + Crypto ArchiveCryptoMetadata `json:"crypto"` + Key string `json:"key"` +} + +func DefaultArchiveKEM() hpke.KEM { + return defaultArchiveKEM() +} + +func defaultArchiveKEM() hpke.KEM { + return hpke.MLKEM1024() +} + +func DefaultArchiveKDF() hpke.KDF { + return defaultArchiveKDF() +} + +func defaultArchiveKDF() hpke.KDF { + return hpke.HKDFSHA512() +} + +func DefaultArchiveAEAD() hpke.AEAD { + return defaultArchiveAEAD() +} + +func defaultArchiveAEAD() hpke.AEAD { + return hpke.AES256GCM() +} + +func DefaultArchiveCryptoMetadata() ArchiveCryptoMetadata { + return defaultArchiveCryptoMetadata() +} + +func defaultArchiveCryptoMetadata() ArchiveCryptoMetadata { + return ArchiveCryptoMetadata{ + Scheme: archiveCryptoScheme, + KEM: archiveKEMName, + KDF: archiveKDFName, + AEAD: archiveAEADName, + } +} + +func GenerateArchiveKeyPair() (hpke.PrivateKey, hpke.PublicKey, error) { + privateKey, err := defaultArchiveKEM().GenerateKey() + if err != nil { + return nil, nil, fmt.Errorf("generate archive key: %w", err) + } + + return privateKey, privateKey.PublicKey(), nil +} + +func GenerateArchiveKeyFiles(privatePath, publicPath string) error { + return generateArchiveKeyFiles(privatePath, publicPath) +} + +func Keygen(options KeygenOptions) error { + if err := options.validate(); err != nil { + return err + } + + return generateArchiveKeyFiles(options.PrivatePath, options.PublicPath) +} + +func generateArchiveKeyFiles(privatePath, publicPath string) error { + privatePath = strings.TrimSpace(privatePath) + publicPath = strings.TrimSpace(publicPath) + if privatePath == "" { + return fmt.Errorf("private key path is required; pass -private") + } + + if publicPath == "" { + return fmt.Errorf("public key path is required; pass -public") + } + + if sameCleanPath(privatePath, publicPath) { + return fmt.Errorf("private and public key paths must be different") + } + + if err := requirePathDoesNotExist(privatePath); err != nil { + return err + } + + if err := requirePathDoesNotExist(publicPath); err != nil { + return err + } + + privateKey, publicKey, err := GenerateArchiveKeyPair() + if err != nil { + return err + } + + if err := writeExclusiveArchivePrivateKeyFile(privatePath, privateKey); err != nil { + return err + } + + if err := writeExclusiveArchivePublicKeyFile(publicPath, publicKey); err != nil { + _ = os.Remove(privatePath) + return err + } + + return nil +} + +func NewArchivePrivateKeyEnvelope(privateKey hpke.PrivateKey) (ArchiveKeyEnvelope, error) { + if privateKey == nil { + return ArchiveKeyEnvelope{}, fmt.Errorf("private key is required") + } + + privateKeyBytes, err := privateKey.Bytes() + if err != nil { + return ArchiveKeyEnvelope{}, fmt.Errorf("serialize private key: %w", err) + } + + return ArchiveKeyEnvelope{ + Format: archiveKeyEnvelopeFormat, + Type: archiveKeyTypePrivate, + Crypto: defaultArchiveCryptoMetadata(), + Key: base64.StdEncoding.EncodeToString(privateKeyBytes), + }, nil +} + +func NewArchivePublicKeyEnvelope(publicKey hpke.PublicKey) (ArchiveKeyEnvelope, error) { + if publicKey == nil { + return ArchiveKeyEnvelope{}, fmt.Errorf("public key is required") + } + + return ArchiveKeyEnvelope{ + Format: archiveKeyEnvelopeFormat, + Type: archiveKeyTypePublic, + Crypto: defaultArchiveCryptoMetadata(), + Key: base64.StdEncoding.EncodeToString(publicKey.Bytes()), + }, nil +} + +func WriteArchivePrivateKey(writer io.Writer, privateKey hpke.PrivateKey) error { + envelope, err := NewArchivePrivateKeyEnvelope(privateKey) + if err != nil { + return err + } + + if err := writeArchiveKeyEnvelope(writer, envelope); err != nil { + return fmt.Errorf("write private key: %w", err) + } + + return nil +} + +func WriteArchivePublicKey(writer io.Writer, publicKey hpke.PublicKey) error { + envelope, err := NewArchivePublicKeyEnvelope(publicKey) + if err != nil { + return err + } + + if err := writeArchiveKeyEnvelope(writer, envelope); err != nil { + return fmt.Errorf("write public key: %w", err) + } + + return nil +} + +func LoadArchivePublicKey(path string) (hpke.PublicKey, error) { + return loadArchivePublicKey(path) +} + +func loadArchivePublicKey(path string) (hpke.PublicKey, error) { + path = strings.TrimSpace(path) + if path == "" { + return nil, fmt.Errorf("public key path is required") + } + + file, err := os.Open(path) + if err != nil { + return nil, fmt.Errorf("read public key: %w", err) + } + defer file.Close() + + return ReadArchivePublicKey(file) +} + +func ReadArchivePublicKey(reader io.Reader) (hpke.PublicKey, error) { + keyBytes, err := readArchiveKeyBytes(reader, archiveKeyTypePublic) + if err != nil { + return nil, err + } + + publicKey, err := defaultArchiveKEM().NewPublicKey(keyBytes) + if err != nil { + return nil, fmt.Errorf("parse public key: %w", err) + } + + return publicKey, nil +} + +func LoadArchivePrivateKey(path string) (hpke.PrivateKey, error) { + return loadArchivePrivateKey(path) +} + +func loadArchivePrivateKey(path string) (hpke.PrivateKey, error) { + path = strings.TrimSpace(path) + if path == "" { + return nil, fmt.Errorf("private key path is required") + } + + file, err := os.Open(path) + if err != nil { + return nil, fmt.Errorf("read private key: %w", err) + } + defer file.Close() + + return ReadArchivePrivateKey(file) +} + +func ReadArchivePrivateKey(reader io.Reader) (hpke.PrivateKey, error) { + keyBytes, err := readArchiveKeyBytes(reader, archiveKeyTypePrivate) + if err != nil { + return nil, err + } + + privateKey, err := defaultArchiveKEM().NewPrivateKey(keyBytes) + if err != nil { + return nil, fmt.Errorf("parse private key: %w", err) + } + + return privateKey, nil +} + +func readArchiveKeyBytes(reader io.Reader, expectedType string) ([]byte, error) { + if reader == nil { + return nil, fmt.Errorf("%s key reader is required", expectedType) + } + + var envelope ArchiveKeyEnvelope + if err := json.NewDecoder(reader).Decode(&envelope); err != nil { + return nil, fmt.Errorf("decode %s key envelope: %w", expectedType, err) + } + + if err := validateArchiveKeyEnvelope(envelope, expectedType); err != nil { + return nil, err + } + + keyBytes, err := base64.StdEncoding.DecodeString(envelope.Key) + if err != nil { + return nil, fmt.Errorf("decode %s key material: %w", expectedType, err) + } + + if len(keyBytes) == 0 { + return nil, fmt.Errorf("%s key material is empty", expectedType) + } + + return keyBytes, nil +} + +func writeArchiveKeyEnvelope(writer io.Writer, envelope ArchiveKeyEnvelope) error { + if writer == nil { + return fmt.Errorf("key writer is required") + } + + encoder := json.NewEncoder(writer) + encoder.SetIndent("", " ") + + return encoder.Encode(envelope) +} + +func validateArchiveKeyEnvelope(envelope ArchiveKeyEnvelope, expectedType string) error { + if envelope.Format != archiveKeyEnvelopeFormat { + return fmt.Errorf("unsupported archive key format %q", envelope.Format) + } + + if envelope.Type != expectedType { + return fmt.Errorf("expected %s key envelope, got %q", expectedType, envelope.Type) + } + + if err := validateArchiveCryptoMetadata(envelope.Crypto); err != nil { + return err + } + + if strings.TrimSpace(envelope.Key) == "" { + return fmt.Errorf("%s key envelope is missing key material", expectedType) + } + + return nil +} + +func validateArchiveCryptoMetadata(metadata ArchiveCryptoMetadata) error { + if metadata.Scheme != archiveCryptoScheme { + return fmt.Errorf("unsupported archive crypto scheme %q", metadata.Scheme) + } + + if metadata.KEM != archiveKEMName { + return fmt.Errorf("unsupported archive KEM %q", metadata.KEM) + } + + if metadata.KDF != archiveKDFName { + return fmt.Errorf("unsupported archive KDF %q", metadata.KDF) + } + + if metadata.AEAD != archiveAEADName { + return fmt.Errorf("unsupported archive AEAD %q", metadata.AEAD) + } + + return nil +} + +func writeExclusiveArchivePrivateKeyFile(path string, privateKey hpke.PrivateKey) error { + return writeExclusiveArchiveKeyFile(path, 0o600, func(writer io.Writer) error { + return WriteArchivePrivateKey(writer, privateKey) + }) +} + +func writeExclusiveArchivePublicKeyFile(path string, publicKey hpke.PublicKey) error { + return writeExclusiveArchiveKeyFile(path, 0o644, func(writer io.Writer) error { + return WriteArchivePublicKey(writer, publicKey) + }) +} + +func writeExclusiveArchiveKeyFile(path string, mode os.FileMode, write func(io.Writer) error) error { + if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil { + return fmt.Errorf("create parent directory for %s: %w", path, err) + } + + file, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_EXCL, mode) + if err != nil { + return fmt.Errorf("create %s: %w", path, err) + } + + writeErr := write(file) + closeErr := file.Close() + + if writeErr != nil { + _ = os.Remove(path) + return writeErr + } + + if closeErr != nil { + _ = os.Remove(path) + return fmt.Errorf("close %s: %w", path, closeErr) + } + + return nil +} + +func requirePathDoesNotExist(path string) error { + if _, err := os.Stat(path); err == nil { + return fmt.Errorf("%s already exists", path) + } else if !os.IsNotExist(err) { + return fmt.Errorf("inspect %s: %w", path, err) + } + + return nil +} + +func sameCleanPath(left, right string) bool { + leftAbs, leftErr := filepath.Abs(left) + rightAbs, rightErr := filepath.Abs(right) + if leftErr == nil && rightErr == nil { + return filepath.Clean(leftAbs) == filepath.Clean(rightAbs) + } + + return filepath.Clean(left) == filepath.Clean(right) +} diff --git a/retriever/archive_keys_test.go b/retriever/archive_keys_test.go new file mode 100644 index 00000000..e25c40ec --- /dev/null +++ b/retriever/archive_keys_test.go @@ -0,0 +1,151 @@ +package retriever + +import ( + "bytes" + "os" + "path/filepath" + "strings" + "testing" +) + +func TestArchiveKeygenAndLoad(t *testing.T) { + dir := t.TempDir() + privatePath := filepath.Join(dir, "retriever-private.key") + publicPath := filepath.Join(dir, "retriever-public.key") + + if err := generateArchiveKeyFiles(privatePath, publicPath); err != nil { + t.Fatalf("generate archive keys: %v", err) + } + + privateInfo, err := os.Stat(privatePath) + if err != nil { + t.Fatalf("stat private key: %v", err) + } + + if got := privateInfo.Mode().Perm(); got != 0o600 { + t.Fatalf("private key mode = %v, want 0600", got) + } + + if _, err := os.Stat(publicPath); err != nil { + t.Fatalf("stat public key: %v", err) + } + + publicKey, err := loadArchivePublicKey(publicPath) + if err != nil { + t.Fatalf("load public key: %v", err) + } + + privateKey, err := loadArchivePrivateKey(privatePath) + if err != nil { + t.Fatalf("load private key: %v", err) + } + + payload := []byte("archive payload") + ciphertext := encryptArchiveBytes(t, payload, publicKey) + plaintext, err := decryptArchiveBytes(ciphertext, privateKey) + if err != nil { + t.Fatalf("decrypt with generated keys: %v", err) + } + + if !bytes.Equal(plaintext, payload) { + t.Fatalf("decrypted payload = %q, want %q", plaintext, payload) + } + + if err := generateArchiveKeyFiles(privatePath, publicPath); err == nil { + t.Fatalf("expected keygen to refuse existing key files") + } + + if _, err := loadArchivePrivateKey(publicPath); err == nil { + t.Fatalf("expected public key envelope to be rejected as private key") + } +} + +func TestArchiveKeyStreamRoundTrip(t *testing.T) { + privateKey, publicKey, err := GenerateArchiveKeyPair() + if err != nil { + t.Fatalf("generate archive key pair: %v", err) + } + + var privateBuffer bytes.Buffer + if err := WriteArchivePrivateKey(&privateBuffer, privateKey); err != nil { + t.Fatalf("write private key: %v", err) + } + + readPrivateKey, err := ReadArchivePrivateKey(bytes.NewReader(privateBuffer.Bytes())) + if err != nil { + t.Fatalf("read private key: %v", err) + } + + var publicBuffer bytes.Buffer + + if err := WriteArchivePublicKey(&publicBuffer, publicKey); err != nil { + t.Fatalf("write public key: %v", err) + } + + readPublicKey, err := ReadArchivePublicKey(bytes.NewReader(publicBuffer.Bytes())) + if err != nil { + t.Fatalf("read public key: %v", err) + } + + payload := []byte("stream archive payload") + ciphertext := encryptArchiveBytes(t, payload, readPublicKey) + plaintext, err := decryptArchiveBytes(ciphertext, readPrivateKey) + if err != nil { + t.Fatalf("decrypt with stream keys: %v", err) + } + + if !bytes.Equal(plaintext, payload) { + t.Fatalf("decrypted payload = %q, want %q", plaintext, payload) + } + + if _, err := ReadArchivePrivateKey(bytes.NewReader(publicBuffer.Bytes())); err == nil { + t.Fatalf("expected public key envelope to be rejected as private key") + } + + if err := WriteArchivePrivateKey(&privateBuffer, nil); err == nil { + t.Fatalf("expected nil private key error") + } +} + +func TestKeygenOptions(t *testing.T) { + dir := t.TempDir() + privatePath := filepath.Join(dir, "private.key") + publicPath := filepath.Join(dir, "public.key") + + if err := Keygen(DefaultKeygenOptions(privatePath, publicPath)); err != nil { + t.Fatalf("keygen with options: %v", err) + } + + if _, err := LoadArchivePrivateKey(privatePath); err != nil { + t.Fatalf("load generated private key: %v", err) + } + + if _, err := LoadArchivePublicKey(publicPath); err != nil { + t.Fatalf("load generated public key: %v", err) + } + + if err := Keygen(DefaultKeygenOptions(privatePath, publicPath)); err == nil { + t.Fatalf("expected keygen options to refuse existing files") + } + + if err := Keygen(DefaultKeygenOptions("", publicPath)); err == nil { + t.Fatalf("expected keygen options to validate paths") + } +} + +func TestArchiveKeygenValidation(t *testing.T) { + dir := t.TempDir() + samePath := filepath.Join(dir, "same.key") + + if err := generateArchiveKeyFiles("", filepath.Join(dir, "public.key")); err == nil || !strings.Contains(err.Error(), "private key path") { + t.Fatalf("expected missing private path error, got %v", err) + } + + if err := generateArchiveKeyFiles(filepath.Join(dir, "private.key"), ""); err == nil || !strings.Contains(err.Error(), "public key path") { + t.Fatalf("expected missing public path error, got %v", err) + } + + if err := generateArchiveKeyFiles(samePath, samePath); err == nil || !strings.Contains(err.Error(), "must be different") { + t.Fatalf("expected same path error, got %v", err) + } +} diff --git a/retriever/archive_tar.go b/retriever/archive_tar.go new file mode 100644 index 00000000..64b5cad6 --- /dev/null +++ b/retriever/archive_tar.go @@ -0,0 +1,403 @@ +package retriever + +import ( + "archive/tar" + "fmt" + "io" + "log/slog" + "os" + "path" + "path/filepath" + "sort" + "strings" + "time" +) + +const archiveFileMode int64 = 0o600 + +var archiveZeroTime = time.Unix(0, 0).UTC() + +func WriteCollectionTar(writer io.Writer, dumpDir string, excludePaths ...string) error { + return WriteCollectionTarWithOptions(writer, dumpDir, ArchiveOptions{}, excludePaths...) +} + +func WriteCollectionTarWithOptions(writer io.Writer, dumpDir string, options ArchiveOptions, excludePaths ...string) error { + slog.Info("retriever archive walking started", + slog.String("input_dir", dumpDir), + ) + options.Progress.emit(ProgressEvent{ + Operation: OperationArchive, + Message: "retriever archive walking started", + InputDir: dumpDir, + }) + + nextManifest, err := readManifest(dumpDir) + if err != nil { + return err + } + + archivePaths, err := archivePathsFromManifest(nextManifest) + if err != nil { + return err + } + + if err := validateArchiveExcludesDoNotReplaceCollectionFiles(dumpDir, archivePaths, excludePaths); err != nil { + return err + } + + slog.Info("retriever archive walking completed", + slog.String("input_dir", dumpDir), + slog.Int("file_count", len(archivePaths)), + ) + options.Progress.emit(ProgressEvent{ + Operation: OperationArchive, + Message: "retriever archive walking completed", + InputDir: dumpDir, + FileCount: len(archivePaths), + }) + + tarWriter := tar.NewWriter(writer) + + slog.Info("retriever archive tar streaming started", + slog.String("input_dir", dumpDir), + slog.Int("file_count", len(archivePaths)), + ) + options.Progress.emit(ProgressEvent{ + Operation: OperationArchive, + Message: "retriever archive tar streaming started", + InputDir: dumpDir, + FileCount: len(archivePaths), + Planned: int64(len(archivePaths)), + }) + + startedAt := time.Now() + for index, relativePath := range archivePaths { + if err := writeCollectionTarFile(tarWriter, dumpDir, relativePath); err != nil { + _ = tarWriter.Close() + return err + } + + processed := int64(index + 1) + options.Progress.emit(ProgressEvent{ + Operation: OperationArchive, + Message: "retriever archive tar file streamed", + InputDir: dumpDir, + FileCount: len(archivePaths), + Processed: processed, + Planned: int64(len(archivePaths)), + Elapsed: time.Since(startedAt), + EntitiesPerSecond: perSecond(processed, time.Since(startedAt)), + }) + } + + if err := tarWriter.Close(); err != nil { + return fmt.Errorf("finish tar stream: %w", err) + } + + slog.Info("retriever archive tar streaming completed", + slog.String("input_dir", dumpDir), + slog.Int("file_count", len(archivePaths)), + ) + options.Progress.emit(ProgressEvent{ + Operation: OperationArchive, + Message: "retriever archive tar streaming completed", + InputDir: dumpDir, + FileCount: len(archivePaths), + Processed: int64(len(archivePaths)), + Planned: int64(len(archivePaths)), + Elapsed: time.Since(startedAt), + }) + return nil +} + +func writeCollectionTar(writer io.Writer, dumpDir string, excludePaths ...string) error { + return WriteCollectionTar(writer, dumpDir, excludePaths...) +} + +func ArchivePathsFromManifest(value Manifest) ([]string, error) { + return archivePathsFromManifest(value) +} + +func archivePathsFromManifest(value Manifest) ([]string, error) { + paths := []string{manifestFileName} + for _, graphEntry := range value.Graphs { + for _, fileEntry := range graphEntry.Files { + paths = append(paths, fileEntry.Path) + } + } + + seen := make(map[string]struct{}, len(paths)) + normalized := make([]string, 0, len(paths)) + + for _, candidate := range paths { + relativePath, err := sanitizeArchivePath(candidate) + if err != nil { + return nil, err + } + + if _, ok := seen[relativePath]; ok { + return nil, fmt.Errorf("archive contains duplicate path %q", relativePath) + } + + seen[relativePath] = struct{}{} + normalized = append(normalized, relativePath) + } + + sort.Strings(normalized) + return normalized, nil +} + +func validateArchiveExcludesDoNotReplaceCollectionFiles(dumpDir string, archivePaths []string, excludePaths []string) error { + if len(excludePaths) == 0 { + return nil + } + + collectionPaths := make(map[string]string, len(archivePaths)) + for _, relativePath := range archivePaths { + absolutePath, err := filepath.Abs(filepath.Join(dumpDir, filepath.FromSlash(relativePath))) + if err != nil { + return fmt.Errorf("resolve archive path %q: %w", relativePath, err) + } + + collectionPaths[filepath.Clean(absolutePath)] = relativePath + } + + for _, excludePath := range excludePaths { + if strings.TrimSpace(excludePath) == "" { + continue + } + + absolutePath, err := filepath.Abs(excludePath) + if err != nil { + return fmt.Errorf("resolve archive output path %q: %w", excludePath, err) + } + + if relativePath, ok := collectionPaths[filepath.Clean(absolutePath)]; ok { + return fmt.Errorf("archive output path %q would replace collection file %q", excludePath, relativePath) + } + } + + return nil +} + +func writeCollectionTarFile(tarWriter *tar.Writer, dumpDir, relativePath string) error { + absolutePath := filepath.Join(dumpDir, filepath.FromSlash(relativePath)) + info, err := os.Lstat(absolutePath) + if err != nil { + return fmt.Errorf("inspect archive file %q: %w", relativePath, err) + } + + if info.Mode()&os.ModeSymlink != 0 { + return fmt.Errorf("archive file %q is a symlink; refusing to package it", relativePath) + } + + if !info.Mode().IsRegular() { + return fmt.Errorf("archive file %q is not a regular file", relativePath) + } + + header := &tar.Header{ + Typeflag: tar.TypeReg, + Name: relativePath, + Mode: archiveFileMode, + Uid: 0, + Gid: 0, + Uname: "", + Gname: "", + Size: info.Size(), + ModTime: archiveZeroTime, + Format: tar.FormatUSTAR, + } + if err := tarWriter.WriteHeader(header); err != nil { + return fmt.Errorf("write tar header for %q: %w", relativePath, err) + } + + file, err := os.Open(absolutePath) + if err != nil { + return fmt.Errorf("open archive file %q: %w", relativePath, err) + } + defer file.Close() + + openInfo, err := file.Stat() + if err != nil { + return fmt.Errorf("stat archive file %q: %w", relativePath, err) + } + + if !openInfo.Mode().IsRegular() { + return fmt.Errorf("archive file %q is not a regular file", relativePath) + } + + if !os.SameFile(info, openInfo) { + return fmt.Errorf("archive file %q changed while packaging", relativePath) + } + + copied, err := io.Copy(tarWriter, file) + if err != nil { + return fmt.Errorf("write tar file %q: %w", relativePath, err) + } + + if copied != info.Size() { + return fmt.Errorf("archive file %q changed while packaging: expected %d bytes, copied %d", relativePath, info.Size(), copied) + } + + return nil +} + +func UnpackTar(reader io.Reader, outputDir string, force bool) error { + return UnpackTarWithOptions(reader, outputDir, force, ArchiveOptions{}) +} + +func UnpackTarWithOptions(reader io.Reader, outputDir string, force bool, options ArchiveOptions) error { + slog.Info("retriever archive unpacking started", + slog.String("output_dir", outputDir), + slog.Bool("force", force), + ) + options.Progress.emit(ProgressEvent{ + Operation: OperationUnpack, + Message: "retriever archive unpacking started", + OutputDir: outputDir, + }) + + if err := prepareOutputDirectory(outputDir, force); err != nil { + return err + } + + tarReader := tar.NewReader(reader) + seen := map[string]struct{}{} + var fileCount int + startedAt := time.Now() + + for { + header, err := tarReader.Next() + if err == io.EOF { + break + } + + if err != nil { + return fmt.Errorf("read tar entry: %w", err) + } + + relativePath, err := sanitizeArchivePath(header.Name) + if err != nil { + return err + } + + if _, ok := seen[relativePath]; ok { + return fmt.Errorf("tar archive contains duplicate path %q", relativePath) + } + + seen[relativePath] = struct{}{} + + if header.Typeflag != tar.TypeReg && header.Typeflag != tar.TypeRegA { + return fmt.Errorf("tar archive path %q is not a regular file", relativePath) + } + + if header.Size < 0 { + return fmt.Errorf("tar archive path %q has negative size", relativePath) + } + + if err := unpackTarFile(tarReader, outputDir, relativePath, header.Size); err != nil { + return err + } + + fileCount++ + + options.Progress.emit(ProgressEvent{ + Operation: OperationUnpack, + Message: "retriever archive file unpacked", + OutputDir: outputDir, + FileCount: fileCount, + Processed: int64(fileCount), + Elapsed: time.Since(startedAt), + EntitiesPerSecond: perSecond(int64(fileCount), time.Since(startedAt)), + }) + } + + slog.Info("retriever archive unpacking completed", + slog.String("output_dir", outputDir), + slog.Int("file_count", fileCount), + ) + options.Progress.emit(ProgressEvent{ + Operation: OperationUnpack, + Message: "retriever archive unpacking completed", + OutputDir: outputDir, + FileCount: fileCount, + Processed: int64(fileCount), + Elapsed: time.Since(startedAt), + }) + + return nil +} + +func unpackTar(reader io.Reader, outputDir string, force bool) error { + return UnpackTar(reader, outputDir, force) +} + +func unpackTarFile(reader io.Reader, outputDir, relativePath string, expectedSize int64) error { + absolutePath := filepath.Join(outputDir, filepath.FromSlash(relativePath)) + if err := os.MkdirAll(filepath.Dir(absolutePath), 0o755); err != nil { + return fmt.Errorf("create parent directory for %q: %w", relativePath, err) + } + + file, err := os.OpenFile(absolutePath, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o600) + if err != nil { + return fmt.Errorf("create unpacked file %q: %w", relativePath, err) + } + + copied, copyErr := io.Copy(file, reader) + closeErr := file.Close() + + if copyErr != nil { + _ = os.Remove(absolutePath) + return fmt.Errorf("write unpacked file %q: %w", relativePath, copyErr) + } + + if closeErr != nil { + _ = os.Remove(absolutePath) + return fmt.Errorf("close unpacked file %q: %w", relativePath, closeErr) + } + + if copied != expectedSize { + _ = os.Remove(absolutePath) + return fmt.Errorf("unpacked file %q size mismatch: expected %d, wrote %d", relativePath, expectedSize, copied) + } + + return nil +} + +func sanitizeArchivePath(value string) (string, error) { + trimmed := strings.TrimSpace(value) + if trimmed == "" { + return "", fmt.Errorf("archive path is empty") + } + + if strings.Contains(trimmed, "\\") { + return "", fmt.Errorf("archive path %q must use slash separators", value) + } + + if path.IsAbs(trimmed) || filepath.IsAbs(trimmed) || hasWindowsVolumeName(trimmed) { + return "", fmt.Errorf("archive path %q must be relative", value) + } + + for _, part := range strings.Split(trimmed, "/") { + if part == ".." { + return "", fmt.Errorf("archive path %q contains path traversal", value) + } + } + + cleaned := path.Clean(trimmed) + if cleaned == "." || cleaned == ".." || strings.HasPrefix(cleaned, "../") { + return "", fmt.Errorf("archive path %q is invalid", value) + } + + return cleaned, nil +} + +func hasWindowsVolumeName(value string) bool { + if len(value) < 2 { + return false + } + + first := value[0] + + return value[1] == ':' && ((first >= 'A' && first <= 'Z') || (first >= 'a' && first <= 'z')) +} diff --git a/retriever/archive_tar_test.go b/retriever/archive_tar_test.go new file mode 100644 index 00000000..1c8b1a34 --- /dev/null +++ b/retriever/archive_tar_test.go @@ -0,0 +1,295 @@ +package retriever + +import ( + "archive/tar" + "bytes" + "io" + "os" + "path/filepath" + "sort" + "strings" + "testing" +) + +func TestCollectionTarDeterministicAndSorted(t *testing.T) { + dir := writeArchiveFixture(t) + value, err := readManifest(dir) + if err != nil { + t.Fatalf("read manifest: %v", err) + } + + archivePaths, err := ArchivePathsFromManifest(value) + if err != nil { + t.Fatalf("archive paths from manifest: %v", err) + } + + if len(archivePaths) != 3 { + t.Fatalf("archive path count = %d, want 3: %v", len(archivePaths), archivePaths) + } + + var first bytes.Buffer + + if err := writeCollectionTar(&first, dir); err != nil { + t.Fatalf("write first tar: %v", err) + } + + var second bytes.Buffer + + if err := writeCollectionTar(&second, dir); err != nil { + t.Fatalf("write second tar: %v", err) + } + + if !bytes.Equal(first.Bytes(), second.Bytes()) { + t.Fatalf("tar output is not deterministic") + } + + names := tarEntryNames(t, first.Bytes()) + + if !sort.StringsAreSorted(names) { + t.Fatalf("tar paths are not sorted: %v", names) + } + + if len(names) != 3 { + t.Fatalf("tar entry count = %d, want 3: %v", len(names), names) + } +} + +func TestCollectionTarStableMetadata(t *testing.T) { + dir := writeArchiveFixture(t) + var buffer bytes.Buffer + if err := writeCollectionTar(&buffer, dir); err != nil { + t.Fatalf("write tar: %v", err) + } + + reader := tar.NewReader(bytes.NewReader(buffer.Bytes())) + for { + header, err := reader.Next() + if err == io.EOF { + break + } + if err != nil { + t.Fatalf("read tar header: %v", err) + } + + if header.Uid != 0 || header.Gid != 0 || header.Uname != "" || header.Gname != "" { + t.Fatalf("tar header leaked host ownership: %+v", header) + } + + if header.Mode != archiveFileMode { + t.Fatalf("tar header mode for %s = %o, want %o", header.Name, header.Mode, archiveFileMode) + } + + if !header.ModTime.Equal(archiveZeroTime) { + t.Fatalf("tar header mtime for %s = %s, want %s", header.Name, header.ModTime, archiveZeroTime) + } + + if header.Typeflag != tar.TypeReg { + t.Fatalf("tar header type for %s = %d, want regular file", header.Name, header.Typeflag) + } + } +} + +func TestCollectionTarRejectsUnsafeInputs(t *testing.T) { + for name, archivePath := range map[string]string{ + "absolute": "/tmp/fragment.gz", + "parent": "../fragment.gz", + "nested parent": "graphs/../fragment.gz", + "back separator": `graphs\default\nodes-000001.ogfrag.gz`, + } { + t.Run(name, func(t *testing.T) { + dir := writeManifestWithArchivePath(t, archivePath) + var buffer bytes.Buffer + if err := writeCollectionTar(&buffer, dir); err == nil { + t.Fatalf("expected unsafe path %q to be rejected", archivePath) + } + }) + } +} + +func TestCollectionTarRejectsSymlinks(t *testing.T) { + dir := writeArchiveFixture(t) + value, err := readManifest(dir) + if err != nil { + t.Fatalf("read manifest: %v", err) + } + + relativePath := value.Graphs[0].Files[0].Path + absolutePath := filepath.Join(dir, filepath.FromSlash(relativePath)) + + if err := os.Remove(absolutePath); err != nil { + t.Fatalf("remove fragment: %v", err) + } + + if err := os.Symlink("target", absolutePath); err != nil { + t.Skipf("symlink unavailable: %v", err) + } + + var buffer bytes.Buffer + if err := writeCollectionTar(&buffer, dir); err == nil || !strings.Contains(err.Error(), "symlink") { + t.Fatalf("expected symlink rejection, got %v", err) + } +} + +func TestUnpackTarRejectsUnsafeEntries(t *testing.T) { + for name, build := range map[string]func(*testing.T) []byte{ + "path traversal": func(t *testing.T) []byte { + return buildTarPayload(t, &tar.Header{ + Typeflag: tar.TypeReg, + Name: "../evil", + Mode: 0o600, + Size: int64(len("bad")), + }, []byte("bad")) + }, + "symlink": func(t *testing.T) []byte { + return buildTarPayload(t, &tar.Header{ + Typeflag: tar.TypeSymlink, + Name: "Manifest.json", + Linkname: "target", + }, nil) + }, + } { + t.Run(name, func(t *testing.T) { + err := unpackTar(bytes.NewReader(build(t)), filepath.Join(t.TempDir(), "out"), false) + if err == nil { + t.Fatalf("expected unsafe tar entry error") + } + }) + } +} + +func TestUnpackTarForceSemantics(t *testing.T) { + dir := writeArchiveFixture(t) + var buffer bytes.Buffer + if err := writeCollectionTar(&buffer, dir); err != nil { + t.Fatalf("write tar: %v", err) + } + + outputDir := t.TempDir() + if err := os.WriteFile(filepath.Join(outputDir, "old"), []byte("old"), 0o600); err != nil { + t.Fatalf("write old file: %v", err) + } + + if err := unpackTar(bytes.NewReader(buffer.Bytes()), outputDir, false); err == nil { + t.Fatalf("expected non-empty output directory error") + } + + if err := unpackTar(bytes.NewReader(buffer.Bytes()), outputDir, true); err != nil { + t.Fatalf("unpack with force: %v", err) + } + + if _, err := os.Stat(filepath.Join(outputDir, "old")); !os.IsNotExist(err) { + t.Fatalf("expected force unpack to remove old file, got %v", err) + } + + if _, err := readManifest(outputDir); err != nil { + t.Fatalf("read unpacked Manifest: %v", err) + } +} + +func writeArchiveFixture(t *testing.T) string { + t.Helper() + dir := t.TempDir() + options := DumpOptions{ + OutputDir: dir, + Compression: CompressionGzip, + ZstdLevel: DefaultZstdLevel, + } + nodeEntry, err := writeNodeFragment(dir, "secret-graph", 1, options, []FragmentNode{{ + ID: "1", + Kinds: []string{"User"}, + Properties: map[string]any{"name": "alice"}, + }}, nil) + if err != nil { + t.Fatalf("write node fragment: %v", err) + } + + edgeEntry, err := writeEdgeFragment(dir, "secret-graph", 1, options, []FragmentEdge{{ + StartID: "1", + EndID: "1", + Kind: "MemberOf", + }}, nil) + if err != nil { + t.Fatalf("write edge fragment: %v", err) + } + + value := newValidTestManifest(1) + value.Schema.Graphs = []GraphSchemaMetadata{{ + Name: "secret-graph", + NodeKinds: []string{"User"}, + EdgeKinds: []string{"MemberOf"}, + }} + value.Graphs = []GraphManifest{{ + Name: "secret-graph", + NodeCount: 1, + EdgeCount: 1, + Files: []FileManifest{nodeEntry, edgeEntry}, + }} + if err := writeManifest(dir, value); err != nil { + t.Fatalf("write Manifest: %v", err) + } + + return dir +} + +func writeManifestWithArchivePath(t *testing.T, archivePath string) string { + t.Helper() + dir := t.TempDir() + value := newValidTestManifest(1) + value.Graphs = []GraphManifest{{ + Name: "default", + NodeCount: 0, + EdgeCount: 0, + Files: []FileManifest{{ + Phase: PhaseNodes, + Path: archivePath, + Count: 0, + CompressedBytes: 0, + SHA256: "abc", + }}, + }} + if err := writeManifest(dir, value); err != nil { + t.Fatalf("write Manifest: %v", err) + } + + return dir +} + +func tarEntryNames(t *testing.T, payload []byte) []string { + t.Helper() + reader := tar.NewReader(bytes.NewReader(payload)) + var names []string + for { + header, err := reader.Next() + if err == io.EOF { + break + } + if err != nil { + t.Fatalf("read tar entry: %v", err) + } + + names = append(names, header.Name) + } + + return names +} + +func buildTarPayload(t *testing.T, header *tar.Header, payload []byte) []byte { + t.Helper() + var buffer bytes.Buffer + writer := tar.NewWriter(&buffer) + if err := writer.WriteHeader(header); err != nil { + t.Fatalf("write tar header: %v", err) + } + + if len(payload) > 0 { + if _, err := writer.Write(payload); err != nil { + t.Fatalf("write tar payload: %v", err) + } + } + + if err := writer.Close(); err != nil { + t.Fatalf("close tar writer: %v", err) + } + + return buffer.Bytes() +} diff --git a/retriever/compression.go b/retriever/compression.go new file mode 100644 index 00000000..6276ace0 --- /dev/null +++ b/retriever/compression.go @@ -0,0 +1,237 @@ +package retriever + +import ( + "bytes" + "compress/gzip" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "fmt" + "hash" + "io" + "os" + "path/filepath" + + "github.com/klauspost/compress/zstd" +) + +const DefaultZstdLevel = 11 + +type countingWriter struct { + writer io.Writer + count int64 +} + +func (s *countingWriter) Write(p []byte) (int, error) { + n, err := s.writer.Write(p) + s.count += int64(n) + return n, err +} + +func ValidateCompression(codec CompressionCodec) error { + switch codec { + case CompressionGzip, CompressionZstd: + return nil + default: + return ValidationError{Message: fmt.Sprintf("unsupported compression codec %q", codec)} + } +} + +func validateCompression(codec CompressionCodec) error { + return ValidateCompression(codec) +} + +func compressionExtension(codec CompressionCodec) (string, error) { + switch codec { + case CompressionGzip: + return ".gz", nil + case CompressionZstd: + return ".zst", nil + default: + return "", fmt.Errorf("unsupported compression codec %q", codec) + } +} + +func newCompressionWriter(writer io.Writer, codec CompressionCodec, zstdLevel int) (io.WriteCloser, error) { + switch codec { + case CompressionGzip: + return gzip.NewWriterLevel(writer, gzip.BestCompression) + case CompressionZstd: + return zstd.NewWriter(writer, zstd.WithEncoderLevel(zstd.EncoderLevelFromZstd(zstdLevel))) + default: + return nil, fmt.Errorf("unsupported compression codec %q", codec) + } +} + +func newDecompressionReader(reader io.Reader, codec CompressionCodec) (io.ReadCloser, error) { + switch codec { + case CompressionGzip: + return gzip.NewReader(reader) + case CompressionZstd: + decoder, err := zstd.NewReader(reader) + if err != nil { + return nil, err + } + + return decoder.IOReadCloser(), nil + default: + return nil, fmt.Errorf("unsupported compression codec %q", codec) + } +} + +func encodeCompactJSON(value any) ([]byte, error) { + var buffer bytes.Buffer + encoder := json.NewEncoder(&buffer) + encoder.SetEscapeHTML(false) + + if err := encoder.Encode(value); err != nil { + return nil, err + } + + return bytes.TrimRight(buffer.Bytes(), "\n"), nil +} + +func writeCompressedJSON(path string, codec CompressionCodec, zstdLevel int, value any) (FileManifest, error) { + payload, err := encodeCompactJSON(value) + if err != nil { + return FileManifest{}, fmt.Errorf("encode fragment: %w", err) + } + + if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil { + return FileManifest{}, fmt.Errorf("create fragment directory: %w", err) + } + + tempPath := path + ".tmp" + file, err := os.OpenFile(tempPath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0o600) + if err != nil { + return FileManifest{}, fmt.Errorf("open fragment temp file: %w", err) + } + + hasher := sha256.New() + counter := &countingWriter{ + writer: io.MultiWriter(file, hasher), + } + compressor, err := newCompressionWriter(counter, codec, zstdLevel) + if err != nil { + file.Close() + os.Remove(tempPath) + + return FileManifest{}, err + } + + _, writeErr := compressor.Write(payload) + closeCompressionErr := compressor.Close() + closeFileErr := file.Close() + + if writeErr != nil { + os.Remove(tempPath) + return FileManifest{}, fmt.Errorf("write compressed fragment: %w", writeErr) + } + + if closeCompressionErr != nil { + os.Remove(tempPath) + return FileManifest{}, fmt.Errorf("finish compressed fragment: %w", closeCompressionErr) + } + + if closeFileErr != nil { + os.Remove(tempPath) + return FileManifest{}, fmt.Errorf("close fragment file: %w", closeFileErr) + } + + if err := os.Rename(tempPath, path); err != nil { + os.Remove(tempPath) + return FileManifest{}, fmt.Errorf("rename fragment: %w", err) + } + + return FileManifest{ + CompressedBytes: counter.count, + UncompressedBytes: int64(len(payload)), + SHA256: hex.EncodeToString(hasher.Sum(nil)), + }, nil +} + +func readCompressedJSON(path string, codec CompressionCodec, target any) error { + file, err := os.Open(path) + if err != nil { + return fmt.Errorf("open fragment: %w", err) + } + defer file.Close() + + reader, err := newDecompressionReader(file, codec) + if err != nil { + return fmt.Errorf("open compressed fragment: %w", err) + } + defer reader.Close() + + decoder := json.NewDecoder(reader) + if err := decoder.Decode(target); err != nil { + return fmt.Errorf("decode fragment: %w", err) + } + + return nil +} + +func verifyChecksum(path string, expectedSHA256 string, expectedCompressedBytes int64) error { + file, err := os.Open(path) + if err != nil { + return fmt.Errorf("open checksum target: %w", err) + } + defer file.Close() + + hasher := sha256.New() + copied, err := copyHash(hasher, file) + if err != nil { + return fmt.Errorf("hash checksum target: %w", err) + } + + if expectedCompressedBytes >= 0 && copied != expectedCompressedBytes { + return ByteCountMismatchError{ + Path: path, + ExpectedBytes: expectedCompressedBytes, + ActualBytes: copied, + } + } + + actual := hex.EncodeToString(hasher.Sum(nil)) + if actual != expectedSHA256 { + return ChecksumMismatchError{ + Path: path, + ExpectedSHA256: expectedSHA256, + ActualSHA256: actual, + } + } + + return nil +} + +func copyHash(hasher hash.Hash, reader io.Reader) (int64, error) { + return io.Copy(hasher, reader) +} + +func compressedJSONSize(codec CompressionCodec, zstdLevel int, value any) (int64, int64, error) { + payload, err := encodeCompactJSON(value) + if err != nil { + return 0, 0, err + } + + var buffer bytes.Buffer + compressor, err := newCompressionWriter(&buffer, codec, zstdLevel) + if err != nil { + return 0, 0, err + } + + if _, err := compressor.Write(payload); err != nil { + compressor.Close() + return 0, 0, err + } + + if err := compressor.Close(); err != nil { + return 0, 0, err + } + + return int64(len(payload)), int64(buffer.Len()), nil +} + +func CompressedJSONSize(codec CompressionCodec, zstdLevel int, value any) (int64, int64, error) { + return compressedJSONSize(codec, zstdLevel, value) +} diff --git a/retriever/compression_test.go b/retriever/compression_test.go new file mode 100644 index 00000000..bfb1b340 --- /dev/null +++ b/retriever/compression_test.go @@ -0,0 +1,175 @@ +package retriever + +import ( + "bytes" + "errors" + "os" + "path/filepath" + "strings" + "testing" +) + +func TestCompressedJSONRoundTrip(t *testing.T) { + for _, codec := range []CompressionCodec{CompressionGzip, CompressionZstd} { + t.Run(string(codec), func(t *testing.T) { + path := filepath.Join(t.TempDir(), "fragment") + entry, err := writeCompressedJSON(path, codec, DefaultZstdLevel, NodeFragment{ + Phase: PhaseNodes, + Items: []FragmentNode{{ + ID: "1", + Kinds: []string{"User"}, + Properties: map[string]any{"name": "alice"}, + }}, + }) + if err != nil { + t.Fatalf("write compressed JSON: %v", err) + } + + if entry.SHA256 == "" { + t.Fatalf("expected checksum") + } + + if entry.CompressedBytes <= 0 || entry.UncompressedBytes <= 0 { + t.Fatalf("expected positive byte counts: %+v", entry) + } + + if err := verifyChecksum(path, entry.SHA256, entry.CompressedBytes); err != nil { + t.Fatalf("verify checksum: %v", err) + } + + var decoded NodeFragment + if err := readCompressedJSON(path, codec, &decoded); err != nil { + t.Fatalf("read compressed JSON: %v", err) + } + + if decoded.Phase != PhaseNodes || len(decoded.Items) != 1 || decoded.Items[0].ID != "1" { + t.Fatalf("unexpected decoded fragment: %+v", decoded) + } + }) + } +} + +func TestVerifyChecksumFailure(t *testing.T) { + path := filepath.Join(t.TempDir(), "fragment") + entry, err := writeCompressedJSON(path, CompressionGzip, DefaultZstdLevel, EdgeFragment{ + Phase: PhaseEdges, + Items: []FragmentEdge{{ + StartID: "1", + EndID: "2", + Kind: "AdminTo", + }}, + }) + if err != nil { + t.Fatalf("write compressed JSON: %v", err) + } + + if err := verifyChecksum(path, entry.SHA256, entry.CompressedBytes+1); err == nil { + t.Fatalf("expected byte-count mismatch") + } else { + var mismatch ByteCountMismatchError + + if !errors.As(err, &mismatch) { + t.Fatalf("expected ByteCountMismatchError, got %T: %v", err, err) + } + + if mismatch.Path != path || mismatch.ExpectedBytes != entry.CompressedBytes+1 || mismatch.ActualBytes != entry.CompressedBytes { + t.Fatalf("unexpected byte-count mismatch: %+v", mismatch) + } + } + + if err := verifyChecksum(path, "not-the-real-checksum", entry.CompressedBytes); err == nil { + t.Fatalf("expected checksum mismatch") + } else { + var mismatch ChecksumMismatchError + + if !errors.As(err, &mismatch) { + t.Fatalf("expected ChecksumMismatchError, got %T: %v", err, err) + } + + if mismatch.Path != path || mismatch.ExpectedSHA256 != "not-the-real-checksum" || mismatch.ActualSHA256 == "" { + t.Fatalf("unexpected checksum mismatch: %+v", mismatch) + } + } +} + +func TestCompressionUnsupportedCodec(t *testing.T) { + if err := validateCompression(CompressionCodec("zip")); err == nil { + t.Fatalf("expected unsupported codec") + } + if _, err := compressionExtension(CompressionCodec("zip")); err == nil { + t.Fatalf("expected unsupported extension codec") + } + + if _, err := newCompressionWriter(&bytes.Buffer{}, CompressionCodec("zip"), DefaultZstdLevel); err == nil { + t.Fatalf("expected unsupported writer codec") + } + + if _, err := newDecompressionReader(strings.NewReader("bad"), CompressionCodec("zip")); err == nil { + t.Fatalf("expected unsupported reader codec") + } +} + +func TestReadCompressedJSONRejectsCorruptPayload(t *testing.T) { + path := filepath.Join(t.TempDir(), "bad.gz") + if err := os.WriteFile(path, []byte("not gzip"), 0o600); err != nil { + t.Fatalf("write corrupt payload: %v", err) + } + + var fragment NodeFragment + if err := readCompressedJSON(path, CompressionGzip, &fragment); err == nil { + t.Fatalf("expected corrupt gzip error") + } +} + +func TestEmptyFragmentRoundTrip(t *testing.T) { + path := filepath.Join(t.TempDir(), "empty.zst") + _, err := writeCompressedJSON(path, CompressionZstd, DefaultZstdLevel, EdgeFragment{ + Phase: PhaseEdges, + Items: nil, + }) + if err != nil { + t.Fatalf("write empty fragment: %v", err) + } + + var decoded EdgeFragment + + if err := readCompressedJSON(path, CompressionZstd, &decoded); err != nil { + t.Fatalf("read empty fragment: %v", err) + } + + if decoded.Phase != PhaseEdges || len(decoded.Items) != 0 { + t.Fatalf("unexpected empty fragment decode: %+v", decoded) + } +} + +func TestCompressedJSONSizeAndCompactEncoding(t *testing.T) { + fragment := NodeFragment{ + Phase: PhaseNodes, + Items: []FragmentNode{{ + ID: "1", + Kinds: []string{"User"}, + }}, + } + + payload, err := encodeCompactJSON(fragment) + if err != nil { + t.Fatalf("encode compact JSON: %v", err) + } + + if strings.Contains(string(payload), "\n") || strings.Contains(string(payload), " ") { + t.Fatalf("expected compact JSON, got %q", payload) + } + + uncompressedBytes, compressedBytes, err := compressedJSONSize(CompressionGzip, DefaultZstdLevel, fragment) + if err != nil { + t.Fatalf("compressed JSON size: %v", err) + } + + if uncompressedBytes != int64(len(payload)) { + t.Fatalf("uncompressed bytes = %d, want %d", uncompressedBytes, len(payload)) + } + + if compressedBytes <= 0 { + t.Fatalf("expected compressed bytes > 0") + } +} diff --git a/retriever/defaults.toml b/retriever/defaults.toml new file mode 100644 index 00000000..9145b903 --- /dev/null +++ b/retriever/defaults.toml @@ -0,0 +1,63 @@ +[scrub] +fake_domain = "example.invalid" +timestamp_shift_days = 17 +redaction_marker = "[REDACTED]" + +[scrub.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/retriever/doc.go b/retriever/doc.go new file mode 100644 index 00000000..a88f0b53 --- /dev/null +++ b/retriever/doc.go @@ -0,0 +1,18 @@ +// Package retriever exports and imports Dawgs graph databases using the +// manifest-based retriever collection format. +// +// The primary database operations accept an already-open graph.Database. The +// package also exposes the collection manifest and fragment structs for callers +// that need to inspect or transform collection contents directly. +// +// Default option constructors mirror the CLI defaults. Callers can attach a +// ProgressFunc to receive structured operation, phase, archive, and periodic +// scan progress events. Validation, compatibility, checksum, metrics, and +// count mismatch failures have typed error values for errors.As checks. Dump, +// Load, Verify, Unpack, and Keygen provide option-based entry points. +// +// Archive APIs are available in both stream-oriented and path-oriented forms. +// Stream-oriented encrypted archive functions use the standard library +// crypto/hpke PublicKey and PrivateKey interfaces. Archive key envelopes can +// also be read from and written to any io.Reader or io.Writer. +package retriever diff --git a/retriever/dump.go b/retriever/dump.go new file mode 100644 index 00000000..45a6e870 --- /dev/null +++ b/retriever/dump.go @@ -0,0 +1,766 @@ +package retriever + +import ( + "context" + "fmt" + "log/slog" + "os" + "path" + "path/filepath" + "sort" + "time" + + "github.com/specterops/dawgs/graph" + "github.com/specterops/dawgs/query" +) + +type DumpResult struct { + Manifest Manifest + ManifestPath string + NodeCount int64 + EdgeCount int64 +} + +type graphEntitySnapshot struct { + NodeCount int64 + EdgeCount int64 +} + +func Dump(ctx context.Context, db graph.Database, driverName string, targets []GraphTarget, options DumpOptions) (DumpResult, error) { + if err := options.validate(); err != nil { + return DumpResult{}, err + } + + if len(targets) == 0 { + return DumpResult{}, fmt.Errorf("at least one graph target is required") + } + + startedAt := time.Now() + slog.Info("retriever dump started", + slog.String("driver", driverName), + slog.Int("graph_count", len(targets)), + slog.String("output_dir", options.OutputDir), + slog.Int("batch_size", options.BatchSize), + slog.Int("shard_size", options.ShardSize), + slog.String("compression", string(options.Compression)), + slog.String("scrub", string(options.Scrub)), + ) + options.Progress.emit(ProgressEvent{ + Operation: OperationDump, + Message: "retriever dump started", + Driver: driverName, + GraphCount: len(targets), + OutputDir: options.OutputDir, + BatchSize: options.BatchSize, + ShardSize: options.ShardSize, + Compression: options.Compression, + Scrub: options.Scrub, + }) + + var ( + activeScrubber *scrubber + + scrubInfo = ScrubMetadata{ + Mode: ScrubNone, + NodeActionCounts: map[string]int{}, + EdgeActionCounts: map[string]int{}, + } + ) + if options.Scrub == ScrubFull { + nextScrubber, err := newScrubber(options.ScrubConfig, options.Salt) + if err != nil { + return DumpResult{}, err + } + activeScrubber = nextScrubber + scrubInfo = activeScrubber.metadata() + } + + slog.Info("retriever dump preparing output directory", + slog.String("output_dir", options.OutputDir), + slog.Bool("force", options.Force), + ) + if err := prepareOutputDirectory(options.OutputDir, options.Force); err != nil { + return DumpResult{}, err + } + + slog.Info("retriever dump output directory ready", + slog.String("output_dir", options.OutputDir), + ) + options.Progress.emit(ProgressEvent{ + Operation: OperationDump, + Message: "retriever dump output directory ready", + OutputDir: options.OutputDir, + }) + + nextManifest := newManifest(driverName, options.Compression, options.ZstdLevel, scrubInfo, len(targets)) + nextMetrics := newMetricsManifest(len(targets)) + var totalNodes, totalEdges int64 + + for targetIndex, target := range targets { + graphStartedAt := time.Now() + slog.Info("retriever dump graph started", + slog.String("graph", target.Name), + slog.Int("graph_index", targetIndex+1), + slog.Int("graph_count", len(targets)), + ) + options.Progress.emit(ProgressEvent{ + Operation: OperationDump, + Message: "retriever dump graph started", + Graph: target.Name, + GraphIndex: targetIndex + 1, + GraphCount: len(targets), + OutputDir: options.OutputDir, + }) + + graphEntry, schemaEntry, metricsEntry, err := dumpGraph(ctx, db, target, options, activeScrubber) + if err != nil { + return DumpResult{}, err + } + + nextManifest.Graphs = append(nextManifest.Graphs, graphEntry) + nextManifest.Schema.Graphs = append(nextManifest.Schema.Graphs, schemaEntry) + nextMetrics.Graphs = append(nextMetrics.Graphs, metricsEntry) + + addActionCounts(nextManifest.Scrub.NodeActionCounts, graphEntry.NodeActionCounts) + addActionCounts(nextManifest.Scrub.EdgeActionCounts, graphEntry.EdgeActionCounts) + + totalNodes += graphEntry.NodeCount + totalEdges += graphEntry.EdgeCount + + slog.Info("retriever dump graph completed", + slog.String("graph", target.Name), + slog.Int64("node_count", graphEntry.NodeCount), + slog.Int64("edge_count", graphEntry.EdgeCount), + slog.Int("file_count", len(graphEntry.Files)), + slog.Duration("wall_elapsed", time.Since(graphStartedAt)), + ) + options.Progress.emit(ProgressEvent{ + Operation: OperationDump, + Message: "retriever dump graph completed", + Graph: target.Name, + GraphIndex: targetIndex + 1, + GraphCount: len(targets), + OutputDir: options.OutputDir, + FileCount: len(graphEntry.Files), + NodeCount: graphEntry.NodeCount, + EdgeCount: graphEntry.EdgeCount, + Elapsed: time.Since(graphStartedAt), + }) + } + + nextManifest.Metrics = &nextMetrics + + slog.Info("retriever dump writing manifest", + slog.String("output_dir", options.OutputDir), + slog.Int64("node_count", totalNodes), + slog.Int64("edge_count", totalEdges), + ) + if err := writeManifest(options.OutputDir, nextManifest); err != nil { + return DumpResult{}, err + } + + manifestPath := filepath.Join(options.OutputDir, manifestFileName) + + slog.Info("retriever dump completed", + slog.String("driver", driverName), + slog.Int("graph_count", len(targets)), + slog.String("manifest", manifestPath), + slog.Int64("node_count", totalNodes), + slog.Int64("edge_count", totalEdges), + slog.Duration("wall_elapsed", time.Since(startedAt)), + ) + options.Progress.emit(ProgressEvent{ + Operation: OperationDump, + Message: "retriever dump completed", + Driver: driverName, + GraphCount: len(targets), + OutputDir: options.OutputDir, + NodeCount: totalNodes, + EdgeCount: totalEdges, + Elapsed: time.Since(startedAt), + }) + + return DumpResult{ + Manifest: nextManifest, + ManifestPath: manifestPath, + NodeCount: totalNodes, + EdgeCount: totalEdges, + }, nil +} + +func prepareOutputDirectory(outputDir string, force bool) error { + info, err := os.Stat(outputDir) + if err == nil { + if !info.IsDir() { + return fmt.Errorf("output path %q exists and is not a directory", outputDir) + } + + entries, err := os.ReadDir(outputDir) + if err != nil { + return fmt.Errorf("read output directory: %w", err) + } + + if len(entries) > 0 { + if !force { + return fmt.Errorf("output directory %q is not empty; pass -force to replace it", outputDir) + } + + if err := os.RemoveAll(outputDir); err != nil { + return fmt.Errorf("replace output directory: %w", err) + } + } + } else if !os.IsNotExist(err) { + return fmt.Errorf("inspect output directory: %w", err) + } + + if err := os.MkdirAll(outputDir, 0o755); err != nil { + return fmt.Errorf("create output directory: %w", err) + } + + return nil +} + +func dumpGraph(ctx context.Context, db graph.Database, target GraphTarget, options DumpOptions, activeScrubber *scrubber) (GraphManifest, GraphSchemaMetadata, GraphMetrics, error) { + targetGraph := graph.Graph{ + Name: target.Name, + } + + countStartedAt := time.Now() + slog.Info("retriever dump counting graph entities", + slog.String("graph", target.Name), + ) + entitySnapshot, err := countGraphEntitySnapshot(ctx, db, targetGraph) + if err != nil { + return GraphManifest{}, GraphSchemaMetadata{}, GraphMetrics{}, err + } + + slog.Info("retriever dump graph counts ready", + slog.String("graph", target.Name), + slog.Int64("node_count", entitySnapshot.NodeCount), + slog.Int64("edge_count", entitySnapshot.EdgeCount), + slog.Duration("wall_elapsed", time.Since(countStartedAt)), + ) + options.Progress.emit(ProgressEvent{ + Operation: OperationDump, + Message: "retriever dump graph counts ready", + Graph: target.Name, + NodeCount: entitySnapshot.NodeCount, + EdgeCount: entitySnapshot.EdgeCount, + Elapsed: time.Since(countStartedAt), + }) + + if activeScrubber != nil { + scrubStartedAt := time.Now() + slog.Info("retriever dump scrub pre-pass started", + slog.String("graph", target.Name), + slog.Int64("node_count", entitySnapshot.NodeCount), + slog.Int("batch_size", options.BatchSize), + ) + options.Progress.emit(ProgressEvent{ + Operation: OperationDump, + Message: "retriever dump scrub pre-pass started", + Graph: target.Name, + Phase: PhaseNodes, + Planned: entitySnapshot.NodeCount, + BatchSize: options.BatchSize, + }) + + observedNodes, err := collectScrubRegistry(ctx, db, targetGraph, options.BatchSize, activeScrubber, entitySnapshot, options.Progress, options.ProgressInterval) + if err != nil { + return GraphManifest{}, GraphSchemaMetadata{}, GraphMetrics{}, err + } + + slog.Info("retriever dump scrub pre-pass completed", + slog.String("graph", target.Name), + slog.Int64("processed", observedNodes), + slog.Duration("wall_elapsed", time.Since(scrubStartedAt)), + ) + options.Progress.emit(ProgressEvent{ + Operation: OperationDump, + Message: "retriever dump scrub pre-pass completed", + Graph: target.Name, + Phase: PhaseNodes, + Processed: observedNodes, + Planned: entitySnapshot.NodeCount, + Elapsed: time.Since(scrubStartedAt), + EntitiesPerSecond: perSecond(observedNodes, time.Since(scrubStartedAt)), + }) + } + + graphEntry := GraphManifest{ + Name: target.Name, + NodeCount: entitySnapshot.NodeCount, + EdgeCount: entitySnapshot.EdgeCount, + NodeActionCounts: map[string]int{}, + EdgeActionCounts: map[string]int{}, + } + + nodeKinds := map[string]struct{}{} + edgeKinds := map[string]struct{}{} + metricsBuilder := newMetricsBuilder(target.Name, entitySnapshot.NodeCount) + + nodeStartedAt := time.Now() + slog.Info("retriever dump node phase started", + slog.String("graph", target.Name), + slog.Int64("node_count", entitySnapshot.NodeCount), + slog.Int("batch_size", options.BatchSize), + slog.Int("shard_size", options.ShardSize), + ) + options.Progress.emit(ProgressEvent{ + Operation: OperationDump, + Message: "retriever dump node phase started", + Graph: target.Name, + Phase: PhaseNodes, + Planned: entitySnapshot.NodeCount, + BatchSize: options.BatchSize, + ShardSize: options.ShardSize, + }) + + nodeFiles, err := dumpNodePhase(ctx, db, targetGraph, options, activeScrubber, nodeKinds, graphEntry.NodeActionCounts, entitySnapshot, metricsBuilder) + if err != nil { + return GraphManifest{}, GraphSchemaMetadata{}, GraphMetrics{}, err + } + + graphEntry.Files = append(graphEntry.Files, nodeFiles...) + + slog.Info("retriever dump node phase completed", + slog.String("graph", target.Name), + slog.Int64("processed", fileTotal(nodeFiles)), + slog.Int("file_count", len(nodeFiles)), + slog.Duration("wall_elapsed", time.Since(nodeStartedAt)), + ) + options.Progress.emit(ProgressEvent{ + Operation: OperationDump, + Message: "retriever dump node phase completed", + Graph: target.Name, + Phase: PhaseNodes, + Processed: fileTotal(nodeFiles), + Planned: entitySnapshot.NodeCount, + FileCount: len(nodeFiles), + Elapsed: time.Since(nodeStartedAt), + EntitiesPerSecond: perSecond(fileTotal(nodeFiles), time.Since(nodeStartedAt)), + }) + + edgeStartedAt := time.Now() + slog.Info("retriever dump edge phase started", + slog.String("graph", target.Name), + slog.Int64("edge_count", entitySnapshot.EdgeCount), + slog.Int("batch_size", options.BatchSize), + slog.Int("shard_size", options.ShardSize), + ) + options.Progress.emit(ProgressEvent{ + Operation: OperationDump, + Message: "retriever dump edge phase started", + Graph: target.Name, + Phase: PhaseEdges, + Planned: entitySnapshot.EdgeCount, + BatchSize: options.BatchSize, + ShardSize: options.ShardSize, + }) + + edgeFiles, err := dumpEdgePhase(ctx, db, targetGraph, options, activeScrubber, edgeKinds, graphEntry.EdgeActionCounts, entitySnapshot, metricsBuilder) + if err != nil { + return GraphManifest{}, GraphSchemaMetadata{}, GraphMetrics{}, err + } + + graphEntry.Files = append(graphEntry.Files, edgeFiles...) + + slog.Info("retriever dump edge phase completed", + slog.String("graph", target.Name), + slog.Int64("processed", fileTotal(edgeFiles)), + slog.Int("file_count", len(edgeFiles)), + slog.Duration("wall_elapsed", time.Since(edgeStartedAt)), + ) + options.Progress.emit(ProgressEvent{ + Operation: OperationDump, + Message: "retriever dump edge phase completed", + Graph: target.Name, + Phase: PhaseEdges, + Processed: fileTotal(edgeFiles), + Planned: entitySnapshot.EdgeCount, + FileCount: len(edgeFiles), + Elapsed: time.Since(edgeStartedAt), + EntitiesPerSecond: perSecond(fileTotal(edgeFiles), time.Since(edgeStartedAt)), + }) + + if fileTotal(nodeFiles) != entitySnapshot.NodeCount { + return GraphManifest{}, GraphSchemaMetadata{}, GraphMetrics{}, EntityCountMismatchError{ + Operation: OperationDump, + Graph: target.Name, + Phase: PhaseNodes, + Expected: entitySnapshot.NodeCount, + Actual: fileTotal(nodeFiles), + Message: fmt.Sprintf("dumped %d nodes for graph %q but counted %d at scan start; source graph changed during dump or the ID scan was inconsistent", fileTotal(nodeFiles), target.Name, entitySnapshot.NodeCount), + } + } + + if fileTotal(edgeFiles) != entitySnapshot.EdgeCount { + return GraphManifest{}, GraphSchemaMetadata{}, GraphMetrics{}, EntityCountMismatchError{ + Operation: OperationDump, + Graph: target.Name, + Phase: PhaseEdges, + Expected: entitySnapshot.EdgeCount, + Actual: fileTotal(edgeFiles), + Message: fmt.Sprintf("dumped %d relationships for graph %q but counted %d at scan start; source graph changed during dump or the ID scan was inconsistent", fileTotal(edgeFiles), target.Name, entitySnapshot.EdgeCount), + } + } + + metricsEntry := metricsBuilder.finalize() + + slog.Info("retriever dump metrics fingerprint computed", + slog.String("graph", target.Name), + slog.String("fingerprint", metricsEntry.Fingerprint), + slog.Int64("node_count", metricsEntry.NodeCount), + slog.Int64("edge_count", metricsEntry.EdgeCount), + ) + + schemaEntry := GraphSchemaMetadata{ + Name: target.Name, + NodeKinds: stringsFromKindSet(nodeKinds), + EdgeKinds: stringsFromKindSet(edgeKinds), + } + + return graphEntry, schemaEntry, metricsEntry, nil +} + +func collectScrubRegistry(ctx context.Context, db graph.Database, targetGraph graph.Graph, batchSize int, activeScrubber *scrubber, entitySnapshot graphEntitySnapshot, progress ProgressFunc, progressInterval int64) (int64, error) { + processed, err := scanDatabaseNodesWithProgressInterval(ctx, db, targetGraph, entitySnapshot.NodeCount, batchSize, progressInterval, func(nodes []*graph.Node) error { + for _, node := range nodes { + activeScrubber.observeNode(node.Properties.MapOrEmpty()) + } + return nil + }, func(processed int64, startedAt time.Time, nextProgressAt int64) int64 { + return logRetrieverEntityProgressInterval("retriever dump scrub pre-pass progress", targetGraph.Name, PhaseNodes, processed, entitySnapshot.NodeCount, startedAt, nextProgressAt, progress, progressInterval) + }) + if err != nil { + return processed, fmt.Errorf("scrub pre-pass: %w", err) + } + + return processed, nil +} + +func countGraphEntities(ctx context.Context, db graph.Database, targetGraph graph.Graph) (int64, int64, error) { + entitySnapshot, err := countGraphEntitySnapshot(ctx, db, targetGraph) + if err != nil { + return 0, 0, err + } + + return entitySnapshot.NodeCount, entitySnapshot.EdgeCount, nil +} + +func countGraphEntitySnapshot(ctx context.Context, db graph.Database, targetGraph graph.Graph) (graphEntitySnapshot, error) { + var entitySnapshot graphEntitySnapshot + if err := db.ReadTransaction(ctx, func(tx graph.Transaction) error { + tx = tx.WithGraph(targetGraph) + + var err error + if entitySnapshot.NodeCount, err = tx.Nodes().Count(); err != nil { + return fmt.Errorf("count nodes: %w", err) + } + + if entitySnapshot.EdgeCount, err = tx.Relationships().Count(); err != nil { + return fmt.Errorf("count relationships: %w", err) + } + + return nil + }); err != nil { + return graphEntitySnapshot{}, err + } + + return entitySnapshot, nil +} + +func dumpNodePhase(ctx context.Context, db graph.Database, targetGraph graph.Graph, options DumpOptions, activeScrubber *scrubber, nodeKinds map[string]struct{}, graphActionCounts map[string]int, entitySnapshot graphEntitySnapshot, metricsBuilder *metricsBuilder) ([]FileManifest, error) { + if entitySnapshot.NodeCount == 0 { + return nil, nil + } + + var ( + files []FileManifest + items []FragmentNode + + shardActionCounts = map[string]int{} + shardNumber = 1 + ) + + flush := func() error { + if len(items) == 0 { + return nil + } + fileEntry, err := writeNodeFragment(options.OutputDir, targetGraph.Name, shardNumber, options, items, shardActionCounts) + if err != nil { + return err + } + + files = append(files, fileEntry) + items = nil + shardActionCounts = map[string]int{} + shardNumber++ + + return nil + } + + if _, err := scanDatabaseNodesWithProgressInterval(ctx, db, targetGraph, entitySnapshot.NodeCount, options.BatchSize, options.ProgressInterval, func(nodes []*graph.Node) error { + for _, node := range nodes { + kinds := node.Kinds.Strings() + sort.Strings(kinds) + addKindsToSet(nodeKinds, kinds) + + properties := node.Properties.MapOrEmpty() + if activeScrubber != nil { + var actionCounts map[string]int + properties, actionCounts = activeScrubber.scrubProperties(properties) + addActionCounts(shardActionCounts, actionCounts) + addActionCounts(graphActionCounts, actionCounts) + } + + item := FragmentNode{ + ID: node.ID.String(), + Kinds: kinds, + Properties: properties, + } + + if err := metricsBuilder.observeFragmentNode(item); err != nil { + return err + } + + items = append(items, item) + + if len(items) >= options.ShardSize { + if err := flush(); err != nil { + return err + } + } + } + + return nil + }, func(processed int64, startedAt time.Time, nextProgressAt int64) int64 { + return logRetrieverEntityProgressInterval("retriever dump node phase progress", targetGraph.Name, PhaseNodes, processed, entitySnapshot.NodeCount, startedAt, nextProgressAt, options.Progress, options.ProgressInterval) + }); err != nil { + return nil, err + } + + if err := flush(); err != nil { + return nil, err + } + + return files, nil +} + +func dumpEdgePhase(ctx context.Context, db graph.Database, targetGraph graph.Graph, options DumpOptions, activeScrubber *scrubber, edgeKinds map[string]struct{}, graphActionCounts map[string]int, entitySnapshot graphEntitySnapshot, metricsBuilder *metricsBuilder) ([]FileManifest, error) { + if entitySnapshot.EdgeCount == 0 { + return nil, nil + } + + var ( + files []FileManifest + items []FragmentEdge + + shardActionCounts = map[string]int{} + shardNumber = 1 + ) + + flush := func() error { + if len(items) == 0 { + return nil + } + fileEntry, err := writeEdgeFragment(options.OutputDir, targetGraph.Name, shardNumber, options, items, shardActionCounts) + if err != nil { + return err + } + + files = append(files, fileEntry) + items = nil + shardActionCounts = map[string]int{} + shardNumber++ + + return nil + } + + if _, err := scanDatabaseRelationshipsWithProgressInterval(ctx, db, targetGraph, entitySnapshot.EdgeCount, options.BatchSize, options.ProgressInterval, func(relationships []*graph.Relationship) error { + for _, relationship := range relationships { + kind := "" + if relationship.Kind != nil { + kind = relationship.Kind.String() + edgeKinds[kind] = struct{}{} + } + + properties := relationship.Properties.MapOrEmpty() + if activeScrubber != nil { + var actionCounts map[string]int + properties, actionCounts = activeScrubber.scrubProperties(properties) + addActionCounts(shardActionCounts, actionCounts) + addActionCounts(graphActionCounts, actionCounts) + } + + item := FragmentEdge{ + StartID: relationship.StartID.String(), + EndID: relationship.EndID.String(), + Kind: kind, + Properties: properties, + } + + if err := metricsBuilder.observeFragmentEdge(item); err != nil { + return err + } + + items = append(items, item) + + if len(items) >= options.ShardSize { + if err := flush(); err != nil { + return err + } + } + } + + return nil + }, func(processed int64, startedAt time.Time, nextProgressAt int64) int64 { + return logRetrieverEntityProgressInterval("retriever dump edge phase progress", targetGraph.Name, PhaseEdges, processed, entitySnapshot.EdgeCount, startedAt, nextProgressAt, options.Progress, options.ProgressInterval) + }); err != nil { + return nil, err + } + + if err := flush(); err != nil { + return nil, err + } + + return files, nil +} + +func writeNodeFragment(outputDir, graphName string, shardNumber int, options DumpOptions, items []FragmentNode, actionCounts map[string]int) (FileManifest, error) { + relativePath, err := fragmentPath(graphName, PhaseNodes, shardNumber, options.Compression) + if err != nil { + return FileManifest{}, err + } + + absolutePath := filepath.Join(outputDir, filepath.FromSlash(relativePath)) + fileEntry, err := writeCompressedJSON(absolutePath, options.Compression, options.ZstdLevel, NodeFragment{ + Phase: PhaseNodes, + Items: items, + }) + if err != nil { + return FileManifest{}, err + } + + fileEntry.Phase = PhaseNodes + fileEntry.Path = relativePath + fileEntry.Count = len(items) + fileEntry.ActionCounts = cloneActionCounts(actionCounts) + + return fileEntry, nil +} + +func writeEdgeFragment(outputDir, graphName string, shardNumber int, options DumpOptions, items []FragmentEdge, actionCounts map[string]int) (FileManifest, error) { + relativePath, err := fragmentPath(graphName, PhaseEdges, shardNumber, options.Compression) + if err != nil { + return FileManifest{}, err + } + + absolutePath := filepath.Join(outputDir, filepath.FromSlash(relativePath)) + fileEntry, err := writeCompressedJSON(absolutePath, options.Compression, options.ZstdLevel, EdgeFragment{ + Phase: PhaseEdges, + Items: items, + }) + if err != nil { + return FileManifest{}, err + } + + fileEntry.Phase = PhaseEdges + fileEntry.Path = relativePath + fileEntry.Count = len(items) + fileEntry.ActionCounts = cloneActionCounts(actionCounts) + + return fileEntry, nil +} + +func fragmentPath(graphName string, fragmentPhase Phase, shardNumber int, codec CompressionCodec) (string, error) { + if shardNumber <= 0 { + return "", fmt.Errorf("shard number must be > 0") + } + + extension, err := compressionExtension(codec) + if err != nil { + return "", err + } + + var prefix string + switch fragmentPhase { + case PhaseNodes: + prefix = "nodes" + case PhaseEdges: + prefix = "edges" + default: + return "", fmt.Errorf("unsupported fragment phase %q", fragmentPhase) + } + + return path.Join("graphs", graphDirectoryName(graphName), fmt.Sprintf("%s-%06d.ogfrag%s", prefix, shardNumber, extension)), nil +} + +func fileTotal(files []FileManifest) int64 { + var total int64 + for _, fileEntry := range files { + total += int64(fileEntry.Count) + } + + return total +} + +func readDatabaseNodes(ctx context.Context, db graph.Database, targetGraph graph.Graph, afterID graph.ID, hasAfterID bool, batchSize int) ([]*graph.Node, error) { + var nodes []*graph.Node + if err := db.ReadTransaction(ctx, func(tx graph.Transaction) error { + tx = tx.WithGraph(targetGraph) + nodeQuery := tx.Nodes(). + OrderBy(query.NodeID()). + Limit(batchSize) + + if hasAfterID { + nodeQuery = nodeQuery.Filter(query.GreaterThan(query.NodeID(), afterID)) + } + + return nodeQuery.Fetch(func(cursor graph.Cursor[*graph.Node]) error { + for node := range cursor.Chan() { + nodes = append(nodes, node) + } + + return cursor.Error() + }) + }); err != nil { + if hasAfterID { + return nil, fmt.Errorf("read node batch after ID %d: %w", afterID.Uint64(), err) + } + + return nil, fmt.Errorf("read initial node batch: %w", err) + } + + return nodes, nil +} + +func readDatabaseRelationships(ctx context.Context, db graph.Database, targetGraph graph.Graph, afterID graph.ID, hasAfterID bool, batchSize int) ([]*graph.Relationship, error) { + var relationships []*graph.Relationship + if err := db.ReadTransaction(ctx, func(tx graph.Transaction) error { + tx = tx.WithGraph(targetGraph) + relationshipQuery := tx.Relationships(). + OrderBy(query.RelationshipID()). + Limit(batchSize) + + if hasAfterID { + relationshipQuery = relationshipQuery.Filter(query.GreaterThan(query.RelationshipID(), afterID)) + } + + return relationshipQuery.Fetch(func(cursor graph.Cursor[*graph.Relationship]) error { + for relationship := range cursor.Chan() { + relationships = append(relationships, relationship) + } + + return cursor.Error() + }) + }); err != nil { + if hasAfterID { + return nil, fmt.Errorf("read relationship batch after ID %d: %w", afterID.Uint64(), err) + } + + return nil, fmt.Errorf("read initial relationship batch: %w", err) + } + + return relationships, nil +} diff --git a/retriever/dump_test.go b/retriever/dump_test.go new file mode 100644 index 00000000..c366b5b1 --- /dev/null +++ b/retriever/dump_test.go @@ -0,0 +1,101 @@ +package retriever + +import ( + "path/filepath" + "testing" +) + +func TestFragmentPath(t *testing.T) { + nodePath, err := fragmentPath("graph/name", PhaseNodes, 7, CompressionZstd) + if err != nil { + t.Fatalf("node fragment path: %v", err) + } + if nodePath != "graphs/graph%2Fname/nodes-000007.ogfrag.zst" { + t.Fatalf("unexpected node fragment path %q", nodePath) + } + + edgePath, err := fragmentPath("default", PhaseEdges, 3, CompressionGzip) + if err != nil { + t.Fatalf("edge fragment path: %v", err) + } + if edgePath != "graphs/default/edges-000003.ogfrag.gz" { + t.Fatalf("unexpected edge fragment path %q", edgePath) + } + + if _, err := fragmentPath("default", Phase("bad"), 1, CompressionGzip); err == nil { + t.Fatalf("expected unsupported Phase error") + } + if _, err := fragmentPath("default", PhaseNodes, 0, CompressionGzip); err == nil { + t.Fatalf("expected invalid shard number error") + } +} + +func TestWriteFragmentMetadata(t *testing.T) { + options := DumpOptions{ + OutputDir: t.TempDir(), + Compression: CompressionGzip, + ZstdLevel: DefaultZstdLevel, + } + + fileEntry, err := writeNodeFragment(options.OutputDir, "default", 1, options, []FragmentNode{{ + ID: "1", + Kinds: []string{"User"}, + Properties: map[string]any{"name": "alice"}, + }}, map[string]int{"pseudonymize": 1}) + if err != nil { + t.Fatalf("write node fragment: %v", err) + } + if fileEntry.Phase != PhaseNodes || fileEntry.Path != "graphs/default/nodes-000001.ogfrag.gz" || fileEntry.Count != 1 { + t.Fatalf("unexpected node file Manifest: %+v", fileEntry) + } + if fileEntry.ActionCounts["pseudonymize"] != 1 { + t.Fatalf("missing action count: %+v", fileEntry.ActionCounts) + } + if _, err := readManifest(filepath.Join(options.OutputDir, "graphs")); err == nil { + t.Fatalf("fragment write should not create Manifest") + } + + edgeEntry, err := writeEdgeFragment(options.OutputDir, "default", 2, options, []FragmentEdge{{ + StartID: "1", + EndID: "2", + Kind: "AdminTo", + }}, nil) + if err != nil { + t.Fatalf("write edge fragment: %v", err) + } + if edgeEntry.Phase != PhaseEdges || edgeEntry.Path != "graphs/default/edges-000002.ogfrag.gz" || edgeEntry.Count != 1 { + t.Fatalf("unexpected edge file Manifest: %+v", edgeEntry) + } +} + +func TestKindAndActionHelpers(t *testing.T) { + kinds := map[string]struct{}{} + addKindsToSet(kinds, []string{"User", "", "Computer", "User"}) + if got := stringsFromKindSet(kinds); len(got) != 2 || got[0] != "Computer" || got[1] != "User" { + t.Fatalf("unexpected kinds: %v", got) + } + + target := map[string]int{"preserve": 1} + addActionCounts(target, map[string]int{"preserve": 2, "redact": 3}) + if target["preserve"] != 3 || target["redact"] != 3 { + t.Fatalf("unexpected action counts: %+v", target) + } + + clone := cloneActionCounts(target) + clone["preserve"] = 100 + if target["preserve"] == 100 { + t.Fatalf("expected clone to be independent") + } + + total := fileTotal([]FileManifest{ + { + Count: 2, + }, + { + Count: 3, + }, + }) + if total != 5 { + t.Fatalf("file total = %d", total) + } +} diff --git a/retriever/errors.go b/retriever/errors.go new file mode 100644 index 00000000..c301c0c1 --- /dev/null +++ b/retriever/errors.go @@ -0,0 +1,85 @@ +package retriever + +import ( + "fmt" +) + +type ChecksumMismatchError struct { + Path string + ExpectedSHA256 string + ActualSHA256 string +} + +func (s ChecksumMismatchError) Error() string { + return fmt.Sprintf("sha256 mismatch for %s: manifest has %s, file has %s", s.Path, s.ExpectedSHA256, s.ActualSHA256) +} + +type ByteCountMismatchError struct { + Path string + ExpectedBytes int64 + ActualBytes int64 +} + +func (s ByteCountMismatchError) Error() string { + return fmt.Sprintf("compressed byte mismatch for %s: manifest has %d, file has %d", s.Path, s.ExpectedBytes, s.ActualBytes) +} + +type IncompatibleDriverError struct { + Operation string + Driver string + Reason string +} + +func (s IncompatibleDriverError) Error() string { + if s.Reason != "" { + return s.Reason + } + if s.Operation != "" && s.Driver != "" { + return fmt.Sprintf("driver %q is incompatible with retriever %s", s.Driver, s.Operation) + } + if s.Driver != "" { + return fmt.Sprintf("driver %q is incompatible with retriever", s.Driver) + } + return "driver is incompatible with retriever" +} + +type NonEmptyTargetGraphError struct { + GraphName string + NodeCount int64 + EdgeCount int64 +} + +func (s NonEmptyTargetGraphError) Error() string { + return fmt.Sprintf("target graph %q is not empty; found %d nodes and %d relationships; clear the graph before loading", s.GraphName, s.NodeCount, s.EdgeCount) +} + +type MissingMetricsError struct { + Operation string +} + +func (s MissingMetricsError) Error() string { + return "manifest does not contain graph metrics; create a new dump with metrics support before verifying" +} + +type EntityCountMismatchError struct { + Operation string + Graph string + Phase Phase + Expected int64 + Actual int64 + Message string +} + +func (s EntityCountMismatchError) Error() string { + if s.Message != "" { + return s.Message + } + phase := string(s.Phase) + if phase == "" { + phase = "entities" + } + if s.Operation == "" { + return fmt.Sprintf("%s count mismatch for graph %q: expected %d, actual %d", phase, s.Graph, s.Expected, s.Actual) + } + return fmt.Sprintf("%s %s count mismatch for graph %q: expected %d, actual %d", s.Operation, phase, s.Graph, s.Expected, s.Actual) +} diff --git a/retriever/graph.go b/retriever/graph.go new file mode 100644 index 00000000..493e8232 --- /dev/null +++ b/retriever/graph.go @@ -0,0 +1,12 @@ +package retriever + +import "net/url" + +func graphDirectoryName(name string) string { + escaped := url.PathEscape(name) + if escaped == "" { + return DefaultGraphName + } + + return escaped +} diff --git a/retriever/load.go b/retriever/load.go new file mode 100644 index 00000000..c757da87 --- /dev/null +++ b/retriever/load.go @@ -0,0 +1,655 @@ +package retriever + +import ( + "context" + "fmt" + "log/slog" + "os" + "path/filepath" + "strings" + "time" + + "github.com/specterops/dawgs/drivers/neo4j" + "github.com/specterops/dawgs/graph" +) + +type LoadResult struct { + GraphCount int + NodeCount int64 + EdgeCount int64 +} + +type schemaAssertion struct { + GraphName string + Schema graph.Schema +} + +type resolvedFragmentEdge struct { + StartID graph.ID + EndID graph.ID + Kind graph.Kind + Properties *graph.Properties +} + +func Load(ctx context.Context, db graph.Database, driverName string, options LoadOptions) (LoadResult, error) { + preparedOptions, cleanupInput, err := prepareLoadInput(options) + if err != nil { + return LoadResult{}, err + } + defer cleanupInput() + + options = preparedOptions + + startedAt := time.Now() + slog.Info("retriever load started", + slog.String("driver", driverName), + slog.String("input_dir", options.InputDir), + slog.Int("batch_size", options.BatchSize), + ) + options.Progress.emit(ProgressEvent{ + Operation: OperationLoad, + Message: "retriever load started", + Driver: driverName, + InputDir: options.InputDir, + BatchSize: options.BatchSize, + }) + + readManifestStartedAt := time.Now() + slog.Info("retriever load reading manifest", + slog.String("input_dir", options.InputDir), + ) + nextManifest, err := readLoadManifest(options.InputDir, driverName) + if err != nil { + return LoadResult{}, err + } + + slog.Info("retriever load manifest ready", + slog.String("input_dir", options.InputDir), + slog.Int("graph_count", len(nextManifest.Graphs)), + slog.String("source_driver", nextManifest.Driver), + slog.String("compression", string(nextManifest.Compression)), + slog.Duration("wall_elapsed", time.Since(readManifestStartedAt)), + ) + options.Progress.emit(ProgressEvent{ + Operation: OperationLoad, + Message: "retriever load manifest ready", + InputDir: options.InputDir, + GraphCount: len(nextManifest.Graphs), + FileCount: manifestFileCount(nextManifest), + Compression: nextManifest.Compression, + Elapsed: time.Since(readManifestStartedAt), + }) + + verifyStartedAt := time.Now() + slog.Info("retriever load verifying fragments", + slog.String("input_dir", options.InputDir), + slog.Int("file_count", manifestFileCount(nextManifest)), + ) + if err := verifyLoadFragments(options.InputDir, nextManifest); err != nil { + return LoadResult{}, err + } + + slog.Info("retriever load fragments verified", + slog.String("input_dir", options.InputDir), + slog.Int("file_count", manifestFileCount(nextManifest)), + slog.Duration("wall_elapsed", time.Since(verifyStartedAt)), + ) + options.Progress.emit(ProgressEvent{ + Operation: OperationLoad, + Message: "retriever load fragments verified", + InputDir: options.InputDir, + FileCount: manifestFileCount(nextManifest), + Elapsed: time.Since(verifyStartedAt), + }) + + schemaStartedAt := time.Now() + slog.Info("retriever load asserting schemas", + slog.Int("graph_count", len(nextManifest.Graphs)), + ) + if err := assertManifestSchemas(ctx, db, nextManifest); err != nil { + return LoadResult{}, err + } + + slog.Info("retriever load schemas ready", + slog.Int("graph_count", len(nextManifest.Graphs)), + slog.Duration("wall_elapsed", time.Since(schemaStartedAt)), + ) + options.Progress.emit(ProgressEvent{ + Operation: OperationLoad, + Message: "retriever load schemas ready", + GraphCount: len(nextManifest.Graphs), + Elapsed: time.Since(schemaStartedAt), + }) + + emptyStartedAt := time.Now() + slog.Info("retriever load checking target graphs", + slog.Int("graph_count", len(nextManifest.Graphs)), + ) + if err := requireEmptyLoadTargets(ctx, db, nextManifest.Graphs); err != nil { + return LoadResult{}, err + } + + slog.Info("retriever load target graphs ready", + slog.Int("graph_count", len(nextManifest.Graphs)), + slog.Duration("wall_elapsed", time.Since(emptyStartedAt)), + ) + options.Progress.emit(ProgressEvent{ + Operation: OperationLoad, + Message: "retriever load target graphs ready", + GraphCount: len(nextManifest.Graphs), + Elapsed: time.Since(emptyStartedAt), + }) + + var result LoadResult + result.GraphCount = len(nextManifest.Graphs) + + for graphIndex, graphEntry := range nextManifest.Graphs { + nodeCount, edgeCount, err := loadManifestGraph(ctx, db, options, nextManifest.Compression, graphIndex, len(nextManifest.Graphs), graphEntry) + if err != nil { + return LoadResult{}, err + } + + result.NodeCount += nodeCount + result.EdgeCount += edgeCount + } + + if options.VerifyMetrics { + if err := verifyLoadedMetrics(ctx, db, nextManifest, options.BatchSize, options.Progress, options.ProgressInterval); err != nil { + return LoadResult{}, err + } + } + + slog.Info("retriever load completed", + slog.String("driver", driverName), + slog.Int("graph_count", result.GraphCount), + slog.Int64("node_count", result.NodeCount), + slog.Int64("edge_count", result.EdgeCount), + slog.Duration("wall_elapsed", time.Since(startedAt)), + ) + options.Progress.emit(ProgressEvent{ + Operation: OperationLoad, + Message: "retriever load completed", + Driver: driverName, + InputDir: options.InputDir, + GraphCount: result.GraphCount, + NodeCount: result.NodeCount, + EdgeCount: result.EdgeCount, + Elapsed: time.Since(startedAt), + }) + + return result, nil +} + +func readLoadManifest(inputDir string, driverName string) (Manifest, error) { + nextManifest, err := readManifest(inputDir) + if err != nil { + return Manifest{}, err + } + if driverName == neo4j.DriverName && len(nextManifest.Graphs) > 1 { + return Manifest{}, IncompatibleDriverError{ + Operation: OperationLoad, + Driver: driverName, + Reason: "cannot load a multi-graph collection into neo4j because Dawgs graph names are no-ops for that driver", + } + } + + return nextManifest, nil +} + +func verifyLoadFragments(inputDir string, nextManifest Manifest) error { + return verifyManifestFiles(inputDir, nextManifest) +} + +type graphEntitySnapshotCounter func(context.Context, graph.Database, graph.Graph) (graphEntitySnapshot, error) + +func requireEmptyLoadTargets(ctx context.Context, db graph.Database, graphEntries []GraphManifest) error { + return requireEmptyLoadTargetsWithCounter(ctx, db, graphEntries, countGraphEntitySnapshot) +} + +func requireEmptyLoadTargetsWithCounter(ctx context.Context, db graph.Database, graphEntries []GraphManifest, countSnapshot graphEntitySnapshotCounter) error { + for _, graphEntry := range graphEntries { + entitySnapshot, err := countSnapshot(ctx, db, graph.Graph{ + Name: graphEntry.Name, + }) + if err != nil { + return fmt.Errorf("count existing entities for graph %q: %w", graphEntry.Name, err) + } + + if entitySnapshot.NodeCount > 0 || entitySnapshot.EdgeCount > 0 { + return NonEmptyTargetGraphError{ + GraphName: graphEntry.Name, + NodeCount: entitySnapshot.NodeCount, + EdgeCount: entitySnapshot.EdgeCount, + } + } + } + + return nil +} + +func loadManifestGraph(ctx context.Context, db graph.Database, options LoadOptions, codec CompressionCodec, graphIndex int, graphCount int, graphEntry GraphManifest) (int64, int64, error) { + graphStartedAt := time.Now() + slog.Info("retriever load graph started", + slog.String("graph", graphEntry.Name), + slog.Int("graph_index", graphIndex+1), + slog.Int("graph_count", graphCount), + slog.Int64("node_count", graphEntry.NodeCount), + slog.Int64("edge_count", graphEntry.EdgeCount), + ) + options.Progress.emit(ProgressEvent{ + Operation: OperationLoad, + Message: "retriever load graph started", + Graph: graphEntry.Name, + GraphIndex: graphIndex + 1, + GraphCount: graphCount, + InputDir: options.InputDir, + NodeCount: graphEntry.NodeCount, + EdgeCount: graphEntry.EdgeCount, + }) + + nodeStartedAt := time.Now() + slog.Info("retriever load node phase started", + slog.String("graph", graphEntry.Name), + slog.Int64("node_count", graphEntry.NodeCount), + ) + options.Progress.emit(ProgressEvent{ + Operation: OperationLoad, + Message: "retriever load node phase started", + Graph: graphEntry.Name, + Phase: PhaseNodes, + InputDir: options.InputDir, + Planned: graphEntry.NodeCount, + }) + + nodeMap, nodeCount, err := loadGraphNodes(ctx, db, options.InputDir, codec, graphEntry, options.Progress, options.ProgressInterval) + if err != nil { + return 0, 0, err + } + + slog.Info("retriever load node phase completed", + slog.String("graph", graphEntry.Name), + slog.Int64("processed", nodeCount), + slog.Duration("wall_elapsed", time.Since(nodeStartedAt)), + slog.Float64("entities_per_second", perSecond(nodeCount, time.Since(nodeStartedAt))), + ) + options.Progress.emit(ProgressEvent{ + Operation: OperationLoad, + Message: "retriever load node phase completed", + Graph: graphEntry.Name, + Phase: PhaseNodes, + InputDir: options.InputDir, + Processed: nodeCount, + Planned: graphEntry.NodeCount, + Elapsed: time.Since(nodeStartedAt), + EntitiesPerSecond: perSecond(nodeCount, time.Since(nodeStartedAt)), + }) + + edgeStartedAt := time.Now() + slog.Info("retriever load edge phase started", + slog.String("graph", graphEntry.Name), + slog.Int64("edge_count", graphEntry.EdgeCount), + slog.Int("batch_size", options.BatchSize), + ) + options.Progress.emit(ProgressEvent{ + Operation: OperationLoad, + Message: "retriever load edge phase started", + Graph: graphEntry.Name, + Phase: PhaseEdges, + InputDir: options.InputDir, + Planned: graphEntry.EdgeCount, + BatchSize: options.BatchSize, + }) + + edgeCount, err := loadGraphEdges(ctx, db, options.InputDir, codec, graphEntry, nodeMap, options.BatchSize, options.Progress, options.ProgressInterval) + if err != nil { + return 0, 0, err + } + + slog.Info("retriever load edge phase completed", + slog.String("graph", graphEntry.Name), + slog.Int64("processed", edgeCount), + slog.Duration("wall_elapsed", time.Since(edgeStartedAt)), + slog.Float64("entities_per_second", perSecond(edgeCount, time.Since(edgeStartedAt))), + ) + options.Progress.emit(ProgressEvent{ + Operation: OperationLoad, + Message: "retriever load edge phase completed", + Graph: graphEntry.Name, + Phase: PhaseEdges, + InputDir: options.InputDir, + Processed: edgeCount, + Planned: graphEntry.EdgeCount, + Elapsed: time.Since(edgeStartedAt), + EntitiesPerSecond: perSecond(edgeCount, time.Since(edgeStartedAt)), + }) + + if nodeCount != graphEntry.NodeCount { + return 0, 0, EntityCountMismatchError{ + Operation: OperationLoad, + Graph: graphEntry.Name, + Phase: PhaseNodes, + Expected: graphEntry.NodeCount, + Actual: nodeCount, + Message: fmt.Sprintf("loaded %d nodes for graph %q but manifest expected %d", nodeCount, graphEntry.Name, graphEntry.NodeCount), + } + } + + if edgeCount != graphEntry.EdgeCount { + return 0, 0, EntityCountMismatchError{ + Operation: OperationLoad, + Graph: graphEntry.Name, + Phase: PhaseEdges, + Expected: graphEntry.EdgeCount, + Actual: edgeCount, + Message: fmt.Sprintf("loaded %d relationships for graph %q but manifest expected %d", edgeCount, graphEntry.Name, graphEntry.EdgeCount), + } + } + + slog.Info("retriever load graph completed", + slog.String("graph", graphEntry.Name), + slog.Int64("node_count", nodeCount), + slog.Int64("edge_count", edgeCount), + slog.Duration("wall_elapsed", time.Since(graphStartedAt)), + ) + options.Progress.emit(ProgressEvent{ + Operation: OperationLoad, + Message: "retriever load graph completed", + Graph: graphEntry.Name, + GraphIndex: graphIndex + 1, + GraphCount: graphCount, + InputDir: options.InputDir, + NodeCount: nodeCount, + EdgeCount: edgeCount, + Elapsed: time.Since(graphStartedAt), + }) + + return nodeCount, edgeCount, nil +} + +func verifyLoadedMetrics(ctx context.Context, db graph.Database, nextManifest Manifest, batchSize int, progress ProgressFunc, progressInterval int64) error { + if nextManifest.Metrics == nil { + return MissingMetricsError{Operation: OperationVerify} + } + + verifyStartedAt := time.Now() + slog.Info("retriever load metrics verification started", + slog.Int("graph_count", len(nextManifest.Graphs)), + slog.Int("batch_size", batchSize), + ) + actualMetrics, _, err := collectDatabaseMetrics(ctx, db, nextManifest.Graphs, batchSize, progress, progressInterval) + if err != nil { + return err + } + + differences := compareMetricsManifest(*nextManifest.Metrics, actualMetrics) + if len(differences) > 0 { + slog.Info("retriever load metrics verification failed", + slog.Int("difference_count", len(differences)), + slog.Duration("wall_elapsed", time.Since(verifyStartedAt)), + ) + return MetricsMismatchError{ + Differences: differences, + } + } + + slog.Info("retriever load metrics verification passed", + slog.Int("graph_count", len(nextManifest.Graphs)), + slog.Duration("wall_elapsed", time.Since(verifyStartedAt)), + ) + + return nil +} + +func prepareLoadInput(options LoadOptions) (LoadOptions, func(), error) { + if err := options.validate(); err != nil { + return LoadOptions{}, nil, err + } + + options.InputDir = strings.TrimSpace(options.InputDir) + + if options.ArchiveReader == nil { + return options, func() {}, nil + } + + tempDir, err := os.MkdirTemp("", "retriever-load-archive-*") + if err != nil { + return LoadOptions{}, nil, fmt.Errorf("create load archive temp directory: %w", err) + } + + cleanup := func() { + _ = os.RemoveAll(tempDir) + } + + slog.Info("retriever load archive unpacking started") + options.Progress.emit(ProgressEvent{ + Operation: OperationLoad, + Message: "retriever load archive unpacking started", + OutputDir: tempDir, + }) + + if err := UnpackEncryptedCollectionArchiveWithOptions(options.ArchiveReader, tempDir, options.ArchiveIdentity, ArchiveOptions{Progress: options.Progress}); err != nil { + cleanup() + return LoadOptions{}, nil, err + } + + slog.Info("retriever load archive unpacking completed") + options.Progress.emit(ProgressEvent{ + Operation: OperationLoad, + Message: "retriever load archive unpacking completed", + OutputDir: tempDir, + }) + + options.InputDir = tempDir + options.ArchiveReader = nil + options.ArchiveIdentity = nil + + if err := options.validate(); err != nil { + cleanup() + return LoadOptions{}, nil, err + } + + return options, cleanup, nil +} + +func assertManifestSchemas(ctx context.Context, db graph.Database, value Manifest) error { + assertions, err := schemaAssertionsFromManifest(value) + if err != nil { + return err + } + + for _, assertion := range assertions { + if err := db.AssertSchema(ctx, assertion.Schema); err != nil { + return fmt.Errorf("assert schema for graph %q: %w", assertion.GraphName, err) + } + } + + return nil +} + +func schemaAssertionsFromManifest(value Manifest) ([]schemaAssertion, error) { + schemaByGraph := map[string]GraphSchemaMetadata{} + for _, schemaEntry := range value.Schema.Graphs { + schemaByGraph[schemaEntry.Name] = schemaEntry + } + + assertions := make([]schemaAssertion, 0, len(value.Graphs)) + for _, graphEntry := range value.Graphs { + schemaEntry, ok := schemaByGraph[graphEntry.Name] + if !ok { + return nil, fmt.Errorf("manifest missing schema metadata for graph %q", graphEntry.Name) + } + + graphSchema := graphSchemaFromMetadata(schemaEntry) + + assertions = append(assertions, schemaAssertion{ + GraphName: graphEntry.Name, + Schema: graph.Schema{ + Graphs: []graph.Graph{graphSchema}, + DefaultGraph: graphSchema, + }, + }) + } + + return assertions, nil +} + +func manifestFileCount(value Manifest) int { + var count int + for _, graphEntry := range value.Graphs { + count += len(graphEntry.Files) + } + + return count +} + +func loadGraphNodes(ctx context.Context, db graph.Database, inputDir string, codec CompressionCodec, graphEntry GraphManifest, progress ProgressFunc, progressInterval int64) (map[string]graph.ID, int64, error) { + nodeMap := make(map[string]graph.ID, int(graphEntry.NodeCount)) + var ( + loaded int64 + + startedAt = time.Now() + nextProgressAt = retrieverInitialProgressAtInterval(graphEntry.NodeCount, progressInterval) + ) + + for _, fileEntry := range graphEntry.Files { + if fileEntry.Phase != PhaseNodes { + continue + } + + fragment, err := readNodeFragmentFile(inputDir, codec, fileEntry) + if err != nil { + return nil, loaded, err + } + + if err := db.WriteTransaction(ctx, func(tx graph.Transaction) error { + tx = tx.WithGraph(graph.Graph{ + Name: graphEntry.Name, + }) + + for _, item := range fragment.Items { + if item.ID == "" { + return fmt.Errorf("node fragment %s contains empty source ID", fileEntry.Path) + } + + if _, exists := nodeMap[item.ID]; exists { + return fmt.Errorf("duplicate source node ID %q in graph %q", item.ID, graphEntry.Name) + } + + dbNode, err := tx.CreateNode(graph.AsProperties(item.Properties), graph.StringsToKinds(item.Kinds)...) + if err != nil { + return fmt.Errorf("create node %q: %w", item.ID, err) + } + + nodeMap[item.ID] = dbNode.ID + } + + return nil + }); err != nil { + return nil, loaded, fmt.Errorf("load node fragment %s: %w", fileEntry.Path, err) + } + + loaded += int64(len(fragment.Items)) + nextProgressAt = logRetrieverEntityProgressInterval("retriever load node phase progress", graphEntry.Name, PhaseNodes, loaded, graphEntry.NodeCount, startedAt, nextProgressAt, progress, progressInterval) + } + + return nodeMap, loaded, nil +} + +func loadGraphEdges(ctx context.Context, db graph.Database, inputDir string, codec CompressionCodec, graphEntry GraphManifest, nodeMap map[string]graph.ID, batchSize int, progress ProgressFunc, progressInterval int64) (int64, error) { + var ( + loaded int64 + + startedAt = time.Now() + nextProgressAt = retrieverInitialProgressAtInterval(graphEntry.EdgeCount, progressInterval) + ) + + for _, fileEntry := range graphEntry.Files { + if fileEntry.Phase != PhaseEdges { + continue + } + + fragment, err := readEdgeFragmentFile(inputDir, codec, fileEntry) + if err != nil { + return loaded, err + } + + if err := db.BatchOperation(ctx, func(batch graph.Batch) error { + batch = batch.WithGraph(graph.Graph{ + Name: graphEntry.Name, + }) + + for _, item := range fragment.Items { + resolved, err := resolveFragmentEdge(item, nodeMap) + if err != nil { + return err + } + + if err := batch.CreateRelationshipByIDs(resolved.StartID, resolved.EndID, resolved.Kind, resolved.Properties); err != nil { + return fmt.Errorf("create edge (%s)-[%s]->(%s): %w", item.StartID, item.Kind, item.EndID, err) + } + } + + return nil + }, graph.WithBatchSize(batchSize)); err != nil { + return loaded, fmt.Errorf("load edge fragment %s: %w", fileEntry.Path, err) + } + + loaded += int64(len(fragment.Items)) + nextProgressAt = logRetrieverEntityProgressInterval("retriever load edge phase progress", graphEntry.Name, PhaseEdges, loaded, graphEntry.EdgeCount, startedAt, nextProgressAt, progress, progressInterval) + } + + return loaded, nil +} + +func readNodeFragmentFile(inputDir string, codec CompressionCodec, fileEntry FileManifest) (NodeFragment, error) { + var fragment NodeFragment + if err := readCompressedJSON(filepath.Join(inputDir, filepath.FromSlash(fileEntry.Path)), codec, &fragment); err != nil { + return NodeFragment{}, fmt.Errorf("read node fragment %s: %w", fileEntry.Path, err) + } + + if fragment.Phase != PhaseNodes { + return NodeFragment{}, fmt.Errorf("fragment %s has phase %q, expected nodes", fileEntry.Path, fragment.Phase) + } + + if len(fragment.Items) != fileEntry.Count { + return NodeFragment{}, fmt.Errorf("fragment %s item count %d does not match manifest count %d", fileEntry.Path, len(fragment.Items), fileEntry.Count) + } + + return fragment, nil +} + +func readEdgeFragmentFile(inputDir string, codec CompressionCodec, fileEntry FileManifest) (EdgeFragment, error) { + var fragment EdgeFragment + if err := readCompressedJSON(filepath.Join(inputDir, filepath.FromSlash(fileEntry.Path)), codec, &fragment); err != nil { + return EdgeFragment{}, fmt.Errorf("read edge fragment %s: %w", fileEntry.Path, err) + } + + if fragment.Phase != PhaseEdges { + return EdgeFragment{}, fmt.Errorf("fragment %s has phase %q, expected edges", fileEntry.Path, fragment.Phase) + } + + if len(fragment.Items) != fileEntry.Count { + return EdgeFragment{}, fmt.Errorf("fragment %s item count %d does not match manifest count %d", fileEntry.Path, len(fragment.Items), fileEntry.Count) + } + + return fragment, nil +} + +func resolveFragmentEdge(item FragmentEdge, nodeMap map[string]graph.ID) (resolvedFragmentEdge, error) { + startID, ok := nodeMap[item.StartID] + if !ok { + return resolvedFragmentEdge{}, fmt.Errorf("edge references missing start node %q", item.StartID) + } + + endID, ok := nodeMap[item.EndID] + if !ok { + return resolvedFragmentEdge{}, fmt.Errorf("edge references missing end node %q", item.EndID) + } + + return resolvedFragmentEdge{ + StartID: startID, + EndID: endID, + Kind: graph.StringKind(item.Kind), + Properties: graph.AsProperties(item.Properties), + }, nil +} diff --git a/retriever/load_test.go b/retriever/load_test.go new file mode 100644 index 00000000..b552dfc5 --- /dev/null +++ b/retriever/load_test.go @@ -0,0 +1,296 @@ +package retriever + +import ( + "context" + "errors" + "path/filepath" + "strings" + "testing" + + "github.com/specterops/dawgs/drivers/neo4j" + "github.com/specterops/dawgs/drivers/pg" + "github.com/specterops/dawgs/graph" +) + +func TestRequireEmptyLoadTargets(t *testing.T) { + ctx := context.Background() + graphEntries := []GraphManifest{ + {Name: "empty"}, + {Name: "busy"}, + } + + if err := requireEmptyLoadTargetsWithCounter(ctx, nil, graphEntries[:1], func(context.Context, graph.Database, graph.Graph) (graphEntitySnapshot, error) { + return graphEntitySnapshot{}, nil + }); err != nil { + t.Fatalf("require empty load targets: %v", err) + } + + err := requireEmptyLoadTargetsWithCounter(ctx, nil, graphEntries, func(_ context.Context, _ graph.Database, target graph.Graph) (graphEntitySnapshot, error) { + if target.Name == "busy" { + return graphEntitySnapshot{ + NodeCount: 1, + EdgeCount: 2, + }, nil + } + + return graphEntitySnapshot{}, nil + }) + + if err == nil || !strings.Contains(err.Error(), `target graph "busy" is not empty`) { + t.Fatalf("expected non-empty graph error, got %v", err) + } + + var nonEmptyErr NonEmptyTargetGraphError + if !errors.As(err, &nonEmptyErr) { + t.Fatalf("expected NonEmptyTargetGraphError, got %T: %v", err, err) + } + + if nonEmptyErr.GraphName != "busy" || nonEmptyErr.NodeCount != 1 || nonEmptyErr.EdgeCount != 2 { + t.Fatalf("unexpected non-empty graph error: %+v", nonEmptyErr) + } + + expectedErr := errors.New("count failed") + err = requireEmptyLoadTargetsWithCounter(ctx, nil, graphEntries[:1], func(context.Context, graph.Database, graph.Graph) (graphEntitySnapshot, error) { + return graphEntitySnapshot{}, expectedErr + }) + if !errors.Is(err, expectedErr) { + t.Fatalf("expected counter error, got %v", err) + } +} + +func TestSchemaAssertionsFromManifest(t *testing.T) { + value := newValidTestManifest(2) + value.Graphs = []GraphManifest{ + { + Name: "a", + }, + { + Name: "b", + }, + } + value.Schema.Graphs = []GraphSchemaMetadata{ + { + Name: "a", + NodeKinds: []string{"User"}, + EdgeKinds: []string{"AdminTo"}, + }, + { + Name: "b", + NodeKinds: []string{"Computer"}, + EdgeKinds: []string{"MemberOf"}, + }, + } + + assertions, err := schemaAssertionsFromManifest(value) + if err != nil { + t.Fatalf("schema assertions: %v", err) + } + + if len(assertions) != 2 { + t.Fatalf("assertion count = %d", len(assertions)) + } + + if assertions[0].Schema.DefaultGraph.Name != "a" || assertions[0].Schema.Graphs[0].Nodes[0].String() != "User" { + t.Fatalf("unexpected first schema assertion: %+v", assertions[0]) + } + + value.Schema.Graphs = value.Schema.Graphs[:1] + + if _, err := schemaAssertionsFromManifest(value); err == nil { + t.Fatalf("expected missing schema error") + } +} + +func TestResolveFragmentEdge(t *testing.T) { + nodeMap := map[string]graph.ID{ + "1": 101, + "2": 202, + } + item := FragmentEdge{ + StartID: "1", + EndID: "2", + Kind: "AdminTo", + Properties: map[string]any{"source": "test"}, + } + + resolved, err := resolveFragmentEdge(item, nodeMap) + if err != nil { + t.Fatalf("resolve edge: %v", err) + } + + if resolved.StartID != 101 || resolved.EndID != 202 || resolved.Kind.String() != "AdminTo" { + t.Fatalf("unexpected resolved edge: %+v", resolved) + } + + if resolved.Properties.Get("source").Any() != "test" { + t.Fatalf("unexpected resolved properties: %+v", resolved.Properties.Map) + } + + item.StartID = "missing" + + if _, err := resolveFragmentEdge(item, nodeMap); err == nil { + t.Fatalf("expected missing start node error") + } + + item.StartID = "1" + item.EndID = "missing" + + if _, err := resolveFragmentEdge(item, nodeMap); err == nil { + t.Fatalf("expected missing end node error") + } +} + +func TestReadLoadManifestRejectsNeo4jMultiGraph(t *testing.T) { + dir := t.TempDir() + value := newValidTestManifest(2) + value.Graphs = []GraphManifest{ + {Name: "a"}, + {Name: "b"}, + } + if err := writeManifest(dir, value); err != nil { + t.Fatalf("write Manifest: %v", err) + } + + if _, err := readLoadManifest(dir, pg.DriverName); err != nil { + t.Fatalf("read postgres load Manifest: %v", err) + } + + if _, err := readLoadManifest(dir, neo4j.DriverName); err == nil { + t.Fatalf("expected neo4j multi-graph load rejection") + } else { + var incompatibleErr IncompatibleDriverError + + if !errors.As(err, &incompatibleErr) { + t.Fatalf("expected IncompatibleDriverError, got %T: %v", err, err) + } + + if incompatibleErr.Driver != neo4j.DriverName || incompatibleErr.Operation != OperationLoad { + t.Fatalf("unexpected incompatible driver error: %+v", incompatibleErr) + } + } +} + +func TestVerifyManifestFilesRejectsBadChecksum(t *testing.T) { + dir := t.TempDir() + entry, err := writeCompressedJSON(filepath.Join(dir, "fragment.gz"), CompressionGzip, DefaultZstdLevel, NodeFragment{ + Phase: PhaseNodes, + Items: []FragmentNode{{ + ID: "1", + }}, + }) + if err != nil { + t.Fatalf("write fragment: %v", err) + } + + value := newValidTestManifest(1) + value.Graphs = []GraphManifest{{ + Name: "default", + NodeCount: 1, + Files: []FileManifest{{ + Phase: PhaseNodes, + Path: "fragment.gz", + Count: 1, + CompressedBytes: entry.CompressedBytes, + SHA256: "bad", + }}, + }} + + if err := verifyManifestFiles(dir, value); err == nil { + t.Fatalf("expected checksum verification failure") + } else { + var checksumErr ChecksumMismatchError + + if !errors.As(err, &checksumErr) { + t.Fatalf("expected ChecksumMismatchError, got %T: %v", err, err) + } + } + + value.Graphs[0].Files[0].SHA256 = entry.SHA256 + + if err := verifyManifestFiles(dir, value); err != nil { + t.Fatalf("verify Manifest files: %v", err) + } +} + +func TestReadFragmentFilesValidatePhaseAndCount(t *testing.T) { + dir := t.TempDir() + nodeEntry, err := writeCompressedJSON(filepath.Join(dir, "nodes.gz"), CompressionGzip, DefaultZstdLevel, NodeFragment{ + Phase: PhaseNodes, + Items: []FragmentNode{{ + ID: "1", + }}, + }) + if err != nil { + t.Fatalf("write node fragment: %v", err) + } + + nodeFile := FileManifest{ + Path: "nodes.gz", + Count: 1, + CompressedBytes: nodeEntry.CompressedBytes, + SHA256: nodeEntry.SHA256, + } + + if fragment, err := readNodeFragmentFile(dir, CompressionGzip, nodeFile); err != nil { + t.Fatalf("read node fragment: %v", err) + } else if len(fragment.Items) != 1 || fragment.Items[0].ID != "1" { + t.Fatalf("unexpected node fragment: %+v", fragment) + } + + nodeFile.Count = 2 + + if _, err := readNodeFragmentFile(dir, CompressionGzip, nodeFile); err == nil { + t.Fatalf("expected node count mismatch") + } + + wrongPhaseEntry, err := writeCompressedJSON(filepath.Join(dir, "wrong-Phase.gz"), CompressionGzip, DefaultZstdLevel, EdgeFragment{ + Phase: PhaseEdges, + Items: []FragmentEdge{{ + StartID: "1", + EndID: "2", + }}, + }) + if err != nil { + t.Fatalf("write wrong Phase fragment: %v", err) + } + + if _, err := readNodeFragmentFile(dir, CompressionGzip, FileManifest{ + Path: "wrong-Phase.gz", + Count: 1, + CompressedBytes: wrongPhaseEntry.CompressedBytes, + SHA256: wrongPhaseEntry.SHA256, + }); err == nil { + t.Fatalf("expected node Phase mismatch") + } else if !strings.Contains(err.Error(), "Phase") { + t.Fatalf("expected node Phase mismatch, got %v", err) + } + + edgeEntry, err := writeCompressedJSON(filepath.Join(dir, "edges.gz"), CompressionGzip, DefaultZstdLevel, EdgeFragment{ + Phase: PhaseEdges, + Items: []FragmentEdge{{ + StartID: "1", + EndID: "2", + Kind: "AdminTo", + }}, + }) + if err != nil { + t.Fatalf("write edge fragment: %v", err) + } + + edgeFile := FileManifest{ + Path: "edges.gz", + Count: 1, + CompressedBytes: edgeEntry.CompressedBytes, + SHA256: edgeEntry.SHA256, + } + if fragment, err := readEdgeFragmentFile(dir, CompressionGzip, edgeFile); err != nil { + t.Fatalf("read edge fragment: %v", err) + } else if len(fragment.Items) != 1 || fragment.Items[0].Kind != "AdminTo" { + t.Fatalf("unexpected edge fragment: %+v", fragment) + } + + edgeFile.Count = 2 + if _, err := readEdgeFragmentFile(dir, CompressionGzip, edgeFile); err == nil { + t.Fatalf("expected edge count mismatch") + } +} diff --git a/retriever/manifest.go b/retriever/manifest.go new file mode 100644 index 00000000..11c02ce1 --- /dev/null +++ b/retriever/manifest.go @@ -0,0 +1,77 @@ +package retriever + +import ( + "encoding/json" + "fmt" + "os" + "path/filepath" +) + +const manifestFileName = "manifest.json" + +const ManifestFileName = manifestFileName + +func ReadManifest(inputDir string) (Manifest, error) { + return readManifest(inputDir) +} + +func readManifest(inputDir string) (Manifest, error) { + var value Manifest + + manifestPath := filepath.Join(inputDir, manifestFileName) + if contents, err := os.ReadFile(manifestPath); err != nil { + return value, fmt.Errorf("read manifest: %w", err) + } else if err := json.Unmarshal(contents, &value); err != nil { + return value, fmt.Errorf("decode manifest: %w", err) + } else if err := value.validate(); err != nil { + return value, err + } + + return value, nil +} + +func WriteManifest(outputDir string, value Manifest) error { + return writeManifest(outputDir, value) +} + +func writeManifest(outputDir string, value Manifest) error { + if err := value.validate(); err != nil { + return err + } + + tempPath := filepath.Join(outputDir, manifestFileName+".tmp") + finalPath := filepath.Join(outputDir, manifestFileName) + payload, err := json.MarshalIndent(value, "", " ") + if err != nil { + return fmt.Errorf("encode manifest: %w", err) + } + + payload = append(payload, '\n') + + if err := os.WriteFile(tempPath, payload, 0o600); err != nil { + return fmt.Errorf("write manifest temp file: %w", err) + } + + if err := os.Rename(tempPath, finalPath); err != nil { + os.Remove(tempPath) + return fmt.Errorf("rename manifest: %w", err) + } + + return nil +} + +func VerifyManifestFiles(inputDir string, value Manifest) error { + return verifyManifestFiles(inputDir, value) +} + +func verifyManifestFiles(inputDir string, value Manifest) error { + for _, graphEntry := range value.Graphs { + for _, fileEntry := range graphEntry.Files { + if err := verifyChecksum(filepath.Join(inputDir, filepath.FromSlash(fileEntry.Path)), fileEntry.SHA256, fileEntry.CompressedBytes); err != nil { + return err + } + } + } + + return nil +} diff --git a/retriever/manifest_test.go b/retriever/manifest_test.go new file mode 100644 index 00000000..becbef84 --- /dev/null +++ b/retriever/manifest_test.go @@ -0,0 +1,225 @@ +package retriever + +import ( + "path/filepath" + "testing" +) + +func TestManifestValidateRejectsUnsupportedFormat(t *testing.T) { + value := newValidTestManifest(0) + value.Format = "future" + + if err := value.validate(); err == nil { + t.Fatalf("expected unsupported format error") + } +} + +func TestManifestValidateAcceptsLegacyFormat(t *testing.T) { + value := newValidTestManifest(0) + value.Format = legacyManifestFormat + + if err := value.validate(); err != nil { + t.Fatalf("validate legacy Manifest format: %v", err) + } +} + +func TestPublicManifestHelpers(t *testing.T) { + value := newValidTestManifest(0) + if err := value.Validate(); err != nil { + t.Fatalf("validate manifest through public method: %v", err) + } + + if !IsSupportedManifestFormat(manifestFormat) || !IsSupportedManifestFormat(legacyManifestFormat) { + t.Fatalf("expected public format helper to accept current and legacy formats") + } + + if IsSupportedManifestFormat("future") { + t.Fatalf("expected unsupported format to be rejected") + } + + graphValue := GraphSchemaFromMetadata(GraphSchemaMetadata{ + Name: "default", + NodeKinds: []string{"User"}, + EdgeKinds: []string{"MemberOf"}, + }) + if graphValue.Name != "default" || len(graphValue.Nodes) != 1 || graphValue.Nodes[0].String() != "User" { + t.Fatalf("unexpected graph schema conversion: %+v", graphValue) + } +} + +func TestManifestValidateRequiresNodeFilesBeforeEdgeFiles(t *testing.T) { + value := newValidTestManifest(1) + value.Graphs = []GraphManifest{{ + Name: "default", + NodeCount: 1, + EdgeCount: 1, + Files: []FileManifest{ + { + Phase: PhaseEdges, + Path: "graphs/default/edges-000001.ogfrag.gz", + Count: 1, + SHA256: "abc", + }, + { + Phase: PhaseNodes, + Path: "graphs/default/nodes-000001.ogfrag.gz", + Count: 1, + SHA256: "def", + }, + }, + }} + + if err := value.validate(); err == nil { + t.Fatalf("expected Phase ordering error") + } +} + +func TestManifestValidateRejectsMalformedGraphEntries(t *testing.T) { + validFile := FileManifest{ + Phase: PhaseNodes, + Path: "graphs/default/nodes-000001.ogfrag.gz", + Count: 1, + CompressedBytes: 10, + SHA256: "abc", + } + + cases := map[string]func(Manifest) Manifest{ + "graph count mismatch": func(value Manifest) Manifest { + value.Source.GraphCount = 2 + return value + }, + "duplicate graph": func(value Manifest) Manifest { + value.Source.GraphCount = 2 + value.Graphs = append(value.Graphs, value.Graphs[0]) + return value + }, + "empty graph name": func(value Manifest) Manifest { + value.Graphs[0].Name = "" + return value + }, + "unsupported Phase": func(value Manifest) Manifest { + value.Graphs[0].Files[0].Phase = Phase("bad") + return value + }, + "missing path": func(value Manifest) Manifest { + value.Graphs[0].Files[0].Path = "" + return value + }, + "missing checksum": func(value Manifest) Manifest { + value.Graphs[0].Files[0].SHA256 = "" + return value + }, + "negative count": func(value Manifest) Manifest { + value.Graphs[0].Files[0].Count = -1 + return value + }, + "negative bytes": func(value Manifest) Manifest { + value.Graphs[0].Files[0].CompressedBytes = -1 + return value + }, + "node count mismatch": func(value Manifest) Manifest { + value.Graphs[0].NodeCount = 2 + return value + }, + "edge count mismatch": func(value Manifest) Manifest { + value.Graphs[0].EdgeCount = 1 + return value + }, + "unsupported codec": func(value Manifest) Manifest { + value.Compression = CompressionCodec("zip") + return value + }, + "unsupported scrub mode": func(value Manifest) Manifest { + value.Scrub.Mode = ScrubMode("partial") + return value + }, + "unsupported id strategy": func(value Manifest) Manifest { + value.IDStrategy = "other" + return value + }, + } + + for name, mutate := range cases { + t.Run(name, func(t *testing.T) { + value := newValidTestManifest(1) + value.Graphs = []GraphManifest{{ + Name: "default", + NodeCount: 1, + Files: []FileManifest{validFile}, + }} + + if err := mutate(value).validate(); err == nil { + t.Fatalf("expected validation error") + } + }) + } +} + +func TestWriteReadManifest(t *testing.T) { + dir := t.TempDir() + value := newValidTestManifest(1) + value.Schema.Graphs = []GraphSchemaMetadata{{ + Name: "default", + NodeKinds: []string{"User"}, + EdgeKinds: []string{"AdminTo"}, + }} + value.Graphs = []GraphManifest{{ + Name: "default", + NodeCount: 0, + EdgeCount: 0, + }} + + if err := writeManifest(dir, value); err != nil { + t.Fatalf("write Manifest: %v", err) + } + + if _, err := readManifest(dir); err != nil { + t.Fatalf("read manifest: %v", err) + } + + if _, err := readManifest(filepath.Join(dir, "missing")); err == nil { + t.Fatalf("expected missing Manifest error") + } +} + +func TestManifestValidateAcceptsMetrics(t *testing.T) { + value := newValidTestManifest(1) + metrics := buildFingerprintFixture(t, []FragmentNode{{ + ID: "a", + Kinds: []string{"User"}, + }}, nil) + value.Graphs = []GraphManifest{{ + Name: "default", + NodeCount: 1, + EdgeCount: 0, + Files: []FileManifest{{ + Phase: PhaseNodes, + Path: "graphs/default/nodes-000001.ogfrag.gz", + Count: 1, + CompressedBytes: 1, + SHA256: "abc", + }}, + }} + value.Metrics = &MetricsManifest{ + Version: metricsVersion, + Graphs: []GraphMetrics{metrics}, + } + + if err := value.validate(); err != nil { + t.Fatalf("validate Manifest metrics: %v", err) + } + + value.Metrics.Graphs[0].Fingerprint = "sha256:bad" + + if err := value.validate(); err == nil { + t.Fatalf("expected metrics fingerprint validation error") + } +} + +func newValidTestManifest(graphCount int) Manifest { + return newManifest("pg", CompressionGzip, DefaultZstdLevel, ScrubMetadata{ + Mode: ScrubNone, + NodeActionCounts: map[string]int{}, + EdgeActionCounts: map[string]int{}, + }, graphCount) +} diff --git a/retriever/metrics.go b/retriever/metrics.go new file mode 100644 index 00000000..e05b84ac --- /dev/null +++ b/retriever/metrics.go @@ -0,0 +1,530 @@ +package retriever + +import ( + "crypto/sha256" + "encoding/hex" + "encoding/json" + "fmt" + "sort" + "strconv" + "strings" +) + +const ( + metricsVersion = "retriever-metrics-v1" + maxMetricDiffs = 20 + metricDiffsOmittedMessage = "... additional differences omitted" + metricsNoneKind = "0:" +) + +type MetricsManifest struct { + Version string `json:"version"` + Graphs []GraphMetrics `json:"graphs"` +} + +type GraphMetrics struct { + Name string `json:"name"` + NodeCount int64 `json:"node_count"` + EdgeCount int64 `json:"edge_count"` + NodeKindHistogram map[string]int64 `json:"node_kind_histogram"` + EdgeKindHistogram map[string]int64 `json:"edge_kind_histogram"` + InDegreeHistogram map[string]int64 `json:"in_degree_histogram"` + OutDegreeHistogram map[string]int64 `json:"out_degree_histogram"` + TotalDegreeHistogram map[string]int64 `json:"total_degree_histogram"` + EndpointKindHistogram map[string]int64 `json:"endpoint_kind_histogram"` + Fingerprint string `json:"fingerprint"` +} + +type MetricHistogramEntry struct { + Key string `json:"key"` + Count int64 `json:"count"` +} + +type canonicalGraphMetrics struct { + Name string `json:"name"` + NodeCount int64 `json:"node_count"` + EdgeCount int64 `json:"edge_count"` + NodeKindHistogram []MetricHistogramEntry `json:"node_kind_histogram"` + EdgeKindHistogram []MetricHistogramEntry `json:"edge_kind_histogram"` + InDegreeHistogram []MetricHistogramEntry `json:"in_degree_histogram"` + OutDegreeHistogram []MetricHistogramEntry `json:"out_degree_histogram"` + TotalDegreeHistogram []MetricHistogramEntry `json:"total_degree_histogram"` + EndpointKindHistogram []MetricHistogramEntry `json:"endpoint_kind_histogram"` +} + +type metricsBuilder struct { + graphName string + nodeCount int64 + edgeCount int64 + nodeKindByID map[string]string + inDegreeByID map[string]uint64 + outDegreeByID map[string]uint64 + nodeKindHistogram map[string]int64 + edgeKindHistogram map[string]int64 + endpointKindHistogram map[string]int64 +} + +func newMetricsManifest(graphCount int) MetricsManifest { + return MetricsManifest{ + Version: metricsVersion, + Graphs: make([]GraphMetrics, 0, graphCount), + } +} + +func newMetricsBuilder(graphName string, expectedNodeCount int64) *metricsBuilder { + if expectedNodeCount < 0 { + expectedNodeCount = 0 + } + + return &metricsBuilder{ + graphName: graphName, + nodeKindByID: make(map[string]string, expectedNodeCount), + inDegreeByID: make(map[string]uint64, expectedNodeCount), + outDegreeByID: make(map[string]uint64, expectedNodeCount), + nodeKindHistogram: map[string]int64{}, + edgeKindHistogram: map[string]int64{}, + endpointKindHistogram: map[string]int64{}, + } +} + +func (s *metricsBuilder) observeFragmentNode(node FragmentNode) error { + return s.observeNode(node.ID, node.Kinds) +} + +func (s *metricsBuilder) observeNode(id string, kinds []string) error { + if strings.TrimSpace(id) == "" { + return fmt.Errorf("metrics node observation has empty ID") + } + + if _, seen := s.nodeKindByID[id]; seen { + return fmt.Errorf("metrics node observation has duplicate ID %q", id) + } + + kindKey := metricKindSetKey(kinds) + s.nodeKindByID[id] = kindKey + s.nodeKindHistogram[kindKey]++ + s.nodeCount++ + + return nil +} + +func (s *metricsBuilder) observeFragmentEdge(edge FragmentEdge) error { + return s.observeRelationship(edge.StartID, edge.EndID, edge.Kind) +} + +func (s *metricsBuilder) observeRelationship(startID, endID, kind string) error { + startKindKey, startFound := s.nodeKindByID[startID] + endKindKey, endFound := s.nodeKindByID[endID] + if !startFound || !endFound { + return fmt.Errorf("metrics relationship observation references an endpoint missing from the node scan") + } + + edgeKindKey := metricKindKey(kind) + s.edgeKindHistogram[edgeKindKey]++ + s.endpointKindHistogram[metricEndpointKindKey(startKindKey, edgeKindKey, endKindKey)]++ + s.outDegreeByID[startID]++ + s.inDegreeByID[endID]++ + s.edgeCount++ + + return nil +} + +func (s *metricsBuilder) finalize() GraphMetrics { + value := GraphMetrics{ + Name: s.graphName, + NodeCount: s.nodeCount, + EdgeCount: s.edgeCount, + NodeKindHistogram: cloneMetricHistogram(s.nodeKindHistogram), + EdgeKindHistogram: cloneMetricHistogram(s.edgeKindHistogram), + InDegreeHistogram: map[string]int64{}, + OutDegreeHistogram: map[string]int64{}, + TotalDegreeHistogram: map[string]int64{}, + EndpointKindHistogram: cloneMetricHistogram(s.endpointKindHistogram), + } + + for id := range s.nodeKindByID { + inDegree := s.inDegreeByID[id] + outDegree := s.outDegreeByID[id] + + value.InDegreeHistogram[metricDegreeKey(inDegree)]++ + value.OutDegreeHistogram[metricDegreeKey(outDegree)]++ + value.TotalDegreeHistogram[metricDegreeKey(inDegree+outDegree)]++ + } + + value.Fingerprint = fingerprintGraphMetrics(value) + + return value +} + +func fingerprintGraphMetrics(value GraphMetrics) string { + payload, err := json.Marshal(canonicalizeGraphMetrics(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 FingerprintGraphMetrics(value GraphMetrics) string { + return fingerprintGraphMetrics(value) +} + +func canonicalizeGraphMetrics(value GraphMetrics) canonicalGraphMetrics { + return canonicalGraphMetrics{ + Name: value.Name, + NodeCount: value.NodeCount, + EdgeCount: value.EdgeCount, + NodeKindHistogram: canonicalMetricHistogram(value.NodeKindHistogram), + EdgeKindHistogram: canonicalMetricHistogram(value.EdgeKindHistogram), + InDegreeHistogram: canonicalMetricHistogram(value.InDegreeHistogram), + OutDegreeHistogram: canonicalMetricHistogram(value.OutDegreeHistogram), + TotalDegreeHistogram: canonicalMetricHistogram(value.TotalDegreeHistogram), + EndpointKindHistogram: canonicalMetricHistogram(value.EndpointKindHistogram), + } +} + +func canonicalMetricHistogram(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 +} + +func compareMetricsManifest(expected, actual MetricsManifest) []string { + var differences []string + if expected.Version != actual.Version { + differences = appendMetricDiff(differences, fmt.Sprintf("metrics.version: expected %q, actual %q", expected.Version, actual.Version)) + } + + actualByGraph := map[string]GraphMetrics{} + for _, actualGraph := range actual.Graphs { + actualByGraph[actualGraph.Name] = actualGraph + } + + expectedGraphNames := make([]string, 0, len(expected.Graphs)) + expectedByGraph := map[string]struct{}{} + for _, expectedGraph := range expected.Graphs { + expectedGraphNames = append(expectedGraphNames, expectedGraph.Name) + expectedByGraph[expectedGraph.Name] = struct{}{} + } + + sort.Strings(expectedGraphNames) + + for _, graphName := range expectedGraphNames { + expectedGraph := findGraphMetrics(expected.Graphs, graphName) + actualGraph, ok := actualByGraph[graphName] + if !ok { + differences = appendMetricDiff(differences, fmt.Sprintf("graph %q missing from actual metrics", graphName)) + continue + } + + differences = appendMetricDiffs(differences, compareGraphMetrics(expectedGraph, actualGraph)...) + } + + actualGraphNames := make([]string, 0, len(actual.Graphs)) + for _, actualGraph := range actual.Graphs { + if _, ok := expectedByGraph[actualGraph.Name]; !ok { + actualGraphNames = append(actualGraphNames, actualGraph.Name) + } + } + + sort.Strings(actualGraphNames) + + for _, graphName := range actualGraphNames { + differences = appendMetricDiff(differences, fmt.Sprintf("graph %q only present in actual metrics", graphName)) + } + + return differences +} + +func CompareMetricsManifest(expected, actual MetricsManifest) []string { + return compareMetricsManifest(expected, actual) +} + +func compareGraphMetrics(expected, actual GraphMetrics) []string { + var differences []string + graphPrefix := fmt.Sprintf("graph %q", expected.Name) + + if expected.Name != actual.Name { + differences = appendMetricDiff(differences, fmt.Sprintf("graph name: expected %q, actual %q", expected.Name, actual.Name)) + } + + if expected.NodeCount != actual.NodeCount { + differences = appendMetricDiff(differences, fmt.Sprintf("%s node_count: expected %d, actual %d", graphPrefix, expected.NodeCount, actual.NodeCount)) + } + + if expected.EdgeCount != actual.EdgeCount { + differences = appendMetricDiff(differences, fmt.Sprintf("%s edge_count: expected %d, actual %d", graphPrefix, expected.EdgeCount, actual.EdgeCount)) + } + + differences = appendMetricDiffs(differences, compareMetricHistogram(graphPrefix, "node_kind_histogram", expected.NodeKindHistogram, actual.NodeKindHistogram)...) + differences = appendMetricDiffs(differences, compareMetricHistogram(graphPrefix, "edge_kind_histogram", expected.EdgeKindHistogram, actual.EdgeKindHistogram)...) + differences = appendMetricDiffs(differences, compareMetricHistogram(graphPrefix, "in_degree_histogram", expected.InDegreeHistogram, actual.InDegreeHistogram)...) + differences = appendMetricDiffs(differences, compareMetricHistogram(graphPrefix, "out_degree_histogram", expected.OutDegreeHistogram, actual.OutDegreeHistogram)...) + differences = appendMetricDiffs(differences, compareMetricHistogram(graphPrefix, "total_degree_histogram", expected.TotalDegreeHistogram, actual.TotalDegreeHistogram)...) + differences = appendMetricDiffs(differences, compareMetricHistogram(graphPrefix, "endpoint_kind_histogram", expected.EndpointKindHistogram, actual.EndpointKindHistogram)...) + + if len(differences) == 0 && expected.Fingerprint != actual.Fingerprint { + differences = appendMetricDiff(differences, fmt.Sprintf("%s fingerprint: expected %q, actual %q", graphPrefix, expected.Fingerprint, actual.Fingerprint)) + } + + return differences +} + +func compareMetricHistogram(prefix, field string, expected, actual map[string]int64) []string { + keys := map[string]struct{}{} + 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) + + var differences []string + for _, key := range sortedKeys { + expectedCount := expected[key] + actualCount := actual[key] + + if expectedCount != actualCount { + differences = appendMetricDiff(differences, fmt.Sprintf("%s %s[%q]: expected %d, actual %d", prefix, field, key, expectedCount, actualCount)) + } + } + + return differences +} + +func appendMetricDiffs(target []string, values ...string) []string { + for _, value := range values { + target = appendMetricDiff(target, value) + } + + return target +} + +func appendMetricDiff(target []string, value string) []string { + if len(target) > maxMetricDiffs { + return target + } + + if len(target) == maxMetricDiffs { + return append(target, metricDiffsOmittedMessage) + } + + return append(target, value) +} + +func findGraphMetrics(values []GraphMetrics, graphName string) GraphMetrics { + for _, value := range values { + if value.Name == graphName { + return value + } + } + + return GraphMetrics{} +} + +func validateMetricsManifest(value MetricsManifest, graphEntries []GraphManifest) error { + if value.Version != metricsVersion { + return fmt.Errorf("unsupported metrics version %q", value.Version) + } + + if len(value.Graphs) != len(graphEntries) { + return fmt.Errorf("metrics graph count %d does not match %d manifest graph entries", len(value.Graphs), len(graphEntries)) + } + + expectedGraphs := map[string]GraphManifest{} + for _, graphEntry := range graphEntries { + expectedGraphs[graphEntry.Name] = graphEntry + } + + seenGraphs := map[string]struct{}{} + for _, graphEntry := range value.Graphs { + if graphEntry.Name == "" { + return fmt.Errorf("metrics graph entry has empty name") + } + + if _, seen := seenGraphs[graphEntry.Name]; seen { + return fmt.Errorf("metrics contains duplicate graph %q", graphEntry.Name) + } + + seenGraphs[graphEntry.Name] = struct{}{} + + manifestGraph, expected := expectedGraphs[graphEntry.Name] + if !expected { + return fmt.Errorf("metrics contains graph %q not present in manifest graphs", graphEntry.Name) + } + + if graphEntry.NodeCount != manifestGraph.NodeCount { + return fmt.Errorf("metrics graph %q node_count %d does not match manifest node_count %d", graphEntry.Name, graphEntry.NodeCount, manifestGraph.NodeCount) + } + + if graphEntry.EdgeCount != manifestGraph.EdgeCount { + return fmt.Errorf("metrics graph %q edge_count %d does not match manifest edge_count %d", graphEntry.Name, graphEntry.EdgeCount, manifestGraph.EdgeCount) + } + + if err := validateGraphMetricHistograms(graphEntry); err != nil { + return err + } + + if expectedFingerprint := fingerprintGraphMetrics(graphEntry); graphEntry.Fingerprint != expectedFingerprint { + return fmt.Errorf("metrics graph %q fingerprint %q does not match computed fingerprint %q", graphEntry.Name, graphEntry.Fingerprint, expectedFingerprint) + } + } + + return nil +} + +func ValidateMetricsManifest(value MetricsManifest, graphEntries []GraphManifest) error { + return validateMetricsManifest(value, graphEntries) +} + +func validateGraphMetricHistograms(value GraphMetrics) error { + if value.NodeCount < 0 { + return fmt.Errorf("metrics graph %q has negative node_count %d", value.Name, value.NodeCount) + } + + if value.EdgeCount < 0 { + return fmt.Errorf("metrics graph %q has negative edge_count %d", value.Name, value.EdgeCount) + } + + for _, histogram := range []struct { + name string + values map[string]int64 + expected int64 + }{ + { + name: "node_kind_histogram", + values: value.NodeKindHistogram, + expected: value.NodeCount, + }, + { + name: "in_degree_histogram", + values: value.InDegreeHistogram, + expected: value.NodeCount, + }, + { + name: "out_degree_histogram", + values: value.OutDegreeHistogram, + expected: value.NodeCount, + }, + { + name: "total_degree_histogram", + values: value.TotalDegreeHistogram, + expected: value.NodeCount, + }, + { + name: "edge_kind_histogram", + values: value.EdgeKindHistogram, + expected: value.EdgeCount, + }, + { + name: "endpoint_kind_histogram", + values: value.EndpointKindHistogram, + expected: value.EdgeCount, + }, + } { + if err := validateMetricHistogramSum(value.Name, histogram.name, histogram.values, histogram.expected); err != nil { + return err + } + } + + return nil +} + +func validateMetricHistogramSum(graphName string, histogramName string, histogram map[string]int64, expected int64) error { + var total int64 + for key, count := range histogram { + if count < 0 { + return fmt.Errorf("metrics graph %q %s[%q] has negative count %d", graphName, histogramName, key, count) + } + + total += count + } + + if total != expected { + return fmt.Errorf("metrics graph %q %s total %d does not match expected count %d", graphName, histogramName, total, expected) + } + + return nil +} + +func cloneMetricHistogram(source map[string]int64) map[string]int64 { + target := make(map[string]int64, len(source)) + for key, count := range source { + target[key] = count + } + + return target +} + +func metricKindSetKey(kinds []string) string { + seen := map[string]struct{}{} + for _, kind := range kinds { + if kind != "" { + seen[kind] = struct{}{} + } + } + + values := make([]string, 0, len(seen)) + for kind := range seen { + values = append(values, kind) + } + + if len(values) == 0 { + return metricsNoneKind + } + + sort.Strings(values) + + parts := make([]string, 0, len(values)) + for _, value := range values { + parts = append(parts, metricKeyPart(value)) + } + + return strings.Join(parts, "+") +} + +func metricKindKey(kind string) string { + if kind == "" { + return metricsNoneKind + } + + return metricKeyPart(kind) +} + +func metricEndpointKindKey(startKindKey, edgeKindKey, endKindKey string) string { + return strings.Join([]string{ + metricKeyPart(startKindKey), + metricKeyPart(edgeKindKey), + metricKeyPart(endKindKey), + }, "|") +} + +func metricDegreeKey(degree uint64) string { + return strconv.FormatUint(degree, 10) +} + +func metricKeyPart(value string) string { + return strconv.Itoa(len(value)) + ":" + value +} diff --git a/retriever/metrics_test.go b/retriever/metrics_test.go new file mode 100644 index 00000000..6e355e54 --- /dev/null +++ b/retriever/metrics_test.go @@ -0,0 +1,367 @@ +package retriever + +import ( + "encoding/json" + "strings" + "testing" +) + +func TestMetricsBuilderFinalizesGraphShape(t *testing.T) { + builder := newMetricsBuilder("default", 3) + for _, node := range []FragmentNode{ + { + ID: "node-secret-1", + Kinds: []string{"User"}, + Properties: map[string]any{ + "objectid": "S-1-5-21-secret", + "name": "alice", + }, + }, + { + ID: "node-secret-2", + Kinds: []string{"Group"}, + }, + { + ID: "node-secret-3", + }, + } { + if err := builder.observeFragmentNode(node); err != nil { + t.Fatalf("observe node: %v", err) + } + } + + for _, edge := range []FragmentEdge{ + { + StartID: "node-secret-1", + EndID: "node-secret-2", + Kind: "MemberOf", + Properties: map[string]any{ + "source": "secret-source", + }, + }, + { + StartID: "node-secret-2", + EndID: "node-secret-1", + Kind: "AdminTo", + }, + } { + if err := builder.observeFragmentEdge(edge); err != nil { + t.Fatalf("observe edge: %v", err) + } + } + + metrics := builder.finalize() + + if metrics.Name != "default" { + t.Fatalf("graph name = %q", metrics.Name) + } + + if metrics.NodeCount != 3 { + t.Fatalf("node count = %d", metrics.NodeCount) + } + + if metrics.EdgeCount != 2 { + t.Fatalf("edge count = %d", metrics.EdgeCount) + } + + if got := metrics.NodeKindHistogram[metricKindSetKey([]string{"User"})]; got != 1 { + t.Fatalf("user node kind count = %d", got) + } + + if got := metrics.NodeKindHistogram[metricKindSetKey([]string{"Group"})]; got != 1 { + t.Fatalf("group node kind count = %d", got) + } + + if got := metrics.NodeKindHistogram[metricKindSetKey(nil)]; got != 1 { + t.Fatalf("empty node kind count = %d", got) + } + + if got := metrics.EdgeKindHistogram[metricKindKey("MemberOf")]; got != 1 { + t.Fatalf("member edge kind count = %d", got) + } + + if got := metrics.EdgeKindHistogram[metricKindKey("AdminTo")]; got != 1 { + t.Fatalf("admin edge kind count = %d", got) + } + + if got := metrics.InDegreeHistogram[metricDegreeKey(0)]; got != 1 { + t.Fatalf("zero in-degree count = %d", got) + } + + if got := metrics.OutDegreeHistogram[metricDegreeKey(0)]; got != 1 { + t.Fatalf("zero out-degree count = %d", got) + } + + if got := metrics.TotalDegreeHistogram[metricDegreeKey(0)]; got != 1 { + t.Fatalf("zero total-degree count = %d", got) + } + + if got := metrics.TotalDegreeHistogram[metricDegreeKey(2)]; got != 2 { + t.Fatalf("total-degree two count = %d", got) + } + + endpointKey := metricEndpointKindKey(metricKindSetKey([]string{"User"}), metricKindKey("MemberOf"), metricKindSetKey([]string{"Group"})) + + if got := metrics.EndpointKindHistogram[endpointKey]; got != 1 { + t.Fatalf("endpoint kind count = %d", got) + } + + if metrics.Fingerprint == "" || !strings.HasPrefix(metrics.Fingerprint, "sha256:") { + t.Fatalf("missing fingerprint: %q", metrics.Fingerprint) + } +} + +func TestMetricsFingerprintStableAcrossObservationOrder(t *testing.T) { + first := buildNamedFingerprintFixture(t, "default", []FragmentNode{ + {ID: "a", Kinds: []string{"User", "Person"}}, + {ID: "b", Kinds: []string{"Group"}}, + }, []FragmentEdge{ + {StartID: "a", EndID: "b", Kind: "MemberOf"}, + }) + second := buildNamedFingerprintFixture(t, "default", []FragmentNode{ + {ID: "b", Kinds: []string{"Group"}}, + {ID: "a", Kinds: []string{"Person", "User"}}, + }, []FragmentEdge{ + {StartID: "a", EndID: "b", Kind: "MemberOf"}, + }) + + if first.Fingerprint != second.Fingerprint { + t.Fatalf("fingerprint changed with observation order: %q != %q", first.Fingerprint, second.Fingerprint) + } +} + +func TestMetricKindSetKeyDeduplicatesKinds(t *testing.T) { + once := metricKindSetKey([]string{"User", "Person"}) + duplicated := metricKindSetKey([]string{"Person", "User", "User"}) + if once != duplicated { + t.Fatalf("duplicate kind changed kind set key: %q != %q", once, duplicated) + } +} + +func TestMetricsIgnorePropertiesAndIDsInFinalForm(t *testing.T) { + builder := newMetricsBuilder("default", 1) + if err := builder.observeFragmentNode(FragmentNode{ + ID: "node-secret-id", + Kinds: []string{"User"}, + Properties: map[string]any{ + "objectid": "S-1-5-21-secret", + "name": "alice@example.test", + }, + }); err != nil { + t.Fatalf("observe node: %v", err) + } + + metrics := builder.finalize() + payload, err := json.Marshal(metrics) + if err != nil { + t.Fatalf("marshal metrics: %v", err) + } + + serialized := string(payload) + + for _, forbidden := range []string{"node-secret-id", "objectid", "S-1-5-21-secret", "alice@example.test"} { + if strings.Contains(serialized, forbidden) { + t.Fatalf("serialized metrics contains attribution %q: %s", forbidden, serialized) + } + } +} + +func TestMetricsComparisonReportsDeterministicDiffs(t *testing.T) { + expected := buildNamedFingerprintFixture(t, "default", []FragmentNode{ + {ID: "a", Kinds: []string{"User"}}, + }, nil) + actual := expected + actual.NodeCount = 2 + actual.NodeKindHistogram = map[string]int64{ + metricKindSetKey([]string{"Group"}): 1, + metricKindSetKey([]string{"User"}): 1, + } + actual.Fingerprint = fingerprintGraphMetrics(actual) + + differences := compareGraphMetrics(expected, actual) + + if len(differences) != 2 { + t.Fatalf("diff count = %d: %v", len(differences), differences) + } + + if !strings.Contains(differences[0], "node_count") { + t.Fatalf("first diff should be node count, got %q", differences[0]) + } + + if !strings.Contains(differences[1], "node_kind_histogram") { + t.Fatalf("second diff should be kind histogram, got %q", differences[1]) + } +} + +func TestMetricsManifestComparisonReportsGraphSetDiffs(t *testing.T) { + expected := MetricsManifest{ + Version: metricsVersion, + Graphs: []GraphMetrics{ + buildNamedFingerprintFixture(t, "b", []FragmentNode{{ID: "b", Kinds: []string{"User"}}}, nil), + buildNamedFingerprintFixture(t, "a", []FragmentNode{{ID: "a", Kinds: []string{"User"}}}, nil), + }, + } + actual := MetricsManifest{ + Version: "future", + Graphs: []GraphMetrics{ + buildNamedFingerprintFixture(t, "c", []FragmentNode{{ID: "c", Kinds: []string{"User"}}}, nil), + }, + } + + differences := compareMetricsManifest(expected, actual) + + for _, expectedDiff := range []string{ + `metrics.version: expected "retriever-metrics-v1", actual "future"`, + `graph "a" missing from actual metrics`, + `graph "b" missing from actual metrics`, + `graph "c" only present in actual metrics`, + } { + if !containsExact(differences, expectedDiff) { + t.Fatalf("missing diff %q in %v", expectedDiff, differences) + } + } +} + +func TestPublicMetricsHelpers(t *testing.T) { + expectedGraph := buildNamedFingerprintFixture(t, "default", []FragmentNode{{ID: "a", Kinds: []string{"User"}}}, nil) + expected := MetricsManifest{ + Version: metricsVersion, + Graphs: []GraphMetrics{expectedGraph}, + } + actualGraph := expectedGraph + actualGraph.NodeCount = 2 + actualGraph.Fingerprint = FingerprintGraphMetrics(actualGraph) + actual := MetricsManifest{ + Version: metricsVersion, + Graphs: []GraphMetrics{actualGraph}, + } + + differences := CompareMetricsManifest(expected, actual) + if len(differences) == 0 { + t.Fatalf("expected public compare helper to report differences") + } + + if err := ValidateMetricsManifest(expected, []GraphManifest{{ + Name: "default", + NodeCount: 1, + EdgeCount: 0, + }}); err != nil { + t.Fatalf("validate metrics through public helper: %v", err) + } +} + +func TestAppendMetricDiffReportsTruncation(t *testing.T) { + var differences []string + for i := 0; i < maxMetricDiffs+5; i++ { + differences = appendMetricDiff(differences, "diff") + } + + if len(differences) != maxMetricDiffs+1 { + t.Fatalf("diff count = %d", len(differences)) + } + + if differences[len(differences)-1] != metricDiffsOmittedMessage { + t.Fatalf("missing omitted marker: %v", differences) + } +} + +func TestMetricsValidationRejectsMismatchedFingerprint(t *testing.T) { + metrics := buildNamedFingerprintFixture(t, "default", []FragmentNode{ + {ID: "a", Kinds: []string{"User"}}, + }, nil) + metrics.Fingerprint = "sha256:bad" + + err := validateMetricsManifest(MetricsManifest{ + Version: metricsVersion, + Graphs: []GraphMetrics{metrics}, + }, []GraphManifest{{ + Name: "default", + NodeCount: 1, + EdgeCount: 0, + }}) + if err == nil || !strings.Contains(err.Error(), "fingerprint") { + t.Fatalf("expected fingerprint validation error, got %v", err) + } +} + +func TestMetricsValidationRejectsHistogramInvariantViolations(t *testing.T) { + base := buildNamedFingerprintFixture(t, "default", []FragmentNode{ + {ID: "a", Kinds: []string{"User"}}, + {ID: "b", Kinds: []string{"Group"}}, + }, []FragmentEdge{ + {StartID: "a", EndID: "b", Kind: "MemberOf"}, + }) + + cases := map[string]func(GraphMetrics) GraphMetrics{ + "node kind sum mismatch": func(value GraphMetrics) GraphMetrics { + value.NodeKindHistogram = map[string]int64{ + metricKindSetKey([]string{"User"}): 1, + } + value.Fingerprint = fingerprintGraphMetrics(value) + return value + }, + "edge kind sum mismatch": func(value GraphMetrics) GraphMetrics { + value.EdgeKindHistogram = map[string]int64{} + value.Fingerprint = fingerprintGraphMetrics(value) + return value + }, + "negative count": func(value GraphMetrics) GraphMetrics { + value.OutDegreeHistogram = map[string]int64{ + metricDegreeKey(0): -1, + metricDegreeKey(1): 3, + } + value.Fingerprint = fingerprintGraphMetrics(value) + return value + }, + } + + for name, mutate := range cases { + t.Run(name, func(t *testing.T) { + err := validateMetricsManifest(MetricsManifest{ + Version: metricsVersion, + Graphs: []GraphMetrics{mutate(base)}, + }, []GraphManifest{{ + Name: "default", + NodeCount: 2, + EdgeCount: 1, + }}) + if err == nil { + t.Fatalf("expected invariant validation error") + } + }) + } +} + +func buildFingerprintFixture(t *testing.T, nodes []FragmentNode, edges []FragmentEdge) GraphMetrics { + t.Helper() + return buildNamedFingerprintFixture(t, "default", nodes, edges) +} + +func buildNamedFingerprintFixture(t *testing.T, graphName string, nodes []FragmentNode, edges []FragmentEdge) GraphMetrics { + t.Helper() + + builder := newMetricsBuilder(graphName, int64(len(nodes))) + for _, node := range nodes { + if err := builder.observeFragmentNode(node); err != nil { + t.Fatalf("observe node: %v", err) + } + } + + for _, edge := range edges { + if err := builder.observeFragmentEdge(edge); err != nil { + t.Fatalf("observe edge: %v", err) + } + } + + return builder.finalize() +} + +func containsExact(values []string, expected string) bool { + for _, value := range values { + if value == expected { + return true + } + } + + return false +} diff --git a/retriever/options.go b/retriever/options.go new file mode 100644 index 00000000..4eaba173 --- /dev/null +++ b/retriever/options.go @@ -0,0 +1,293 @@ +package retriever + +import ( + "crypto/hpke" + "fmt" + "io" + "strings" + "time" +) + +const ( + DefaultShardSize = 100_000 + DefaultBatchSize = 10_000 + DefaultProgressInterval = 250_000 +) + +const DefaultGraphName = "default" + +const ( + OperationDump = "dump" + OperationLoad = "load" + OperationVerify = "verify" + OperationArchive = "archive" + OperationUnpack = "unpack" + OperationKeygen = "keygen" +) + +type GraphTarget struct { + Name string +} + +type ProgressEvent struct { + Operation string + Message string + Driver string + Graph string + Phase Phase + InputDir string + OutputDir string + ArchivePath string + GraphIndex int + GraphCount int + FileCount int + BatchSize int + ShardSize int + Processed int64 + Planned int64 + NodeCount int64 + EdgeCount int64 + Compression CompressionCodec + Scrub ScrubMode + Elapsed time.Duration + EntitiesPerSecond float64 +} + +type ProgressFunc func(ProgressEvent) + +func (s ProgressFunc) emit(event ProgressEvent) { + if s != nil { + s(event) + } +} + +type DumpOptions struct { + OutputDir string + Force bool + Scrub ScrubMode + Salt string + ScrubConfig io.Reader + Compression CompressionCodec + ZstdLevel int + ShardSize int + BatchSize int + ProgressInterval int64 + Progress ProgressFunc +} + +func DefaultDumpOptions(outputDir string) DumpOptions { + return DumpOptions{ + OutputDir: outputDir, + Scrub: ScrubNone, + Compression: CompressionZstd, + ZstdLevel: DefaultZstdLevel, + ShardSize: DefaultShardSize, + BatchSize: DefaultBatchSize, + ProgressInterval: DefaultProgressInterval, + } +} + +func (s DumpOptions) Validate() error { + if strings.TrimSpace(s.OutputDir) == "" { + return ValidationError{Message: "output directory is required; pass -out"} + } + + if err := ValidateCompression(s.Compression); err != nil { + return err + } + + if s.ZstdLevel <= 0 { + return ValidationError{Message: "zstd-level must be > 0"} + } + + if s.ShardSize <= 0 { + return ValidationError{Message: "shard-size must be > 0"} + } + + if s.BatchSize <= 0 { + return ValidationError{Message: "batch-size must be > 0"} + } + + switch s.Scrub { + case ScrubNone: + return nil + case ScrubFull: + if strings.TrimSpace(s.Salt) == "" { + return ValidationError{Message: "-scrub full requires -salt, RETRIEVER_SCRUB_SALT, or legacy RETRIEVR_SCRUB_SALT; refusing to write scrubbed output without deterministic pseudonymization"} + } + return nil + default: + return ValidationError{Message: fmt.Sprintf("unsupported scrub mode %q", s.Scrub)} + } +} + +func (s DumpOptions) validate() error { + return s.Validate() +} + +type LoadOptions struct { + InputDir string + ArchiveReader io.Reader + ArchiveIdentity hpke.PrivateKey + BatchSize int + ProgressInterval int64 + VerifyMetrics bool + Progress ProgressFunc +} + +func DefaultLoadOptions(inputDir string) LoadOptions { + return LoadOptions{ + InputDir: inputDir, + BatchSize: DefaultBatchSize, + ProgressInterval: DefaultProgressInterval, + } +} + +func (s LoadOptions) Validate() error { + inputDir := strings.TrimSpace(s.InputDir) + hasArchive := s.ArchiveReader != nil + + if inputDir != "" && hasArchive { + return ValidationError{Message: "load accepts either an input directory or archive reader, not both"} + } + + if inputDir == "" && !hasArchive { + return ValidationError{Message: "input directory or archive reader is required"} + } + + if hasArchive && s.ArchiveIdentity == nil { + return ValidationError{Message: "archive reader requires archive identity"} + } + + if !hasArchive && s.ArchiveIdentity != nil { + return ValidationError{Message: "archive identity requires archive reader"} + } + + if s.BatchSize <= 0 { + return ValidationError{Message: "batch-size must be > 0"} + } + + return nil +} + +func (s LoadOptions) validate() error { + return s.Validate() +} + +type UnpackOptions struct { + ArchiveReader io.Reader + ArchiveIdentity hpke.PrivateKey + OutputDir string + Force bool + Progress ProgressFunc +} + +func DefaultUnpackOptions(outputDir string) UnpackOptions { + return UnpackOptions{ + OutputDir: outputDir, + } +} + +func (s UnpackOptions) Validate() error { + if s.ArchiveReader == nil { + return ValidationError{Message: "archive reader is required"} + } + + if s.ArchiveIdentity == nil { + return ValidationError{Message: "archive identity is required"} + } + + if strings.TrimSpace(s.OutputDir) == "" { + return ValidationError{Message: "output directory is required; pass -out"} + } + + return nil +} + +func (s UnpackOptions) validate() error { + return s.Validate() +} + +type KeygenOptions struct { + PrivatePath string + PublicPath string +} + +func DefaultKeygenOptions(privatePath, publicPath string) KeygenOptions { + return KeygenOptions{ + PrivatePath: privatePath, + PublicPath: publicPath, + } +} + +func (s KeygenOptions) Validate() error { + if strings.TrimSpace(s.PrivatePath) == "" { + return ValidationError{Message: "private key path is required; pass -private"} + } + + if strings.TrimSpace(s.PublicPath) == "" { + return ValidationError{Message: "public key path is required; pass -public"} + } + + if sameCleanPath(s.PrivatePath, s.PublicPath) { + return ValidationError{Message: "private and public key paths must be different"} + } + + return nil +} + +func (s KeygenOptions) validate() error { + return s.Validate() +} + +type VerifyOptions struct { + InputDir string + BatchSize int + ProgressInterval int64 + Progress ProgressFunc +} + +func DefaultVerifyOptions(inputDir string) VerifyOptions { + return VerifyOptions{ + InputDir: inputDir, + BatchSize: DefaultBatchSize, + ProgressInterval: DefaultProgressInterval, + } +} + +func (s VerifyOptions) Validate() error { + if strings.TrimSpace(s.InputDir) == "" { + return ValidationError{Message: "input directory is required; pass -in"} + } + + if s.BatchSize <= 0 { + return ValidationError{Message: "batch-size must be > 0"} + } + + return nil +} + +func (s VerifyOptions) validate() error { + return s.Validate() +} + +type ValidationError struct { + Message string + Err error +} + +func (s ValidationError) Error() string { + if s.Err != nil { + return s.Message + ": " + s.Err.Error() + } + + return s.Message +} + +func (s ValidationError) Unwrap() error { + return s.Err +} + +type ArchiveOptions struct { + Progress ProgressFunc +} diff --git a/retriever/options_test.go b/retriever/options_test.go new file mode 100644 index 00000000..89e096a4 --- /dev/null +++ b/retriever/options_test.go @@ -0,0 +1,187 @@ +package retriever + +import ( + "bytes" + "crypto/hpke" + "path/filepath" + "testing" +) + +func TestDumpOptionsScrubFullRequiresSalt(t *testing.T) { + options := DumpOptions{ + OutputDir: t.TempDir(), + Scrub: ScrubFull, + Compression: CompressionZstd, + ZstdLevel: DefaultZstdLevel, + ShardSize: DefaultShardSize, + BatchSize: DefaultBatchSize, + } + + if err := options.Validate(); err == nil { + t.Fatalf("expected missing salt error") + } +} + +func TestOptionsValidate(t *testing.T) { + dump := DumpOptions{ + OutputDir: t.TempDir(), + Scrub: ScrubNone, + Compression: CompressionGzip, + ZstdLevel: DefaultZstdLevel, + ShardSize: DefaultShardSize, + BatchSize: DefaultBatchSize, + } + if err := dump.Validate(); err != nil { + t.Fatalf("valid dump options: %v", err) + } + + dump.Compression = CompressionCodec("zip") + + if err := dump.Validate(); err == nil { + t.Fatalf("expected invalid compression") + } + + load := LoadOptions{ + InputDir: t.TempDir(), + BatchSize: 1, + } + if err := load.Validate(); err != nil { + t.Fatalf("valid load options: %v", err) + } + + load.InputDir = "" + + if err := load.Validate(); err == nil { + t.Fatalf("expected missing input") + } + + load.ArchiveReader = bytes.NewReader(nil) + + if err := load.Validate(); err == nil { + t.Fatalf("expected missing archive identity") + } + + privateKey, _, err := generateArchiveKeyPairForOptions() + if err != nil { + t.Fatalf("generate archive key: %v", err) + } + + load.ArchiveIdentity = privateKey + + if err := load.Validate(); err != nil { + t.Fatalf("valid archive load options: %v", err) + } + + load.InputDir = t.TempDir() + + if err := load.Validate(); err == nil { + t.Fatalf("expected mutually exclusive load input error") + } + + unpack := UnpackOptions{ + ArchiveReader: bytes.NewReader(nil), + ArchiveIdentity: privateKey, + OutputDir: t.TempDir(), + } + if err := unpack.Validate(); err != nil { + t.Fatalf("valid unpack options: %v", err) + } + + unpack.ArchiveIdentity = nil + + if err := unpack.Validate(); err == nil { + t.Fatalf("expected missing identity") + } + + keygen := KeygenOptions{ + PrivatePath: filepath.Join(t.TempDir(), "private.key"), + PublicPath: filepath.Join(t.TempDir(), "public.key"), + } + if err := keygen.Validate(); err != nil { + t.Fatalf("valid keygen options: %v", err) + } + + keygen.PublicPath = keygen.PrivatePath + + if err := keygen.Validate(); err == nil { + t.Fatalf("expected duplicate key path validation error") + } + + verify := VerifyOptions{ + InputDir: t.TempDir(), + BatchSize: 1, + } + if err := verify.Validate(); err != nil { + t.Fatalf("valid verify options: %v", err) + } + + verify.BatchSize = 0 + + if err := verify.Validate(); err == nil { + t.Fatalf("expected invalid verify batch size") + } +} + +func TestDefaultOptions(t *testing.T) { + dump := DefaultDumpOptions(t.TempDir()) + if dump.Scrub != ScrubNone || dump.Compression != CompressionZstd || dump.ZstdLevel != DefaultZstdLevel { + t.Fatalf("unexpected dump defaults: %+v", dump) + } + if dump.ShardSize != DefaultShardSize || dump.BatchSize != DefaultBatchSize || dump.ProgressInterval != DefaultProgressInterval { + t.Fatalf("unexpected dump sizing defaults: %+v", dump) + } + + if err := dump.Validate(); err != nil { + t.Fatalf("validate default dump options: %v", err) + } + + load := DefaultLoadOptions(t.TempDir()) + if load.BatchSize != DefaultBatchSize || load.ProgressInterval != DefaultProgressInterval { + t.Fatalf("unexpected load defaults: %+v", load) + } + + if err := load.Validate(); err != nil { + t.Fatalf("validate default load options: %v", err) + } + + verify := DefaultVerifyOptions(t.TempDir()) + if verify.BatchSize != DefaultBatchSize || verify.ProgressInterval != DefaultProgressInterval { + t.Fatalf("unexpected verify defaults: %+v", verify) + } + + if err := verify.Validate(); err != nil { + t.Fatalf("validate default verify options: %v", err) + } + + unpack := DefaultUnpackOptions(t.TempDir()) + if unpack.OutputDir == "" { + t.Fatalf("expected unpack output directory") + } + + keygen := DefaultKeygenOptions("private.key", "public.key") + if keygen.PrivatePath != "private.key" || keygen.PublicPath != "public.key" { + t.Fatalf("unexpected keygen defaults: %+v", keygen) + } +} + +func TestPrepareOutputDirectoryForce(t *testing.T) { + dir := t.TempDir() + writeTestFile(t, filepath.Join(dir, "old"), []byte("old")) + + if err := prepareOutputDirectory(dir, false); err == nil { + t.Fatalf("expected non-empty directory error") + } + + if err := prepareOutputDirectory(dir, true); err != nil { + t.Fatalf("prepare output directory with force: %v", err) + } +} + +func generateArchiveKeyPairForOptions() (hpke.PrivateKey, hpke.PublicKey, error) { + privateKey, err := defaultArchiveKEM().GenerateKey() + if err != nil { + return nil, nil, err + } + + return privateKey, privateKey.PublicKey(), nil +} diff --git a/retriever/progress.go b/retriever/progress.go new file mode 100644 index 00000000..59fcddbc --- /dev/null +++ b/retriever/progress.go @@ -0,0 +1,97 @@ +package retriever + +import ( + "log/slog" + "time" +) + +const retrieverProgressEntityInterval int64 = DefaultProgressInterval + +func retrieverInitialProgressAt(planned int64) int64 { + return retrieverInitialProgressAtInterval(planned, retrieverProgressEntityInterval) +} + +func retrieverInitialProgressAtInterval(planned int64, interval int64) int64 { + interval = normalizedProgressInterval(interval) + if planned <= interval { + return 0 + } + + return interval +} + +func retrieverBatchLimit(remaining int64, batchSize int) int { + if remaining <= 0 { + return 0 + } + + if int64(batchSize) > remaining { + return int(remaining) + } + + return batchSize +} + +func logRetrieverEntityProgress(message string, graphName string, phaseName Phase, processed int64, planned int64, startedAt time.Time, nextProgressAt int64, progress ProgressFunc) int64 { + return logRetrieverEntityProgressInterval(message, graphName, phaseName, processed, planned, startedAt, nextProgressAt, progress, retrieverProgressEntityInterval) +} + +func logRetrieverEntityProgressInterval(message string, graphName string, phaseName Phase, processed int64, planned int64, startedAt time.Time, nextProgressAt int64, progress ProgressFunc, interval int64) int64 { + if nextProgressAt == 0 || processed < nextProgressAt || processed >= planned { + return nextProgressAt + } + + elapsed := time.Since(startedAt) + slog.Info(message, + slog.String("graph", graphName), + slog.String("phase", string(phaseName)), + slog.Int64("processed", processed), + slog.Int64("planned_count", planned), + slog.Duration("wall_elapsed", elapsed), + slog.Float64("entities_per_second", perSecond(processed, elapsed)), + ) + progress.emit(ProgressEvent{ + Message: message, + Graph: graphName, + Phase: phaseName, + Processed: processed, + Planned: planned, + Elapsed: elapsed, + EntitiesPerSecond: perSecond(processed, elapsed), + }) + + return retrieverNextProgressAtInterval(processed, planned, nextProgressAt, interval) +} + +func retrieverNextProgressAt(processed int64, planned int64, nextProgressAt int64) int64 { + return retrieverNextProgressAtInterval(processed, planned, nextProgressAt, retrieverProgressEntityInterval) +} + +func retrieverNextProgressAtInterval(processed int64, planned int64, nextProgressAt int64, interval int64) int64 { + interval = normalizedProgressInterval(interval) + if nextProgressAt <= processed { + nextProgressAt += ((processed-nextProgressAt)/interval + 1) * interval + } + + if nextProgressAt >= planned { + return 0 + } + + return nextProgressAt +} + +func normalizedProgressInterval(interval int64) int64 { + if interval <= 0 { + return retrieverProgressEntityInterval + } + + return interval +} + +func perSecond(count int64, elapsed time.Duration) float64 { + if count == 0 || elapsed <= 0 { + return 0 + } + + return float64(count) / elapsed.Seconds() +} diff --git a/retriever/progress_test.go b/retriever/progress_test.go new file mode 100644 index 00000000..ee70a9c9 --- /dev/null +++ b/retriever/progress_test.go @@ -0,0 +1,64 @@ +package retriever + +import "testing" + +func TestRetrieverNextProgressAtInterval(t *testing.T) { + cases := []struct { + name string + processed int64 + planned int64 + nextProgressAt int64 + interval int64 + want int64 + }{ + { + name: "before threshold", + processed: 9, + planned: 100, + nextProgressAt: 10, + interval: 10, + want: 10, + }, + { + name: "at threshold", + processed: 10, + planned: 100, + nextProgressAt: 10, + interval: 10, + want: 20, + }, + { + name: "large jump", + processed: 95, + planned: 200, + nextProgressAt: 10, + interval: 10, + want: 100, + }, + { + name: "planned boundary", + processed: 95, + planned: 100, + nextProgressAt: 10, + interval: 10, + want: 0, + }, + { + name: "default interval", + processed: retrieverProgressEntityInterval * 2, + planned: retrieverProgressEntityInterval * 4, + nextProgressAt: retrieverProgressEntityInterval, + interval: 0, + want: retrieverProgressEntityInterval * 3, + }, + } + + for _, testCase := range cases { + t.Run(testCase.name, func(t *testing.T) { + got := retrieverNextProgressAtInterval(testCase.processed, testCase.planned, testCase.nextProgressAt, testCase.interval) + if got != testCase.want { + t.Fatalf("next progress = %d, want %d", got, testCase.want) + } + }) + } +} diff --git a/retriever/scan.go b/retriever/scan.go new file mode 100644 index 00000000..938eda05 --- /dev/null +++ b/retriever/scan.go @@ -0,0 +1,170 @@ +package retriever + +import ( + "context" + "fmt" + "time" + + "github.com/specterops/dawgs/graph" +) + +type entityBatchReader[T any] func(afterID graph.ID, hasAfterID bool, batchSize int) ([]T, error) +type entityBatchHandler[T any] func([]T) error +type entityIDFunc[T any] func(T) graph.ID +type entityProgressLogger func(processed int64, startedAt time.Time, nextProgressAt int64) int64 + +type entityScanOptions[T any] struct { + Total int64 + BatchSize int + ProgressInterval int64 + EntityName string + Read entityBatchReader[T] + ID entityIDFunc[T] + Handle entityBatchHandler[T] + LogProgress entityProgressLogger +} + +func scanEntityBatches[T any](options entityScanOptions[T]) (int64, error) { + entityName := options.EntityName + if entityName == "" { + entityName = "entity" + } + + if options.Total <= 0 { + return 0, nil + } + + if options.BatchSize <= 0 { + return 0, fmt.Errorf("%s batch size must be > 0", entityName) + } + + if options.Read == nil { + return 0, fmt.Errorf("%s batch reader is required", entityName) + } + + if options.ID == nil { + return 0, fmt.Errorf("%s ID accessor is required", entityName) + } + + var ( + lastID graph.ID + hasLastID bool + processed int64 + + startedAt = time.Now() + nextProgressAt = retrieverInitialProgressAtInterval(options.Total, options.ProgressInterval) + ) + + for processed < options.Total { + remaining := options.Total - processed + batch, err := options.Read(lastID, hasLastID, retrieverBatchLimit(remaining, options.BatchSize)) + if err != nil { + return processed, err + } + + if int64(len(batch)) > remaining { + batch = batch[:int(remaining)] + } + + if len(batch) == 0 { + break + } + + nextID := options.ID(batch[len(batch)-1]) + if hasLastID && nextID <= lastID { + return processed, fmt.Errorf("%s keyset scan did not advance after ID %d", entityName, lastID.Uint64()) + } + + if options.Handle != nil { + if err := options.Handle(batch); err != nil { + return processed, err + } + } + + processed += int64(len(batch)) + lastID = nextID + hasLastID = true + + if options.LogProgress != nil { + nextProgressAt = options.LogProgress(processed, startedAt, nextProgressAt) + } + } + + return processed, nil +} + +func scanDatabaseNodes( + ctx context.Context, + db graph.Database, + targetGraph graph.Graph, + total int64, + batchSize int, + handle entityBatchHandler[*graph.Node], + logProgress entityProgressLogger, +) (int64, error) { + return scanDatabaseNodesWithProgressInterval(ctx, db, targetGraph, total, batchSize, 0, handle, logProgress) +} + +func scanDatabaseNodesWithProgressInterval( + ctx context.Context, + db graph.Database, + targetGraph graph.Graph, + total int64, + batchSize int, + progressInterval int64, + handle entityBatchHandler[*graph.Node], + logProgress entityProgressLogger, +) (int64, error) { + return scanEntityBatches(entityScanOptions[*graph.Node]{ + Total: total, + BatchSize: batchSize, + ProgressInterval: progressInterval, + EntityName: "node", + Read: func(afterID graph.ID, hasAfterID bool, limit int) ([]*graph.Node, error) { + return readDatabaseNodes(ctx, db, targetGraph, afterID, hasAfterID, limit) + }, + ID: func(node *graph.Node) graph.ID { + return node.ID + }, + Handle: handle, + LogProgress: logProgress, + }) +} + +func scanDatabaseRelationships( + ctx context.Context, + db graph.Database, + targetGraph graph.Graph, + total int64, + batchSize int, + handle entityBatchHandler[*graph.Relationship], + logProgress entityProgressLogger, +) (int64, error) { + return scanDatabaseRelationshipsWithProgressInterval(ctx, db, targetGraph, total, batchSize, 0, handle, logProgress) +} + +func scanDatabaseRelationshipsWithProgressInterval( + ctx context.Context, + db graph.Database, + targetGraph graph.Graph, + total int64, + batchSize int, + progressInterval int64, + handle entityBatchHandler[*graph.Relationship], + logProgress entityProgressLogger, +) (int64, error) { + return scanEntityBatches(entityScanOptions[*graph.Relationship]{ + Total: total, + BatchSize: batchSize, + ProgressInterval: progressInterval, + EntityName: "relationship", + Read: func(afterID graph.ID, hasAfterID bool, limit int) ([]*graph.Relationship, error) { + return readDatabaseRelationships(ctx, db, targetGraph, afterID, hasAfterID, limit) + }, + ID: func(relationship *graph.Relationship) graph.ID { + return relationship.ID + }, + Handle: handle, + LogProgress: logProgress, + }) +} diff --git a/retriever/scan_test.go b/retriever/scan_test.go new file mode 100644 index 00000000..d502c334 --- /dev/null +++ b/retriever/scan_test.go @@ -0,0 +1,221 @@ +package retriever + +import ( + "errors" + "reflect" + "strings" + "testing" + "time" + + "github.com/specterops/dawgs/graph" +) + +type scanTestEntity struct { + id graph.ID +} + +type scanReadCall struct { + afterID graph.ID + hasAfterID bool + limit int +} + +func TestScanEntityBatchesUsesKeysetAndRemainingLimit(t *testing.T) { + entities := []scanTestEntity{{id: 1}, {id: 2}, {id: 3}, {id: 4}} + var ( + calls []scanReadCall + handled []graph.ID + progress []int64 + ) + + processed, err := scanEntityBatches(entityScanOptions[scanTestEntity]{ + Total: 3, + BatchSize: 2, + EntityName: "node", + Read: func(afterID graph.ID, hasAfterID bool, limit int) ([]scanTestEntity, error) { + calls = append(calls, scanReadCall{ + afterID: afterID, + hasAfterID: hasAfterID, + limit: limit, + }) + var batch []scanTestEntity + for _, entity := range entities { + if !hasAfterID || entity.id > afterID { + batch = append(batch, entity) + } + } + + if len(batch) > limit { + batch = batch[:limit] + } + + return batch, nil + }, + ID: func(entity scanTestEntity) graph.ID { + return entity.id + }, + Handle: func(batch []scanTestEntity) error { + for _, entity := range batch { + handled = append(handled, entity.id) + } + + return nil + }, + LogProgress: func(processed int64, _ time.Time, nextProgressAt int64) int64 { + progress = append(progress, processed) + return nextProgressAt + }, + }) + if err != nil { + t.Fatalf("scan entity batches: %v", err) + } + + if processed != 3 { + t.Fatalf("processed = %d", processed) + } + + if !reflect.DeepEqual(handled, []graph.ID{1, 2, 3}) { + t.Fatalf("handled IDs = %v", handled) + } + + expectedCalls := []scanReadCall{ + {limit: 2}, + {afterID: 2, hasAfterID: true, limit: 1}, + } + + if !reflect.DeepEqual(calls, expectedCalls) { + t.Fatalf("read calls = %+v", calls) + } + + if !reflect.DeepEqual(progress, []int64{2, 3}) { + t.Fatalf("progress = %v", progress) + } +} + +func TestScanEntityBatchesTruncatesOverfullFinalBatch(t *testing.T) { + var handled []graph.ID + processed, err := scanEntityBatches(entityScanOptions[scanTestEntity]{ + Total: 2, + BatchSize: 10, + EntityName: "node", + Read: func(graph.ID, bool, int) ([]scanTestEntity, error) { + return []scanTestEntity{{id: 1}, {id: 2}, {id: 3}}, nil + }, + ID: func(entity scanTestEntity) graph.ID { + return entity.id + }, + Handle: func(batch []scanTestEntity) error { + for _, entity := range batch { + handled = append(handled, entity.id) + } + + return nil + }, + }) + + if err != nil { + t.Fatalf("scan entity batches: %v", err) + } + + if processed != 2 || !reflect.DeepEqual(handled, []graph.ID{1, 2}) { + t.Fatalf("processed=%d handled=%v", processed, handled) + } +} + +func TestScanEntityBatchesUsesProgressInterval(t *testing.T) { + var firstNextProgressAt int64 + processed, err := scanEntityBatches(entityScanOptions[scanTestEntity]{ + Total: 10, + BatchSize: 1, + ProgressInterval: 3, + EntityName: "node", + Read: func(afterID graph.ID, hasAfterID bool, _ int) ([]scanTestEntity, error) { + if !hasAfterID { + return []scanTestEntity{{id: 1}}, nil + } + + return []scanTestEntity{{id: afterID + 1}}, nil + }, + ID: func(entity scanTestEntity) graph.ID { + return entity.id + }, + LogProgress: func(_ int64, _ time.Time, nextProgressAt int64) int64 { + if firstNextProgressAt == 0 { + firstNextProgressAt = nextProgressAt + } + return nextProgressAt + }, + }) + if err != nil { + t.Fatalf("scan entity batches: %v", err) + } + + if processed != 10 { + t.Fatalf("processed = %d", processed) + } + + if firstNextProgressAt != 3 { + t.Fatalf("first next progress threshold = %d, want 3", firstNextProgressAt) + } +} + +func TestScanEntityBatchesReportsNonAdvancingCursor(t *testing.T) { + processed, err := scanEntityBatches(entityScanOptions[scanTestEntity]{ + Total: 2, + BatchSize: 1, + EntityName: "relationship", + Read: func(afterID graph.ID, hasAfterID bool, _ int) ([]scanTestEntity, error) { + if hasAfterID { + return []scanTestEntity{{id: afterID}}, nil + } + + return []scanTestEntity{{id: 7}}, nil + }, + ID: func(entity scanTestEntity) graph.ID { + return entity.id + }, + }) + if err == nil || !strings.Contains(err.Error(), "relationship keyset scan did not advance after ID 7") { + t.Fatalf("expected non-advancing cursor error, got %v", err) + } + + if processed != 1 { + t.Fatalf("processed = %d", processed) + } +} + +func TestScanEntityBatchesValidationAndReaderErrors(t *testing.T) { + if processed, err := scanEntityBatches(entityScanOptions[scanTestEntity]{Total: 0}); err != nil || processed != 0 { + t.Fatalf("empty scan processed=%d err=%v", processed, err) + } + + if _, err := scanEntityBatches(entityScanOptions[scanTestEntity]{Total: 1, BatchSize: 0}); err == nil { + t.Fatalf("expected batch size validation error") + } + + if _, err := scanEntityBatches(entityScanOptions[scanTestEntity]{ + Total: 1, + BatchSize: 1, + Read: func(graph.ID, bool, int) ([]scanTestEntity, error) { + return nil, nil + }, + }); err == nil { + t.Fatalf("expected ID accessor validation error") + } + + readerErr := errors.New("reader failed") + processed, err := scanEntityBatches(entityScanOptions[scanTestEntity]{ + Total: 1, + BatchSize: 1, + EntityName: "node", + Read: func(graph.ID, bool, int) ([]scanTestEntity, error) { + return nil, readerErr + }, + ID: func(entity scanTestEntity) graph.ID { + return entity.id + }, + }) + if !errors.Is(err, readerErr) || processed != 0 { + t.Fatalf("processed=%d err=%v", processed, err) + } +} diff --git a/retriever/scrubber.go b/retriever/scrubber.go new file mode 100644 index 00000000..bb877734 --- /dev/null +++ b/retriever/scrubber.go @@ -0,0 +1,675 @@ +package retriever + +import ( + "crypto/hmac" + "crypto/sha256" + _ "embed" + "encoding/hex" + "fmt" + "io" + "regexp" + "strings" + "time" + + "github.com/pelletier/go-toml/v2" +) + +const scrubRulesVersion = "retriever-scrub-v2" + +type PropertyAction string + +const ( + actionPreserve PropertyAction = "preserve" + actionPseudonymize PropertyAction = "pseudonymize" + actionRedact PropertyAction = "redact" + actionShiftTimestamp PropertyAction = "shift_timestamp" +) + +type PropertyPlan struct { + Key string + Action PropertyAction + Shape string +} + +type ScrubberConfig struct { + Salt string `toml:"salt"` + 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"` +} + +type ScrubberFileConfig struct { + Scrub ScrubberConfig `toml:"scrub"` + Classifier ClassifierConfig `toml:"classifier"` +} + +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"` +} + +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"` +} + +type ValueShapeConfig struct { + Name string `toml:"name"` + Pattern string `toml:"pattern"` +} + +type compiledShape struct { + name string + pattern *regexp.Regexp +} + +type scrubber struct { + config ScrubberConfig + preserveKeys map[string]struct{} + referenceKeys map[string]struct{} + shapeRules []compiledShape + identifierLookup map[string]string +} + +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)`) +) + +//go:embed defaults.toml +var defaultScrubberConfigTOML []byte + +func DefaultScrubberConfig() ScrubberConfig { + cfg, err := decodeScrubberFileConfig(defaultScrubberConfigTOML, ScrubberConfig{}) + if err != nil { + panic(fmt.Sprintf("parse embedded scrub defaults: %v", err)) + } + + return cfg +} + +func defaultScrubberConfig() ScrubberConfig { + return DefaultScrubberConfig() +} + +func ReadScrubberConfig(reader io.Reader, base ScrubberConfig) (ScrubberConfig, error) { + contents, err := io.ReadAll(reader) + if err != nil { + return ScrubberConfig{}, fmt.Errorf("read scrub config: %w", err) + } + + return decodeScrubberFileConfig(contents, base) +} + +func decodeScrubberFileConfig(contents []byte, base ScrubberConfig) (ScrubberConfig, error) { + fileCfg := ScrubberFileConfig{ + Scrub: base, + Classifier: base.Classifier, + } + if err := toml.Unmarshal(contents, &fileCfg); err != nil { + return ScrubberConfig{}, err + } + + cfg := fileCfg.Scrub + cfg.Classifier = fileCfg.Classifier + + return cfg, nil +} + +func newScrubber(configReader io.Reader, salt string) (*scrubber, error) { + cfg := defaultScrubberConfig() + if configReader != nil { + if decoded, err := ReadScrubberConfig(configReader, cfg); err != nil { + return nil, fmt.Errorf("parse scrub config: %w", err) + } else { + cfg = decoded + } + } + + cfg.Salt = strings.TrimSpace(salt) + cfg.FakeDomain = strings.Trim(strings.ToLower(strings.TrimSpace(cfg.FakeDomain)), ".") + cfg.RedactionMarker = strings.TrimSpace(cfg.RedactionMarker) + + if cfg.RedactionMarker == "" { + cfg.RedactionMarker = "[REDACTED]" + } + + if cfg.Classifier.LongTextThreshold <= 0 { + cfg.Classifier.LongTextThreshold = 512 + } + + if cfg.TimestampShiftDays == 0 { + cfg.TimestampShiftDays = 17 + } + + preserveKeys := make(map[string]struct{}, len(cfg.Classifier.PreserveKeys)) + for _, key := range cfg.Classifier.PreserveKeys { + preserveKeys[normalizeKey(key)] = struct{}{} + } + + referenceKeys := map[string]struct{}{} + for _, keys := range [][]string{ + cfg.GraphRules.DomainSIDReferenceKeys, + cfg.GraphRules.ObjectIDReferenceKeys, + cfg.GraphRules.SelfObjectIDAliasKeys, + cfg.GraphRules.DomainNameReferenceKeys, + {cfg.GraphRules.ObjectIDKey, cfg.GraphRules.DomainNameKey}, + } { + for _, key := range keys { + if normalized := normalizeKey(key); normalized != "" { + referenceKeys[normalized] = struct{}{} + } + } + } + + shapeRules := make([]compiledShape, 0, len(cfg.Classifier.ValueShapePatterns)) + for _, shape := range cfg.Classifier.ValueShapePatterns { + if strings.TrimSpace(shape.Name) == "" || strings.TrimSpace(shape.Pattern) == "" { + continue + } + + if pattern, err := regexp.Compile(shape.Pattern); err != nil { + return nil, fmt.Errorf("compile value shape %q: %w", shape.Name, err) + } else { + shapeRules = append(shapeRules, compiledShape{ + name: shape.Name, + pattern: pattern, + }) + } + } + + return &scrubber{ + config: cfg, + preserveKeys: preserveKeys, + referenceKeys: referenceKeys, + shapeRules: shapeRules, + identifierLookup: map[string]string{}, + }, nil +} + +func (s *scrubber) metadata() ScrubMetadata { + return ScrubMetadata{ + Mode: ScrubFull, + RulesVersion: scrubRulesVersion, + SaltProvided: strings.TrimSpace(s.config.Salt) != "", + NodeActionCounts: map[string]int{}, + EdgeActionCounts: map[string]int{}, + } +} + +func (s *scrubber) observeNode(properties map[string]any) { + for key, value := range properties { + normalized := normalizeKey(key) + if _, ok := s.referenceKeys[normalized]; !ok { + continue + } + + s.observeIdentifier(value) + } +} + +func (s *scrubber) observeIdentifier(value any) { + switch typed := value.(type) { + case string: + trimmed := strings.TrimSpace(typed) + if trimmed != "" { + s.identifierLookup[trimmed] = s.pseudonymizeString(trimmed, s.classifyString(trimmed)) + } + + case []any: + for _, item := range typed { + s.observeIdentifier(item) + } + + case []string: + for _, item := range typed { + s.observeIdentifier(item) + } + } +} + +func (s *scrubber) scrubProperties(properties map[string]any) (map[string]any, map[string]int) { + if properties == nil { + return map[string]any{}, map[string]int{} + } + + scrubbed := make(map[string]any, len(properties)) + actionCounts := map[string]int{} + for key, value := range properties { + nextValue, plan := s.scrubProperty(key, value) + scrubbed[key] = nextValue + actionCounts[string(plan.Action)] += 1 + } + + return scrubbed, actionCounts +} + +func (s *scrubber) scrubProperty(key string, value any) (any, PropertyPlan) { + plan := s.planProperty(key, value) + return s.scrubWithPlan(key, value, plan), plan +} + +func (s *scrubber) planProperty(key string, value any) PropertyPlan { + normalizedKey := normalizeKey(key) + plan := PropertyPlan{ + Key: key, + Action: actionPreserve, + } + + if _, ok := s.referenceKeys[normalizedKey]; ok && isStringLike(value) { + plan.Action = actionPseudonymize + plan.Shape = s.classifyValue(value) + + return plan + } + + if _, ok := s.preserveKeys[normalizedKey]; ok { + return plan + } + + if isTimestampKey(normalizedKey) { + plan.Action = actionShiftTimestamp + + return plan + } + + if isFreeTextKey(normalizedKey) { + plan.Action = actionRedact + + return plan + } + + if isPathKey(normalizedKey) { + plan.Action = actionPseudonymize + plan.Shape = s.classifyValue(value) + + return plan + } + + if isScriptKey(normalizedKey) { + plan.Action = actionPseudonymize + plan.Shape = s.classifyValue(value) + + return plan + } + + if s.shouldRedact(normalizedKey, value) { + plan.Action = actionRedact + + return plan + } + + if shape := s.classifyValue(value); shape != "" { + plan.Action = actionPseudonymize + plan.Shape = shape + + return plan + } + + if s.isSensitiveKey(normalizedKey) { + plan.Action = actionPseudonymize + + return plan + } + + if isSemanticOrgKey(normalizedKey) || isStringLike(value) { + plan.Action = actionPseudonymize + plan.Shape = s.classifyValue(value) + + return plan + } + + return plan +} + +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) scrubWithPlan(key string, value any, plan PropertyPlan) any { + switch plan.Action { + case actionPreserve: + return value + case actionRedact: + return s.redact(value) + case actionShiftTimestamp: + return s.shiftTimestamp(value) + case actionPseudonymize: + return s.pseudonymizeValue(key, value, plan.Shape) + 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.config.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) isSensitiveKey(normalizedKey string) bool { + for _, marker := range s.config.Classifier.SensitiveKeyMarks { + if marker = normalizeKey(marker); marker != "" && strings.Contains(normalizedKey, marker) { + return true + } + } + + return false +} + +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 + } +} + +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) redact(value any) any { + switch typed := value.(type) { + case []any: + values := make([]any, len(typed)) + for idx := range values { + values[idx] = s.config.RedactionMarker + } + + return values + + case []string: + values := make([]string, len(typed)) + for idx := range values { + values[idx] = s.config.RedactionMarker + } + + return values + + case map[string]any: + values := make(map[string]any, len(typed)) + for key := range typed { + values[key] = s.config.RedactionMarker + } + + return values + + default: + return s.config.RedactionMarker + } +} + +func (s *scrubber) shiftTimestamp(value any) any { + shift := time.Duration(s.config.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 { + shifted := s.shiftTimestamp(item) + if stringValue, ok := shifted.(string); ok { + values = append(values, stringValue) + } + } + + 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 replacement, ok := s.identifierLookup[strings.TrimSpace(typed)]; ok { + return replacement + } + + 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) 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.config.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.config.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 (s *scrubber) digest(value string) string { + mac := hmac.New(sha256.New, []byte(s.config.Salt)) + mac.Write([]byte(value)) + + return hex.EncodeToString(mac.Sum(nil)) +} + +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, "_", "") + value = strings.ReplaceAll(value, " ", "") + return 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") +} diff --git a/retriever/scrubber_test.go b/retriever/scrubber_test.go new file mode 100644 index 00000000..86e111dc --- /dev/null +++ b/retriever/scrubber_test.go @@ -0,0 +1,403 @@ +package retriever + +import ( + "os" + "path/filepath" + "regexp" + "strings" + "testing" + "time" +) + +func TestScrubberPseudonymizesSensitiveValues(t *testing.T) { + scrubber, err := newScrubber(nil, "test-salt") + if err != nil { + t.Fatalf("new scrubber: %v", err) + } + + scrubbed, counts := scrubber.scrubProperties(map[string]any{ + "name": "Alice", + "email": "alice@example.com", + "password": "super-secret", + }) + + for key, raw := range map[string]string{ + "name": "Alice", + "email": "alice@example.com", + "password": "super-secret", + } { + if scrubbed[key] == raw { + t.Fatalf("expected %s to be scrubbed", key) + } + } + + if counts[string(actionPseudonymize)] == 0 { + t.Fatalf("expected pseudonymize action count") + } + + if counts[string(actionRedact)] == 0 { + t.Fatalf("expected redact action count") + } +} + +func TestScrubberDeterministic(t *testing.T) { + left, err := newScrubber(nil, "test-salt") + if err != nil { + t.Fatalf("new left scrubber: %v", err) + } + + right, err := newScrubber(nil, "test-salt") + if err != nil { + t.Fatalf("new right scrubber: %v", err) + } + + leftValue, _ := left.scrubProperty("email", "alice@example.com") + rightValue, _ := right.scrubProperty("email", "alice@example.com") + + if leftValue != rightValue { + t.Fatalf("expected deterministic scrub output, got %q and %q", leftValue, rightValue) + } +} + +func TestScrubberConfigFileShape(t *testing.T) { + configPath := filepath.Join(t.TempDir(), "retriever.toml") + writeTestFile(t, configPath, []byte(` +[scrub] +fake_domain = "scrub.example" +redaction_marker = "[X]" + +[classifier] +long_text_threshold = 8 +`)) + + configFile, err := os.Open(configPath) + if err != nil { + t.Fatalf("open scrub config: %v", err) + } + defer configFile.Close() + + scrubber, err := newScrubber(configFile, "test-salt") + if err != nil { + t.Fatalf("new scrubber: %v", err) + } + + if scrubber.config.FakeDomain != "scrub.example" { + t.Fatalf("fake domain = %q", scrubber.config.FakeDomain) + } + + if scrubber.config.RedactionMarker != "[X]" { + t.Fatalf("redaction marker = %q", scrubber.config.RedactionMarker) + } + + if scrubber.config.Classifier.LongTextThreshold != 8 { + t.Fatalf("long text threshold = %d", scrubber.config.Classifier.LongTextThreshold) + } +} + +func TestScrubberShapeSpecificPseudonyms(t *testing.T) { + scrubber, err := newScrubber(nil, "test-salt") + if err != nil { + t.Fatalf("new scrubber: %v", err) + } + + cases := []struct { + name string + value string + shape string + pattern *regexp.Regexp + }{ + { + name: "email", + value: "alice@example.com", + shape: "email", + pattern: regexp.MustCompile(`^user-[0-9a-f]{12}@example\.invalid$`), + }, + { + name: "uuid", + value: "00112233-4455-6677-8899-aabbccddeeff", + shape: "uuid", + pattern: regexp.MustCompile(`^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$`), + }, + { + name: "domain sid", + value: "S-1-5-21-1-2-3", + shape: "domain_sid", + pattern: regexp.MustCompile(`^S-1-5-21-\d{9}-\d{9}-\d{9}$`), + }, + { + name: "object sid", + value: "S-1-5-21-1-2-3-500", + shape: "object_sid", + pattern: regexp.MustCompile(`^S-1-5-21-\d{9}-\d{9}-\d{9}-500$`), + }, + { + name: "object sid shape fallback", + value: "not-a-sid", + shape: "object_sid", + pattern: regexp.MustCompile(`^value-[0-9a-f]{16}$`), + }, + { + name: "ipv4", + value: "192.0.2.10", + shape: "ipv4", + pattern: regexp.MustCompile(`^10\.\d+\.\d+\.\d+$`), + }, + { + name: "host", + value: "server.example.com", + shape: "host", + pattern: regexp.MustCompile(`^host-[0-9a-f]{12}\.example\.invalid$`), + }, + { + name: "generic", + value: "Alice", + shape: "", + pattern: regexp.MustCompile(`^value-[0-9a-f]{16}$`), + }, + } + + for _, testCase := range cases { + t.Run(testCase.name, func(t *testing.T) { + got := scrubber.pseudonymizeString(testCase.value, testCase.shape) + if got == testCase.value { + t.Fatalf("value was not pseudonymized") + } + + if !testCase.pattern.MatchString(got) { + t.Fatalf("pseudonym %q did not match %s", got, testCase.pattern) + } + }) + } +} + +func TestScrubberTimestampAndRedactionBranches(t *testing.T) { + scrubber, err := newScrubber(nil, "test-salt") + if err != nil { + t.Fatalf("new scrubber: %v", err) + } + + shifted := scrubber.shiftTimestamp("2026-01-01T00:00:00Z") + if shifted != "2026-01-18T00:00:00Z" { + t.Fatalf("shifted timestamp = %v", shifted) + } + + shiftedTime := scrubber.shiftTimestamp(time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC)) + if shiftedTime != "2026-01-18T00:00:00Z" { + t.Fatalf("shifted time.Time = %v", shiftedTime) + } + + shiftedSlice := scrubber.shiftTimestamp([]string{"2026-01-01T00:00:00Z"}) + if values, ok := shiftedSlice.([]string); !ok || len(values) != 1 || values[0] != "2026-01-18T00:00:00Z" { + t.Fatalf("shifted slice = %#v", shiftedSlice) + } + + redacted := scrubber.redact(map[string]any{"a": "secret", "b": "secret"}) + redactedMap, ok := redacted.(map[string]any) + + if !ok || redactedMap["a"] != scrubber.config.RedactionMarker || redactedMap["b"] != scrubber.config.RedactionMarker { + t.Fatalf("redacted map = %#v", redacted) + } + + scrubbed, counts := scrubber.scrubProperties(map[string]any{ + "description": strings.Repeat("x", scrubber.config.Classifier.LongTextThreshold+1), + "email_map": map[string]any{"primary": "alice@example.com"}, + "seen_at": "2026-01-01T00:00:00Z", + }) + if counts[string(actionRedact)] == 0 || counts[string(actionShiftTimestamp)] == 0 || counts[string(actionPseudonymize)] == 0 { + t.Fatalf("expected redact, shift, and pseudonymize counts, got %+v", counts) + } + + if scrubbed["description"] != scrubber.config.RedactionMarker { + t.Fatalf("long text was not redacted: %#v", scrubbed["description"]) + } + + if scrubbed["seen_at"] != "2026-01-18T00:00:00Z" { + t.Fatalf("timestamp property was not shifted: %#v", scrubbed["seen_at"]) + } +} + +func TestScrubberTimestampKeyHeuristics(t *testing.T) { + timestampKeys := []string{ + "timestamp", + "created_at", + "updated_at", + "deleted_at", + "modified_at", + "seen_at", + } + + for _, key := range timestampKeys { + if !isTimestampKey(normalizeKey(key)) { + t.Fatalf("expected %q to be a timestamp key", key) + } + } + + nonTimestampKeys := []string{ + "format", + "seat", + "heat", + "coat", + "float", + } + + for _, key := range nonTimestampKeys { + if isTimestampKey(normalizeKey(key)) { + t.Fatalf("expected %q not to be a timestamp key", key) + } + } + + scrubber, err := newScrubber(nil, "test-salt") + if err != nil { + t.Fatalf("new scrubber: %v", err) + } + + if plan := scrubber.planProperty("format", "json"); plan.Action == actionShiftTimestamp { + t.Fatalf("format was planned as timestamp: %+v", plan) + } +} + +func TestScrubberRedactsFreeTextFields(t *testing.T) { + scrubber, err := newScrubber(nil, "test-salt") + if err != nil { + t.Fatalf("new scrubber: %v", err) + } + + scrubbed, counts := scrubber.scrubProperties(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", + }) + + for key, value := range scrubbed { + if value != scrubber.config.RedactionMarker { + t.Fatalf("expected %s to be redacted, got %#v", key, value) + } + } + + if counts[string(actionRedact)] != 3 { + t.Fatalf("redact count = %d, want 3", counts[string(actionRedact)]) + } +} + +func TestScrubberPseudonymizesPathAndScriptFields(t *testing.T) { + scrubber, err := newScrubber(nil, "test-salt") + if err != nil { + t.Fatalf("new scrubber: %v", err) + } + + scrubbed, counts := scrubber.scrubProperties(map[string]any{ + "homedirectory": `\\fileserver01\share\account123`, + "profilepath": `\\profilehost01\profiles\group\account456`, + "logonscript": `startup\login.bat`, + }) + + if counts[string(actionPseudonymize)] != 3 { + t.Fatalf("pseudonymize count = %d, want 3", counts[string(actionPseudonymize)]) + } + + assertScrubbedStringDoesNotContain(t, scrubbed["homedirectory"], "fileserver01", "share", "account123") + assertScrubbedStringDoesNotContain(t, scrubbed["profilepath"], "profilehost01", "profiles", "account456") + assertScrubbedStringDoesNotContain(t, scrubbed["logonscript"], "startup", "login") + + home, ok := scrubbed["homedirectory"].(string) + + if !ok || !strings.HasPrefix(home, "value-") { + t.Fatalf("home directory was not pseudonymized as generic value: %#v", scrubbed["homedirectory"]) + } + + script, ok := scrubbed["logonscript"].(string) + + if !ok || !strings.HasPrefix(script, "value-") { + t.Fatalf("logon script was not pseudonymized as generic value: %#v", scrubbed["logonscript"]) + } +} + +func TestScrubberPseudonymizesUnknownStringsInFullScrub(t *testing.T) { + scrubber, err := newScrubber(nil, "test-salt") + if err != nil { + t.Fatalf("new scrubber: %v", err) + } + + scrubbed, counts := scrubber.scrubProperties(map[string]any{ + "business_justification": "Read only access to placeholder application for example organization", + "enabled": true, + "risk_score": 42, + }) + + if counts[string(actionPseudonymize)] != 1 || counts[string(actionPreserve)] != 2 { + t.Fatalf("unexpected action counts: %+v", counts) + } + + assertScrubbedStringDoesNotContain(t, scrubbed["business_justification"], "placeholder", "application", "organization") + + if got, ok := scrubbed["business_justification"].(string); !ok || !strings.HasPrefix(got, "value-") { + t.Fatalf("unknown string was not pseudonymized as generic value: %#v", scrubbed["business_justification"]) + } + + if scrubbed["enabled"] != true || scrubbed["risk_score"] != 42 { + t.Fatalf("safe scalar values were not preserved: %#v", scrubbed) + } +} + +func TestScrubberPseudonymizesTicketLikeValues(t *testing.T) { + scrubber, err := newScrubber(nil, "test-salt") + if err != nil { + t.Fatalf("new scrubber: %v", err) + } + + scrubbed, counts := scrubber.scrubProperties(map[string]any{ + "request_id": "WORKITEM-12345", + }) + + if counts[string(actionPseudonymize)] != 1 { + t.Fatalf("expected pseudonymize count, got %+v", counts) + } + + if got, ok := scrubbed["request_id"].(string); !ok || !strings.HasPrefix(got, "value-") { + t.Fatalf("ticket was not pseudonymized as generic value: %#v", scrubbed["request_id"]) + } +} + +func TestScrubberIdentifierRegistryConsistency(t *testing.T) { + scrubber, err := newScrubber(nil, "test-salt") + if err != nil { + t.Fatalf("new scrubber: %v", err) + } + + sourceSID := "S-1-5-21-1-2-3-500" + scrubber.observeNode(map[string]any{"objectid": sourceSID}) + + scrubbed, _ := scrubber.scrubProperties(map[string]any{ + "objectid": sourceSID, + "owner_sid": sourceSID, + }) + + if scrubbed["objectid"] != scrubbed["owner_sid"] { + t.Fatalf("registry rewrite mismatch: objectid=%#v owner_sid=%#v", scrubbed["objectid"], scrubbed["owner_sid"]) + } + + if scrubbed["objectid"] == sourceSID { + t.Fatalf("identifier was not scrubbed") + } +} + +func assertScrubbedStringDoesNotContain(t *testing.T, value any, forbidden ...string) { + t.Helper() + stringValue, ok := value.(string) + if !ok { + t.Fatalf("expected scrubbed string, got %#v", value) + } + + for _, next := range forbidden { + if strings.Contains(stringValue, next) { + t.Fatalf("scrubbed value %q still contains %q", stringValue, next) + } + } +} + +func writeTestFile(t *testing.T, path string, contents []byte) { + t.Helper() + if err := os.WriteFile(path, contents, 0o600); err != nil { + t.Fatalf("write test file: %v", err) + } +} diff --git a/retriever/types.go b/retriever/types.go new file mode 100644 index 00000000..0f1c6499 --- /dev/null +++ b/retriever/types.go @@ -0,0 +1,290 @@ +package retriever + +import ( + "fmt" + "sort" + "time" + + "github.com/specterops/dawgs/graph" +) + +const ( + manifestFormat = "retriever-opengraph-collection-v1" + legacyManifestFormat = "retrievr-opengraph-collection-v1" + idStrategy = "source_database_id_string" +) + +type Phase string + +const ( + PhaseNodes Phase = "nodes" + PhaseEdges Phase = "edges" +) + +type CompressionCodec string + +const ( + CompressionDisabled CompressionCodec = "" + CompressionGzip CompressionCodec = "gzip" + CompressionZstd CompressionCodec = "zstd" +) + +type ScrubMode string + +const ( + ScrubNone ScrubMode = "none" + ScrubFull ScrubMode = "full" +) + +type Manifest struct { + Format string `json:"format"` + GeneratedAt time.Time `json:"generated_at"` + RetrieverVersion string `json:"retriever_version,omitempty"` + Driver string `json:"driver"` + Source SourceMetadata `json:"source"` + Compression CompressionCodec `json:"compression"` + CompressionLevel int `json:"compression_level,omitempty"` + Scrub ScrubMetadata `json:"scrub"` + IDStrategy string `json:"id_strategy"` + Schema SchemaMetadata `json:"schema"` + Graphs []GraphManifest `json:"graphs"` + Metrics *MetricsManifest `json:"metrics,omitempty"` + Warnings []string `json:"warnings,omitempty"` +} + +type SourceMetadata struct { + GraphCount int `json:"graph_count"` +} + +type ScrubMetadata struct { + Mode ScrubMode `json:"mode"` + RulesVersion string `json:"rules_version,omitempty"` + SaltProvided bool `json:"salt_provided"` + NodeActionCounts map[string]int `json:"node_action_counts"` + EdgeActionCounts map[string]int `json:"edge_action_counts"` +} + +type SchemaMetadata struct { + Graphs []GraphSchemaMetadata `json:"graphs"` +} + +type GraphSchemaMetadata struct { + Name string `json:"name"` + NodeKinds []string `json:"node_kinds"` + EdgeKinds []string `json:"edge_kinds"` +} + +type GraphManifest struct { + Name string `json:"name"` + NodeCount int64 `json:"node_count"` + EdgeCount int64 `json:"edge_count"` + NodeActionCounts map[string]int `json:"node_action_counts"` + EdgeActionCounts map[string]int `json:"edge_action_counts"` + Files []FileManifest `json:"files"` +} + +type FileManifest struct { + Phase Phase `json:"phase"` + Path string `json:"path"` + Count int `json:"count"` + CompressedBytes int64 `json:"compressed_bytes"` + UncompressedBytes int64 `json:"uncompressed_bytes"` + SHA256 string `json:"sha256"` + ActionCounts map[string]int `json:"action_counts"` +} + +type NodeFragment struct { + Phase Phase `json:"phase"` + Items []FragmentNode `json:"items"` +} + +type FragmentNode struct { + ID string `json:"id"` + Kinds []string `json:"kinds"` + Properties map[string]any `json:"properties,omitempty"` +} + +type EdgeFragment struct { + Phase Phase `json:"phase"` + Items []FragmentEdge `json:"items"` +} + +type FragmentEdge struct { + StartID string `json:"start_id"` + EndID string `json:"end_id"` + Kind string `json:"kind"` + Properties map[string]any `json:"properties,omitempty"` +} + +func newManifest(driverName string, codec CompressionCodec, compressionLevel int, scrub ScrubMetadata, graphCount int) Manifest { + return Manifest{ + Format: manifestFormat, + GeneratedAt: time.Now().UTC(), + Driver: driverName, + Source: SourceMetadata{ + GraphCount: graphCount, + }, + Compression: codec, + CompressionLevel: compressionLevel, + Scrub: scrub, + IDStrategy: idStrategy, + Graphs: make([]GraphManifest, 0, graphCount), + } +} + +func (s Manifest) validate() error { + if !isSupportedManifestFormat(s.Format) { + return fmt.Errorf("unsupported manifest format %q", s.Format) + } + + if s.IDStrategy != idStrategy { + return fmt.Errorf("unsupported ID strategy %q", s.IDStrategy) + } + + if err := validateCompression(s.Compression); err != nil { + return err + } + + switch s.Scrub.Mode { + case ScrubNone, ScrubFull: + default: + return fmt.Errorf("unsupported scrub mode %q", s.Scrub.Mode) + } + + if s.Source.GraphCount != len(s.Graphs) { + return fmt.Errorf("manifest graph_count %d does not match %d graph entries", s.Source.GraphCount, len(s.Graphs)) + } + + seenGraphs := map[string]struct{}{} + for _, graphEntry := range s.Graphs { + if graphEntry.Name == "" { + return fmt.Errorf("manifest graph entry has empty name") + } + + if _, seen := seenGraphs[graphEntry.Name]; seen { + return fmt.Errorf("manifest contains duplicate graph %q", graphEntry.Name) + } + + seenGraphs[graphEntry.Name] = struct{}{} + + var nodeFiles, edgeFiles int64 + seenEdgePhase := false + + for _, fileEntry := range graphEntry.Files { + if fileEntry.Path == "" { + return fmt.Errorf("manifest graph %q contains empty file path", graphEntry.Name) + } + + switch fileEntry.Phase { + case PhaseNodes: + if seenEdgePhase { + return fmt.Errorf("manifest graph %q lists node file after edge file", graphEntry.Name) + } + + nodeFiles += int64(fileEntry.Count) + + case PhaseEdges: + seenEdgePhase = true + edgeFiles += int64(fileEntry.Count) + + default: + return fmt.Errorf("manifest graph %q contains unsupported phase %q", graphEntry.Name, fileEntry.Phase) + } + + if fileEntry.Count < 0 { + return fmt.Errorf("manifest file %q has negative count", fileEntry.Path) + } + + if fileEntry.CompressedBytes < 0 || fileEntry.UncompressedBytes < 0 { + return fmt.Errorf("manifest file %q has negative byte size", fileEntry.Path) + } + + if fileEntry.SHA256 == "" { + return fmt.Errorf("manifest file %q is missing sha256", fileEntry.Path) + } + } + + if graphEntry.NodeCount != nodeFiles { + return fmt.Errorf("manifest graph %q node_count %d does not match node file total %d", graphEntry.Name, graphEntry.NodeCount, nodeFiles) + } + + if graphEntry.EdgeCount != edgeFiles { + return fmt.Errorf("manifest graph %q edge_count %d does not match edge file total %d", graphEntry.Name, graphEntry.EdgeCount, edgeFiles) + } + } + + if s.Metrics != nil { + if err := validateMetricsManifest(*s.Metrics, s.Graphs); err != nil { + return err + } + } + + return nil +} + +func (s Manifest) Validate() error { + return s.validate() +} + +func isSupportedManifestFormat(format string) bool { + switch format { + case manifestFormat, legacyManifestFormat: + return true + default: + return false + } +} + +func IsSupportedManifestFormat(format string) bool { + return isSupportedManifestFormat(format) +} + +func graphSchemaFromMetadata(metadata GraphSchemaMetadata) graph.Graph { + return graph.Graph{ + Name: metadata.Name, + Nodes: graph.StringsToKinds(metadata.NodeKinds), + Edges: graph.StringsToKinds(metadata.EdgeKinds), + } +} + +func GraphSchemaFromMetadata(metadata GraphSchemaMetadata) graph.Graph { + return graphSchemaFromMetadata(metadata) +} + +func stringsFromKindSet(set map[string]struct{}) []string { + values := make([]string, 0, len(set)) + for value := range set { + values = append(values, value) + } + + sort.Strings(values) + + return values +} + +func addKindsToSet(set map[string]struct{}, values []string) { + for _, value := range values { + if value != "" { + set[value] = struct{}{} + } + } +} + +func addActionCounts(target map[string]int, source map[string]int) { + for key, count := range source { + target[key] += count + } +} + +func cloneActionCounts(source map[string]int) map[string]int { + if len(source) == 0 { + return map[string]int{} + } + + target := make(map[string]int, len(source)) + for key, count := range source { + target[key] = count + } + + return target +} diff --git a/retriever/verify.go b/retriever/verify.go new file mode 100644 index 00000000..0147e0df --- /dev/null +++ b/retriever/verify.go @@ -0,0 +1,357 @@ +package retriever + +import ( + "context" + "fmt" + "log/slog" + "strings" + "time" + + "github.com/specterops/dawgs/drivers/neo4j" + "github.com/specterops/dawgs/graph" +) + +type VerifyResult struct { + GraphCount int + NodeCount int64 + EdgeCount int64 +} + +type MetricsMismatchError struct { + Differences []string +} + +func (s MetricsMismatchError) Error() string { + if len(s.Differences) == 0 { + return "graph metrics mismatch" + } + return "graph metrics mismatch:\n " + strings.Join(s.Differences, "\n ") +} + +func Verify(ctx context.Context, db graph.Database, driverName string, options VerifyOptions) (VerifyResult, error) { + if err := options.validate(); err != nil { + return VerifyResult{}, err + } + + startedAt := time.Now() + slog.Info("retriever verify started", + slog.String("driver", driverName), + slog.String("input_dir", options.InputDir), + slog.Int("batch_size", options.BatchSize), + ) + options.Progress.emit(ProgressEvent{ + Operation: OperationVerify, + Message: "retriever verify started", + Driver: driverName, + InputDir: options.InputDir, + BatchSize: options.BatchSize, + }) + + readManifestStartedAt := time.Now() + slog.Info("retriever verify reading manifest", + slog.String("input_dir", options.InputDir), + ) + nextManifest, err := readManifest(options.InputDir) + if err != nil { + return VerifyResult{}, err + } + + if nextManifest.Metrics == nil { + return VerifyResult{}, MissingMetricsError{Operation: OperationVerify} + } + + if driverName == neo4j.DriverName && len(nextManifest.Graphs) > 1 { + return VerifyResult{}, IncompatibleDriverError{ + Operation: OperationVerify, + Driver: driverName, + Reason: "cannot verify a multi-graph collection against neo4j because Dawgs graph names are no-ops for that driver", + } + } + + slog.Info("retriever verify manifest ready", + slog.String("input_dir", options.InputDir), + slog.Int("graph_count", len(nextManifest.Graphs)), + slog.Duration("wall_elapsed", time.Since(readManifestStartedAt)), + ) + options.Progress.emit(ProgressEvent{ + Operation: OperationVerify, + Message: "retriever verify manifest ready", + InputDir: options.InputDir, + GraphCount: len(nextManifest.Graphs), + FileCount: manifestFileCount(nextManifest), + Elapsed: time.Since(readManifestStartedAt), + }) + + actualMetrics, result, err := collectDatabaseMetrics(ctx, db, nextManifest.Graphs, options.BatchSize, options.Progress, options.ProgressInterval) + if err != nil { + return VerifyResult{}, err + } + + differences := compareMetricsManifest(*nextManifest.Metrics, actualMetrics) + if len(differences) > 0 { + slog.Info("retriever verify failed", + slog.Int("graph_count", result.GraphCount), + slog.Int("difference_count", len(differences)), + slog.Duration("wall_elapsed", time.Since(startedAt)), + ) + return result, MetricsMismatchError{ + Differences: differences, + } + } + + slog.Info("retriever verify passed", + slog.Int("graph_count", result.GraphCount), + slog.Int64("node_count", result.NodeCount), + slog.Int64("edge_count", result.EdgeCount), + slog.Duration("wall_elapsed", time.Since(startedAt)), + ) + options.Progress.emit(ProgressEvent{ + Operation: OperationVerify, + Message: "retriever verify passed", + Driver: driverName, + InputDir: options.InputDir, + GraphCount: result.GraphCount, + NodeCount: result.NodeCount, + EdgeCount: result.EdgeCount, + Elapsed: time.Since(startedAt), + }) + + return result, nil +} + +func collectDatabaseMetrics(ctx context.Context, db graph.Database, graphEntries []GraphManifest, batchSize int, progress ProgressFunc, progressInterval int64) (MetricsManifest, VerifyResult, error) { + result := VerifyResult{ + GraphCount: len(graphEntries), + } + nextMetrics := newMetricsManifest(len(graphEntries)) + + slog.Info("retriever metrics collection started", + slog.Int("graph_count", len(graphEntries)), + slog.Int("batch_size", batchSize), + ) + + for graphIndex, graphEntry := range graphEntries { + graphStartedAt := time.Now() + + slog.Info("retriever metrics graph started", + slog.String("graph", graphEntry.Name), + slog.Int("graph_index", graphIndex+1), + slog.Int("graph_count", len(graphEntries)), + ) + progress.emit(ProgressEvent{ + Operation: OperationVerify, + Message: "retriever metrics graph started", + Graph: graphEntry.Name, + GraphIndex: graphIndex + 1, + GraphCount: len(graphEntries), + }) + + metricsEntry, err := collectDatabaseGraphMetrics(ctx, db, graphEntry.Name, batchSize, progress, progressInterval) + if err != nil { + return MetricsManifest{}, VerifyResult{}, err + } + + nextMetrics.Graphs = append(nextMetrics.Graphs, metricsEntry) + result.NodeCount += metricsEntry.NodeCount + result.EdgeCount += metricsEntry.EdgeCount + + slog.Info("retriever metrics graph completed", + slog.String("graph", graphEntry.Name), + slog.String("fingerprint", metricsEntry.Fingerprint), + slog.Int64("node_count", metricsEntry.NodeCount), + slog.Int64("edge_count", metricsEntry.EdgeCount), + slog.Duration("wall_elapsed", time.Since(graphStartedAt)), + ) + progress.emit(ProgressEvent{ + Operation: OperationVerify, + Message: "retriever metrics graph completed", + Graph: graphEntry.Name, + GraphIndex: graphIndex + 1, + GraphCount: len(graphEntries), + NodeCount: metricsEntry.NodeCount, + EdgeCount: metricsEntry.EdgeCount, + Elapsed: time.Since(graphStartedAt), + }) + } + + return nextMetrics, result, nil +} + +func collectDatabaseGraphMetrics(ctx context.Context, db graph.Database, graphName string, batchSize int, progress ProgressFunc, progressInterval int64) (GraphMetrics, error) { + targetGraph := graph.Graph{ + Name: graphName, + } + + countStartedAt := time.Now() + slog.Info("retriever metrics counting graph entities", + slog.String("graph", graphName), + ) + entitySnapshot, err := countGraphEntitySnapshot(ctx, db, targetGraph) + if err != nil { + return GraphMetrics{}, err + } + + slog.Info("retriever metrics graph counts ready", + slog.String("graph", graphName), + slog.Int64("node_count", entitySnapshot.NodeCount), + slog.Int64("edge_count", entitySnapshot.EdgeCount), + slog.Duration("wall_elapsed", time.Since(countStartedAt)), + ) + progress.emit(ProgressEvent{ + Operation: OperationVerify, + Message: "retriever metrics graph counts ready", + Graph: graphName, + NodeCount: entitySnapshot.NodeCount, + EdgeCount: entitySnapshot.EdgeCount, + Elapsed: time.Since(countStartedAt), + }) + + metricsBuilder := newMetricsBuilder(graphName, entitySnapshot.NodeCount) + if err := collectDatabaseNodeMetrics(ctx, db, targetGraph, batchSize, entitySnapshot, metricsBuilder, progress, progressInterval); err != nil { + return GraphMetrics{}, err + } + + if err := collectDatabaseRelationshipMetrics(ctx, db, targetGraph, batchSize, entitySnapshot, metricsBuilder, progress, progressInterval); err != nil { + return GraphMetrics{}, err + } + + metricsEntry := metricsBuilder.finalize() + + slog.Info("retriever metrics fingerprint computed", + slog.String("graph", graphName), + slog.String("fingerprint", metricsEntry.Fingerprint), + slog.Int64("node_count", metricsEntry.NodeCount), + slog.Int64("edge_count", metricsEntry.EdgeCount), + ) + + return metricsEntry, nil +} + +func collectDatabaseNodeMetrics(ctx context.Context, db graph.Database, targetGraph graph.Graph, batchSize int, entitySnapshot graphEntitySnapshot, metricsBuilder *metricsBuilder, progress ProgressFunc, progressInterval int64) error { + slog.Info("retriever metrics node phase started", + slog.String("graph", targetGraph.Name), + slog.Int64("node_count", entitySnapshot.NodeCount), + slog.Int("batch_size", batchSize), + ) + progress.emit(ProgressEvent{ + Operation: OperationVerify, + Message: "retriever metrics node phase started", + Graph: targetGraph.Name, + Phase: PhaseNodes, + Planned: entitySnapshot.NodeCount, + BatchSize: batchSize, + }) + + startedAt := time.Now() + processed, err := scanDatabaseNodesWithProgressInterval(ctx, db, targetGraph, entitySnapshot.NodeCount, batchSize, progressInterval, func(nodes []*graph.Node) error { + for _, node := range nodes { + if err := metricsBuilder.observeNode(node.ID.String(), node.Kinds.Strings()); err != nil { + return err + } + } + + return nil + }, func(processed int64, startedAt time.Time, nextProgressAt int64) int64 { + return logRetrieverEntityProgressInterval("retriever metrics node phase progress", targetGraph.Name, PhaseNodes, processed, entitySnapshot.NodeCount, startedAt, nextProgressAt, progress, progressInterval) + }) + if err != nil { + return err + } + + if processed != entitySnapshot.NodeCount { + return EntityCountMismatchError{ + Operation: OperationVerify, + Graph: targetGraph.Name, + Phase: PhaseNodes, + Expected: entitySnapshot.NodeCount, + Actual: processed, + Message: fmt.Sprintf("collected %d nodes for graph %q but counted %d at scan start; destination graph changed during metrics collection or the ID scan was inconsistent", processed, targetGraph.Name, entitySnapshot.NodeCount), + } + } + + slog.Info("retriever metrics node phase completed", + slog.String("graph", targetGraph.Name), + slog.Int64("processed", processed), + slog.Duration("wall_elapsed", time.Since(startedAt)), + slog.Float64("entities_per_second", perSecond(processed, time.Since(startedAt))), + ) + progress.emit(ProgressEvent{ + Operation: OperationVerify, + Message: "retriever metrics node phase completed", + Graph: targetGraph.Name, + Phase: PhaseNodes, + Processed: processed, + Planned: entitySnapshot.NodeCount, + Elapsed: time.Since(startedAt), + EntitiesPerSecond: perSecond(processed, time.Since(startedAt)), + }) + + return nil +} + +func collectDatabaseRelationshipMetrics(ctx context.Context, db graph.Database, targetGraph graph.Graph, batchSize int, entitySnapshot graphEntitySnapshot, metricsBuilder *metricsBuilder, progress ProgressFunc, progressInterval int64) error { + slog.Info("retriever metrics relationship phase started", + slog.String("graph", targetGraph.Name), + slog.Int64("edge_count", entitySnapshot.EdgeCount), + slog.Int("batch_size", batchSize), + ) + progress.emit(ProgressEvent{ + Operation: OperationVerify, + Message: "retriever metrics relationship phase started", + Graph: targetGraph.Name, + Phase: PhaseEdges, + Planned: entitySnapshot.EdgeCount, + BatchSize: batchSize, + }) + + startedAt := time.Now() + processed, err := scanDatabaseRelationshipsWithProgressInterval(ctx, db, targetGraph, entitySnapshot.EdgeCount, batchSize, progressInterval, func(relationships []*graph.Relationship) error { + for _, relationship := range relationships { + kind := "" + if relationship.Kind != nil { + kind = relationship.Kind.String() + } + + if err := metricsBuilder.observeRelationship(relationship.StartID.String(), relationship.EndID.String(), kind); err != nil { + return err + } + } + + return nil + }, func(processed int64, startedAt time.Time, nextProgressAt int64) int64 { + return logRetrieverEntityProgressInterval("retriever metrics relationship phase progress", targetGraph.Name, PhaseEdges, processed, entitySnapshot.EdgeCount, startedAt, nextProgressAt, progress, progressInterval) + }) + if err != nil { + return err + } + + if processed != entitySnapshot.EdgeCount { + return EntityCountMismatchError{ + Operation: OperationVerify, + Graph: targetGraph.Name, + Phase: PhaseEdges, + Expected: entitySnapshot.EdgeCount, + Actual: processed, + Message: fmt.Sprintf("collected %d relationships for graph %q but counted %d at scan start; destination graph changed during metrics collection or the ID scan was inconsistent", processed, targetGraph.Name, entitySnapshot.EdgeCount), + } + } + + slog.Info("retriever metrics relationship phase completed", + slog.String("graph", targetGraph.Name), + slog.Int64("processed", processed), + slog.Duration("wall_elapsed", time.Since(startedAt)), + slog.Float64("entities_per_second", perSecond(processed, time.Since(startedAt))), + ) + progress.emit(ProgressEvent{ + Operation: OperationVerify, + Message: "retriever metrics relationship phase completed", + Graph: targetGraph.Name, + Phase: PhaseEdges, + Processed: processed, + Planned: entitySnapshot.EdgeCount, + Elapsed: time.Since(startedAt), + EntitiesPerSecond: perSecond(processed, time.Since(startedAt)), + }) + + return nil +}