diff --git a/cpp/CLAUDE.md b/cpp/CLAUDE.md index 221a84ef..713001a4 100644 --- a/cpp/CLAUDE.md +++ b/cpp/CLAUDE.md @@ -93,14 +93,18 @@ instead, configure with `-DZEROBUS_FFI_LIBRARY=` 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` 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*` 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` 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. diff --git a/cpp/NEXT_CHANGELOG.md b/cpp/NEXT_CHANGELOG.md index 8f30b5f5..7fcd290b 100644 --- a/cpp/NEXT_CHANGELOG.md +++ b/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` — 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 / diff --git a/cpp/include/zerobus/arrow_stream.hpp b/cpp/include/zerobus/arrow_stream.hpp index ac820101..e30a0292 100644 --- a/cpp/include/zerobus/arrow_stream.hpp +++ b/cpp/include/zerobus/arrow_stream.hpp @@ -86,11 +86,12 @@ class ArrowStream { private: friend class Sdk; - ArrowStream(CArrowStream* handle, std::shared_ptr provider) - : handle_(handle), provider_(std::move(provider)) {} + ArrowStream(CArrowStream* handle, std::shared_ptr /*unused*/) + : handle_(handle) {} CArrowStream* handle_; - std::shared_ptr provider_; + // The headers provider (if any) is owned by the FFI, not the ArrowStream (see + // Sdk::create_arrow_stream / headers_provider.hpp). }; } // namespace zerobus diff --git a/cpp/include/zerobus/headers_provider.hpp b/cpp/include/zerobus/headers_provider.hpp index ca972d82..d0c5b604 100644 --- a/cpp/include/zerobus/headers_provider.hpp +++ b/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; diff --git a/cpp/include/zerobus/stream.hpp b/cpp/include/zerobus/stream.hpp index 5fc9ac1b..1319e8c2 100644 --- a/cpp/include/zerobus/stream.hpp +++ b/cpp/include/zerobus/stream.hpp @@ -131,15 +131,15 @@ class Stream { private: friend class Sdk; - Stream(CZerobusStream* handle, std::shared_ptr provider, + Stream(CZerobusStream* handle, std::shared_ptr /*unused*/, std::shared_ptr 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 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. diff --git a/cpp/src/arrow_stream.cpp b/cpp/src/arrow_stream.cpp index 45ecf302..d5b35509 100644 --- a/cpp/src/arrow_stream.cpp +++ b/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; diff --git a/cpp/src/detail/headers_callback.hpp b/cpp/src/detail/headers_callback.hpp index 086619d2..29292f2a 100644 --- a/cpp/src/detail/headers_callback.hpp +++ b/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` (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` 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 diff --git a/cpp/src/headers_callback.cpp b/cpp/src/headers_callback.cpp index ffee5aff..428a4ca9 100644 --- a/cpp/src/headers_callback.cpp +++ b/cpp/src/headers_callback.cpp @@ -15,6 +15,7 @@ #include #include +#include #include #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(user_data); + // user_data is a heap-owned shared_ptr (see + // zerobus_cpp_headers_free); deref it to reach the provider. + auto& provider_sp = + *static_cast*>(user_data); + HeadersProvider* provider = provider_sp.get(); + if (provider == nullptr) { + return make_error("null headers provider"); + } std::map 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*>(user_data); +} + } // namespace detail } // namespace zerobus diff --git a/cpp/src/sdk.cpp b/cpp/src/sdk.cpp index 9c5e868c..51a582b0 100644 --- a/cpp/src/sdk.cpp +++ b/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 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(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& 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(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 diff --git a/cpp/src/stream.cpp b/cpp/src/stream.cpp index d7d4a287..766368bb 100644 --- a/cpp/src/stream.cpp +++ b/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; } diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 94b7acdd..c819737e 100644 --- a/cpp/tests/CMakeLists.txt +++ b/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) diff --git a/cpp/tests/headers_callback_test.cpp b/cpp/tests/headers_callback_test.cpp index 8a9f839d..0dcb9562 100644 --- a/cpp/tests/headers_callback_test.cpp +++ b/cpp/tests/headers_callback_test.cpp @@ -9,6 +9,7 @@ #include #include #include +#include #include #include @@ -25,6 +26,17 @@ void fail(const char* msg) { using zerobus::HeadersProvider; +// The trampoline's user_data is a heap std::shared_ptr*, 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 provider) { + auto* owned = new std::shared_ptr(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(); + 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()); // 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()); 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(); + 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"); } diff --git a/cpp/tests/headers_provider_test.cpp b/cpp/tests/headers_provider_test.cpp new file mode 100644 index 00000000..4519de16 --- /dev/null +++ b/cpp/tests/headers_provider_test.cpp @@ -0,0 +1,98 @@ +// Verifies the C++ headers-provider FFI wiring: the trampoline derefs a +// heap-owned shared_ptr and dispatches get_headers(), the free +// trampoline releases that heap shared_ptr (so the provider is destroyed), and +// both trampolines tolerate null user_data. Dependency-free, like the other +// tests: returns non-zero on failure. +// +// This covers the ownership-transfer mechanism in isolation (free-once, correct +// dispatch, null tolerance) — the FFI, not the Stream, owns the provider. The +// live recovery-vs-teardown race that motivates it is driven from the Rust +// core, where the supervisor task holds the provider Arc across the callback. + +#include "zerobus/headers_provider.hpp" + +#include +#include +#include +#include + +#include "detail/headers_callback.hpp" + +namespace { + +int g_failures = 0; + +void check(const char* what, bool ok) { + if (!ok) { + std::fprintf(stderr, "FAIL: %s\n", what); + ++g_failures; + } +} + +// Records whether it was destroyed (via a shared flag) and what it returns. +class RecordingProvider : public zerobus::HeadersProvider { + public: + explicit RecordingProvider(std::shared_ptr destroyed) + : destroyed_(std::move(destroyed)) {} + ~RecordingProvider() override { *destroyed_ = true; } + + std::map get_headers() override { + ++calls; + return {{"authorization", "Bearer test"}}; + } + + int calls = 0; + + private: + std::shared_ptr destroyed_; +}; + +// The trampoline derefs the heap shared_ptr and dispatches get_headers(); the +// free trampoline then destroys the provider. +void test_trampoline_dispatches_and_free_destroys() { + auto destroyed = std::make_shared(false); + auto provider = std::make_shared(destroyed); + RecordingProvider* raw = provider.get(); + + // Mirror Sdk::create_*: hand a heap shared_ptr to the FFI as user_data. + void* user_data = new std::shared_ptr(provider); + provider.reset(); // Only the heap copy keeps it alive now. + check("provider alive while FFI owns it", *destroyed == false); + + zerobus::CHeaders headers = + zerobus::detail::zerobus_cpp_headers_trampoline(user_data); + check("trampoline reported no error", headers.error_message == nullptr); + check("trampoline returned one header", headers.count == 1); + check("get_headers was called", raw->calls == 1); + zerobus::zerobus_free_headers(headers); + + // The FFI's free callback destroys the provider (last owner dropped). + zerobus::detail::zerobus_cpp_headers_free(user_data); + check("free destroyed the provider", *destroyed == true); +} + +// A null user_data must be handled, not dereferenced. +void test_null_user_data() { + zerobus::CHeaders headers = + zerobus::detail::zerobus_cpp_headers_trampoline(nullptr); + check("null user_data => error reported", headers.error_message != nullptr); + zerobus::zerobus_free_headers(headers); + + // Freeing null is a no-op (delete nullptr); reaching the end is the + // assertion. + zerobus::detail::zerobus_cpp_headers_free(nullptr); +} + +} // namespace + +int main() { + test_trampoline_dispatches_and_free_destroys(); + test_null_user_data(); + + if (g_failures != 0) { + std::fprintf(stderr, "%d headers-provider check(s) failed.\n", g_failures); + return 1; + } + std::printf("headers provider wiring OK\n"); + return 0; +} diff --git a/dotnet/NEXT_CHANGELOG.md b/dotnet/NEXT_CHANGELOG.md index 3406a332..533689cf 100644 --- a/dotnet/NEXT_CHANGELOG.md +++ b/dotnet/NEXT_CHANGELOG.md @@ -8,6 +8,14 @@ ### Bug Fixes +- Fixed a use-after-free in which a custom `IHeadersProvider` could be freed + while the Rust core was still inside a `GetHeaders()` call into it during + connection recovery. Provider ownership is now handed to the FFI via the new + `free_user_data` destroy callback, which releases the provider's `GCHandle` + only after any in-flight `GetHeaders()` has returned; the stream no longer + frees the handle on dispose. Tracks the FFI signature change to + `zerobus_sdk_create_stream_with_headers_provider`. No public API change. + ### Documentation ### Internal Changes diff --git a/dotnet/src/Zerobus/IHeadersProvider.cs b/dotnet/src/Zerobus/IHeadersProvider.cs index 07fa0ec0..1dd1e034 100644 --- a/dotnet/src/Zerobus/IHeadersProvider.cs +++ b/dotnet/src/Zerobus/IHeadersProvider.cs @@ -5,6 +5,14 @@ namespace Databricks.Zerobus; /// Implement this interface to supply custom authentication logic /// (e.g. fetching tokens from a vault, using managed identity, etc.). /// +/// +/// Lifetime: the SDK owns the provider for the stream's lifetime and releases it +/// only after any in-flight call (including one during +/// connection recovery) has returned, so you do not need to keep your own +/// reference alive past stream creation. may be invoked +/// from an internal SDK worker thread, not the thread that created the stream, so +/// implementations must be safe to call from any thread. +/// /// /// /// public class CustomHeadersProvider : IHeadersProvider diff --git a/dotnet/src/Zerobus/Native/HeadersProviderBridge.cs b/dotnet/src/Zerobus/Native/HeadersProviderBridge.cs index 557f867a..6c64ada7 100644 --- a/dotnet/src/Zerobus/Native/HeadersProviderBridge.cs +++ b/dotnet/src/Zerobus/Native/HeadersProviderBridge.cs @@ -1,6 +1,7 @@ // Bridges the managed IHeadersProvider interface to the native callback expected by Rust. // This is the .NET equivalent of the goGetHeaders / cHeadersCallback pattern in ffi.go. +using System.Runtime.InteropServices; using System.Text; namespace Databricks.Zerobus.Native; @@ -13,9 +14,49 @@ internal sealed class HeadersProviderBridge { private readonly IHeadersProvider _provider; + // Delegate instances are held here so their native thunks stay alive exactly + // as long as this bridge does. On the ownership-transfer (sync) path the + // bridge is rooted only by the GCHandle handed to the FFI, so rooting the + // delegates through the bridge keeps both callbacks valid until the FFI + // releases that handle via NativeFree. + public HeadersProviderCallback Callback { get; } + public HeadersProviderFreeCallback FreeCallback { get; } + public HeadersProviderBridge(IHeadersProvider provider) { _provider = provider; + Callback = NativeCallback; + FreeCallback = NativeFree; + } + + /// + /// Native destroy callback matching . + /// The FFI owns (a to the + /// bridge) and invokes this exactly once — after any in-flight + /// has returned — so freeing the handle here + /// cannot race a live callback. This is what closes the recovery-vs-teardown + /// use-after-free; it is the .NET equivalent of goFreeHeadersProvider / + /// zerobus_cpp_headers_free. Must not throw across the native boundary. + /// + public static void NativeFree(IntPtr userData) + { + if (userData == IntPtr.Zero) + { + return; + } + + try + { + var handle = GCHandle.FromIntPtr(userData); + if (handle.IsAllocated) + { + handle.Free(); + } + } + catch + { + // Never let an exception unwind across the native boundary. + } } /// diff --git a/dotnet/src/Zerobus/Native/NativeBindings.cs b/dotnet/src/Zerobus/Native/NativeBindings.cs index e445eb8b..1cf6abdd 100644 --- a/dotnet/src/Zerobus/Native/NativeBindings.cs +++ b/dotnet/src/Zerobus/Native/NativeBindings.cs @@ -124,6 +124,19 @@ internal struct CRecordArray [UnmanagedFunctionPointer(CallingConvention.Cdecl)] internal delegate CHeaders HeadersProviderCallback(IntPtr userData); +/// +/// Callback that releases the headers provider's user_data. +/// Matches: void (*HeadersProviderFreeCallback)(void* user_data) +/// +/// +/// The FFI owns user_data and invokes this exactly once, after any in-flight +/// GetHeaders callback has returned — so it is safe to free the provider here. +/// May run on an internal SDK worker thread, so it must not throw across the +/// native boundary. +/// +[UnmanagedFunctionPointer(CallingConvention.Cdecl)] +internal delegate void HeadersProviderFreeCallback(IntPtr userData); + /// /// Callback for async stream creation completion. /// Matches: void (*CreateStreamAsyncCallback)(CZerobusStream* stream, const CResult* result, void* user_data) @@ -226,6 +239,7 @@ public static extern unsafe IntPtr SdkCreateStreamWithHeadersProvider( nuint descriptorProtoLen, HeadersProviderCallback headersCallback, IntPtr userData, + HeadersProviderFreeCallback? freeUserData, ref CStreamConfigurationOptions options, ref CResult result); diff --git a/dotnet/src/Zerobus/Native/NativeInterop.cs b/dotnet/src/Zerobus/Native/NativeInterop.cs index 67ef06e6..362e6b2f 100644 --- a/dotnet/src/Zerobus/Native/NativeInterop.cs +++ b/dotnet/src/Zerobus/Native/NativeInterop.cs @@ -208,6 +208,7 @@ public static unsafe IntPtr SdkCreateStreamWithHeadersProvider( ReadOnlySpan descriptorProto, HeadersProviderCallback callback, IntPtr userData, + HeadersProviderFreeCallback freeUserData, ref CStreamConfigurationOptions options) { var result = new CResult(); @@ -222,6 +223,7 @@ public static unsafe IntPtr SdkCreateStreamWithHeadersProvider( (nuint)descriptorProto.Length, callback, userData, + freeUserData, ref options, ref result); } diff --git a/dotnet/src/Zerobus/ZerobusSdk.cs b/dotnet/src/Zerobus/ZerobusSdk.cs index 712430a0..7c6ead1d 100644 --- a/dotnet/src/Zerobus/ZerobusSdk.cs +++ b/dotnet/src/Zerobus/ZerobusSdk.cs @@ -417,11 +417,16 @@ private ZerobusStream CreateStreamWithHeadersProviderCore( var nativeOpts = NativeInterop.ConvertConfig(options); - // Create the callback bridge that the native code will invoke. + // Create the callback bridge that the native code will invoke. The bridge + // holds both delegate instances, so rooting the bridge keeps their native + // thunks alive. var bridge = new HeadersProviderBridge(headersProvider); - var callback = new HeadersProviderCallback(bridge.NativeCallback); - // GCHandle keeps the bridge + callback alive for the lifetime of the stream. + // Ownership transfer: the GCHandle to the bridge is handed to the FFI as + // user_data. The FFI releases it via the free callback (NativeFree) + // exactly once — after any in-flight GetHeaders has returned — which + // closes the recovery-vs-teardown use-after-free. The stream therefore + // owns nothing and never frees this handle itself. var handle = GCHandle.Alloc(bridge); IntPtr streamPtr; @@ -431,17 +436,21 @@ private ZerobusStream CreateStreamWithHeadersProviderCore( _ptr, tableProperties.TableName, tableProperties.DescriptorProto ?? [], - callback, + bridge.Callback, GCHandle.ToIntPtr(handle), + bridge.FreeCallback, ref nativeOpts); } catch { - handle.Free(); + // On a thrown failure the FFI already invoked the free callback (the + // provider is built before any fallible work), so the handle is + // already freed — do NOT free it again here (double-free). throw; } - return new ZerobusStream(streamPtr, handle, callback); + // The FFI owns the provider handle now; the stream keeps no reference. + return new ZerobusStream(streamPtr); } private async Task CreateStreamWithHeadersProviderCoreAsync( diff --git a/dotnet/src/Zerobus/ZerobusStream.cs b/dotnet/src/Zerobus/ZerobusStream.cs index db7c38f3..16e39ee9 100644 --- a/dotnet/src/Zerobus/ZerobusStream.cs +++ b/dotnet/src/Zerobus/ZerobusStream.cs @@ -35,7 +35,11 @@ public sealed class ZerobusStream : IDisposable, IAsyncDisposable private int _inflightAsyncOperations; private readonly ManualResetEventSlim _asyncOperationsDrained = new(initialState: true); - // Prevent the GCHandle / delegate from being collected while the native code holds a reference. + // Only set on the ASYNC headers-provider path, which still borrows the + // provider: the stream owns the GCHandle / delegate and frees them on + // dispose. The sync path transfers provider ownership to the FFI (freed via + // a destroy callback after any in-flight GetHeaders returns), so those + // streams leave these at default/null and never free the handle here. // Not readonly: GCHandle is not a readonly struct, so calling Free() on a readonly field // creates a defensive copy — the field is never actually mutated, causing double-free. private GCHandle _bridgeHandle; @@ -46,6 +50,7 @@ internal ZerobusStream(IntPtr ptr) _ptr = ptr; } + // Async headers-provider path: the stream owns the provider handle/delegate. internal ZerobusStream(IntPtr ptr, GCHandle bridgeHandle, HeadersProviderCallback callbackRef) { _ptr = ptr; @@ -420,8 +425,10 @@ internal IntPtr NativePointer } /// - /// Attempts to retrieve the bridge handle and callback reference for the stream. - /// Internal use only. + /// Transfers this stream's native pointer and any provider handle to a new + /// wrapper around , disposing this one. On the async + /// headers-provider path the owned bridge handle moves to the new wrapper; the + /// sync path owns no handle, so only the pointer moves. Internal use only. /// internal ZerobusStream Recreate(IntPtr newPtr) { diff --git a/dotnet/tests/Zerobus.Tests/HeadersProviderBridgeTests.cs b/dotnet/tests/Zerobus.Tests/HeadersProviderBridgeTests.cs new file mode 100644 index 00000000..721e414f --- /dev/null +++ b/dotnet/tests/Zerobus.Tests/HeadersProviderBridgeTests.cs @@ -0,0 +1,64 @@ +using System.Runtime.CompilerServices; +using System.Runtime.InteropServices; +using Databricks.Zerobus; +using Databricks.Zerobus.Native; +using NUnit.Framework; + +namespace Databricks.Zerobus.Tests; + +// Verifies the .NET side of the headers-provider ownership transfer: the free +// callback (NativeFree) releases the GCHandle it is handed, and tolerates null. +// This is the mechanism that closes the recovery-vs-teardown use-after-free — +// the FFI now owns the provider handle and releases it via NativeFree after any +// in-flight GetHeaders returns, instead of the stream freeing it on dispose. +// The live race itself is driven from the Rust core (the supervisor holds the +// provider Arc across the callback). Pure-managed, so no native lib is needed. +[TestFixture] +public class HeadersProviderBridgeTests +{ + private sealed class MapProvider : IHeadersProvider + { + public IDictionary GetHeaders() => + new Dictionary(); + } + + [Test] + public void NativeFree_ReleasesHandle() + { + // The GCHandle is the only strong root keeping the bridge alive, mirroring + // ownership transfer to the FFI. After NativeFree releases that root the + // bridge must become collectable. + // + // Note: we cannot assert on `handle.IsAllocated` here — GCHandle is a + // struct wrapping an IntPtr, so NativeFree frees the runtime handle via + // its own reconstructed copy (GCHandle.FromIntPtr) and cannot zero this + // local's field. A WeakReference observes the actual release instead. + var (userData, weak) = AllocBridgeHandle(); + + HeadersProviderBridge.NativeFree(userData); + + GC.Collect(); + GC.WaitForPendingFinalizers(); + GC.Collect(); + + Assert.That(weak.IsAlive, Is.False, + "NativeFree must free the GCHandle so the bridge can be collected"); + } + + // Separate non-inlined method so the bridge has no lingering local/JIT root + // on the caller's frame once it returns. + [MethodImpl(MethodImplOptions.NoInlining)] + private static (IntPtr userData, WeakReference weak) AllocBridgeHandle() + { + var bridge = new HeadersProviderBridge(new MapProvider()); + var handle = GCHandle.Alloc(bridge); + return (GCHandle.ToIntPtr(handle), new WeakReference(bridge)); + } + + [Test] + public void NativeFree_NullUserData_IsNoOp() + { + // Freeing IntPtr.Zero must not throw (mirrors delete nullptr / nil guard). + Assert.DoesNotThrow(() => HeadersProviderBridge.NativeFree(IntPtr.Zero)); + } +} diff --git a/go/CLAUDE.md b/go/CLAUDE.md index 41277c3b..87728155 100644 --- a/go/CLAUDE.md +++ b/go/CLAUDE.md @@ -73,18 +73,23 @@ Go's GC can relocate heap objects. When passing Go memory to C: - **Never remove `runtime.Pinner` calls** — doing so causes "cgo argument has Go pointer to unpinned Go pointer" panics. - Requires Go 1.21+ (the `runtime.Pinner` API). -### Handle registry for callbacks +### Handle ownership for callbacks When using custom `HeadersProvider` (instead of default OAuth): - A `cgo.Handle` wraps the Go interface value and prevents GC collection. -- Handles are stored in `streamHandleRegistry` (mutex-protected map, keyed by stream pointer). -- **Cleanup sequence**: lock registry → delete handle → remove from map → unlock → free C stream. +- The handle is passed to the FFI as `user_data`, and its **ownership is handed to + the FFI** along with a `free_user_data` destroy callback (`goFreeHeadersProvider`). + The FFI invokes it exactly once — after any in-flight `get_headers` returns — and + that is where `handle.Delete()` runs. On a failed create (before the provider is + installed), the Go side deletes the handle itself. +- This replaces the older per-stream handle registry: freeing the handle on `close()` + could race a recovery `get_headers` still running on a worker thread (use-after-free). - Leaking a handle leaks the Go `HeadersProvider` object and any resources it holds. ### Arrow batch cleanup - `zerobus_arrow_free_batch_array()` must be called after reading unacknowledged batches. -- Same handle registry pattern applies for Arrow streams with custom headers. +- Arrow streams with custom headers use the same FFI-owned handle destroy path. ## Breaking change rules @@ -103,7 +108,7 @@ Public API is everything exported (capitalized) in the `go/` package: ## Thread safety -Go SDK is safe for concurrent use from multiple goroutines. Internal synchronization handles concurrent `Ingest` calls. The handle registry uses a mutex. +Go SDK is safe for concurrent use from multiple goroutines. Internal synchronization handles concurrent `Ingest` calls. ## Changelog and documentation diff --git a/go/NEXT_CHANGELOG.md b/go/NEXT_CHANGELOG.md index 8257b586..cf40fbf1 100644 --- a/go/NEXT_CHANGELOG.md +++ b/go/NEXT_CHANGELOG.md @@ -8,6 +8,13 @@ ### Bug Fixes +- Fixed a use-after-free in which a custom `HeadersProvider` could be freed while + a background worker was still calling into it during connection recovery. The + provider's `cgo.Handle` ownership is now handed to the FFI, which releases it + (via a new destroy callback) only after any in-flight `GetHeaders` has + returned, instead of deleting it on stream close. This removes the per-stream + handle registry. No public API change. + ### Documentation - Clarified throughput guidance in the README, godoc, and examples: ingest records in a loop without waiting and call `Flush()` once, rather than calling `WaitForOffset()` after every record. Documented that the ack watermark is monotonic, so waiting on the last offset confirms all prior records. diff --git a/go/arrow_ffi.go b/go/arrow_ffi.go index bbf15651..ff4214dc 100644 --- a/go/arrow_ffi.go +++ b/go/arrow_ffi.go @@ -36,6 +36,7 @@ typedef struct CHeaders { } CHeaders; typedef CHeaders (*HeadersProviderCallback)(void *user_data); +typedef void (*HeadersProviderFreeCallback)(void *user_data); // Arrow stream opaque handle. typedef struct CArrowStream { uint8_t _private[0]; } CArrowStream; @@ -62,8 +63,9 @@ typedef struct CArrowBatchArray { uintptr_t count; } CArrowBatchArray; -// Forward declare the Go-exported headers callback (defined in ffi.go). +// Forward declare the Go-exported callbacks (defined in ffi.go). extern void goGetHeaders(void *userData, CHeader **headers, uintptr_t *count, char **error); +extern void goFreeHeadersProvider(void *userData); // C callback wrapper for arrow streams — calls the same Go function as the // regular stream callback. We define it here as a static so it exists in this @@ -84,6 +86,16 @@ static HeadersProviderCallback getArrowHeadersCallback() { return (HeadersProviderCallback)cArrowHeadersCallback; } +// Same as ffi.go's cFreeHeadersProvider; a distinct static keeps this +// translation unit self-contained. +static void cArrowFreeHeadersProvider(void *userData) { + goFreeHeadersProvider(userData); +} + +static HeadersProviderFreeCallback getArrowFreeHeadersProviderCallback() { + return (HeadersProviderFreeCallback)cArrowFreeHeadersProvider; +} + // Arrow FFI function declarations. extern CArrowStream *zerobus_sdk_create_arrow_stream( CZerobusSdk *sdk, @@ -102,6 +114,7 @@ extern CArrowStream *zerobus_sdk_create_arrow_stream_with_headers_provider( uintptr_t schema_ipc_len, HeadersProviderCallback headers_callback, void *user_data, + HeadersProviderFreeCallback free_user_data, const CArrowStreamConfigurationOptions *options, CResult *result); @@ -122,17 +135,9 @@ import "C" import ( "runtime" "runtime/cgo" - "sync" "unsafe" ) -// arrowStreamHandleRegistry holds cgo.Handles for arrow streams created with a -// custom headers provider so the Go object stays alive for the stream lifetime. -var ( - arrowStreamHandleRegistry = make(map[unsafe.Pointer]cgo.Handle) - arrowStreamHandleRegistryMu sync.Mutex -) - // arrowFfiResult converts a C.CResult (from this translation unit) to a Go error. func arrowFfiResult(cres C.CResult) error { if cres.success { @@ -279,6 +284,8 @@ func sdkCreateArrowStreamWithHeadersProvider( cSchema := (*C.uint8_t)(unsafe.SliceData(schemaIpcBytes)) pinner.Pin(cSchema) + // Hand the provider handle's ownership to the FFI; it is released via + // goFreeHeadersProvider once the core is done with it (see ffi.go). handle := cgo.NewHandle(headersProvider) handlePtr := *(*unsafe.Pointer)(unsafe.Pointer(&handle)) @@ -291,32 +298,26 @@ func sdkCreateArrowStreamWithHeadersProvider( C.uintptr_t(len(schemaIpcBytes)), C.getArrowHeadersCallback(), handlePtr, + C.getArrowFreeHeadersProviderCallback(), &cOpts, &cres, ) if ptr == nil { - handle.Delete() + // The FFI owns the handle once create is called and releases it via + // goFreeHeadersProvider on every path — including this failure — so we + // must NOT delete it here (that would double-free the handle). return nil, arrowFfiResult(cres) } - arrowStreamHandleRegistryMu.Lock() - arrowStreamHandleRegistry[unsafe.Pointer(ptr)] = handle - arrowStreamHandleRegistryMu.Unlock() - return unsafe.Pointer(ptr), nil } -// arrowStreamFree frees an Arrow Flight stream and its associated handle. +// arrowStreamFree frees an Arrow Flight stream. The provider handle (if any) is +// released by the FFI via goFreeHeadersProvider, not here. func arrowStreamFree(ptr unsafe.Pointer) { if ptr == nil { return } - arrowStreamHandleRegistryMu.Lock() - if handle, ok := arrowStreamHandleRegistry[ptr]; ok { - handle.Delete() - delete(arrowStreamHandleRegistry, ptr) - } - arrowStreamHandleRegistryMu.Unlock() C.zerobus_arrow_stream_free((*C.CArrowStream)(ptr)) } diff --git a/go/ffi.go b/go/ffi.go index 5dea30e5..b7c2802f 100644 --- a/go/ffi.go +++ b/go/ffi.go @@ -88,6 +88,10 @@ extern CZerobusStream* zerobus_sdk_create_stream(CZerobusSdk* sdk, const char* client_secret, const CStreamConfigurationOptions* options, CResult* result); +// Releases the headers provider user_data. The FFI owns user_data and invokes +// this once, after any in-flight get_headers callback has returned. +typedef void (*HeadersProviderFreeCallback)(void* user_data); + extern CZerobusStream* zerobus_sdk_create_stream_with_headers_provider( CZerobusSdk* sdk, const char* table_name, @@ -95,6 +99,7 @@ extern CZerobusStream* zerobus_sdk_create_stream_with_headers_provider( uintptr_t descriptor_proto_len, HeadersProviderCallback headers_callback, void* user_data, + HeadersProviderFreeCallback free_user_data, const CStreamConfigurationOptions* options, CResult* result); extern void zerobus_stream_free(CZerobusStream* stream); @@ -140,8 +145,9 @@ extern bool zerobus_stream_close(CZerobusStream* stream, CResult* result); extern void zerobus_free_error_message(char* error_message); extern CStreamConfigurationOptions zerobus_get_default_config(); -// Forward declaration of Go function +// Forward declaration of Go functions extern void goGetHeaders(void* userData, CHeader** headers, uintptr_t* count, char** error); +extern void goFreeHeadersProvider(void* userData); // C callback that matches the HeadersProviderCallback signature static CHeaders cHeadersCallback(void* userData) { @@ -163,22 +169,24 @@ static CHeaders cHeadersCallback(void* userData) { static HeadersProviderCallback getHeadersCallback() { return (HeadersProviderCallback)cHeadersCallback; } + +// C callback matching HeadersProviderFreeCallback; deletes the cgo.Handle so the +// Go provider is released once the Rust core is done with it. +static void cFreeHeadersProvider(void* userData) { + goFreeHeadersProvider(userData); +} + +static HeadersProviderFreeCallback getFreeHeadersProviderCallback() { + return (HeadersProviderFreeCallback)cFreeHeadersProvider; +} */ import "C" import ( "runtime" "runtime/cgo" - "sync" "unsafe" ) -// Registry to map stream pointers to their handles for cleanup -// This allows us to properly release cgo.Handle when streams are freed -var ( - streamHandleRegistry = make(map[unsafe.Pointer]cgo.Handle) - streamHandleRegistryMu sync.Mutex -) - // ffiResult converts a C.CResult to a Go error func ffiResult(cres C.CResult) error { if cres.success { @@ -415,6 +423,21 @@ func goGetHeaders(userData unsafe.Pointer, headers **C.CHeader, count *C.uintptr *errorMsg = nil } +//export goFreeHeadersProvider +func goFreeHeadersProvider(userData unsafe.Pointer) { + // The FFI owns the provider handle and calls this once, after any in-flight + // get_headers has returned. Delete the cgo.Handle to release the Go object. + // + // A zero handle can't occur on the ownership-transfer path (cgo.NewHandle + // never returns 0, and we only pass real handles to the FFI), but guard + // anyway: Delete() panics on a zero/invalid handle, which would unwind + // across the C boundary. + if userData == nil { + return + } + cgo.Handle(userData).Delete() +} + // sdkCreateStreamWithHeadersProvider creates a stream with custom headers provider via FFI func sdkCreateStreamWithHeadersProvider( sdkPtr unsafe.Pointer, @@ -437,8 +460,10 @@ func sdkCreateStreamWithHeadersProvider( descriptorLen = C.size_t(len(descriptorProto)) } - // Create a cgo.Handle for the provider - // This keeps it alive and gives us a safe uintptr to pass to C + // Create a cgo.Handle for the provider and hand its ownership to the FFI. + // The FFI releases it via goFreeHeadersProvider once the Rust core is done + // with the provider — after any in-flight get_headers callback returns — + // which closes the recovery-vs-teardown use-after-free window. handle := cgo.NewHandle(headersProvider) // Convert handle to unsafe.Pointer using a pattern the linter accepts handlePtr := *(*unsafe.Pointer)(unsafe.Pointer(&handle)) @@ -453,35 +478,26 @@ func sdkCreateStreamWithHeadersProvider( descriptorLen, C.getHeadersCallback(), handlePtr, + C.getFreeHeadersProviderCallback(), &cOpts, &cres, ) if ptr == nil { - // Clean up handle on error - handle.Delete() + // The FFI owns the handle once create is called and releases it via + // goFreeHeadersProvider on every path — including this failure — so we + // must NOT delete it here (that would double-free the handle). return nil, ffiResult(cres) } - // Store the handle so we can clean it up when the stream is freed - streamHandleRegistryMu.Lock() - streamHandleRegistry[unsafe.Pointer(ptr)] = handle - streamHandleRegistryMu.Unlock() - return unsafe.Pointer(ptr), nil } -// streamFree frees a stream instance +// streamFree frees a stream instance. The provider handle (if any) is released +// by the FFI via goFreeHeadersProvider, not here, so there is no registry to +// clean up. func streamFree(ptr unsafe.Pointer) { if ptr != nil { - // Clean up the associated handle BEFORE freeing the stream - streamHandleRegistryMu.Lock() - if handle, exists := streamHandleRegistry[ptr]; exists { - handle.Delete() // This releases the Go object reference - delete(streamHandleRegistry, ptr) - } - streamHandleRegistryMu.Unlock() - C.zerobus_stream_free((*C.CZerobusStream)(ptr)) } } diff --git a/go/ffi_test.go b/go/ffi_test.go index e49c4473..01bf77a6 100644 --- a/go/ffi_test.go +++ b/go/ffi_test.go @@ -6,113 +6,28 @@ import ( "unsafe" ) -// TestStreamHandleRegistry tests the stream handle registry -func TestStreamHandleRegistry(t *testing.T) { - // Create a test handle - testProvider := &mockHeadersProvider{} - handle := cgo.NewHandle(testProvider) - - // Create a real allocated pointer (not a fake one) - dummyStream := struct{ id int }{1234} - dummyStreamPtr := unsafe.Pointer(&dummyStream) - - // Store in registry - streamHandleRegistryMu.Lock() - streamHandleRegistry[dummyStreamPtr] = handle - streamHandleRegistryMu.Unlock() - - // Verify it's stored - streamHandleRegistryMu.Lock() - storedHandle, exists := streamHandleRegistry[dummyStreamPtr] - streamHandleRegistryMu.Unlock() - - if !exists { - t.Fatal("Handle not found in registry") - } - - if storedHandle != handle { - t.Fatal("Retrieved handle doesn't match stored handle") - } - - // Clean up - streamHandleRegistryMu.Lock() - delete(streamHandleRegistry, dummyStreamPtr) - streamHandleRegistryMu.Unlock() - handle.Delete() -} - -// TestStreamHandleCleanup tests that handles are properly cleaned up -func TestStreamHandleCleanup(t *testing.T) { - testProvider := &mockHeadersProvider{} - handle := cgo.NewHandle(testProvider) - - dummyStream := struct{ id int }{5678} - dummyStreamPtr := unsafe.Pointer(&dummyStream) - - // Store in registry - streamHandleRegistryMu.Lock() - streamHandleRegistry[dummyStreamPtr] = handle - streamHandleRegistryMu.Unlock() - - // Simulate streamFree cleanup logic - streamHandleRegistryMu.Lock() - if h, exists := streamHandleRegistry[dummyStreamPtr]; exists { - h.Delete() - delete(streamHandleRegistry, dummyStreamPtr) - } - streamHandleRegistryMu.Unlock() - - // Verify it's removed - streamHandleRegistryMu.Lock() - _, exists := streamHandleRegistry[dummyStreamPtr] - streamHandleRegistryMu.Unlock() - - if exists { - t.Fatal("Handle should have been removed from registry") - } -} - -// TestHandleConcurrency tests concurrent access to the handle registry -func TestHandleConcurrency(t *testing.T) { - const numGoroutines = 10 - - done := make(chan bool, numGoroutines) - - for i := 0; i < numGoroutines; i++ { - go func(id int) { - testProvider := &mockHeadersProvider{} - handle := cgo.NewHandle(testProvider) - dummyStream := struct{ id int }{1000 + id} - ptr := unsafe.Pointer(&dummyStream) - - // Store - streamHandleRegistryMu.Lock() - streamHandleRegistry[ptr] = handle - streamHandleRegistryMu.Unlock() - - // Retrieve - streamHandleRegistryMu.Lock() - _, exists := streamHandleRegistry[ptr] - streamHandleRegistryMu.Unlock() - - if !exists { - t.Errorf("Handle %d not found", id) - } - - // Clean up - streamHandleRegistryMu.Lock() - delete(streamHandleRegistry, ptr) - streamHandleRegistryMu.Unlock() - handle.Delete() - - done <- true - }(i) - } - - // Wait for all goroutines - for i := 0; i < numGoroutines; i++ { - <-done - } +// TestGoFreeHeadersProviderReleasesHandle verifies the FFI-owned destroy path: +// goFreeHeadersProvider must delete the cgo.Handle it is handed, releasing the +// Go provider. This is the Go side of the ownership transfer that closes the +// recovery-vs-teardown use-after-free. +func TestGoFreeHeadersProviderReleasesHandle(t *testing.T) { + handle := cgo.NewHandle(&mockHeadersProvider{}) + handlePtr := *(*unsafe.Pointer)(unsafe.Pointer(&handle)) + + // The handle resolves to the provider while live. + if _, ok := handle.Value().(HeadersProvider); !ok { + t.Fatal("handle should resolve to the provider before free") + } + + // The destroy callback releases it; the handle must no longer be valid. + goFreeHeadersProvider(handlePtr) + + defer func() { + if r := recover(); r == nil { + t.Fatal("expected handle.Value() to panic after goFreeHeadersProvider") + } + }() + _ = handle.Value() // panics: handle was deleted } // Mock HeadersProvider for testing diff --git a/go/zerobus.go b/go/zerobus.go index 7534033f..8d796fc7 100644 --- a/go/zerobus.go +++ b/go/zerobus.go @@ -249,6 +249,13 @@ func (s *ZerobusSdk) CreateStream( // "x-databricks-zerobus-table-name": "catalog.schema.table", // }, nil // } +// +// The SDK owns the provider for the stream's lifetime and releases it only +// after any in-flight GetHeaders call (including one during connection +// recovery) has returned — so a slow GetHeaders racing stream teardown is +// never invoked on a released provider. GetHeaders may be called from an +// internal SDK worker thread, so implementations must be safe to use from a +// goroutine other than the one that created the stream. type HeadersProvider interface { // GetHeaders returns the headers to be used for authentication. // This method will be called by the SDK when authentication is needed. diff --git a/rust/ffi/NEXT_CHANGELOG.md b/rust/ffi/NEXT_CHANGELOG.md index 6b975212..5b61fa5d 100644 --- a/rust/ffi/NEXT_CHANGELOG.md +++ b/rust/ffi/NEXT_CHANGELOG.md @@ -10,10 +10,23 @@ ### Bug Fixes +- Fixed a use-after-free in which a custom headers provider could be freed by the + wrapper while a Rust worker thread was still inside a synchronous + `get_headers()` callback into it during connection recovery. The FFI now takes + ownership of the provider `user_data`: `CallbackHeadersProvider` invokes a + caller-supplied `free_user_data` destroy callback from its `Drop`, which runs + only after every task that could call `get_headers()` is gone (the supervisor + task holds its own `Arc` across the in-flight call). The provider is + constructed before any fallible work in `create_stream_with_headers_provider` / + `create_arrow_stream_with_headers_provider`, so `free_user_data` is invoked + exactly once on every path — on success after the last reference drops, on a + failed create before returning. + ### Documentation ### Internal Changes +- Add headers-provider ownership tests: `free_user_data` fires once on `Drop` and once on failed create, a null destroy callback is a no-op, and a teardown test that reproduces the recovery-vs-teardown race (a blocking in-flight `get_headers` on one `Arc` clone while another is dropped) asserts the free is deferred until the callback returns. Test-only; no ABI or behavior change. - Add ack-callback live-teardown / use-after-free tests. They drive the real `CallbackAckCallback` bridge over a heap-allocated `user_data` through the real callback-handler task, then tear it down via the production teardown code, asserting no callback fires after teardown returns and that `user_data` is safe to release at that point. All teardown paths are covered: drain-within-`callback_max_wait_time_ms`, wait-indefinitely, and a callback still synchronously in-flight when a bounded budget expires — which the drain aborts but cannot preempt, so the callback outlives `teardown()` and `user_data` must stay alive until it finishes. Test-only; no ABI or behavior change. - Add tests asserting the ack-callback `Arc` is released, not leaked to a background task, when `zerobus_sdk_create_stream` / `zerobus_sdk_create_stream_with_headers_provider` fails. Uses a `#[cfg(test)]` drop hook keyed to a test-only sentinel `user_data`. Test-only; no ABI or behavior change. - Move the header-callback helpers `zerobus_alloc_header_array`, `zerobus_alloc_cstring`, and `zerobus_free_headers` from `arrow.rs` into the shared `common.rs` module, next to the `CHeader`/`CHeaders` types they serve. No behavior, signature, or ABI change; their declarations move ahead of the Arrow functions in the generated `zerobus.h`. @@ -22,6 +35,14 @@ ### Breaking Changes +- `zerobus_sdk_create_stream_with_headers_provider` and + `zerobus_sdk_create_arrow_stream_with_headers_provider` take a new + `free_user_data` parameter (a nullable `void (*)(void *user_data)`) after + `user_data`. Callers must hand ownership of `user_data` across and supply a + destroy callback (or pass null to opt out and manage the lifetime themselves). + This changes the generated `zerobus.h` signatures, so Go and any other C FFI + consumer must update their call sites. + ### Deprecations ### API Changes diff --git a/rust/ffi/src/arrow.rs b/rust/ffi/src/arrow.rs index 8727dc7e..690464cc 100644 --- a/rust/ffi/src/arrow.rs +++ b/rust/ffi/src/arrow.rs @@ -207,6 +207,12 @@ pub extern "C" fn zerobus_sdk_create_arrow_stream( /// Creates an Arrow Flight stream with a custom headers provider callback. /// /// `schema_ipc_bytes` must point to Arrow IPC stream bytes encoding only the schema. +/// +/// Ownership of `user_data` / `free_user_data` follows +/// `zerobus_sdk_create_stream_with_headers_provider` — once called the FFI owns +/// `user_data` and invokes `free_user_data` exactly once on every path (on +/// success after any in-flight `get_headers` returns; on failure before +/// returning null). The caller must never free `user_data` itself. #[no_mangle] pub extern "C" fn zerobus_sdk_create_arrow_stream_with_headers_provider( sdk: *mut CZerobusSdk, @@ -215,10 +221,24 @@ pub extern "C" fn zerobus_sdk_create_arrow_stream_with_headers_provider( schema_ipc_len: usize, headers_callback: HeadersProviderCallback, user_data: *mut std::ffi::c_void, + // Written inline (not via HeadersProviderFreeCallback) so cbindgen emits a + // nullable C function pointer instead of an opaque struct. + free_user_data: Option, options: *const CArrowStreamConfigurationOptions, result: *mut CResult, ) -> *mut CArrowStream { ffi_guard(result, ptr::null_mut(), move || { + // INVARIANT: construct the provider Arc *before* any fallible work (see + // the proto path in stream.rs). It owns `user_data`, so its Drop invokes + // `free_user_data` exactly once on every path — the free-once contract + // the wrappers rely on. Moving this after an early-return would leak + // `user_data`, since the wrappers do not free on failure. + let headers_provider: Arc = Arc::new(CallbackHeadersProvider::new( + headers_callback, + user_data, + free_user_data, + )); + let sdk_ref = match validate_sdk_ptr(sdk) { Ok(s) => s, Err(msg) => { @@ -237,9 +257,6 @@ pub extern "C" fn zerobus_sdk_create_arrow_stream_with_headers_provider( unsafe { std::slice::from_raw_parts(schema_ipc_bytes, schema_ipc_len) }; let schema = ipc_bytes_to_schema(schema_bytes).map_err(|e| e.to_string())?; - let headers_provider: Arc = - Arc::new(CallbackHeadersProvider::new(headers_callback, user_data)); - let mut builder = sdk_ref .stream_builder() .table(table_name_str) diff --git a/rust/ffi/src/common.rs b/rust/ffi/src/common.rs index 0883d2b5..9eb710ca 100644 --- a/rust/ffi/src/common.rs +++ b/rust/ffi/src/common.rs @@ -207,6 +207,19 @@ pub struct CHeaders { /// The caller is responsible for freeing the returned CHeaders using zerobus_free_headers pub type HeadersProviderCallback = extern "C" fn(user_data: *mut std::ffi::c_void) -> CHeaders; +/// Function pointer type for releasing the headers provider's `user_data`. +/// +/// The FFI owns `user_data`: this is invoked exactly once, when the last `Arc` +/// referencing the provider drops. Because the supervisor task holds its own +/// `Arc` clone across an in-flight `get_headers` call, that drop is guaranteed +/// to happen only after any synchronous callback has returned — closing the +/// use-after-free window that a wrapper freeing `user_data` right after +/// `close()` would otherwise leave open. +/// +/// May run on an internal SDK thread (a tokio worker, not necessarily the +/// thread that created the stream), so it must be safe to call from any thread. +pub type HeadersProviderFreeCallback = extern "C" fn(user_data: *mut std::ffi::c_void); + /// Function pointer type for async stream creation completion. /// /// `stream` is non-null on success and null on failure. `result` points to a @@ -244,23 +257,46 @@ pub type BoolAsyncCallback = pub type RecordArrayAsyncCallback = extern "C" fn(records: CRecordArray, result: *const CResult, user_data: *mut std::ffi::c_void); -/// Rust struct that wraps a Go callback and implements HeadersProvider +/// Rust struct that wraps a Go/C callback and implements HeadersProvider. +/// +/// Owns `user_data` when `free_user_data` is set: the destroy callback fires +/// from `Drop`, i.e. once every task that could call `get_headers` is gone. pub(crate) struct CallbackHeadersProvider { callback: HeadersProviderCallback, user_data: *mut std::ffi::c_void, + free_user_data: Option, in_use: AtomicBool, // Track concurrent access to detect thread-safety issues } impl CallbackHeadersProvider { - pub(crate) fn new(callback: HeadersProviderCallback, user_data: *mut std::ffi::c_void) -> Self { + pub(crate) fn new( + callback: HeadersProviderCallback, + user_data: *mut std::ffi::c_void, + free_user_data: Option, + ) -> Self { Self { callback, user_data, + free_user_data, in_use: AtomicBool::new(false), } } } +impl Drop for CallbackHeadersProvider { + fn drop(&mut self) { + if let Some(free) = self.free_user_data { + let user_data = self.user_data; + // Contain panics: unwinding across the FFI boundary is UB. + if catch_unwind(AssertUnwindSafe(|| free(user_data))).is_err() { + tracing::error!( + "headers provider free callback panicked; contained at FFI boundary" + ); + } + } + } +} + // Safety: We assume the Go callback is thread-safe, but we validate at runtime unsafe impl Send for CallbackHeadersProvider {} unsafe impl Sync for CallbackHeadersProvider {} diff --git a/rust/ffi/src/stream.rs b/rust/ffi/src/stream.rs index c8b7a0b5..5354214c 100644 --- a/rust/ffi/src/stream.rs +++ b/rust/ffi/src/stream.rs @@ -72,8 +72,12 @@ enum StreamCreateAuth { client_secret: String, }, HeadersProvider { - headers_callback: HeadersProviderCallback, - user_data: SendPtr, + // A fully-constructed provider. It is built by the FFI entry point + // *before* any fallible work so that it owns `user_data`/`free_user_data` + // and its Drop frees them on every error path (the UAF fix). Carrying + // the Arc here — rather than the raw callback + pointer — keeps that + // guarantee while still routing through the shared builder below. + headers_provider: Arc, }, } @@ -97,18 +101,10 @@ async fn build_stream_from_parts( .stream_builder() .table(table_name) .oauth(client_id, client_secret), - StreamCreateAuth::HeadersProvider { - headers_callback, - user_data, - } => { - let headers_provider: Arc = Arc::new( - CallbackHeadersProvider::new(headers_callback, user_data.get()), - ); - sdk_ref - .stream_builder() - .table(table_name) - .headers_provider(headers_provider) - } + StreamCreateAuth::HeadersProvider { headers_provider } => sdk_ref + .stream_builder() + .table(table_name) + .headers_provider(headers_provider), }; let mut builder = match record_type { @@ -483,7 +479,22 @@ pub extern "C" fn zerobus_sdk_create_stream_async( } /// Create a stream with a custom headers provider callback -/// This allows you to provide custom authentication headers via a Go callback function +/// This allows you to provide custom authentication headers via a Go/C callback function. +/// +/// Ownership: once this function is called, the FFI owns `user_data`. When +/// `free_user_data` is set it is invoked exactly once, on every path: +/// - on success, when the last internal reference to the provider drops — after +/// any in-flight `get_headers` callback has returned (this is what closes the +/// recovery-vs-teardown use-after-free); +/// - on failure (this call returns null), before returning. +/// +/// The caller must therefore hand ownership across and never free `user_data` +/// itself, not even when create fails. Pass a null `free_user_data` to opt out +/// (the caller then owns `user_data` and must keep it alive for the stream's +/// whole lifetime, including in-flight recovery callbacks). +/// +/// `free_user_data` may run on an internal SDK thread, so it must be safe to +/// call from any thread. #[no_mangle] pub extern "C" fn zerobus_sdk_create_stream_with_headers_provider( sdk: *mut CZerobusSdk, @@ -492,10 +503,27 @@ pub extern "C" fn zerobus_sdk_create_stream_with_headers_provider( descriptor_proto_len: usize, headers_callback: HeadersProviderCallback, user_data: *mut std::ffi::c_void, + // Written inline (not via HeadersProviderFreeCallback) so cbindgen emits a + // nullable C function pointer instead of an opaque struct. + free_user_data: Option, options: *const CStreamConfigurationOptions, result: *mut CResult, ) -> *mut CZerobusStream { ffi_guard(result, ptr::null_mut(), move || { + // INVARIANT: construct the provider Arc *before* any fallible work. It + // owns `user_data`, so on every error path below its Drop invokes + // `free_user_data` exactly once. This gives the caller one contract — + // the FFI always frees `user_data` once create was called (success or + // failure) — so wrappers must never free it themselves. Moving this + // after a `?`/early-return would leak on that path; the wrappers rely on + // free-on-failure and do not compensate. Guarded by + // `test_create_stream_with_headers_provider_frees_user_data_on_failure`. + let headers_provider: Arc = Arc::new(CallbackHeadersProvider::new( + headers_callback, + user_data, + free_user_data, + )); + let sdk_ref = match validate_sdk_ptr(sdk) { Ok(s) => s, Err(msg) => { @@ -533,10 +561,10 @@ pub extern "C" fn zerobus_sdk_create_stream_with_headers_provider( sdk_ref, table_name_str, descriptor_proto, - StreamCreateAuth::HeadersProvider { - headers_callback, - user_data: SendPtr::new(user_data), - }, + // Pass the pre-built provider (constructed above, before any + // fallible work) so its Drop still frees `user_data` on every + // error path in here. + StreamCreateAuth::HeadersProvider { headers_provider }, c_opts, ) .await @@ -565,6 +593,15 @@ pub extern "C" fn zerobus_sdk_create_stream_with_headers_provider( /// callback is invoked exactly once with either a non-null stream pointer and a /// success result, or a null stream pointer and a failure result. The SDK /// handle must remain valid until the callback runs. +/// +/// OWNERSHIP: unlike the synchronous `zerobus_sdk_create_stream_with_headers_provider`, +/// this variant exposes no `free_user_data` callback, so it does NOT take +/// ownership of `user_data` — the provider is built with `free_user_data = None` +/// and the caller stays responsible for freeing `user_data` (only after the +/// completion callback has fired, since a recovery `get_headers` may still run). +/// No wrapper currently calls this; before one adopts it, add a `free_user_data` +/// parameter and hand ownership across (as the sync path does), or it will +/// reintroduce the recovery-vs-teardown use-after-free this change fixes. #[no_mangle] pub extern "C" fn zerobus_sdk_create_stream_with_headers_provider_async( sdk: *mut CZerobusSdk, @@ -612,8 +649,17 @@ pub extern "C" fn zerobus_sdk_create_stream_with_headers_provider_async( None }; + // Build the provider before spawning so it owns `user_data` for the whole + // async attempt. This entry point does not expose a `free_user_data` + // callback, so it opts out of ownership (None) — the caller retains + // responsibility for freeing `user_data` after the completion callback. + let headers_provider: Arc = Arc::new(CallbackHeadersProvider::new( + headers_callback, + user_data, + None, + )); + let sdk_ptr = SendPtr::new(sdk); - let stream_user_data = SendPtr::new(user_data); let callback_user_data = SendPtr::new(callback_user_data); RUNTIME.spawn(async move { let callback_result = match validate_sdk_ptr(sdk_ptr.get()) { @@ -621,10 +667,7 @@ pub extern "C" fn zerobus_sdk_create_stream_with_headers_provider_async( sdk_ref, table_name_str, descriptor_proto, - StreamCreateAuth::HeadersProvider { - headers_callback, - user_data: stream_user_data, - }, + StreamCreateAuth::HeadersProvider { headers_provider }, c_opts, ) .await diff --git a/rust/ffi/src/tests.rs b/rust/ffi/src/tests.rs index c9382ab5..3a1f7e98 100644 --- a/rust/ffi/src/tests.rs +++ b/rust/ffi/src/tests.rs @@ -881,7 +881,7 @@ mod tests { } } - let provider = CallbackHeadersProvider::new(test_callback, ptr::null_mut()); + let provider = CallbackHeadersProvider::new(test_callback, ptr::null_mut(), None); // Sequential calls should work fine let rt = tokio::runtime::Runtime::new().unwrap(); @@ -911,7 +911,7 @@ mod tests { } } - let provider = CallbackHeadersProvider::new(test_callback, ptr::null_mut()); + let provider = CallbackHeadersProvider::new(test_callback, ptr::null_mut(), None); let rt = tokio::runtime::Runtime::new().unwrap(); let result = rt.block_on(provider.get_headers()); @@ -922,6 +922,133 @@ mod tests { assert!(headers.contains_key("Authorization")); } + // The provider owns `user_data`: dropping it must invoke `free_user_data` + // exactly once. This is the mechanism that closes the recovery-vs-teardown + // UAF — the destroy fires from Drop, i.e. after every task that could call + // get_headers is gone, not when the wrapper's close() returns. + #[test] + fn test_headers_provider_free_user_data_called_on_drop() { + static FREE_COUNT: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0); + static FREE_SENTINEL: u8 = 0; + + extern "C" fn callback(_user_data: *mut std::ffi::c_void) -> CHeaders { + CHeaders { + headers: ptr::null_mut(), + count: 0, + error_message: ptr::null_mut(), + } + } + extern "C" fn free_user_data(user_data: *mut std::ffi::c_void) { + // Only count the sentinel we own, so nothing else perturbs the count. + if std::ptr::eq(user_data as *const u8, &FREE_SENTINEL) { + FREE_COUNT.fetch_add(1, AtomicOrdering::SeqCst); + } + } + + let user_data = (&FREE_SENTINEL as *const u8) as *mut std::ffi::c_void; + let before = FREE_COUNT.load(AtomicOrdering::SeqCst); + { + let provider = CallbackHeadersProvider::new(callback, user_data, Some(free_user_data)); + assert_eq!( + FREE_COUNT.load(AtomicOrdering::SeqCst), + before, + "free must not run while the provider is alive" + ); + drop(provider); + } + assert_eq!( + FREE_COUNT.load(AtomicOrdering::SeqCst) - before, + 1, + "free_user_data must run exactly once when the provider drops" + ); + } + + // Reproduces the recovery-vs-teardown race the fix targets: one Arc clone + // (the supervisor task) is inside a blocking synchronous get_headers while + // another Arc (the stream's struct field) is dropped by teardown. free must + // NOT fire while the callback is in flight, and must fire exactly once once + // the in-flight clone finally drops. If free ran on the teardown drop, the + // callback would be touching freed user_data — the original UAF. + #[test] + fn test_headers_provider_free_deferred_until_in_flight_callback_returns() { + use std::sync::atomic::AtomicBool; + use std::sync::Arc; + + static ENTERED: AtomicBool = AtomicBool::new(false); + static RELEASE: AtomicBool = AtomicBool::new(false); + static FREE_COUNT: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0); + static SENTINEL: u8 = 0; + + // Blocks in-flight until the main thread releases it, mirroring a slow + // get_headers on a worker thread during recovery. + extern "C" fn blocking_callback(_user_data: *mut std::ffi::c_void) -> CHeaders { + ENTERED.store(true, AtomicOrdering::SeqCst); + while !RELEASE.load(AtomicOrdering::SeqCst) { + std::thread::yield_now(); + } + CHeaders { + headers: ptr::null_mut(), + count: 0, + error_message: ptr::null_mut(), + } + } + extern "C" fn free_user_data(user_data: *mut std::ffi::c_void) { + if std::ptr::eq(user_data as *const u8, &SENTINEL) { + FREE_COUNT.fetch_add(1, AtomicOrdering::SeqCst); + } + } + + let user_data = (&SENTINEL as *const u8) as *mut std::ffi::c_void; + let provider = Arc::new(CallbackHeadersProvider::new( + blocking_callback, + user_data, + Some(free_user_data), + )); + let in_flight = Arc::clone(&provider); // the "supervisor" clone + + let worker = std::thread::spawn(move || { + let rt = tokio::runtime::Runtime::new().unwrap(); + let _ = rt.block_on(in_flight.get_headers()); + // in_flight drops here, after the callback returns. + }); + + // Wait until the callback is in flight, then simulate teardown dropping + // the struct-field Arc while the worker still holds its clone. + while !ENTERED.load(AtomicOrdering::SeqCst) { + std::thread::yield_now(); + } + drop(provider); + assert_eq!( + FREE_COUNT.load(AtomicOrdering::SeqCst), + 0, + "free must not run while a get_headers callback is in flight" + ); + + // Let the callback return; the last Arc now drops on the worker. + RELEASE.store(true, AtomicOrdering::SeqCst); + worker.join().unwrap(); + assert_eq!( + FREE_COUNT.load(AtomicOrdering::SeqCst), + 1, + "free must run exactly once, after the in-flight callback returned" + ); + } + + // A null free_user_data opts out of ownership: Drop must not call anything. + #[test] + fn test_headers_provider_no_free_callback_is_noop_on_drop() { + extern "C" fn callback(_user_data: *mut std::ffi::c_void) -> CHeaders { + CHeaders { + headers: ptr::null_mut(), + count: 0, + error_message: ptr::null_mut(), + } + } + // No free callback: dropping must be a no-op (and must not deref user_data). + let provider = CallbackHeadersProvider::new(callback, ptr::null_mut(), None); + drop(provider); + } + // ======================================================================== // Ack callback bridge tests // ======================================================================== @@ -1390,6 +1517,7 @@ mod tests { 0, empty_headers, ptr::null_mut(), + None, &options as *const crate::CStreamConfigurationOptions, &mut result as *mut CResult, ); @@ -1413,6 +1541,53 @@ mod tests { zerobus_sdk_free(sdk); } + // The provider owns `user_data`, so a failed create must still invoke + // `free_user_data` exactly once (the create-failure path of the free-on- + // every-path contract) — otherwise the wrapper's handle would leak. + #[test] + fn test_create_stream_with_headers_provider_frees_user_data_on_failure() { + static FREE_COUNT: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0); + static FREE_SENTINEL: u8 = 0; + extern "C" fn free_user_data(user_data: *mut std::ffi::c_void) { + if std::ptr::eq(user_data as *const u8, &FREE_SENTINEL) { + FREE_COUNT.fetch_add(1, AtomicOrdering::SeqCst); + } + } + + let (sdk, sdk_result) = + build_via_c_builder("https://workspace.zerobus.databricks.com", "", None, None); + assert!(sdk_result.success); + assert!(!sdk.is_null()); + + let empty_table = CString::new("").unwrap(); + let options = zerobus_get_default_config(); + let mut result = presumed_success_result(); + let user_data = (&FREE_SENTINEL as *const u8) as *mut std::ffi::c_void; + + let before = FREE_COUNT.load(AtomicOrdering::SeqCst); + let stream = zerobus_sdk_create_stream_with_headers_provider( + sdk, + empty_table.as_ptr(), + ptr::null(), + 0, + empty_headers, + user_data, + Some(free_user_data), + &options as *const crate::CStreamConfigurationOptions, + &mut result as *mut CResult, + ); + let after = FREE_COUNT.load(AtomicOrdering::SeqCst); + + assert!(stream.is_null()); + assert_eq!( + after - before, + 1, + "expected free_user_data to run exactly once on failed create" + ); + + zerobus_sdk_free(sdk); + } + // ======================================================================== // Dynamic protobuf schema tests // ======================================================================== diff --git a/rust/ffi/zerobus.h b/rust/ffi/zerobus.h index 3c3aedb1..b39f4a6f 100644 --- a/rust/ffi/zerobus.h +++ b/rust/ffi/zerobus.h @@ -266,6 +266,12 @@ struct CArrowStream *zerobus_sdk_create_arrow_stream(struct CZerobusSdk *sdk, * Creates an Arrow Flight stream with a custom headers provider callback. * * `schema_ipc_bytes` must point to Arrow IPC stream bytes encoding only the schema. + * + * Ownership of `user_data` / `free_user_data` follows + * `zerobus_sdk_create_stream_with_headers_provider` — once called the FFI owns + * `user_data` and invokes `free_user_data` exactly once on every path (on + * success after any in-flight `get_headers` returns; on failure before + * returning null). The caller must never free `user_data` itself. */ struct CArrowStream *zerobus_sdk_create_arrow_stream_with_headers_provider(struct CZerobusSdk *sdk, const char *table_name, @@ -273,6 +279,7 @@ struct CArrowStream *zerobus_sdk_create_arrow_stream_with_headers_provider(struc uintptr_t schema_ipc_len, HeadersProviderCallback headers_callback, void *user_data, + void (*free_user_data)(void *user_data), const struct CArrowStreamConfigurationOptions *options, struct CResult *result); @@ -456,7 +463,22 @@ bool zerobus_sdk_create_stream_async(struct CZerobusSdk *sdk, /** * Create a stream with a custom headers provider callback - * This allows you to provide custom authentication headers via a Go callback function + * This allows you to provide custom authentication headers via a Go/C callback function. + * + * Ownership: once this function is called, the FFI owns `user_data`. When + * `free_user_data` is set it is invoked exactly once, on every path: + * - on success, when the last internal reference to the provider drops — after + * any in-flight `get_headers` callback has returned (this is what closes the + * recovery-vs-teardown use-after-free); + * - on failure (this call returns null), before returning. + * + * The caller must therefore hand ownership across and never free `user_data` + * itself, not even when create fails. Pass a null `free_user_data` to opt out + * (the caller then owns `user_data` and must keep it alive for the stream's + * whole lifetime, including in-flight recovery callbacks). + * + * `free_user_data` may run on an internal SDK thread, so it must be safe to + * call from any thread. */ struct CZerobusStream *zerobus_sdk_create_stream_with_headers_provider(struct CZerobusSdk *sdk, const char *table_name, @@ -464,6 +486,7 @@ struct CZerobusStream *zerobus_sdk_create_stream_with_headers_provider(struct CZ uintptr_t descriptor_proto_len, HeadersProviderCallback headers_callback, void *user_data, + void (*free_user_data)(void *user_data), const struct CStreamConfigurationOptions *options, struct CResult *result); @@ -474,6 +497,15 @@ struct CZerobusStream *zerobus_sdk_create_stream_with_headers_provider(struct CZ * callback is invoked exactly once with either a non-null stream pointer and a * success result, or a null stream pointer and a failure result. The SDK * handle must remain valid until the callback runs. + * + * OWNERSHIP: unlike the synchronous `zerobus_sdk_create_stream_with_headers_provider`, + * this variant exposes no `free_user_data` callback, so it does NOT take + * ownership of `user_data` — the provider is built with `free_user_data = None` + * and the caller stays responsible for freeing `user_data` (only after the + * completion callback has fired, since a recovery `get_headers` may still run). + * No wrapper currently calls this; before one adopts it, add a `free_user_data` + * parameter and hand ownership across (as the sync path does), or it will + * reintroduce the recovery-vs-teardown use-after-free this change fixes. */ bool zerobus_sdk_create_stream_with_headers_provider_async(struct CZerobusSdk *sdk, const char *table_name,