[FFI] Fix headers-provider use-after-free via ownership transfer

cpp/CLAUDE.md

@@ -93,14 +93,18 @@ instead, configure with `-DZEROBUS_FFI_LIBRARY=<path-to-.a>` and
 
 When using a custom `HeadersProvider`:
 - `detail::zerobus_cpp_headers_trampoline` is an `extern "C"` function passed to
-  `zerobus_sdk_create_stream_with_headers_provider`. `user_data` is the raw
-  `HeadersProvider*`.
-- The `Stream` / `ArrowStream` keeps a `std::shared_ptr<HeadersProvider>` alive
-  for its whole lifetime; the handle is freed in the destructor **before** the
-  shared_ptr member is destroyed. Necessary but not always sufficient (see
-  `headers_provider.hpp`): a `get_headers()` call still running when `close()`
-  times out (~1s) can be invoked on a freed provider. Keep `get_headers()` well
-  under that budget, or keep the provider alive past the stream.
+  `zerobus_sdk_create_stream_with_headers_provider`. `user_data` is a
+  heap-allocated `std::shared_ptr<HeadersProvider>*` whose **ownership is handed
+  to the FFI**, alongside `detail::zerobus_cpp_headers_free` as the FFI's
+  provider `free_user_data` destroy callback.
+- The FFI invokes `zerobus_cpp_headers_free` exactly once, on every path (on
+  success after the core's last reference drops — i.e. after any in-flight
+  `get_headers()` has returned; on a failed create before returning). This is
+  what closes the recovery-vs-teardown use-after-free: a slow `get_headers()`
+  racing stream teardown can no longer run on a freed provider. Because the FFI
+  owns the provider, the `Stream` / `ArrowStream` no longer keeps a
+  `shared_ptr<HeadersProvider>` member, and the wrapper must **not** free
+  `user_data` itself (doing so would double-free).
 - The trampoline marshals the returned map into `CHeaders` using the allocators
   the Rust core's `zerobus_free_headers` expects to free. Do not change the
   allocator without checking `zerobus_free_headers` in the FFI crate.

cpp/NEXT_CHANGELOG.md

@@ -14,6 +14,16 @@
 
 ### Bug Fixes
 
+- Fixed a use-after-free in which a custom `HeadersProvider` could be destroyed
+  while the Rust core was still inside a `get_headers()` call into it during
+  connection recovery. Provider ownership is now handed to the FFI as a
+  heap-allocated `shared_ptr` released by a destroy callback
+  (`detail::zerobus_cpp_headers_free`) only after any in-flight `get_headers()`
+  has returned; the `Stream` / `ArrowStream` no longer keeps its own provider
+  `shared_ptr`. Public API is unchanged — `create_stream` /
+  `create_arrow_stream` still take a `std::shared_ptr<HeadersProvider>` — and you
+  no longer need to keep your own reference alive past `create_stream`.
+
 ### Documentation
 
 - Added the C++ SDK `README.md` (build, install, quickstart for JSON / proto /

cpp/include/zerobus/arrow_stream.hpp

@@ -86,11 +86,12 @@ class ArrowStream {
 
  private:
   friend class Sdk;
-  ArrowStream(CArrowStream* handle, std::shared_ptr<HeadersProvider> provider)
-      : handle_(handle), provider_(std::move(provider)) {}
+  ArrowStream(CArrowStream* handle, std::shared_ptr<HeadersProvider> /*unused*/)
+      : handle_(handle) {}
 
   CArrowStream* handle_;
-  std::shared_ptr<HeadersProvider> provider_;
+  // The headers provider (if any) is owned by the FFI, not the ArrowStream (see
+  // Sdk::create_arrow_stream / headers_provider.hpp).
 };
 
 }  // namespace zerobus

cpp/include/zerobus/headers_provider.hpp

