From 3650b8d17b97790644c5fab27dbae02a23d26d9c Mon Sep 17 00:00:00 2001
From: cameronvoell <cameron@xmtp.com>
Date: Fri, 22 Dec 2023 13:43:16 -0500
Subject: [PATCH] fix: sign actions should wait for pre-callback completion
 before proceeding

---
 .../modules/xmtpreactnativesdk/XMTPModule.kt  | 30 +++++++++++++++++++
 src/index.ts                                  |  8 +++++
 src/lib/Client.ts                             |  2 ++
 3 files changed, 40 insertions(+)

diff --git a/android/src/main/java/expo/modules/xmtpreactnativesdk/XMTPModule.kt b/android/src/main/java/expo/modules/xmtpreactnativesdk/XMTPModule.kt
index c531c9016..10b6e2853 100644
--- a/android/src/main/java/expo/modules/xmtpreactnativesdk/XMTPModule.kt
+++ b/android/src/main/java/expo/modules/xmtpreactnativesdk/XMTPModule.kt
@@ -18,6 +18,7 @@ import expo.modules.xmtpreactnativesdk.wrappers.DecryptedLocalAttachment
 import expo.modules.xmtpreactnativesdk.wrappers.EncryptedLocalAttachment
 import expo.modules.xmtpreactnativesdk.wrappers.PreparedLocalMessage
 import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.delay
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.Job
 import kotlinx.coroutines.launch
@@ -126,6 +127,8 @@ class XMTPModule : Module() {
     private val isDebugEnabled = BuildConfig.DEBUG // TODO: consider making this configurable
     private val conversations: MutableMap<String, Conversation> = mutableMapOf()
     private val subscriptions: MutableMap<String, Job> = mutableMapOf()
+    var waitForPreEnableIdentityCallback: Boolean = false
+    var waitForPreCreateIdentityCallback: Boolean = false
 
     override fun definition() = ModuleDefinition {
         Name("XMTP")
@@ -151,6 +154,9 @@ class XMTPModule : Module() {
             logV("auth")
             val reactSigner = ReactNativeSigner(module = this@XMTPModule, address = address)
             signer = reactSigner
+
+            waitForPreEnableIdentityCallback = hasEnableIdentityCallback == true
+            waitForPreCreateIdentityCallback = hasCreateIdentityCallback == true
             val preCreateIdentityCallback: PreEventCallback? =
                 preCreateIdentityCallback.takeIf { hasCreateIdentityCallback == true }
             val preEnableIdentityCallback: PreEventCallback? =
@@ -175,6 +181,9 @@ class XMTPModule : Module() {
         AsyncFunction("createRandom") { environment: String, appVersion: String?, hasCreateIdentityCallback: Boolean?, hasEnableIdentityCallback: Boolean? ->
             logV("createRandom")
             val privateKey = PrivateKeyBuilder()
+
+            waitForPreEnableIdentityCallback = hasEnableIdentityCallback == true
+            waitForPreCreateIdentityCallback = hasCreateIdentityCallback == true
             val preCreateIdentityCallback: PreEventCallback? =
                 preCreateIdentityCallback.takeIf { hasCreateIdentityCallback == true }
             val preEnableIdentityCallback: PreEventCallback? =
@@ -590,6 +599,18 @@ class XMTPModule : Module() {
             val client = clients[clientAddress] ?: throw XMTPException("No client")
             client.contacts.consentList.entries.map { ConsentWrapper.encode(it.value) }
         }
+
+        Function("preEnableIdentityCallbackCompleted") { 
+            logV("preEnableIdentityCallbackCompleted")
+            waitForPreEnableIdentityCallback = false
+            true
+        }
+
+        Function("preCreateIdentityCallbackCompleted") {
+            logV("preCreateIdentityCallbackCompleted")
+            waitForPreCreateIdentityCallback = false
+            true
+        }
     }
 
     //
@@ -715,10 +736,19 @@ class XMTPModule : Module() {
 
     private val preEnableIdentityCallback: suspend () -> Unit = {
         sendEvent("preEnableIdentityCallback")
+        waitForCallback { waitForPreEnableIdentityCallback }
     }
 
     private val preCreateIdentityCallback: suspend () -> Unit = {
         sendEvent("preCreateIdentityCallback")
+        waitForCallback { waitForPreCreateIdentityCallback }
+    }
+
+    // Helper function to wait for a callback
+    private suspend fun waitForCallback(check: () -> Boolean) {
+        while (check()) {
+            delay(100) // Wait for 100ms before checking again
+        }
     }
 }
 
diff --git a/src/index.ts b/src/index.ts
index 5e324ad8b..80c8a5834 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -393,6 +393,14 @@ export async function consentList(
   })
 }
 
+export function preEnableIdentityCallbackCompleted() {
+  XMTPModule.preEnableIdentityCallbackCompleted()
+}
+
+export function preCreateIdentityCallbackCompleted() {
+  XMTPModule.preCreateIdentityCallbackCompleted()
+}
+
 export const emitter = new EventEmitter(XMTPModule ?? NativeModulesProxy.XMTP)
 
 export * from './lib/ContentCodec'
diff --git a/src/lib/Client.ts b/src/lib/Client.ts
index 745e333b2..8ea8b96b5 100644
--- a/src/lib/Client.ts
+++ b/src/lib/Client.ts
@@ -236,6 +236,7 @@ export class Client<ContentTypes> {
       opts,
       async () => {
         await this.executeCallback(opts?.preEnableIdentityCallback)
+        XMTPModule.preEnableIdentityCallbackCompleted()
       }
     )
 
@@ -244,6 +245,7 @@ export class Client<ContentTypes> {
       opts,
       async () => {
         await this.executeCallback(opts?.preCreateIdentityCallback)
+        XMTPModule.preCreateIdentityCallbackCompleted()
       }
     )