feat(swift): transactions run withTaskCancellationHandler [WPB-27117]

.github/actions/make/bindings-swift/action.yml

@@ -15,10 +15,14 @@ runs:
           gh-token: ${{ inputs.gh-token }}
           artifact-generation: ${{ inputs.artifact-generation }}
 
-      - name: fetch ffi library
-        uses: ./.github/actions/make/ffi-library
+      - name: fetch or make Swift FFI library
+        uses: ./.github/actions/make
         with:
+          key: swift-ffi-library
+          make-rule: swift-ffi-library
+          target-path: target/swift-bindgen/release/libcore_crypto_ffi.dylib
           gh-token: ${{ inputs.gh-token }}
+          uses-rust: true
           artifact-generation: ${{ inputs.artifact-generation }}
 
       - uses: ./.github/actions/make

.github/actions/make/ffi-library/action.yml

@@ -1,4 +1,4 @@
-name: Fetch or make ffi library
+name: Fetch or make ffi library for linux
 
 inputs:
   gh-token:
@@ -10,23 +10,11 @@ runs:
   using: composite
 
   steps:
-    - name: ffi library extension
-      run: |
-        os="$(uname -s)"
-        echo "ARTIFACT_NAME=uniffi-${os}" >> $GITHUB_ENV
-        if [ "$os" = "Linux" ]; then
-          library_extension="so"
-        elif [ "$os" = "Darwin" ]; then
-          library_extension="dylib"
-        fi
-        echo "LIBRARY_EXTENSION=${library_extension}" >> $GITHUB_ENV
-      shell: bash -euo pipefail {0}
-
     - uses: ./.github/actions/make
       with:
-        key: libcore_crypto_ffi.${{ env.LIBRARY_EXTENSION }}
+        key: libcore_crypto_ffi.so
         make-rule: ffi-library
-        target-path: target/release/libcore_crypto_ffi.${{ env.LIBRARY_EXTENSION }}
+        target-path: target/release/libcore_crypto_ffi.so
         gh-token: ${{ inputs.gh-token }}
         uses-rust: true
         rust-target: x86_64-unknown-linux-gnu

crypto-ffi/Cargo.toml

@@ -21,6 +21,7 @@ workspace = true
 
 [features]
 default = ["proteus"]
+cancellable-transactions = []
 proteus = ["core-crypto/proteus", "dep:proteus-wasm"]
 wasm = ["dep:uniffi_ubrn"]
 napi = ["dep:uniffi_ubrn"]

crypto-ffi/bindings/swift/WireCoreCrypto/WireCoreCrypto/CoreCrypto.swift

@@ -63,8 +63,19 @@ public final class CoreCrypto: CoreCryptoProtocol, @unchecked Sendable {
         let filePath = dbLocation.map { FilePath(stringLiteral: $0) }
         let transactionExecutor = try TransactionExecutor<Result>(
             keystorePath: filePath, block)
+        let token = CoreCryptoCancellationToken()
+
         do {
-            try await coreCryptoFfi.transactionFfi(command: transactionExecutor)
+            try await withTaskCancellationHandler {
+                try await coreCryptoFfi.transactionFfiCancellable(
+                    command: transactionExecutor,
+                    cancellation: token
+                )
+            } onCancel: {
+                token.cancel()
+            }
+        } catch CoreCryptoError.TransactionCanceled {
+            throw CancellationError()
         } catch {
             throw await transactionExecutor.innerError ?? error
         }

crypto-ffi/bindings/swift/WireCoreCrypto/WireCoreCryptoTests/WireCoreCryptoTests.swift

@@ -294,6 +294,96 @@ final class WireCoreCryptoTests: XCTestCase {
         )
     }
 
