Skip to content

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

Open
zlata-stefanovic-db wants to merge 9 commits into
mainfrom
ffi-headers-provider-uaf-fix
Open

[FFI] Fix headers-provider use-after-free via ownership transfer#529
zlata-stefanovic-db wants to merge 9 commits into
mainfrom
ffi-headers-provider-uaf-fix

Conversation

@zlata-stefanovic-db

Copy link
Copy Markdown
Contributor

Summary

Fixes a use-after-free where a custom HeadersProvider could be freed by the wrapper (Go/C++) while a Rust worker thread was still inside a synchronous get_headers() callback during connection recovery. Cooperative abort() cannot preempt a synchronous FFI call, and close() waited only SHUTDOWN_TIMEOUT_SECS = 1 before aborting, so a slow provider (e.g. one fetching/rotating a token) racing teardown could be invoked on freed memory.

Approach: ownership transfer at the FFI boundary

The C FFI now takes ownership of the provider user_data. zerobus_sdk_create_stream_with_headers_provider / zerobus_sdk_create_arrow_stream_with_headers_provider gain a nullable free_user_data destroy callback. CallbackHeadersProvider invokes it from Drop, which runs only after the last Arc referencing the provider drops — and the supervisor task holds its own Arc clone across the in-flight get_headers() call, so the destroy is guaranteed to run only after any synchronous callback has returned. The provider Arc is constructed before any fallible work, so free_user_data fires exactly once on every path (success after the last reference drops; failed create before returning). Wrappers hand ownership across and never free user_data themselves.

Changes by layer

  • Rust FFI (rust/ffi): free_user_data param + Drop-based destroy (panic-contained via catch_unwind); regenerated zerobus.h; ownership/thread-safety docs.
  • Go (go/ffi.go, go/arrow_ffi.go): hands the cgo.Handle to the FFI via goFreeHeadersProvider; removes the per-stream handle registry.
  • C++ (cpp/): passes a heap std::shared_ptr<HeadersProvider> + zerobus_cpp_headers_free destroy trampoline; Stream/ArrowStream no longer keep the provider.

Scope: why only Go and C++

The bug class is specific to passing a borrowed user_data across the C FFI while the core holds a raw pointer. The other consumers are not affected:

SDK Provider bridge Affected?
Go, C++ borrowed raw pointer / handle freed on close Yes — fixed here
Python HeadersProviderWrapper owns a PyObject, handed as an owned Arc No — Rust's Arc controls lifetime
TypeScript headers resolved once at create → StaticHeadersProvider No — no live JS object across FFI
Java no custom headers-provider feature (OAuth only) No — feature absent

Tests

  • Rust: free_user_data fires once on Drop and on failed create; null destroy is a no-op; and a teardown test reproducing the recovery-vs-teardown race (a blocking in-flight get_headers on one Arc clone while another drops) asserts the free is deferred until the callback returns.
  • Go: goFreeHeadersProvider releases the cgo.Handle.
  • C++: trampoline dispatch + free destroys the provider; null tolerance. Passes under ASan.

All 82 Rust FFI tests and 6 C++ tests pass; changelogs updated for ffi/go/cpp.

Breaking change (FFI)

The two create-with-provider C signatures gain a free_user_data parameter — a non-additive zerobus.h change. Filed under Breaking Changes in rust/ffi/NEXT_CHANGELOG.md.

Follow-up required before Go picks this up

This changes the FFI signature, so an FFI release (ffi/v*) must land and regenerate all five platform libzerobus_ffi.a archives in go/lib/ before the Go SDK can consume it. The committed archives are intentionally left at their pre-PR builds (the locally-rebuilt one exceeded GitHub's 100 MB limit and would have been an inconsistent single-platform artifact). Go CI's cross-SDK job builds the FFI from local Rust source, so it validates this PR without the prebuilt libs.

Draft — open items

  • Consider whether Go should nil-check the provider before cgo.NewHandle (C++ already throws on null) — pre-existing asymmetry, non-blocking.

The FFI now owns the headers provider user_data and frees it via a
destroy callback only after any in-flight get_headers() returns, closing
the recovery-vs-teardown use-after-free. Updates Go and C++ wrappers to
hand ownership across, adds tests, changelogs, and lifetime docs.

Signed-off-by: Zlata Stefanovic <zlata.stefanovic@databricks.com>
…ract

After merging main, headers_callback_test (added there) passed a raw
HeadersProvider* to zerobus_cpp_headers_trampoline, which this branch changed to
expect a heap std::shared_ptr<HeadersProvider>* � causing a segfault. Wrap each
provider in a heap shared_ptr (as Sdk::create_* does) and release it via
zerobus_cpp_headers_free, matching the FFI ownership contract.

Signed-off-by: Zlata Stefanovic <zlata.stefanovic@databricks.com>
@zlata-stefanovic-db
zlata-stefanovic-db marked this pull request as ready for review July 17, 2026 17:31
@zlata-stefanovic-db
zlata-stefanovic-db requested review from a team and irinatomic-db July 17, 2026 17:32
Resolve conflicts in rust/ffi/src/common.rs and rust/ffi/src/stream.rs.

- common.rs: keep both the new HeadersProviderFreeCallback type (this
  branch) and the async-callback types (main), side by side.
- stream.rs: adopt main's build_stream_from_parts / StreamCreateAuth
  refactor, but carry a pre-built Arc<dyn HeadersProvider> through it so
  the provider is still constructed before any fallible work � preserving
  this branch's UAF fix (Drop frees user_data on every error path). The
  async entry point builds the provider with free_user_data = None, since
  it exposes no free callback.

Signed-off-by: Zlata Stefanovic <zlata.stefanovic@databricks.com>
Signed-off-by: Zlata Stefanovic <zlata.stefanovic@databricks.com>
…ship gap

- goFreeHeadersProvider: nil-check before cgo.Handle.Delete() so a zero
  handle can't panic across the C boundary (matches the C++ null tolerance).
- Document that the async create-with-provider variant does not take
  ownership (free_user_data = None) and must gain one before any wrapper
  adopts it, to avoid reintroducing the recovery-vs-teardown UAF.
- Regenerate zerobus.h with the async doc comment (cbindgen; no ABI change).

Signed-off-by: Zlata Stefanovic <zlata.stefanovic@databricks.com>
The FFI now takes ownership of the headers-provider user_data via a
free_user_data destroy callback. Update the .NET sync binding to match
(the async FFI signature is unchanged) and hand the provider's GCHandle
ownership to the FFI: it is released by NativeFree after any in-flight
GetHeaders returns, instead of the stream freeing it on dispose. Fixes
the cross-sdk-dotnet test-host crash from the stale signature and closes
the recovery-vs-teardown use-after-free for the sync path.

Re-scope now-async-only lifetime comments, document the provider
lifetime/threading contract on IHeadersProvider, and add unit tests for
the NativeFree destroy callback.

Signed-off-by: Zlata Stefanovic <zlata.stefanovic@databricks.com>
GCHandle is a struct wrapping an IntPtr, so NativeFree frees the runtime
handle through its own GCHandle.FromIntPtr copy and cannot zero the test's
local handle field � asserting on handle.IsAllocated was wrong and always
saw True. Observe the actual release with a WeakReference instead, from a
non-inlined helper so no stray root keeps the bridge alive.

Signed-off-by: Zlata Stefanovic <zlata.stefanovic@databricks.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[FFI] HeadersProvider can be freed while an in-flight recovery get_headers() call runs (UAF)

1 participant