@@ -19,11 +19,19 @@ namespace zerobus {
 /// Implementations should therefore be thread-safe with respect to their own
 /// state.
 ///
-/// Lifetime: the provider must outlive the `Stream`, which holds a `shared_ptr`
-/// to it. This is necessary but not always sufficient: a `get_headers()` call
-/// still running when `close()` times out (~1s) can be invoked on a freed
-/// provider. Keep `get_headers()` well under that budget, or keep the provider
-/// alive past the `Stream`.
+/// Lifetime: you pass the provider as a `shared_ptr` and its ownership is
+/// handed to the SDK, which keeps it alive until the Rust core is done with it
+/// — after any in-flight `get_headers()` call (including one running during
+/// connection recovery) has returned. You therefore do not need to keep your
+/// own reference alive past `create_stream`, and a slow `get_headers()` racing
+/// stream teardown can no longer be invoked on a freed provider.
+///
+/// Because the SDK owns the provider, its **destructor may run on an internal
+/// SDK worker thread** (whichever thread drops the last reference), not the
+/// thread that called `create_stream` / destroyed the `Stream`. Keep
+/// `~YourProvider()` non-blocking and free of thread-affine work, and do not
+/// let it throw (it runs across the C FFI boundary, where an escaping exception
+/// terminates the process).
 class HeadersProvider {
  public:
   virtual ~HeadersProvider() = default;

cpp/include/zerobus/stream.hpp

@@ -131,15 +131,15 @@ class Stream {
 
  private:
   friend class Sdk;
-  Stream(CZerobusStream* handle, std::shared_ptr<HeadersProvider> provider,
+  Stream(CZerobusStream* handle, std::shared_ptr<HeadersProvider> /*unused*/,
          std::shared_ptr<AckCallback> ack_callback)
-      : handle_(handle),
-        provider_(std::move(provider)),
-        ack_callback_(std::move(ack_callback)) {}
+      : handle_(handle), ack_callback_(std::move(ack_callback)) {}
 
   CZerobusStream* handle_;
-  // Kept alive for the stream's lifetime; the core holds a raw pointer to it.
-  std::shared_ptr<HeadersProvider> provider_;
+  // The headers provider (if any) is owned by the FFI, not the Stream: the core
+  // frees it after any in-flight get_headers returns, so the Stream holds no
+  // reference to it (see Sdk::create_stream / headers_provider.hpp).
+  //
   // Also raw-pointed-to by the core (ack user_data), but with a weaker bound: a
   // callback can still run after close(), so dropping this at ~Stream() can
   // free it mid-call (see AckCallback). May be null.

cpp/src/arrow_stream.cpp

@@ -58,10 +58,10 @@ ArrowStream::~ArrowStream() {
   }
 }
 
-// Move carries the handle and the headers-provider together, nulling the source
-// so the handle is closed/freed exactly once.
+// Move carries the handle, nulling the source so the handle is closed/freed
+// exactly once. The headers provider is owned by the FFI, not the ArrowStream.
 ArrowStream::ArrowStream(ArrowStream&& other) noexcept
-    : handle_(other.handle_), provider_(std::move(other.provider_)) {
+    : handle_(other.handle_) {
   other.handle_ = nullptr;
 }
 
@@ -74,7 +74,6 @@ ArrowStream& ArrowStream::operator=(ArrowStream&& other) noexcept {
       zerobus_arrow_stream_free(handle_);
     }
     handle_ = other.handle_;
-    provider_ = std::move(other.provider_);
     other.handle_ = nullptr;
   }
   return *this;

cpp/src/detail/headers_callback.hpp

@@ -7,13 +7,22 @@ namespace zerobus {
 namespace detail {
 
 /// C trampoline matching `HeadersProviderCallback`. `user_data` must point to a
-/// live `HeadersProvider`. Invokes `get_headers()` and marshals the result into
-/// a `CHeaders` whose buffers are allocated so the Rust core's
+/// heap-allocated `std::shared_ptr<HeadersProvider>` (owned by the FFI; see
+/// `zerobus_cpp_headers_free`). Invokes `get_headers()` and marshals the result
+/// into a `CHeaders` whose buffers are allocated so the Rust core's
 /// `zerobus_free_headers` can release them (keys/values via the C string
 /// allocator, the array via `malloc`). Exceptions are caught and reported via
 /// `CHeaders.error_message`.
 extern "C" CHeaders zerobus_cpp_headers_trampoline(void* user_data);
 
+/// C trampoline matching the FFI's provider `free_user_data` callback. Deletes
+/// the heap-allocated `std::shared_ptr<HeadersProvider>` that `user_data`
+/// points to, releasing the provider. The FFI invokes it exactly once — after
+/// any in-flight `get_headers()` has returned — which is what makes it safe to
+/// destroy the provider (see `headers_provider.hpp`). `noexcept`: it must not
+/// unwind across the C boundary.
+extern "C" void zerobus_cpp_headers_free(void* user_data) noexcept;
+
 }  // namespace detail
 }  // namespace zerobus
 

cpp/src/headers_callback.cpp

@@ -15,6 +15,7 @@
 
 #include <cstdint>
 #include <map>
+#include <memory>
 #include <string>
 
 #include "zerobus/headers_provider.hpp"
@@ -65,7 +66,14 @@ extern "C" CHeaders zerobus_cpp_headers_trampoline(void* user_data) {
   if (user_data == nullptr) {
     return make_error("null headers provider");
   }
-  auto* provider = static_cast<HeadersProvider*>(user_data);
+  // user_data is a heap-owned shared_ptr<HeadersProvider> (see
+  // zerobus_cpp_headers_free); deref it to reach the provider.
+  auto& provider_sp =
+      *static_cast<std::shared_ptr<HeadersProvider>*>(user_data);
+  HeadersProvider* provider = provider_sp.get();
+  if (provider == nullptr) {
+    return make_error("null headers provider");
+  }
 
   std::map<std::string, std::string> headers;
   try {
@@ -115,5 +123,14 @@ extern "C" CHeaders zerobus_cpp_headers_trampoline(void* user_data) {
   return result;
 }
 
+extern "C" void zerobus_cpp_headers_free(void* user_data) noexcept {
+  // Frees the heap shared_ptr that Sdk::create_* allocated with new; the
+  // trampoline dereferences the same pointer. The FFI calls this once, after
+  // any in-flight get_headers has returned, so deleting here cannot race a
+  // callback. delete on a null pointer is a no-op, and ~shared_ptr / provider
+  // destructors are assumed not to throw (noexcept contains it if they do).
+  delete static_cast<std::shared_ptr<HeadersProvider>*>(user_data);
+}
+
 }  // namespace detail
 }  // namespace zerobus

cpp/src/sdk.cpp

@@ -180,10 +180,12 @@ Stream Sdk::create_stream(const TableProperties& table,
   return Stream(stream, nullptr, options.ack_callback);
 }
 
-// Same as above but authenticated by a custom headers provider. Validate the
-// pointer (else the FFI gets a null callback target), pass the raw provider as
-// the trampoline's user_data, and hand the shared_ptr to the Stream so it
-// outlives the core's callbacks.
+// Same as above but authenticated by a custom headers provider. Ownership of
+// the provider is handed to the FFI: it is passed as a heap-allocated
+// shared_ptr (the trampoline's user_data) which the FFI releases via
+// zerobus_cpp_headers_free only after any in-flight get_headers has returned.
+// This closes the recovery-vs-teardown use-after-free, so the Stream no longer
+// keeps its own provider shared_ptr.
 Stream Sdk::create_stream(const TableProperties& table,
                           std::shared_ptr<HeadersProvider> headers_provider,
                           const StreamOptions& options) {
@@ -192,17 +194,25 @@ Stream Sdk::create_stream(const TableProperties& table,
   }
   detail::ResultGuard guard;
   CStreamConfigurationOptions copts = detail::to_c(options);
+  // Validate the table name before allocating `owned` so a throw here can't
+  // leak it (argument evaluation order vs. the `new` is unspecified).
+  const char* c_table = detail::checked_c_str(table.table_name, "table_name");
+  // Heap-allocate the owning shared_ptr; the FFI takes ownership and frees it
+  // via zerobus_cpp_headers_free on every path (success or failure).
+  auto* owned =
+      new std::shared_ptr<HeadersProvider>(std::move(headers_provider));
   CZerobusStream* stream = zerobus_sdk_create_stream_with_headers_provider(
-      handle_, detail::checked_c_str(table.table_name, "table_name"),
-      descriptor_ptr(table.descriptor_proto), table.descriptor_proto.size(),
-      detail::zerobus_cpp_headers_trampoline, headers_provider.get(), &copts,
-      guard.ptr());
+      handle_, c_table, descriptor_ptr(table.descriptor_proto),
+      table.descriptor_proto.size(), detail::zerobus_cpp_headers_trampoline,
+      owned, detail::zerobus_cpp_headers_free, &copts, guard.ptr());
   if (stream == nullptr) {
+    // The FFI already freed `owned` via zerobus_cpp_headers_free on the failure
+    // path, so we must not delete it here.
     guard.throw_if_error();
     throw ZerobusException("failed to create stream", false);
   }
-  // Keep both provider and ack callback alive: the core raw-points to each.
-  return Stream(stream, std::move(headers_provider), options.ack_callback);
+  // Provider ownership now lives in the FFI; keep only the ack callback alive.
+  return Stream(stream, nullptr, options.ack_callback);
 }
 
 // Create an OAuth-authenticated Arrow Flight stream (Beta). The schema IPC
@@ -233,10 +243,10 @@ ArrowStream Sdk::create_arrow_stream(
   return ArrowStream(stream, nullptr);
 }
 
-// Arrow Flight stream (Beta) authenticated by a custom headers provider. Both
-// the provider and the schema bytes are validated up front; as with the proto
-// path, the ArrowStream keeps the provider's shared_ptr alive for the
-// callback's lifetime.
+// Arrow Flight stream (Beta) authenticated by a custom headers provider. As
+// with the proto path, provider ownership is handed to the FFI (a heap
+// shared_ptr freed via zerobus_cpp_headers_free after any in-flight
+// get_headers returns), so the ArrowStream no longer keeps its own shared_ptr.
 ArrowStream Sdk::create_arrow_stream(
     const std::string& table_name,
     const std::vector<std::uint8_t>& schema_ipc_bytes,
@@ -250,17 +260,21 @@ ArrowStream Sdk::create_arrow_stream(
   }
   detail::ResultGuard guard;
   CArrowStreamConfigurationOptions copts = detail::to_c(options);
+  // Validate the table name before allocating `owned` (see the proto path).
+  const char* c_table = detail::checked_c_str(table_name, "table_name");
+  auto* owned =
+      new std::shared_ptr<HeadersProvider>(std::move(headers_provider));
   // Non-empty checked above, so data() is non-null.
   CArrowStream* stream = zerobus_sdk_create_arrow_stream_with_headers_provider(
-      handle_, detail::checked_c_str(table_name, "table_name"),
-      schema_ipc_bytes.data(), schema_ipc_bytes.size(),
-      detail::zerobus_cpp_headers_trampoline, headers_provider.get(), &copts,
-      guard.ptr());
+      handle_, c_table, schema_ipc_bytes.data(), schema_ipc_bytes.size(),
+      detail::zerobus_cpp_headers_trampoline, owned,
+      detail::zerobus_cpp_headers_free, &copts, guard.ptr());
   if (stream == nullptr) {
+    // The FFI already freed `owned` on the failure path; do not delete it here.
     guard.throw_if_error();
     throw ZerobusException("failed to create Arrow stream", false);
   }
-  return ArrowStream(stream, std::move(headers_provider));
+  return ArrowStream(stream, nullptr);
 }
 
 }  // namespace zerobus