+    // swiftlint:disable:next function_body_length
+    func testCancellingTransactionStopsWaitingForLongRunningCallbacks() async throws {
+        // Verifies that cancellation stops waiting for PKI and MLS transport callbacks.
+        let database = try await newDatabase()
+
+        // This struct implements send commit bundle and x509 credential finalization by just waiting.
+        // Also, we're passing in expectations that are used as checkpoints.
+        let sendCommitStarted = expectation(description: "send commit started")
+        let sendCommitExited = expectation(description: "send commit exited")
+        let httpRequestStarted = expectation(description: "http request started")
+        let httpRequestExited = expectation(description: "http request exited")
+        let longRunningCallbacks = LongRunningCallbacks(
+            sencCommitStarted: sendCommitStarted,
+            sencCommitExited: sendCommitExited,
+            httpRequestStarted: httpRequestStarted,
+            httpRequestExited: httpRequestExited
+        )
+
+        let clientId = genClientId()
+        let coreCrypto = try CoreCrypto(database: database)
+
+        let pkiEnvironment = try await PkiEnvironment(
+            hooks: longRunningCallbacks, database: database)
+        await coreCrypto.setPkiEnvironment(pkiEnvironment: pkiEnvironment)
+        let acquisition = try X509CredentialAcquisition(
+            pkiEnvironment: pkiEnvironment,
+            config: X509CredentialAcquisitionConfiguration(
+                acmeDirectoryUrl: "https://acme.example.com/directory",
+                cipherSuite: cipherSuiteDefault(),
+                displayName: "Alice Smith",
+                clientId: clientId,
+                handle: "alice_wire",
+                domain: "world.com",
+                team: nil,
+                validityPeriodSecs: 3600
+            )
+        )
+
+        let conversationId = genConversationId()
+        let credential = try Credential.basic(
+            cipherSuite: cipherSuiteDefault(),
+            clientId: clientId
+        )
+
+        try await coreCrypto.transaction { context in
+            try await context.mlsInit(clientId: clientId, transport: longRunningCallbacks)
+            let credentialRef = try await context.addCredential(credential: credential)
+            try await context.createConversation(
+                conversationId: conversationId,
+                credentialRef: credentialRef,
+                externalSender: nil
+            )
+        }
+
+        let transactionTask = Task {
+            try await coreCrypto.transaction { context in
+                try await context.setData(data: Data("This data should not be committed.".utf8))
+
+                // Start the long-running callbacks concurrently within the cancellable transaction.
+                async let update: () = context.updateKeyingMaterial(
+                    conversationId: conversationId
+                )
+                async let acquiredCredential = acquisition.finalize()
+                _ = try await (update, acquiredCredential)
+            }
+        }
+
+        // Wait until both long-running callbacks have started before cancelling.
+        await fulfillment(of: [sendCommitStarted, httpRequestStarted])
+        transactionTask.cancel()
+
+        // Assert that cancellation is triggered without waiting for either callback result.
+        do {
+            try await transactionTask.value
+            XCTFail("Expected the transaction to be cancelled")
+        } catch is CancellationError {
+            // Expected.
+        } catch {
+            XCTFail("Expected CancellationError, got \(error)")
+        }
+
+        await fulfillment(of: [sendCommitExited, httpRequestExited])
+
+        // Assert that work performed before the callbacks was rolled back.
+        let data = try await coreCrypto.transaction { context in
+            try await context.getData()
+        }
+        XCTAssertNil(data)
+    }
+
     func testParallelTransactionsArePerformedSeriallyAcrossMultipleCoreCryptoInstances()
         async throws
     {
@@ -853,6 +943,79 @@ final class WireCoreCryptoTests: XCTestCase {
         }
     }
 
