bridge: Use bridge_callbacks for KyberPreKeyStore (C bridge only)

Adds support for return values and error results to #[bridge_callbacks],
and then provides the necessary extra *TypeInfo impls to support
KyberPreKeyStore.

This also drops the "lazy clone" optimization previously used by the
store, where a bridged handle owned by Rust can be "borrowed" by the
Swift side and then cloned if still alive when the callback is about
to complete. Now the value is cloned eagerly, and destroyed via Swift
regardless. This is unfortunate, but the Android bridge has never had
this optimization, and Android supports devices with worse
performance, so we should be able to live with it.

rust/bridge/shared/macros/src/ffi.rs

@@ -271,8 +271,10 @@ fn bridge_callback_item(trait_name: &Ident, item: &TraitItem) -> Result<Callback
 
     let sig = &item.sig;
     let req_name = &item.sig.ident;
+    let result_info = ResultInfo::from(&sig.output);
+    let result_ty = result_type(&sig.output);
 
-    // type FfiMyTraitOperation = extern "C" fn(ctx: *mut c_void, foo: ffi_result_type!(u32))
+    // type FfiMyTraitOperation = extern "C" fn(ctx: *mut c_void, foo: ffi_result_type!(u32)) -> c_int
     let callback_ty_name = format_ident!(
         "Ffi{}{}",
         trait_name,
@@ -280,7 +282,10 @@ fn bridge_callback_item(trait_name: &Ident, item: &TraitItem) -> Result<Callback
         span = req_name.span()
     );
     let callback_args = item.sig.inputs.iter().filter_map(|arg| match arg {
-        FnArg::Receiver(_) => None,
+        FnArg::Receiver(_) => match result_info.kind {
+            ResultKind::Regular => Some(quote!(out: *mut ffi_arg_type!(#result_ty))),
+            ResultKind::Void => None,
+        },
         FnArg::Typed(arg) => {
             let Pat::Ident(arg_name) = &*arg.pat else {
                 // We'll error about this elsewhere.
@@ -309,8 +314,12 @@ fn bridge_callback_item(trait_name: &Ident, item: &TraitItem) -> Result<Callback
     //       (self.operation)(self.ctx, ffi::ResultTypeInfo::convert_into(foo).expect("can convert"))
     //   )
     // }
+    let out_ptr_arg = match result_info.kind {
+        ResultKind::Regular => quote!(, __result.as_mut_ptr()), // note the LEADING comma
+        ResultKind::Void => quote!(),
+    };
     let arg_conversions = item.sig.inputs.iter().map(|arg| match arg {
-        FnArg::Receiver(_) => quote!(self.ctx),
+        FnArg::Receiver(_) => quote!(self.ctx #out_ptr_arg),
         FnArg::Typed(arg) => {
             let Pat::Ident(arg_name) = &*arg.pat else {
                 return Error::new(arg.pat.span(), "only simple argument syntax is supported")
@@ -324,14 +333,37 @@ fn bridge_callback_item(trait_name: &Ident, item: &TraitItem) -> Result<Callback
             }
         }
     });
-    let implementation = quote! {
-        // #sig carries everything from `fn` to the return type and possible where-clause.
-        // All we provide is the body.
-        #sig {
-            ffi::CallbackError::log_on_error(
-                stringify!(#req_name),
-                (self.#field_name)(#(#arg_conversions,)*)
-            )
+    let implementation = if result_info.failable {
+        quote! {
+            // #sig carries everything from `fn` to the return type and possible where-clause.
+            // All we provide is the body.
+            #sig {
+                let mut __result = std::mem::MaybeUninit::zeroed();
+                ffi::CallbackError::check((self.#field_name)(#(#arg_conversions,)*))
+                    .map_err(|e| WithContext {
+                        operation: stringify!(#req_name),
+                        inner: e
+                    })?;
+                <<#result_ty as ResultLike>::Success as ffi::CallbackResultTypeInfo>::convert_from_callback(
+                    // SAFETY: if the C function returns 0 (success), they had better initialize this.
+                    // (Exception: a void function has no out-parameter, but `()` is trivially initialized.)
+                    unsafe { __result.assume_init() }
+                ).map_err(|e| WithContext {
+                    operation: stringify!(#req_name),
+                    inner: e
+                }.into())
+            }
+        }
+    } else {
+        quote! {
+            #sig {
+                // Not implemented: callbacks that *do* have a return value but *don't* have a place for errors.
+                // We can handle those some if we need them.
+                ffi::CallbackError::log_on_error(
+                    stringify!(#req_name),
+                    (self.#field_name)(#(#arg_conversions,)*)
+                )
+            }
         }
     };
 

rust/bridge/shared/types/src/ffi/convert.rs

@@ -529,10 +529,25 @@ bridge_trait!(PreKeyStore);
 bridge_trait!(SenderKeyStore);
 bridge_trait!(SessionStore);
 bridge_trait!(SignedPreKeyStore);
-bridge_trait!(KyberPreKeyStore);
+// bridge_trait!(KyberPreKeyStore);
 bridge_trait!(InputStream);
 bridge_trait!(SyncInputStream);
 
+impl<'a> ArgTypeInfo<'a> for &'a mut dyn KyberPreKeyStore {
+    type ArgType = crate::ffi::ConstPointer<FfiBridgeKyberPreKeyStoreStruct>;
+    type StoredType = BridgedStore<OwnedCallbackStruct<FfiBridgeKyberPreKeyStoreStruct>>;
+    #[allow(clippy::not_unsafe_ptr_arg_deref)]
+    fn borrow(foreign: Self::ArgType) -> SignalFfiResult<Self::StoredType> {
+        match unsafe { foreign.into_inner().as_ref() } {
+            None => Err(NullPointerError.into()),
+            Some(store) => Ok(BridgedStore(OwnedCallbackStruct(store.clone()))),
+        }
+    }
+    fn load_from(stored: &'a mut Self::StoredType) -> Self {
+        stored
+    }
+}
+
 impl<'a> ArgTypeInfo<'a> for Box<dyn ChatListener> {
     type ArgType = crate::ffi::ConstPointer<FfiChatListenerStruct>;
     type StoredType = Option<Box<dyn ChatListener>>;
@@ -1244,6 +1259,20 @@ macro_rules! ffi_bridge_handle_fns {
     };
 }
 
+// This can't be done generically because it would conflict with the blanket impl, and it *also*
+// can't be done as part of the common `bridge_as_handle!` macro because when used outside this
+// crate it would violate the orphan rule. Perhaps there'll be a workaround later.
+impl CallbackResultTypeInfo for Option<KyberPreKeyRecord> {
+    type ResultType = MutPointer<KyberPreKeyRecord>;
+
+    fn convert_from_callback(foreign: Self::ResultType) -> SignalFfiResult<Self> {
+        if foreign == MutPointer::null() {
+            return Ok(None);
+        }
+        KyberPreKeyRecord::convert_from_callback(foreign).map(Some)
+    }
+}
+
 macro_rules! trivial {
     ($typ:ty) => {
         impl SimpleArgTypeInfo for $typ {
@@ -1269,13 +1298,13 @@ trivial!(u64);
 trivial!(usize);
 trivial!(bool);
 
-/// Syntactically translates `bridge_fn` argument types to FFI types for `cbindgen`.
+/// Syntactically translates `bridge_fn` argument types (and callback result types) to FFI types for
+/// `cbindgen`.
 ///
 /// This is a syntactic transformation (because that's how Rust macros work), so new argument types
 /// will need to be added here directly even if they already implement [`ArgTypeInfo`]. The default
-/// behavior for references is to pass them through as pointers; the default behavior for
-/// `&mut dyn Foo` is to assume there's a struct called `ffi::FfiFooStruct` and produce a pointer
-/// to that.
+/// behavior for references is to pass them through as pointers; the default behavior for `&mut dyn
+/// Foo` is to assume there's a struct called `ffi::FfiFooStruct` and produce a pointer to that.
 #[macro_export]
 macro_rules! ffi_arg_type {
     (u8) => (u8);
@@ -1330,9 +1359,16 @@ macro_rules! ffi_arg_type {
     // In order to provide a fixed-sized array of the correct length,
     // a serialized type FooBar must have a constant FOO_BAR_LEN that's in scope (and exposed to C).
     (Serialized<$typ:ident>) => (*const [std::ffi::c_uchar; ::paste::paste!([<$typ:snake:upper _LEN>])]);
+
+    // For use in callbacks.
+    (Result<$typ:tt $(, $ignored:ty)?>) => (ffi_arg_type!($typ));
+    (Result<$typ:tt<$($args:tt),+> $(, $ignored:ty)?>) => (ffi_arg_type!($typ<$($args),+>));
+    (Option<$typ:ty>) => (ffi::MutPointer< $typ >);
+    ($typ:ty) => (ffi::MutPointer< $typ >);
 }
 
-/// Syntactically translates `bridge_fn` result types to FFI types for `cbindgen`.
+/// Syntactically translates `bridge_fn` result types (and callback argument types) to FFI types for
+/// `cbindgen`.
 ///
 /// This is a syntactic transformation (because that's how Rust macros work), so new result types
 /// will need to be added here directly even if they already implement [`ResultTypeInfo`]. The

rust/bridge/shared/types/src/ffi/error.rs

@@ -23,7 +23,7 @@ use usernames::{UsernameError, UsernameLinkError};
 use zkgroup::{ZkGroupDeserializationFailure, ZkGroupVerificationFailure};
 
 use super::{FutureCancelled, NullPointerError, UnexpectedPanic};
-use crate::support::{IllegalArgumentError, describe_panic};
+use crate::support::{IllegalArgumentError, WithContext, describe_panic};
 
 #[derive(Debug, Clone, Copy)]
 #[repr(C)]
@@ -1170,3 +1170,22 @@ impl fmt::Display for CallbackError {
 }
 
 impl std::error::Error for CallbackError {}
+
+impl From<WithContext<CallbackError>> for SignalProtocolError {
+    fn from(value: WithContext<CallbackError>) -> Self {
+        let WithContext { operation, inner } = value;
+        SignalProtocolError::for_application_callback(operation)(inner)
+    }
+}
+
+/// This is overly general, but in practice is only used to handle errors converting callback
+/// results.
+impl From<WithContext<SignalFfiError>> for SignalProtocolError {
+    fn from(value: WithContext<SignalFfiError>) -> Self {
+        let WithContext {
+            operation: _,
+            inner,
+        } = value;
+        SignalProtocolError::FfiBindingError(inner.to_string())
+    }
+}

rust/bridge/shared/types/src/ffi/storage.rs

@@ -6,9 +6,12 @@
 use std::ffi::{c_int, c_uint, c_void};
 
 use async_trait::async_trait;
+use libsignal_bridge_macros::bridge_callbacks;
 use uuid::Uuid;
 
 use super::*;
+use crate::ffi;
+use crate::support::{ResultLike, WithContext};
 
 type GetIdentityKeyPair =
     extern "C" fn(store_ctx: *mut c_void, keyp: *mut MutPointer<PrivateKey>) -> c_int;
@@ -263,53 +266,45 @@ impl SignedPreKeyStore for &FfiSignedPreKeyStoreStruct {
     }
 }
 
-type LoadKyberPreKey = extern "C" fn(
-    store_ctx: *mut c_void,
-    recordp: *mut MutPointer<KyberPreKeyRecord>,
-    id: u32,
-) -> c_int;
-type StoreKyberPreKey = extern "C" fn(
-    store_ctx: *mut c_void,
-    id: u32,
-    record: ConstPointer<KyberPreKeyRecord>,
-) -> c_int;
-type MarkKyberPreKeyUsed = extern "C" fn(
-    store_ctx: *mut c_void,
-    id: u32,
-    signed_prekey_id: u32,
-    base_key: ConstPointer<PublicKey>,
-) -> c_int;
-
-#[repr(C)]
-#[derive(Copy, Clone)]
-pub struct FfiKyberPreKeyStoreStruct {
-    ctx: *mut c_void,
-    load_kyber_pre_key: LoadKyberPreKey,
-    store_kyber_pre_key: StoreKyberPreKey,
-    mark_kyber_pre_key_used: MarkKyberPreKeyUsed,
+/// A bridge-friendly version of [`KyberPreKeyStore`].
+#[bridge_callbacks(jni = false, node = false)]
+pub trait BridgeKyberPreKeyStore {
+    fn load_kyber_pre_key(&self, id: u32)
+    -> Result<Option<KyberPreKeyRecord>, SignalProtocolError>;
+    fn store_kyber_pre_key(
+        &self,
+        id: u32,
+        record: KyberPreKeyRecord,
+    ) -> Result<(), SignalProtocolError>;
+    fn mark_kyber_pre_key_used(
+        &self,
+        id: u32,
+        ec_prekey_id: u32,
+        base_key: PublicKey,
+    ) -> Result<(), SignalProtocolError>;
 }
 
+/// A wrapper struct so we can implement e.g. [`KyberPreKeyStore`] for all
+/// [`BridgeKyberPreKeyStore`]s.
+///
+/// Trying to do so directly would violate the [orphan rule][], because rustc doesn't know
+/// `BridgeKyberPreKeyStore` is only implemented by a closed set of types defined in this crate.
+///
+/// [orphan rule]: https://doc.rust-lang.org/book/ch20-02-advanced-traits.html#implementing-external-traits-with-the-newtype-pattern
+pub struct BridgedStore<T>(pub T);
+
+// TODO: This alias is because of the ffi_arg_type macro expecting all bridging structs to use a
+// particular naming scheme; eventually we should be able to remove it.
+pub type FfiKyberPreKeyStoreStruct = FfiBridgeKyberPreKeyStoreStruct;
+
 #[async_trait(?Send)]
-impl KyberPreKeyStore for &FfiKyberPreKeyStoreStruct {
+impl<T: BridgeKyberPreKeyStore> KyberPreKeyStore for BridgedStore<T> {
     async fn get_kyber_pre_key(
         &self,
         id: KyberPreKeyId,
     ) -> Result<KyberPreKeyRecord, SignalProtocolError> {
-        let mut record = MutPointer::null();
-        let result = (self.load_kyber_pre_key)(self.ctx, &mut record, id.into());
-
-        CallbackError::check(result).map_err(SignalProtocolError::for_application_callback(
-            "load_kyber_pre_key",
-        ))?;
-
-        let record = record.into_inner();
-        if record.is_null() {
-            return Err(SignalProtocolError::InvalidKyberPreKeyId);
-        }
-
-        let record = unsafe { Box::from_raw(record) };
-
-        Ok(*record)
+        BridgeKyberPreKeyStore::load_kyber_pre_key(&self.0, id.into())?
+            .ok_or(SignalProtocolError::InvalidKyberPreKeyId)
     }
 
     async fn save_kyber_pre_key(
@@ -317,11 +312,7 @@ impl KyberPreKeyStore for &FfiKyberPreKeyStoreStruct {
         id: KyberPreKeyId,
         record: &KyberPreKeyRecord,
     ) -> Result<(), SignalProtocolError> {
-        let result = (self.store_kyber_pre_key)(self.ctx, id.into(), record.into());
-
-        CallbackError::check(result).map_err(SignalProtocolError::for_application_callback(
-            "store_kyber_pre_key",
-        ))
+        BridgeKyberPreKeyStore::store_kyber_pre_key(&self.0, id.into(), record.clone())
     }
 
     async fn mark_kyber_pre_key_used(
@@ -330,16 +321,12 @@ impl KyberPreKeyStore for &FfiKyberPreKeyStoreStruct {
         ec_prekey_id: SignedPreKeyId,
         base_key: &PublicKey,
     ) -> Result<(), SignalProtocolError> {
-        let result = (self.mark_kyber_pre_key_used)(
-            self.ctx,
+        BridgeKyberPreKeyStore::mark_kyber_pre_key_used(
+            &self.0,
             id.into(),
             ec_prekey_id.into(),
-            base_key.into(),
-        );
-
-        CallbackError::check(result).map_err(SignalProtocolError::for_application_callback(
-            "mark_kyber_pre_key_used",
-        ))
+            *base_key,
+        )
     }
 }
 

rust/bridge/shared/types/src/support/mod.rs

@@ -355,3 +355,22 @@ impl<F: Future<Output: ResultReporter>> AsyncRuntime<F> for NoOpAsyncRuntime {
         CancellationId::NotSupported
     }
 }
+
+/// Attaches context to a value, usually an error.
+///
+/// Intended to be used with `From` implementations, so standard Rust error handling idioms can work
+/// even for error types that want the additional context.
+pub struct WithContext<T> {
+    pub operation: &'static str,
+    pub inner: T,
+}
+
+/// Provides access to [`Result`]'s `Success` and `Error` types using associated type syntax.
+pub trait ResultLike {
+    type Success;
+    type Error;
+}
+impl<T, E> ResultLike for Result<T, E> {
+    type Success = T;
+    type Error = E;
+}

swift/Sources/LibSignalClient/DataStoreUtils.swift

@@ -238,14 +238,13 @@ internal func withKyberPreKeyStore<Result>(
     func ffiShimStoreKyberPreKey(
         storeCtx: UnsafeMutableRawPointer?,
         id: UInt32,
-        record: SignalConstPointerKyberPreKeyRecord
+        record: SignalMutPointerKyberPreKeyRecord
     ) -> Int32 {
         let storeContext = storeCtx!.assumingMemoryBound(
             to: ErrorHandlingContext<(KyberPreKeyStore, StoreContext)>.self
         )
         return storeContext.pointee.catchCallbackErrors { store, context in
-            var record = KyberPreKeyRecord(borrowing: record)
-            defer { cloneOrForgetAsNeeded(&record) }
+            let record = KyberPreKeyRecord(owned: NonNull(record)!)
             try store.storeKyberPreKey(record, id: id, context: context)
             return 0
         }
@@ -270,14 +269,13 @@ internal func withKyberPreKeyStore<Result>(
         storeCtx: UnsafeMutableRawPointer?,
         id: UInt32,
         signedPreKeyId: UInt32,
-        baseKey: SignalConstPointerPublicKey,
+        baseKey: SignalMutPointerPublicKey,
     ) -> Int32 {
         let storeContext = storeCtx!.assumingMemoryBound(
             to: ErrorHandlingContext<(KyberPreKeyStore, StoreContext)>.self
         )
         return storeContext.pointee.catchCallbackErrors { store, context in
-            var baseKey = PublicKey(borrowing: baseKey)
-            defer { cloneOrForgetAsNeeded(&baseKey) }
+            let baseKey = PublicKey(owned: NonNull(baseKey)!)
             try store.markKyberPreKeyUsed(id: id, signedPreKeyId: signedPreKeyId, baseKey: baseKey, context: context)
             return 0
         }
@@ -288,7 +286,8 @@ internal func withKyberPreKeyStore<Result>(
             ctx: $0,
             load_kyber_pre_key: ffiShimLoadKyberPreKey,
             store_kyber_pre_key: ffiShimStoreKyberPreKey,
-            mark_kyber_pre_key_used: ffiShimMarkKyberPreKeyUsed
+            mark_kyber_pre_key_used: ffiShimMarkKyberPreKeyUsed,
+            destroy: { _ in }
         )
         return try withUnsafePointer(to: &ffiStore) {
             try body(SignalConstPointerFfiKyberPreKeyStoreStruct(raw: $0))

swift/Sources/SignalFfi/signal_ffi.h

@@ -867,22 +867,23 @@ typedef struct {
   SignalKyberPreKeyRecord *raw;
 } SignalMutPointerKyberPreKeyRecord;
 
-typedef int (*SignalLoadKyberPreKey)(void *store_ctx, SignalMutPointerKyberPreKeyRecord *recordp, uint32_t id);
+typedef int (*SignalFfiBridgeKyberPreKeyStoreLoadKyberPreKey)(void *ctx, SignalMutPointerKyberPreKeyRecord *out, uint32_t id);
 
-typedef struct {
-  const SignalKyberPreKeyRecord *raw;
-} SignalConstPointerKyberPreKeyRecord;
+typedef int (*SignalFfiBridgeKyberPreKeyStoreStoreKyberPreKey)(void *ctx, uint32_t id, SignalMutPointerKyberPreKeyRecord record);
 
-typedef int (*SignalStoreKyberPreKey)(void *store_ctx, uint32_t id, SignalConstPointerKyberPreKeyRecord record);
+typedef int (*SignalFfiBridgeKyberPreKeyStoreMarkKyberPreKeyUsed)(void *ctx, uint32_t id, uint32_t ec_prekey_id, SignalMutPointerPublicKey base_key);
 
-typedef int (*SignalMarkKyberPreKeyUsed)(void *store_ctx, uint32_t id, uint32_t signed_prekey_id, SignalConstPointerPublicKey base_key);
+typedef void (*SignalFfiBridgeKyberPreKeyStoreDestroy)(void *ctx);
 
 typedef struct {
   void *ctx;
-  SignalLoadKyberPreKey load_kyber_pre_key;
-  SignalStoreKyberPreKey store_kyber_pre_key;
-  SignalMarkKyberPreKeyUsed mark_kyber_pre_key_used;
-} SignalKyberPreKeyStore;
+  SignalFfiBridgeKyberPreKeyStoreLoadKyberPreKey load_kyber_pre_key;
+  SignalFfiBridgeKyberPreKeyStoreStoreKyberPreKey store_kyber_pre_key;
+  SignalFfiBridgeKyberPreKeyStoreMarkKyberPreKeyUsed mark_kyber_pre_key_used;
+  SignalFfiBridgeKyberPreKeyStoreDestroy destroy;
+} SignalFfiBridgeKyberPreKeyStoreStruct;
+
+typedef SignalFfiBridgeKyberPreKeyStoreStruct SignalKyberPreKeyStore;
 
 typedef struct {
   const SignalKyberPreKeyStore *raw;
@@ -1111,6 +1112,10 @@ typedef struct {
   SignalKyberSecretKey *raw;
 } SignalMutPointerKyberSecretKey;
 
+typedef struct {
+  const SignalKyberPreKeyRecord *raw;
+} SignalConstPointerKyberPreKeyRecord;
+
 typedef struct {
   const SignalKyberPublicKey *raw;
 } SignalConstPointerKyberPublicKey;