cpp/src/stream.cpp

@@ -118,12 +118,11 @@ Stream::~Stream() {
   }
 }
 
-// Move transfers the handle plus the provider and ack callback that outlive it,
-// nulling the source so only one Stream closes/frees it.
+// Move transfers the handle plus the ack callback that outlives it, nulling the
+// source so only one Stream closes/frees it. The headers provider is owned by
+// the FFI, not the Stream, so there is nothing to move for it.
 Stream::Stream(Stream&& other) noexcept
-    : handle_(other.handle_),
-      provider_(std::move(other.provider_)),
-      ack_callback_(std::move(other.ack_callback_)) {
+    : handle_(other.handle_), ack_callback_(std::move(other.ack_callback_)) {
   other.handle_ = nullptr;
 }
 
@@ -137,7 +136,6 @@ Stream& Stream::operator=(Stream&& other) noexcept {
       zerobus_stream_free(handle_);
     }
     handle_ = other.handle_;
-    provider_ = std::move(other.provider_);
     ack_callback_ = std::move(other.ack_callback_);
     other.handle_ = nullptr;
   }

cpp/tests/CMakeLists.txt

@@ -22,6 +22,7 @@ zerobus_add_test(arrow_config_defaults_test)
 zerobus_add_test(arrow_config_convert_test)
 zerobus_add_test(embedded_nul_test)
 zerobus_add_test(ack_callback_test)
