Use Handle for passing foreign objects to Rust

Consolidated the various handle maps into a single implementation for
each language.  This handle map works basically the same as all the
others, but it's API is based on the `HandleAlloc` trait.  Handles have
a couple properties:

* All foreign handles are odd, which allows us to distinguish between
  Rust and foreign handles.
* For handles store a map ID that can detect when a handle is used with
  the wrong map.

Made all languages always use the handle maps for passing objects.  No
more trying to leak pointers from to foreign objects.

Started updating the ForeignExecutor code to use handles, but this is
still a WIP while the ForeignExecutor type is in it's limbo state.

Author: bendk

CHANGELOG.md

@@ -19,7 +19,10 @@
 - The `rust_future_continuation_callback_set` FFI function was removed.  `rust_future_poll` now
   inputs the callback pointer.  External bindings authors will need to update their code.
 - The object handle FFI has changed.  External bindings generators will need to update their code to
-  use the new handle system.
+  use the new handle system:
+    * A single `FfiType::Handle` is used for all object handles.
+    * `FfiType::Handle` is always a 64-bit int.
+    * Foreign handles must always set the lowest bit of that int.
 
 ## v0.25.0 (backend crates: v0.25.0) - (_2023-10-18_)
 

fixtures/foreign-executor/tests/bindings/test_foreign_executor.kts

@@ -37,11 +37,3 @@ runBlocking {
     assert(result.delayMs <= 200U)
     tester.close()
 }
-
-// Test that we cleanup when dropping a ForeignExecutor handles
-assert(FfiConverterForeignExecutor.handleCount() == 0)
-val tester = ForeignExecutorTester(coroutineScope)
-val tester2 = ForeignExecutorTester.newFromSequence(listOf(coroutineScope))
-tester.close()
-tester2.close()
-assert(FfiConverterForeignExecutor.handleCount() == 0)

uniffi_bindgen/src/bindings/kotlin/gen_kotlin/mod.rs

@@ -298,12 +298,10 @@ impl KotlinCodeOracle {
             }
             FfiType::ForeignBytes => "ForeignBytes.ByValue".to_string(),
             FfiType::ForeignCallback => "ForeignCallback".to_string(),
-            FfiType::ForeignExecutorHandle => "USize".to_string(),
             FfiType::ForeignExecutorCallback => "UniFfiForeignExecutorCallback".to_string(),
             FfiType::RustFutureContinuationCallback => {
                 "UniFffiRustFutureContinuationCallbackType".to_string()
             }
-            FfiType::RustFutureContinuationData => "USize".to_string(),
         }
     }
 

uniffi_bindgen/src/bindings/kotlin/templates/Async.kt

@@ -3,18 +3,18 @@
 internal const val UNIFFI_RUST_FUTURE_POLL_READY = 0.toShort()
 internal const val UNIFFI_RUST_FUTURE_POLL_MAYBE_READY = 1.toShort()
 