+    private final actor LongRunningCallbacks: MlsTransport, PkiEnvironmentHooks {
+        private let sencCommitStarted: XCTestExpectation
+        private let sendCommitExited: XCTestExpectation
+        private let httpRequestStarted: XCTestExpectation
+        private let httpRequestExited: XCTestExpectation
+
+        init(
+            sencCommitStarted: XCTestExpectation,
+            sencCommitExited: XCTestExpectation,
+            httpRequestStarted: XCTestExpectation,
+            httpRequestExited: XCTestExpectation
+        ) {
+            self.sencCommitStarted = sencCommitStarted
+            self.sendCommitExited = sencCommitExited
+            self.httpRequestStarted = httpRequestStarted
+            self.httpRequestExited = httpRequestExited
+        }
+
+        func sendCommitBundle(commitBundle: CommitBundle) async {
+            sencCommitStarted.fulfill()
+            do {
+                try await Task.sleep(for: .milliseconds(9999))
+                XCTFail("Expected the Task.sleep() to be cancelled")
+            } catch is CancellationError {
+                sendCommitExited.fulfill()
+            } catch {
+                XCTFail("Unexpected error: \(error)")
+            }
+        }
+
+        func prepareForTransport(historySecret: WireCoreCryptoUniffi.HistorySecret) async
+            -> WireCoreCryptoUniffi.MlsTransportData
+        {
+            Data()
+        }
+
+        func httpRequest(
+            method: HttpMethod,
+            url: String,
+            headers: [HttpHeader],
+            body: Data
+        ) async -> HttpResponse {
+            httpRequestStarted.fulfill()
+            do {
+                try await Task.sleep(for: .milliseconds(9999))
+                XCTFail("Expected Task.sleep() to be cancelled")
+            } catch is CancellationError {
+                httpRequestExited.fulfill()
+            } catch {
+                XCTFail("Unexpected error: \(error)")
+            }
+
+            return HttpResponse(status: 200, headers: [], body: Data())
+        }
+
+        func authenticate(
+            idp: String,
+            keyAuth: String,
+            acmeAud: String,
+            acquisitionSnapshot: Data
+        ) async -> String {
+            "mock-id-token"
+        }
+
+        func getBackendNonce() async -> String {
+            "mock-backend-nonce"
+        }
+
+        func fetchBackendAccessToken(dpop: String) async -> String {
+            "mock-backend-access-token"
+        }
+    }
+
     private func newDatabase() async throws -> Database {
         let root = FileManager.default.temporaryDirectory.appending(path: "mls")
         let keystore = root.appending(path: "keystore-\(UUID().uuidString)")

crypto-ffi/src/cancellation/future.rs

@@ -0,0 +1,43 @@
+use std::{
+    pin::Pin,
+    sync::Arc,
+    task::{Context, Poll},
+};
+
+use crate::CoreCryptoCancellationToken;
+
+pub(crate) struct Cancelled {
+    pub(crate) token: Arc<CoreCryptoCancellationToken>,
+    pub(crate) wakers_index: Option<usize>,
+}
+
+impl Future for Cancelled {
+    type Output = ();
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        let this = self.get_mut();
+
+        if this.token.is_cancelled() {
+            return Poll::Ready(());
+        }
+
+        let mut wakers = this.token.wakers();
+        // Cancellation may have happened while waiting for the wakers lock.
+        if this.token.is_cancelled() {
+            return Poll::Ready(());
+        }
+
+        wakers.insert(&mut this.wakers_index, cx.waker());
+        Poll::Pending
+    }
+}
+
+impl Drop for Cancelled {
+    fn drop(&mut self) {
+        let Some(index) = self.wakers_index.take() else {
+            return;
+        };
+
+        self.token.wakers().remove(index);
+    }
+}

crypto-ffi/src/cancellation/mod.rs

@@ -0,0 +1,7 @@
+mod future;
+mod slot;
+mod token;
+
+pub(crate) use future::Cancelled;
+pub(crate) use slot::CancellationSlot;
+pub use token::CoreCryptoCancellationToken;

crypto-ffi/src/cancellation/slot.rs

@@ -0,0 +1,57 @@
+use std::sync::{Arc, Mutex};
+
+use crate::{CoreCryptoCancellationToken, CoreCryptoError, CoreCryptoResult};
+
+/// Makes the current transaction's cancellation token available to foreign callbacks.
+///
+/// Cancelling the outer transaction does not automatically stop a nested Rust call that is
+/// waiting for a Swift callback. The transport uses this token to stop waiting, which
+/// lets the nested call return and the transaction unwind.
+#[derive(Debug, Default)]
+pub(crate) struct CancellationSlot {
+    current: Mutex<Option<Arc<CoreCryptoCancellationToken>>>,
+}
+
+/// Clears the active cancellation token from the slot when dropped.
+#[derive(Debug)]
+pub(crate) struct CancellationGuard {
+    slot: Arc<CancellationSlot>,
+}
+
+impl CancellationSlot {
+    /// Installs the token for the transaction currently holding the semaphore.
+    ///
+    /// Only one token may be installed at a time.
+    pub(crate) fn enter(
+        self: &Arc<Self>,
+        token: Arc<CoreCryptoCancellationToken>,
+    ) -> CoreCryptoResult<CancellationGuard> {
+        let mut current = self.current.lock().map_err(CoreCryptoError::ad_hoc)?;
+
+        assert!(
+            current.is_none(),
+            "only one transaction cancellation token may be in the slot; correct wrapper implementation never hits this"
+        );
+
+        *current = Some(token);
+        Ok(CancellationGuard { slot: self.clone() })
+    }
+
+    pub(crate) fn current(&self) -> CoreCryptoResult<Option<Arc<CoreCryptoCancellationToken>>> {
+        self.current
+            .lock()
+            .map_err(CoreCryptoError::ad_hoc)
+            .map(|guard| guard.clone())
+    }
+}
+
+impl Drop for CancellationGuard {
+    fn drop(&mut self) {
+        let mut current = self.slot.current.lock().unwrap_or_else(|poisoned| {
+            log::warn!("recovering poisoned cancellation slot");
+            poisoned.into_inner()
+        });
+
+        *current = None;
+    }
+}