+zerobus_add_test(headers_provider_test)
 zerobus_add_test(error_test)
 zerobus_add_test(record_test)
 zerobus_add_test(version_test)

cpp/tests/headers_callback_test.cpp

@@ -9,6 +9,7 @@
 #include <cstddef>
 #include <cstdio>
 #include <map>
+#include <memory>
 #include <stdexcept>
 #include <string>
 
@@ -25,6 +26,17 @@ void fail(const char* msg) {
 
 using zerobus::HeadersProvider;
 
+// The trampoline's user_data is a heap std::shared_ptr<HeadersProvider>*, owned
+// by the FFI (see Sdk::create_*). Mirror that here: run the trampoline against
+// a provider wrapped that way, then release it via the FFI's free trampoline.
+zerobus::CHeaders run_trampoline(std::shared_ptr<HeadersProvider> provider) {
+  auto* owned = new std::shared_ptr<HeadersProvider>(std::move(provider));
+  zerobus::CHeaders out =
+      zerobus::detail::zerobus_cpp_headers_trampoline(owned);
+  zerobus::detail::zerobus_cpp_headers_free(owned);
+  return out;
+}
+
 // Returns a caller-set map.
 class MapProvider : public HeadersProvider {
  public:
@@ -49,10 +61,10 @@ int main() {
 
   // Populated map: key/value pairs round-trip, no error.
   {
-    MapProvider provider;
-    provider.headers = {{"Authorization", "Bearer abc"}, {"X-Custom", "v1"}};
+    auto provider = std::make_shared<MapProvider>();
+    provider->headers = {{"Authorization", "Bearer abc"}, {"X-Custom", "v1"}};
 
-    CHeaders out = zerobus_cpp_headers_trampoline(&provider);
+    CHeaders out = run_trampoline(provider);
     if (out.error_message != nullptr) {
       fail("marshalling valid headers set an error_message");
     }
@@ -77,8 +89,8 @@ int main() {
 
   // Empty map: null array, zero count, no error.
   {
-    MapProvider provider;  // no headers
-    CHeaders out = zerobus_cpp_headers_trampoline(&provider);
+    CHeaders out =
+        run_trampoline(std::make_shared<MapProvider>());  // no headers
     if (out.count != 0 || out.headers != nullptr ||
         out.error_message != nullptr) {
       fail("empty provider did not yield an empty, error-free CHeaders");
@@ -89,8 +101,7 @@ int main() {
   // Throwing provider: message surfaces via error_message, no exception
   // escapes.
   {
-    ThrowingProvider provider;
-    CHeaders out = zerobus_cpp_headers_trampoline(&provider);
+    CHeaders out = run_trampoline(std::make_shared<ThrowingProvider>());
     if (out.error_message == nullptr) {
       fail("throwing provider did not set an error_message");
     } else if (std::string(out.error_message).find("provider boom") ==
@@ -105,9 +116,9 @@ int main() {
 
   // Embedded-NUL header value is rejected, not truncated.
   {
-    MapProvider provider;
-    provider.headers = {{"X-Bad", std::string("a\0b", 3)}};
-    CHeaders out = zerobus_cpp_headers_trampoline(&provider);
+    auto provider = std::make_shared<MapProvider>();
+    provider->headers = {{"X-Bad", std::string("a\0b", 3)}};
+    CHeaders out = run_trampoline(provider);
     if (out.error_message == nullptr) {
       fail("embedded-NUL header value was not rejected");
     }