From 55c30a736f47b6931de68331c96da743b4d62f7a Mon Sep 17 00:00:00 2001 From: Rusty Conover Date: Thu, 10 Sep 2026 21:16:56 -0400 Subject: [PATCH 1/6] Expose split token TTL metadata --- include/vgi/types.h | 4 ++++ src/catalog.cpp | 3 +++ 2 files changed, 7 insertions(+) diff --git a/include/vgi/types.h b/include/vgi/types.h index 34a0fc7..d14214a 100644 --- a/include/vgi/types.h +++ b/include/vgi/types.h @@ -288,6 +288,10 @@ struct FunctionMetadata { // and still restore the original order, by sorting on the tag rather than // on arrival. A function that declares it and omits the tag is rejected. bool supports_batch_index = false; + // How long tokens minted for this split-capable function remain usable. + // Absent lets the engine use its normal planning horizon; a declared + // value below that horizon is refused before any work is scheduled. + std::optional split_token_ttl_seconds; // What order the engine may assume of this function's rows. Empty leaves // the engine's own default in place; the values are in `src/enums.h`. std::string order_preservation; diff --git a/src/catalog.cpp b/src/catalog.cpp index 9522b86..e2ae45c 100644 --- a/src/catalog.cpp +++ b/src/catalog.cpp @@ -1192,6 +1192,9 @@ std::string Dispatcher::encode_table_function_info(const TableFunction& fn, // that overrides `plan()` has said so in the only place that cannot // drift from the implementation. .set_bool("supports_splits", fn.supports_splits()); + if (metadata.split_token_ttl_seconds) { + builder.set_int64("split_token_ttl_seconds", *metadata.split_token_ttl_seconds); + } // Only when the function says so: the field is nullable, and writing a // value unconditionally would replace the engine's default for every // table function that never thought about ordering. From b18ffb885558797e6e98687ab2276aebe8d2bd95 Mon Sep 17 00:00:00 2001 From: Rusty Conover Date: Thu, 10 Sep 2026 21:16:56 -0400 Subject: [PATCH 2/6] Add non-filter integration fixtures --- example-worker/CMakeLists.txt | 1 + example-worker/buffering/buffer_input.cpp | 6 +- example-worker/catalog_def.cpp | 19 ++ example-worker/main.cpp | 1 + example-worker/registry.h | 1 + example-worker/table/same_name.cpp | 70 +++++ example-worker/table/splits.cpp | 324 +++++++++++++++++++++- example-worker/table_in_out/substream.cpp | 74 +++++ scripts/run_tests.sh | 20 +- src/function_dispatch.cpp | 17 +- src/split_token.cpp | 22 +- src/split_token.h | 16 +- src/worker.cpp | 51 +++- 13 files changed, 593 insertions(+), 29 deletions(-) create mode 100644 example-worker/table/same_name.cpp diff --git a/example-worker/CMakeLists.txt b/example-worker/CMakeLists.txt index bdda7b1..ff29e8b 100644 --- a/example-worker/CMakeLists.txt +++ b/example-worker/CMakeLists.txt @@ -16,6 +16,7 @@ add_executable(vgi-example-worker scalar/secrets.cpp scalar/cached.cpp table/sequence.cpp + table/same_name.cpp table/cache.cpp table/more.cpp table/generators.cpp diff --git a/example-worker/buffering/buffer_input.cpp b/example-worker/buffering/buffer_input.cpp index 0e2a4c3..14919f8 100644 --- a/example-worker/buffering/buffer_input.cpp +++ b/example-worker/buffering/buffer_input.cpp @@ -3,9 +3,9 @@ #include #include -#include #include #include +#include #include #include @@ -167,8 +167,8 @@ class BufferInput : public vgi::TableBufferingFunction { (void)id; ordered.push_back(decode_indexed(blob)); } - std::sort(ordered.begin(), ordered.end(), - [](const auto& a, const auto& b) { return a.first < b.first; }); + std::stable_sort(ordered.begin(), ordered.end(), + [](const auto& a, const auto& b) { return a.first < b.first; }); for (const auto& [index, bytes] : ordered) { (void)index; params.storage->append(params.execution_id, kNamespace, "", bytes); diff --git a/example-worker/catalog_def.cpp b/example-worker/catalog_def.cpp index 789d8ca..fd417ed 100644 --- a/example-worker/catalog_def.cpp +++ b/example-worker/catalog_def.cpp @@ -70,6 +70,15 @@ vgi::CatalogBranch sequence_branch(int64_t count, return branch; } +vgi::CatalogBranch split_sequence_branch(int64_t count, int64_t splits) { + vgi::CatalogBranch branch; + branch.function_name = "split_sequence"; + branch.scan_arguments = + vgi::serialize_scan_arguments({}, {{"n", int64_arg(count)}, {"splits", int64_arg(splits)}}); + branch.schema_path = vgi::SchemaPath{"data"}; + return branch; +} + vgi::CatalogTable multi_branch(std::string name, std::vector branches, std::string comment = {}) { vgi::CatalogTable table; @@ -391,6 +400,10 @@ void declare_catalog(vgi::Worker& worker) { data.tables.push_back(multi_branch("multi_branch_empty", {}, "Multi-branch: empty branches list — used by " "multi_branch_empty_branches.test")); + data.tables.push_back( + multi_branch("multi_branch_split", {split_sequence_branch(30, 6), sequence_branch(20)}, + "Multi-branch: split_sequence(30, splits=6) + sequence(20) — used by " + "splits/multi_branch.test")); // Heterogeneous branches: one arm is this worker, the others are DuckDB's // own readers over files the test writes first. What they probe is that a @@ -526,6 +539,12 @@ void declare_catalog(vgi::Worker& worker) { // and the suite names `main` for these two. auto& main = worker.catalog().schema("main"); main.comment = "Example functions for testing VGI"; + main.tables.push_back(backed_by("test_same_name_table", "test_same_name_table_scan", + columns({{"tag", arrow::utf8()}}), + "Schema-disambiguation probe; the main-schema table")); + data.tables.push_back(backed_by("test_same_name_table", "test_same_name_table_scan", + columns({{"tag", arrow::utf8()}}), + "Schema-disambiguation probe; the data-schema table")); // Macros never reach the worker at run time — the engine substitutes the // text — so declaring them is the whole implementation. main.macros.push_back({"vgi_multiply", diff --git a/example-worker/main.cpp b/example-worker/main.cpp index ce4f43b..56b8080 100644 --- a/example-worker/main.cpp +++ b/example-worker/main.cpp @@ -62,6 +62,7 @@ int main(int argc, char** argv) { example::register_secret_fixtures(worker); example::register_series(worker); example::register_splits(worker); + if (composite) example::register_same_name_tables(worker); example::register_filter_fixtures(worker); example::register_logging_fixtures(worker); example::register_global_probes(worker); diff --git a/example-worker/registry.h b/example-worker/registry.h index d0705f1..2e8ca48 100644 --- a/example-worker/registry.h +++ b/example-worker/registry.h @@ -36,6 +36,7 @@ void register_settings_tables(vgi::Worker& worker); void register_secret_fixtures(vgi::Worker& worker); void register_series(vgi::Worker& worker); void register_splits(vgi::Worker& worker); +void register_same_name_tables(vgi::Worker& worker); void register_filter_fixtures(vgi::Worker& worker); void register_logging_fixtures(vgi::Worker& worker); void register_global_probes(vgi::Worker& worker); diff --git a/example-worker/table/same_name.cpp b/example-worker/table/same_name.cpp new file mode 100644 index 0000000..135baad --- /dev/null +++ b/example-worker/table/same_name.cpp @@ -0,0 +1,70 @@ +// © Copyright 2025, 2026 Query Farm LLC - https://query.farm + +// Two identically-named scans in different schemas. Each backs a declarative +// table in its own schema, so dispatch has to retain the full schema path. + +#include +#include + +#include +#include +#include + +#include + +namespace example { +namespace { + +class SameNameTableScan : public vgi::TableFunction { +public: + explicit SameNameTableScan(std::string schema) : schema_(std::move(schema)) {} + + std::string name() const override { return "test_same_name_table_scan"; } + + vgi::FunctionMetadata metadata() const override { + vgi::FunctionMetadata md; + md.description = "Schema-disambiguation probe; the " + schema_ + "-schema producer"; + md.categories = {"generator", "testing"}; + return md; + } + + std::vector argument_specs() const override { return {}; } + + std::shared_ptr bind(const vgi::BindParams&) const override { + return arrow::schema({arrow::field("tag", arrow::utf8(), /*nullable=*/true)}); + } + + std::unique_ptr init(const vgi::ProcessParams& params) const override { + arrow::StringBuilder builder; + (void)builder.Append(schema_); + std::shared_ptr tag; + (void)builder.Finish(&tag); + return std::make_unique(arrow::RecordBatch::Make(params.output_schema, 1, {tag})); + } + +private: + class Producer : public vgi::TableProducer { + public: + explicit Producer(std::shared_ptr batch) : batch_(std::move(batch)) {} + + std::shared_ptr next_batch() override { + auto result = batch_; + batch_ = nullptr; + return result; + } + + private: + std::shared_ptr batch_; + }; + + std::string schema_; +}; + +} // namespace + +void register_same_name_tables(vgi::Worker& worker) { + worker.register_table_in("example", "main", std::make_shared("main")); + worker.register_table_in("example", "data", std::make_shared("data")); +} + +} // namespace example diff --git a/example-worker/table/splits.cpp b/example-worker/table/splits.cpp index db05478..abf3711 100644 --- a/example-worker/table/splits.cpp +++ b/example-worker/table/splits.cpp @@ -14,9 +14,11 @@ // planning off, which is what `splits/rollback.test` pins. #include +#include #include #include #include +#include #include #include @@ -25,6 +27,7 @@ #include #include +#include #include #include "scalar/util.h" @@ -60,6 +63,19 @@ std::optional> decode_range(const std::string& paylo return std::pair{read(0), read(8)}; } +std::string encode_indexed_range(int64_t ordinal, int64_t begin, int64_t end) { + return encode_range(ordinal, begin) + encode_range(end, 0).substr(0, 8); +} + +std::optional> decode_indexed_range( + const std::string& payload) { + if (payload.size() != 24) return std::nullopt; + auto ordinal_begin = decode_range(payload.substr(0, 16)); + auto end_unused = decode_range(payload.substr(8, 16)); + if (!ordinal_begin || !end_unused) return std::nullopt; + return std::tuple{ordinal_begin->first, ordinal_begin->second, end_unused->second}; +} + std::shared_ptr n_schema() { return arrow::schema({arrow::field("n", arrow::int64(), /*nullable=*/false)}); } @@ -89,8 +105,9 @@ std::string render_filter_bounds(const vgi::PushdownFilters& filters) { class RangeProducer : public vgi::TableProducer { public: RangeProducer(std::shared_ptr schema, - std::vector> ranges) - : schema_(std::move(schema)), ranges_(std::move(ranges)) {} + std::vector> ranges, + std::map metadata = {}) + : schema_(std::move(schema)), ranges_(std::move(ranges)), metadata_(std::move(metadata)) {} std::shared_ptr next_batch() override { // An empty range is not the end of the scan: `split_empty_ranges` @@ -114,10 +131,13 @@ class RangeProducer : public vgi::TableProducer { return batch; } + std::map last_metadata() const override { return metadata_; } + private: static constexpr int64_t kBatchRows = 2048; std::shared_ptr schema_; std::vector> ranges_; + std::map metadata_; size_t at_ = 0; int64_t cursor_ = 0; }; @@ -128,8 +148,16 @@ class SplitFunction : public vgi::TableFunction { public: enum class Shape { Even, EmptyRanges, Zero, Skewed, Many }; - SplitFunction(std::string name, Shape shape, std::string description) - : name_(std::move(name)), shape_(shape), description_(std::move(description)) {} + SplitFunction(std::string name, Shape shape, std::string description, + std::optional catalog_version = std::nullopt, + std::optional split_token_ttl_seconds = std::nullopt, + bool cacheable = false) + : name_(std::move(name)), + shape_(shape), + description_(std::move(description)), + catalog_version_(catalog_version), + split_token_ttl_seconds_(split_token_ttl_seconds), + cacheable_(cacheable) {} std::string name() const override { return name_; } @@ -137,6 +165,7 @@ class SplitFunction : public vgi::TableFunction { vgi::FunctionMetadata md; md.description = description_; md.categories = {"generator"}; + md.split_token_ttl_seconds = split_token_ttl_seconds_; return md; } @@ -159,6 +188,7 @@ class SplitFunction : public vgi::TableFunction { vgi::PlanResult result; result.estimated_total_rows = shape_ == Shape::Zero ? 0 : rows; result.estimated_total_splits = want; + result.catalog_version = catalog_version_; for (const auto& range : divide(rows, want)) { vgi::ScanSplit split; split.payload = encode_range(range.first, range.second); @@ -190,8 +220,15 @@ class SplitFunction : public vgi::TableFunction { } ranges.push_back(*range); } + std::map metadata; + if (cacheable_) { + vgi::CacheControl control; + control.ttl_seconds = 300; + metadata = control.to_metadata(); + } return std::make_unique( - params.output_schema ? params.output_schema : n_schema(), std::move(ranges)); + params.output_schema ? params.output_schema : n_schema(), std::move(ranges), + std::move(metadata)); } private: @@ -259,6 +296,270 @@ class SplitFunction : public vgi::TableFunction { std::string name_; Shape shape_; std::string description_; + std::optional catalog_version_; + std::optional split_token_ttl_seconds_; + bool cacheable_; +}; + +// `split_paginated(n, splits)` — four disjoint splits per planning page. +class SplitPaginated : public vgi::TableFunction { +public: + std::string name() const override { return "split_paginated"; } + + vgi::FunctionMetadata metadata() const override { + vgi::FunctionMetadata md; + md.description = "Split scan whose plan is enumerated across cursor pages"; + md.categories = {"generator"}; + return md; + } + + std::vector argument_specs() const override { + return {vgi::ArgSpec::named("n", "int64", "How many rows to generate"), + vgi::ArgSpec::named("splits", "int64", "How many splits")}; + } + + std::shared_ptr bind(const vgi::BindParams&) const override { + return n_schema(); + } + + bool supports_splits() const override { return true; } + + vgi::PlanResult plan(const vgi::BindParams& params, + const vgi::PlanParams& request) const override { + const int64_t rows = std::max(0, params.arguments.named_int64("n").value_or(0)); + const int64_t count = + std::max(1, params.arguments.named_int64("splits").value_or(1)); + int64_t page = 0; + if (request.cursor && request.cursor->size() == sizeof(int64_t)) { + auto decoded = decode_range(*request.cursor + std::string(sizeof(int64_t), '\0')); + if (decoded) page = decoded->first; + } + + constexpr int64_t kPerPage = 4; + const int64_t first = page * kPerPage; + const int64_t last = std::min(count, first + kPerPage); + vgi::PlanResult result; + result.estimated_total_splits = count; + result.estimated_total_rows = rows; + for (int64_t i = first; i < last; ++i) { + vgi::ScanSplit split; + split.payload = encode_range(rows * i / count, rows * (i + 1) / count); + split.estimated_rows = rows * (i + 1) / count - rows * i / count; + split.rows_exact = true; + result.splits.push_back(std::move(split)); + } + if (last < count) result.next_cursor = encode_range(page + 1, 0).substr(0, 8); + return result; + } + + std::unique_ptr init(const vgi::ProcessParams& params) const override { + if (!params.split_payloads) { + throw std::runtime_error("table function 'split_paginated' is split-only"); + } + std::vector> ranges; + for (const auto& payload : *params.split_payloads) { + auto range = decode_range(payload); + if (!range) throw std::runtime_error("split_paginated: unrecognized split payload"); + ranges.push_back(*range); + } + return std::make_unique( + params.output_schema ? params.output_schema : n_schema(), std::move(ranges)); + } +}; + +// A batch index derived from the split ordinal remains monotonic when a reader +// claims several splits in ascending order. +class SplitBatchIndex : public vgi::TableFunction { +public: + std::string name() const override { return "split_batch_index"; } + + vgi::FunctionMetadata metadata() const override { + vgi::FunctionMetadata md; + md.description = "Split scan with batch indices monotonic across split boundaries"; + md.categories = {"generator", "ordering"}; + md.supports_batch_index = true; + return md; + } + + std::vector argument_specs() const override { + return {vgi::ArgSpec::named("n", "int64", "How many rows to generate"), + vgi::ArgSpec::named("splits", "int64", "How many splits")}; + } + + std::shared_ptr bind(const vgi::BindParams&) const override { + return n_schema(); + } + + bool supports_splits() const override { return true; } + + vgi::PlanResult plan(const vgi::BindParams& params, const vgi::PlanParams&) const override { + const int64_t rows = std::max(0, params.arguments.named_int64("n").value_or(0)); + const int64_t count = + std::max(1, params.arguments.named_int64("splits").value_or(1)); + vgi::PlanResult result; + result.estimated_total_splits = count; + result.estimated_total_rows = rows; + for (int64_t i = 0; i < count; ++i) { + vgi::ScanSplit split; + const int64_t begin = rows * i / count; + const int64_t end = rows * (i + 1) / count; + split.payload = encode_indexed_range(i, begin, end); + split.estimated_rows = end - begin; + split.rows_exact = true; + result.splits.push_back(std::move(split)); + } + return result; + } + + std::unique_ptr init(const vgi::ProcessParams& params) const override { + if (!params.split_payloads) { + throw std::runtime_error("table function 'split_batch_index' is split-only"); + } + std::vector> ranges; + for (const auto& payload : *params.split_payloads) { + auto range = decode_indexed_range(payload); + if (!range) throw std::runtime_error("split_batch_index: unrecognized split payload"); + ranges.push_back(*range); + } + return std::make_unique(params.output_schema ? params.output_schema : n_schema(), + std::move(ranges)); + } + +private: + class Producer : public vgi::TableProducer { + public: + Producer(std::shared_ptr schema, + std::vector> ranges) + : schema_(std::move(schema)), ranges_(std::move(ranges)) {} + + std::shared_ptr next_batch() override { + while (at_ < ranges_.size() && cursor_ >= std::get<2>(ranges_[at_])) { + ++at_; + local_batch_ = 0; + if (at_ < ranges_.size()) cursor_ = std::get<1>(ranges_[at_]); + } + if (at_ >= ranges_.size()) return nullptr; + if (cursor_ < std::get<1>(ranges_[at_])) cursor_ = std::get<1>(ranges_[at_]); + + const int64_t end = std::min(cursor_ + kBatchRows, std::get<2>(ranges_[at_])); + arrow::Int64Builder builder; + for (int64_t value = cursor_; value < end; ++value) (void)builder.Append(value); + std::shared_ptr values; + (void)builder.Finish(&values); + metadata_["vgi_batch_index"] = + std::to_string(std::get<0>(ranges_[at_]) * kStride + local_batch_++); + cursor_ = end; + return arrow::RecordBatch::Make(schema_, values->length(), {values}); + } + + std::map last_metadata() const override { return metadata_; } + + private: + static constexpr int64_t kBatchRows = 8; + static constexpr int64_t kStride = 1000; + std::shared_ptr schema_; + std::vector> ranges_; + size_t at_ = 0; + int64_t cursor_ = 0; + int64_t local_batch_ = 0; + std::map metadata_; + }; +}; + +// One split per country, with each emitted batch retaining that partition's +// single value and distinct sales range. +class SplitPartitioned : public vgi::TableFunction { +public: + std::string name() const override { return "split_partitioned"; } + + vgi::FunctionMetadata metadata() const override { + vgi::FunctionMetadata md; + md.description = "One split per partition value"; + md.categories = {"generator", "partitioning"}; + md.partition_kind = vgi::partition_kinds::kSingleValuePartitions; + return md; + } + + std::vector argument_specs() const override { + return {vgi::ArgSpec::named("rows_per_country", "int64", "Rows per country")}; + } + + std::shared_ptr bind(const vgi::BindParams&) const override { + return arrow::schema({vgi::partition_field("country", arrow::utf8()), + arrow::field("sales", arrow::int64(), /*nullable=*/true)}); + } + + bool supports_splits() const override { return true; } + + vgi::PlanResult plan(const vgi::BindParams& params, const vgi::PlanParams&) const override { + const int64_t rows = + std::max(0, params.arguments.named_int64("rows_per_country").value_or(0)); + vgi::PlanResult result; + result.estimated_total_splits = static_cast(kCountries.size()); + result.estimated_total_rows = rows * static_cast(kCountries.size()); + for (int64_t i = 0; i < static_cast(kCountries.size()); ++i) { + vgi::ScanSplit split; + split.payload = encode_range(i, rows); + split.estimated_rows = rows; + split.rows_exact = true; + result.splits.push_back(std::move(split)); + } + return result; + } + + std::unique_ptr init(const vgi::ProcessParams& params) const override { + if (!params.split_payloads) { + throw std::runtime_error("table function 'split_partitioned' is split-only"); + } + std::vector> partitions; + for (const auto& payload : *params.split_payloads) { + auto decoded = decode_range(payload); + if (!decoded || decoded->first < 0 || + decoded->first >= static_cast(kCountries.size())) { + throw std::runtime_error("split_partitioned: unrecognized split payload"); + } + partitions.push_back(*decoded); + } + return std::make_unique(params.output_schema ? params.output_schema : bind({}), + std::move(partitions)); + } + +private: + inline static const std::vector kCountries{"US", "DE", "JP", "BR"}; + + class Producer : public vgi::TableProducer { + public: + Producer(std::shared_ptr schema, + std::vector> partitions) + : schema_(std::move(schema)), partitions_(std::move(partitions)) {} + + std::shared_ptr next_batch() override { + while (at_ < partitions_.size() && partitions_[at_].second <= 0) ++at_; + if (at_ >= partitions_.size()) return nullptr; + const auto [country_index, rows] = partitions_[at_++]; + arrow::StringBuilder countries; + arrow::Int64Builder sales; + for (int64_t row = 1; row <= rows; ++row) { + (void)countries.Append(kCountries[static_cast(country_index)]); + (void)sales.Append(country_index * 100 + row); + } + std::shared_ptr country_array; + std::shared_ptr sales_array; + (void)countries.Finish(&country_array); + (void)sales.Finish(&sales_array); + auto batch = arrow::RecordBatch::Make(schema_, rows, {country_array, sales_array}); + metadata_ = vgi::partition_metadata(schema_, batch); + return batch; + } + + std::map last_metadata() const override { return metadata_; } + + private: + std::shared_ptr schema_; + std::vector> partitions_; + size_t at_ = 0; + std::map metadata_; + }; }; // `split_fail_at(n, splits, fail_at, fail_in_init)` — a scan that dies where it @@ -678,6 +979,19 @@ void register_splits(vgi::Worker& worker) { "split_skewed", Shape::Skewed, "Integers 0..n-1, divided very unevenly")); worker.register_table(std::make_shared( "split_many", Shape::Many, "Integers 0..n-1, divided into many more splits than threads")); + worker.register_table(std::make_shared( + "split_stale_plan", Shape::Even, "A plan pinned to a stale catalog version", + /*catalog_version=*/987654321)); + worker.register_table(std::make_shared( + "split_short_ttl", Shape::Even, "A split plan with an unusably short token lifetime", + /*catalog_version=*/std::nullopt, /*split_token_ttl_seconds=*/1)); + worker.register_table(std::make_shared( + "split_cacheable", Shape::Even, "A split-capable result-cache candidate", + /*catalog_version=*/std::nullopt, /*split_token_ttl_seconds=*/std::nullopt, + /*cacheable=*/true)); + worker.register_table(std::make_shared()); + worker.register_table(std::make_shared()); + worker.register_table(std::make_shared()); worker.register_table(std::make_shared()); worker.register_table(std::make_shared()); worker.register_table(std::make_shared()); diff --git a/example-worker/table_in_out/substream.cpp b/example-worker/table_in_out/substream.cpp index 0ad1003..7c08ce4 100644 --- a/example-worker/table_in_out/substream.cpp +++ b/example-worker/table_in_out/substream.cpp @@ -54,6 +54,10 @@ int64_t decode_i64(const std::string& bytes) { return static_cast(value); } +std::string encode_i64_pair(int64_t first, int64_t second) { + return encode_i64(first) + encode_i64(second); +} + class SubstreamPartialSum : public vgi::TableInOutFunction { public: std::string name() const override { return "substream_partial_sum"; } @@ -117,10 +121,80 @@ class SubstreamPartialSum : public vgi::TableInOutFunction { } }; +// Accumulates a substream and emits one one-row batch per input row at +// finalize. More than one batch is intentional: it exercises continuation of +// a finalize response over transports that emit one batch per turn. +class MultiBatchFinish : public vgi::TableInOutFunction { +public: + std::string name() const override { return "multi_batch_finish"; } + + vgi::FunctionMetadata metadata() const override { + vgi::FunctionMetadata md; + md.description = "Streaming finalize that emits one batch per input row"; + md.categories = {"testing", "aggregation"}; + return md; + } + + std::vector argument_specs() const override { + return {vgi::ArgSpec::table("data", 0, "Input relation")}; + } + + std::shared_ptr bind(const vgi::BindParams& params) const override { + if (!params.input_schema || params.input_schema->num_fields() == 0) { + throw std::invalid_argument("multi_batch_finish requires an input column"); + } + return arrow::schema({arrow::field(params.input_schema->field(0)->name(), arrow::int64(), + /*nullable=*/true)}); + } + + bool has_finish() const override { return true; } + + std::vector process( + const vgi::ProcessParams& params, + const std::shared_ptr& batch) const override { + if (!batch || batch->num_columns() == 0) return {}; + auto values = + std::static_pointer_cast(cast_to(batch->column(0), arrow::int64())); + int64_t total = 0; + for (int64_t row = 0; row < values->length(); ++row) { + if (!values->IsNull(row)) total += values->Value(row); + } + params.storage->append(state_scope(params), kMultiBatchNamespace, "", + encode_i64_pair(total, batch->num_rows())); + return {}; + } + + std::vector finish(const vgi::ProcessParams& params) const override { + int64_t total = 0; + int64_t rows = 0; + for (const auto& [id, blob] : + params.storage->scan(state_scope(params), kMultiBatchNamespace, "", 0, SIZE_MAX)) { + (void)id; + total += decode_i64(blob.substr(0, sizeof(int64_t))); + rows += decode_i64(blob.substr(sizeof(int64_t), sizeof(int64_t))); + } + + std::vector result; + result.reserve(static_cast(rows)); + for (int64_t row = 0; row < rows; ++row) { + arrow::Int64Builder builder; + (void)builder.Append(row == 0 ? total : 0); + std::shared_ptr value; + (void)builder.Finish(&value); + result.emplace_back(arrow::RecordBatch::Make(params.output_schema, 1, {value})); + } + return result; + } + +private: + static constexpr const char* kMultiBatchNamespace = "substream.multi_batch"; +}; + } // namespace void register_substream_finalize(vgi::Worker& worker) { worker.register_table_in_out(std::make_shared()); + worker.register_table_in_out(std::make_shared()); } } // namespace example diff --git a/scripts/run_tests.sh b/scripts/run_tests.sh index 6ee7e9f..23116a6 100755 --- a/scripts/run_tests.sh +++ b/scripts/run_tests.sh @@ -29,6 +29,8 @@ mkdir -p "$BRANCH_DIR" BUILD=1 if [[ "${1:-}" == "--no-build" ]]; then BUILD=0; shift; fi +FULL_RUN=0 +if [[ $# -eq 0 ]]; then FULL_RUN=1; fi if [[ ! -x "$UNITTEST" ]]; then echo "[harness] $UNITTEST missing — build the extension first:" @@ -143,6 +145,15 @@ if [[ "${VGI_HTTP:-0}" == "1" ]]; then ) fi +# The bearer fixture needs a protected HTTP worker while the rest of the suite +# needs an anonymous worker. Run it separately during a full suite instead of +# exporting its token into the launcher run, where bearer_token is invalid. +H_BEARER="" +if [[ $FULL_RUN == 1 ]]; then + H_BEARER=$(start_http_worker bearer example \ + VGI_BEARER_TOKENS=test-secret-token=test-principal) || exit 1 +fi + ARGS=() if [[ $# -ge 1 ]]; then case "$1" in @@ -162,10 +173,17 @@ echo "[harness] running: ${ARGS[*]}" VGI_VERSIONED_TABLES_WORKER="$W_VERSIONED_TABLES" \ VGI_ATTACH_OPTIONS_WORKER="$W_ATTACH_OPTIONS" \ VGI_BAD_PROTOCOL_WORKER="$W_BAD_PROTOCOL" \ - VGI_TEST_BEARER_TOKEN="test-secret-token" \ "$UNITTEST" "${ARGS[@]}" ) > "$CACHE/run.log" 2>&1 RC=$? +if [[ $FULL_RUN == 1 ]]; then + ( cd "$VGI_EXT" && env \ + VGI_TEST_WORKER="$H_BEARER" \ + VGI_TEST_BEARER_TOKEN="test-secret-token" \ + "$UNITTEST" "test/sql/integration/bearer_auth/*" ) >> "$CACHE/run.log" 2>&1 + RC=$(( RC | $? )) +fi + grep -oE 'test/sql/integration/[A-Za-z0-9_/]+\.test(_slow)?' "$CACHE/run.log" \ | sort -u > "$CACHE/allmentioned" 2>/dev/null awk '/unexpectedly|FAILED:|Mismatch on/{print}' "$CACHE/run.log" \ diff --git a/src/function_dispatch.cpp b/src/function_dispatch.cpp index 3cbcc26..5fb8503 100644 --- a/src/function_dispatch.cpp +++ b/src/function_dispatch.cpp @@ -1314,13 +1314,18 @@ vgi_rpc::Stream Dispatcher::init(const vgi_rpc::Request& request) { payloads.reserve(tokens.size()); for (const auto& token : tokens) { auto opened = split_token::open(token, fingerprint, anchor); - if (!opened) { - throw std::runtime_error( - "init: split token for '" + function_name + - "' is not redeemable here — it was minted for a different bind, or " - "against a snapshot this worker no longer serves"); + if (!opened.payload) { + if (opened.error == split_token::OpenError::SnapshotExpired) { + throw std::runtime_error("SPLIT_SNAPSHOT_EXPIRED: split token for '" + + function_name + + "' names a snapshot this worker no longer serves; " + "re-run the query to plan against the current " + "snapshot"); + } + throw std::runtime_error("SPLIT_TOKEN_INVALID: split token for '" + function_name + + "' is malformed or bound elsewhere"); } - payloads.push_back(std::move(*opened)); + payloads.push_back(std::move(*opened.payload)); } params.split_payloads = std::move(payloads); } diff --git a/src/split_token.cpp b/src/split_token.cpp index 75aa3d1..aabbce7 100644 --- a/src/split_token.cpp +++ b/src/split_token.cpp @@ -70,32 +70,36 @@ std::string build(const std::string& payload, const std::string& fingerprint, return out; } -std::optional open(const std::string& token, const std::string& expected_fingerprint, - const std::string& current_anchor) { - if (token.size() < kHeaderLen) return std::nullopt; - if (static_cast(token[0]) != kFormatVersion) return std::nullopt; +OpenResult open(const std::string& token, const std::string& expected_fingerprint, + const std::string& current_anchor) { + const OpenResult invalid{std::nullopt, OpenError::Invalid}; + if (token.size() < kHeaderLen) return invalid; + if (static_cast(token[0]) != kFormatVersion) return invalid; const auto flags = static_cast(token[1]); // Every bit is reserved here, `payload_sealed` included: this SDK holds no // key, so a token claiming to be sealed is one we cannot open, and a token // setting a reserved bit is from a future this build does not speak. - if (flags != 0) return std::nullopt; + if (flags != 0) return invalid; const auto anchor_len = static_cast(static_cast(token[2])) | (static_cast(static_cast(token[3])) << 8); const size_t end_of_anchor = kHeaderLen + anchor_len; - if (token.size() < end_of_anchor) return std::nullopt; + if (token.size() < end_of_anchor) return invalid; const auto fingerprint = token.substr(4, kFingerprintLen); if (!vgi_rpc::crypto::constant_time_equal(fingerprint, expected_fingerprint)) { - return std::nullopt; + return invalid; } // Checked after the bind check and kept distinct in the caller's error // text: "this snapshot moved" is a different situation for a client from // "this token is not yours", and only one of them means re-plan. - if (token.compare(kHeaderLen, anchor_len, current_anchor) != 0) return std::nullopt; + if (anchor_len != current_anchor.size() || + token.compare(kHeaderLen, anchor_len, current_anchor) != 0) { + return {std::nullopt, OpenError::SnapshotExpired}; + } - return token.substr(end_of_anchor); + return {token.substr(end_of_anchor), OpenError::None}; } } // namespace vgi::split_token diff --git a/src/split_token.h b/src/split_token.h index 4696693..40f3085 100644 --- a/src/split_token.h +++ b/src/split_token.h @@ -47,10 +47,18 @@ std::string bind_fingerprint(const SchemaPath& schema_path, const std::string& f std::string build(const std::string& payload, const std::string& fingerprint, const std::string& anchor); -// Verify a token and return the payload, or nothing when it is malformed, was -// minted for a different bind, or names a snapshot that has moved on. -std::optional open(const std::string& token, const std::string& expected_fingerprint, - const std::string& current_anchor); +enum class OpenError { None, Invalid, SnapshotExpired }; + +struct OpenResult { + std::optional payload; + OpenError error = OpenError::None; +}; + +// Verify a token while preserving the one actionable distinction: an expired +// snapshot can be retried after replanning, while a malformed or wrongly-bound +// token cannot. +OpenResult open(const std::string& token, const std::string& expected_fingerprint, + const std::string& current_anchor); // The consistency anchor for a catalog version: int64, little-endian, and an // absent version is zero — the same spelling the reference uses. diff --git a/src/worker.cpp b/src/worker.cpp index a464e75..23740de 100644 --- a/src/worker.cpp +++ b/src/worker.cpp @@ -3,6 +3,7 @@ #include #include +#include #include #include #include @@ -46,6 +47,50 @@ std::pair parse_tcp_bind(const std::string& value, const char* bool is_loopback_bind(const std::string& host) { return host == "127.0.0.1" || host == "::1" || host == "localhost"; } + +std::map bearer_tokens_from_env() { + std::map tokens; + const char* configured = std::getenv("VGI_BEARER_TOKENS"); + if (!configured || !*configured) return tokens; + std::string entries(configured); + size_t begin = 0; + while (begin <= entries.size()) { + const auto end = entries.find(',', begin); + const auto entry = entries.substr(begin, end == std::string::npos ? end : end - begin); + const auto separator = entry.find('='); + if (separator == std::string::npos || separator == 0 || separator + 1 == entry.size()) { + throw std::invalid_argument( + "VGI_BEARER_TOKENS must contain comma-separated token=principal entries"); + } + tokens.emplace(entry.substr(0, separator), entry.substr(separator + 1)); + if (end == std::string::npos) break; + begin = end + 1; + } + return tokens; +} + +void configure_bearer_auth(vgi_rpc::HttpConfig& config, + const std::map& tokens) { + if (tokens.empty()) return; + config.peer_identity_providers.push_back( + [tokens](const vgi_rpc::PeerResolutionContext& context) { + const auto authorization = context.header("authorization"); + constexpr const char* kPrefix = "Bearer "; + if (!authorization || authorization->rfind(kPrefix, 0) != 0) { + throw vgi_rpc::PeerIdentityRejected("missing bearer credential"); + } + const auto found = + tokens.find(authorization->substr(std::char_traits::length(kPrefix))); + if (found == tokens.end()) { + throw vgi_rpc::PeerIdentityRejected("invalid bearer credential"); + } + return vgi_rpc::PeerIdentityResult::available(vgi_rpc::PeerIdentity( + "bearer", "authorization", vgi_rpc::IdentityAssurance::CONFIGURED_PROXY, + "vgi-worker", "http", vgi_rpc::PeerSubjectKind::USER, found->second, + vgi_rpc::SubjectStability::STABLE, /*subject_verified=*/true)); + }); + config.peer_authentication_policy = vgi_rpc::peer_identity_primary("bearer"); +} } // namespace Worker::Worker() : disp_(std::make_unique()) {} @@ -252,7 +297,11 @@ void Worker::run(int argc, char** argv) { port = parse_http_port(args[i + 1]); } if (iroh_issuer.empty()) { - server->serve_http(http_host, port); + vgi_rpc::HttpConfig config; + config.host = http_host; + config.port = port; + configure_bearer_auth(config, bearer_tokens_from_env()); + server->serve_http(config); } else { if (!is_loopback_bind(http_host)) { refuse( From 09741ad7f8282467d26bcd3b9173643b61afb22a Mon Sep 17 00:00:00 2001 From: Rusty Conover Date: Thu, 10 Sep 2026 21:19:14 -0400 Subject: [PATCH 3/6] Track HTTP test workers for cleanup --- scripts/run_tests.sh | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/scripts/run_tests.sh b/scripts/run_tests.sh index 23116a6..6812cfb 100755 --- a/scripts/run_tests.sh +++ b/scripts/run_tests.sh @@ -123,7 +123,7 @@ start_http_worker() { # name catalog [extra-env...] echo "[harness] HTTP worker '$name' never reported a port; see $CACHE/worker.log" >&2 return 1 fi - echo "http://127.0.0.1:$port" + STARTED_HTTP_URL="http://127.0.0.1:$port" } # Opt-in, because each one is a process held open for the whole run and the @@ -131,9 +131,12 @@ start_http_worker() { # name catalog [extra-env...] HTTP_ENV=() if [[ "${VGI_HTTP:-0}" == "1" ]]; then echo "[harness] starting HTTP workers..." - H_EXAMPLE=$(start_http_worker example example) || exit 1 - H_VERSIONED=$(start_http_worker versioned versioned) || exit 1 - H_VERSIONED_TABLES=$(start_http_worker versioned_tables versioned_tables) || exit 1 + start_http_worker example example || exit 1 + H_EXAMPLE=$STARTED_HTTP_URL + start_http_worker versioned versioned || exit 1 + H_VERSIONED=$STARTED_HTTP_URL + start_http_worker versioned_tables versioned_tables || exit 1 + H_VERSIONED_TABLES=$STARTED_HTTP_URL echo "[harness] example=$H_EXAMPLE versioned=$H_VERSIONED tables=$H_VERSIONED_TABLES" # VGI_HTTP_TRANSPORT is a flag: it says VGI_TEST_WORKER is itself a URL, so # the whole suite runs over HTTP rather than by spawning a subprocess. @@ -150,8 +153,9 @@ fi # exporting its token into the launcher run, where bearer_token is invalid. H_BEARER="" if [[ $FULL_RUN == 1 ]]; then - H_BEARER=$(start_http_worker bearer example \ - VGI_BEARER_TOKENS=test-secret-token=test-principal) || exit 1 + start_http_worker bearer example \ + VGI_BEARER_TOKENS=test-secret-token=test-principal || exit 1 + H_BEARER=$STARTED_HTTP_URL fi ARGS=() From e20ff56037401a467463372ac8f4d21e62f8ce6e Mon Sep 17 00:00:00 2001 From: Rusty Conover Date: Thu, 10 Sep 2026 21:40:28 -0400 Subject: [PATCH 4/6] Prevent concurrent storage ID reuse --- src/storage.cpp | 26 ++++++++++++++++++++------ 1 file changed, 20 insertions(+), 6 deletions(-) diff --git a/src/storage.cpp b/src/storage.cpp index a988bf8..8ee765c 100644 --- a/src/storage.cpp +++ b/src/storage.cpp @@ -298,13 +298,27 @@ class FilesystemStorage : public FunctionStorage { if (!vgi::portable::TryClaimFile(reservation.string(), "vgi storage")) { continue; // another worker took this id } - write_file_atomically(dir / (pad(id) + ".entry"), value); - // Dropped once the payload is there, so the marker means exactly - // one thing to a reader: this id is reserved and its entry has not - // landed yet. A popper uses that to tell "someone already took it" - // from "it is still on its way". + const auto entry = dir / (pad(id) + ".entry"); + const auto used = dir / (pad(id) + ".used"); std::error_code ec; - fs::remove(reservation, ec); + // `from` came from an unlocked directory scan. A publisher can + // finish between that scan and this claim, leaving the permanent + // tombstone as the authority that this id is already spent. + if (fs::exists(entry, ec) || fs::exists(used, ec)) { + fs::remove(reservation, ec); + continue; + } + write_file_atomically(entry, value); + // Keep a tombstone after publication. Removing the reservation + // made this id reusable as soon as a racing appender's directory + // scan missed the new entry; that process could then publish over + // the existing file and silently lose a batch. The queue also + // needs ids to remain spent after their entry has been popped. + fs::rename(reservation, used, ec); + if (ec) { + throw std::runtime_error("vgi storage: cannot finalize log id under " + + dir.string()); + } return id; } throw std::runtime_error("vgi storage: could not claim a log id under " + dir.string()); From d5a21e018db8a8e9002b8bc7390222560946183c Mon Sep 17 00:00:00 2001 From: Rusty Conover Date: Thu, 10 Sep 2026 21:53:16 -0400 Subject: [PATCH 5/6] Secure split token redemption over HTTP --- example-worker/table/splits.cpp | 3 +- src/catalog.cpp | 15 +++- src/dispatcher.cpp | 6 +- src/dispatcher.h | 16 +++- src/function_dispatch.cpp | 18 ++-- src/split_token.cpp | 51 ++++++++--- src/split_token.h | 25 ++++-- src/worker.cpp | 20 +++++ tests/function_test.cpp | 152 ++++++++++++++++++++++++++++++++ 9 files changed, 270 insertions(+), 36 deletions(-) diff --git a/example-worker/table/splits.cpp b/example-worker/table/splits.cpp index abf3711..3350ae9 100644 --- a/example-worker/table/splits.cpp +++ b/example-worker/table/splits.cpp @@ -969,7 +969,8 @@ class SplitEchoFilters : public vgi::TableFunction { void register_splits(vgi::Worker& worker) { using Shape = SplitFunction::Shape; worker.register_table(std::make_shared( - "split_sequence", Shape::Even, "Integers 0..n-1, divided into n contiguous splits")); + "split_sequence", Shape::Even, "Integers 0..n-1, divided into n contiguous splits", + /*catalog_version=*/1)); worker.register_table( std::make_shared("split_empty_ranges", Shape::EmptyRanges, "Integers 0..n-1, where every other split names no rows")); diff --git a/src/catalog.cpp b/src/catalog.cpp index e2ae45c..1329030 100644 --- a/src/catalog.cpp +++ b/src/catalog.cpp @@ -533,7 +533,7 @@ vgi_rpc::Result Dispatcher::catalog_attach(const vgi_rpc::Request& request) { .set_bool("supports_transactions", model.supports_transactions) .set_bool("supports_time_travel", supports_time_travel(model)) .set_bool("catalog_version_frozen", true) - .set_int64("catalog_version", 1) + .set_int64("catalog_version", *current_catalog_version(request)) .set_bool("attach_opaque_data_required", true) .set_string("default_schema", "main") // On, and per-table from here: the flag gates every @@ -555,9 +555,10 @@ vgi_rpc::Result Dispatcher::catalog_attach(const vgi_rpc::Request& request) { return envelope(batch.fill_defaults().finish()); } -vgi_rpc::Result Dispatcher::catalog_version(const vgi_rpc::Request&) { - return envelope( - wire::ResultBuilder(payload_schema_of("catalog_version")).set_int64("version", 1).finish()); +vgi_rpc::Result Dispatcher::catalog_version(const vgi_rpc::Request& request) { + return envelope(wire::ResultBuilder(payload_schema_of("catalog_version")) + .set_int64("version", *current_catalog_version(request)) + .finish()); } void Dispatcher::catalog_detach(const vgi_rpc::Request&) { @@ -571,6 +572,12 @@ void Dispatcher::catalog_detach(const vgi_rpc::Request&) { namespace {} // namespace +std::optional Dispatcher::current_catalog_version(const vgi_rpc::Request&) const { + // One source for ATTACH, catalog_version, split minting, and redemption. + // When catalogs become mutable, their live version lookup belongs here. + return 1; +} + vgi_rpc::Result Dispatcher::catalog_transaction_begin(const vgi_rpc::Request&) { // Minted here rather than by the engine, and unique across processes: the // pool hands a different worker to each RPC of one transaction, so two diff --git a/src/dispatcher.cpp b/src/dispatcher.cpp index 8cfdb4c..a053243 100644 --- a/src/dispatcher.cpp +++ b/src/dispatcher.cpp @@ -183,7 +183,6 @@ void Dispatcher::install(vgi_rpc::ServerBuilder& builder) { // still registered — see the note above — but refuses when called. const std::unordered_map unary = { {"bind", &Dispatcher::bind}, - {"table_function_plan", &Dispatcher::table_function_plan}, {"table_function_cardinality", &Dispatcher::table_function_cardinality}, {"table_function_statistics", &Dispatcher::table_function_statistics}, {"table_function_dynamic_to_string", &Dispatcher::table_function_dynamic_to_string}, @@ -222,6 +221,7 @@ void Dispatcher::install(vgi_rpc::ServerBuilder& builder) { }; // The handlers that also take the call's log channel. const std::unordered_map unary_with_context = { + {"table_function_plan", &Dispatcher::table_function_plan}, {"table_buffering_process", &Dispatcher::table_buffering_process}, {"table_buffering_combine", &Dispatcher::table_buffering_combine}, }; @@ -242,9 +242,9 @@ void Dispatcher::install(vgi_rpc::ServerBuilder& builder) { // only placeholders — the factory returns the real pair. builder.add_exchange( name, spec.params, arrow::schema({}), arrow::schema({}), - [this, name](const vgi_rpc::Request& req, vgi_rpc::CallContext&) { + [this, name](const vgi_rpc::Request& req, vgi_rpc::CallContext& ctx) { trace(name); - return this->init(req); + return this->init(req, ctx); }, "", global_init_response_schema()); continue; diff --git a/src/dispatcher.h b/src/dispatcher.h index e500962..2a2bef9 100644 --- a/src/dispatcher.h +++ b/src/dispatcher.h @@ -21,6 +21,7 @@ #include "vgi/copy_from.h" #include "vgi/copy_to.h" #include "vgi/table_in_out.h" +#include "split_token.h" namespace vgi { @@ -101,6 +102,13 @@ class Dispatcher { // Register every VGI method on `builder`. void install(vgi_rpc::ServerBuilder& builder); + // HTTP tokens are authenticated with the transport's process key. Raw + // transports deliberately remain keyless inside their existing trust + // boundary. + void set_split_token_signing_key(split_token::SigningKey key) { + split_token_signing_key_ = std::move(key); + } + using UnaryHandler = vgi_rpc::Result (Dispatcher::*)(const vgi_rpc::Request&); // The few handlers that need the call's own channel back to the client. // Kept separate rather than widening every signature: only a method that @@ -117,7 +125,8 @@ class Dispatcher { // than in a 5,000-line switch. vgi_rpc::Result bind(const vgi_rpc::Request& request); - vgi_rpc::Result table_function_plan(const vgi_rpc::Request& request); + vgi_rpc::Result table_function_plan(const vgi_rpc::Request& request, + vgi_rpc::CallContext& context); vgi_rpc::Result table_function_cardinality(const vgi_rpc::Request& request); vgi_rpc::Result table_function_statistics(const vgi_rpc::Request& request); vgi_rpc::Result table_function_dynamic_to_string(const vgi_rpc::Request& request); @@ -138,7 +147,7 @@ class Dispatcher { vgi_rpc::Result table_buffering_combine(const vgi_rpc::Request& request, vgi_rpc::CallContext& context); vgi_rpc::Result table_buffering_destructor(const vgi_rpc::Request& request); - vgi_rpc::Stream init(const vgi_rpc::Request& request); + vgi_rpc::Stream init(const vgi_rpc::Request& request, vgi_rpc::CallContext& context); vgi_rpc::Result catalog_attach(const vgi_rpc::Request& request); vgi_rpc::Result catalog_schemas(const vgi_rpc::Request& request); @@ -164,6 +173,8 @@ class Dispatcher { void catalog_transaction_rollback(const vgi_rpc::Request& request); private: + std::optional current_catalog_version(const vgi_rpc::Request& request) const; + // The fields every FunctionInfo carries, whatever kind of function it is. // Each `encode_*_info` starts here and appends only what is its own; the // five of them spelling the shared set out separately is how a field ends @@ -309,6 +320,7 @@ class Dispatcher { // `catalog()` mean. Held indirectly so a reference handed out by // `catalog(name)` survives a later addition. std::vector> catalogs_; + std::optional split_token_signing_key_; std::set hidden_; std::vector> scalars_; // Parallel to scalars_: where each one is declared. diff --git a/src/function_dispatch.cpp b/src/function_dispatch.cpp index 5fb8503..d07e258 100644 --- a/src/function_dispatch.cpp +++ b/src/function_dispatch.cpp @@ -1039,7 +1039,8 @@ std::string encode_scan_split(const ScanSplit& split, const std::string& token) } // namespace -vgi_rpc::Result Dispatcher::table_function_plan(const vgi_rpc::Request& request) { +vgi_rpc::Result Dispatcher::table_function_plan(const vgi_rpc::Request& request, + vgi_rpc::CallContext& context) { auto plan_request = wire::get_ipc(request.batch(), "request"); if (!plan_request) throw std::runtime_error("plan: empty request"); @@ -1085,13 +1086,15 @@ vgi_rpc::Result Dispatcher::table_function_plan(const vgi_rpc::Request& request) bind_params.schema_path, function_name, wire::get_optional_binary(bind_call, "arguments").value_or(std::string{}), wire::get_optional_binary(bind_call, "settings").value_or(std::string{})); - const auto anchor = split_token::anchor_for(result.catalog_version); + const auto anchor = split_token::anchor_for( + result.catalog_version ? result.catalog_version : current_catalog_version(request)); std::vector splits; splits.reserve(result.splits.size()); for (const auto& split : result.splits) { splits.push_back( - encode_scan_split(split, split_token::build(split.payload, fingerprint, anchor))); + encode_scan_split(split, split_token::build(split.payload, fingerprint, anchor, + split_token_signing_key_, context.auth()))); } auto payload = wire::ResultBuilder(payload_schema_of("table_function_plan")); @@ -1227,7 +1230,7 @@ vgi_rpc::Result Dispatcher::table_function_dynamic_to_string(const vgi_rpc::Requ .finish()); } -vgi_rpc::Stream Dispatcher::init(const vgi_rpc::Request& request) { +vgi_rpc::Stream Dispatcher::init(const vgi_rpc::Request& request, vgi_rpc::CallContext& context) { auto init_request = wire::get_ipc(request.batch(), "request"); if (!init_request) throw std::runtime_error("init: empty request"); @@ -1306,14 +1309,13 @@ vgi_rpc::Stream Dispatcher::init(const vgi_rpc::Request& request) { bind_params.schema_path, function_name, wire::get_optional_binary(bind_call, "arguments").value_or(std::string{}), wire::get_optional_binary(bind_call, "settings").value_or(std::string{})); - // The anchor a plan sealed in. Nothing here time-travels, so the - // current anchor is the one a plan would mint now. - const auto anchor = split_token::anchor_for(std::nullopt); + const auto anchor = split_token::anchor_for(current_catalog_version(request)); std::vector payloads; payloads.reserve(tokens.size()); for (const auto& token : tokens) { - auto opened = split_token::open(token, fingerprint, anchor); + auto opened = split_token::open(token, fingerprint, anchor, split_token_signing_key_, + context.auth()); if (!opened.payload) { if (opened.error == split_token::OpenError::SnapshotExpired) { throw std::runtime_error("SPLIT_SNAPSHOT_EXPIRED: split token for '" + diff --git a/src/split_token.cpp b/src/split_token.cpp index aabbce7..9552882 100644 --- a/src/split_token.cpp +++ b/src/split_token.cpp @@ -15,6 +15,7 @@ namespace { // without the delimiters, a function named "a" in schema "bc" and one named // "ab" in schema "c" would be the same bytes. constexpr const char* kAadPrefix = "vgi.split_token.v1"; +constexpr char kSealVersion = '\x01'; void feed(vgi_rpc::crypto::Sha256& h, const char* label, const std::string& value) { h.update(label); @@ -28,6 +29,19 @@ void put_u16_le(std::string& out, uint16_t v) { out.push_back(static_cast((v >> 8) & 0xFF)); } +std::string identity_tail(const vgi_rpc::AuthContext& auth) { + if (!auth.authenticated) return std::string("\0anonymous", 10); + std::string out(1, '\1'); + out.append(auth.domain); + out.push_back('\0'); + if (auth.principal) out.append(*auth.principal); + return out; +} + +std::string token_aad(const std::string& body, const vgi_rpc::AuthContext& auth) { + return body + identity_tail(auth); +} + } // namespace std::string bind_fingerprint(const SchemaPath& schema_path, const std::string& function_name, @@ -52,7 +66,8 @@ std::string anchor_for(std::optional catalog_version) { } std::string build(const std::string& payload, const std::string& fingerprint, - const std::string& anchor) { + const std::string& anchor, const std::optional& signing_key, + const vgi_rpc::AuthContext& auth) { if (fingerprint.size() != kFingerprintLen) { throw std::invalid_argument("split token: fingerprint must be 16 bytes"); } @@ -60,27 +75,36 @@ std::string build(const std::string& payload, const std::string& fingerprint, throw std::invalid_argument("split token: consistency anchor exceeds u16"); } std::string out; - out.reserve(kHeaderLen + anchor.size() + payload.size()); + out.reserve( + kHeaderLen + anchor.size() + payload.size() + + (signing_key ? 1 + vgi_rpc::crypto::kAeadNonceBytes + vgi_rpc::crypto::kAeadTagBytes : 0)); out.push_back(static_cast(kFormatVersion)); - out.push_back(0); // no key on these transports, so nothing is sealed + out.push_back(signing_key ? static_cast(kFlagPayloadSealed) : 0); put_u16_le(out, static_cast(anchor.size())); out.append(fingerprint); out.append(anchor); - out.append(payload); + if (signing_key) { + out.push_back(kSealVersion); + out.append(vgi_rpc::crypto::aead_seal(*signing_key, payload, token_aad(out, auth))); + } else { + out.append(payload); + } return out; } OpenResult open(const std::string& token, const std::string& expected_fingerprint, - const std::string& current_anchor) { + const std::string& current_anchor, const std::optional& signing_key, + const vgi_rpc::AuthContext& auth) { const OpenResult invalid{std::nullopt, OpenError::Invalid}; if (token.size() < kHeaderLen) return invalid; if (static_cast(token[0]) != kFormatVersion) return invalid; const auto flags = static_cast(token[1]); - // Every bit is reserved here, `payload_sealed` included: this SDK holds no - // key, so a token claiming to be sealed is one we cannot open, and a token - // setting a reserved bit is from a future this build does not speak. - if (flags != 0) return invalid; + if ((flags & ~kFlagPayloadSealed) != 0) return invalid; + const bool sealed = (flags & kFlagPayloadSealed) != 0; + // The worker's key state is authoritative. Trusting the plaintext flag to + // select the keyless path would let an attacker clear it and forge work. + if (signing_key.has_value() != sealed) return invalid; const auto anchor_len = static_cast(static_cast(token[2])) | (static_cast(static_cast(token[3])) << 8); @@ -99,7 +123,14 @@ OpenResult open(const std::string& token, const std::string& expected_fingerprin return {std::nullopt, OpenError::SnapshotExpired}; } - return {token.substr(end_of_anchor), OpenError::None}; + const auto payload = token.substr(end_of_anchor); + if (!sealed) return {payload, OpenError::None}; + + if (payload.empty() || payload.front() != kSealVersion) return invalid; + auto opened = vgi_rpc::crypto::aead_open(*signing_key, payload.substr(1), + token_aad(token.substr(0, end_of_anchor), auth)); + if (!opened) return invalid; + return {std::move(*opened), OpenError::None}; } } // namespace vgi::split_token diff --git a/src/split_token.h b/src/split_token.h index 40f3085..cfaaee2 100644 --- a/src/split_token.h +++ b/src/split_token.h @@ -18,17 +18,21 @@ // cross-SDK fixture covers it. That is why it can hash C++ spellings of the // bind fields rather than reproducing the reference's Python `repr`. // -// Nothing here seals: this SDK's transports carry no signing key, and the -// reference's own header explains why the header must stay plaintext where -// DuckDB runs. If a key ever arrives, the keyed/keyless decision has to come -// from the worker's key state and never from the token's own `flags` byte — -// trusting that byte is `alg:none`. +// The header is always plaintext. HTTP supplies its process token key and +// seals the payload; raw transports remain keyless inside their existing +// trust boundary. The keyed/keyless decision comes from worker key state, +// never from the attacker-controlled flags byte — trusting that bit alone is +// `alg:none`. #pragma once +#include #include #include #include +#include +#include + #include "vgi/types.h" namespace vgi::split_token { @@ -43,9 +47,13 @@ inline constexpr size_t kHeaderLen = 4 + kFingerprintLen; std::string bind_fingerprint(const SchemaPath& schema_path, const std::string& function_name, const std::string& arguments, const std::string& settings); -// Stamp a payload into a token. +using SigningKey = std::array; + +// Stamp a payload into a token. A keyed worker seals the payload and binds it +// to the caller identity; a keyless worker leaves it plaintext. std::string build(const std::string& payload, const std::string& fingerprint, - const std::string& anchor); + const std::string& anchor, const std::optional& signing_key, + const vgi_rpc::AuthContext& auth); enum class OpenError { None, Invalid, SnapshotExpired }; @@ -58,7 +66,8 @@ struct OpenResult { // snapshot can be retried after replanning, while a malformed or wrongly-bound // token cannot. OpenResult open(const std::string& token, const std::string& expected_fingerprint, - const std::string& current_anchor); + const std::string& current_anchor, const std::optional& signing_key, + const vgi_rpc::AuthContext& auth); // The consistency anchor for a catalog version: int64, little-endian, and an // absent version is zero — the same spelling the reference uses. diff --git a/src/worker.cpp b/src/worker.cpp index 23740de..b1b61e8 100644 --- a/src/worker.cpp +++ b/src/worker.cpp @@ -1,6 +1,7 @@ // © Copyright 2025, 2026 Query Farm LLC - https://query.farm #include "vgi/worker.h" +#include #include #include #include @@ -10,6 +11,7 @@ #include #include +#include #include #include #include @@ -91,6 +93,20 @@ void configure_bearer_auth(vgi_rpc::HttpConfig& config, }); config.peer_authentication_policy = vgi_rpc::peer_identity_primary("bearer"); } + +void configure_signing_key(vgi_rpc::HttpConfig& config) { + const char* configured = std::getenv("VGI_SIGNING_KEY"); + if (!configured || !*configured) return; + + const std::string raw(configured); + if (raw.size() == config.token_key.size()) { + std::copy(raw.begin(), raw.end(), config.token_key.begin()); + return; + } + vgi_rpc::crypto::Sha256 hash; + hash.update(raw); + config.token_key = hash.digest(); +} } // namespace Worker::Worker() : disp_(std::make_unique()) {} @@ -300,7 +316,9 @@ void Worker::run(int argc, char** argv) { vgi_rpc::HttpConfig config; config.host = http_host; config.port = port; + configure_signing_key(config); configure_bearer_auth(config, bearer_tokens_from_env()); + disp_->set_split_token_signing_key(config.token_key); server->serve_http(config); } else { if (!is_loopback_bind(http_host)) { @@ -312,11 +330,13 @@ void Worker::run(int argc, char** argv) { vgi_rpc::HttpConfig config; config.host = http_host; config.port = port; + configure_signing_key(config); config.peer_identity_providers.push_back(vgi_rpc::iroh_forwarded_header_provider( {std::move(iroh_issuer), std::move(iroh_trusted_proxies)})); config.peer_authentication_policy = iroh_observe ? vgi_rpc::observe_peer_identity : vgi_rpc::peer_identity_primary("iroh"); + disp_->set_split_token_signing_key(config.token_key); server->serve_http(config); } std::exit(0); diff --git a/tests/function_test.cpp b/tests/function_test.cpp index 02bf2f7..4dc8b90 100644 --- a/tests/function_test.cpp +++ b/tests/function_test.cpp @@ -1,9 +1,12 @@ // © Copyright 2025, 2026 Query Farm LLC - https://query.farm #include +#include #include #include #include +#include +#include #include #include @@ -13,6 +16,7 @@ #include "vgi/catalog.h" #include "vgi/generated/vgi_protocol_schemas.hpp" #include "vgi/pushdown.h" +#include "split_token.h" #include "wire.h" namespace { @@ -115,6 +119,23 @@ std::shared_ptr null_value(const std::shared_ptr& return result.MoveValueUnsafe(); } +std::string from_hex(const std::string& hex) { + std::string out; + out.reserve(hex.size() / 2); + for (size_t i = 0; i < hex.size(); i += 2) { + out.push_back(static_cast(std::stoul(hex.substr(i, 2), nullptr, 16))); + } + return out; +} + +vgi_rpc::AuthContext authenticated_as(const std::string& principal) { + auto auth = vgi_rpc::AuthContext::anonymous(); + auth.domain = "test"; + auth.authenticated = true; + auth.principal = principal; + return auth; +} + } // namespace TEST_CASE("a fixed return type binds without an override", "[function]") { @@ -463,3 +484,134 @@ TEST_CASE("Filter v2 list_contains uses DuckDB nested and NaN equality", "[filte filter_batch(document, int64_list({1, std::nullopt})), {}, nested_input->schema()); REQUIRE(nested_filter.apply(nested_input)->num_rows() == 1); } + +TEST_CASE("split token shared vectors reach their declared verdict", "[split-token]") { + // Byte-for-byte copies of the canonical vectors generated by + // vgi-python/tests/data/split_tokens/generate.py. Keeping the raw bytes in + // the test catches a parser that is self-consistent but disagrees with the + // other SDKs about the envelope layout. + struct Vector { + const char* name; + const char* token_hex; + vgi::split_token::OpenError verdict; + bool keyed; + }; + const std::vector vectors = { + {"valid_unsealed", + "01000800000102030405060708090a0b0c0d0e0f2f0000000000000066696c653d333b763d3437", + vgi::split_token::OpenError::None, false}, + {"valid_sealed", + "01010800000102030405060708090a0b0c0d0e0f2f0000000000000001551ed7558e77fdcea243a" + "2394af27e800146335fb36c7c1f99ae3b13773a7466926cc4f5c1e3701d1d9234182a3b7588852364", + vgi::split_token::OpenError::None, true}, + {"bad_flags_unsealed_but_key_present", + "01000800000102030405060708090a0b0c0d0e0f2f000000000000004f544845522054454e414e54204441544" + "1", + vgi::split_token::OpenError::Invalid, true}, + {"bad_fingerprint", + "01000800eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee2f0000000000000066696c653d333b763d3437", + vgi::split_token::OpenError::Invalid, false}, + {"stale_anchor", + "01000800000102030405060708090a0b0c0d0e0f010000000000000066696c653d333b763d3437", + vgi::split_token::OpenError::SnapshotExpired, false}, + {"truncated", "01000800000102030405", vgi::split_token::OpenError::Invalid, false}, + {"reserved_flag_bit", + "01020800000102030405060708090a0b0c0d0e0f2f0000000000000066696c653d333b763d3437", + vgi::split_token::OpenError::Invalid, false}, + {"bad_version", + "09000800000102030405060708090a0b0c0d0e0f2f0000000000000066696c653d333b763d3437", + vgi::split_token::OpenError::Invalid, false}, + {"anchor_len_overrun", + "01000f27000102030405060708090a0b0c0d0e0f2f0000000000000066696c653d333b763d3437", + vgi::split_token::OpenError::Invalid, false}, + }; + + vgi::split_token::SigningKey key{}; + for (size_t i = 0; i < key.size(); ++i) key[i] = static_cast(i); + const auto fingerprint = from_hex("000102030405060708090a0b0c0d0e0f"); + const auto anchor = from_hex("2f00000000000000"); + const auto anonymous = vgi_rpc::AuthContext::anonymous(); + + for (const auto& vector : vectors) { + CAPTURE(vector.name); + const std::optional signing_key = + vector.keyed ? std::optional(key) : std::nullopt; + const auto opened = vgi::split_token::open(from_hex(vector.token_hex), fingerprint, anchor, + signing_key, anonymous); + REQUIRE(opened.error == vector.verdict); + if (vector.verdict == vgi::split_token::OpenError::None) { + REQUIRE(opened.payload == std::optional("file=3;v=47")); + } else { + REQUIRE_FALSE(opened.payload.has_value()); + } + } +} + +TEST_CASE("split token stamping matches the shared unsealed vector", "[split-token]") { + const auto fingerprint = from_hex("000102030405060708090a0b0c0d0e0f"); + const auto anchor = vgi::split_token::anchor_for(47); + const auto token = vgi::split_token::build("file=3;v=47", fingerprint, anchor, std::nullopt, + vgi_rpc::AuthContext::anonymous()); + REQUIRE( + token == + from_hex("01000800000102030405060708090a0b0c0d0e0f2f0000000000000066696c653d333b763d3437")); +} + +TEST_CASE("sealed split tokens are bound to key, header, and identity", "[split-token]") { + vgi::split_token::SigningKey key{}; + key.fill(0x11); + auto wrong_key = key; + wrong_key[0] ^= 0xff; + const std::string fingerprint(16, '\x05'); + const auto anchor = vgi::split_token::anchor_for(1); + const auto alice = authenticated_as("alice"); + const auto bob = authenticated_as("bob"); + + auto token = vgi::split_token::build("tenant=alice", fingerprint, anchor, key, alice); + REQUIRE((static_cast(token[1]) & vgi::split_token::kFlagPayloadSealed) != 0); + REQUIRE(token.find("tenant=alice") == std::string::npos); + REQUIRE(vgi::split_token::open(token, fingerprint, anchor, key, alice).payload == + std::optional("tenant=alice")); + REQUIRE(vgi::split_token::open(token, fingerprint, anchor, key, bob).error == + vgi::split_token::OpenError::Invalid); + REQUIRE(vgi::split_token::open(token, fingerprint, anchor, wrong_key, alice).error == + vgi::split_token::OpenError::Invalid); + + token[4] ^= static_cast(0xff); + REQUIRE(vgi::split_token::open(token, token.substr(4, 16), anchor, key, alice).error == + vgi::split_token::OpenError::Invalid); +} + +TEST_CASE("a keyed worker refuses alg-none split tokens", "[split-token]") { + vgi::split_token::SigningKey key{}; + key.fill(0x2a); + const std::string fingerprint(16, '\x07'); + const auto anchor = vgi::split_token::anchor_for(47); + const auto anonymous = vgi_rpc::AuthContext::anonymous(); + + const auto forged = + vgi::split_token::build("file=evil", fingerprint, anchor, std::nullopt, anonymous); + REQUIRE(vgi::split_token::open(forged, fingerprint, anchor, key, anonymous).error == + vgi::split_token::OpenError::Invalid); + + const auto sealed = vgi::split_token::build("file=ok", fingerprint, anchor, key, anonymous); + REQUIRE(vgi::split_token::open(sealed, fingerprint, anchor, std::nullopt, anonymous).error == + vgi::split_token::OpenError::Invalid); +} + +TEST_CASE("split token bind failures precede stale anchors", "[split-token]") { + const std::string fingerprint(16, '\x09'); + const std::string other_fingerprint(16, '\x0a'); + const auto old_anchor = vgi::split_token::anchor_for(47); + const auto current_anchor = vgi::split_token::anchor_for(48); + const auto anonymous = vgi_rpc::AuthContext::anonymous(); + const auto token = + vgi::split_token::build("file=1", fingerprint, old_anchor, std::nullopt, anonymous); + + REQUIRE( + vgi::split_token::open(token, other_fingerprint, current_anchor, std::nullopt, anonymous) + .error == vgi::split_token::OpenError::Invalid); + REQUIRE( + vgi::split_token::open(token, fingerprint, current_anchor, std::nullopt, anonymous).error == + vgi::split_token::OpenError::SnapshotExpired); +} From 2e112b99f321326301647eedcea1eefa5147a8ef Mon Sep 17 00:00:00 2001 From: Rusty Conover Date: Thu, 10 Sep 2026 21:56:16 -0400 Subject: [PATCH 6/6] Align split token seal AAD with redemption --- src/split_token.cpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/split_token.cpp b/src/split_token.cpp index 9552882..a44c961 100644 --- a/src/split_token.cpp +++ b/src/split_token.cpp @@ -84,8 +84,9 @@ std::string build(const std::string& payload, const std::string& fingerprint, out.append(fingerprint); out.append(anchor); if (signing_key) { + const auto aad = token_aad(out, auth); out.push_back(kSealVersion); - out.append(vgi_rpc::crypto::aead_seal(*signing_key, payload, token_aad(out, auth))); + out.append(vgi_rpc::crypto::aead_seal(*signing_key, payload, aad)); } else { out.append(payload); }