-internal val uniffiContinuationHandleMap = UniFfiHandleMap<CancellableContinuation<Short>>()
+internal val uniffiContinuationHandleMap = UniffiHandleMap<CancellableContinuation<Short>>()
 
 // FFI type for Rust future continuations
 internal object uniffiRustFutureContinuationCallback: UniFffiRustFutureContinuationCallbackType {
-    override fun callback(continuationHandle: USize, pollResult: Short) {
-        uniffiContinuationHandleMap.remove(continuationHandle)?.resume(pollResult)
+    override fun callback(continuationHandle: UniffiHandle, pollResult: Short) {
+        uniffiContinuationHandleMap.consumeHandle(continuationHandle).resume(pollResult)
     }
 }
 
 internal suspend fun<T, F, E: Exception> uniffiRustCallAsync(
     rustFuture: UniffiHandle,
-    pollFunc: (UniffiHandle, UniFffiRustFutureContinuationCallbackType, USize) -> Unit,
+    pollFunc: (UniffiHandle, UniFffiRustFutureContinuationCallbackType, UniffiHandle) -> Unit,
     completeFunc: (UniffiHandle, RustCallStatus) -> F,
     freeFunc: (UniffiHandle) -> Unit,
     liftFunc: (F) -> T,
@@ -26,7 +26,7 @@ internal suspend fun<T, F, E: Exception> uniffiRustCallAsync(
                 pollFunc(
                     rustFuture,
                     uniffiRustFutureContinuationCallback,
-                    uniffiContinuationHandleMap.insert(continuation)
+                    uniffiContinuationHandleMap.newHandle(continuation)
                 )
             }
         } while (pollResult != UNIFFI_RUST_FUTURE_POLL_READY);

uniffi_bindgen/src/bindings/kotlin/templates/CallbackInterfaceImpl.kt

@@ -3,11 +3,11 @@
 // Implement the foreign callback handler for {{ interface_name }}
 internal class {{ callback_handler_class }} : ForeignCallback {
     @Suppress("TooGenericExceptionCaught")
-    override fun invoke(handle: Handle, method: Int, argsData: Pointer, argsLen: Int, outBuf: RustBufferByReference): Int {
+    override fun invoke(handle: UniffiHandle, method: Int, argsData: Pointer, argsLen: Int, outBuf: RustBufferByReference): Int {
         val cb = {{ ffi_converter_name }}.handleMap.get(handle)
         return when (method) {
             IDX_CALLBACK_FREE -> {
-                {{ ffi_converter_name }}.handleMap.remove(handle)
+                {{ ffi_converter_name }}.handleMap.consumeHandle(handle)
 
                 // Successful return
                 // See docs of ForeignCallback in `uniffi_core/src/ffi/foreigncallbacks.rs`

uniffi_bindgen/src/bindings/kotlin/templates/CallbackInterfaceRuntime.kt

@@ -1,39 +1,8 @@
 {{- self.add_import("java.util.concurrent.atomic.AtomicLong") }}
 {{- self.add_import("java.util.concurrent.locks.ReentrantLock") }}
-{{- self.add_import("kotlin.concurrent.withLock") }}
-
-internal typealias Handle = Long
-internal class ConcurrentHandleMap<T>(
-    private val leftMap: MutableMap<Handle, T> = mutableMapOf(),
-) {
-    private val lock = java.util.concurrent.locks.ReentrantLock()
-    private val currentHandle = AtomicLong(0L)
-    private val stride = 1L
-
-    fun insert(obj: T): Handle =
-        lock.withLock {
-            currentHandle.getAndAdd(stride)
-                .also { handle ->
-                    leftMap[handle] = obj
-                }
-            }
-
-    fun get(handle: Handle) = lock.withLock {
-        leftMap[handle] ?: throw InternalException("No callback in handlemap; this is a Uniffi bug")
-    }
-
-    fun delete(handle: Handle) {
-        this.remove(handle)
-    }
-
-    fun remove(handle: Handle): T? =
-        lock.withLock {
-            leftMap.remove(handle)
-        }
-}
 
 interface ForeignCallback : com.sun.jna.Callback {
-    public fun invoke(handle: Handle, method: Int, argsData: Pointer, argsLen: Int, outBuf: RustBufferByReference): Int
+    public fun invoke(handle: UniffiHandle, method: Int, argsData: Pointer, argsLen: Int, outBuf: RustBufferByReference): Int
 }
 
 // Magic number for the Rust proxy to call using the same mechanism as every other method,
@@ -44,20 +13,16 @@ internal const val UNIFFI_CALLBACK_SUCCESS = 0
 internal const val UNIFFI_CALLBACK_ERROR = 1
 internal const val UNIFFI_CALLBACK_UNEXPECTED_ERROR = 2
 
-public abstract class FfiConverterCallbackInterface<CallbackInterface>: FfiConverter<CallbackInterface, Handle> {
-    internal val handleMap = ConcurrentHandleMap<CallbackInterface>()
-
-    internal fun drop(handle: Handle) {
-        handleMap.remove(handle)
-    }
+public abstract class FfiConverterCallbackInterface<CallbackInterface>: FfiConverter<CallbackInterface, UniffiHandle> {
+    internal val handleMap = UniffiHandleMap<CallbackInterface>()
 
-    override fun lift(value: Handle): CallbackInterface {
+    override fun lift(value: UniffiHandle): CallbackInterface {
         return handleMap.get(value)
     }
 
     override fun read(buf: ByteBuffer) = lift(buf.getLong())
 
-    override fun lower(value: CallbackInterface) = handleMap.insert(value)
+    override fun lower(value: CallbackInterface) = handleMap.newHandle(value)
 
     override fun allocationSize(value: CallbackInterface) = 8
 

uniffi_bindgen/src/bindings/kotlin/templates/ForeignExecutorTemplate.kt

@@ -16,7 +16,7 @@ internal interface UniFfiRustTaskCallback : com.sun.jna.Callback {
 }
 
 internal object UniFfiForeignExecutorCallback : com.sun.jna.Callback {
-    fun callback(handle: USize, delayMs: Int, rustTask: UniFfiRustTaskCallback?, rustTaskData: Pointer?) : Byte {
+    fun callback(handle: UniffiHandle, delayMs: Int, rustTask: UniFfiRustTaskCallback?, rustTaskData: Pointer?) : Byte {
         if (rustTask == null) {
             FfiConverterForeignExecutor.drop(handle)
             return UNIFFI_FOREIGN_EXECUTOR_CALLBACK_SUCCESS
@@ -42,11 +42,11 @@ internal object UniFfiForeignExecutorCallback : com.sun.jna.Callback {
     }
 }
 
-public object FfiConverterForeignExecutor: FfiConverter<CoroutineScope, USize> {
-    internal val handleMap = UniFfiHandleMap<CoroutineScope>()
+public object FfiConverterForeignExecutor: FfiConverter<CoroutineScope, UniffiHandle> {
+    internal val handleMap = UniffiHandleMap<CoroutineScope>()
 
-    internal fun drop(handle: USize) {
-        handleMap.remove(handle)
+    internal fun drop(handle: UniffiHandle) {
+        handleMap.consumeHandle(handle)
     }
 
     internal fun register(lib: _UniFFILib) {
@@ -58,26 +58,21 @@ public object FfiConverterForeignExecutor: FfiConverter<CoroutineScope, USize> {
         {% endmatch %}
     }
 
-    // Number of live handles, exposed so we can test the memory management
-    public fun handleCount() : Int {
-        return handleMap.size
-    }
-
-    override fun allocationSize(value: CoroutineScope) = USize.size
+    override fun allocationSize(value: CoroutineScope) = 8
 
-    override fun lift(value: USize): CoroutineScope {
-        return handleMap.get(value) ?: throw RuntimeException("unknown handle in FfiConverterForeignExecutor.lift")
+    override fun lift(value: UniffiHandle): CoroutineScope {
+        return handleMap.get(value)
     }
 
     override fun read(buf: ByteBuffer): CoroutineScope {
-        return lift(USize.readFromBuffer(buf))
+        return lift(buf.getLong())
     }
 
-    override fun lower(value: CoroutineScope): USize {
-        return handleMap.insert(value)
+    override fun lower(value: CoroutineScope): UniffiHandle {
+        return handleMap.newHandle(value)
     }
 
     override fun write(value: CoroutineScope, buf: ByteBuffer) {
-        lower(value).writeToBuffer(buf)
+        buf.putLong(lower(value))
     }
 }

uniffi_bindgen/src/bindings/kotlin/templates/HandleMap.kt

@@ -0,0 +1,53 @@
+internal class UniffiHandleMap<T> {
+    private val lock = ReentrantReadWriteLock()
+    private var mapId: Long = UniffiHandleMap.nextMapId()
+    private val map: MutableMap<Long, T> = mutableMapOf()
+    // Note: Foreign handles are always odd
+    private var keyCounter = 1L
+
+    private fun nextKey(): Long = keyCounter.also {
+        keyCounter = (keyCounter + 2L).and(0xFFFF_FFFF_FFFFL)
+    }
+
+    private fun makeHandle(key: Long): UniffiHandle = key.or(mapId)
+
+    private fun key(handle: UniffiHandle): Long {
+        if (handle.and(0x7FFF_0000_0000_0000L) != mapId) {
+            throw InternalException("Handle map ID mismatch")
+        }
+        return handle.and(0xFFFF_FFFF_FFFFL)
+    }
+ 
+    fun newHandle(obj: T): UniffiHandle = lock.writeLock().withLock {
+        val key = nextKey()
+        map[key] = obj
+        makeHandle(key)
+    }
+ 
+    fun get(handle: UniffiHandle) = lock.readLock().withLock {
+        map[key(handle)] ?: throw InternalException("Missing key in handlemap: was the handle used after being freed?")
+    }
+ 
+    fun cloneHandle(handle: UniffiHandle): UniffiHandle = lock.writeLock().withLock {
+        val obj = map[key(handle)] ?: throw InternalException("Missing key in handlemap: was the handle used after being freed?")
+        val clone = nextKey()
+        map[clone] = obj
+        makeHandle(clone)
+    }
+ 
+    fun consumeHandle(handle: UniffiHandle): T = lock.writeLock().withLock {
+        map.remove(key(handle)) ?: throw InternalException("Missing key in handlemap: was the handle used after being freed?")
+    }
+
+    companion object {
+        // Generate map IDs that are likely to be unique
+        private var mapIdCounter: Long = {{ ci.namespace_hash() }}.and(0x7FFF)
+
+        // Map ID, shifted into the top 16 bits
+        internal fun nextMapId(): Long = mapIdCounter.shl(48).also {
+            // On Kotlin, map ids are only 15 bits to get around signed/unsigned issues
+            mapIdCounter = ((mapIdCounter + 1).and(0x7FFF))
+        }
+    }
+}
+

uniffi_bindgen/src/bindings/kotlin/templates/Helpers.kt

@@ -78,89 +78,7 @@ private inline fun <U> rustCall(callback: (RustCallStatus) -> U): U {
     return rustCallWithError(NullCallStatusErrorHandler, callback);
 }
 
-// IntegerType that matches Rust's `usize` / C's `size_t`
-public class USize(value: Long = 0) : IntegerType(Native.SIZE_T_SIZE, value, true) {
-    // This is needed to fill in the gaps of IntegerType's implementation of Number for Kotlin.
-    override fun toByte() = toInt().toByte()
-    // Needed until https://youtrack.jetbrains.com/issue/KT-47902 is fixed.
-    @Deprecated("`toInt().toChar()` is deprecated")
-    override fun toChar() = toInt().toChar()
-    override fun toShort() = toInt().toShort()
-
-    fun writeToBuffer(buf: ByteBuffer) {
-        // Make sure we always write usize integers using native byte-order, since they may be
-        // casted to pointer values
-        buf.order(ByteOrder.nativeOrder())
-        try {
-            when (Native.SIZE_T_SIZE) {
-                4 -> buf.putInt(toInt())
-                8 -> buf.putLong(toLong())
-                else -> throw RuntimeException("Invalid SIZE_T_SIZE: ${Native.SIZE_T_SIZE}")
-            }
-        } finally {
-            buf.order(ByteOrder.BIG_ENDIAN)
-        }
-    }
-
-    companion object {
-        val size: Int
-            get() = Native.SIZE_T_SIZE
-
-        fun readFromBuffer(buf: ByteBuffer) : USize {
-            // Make sure we always read usize integers using native byte-order, since they may be
-            // casted from pointer values
-            buf.order(ByteOrder.nativeOrder())
-            try {
-                return when (Native.SIZE_T_SIZE) {
-                    4 -> USize(buf.getInt().toLong())
-                    8 -> USize(buf.getLong())
-                    else -> throw RuntimeException("Invalid SIZE_T_SIZE: ${Native.SIZE_T_SIZE}")
-                }
-            } finally {
-                buf.order(ByteOrder.BIG_ENDIAN)
-            }
-        }
-    }
-}
-
-
-// Map handles to objects
-//
-// This is used when the Rust code expects an opaque pointer to represent some foreign object.
-// Normally we would pass a pointer to the object, but JNA doesn't support getting a pointer from an
-// object reference , nor does it support leaking a reference to Rust.
-//
-// Instead, this class maps USize values to objects so that we can pass a pointer-sized type to
-// Rust when it needs an opaque pointer.
-//
-// TODO: refactor callbacks to use this class
-internal class UniFfiHandleMap<T: Any> {
-    private val map = ConcurrentHashMap<USize, T>()
-    // Use AtomicInteger for our counter, since we may be on a 32-bit system.  4 billion possible
-    // values seems like enough. If somehow we generate 4 billion handles, then this will wrap
-    // around back to zero and we can assume the first handle generated will have been dropped by
-    // then.
-    private val counter = java.util.concurrent.atomic.AtomicInteger(0)
-
-    val size: Int
-        get() = map.size
-
-    fun insert(obj: T): USize {
-        val handle = USize(counter.getAndAdd(1).toLong())
-        map.put(handle, obj)
-        return handle
-    }
-
-    fun get(handle: USize): T? {
-        return map.get(handle)
-    }
-
-    fun remove(handle: USize): T? {
-        return map.remove(handle)
-    }
-}
-
 // FFI type for Rust future continuations
 internal interface UniFffiRustFutureContinuationCallbackType : com.sun.jna.Callback {
-    fun callback(continuationHandle: USize, pollResult: Short);
+    fun callback(continuationHandle: Long, pollResult: Short);
 }

uniffi_bindgen/src/bindings/kotlin/templates/ObjectRuntime.kt

@@ -30,6 +30,16 @@ inline fun <T : Disposable?, R> T.use(block: (T) -> R) =
         }
     }
 
+// Wraps `UniffiHandle` to pass to object constructors.
+//
+// This class exists because `UniffiHandle` is a typealias to `Long`.  If the object constructor
+// inputs `UniffiHandle` directly and the user defines a primary constructor than inputs a single
+// `Long` or `ULong` input, then we get JVM signature conflicts.  To avoid this, we pass this type
+// in instead.
+//
+// Let's try to remove this when we update the code based on ADR-0008.
+data class UniffiHandleWrapper(val handle: UniffiHandle)
+
 // The base class for all UniFFI Object types.
 //
 // This class provides core operations for working with the Rust handle to the live Rust struct on