From 488360db951fc805061b00600a222bf697f65f61 Mon Sep 17 00:00:00 2001
From: Nicklas Warming Jacobsen <nicklas.jacobsen@ocrlabs.com>
Date: Tue, 30 Jan 2024 12:34:54 +0100
Subject: [PATCH] Fixes: #462: In
 com.android.identity.wallet.transfer.Communication sendResponse If
 `closeAfterSending` is true, only close the connection after we know the
 `deviceResponse` payload have been sent.

---
 .../fragment/TransferDocumentFragment.kt      |  4 ++--
 .../identity/wallet/transfer/Communication.kt | 23 ++++++++-----------
 .../wallet/transfer/TransferManager.kt        | 12 +++++-----
 3 files changed, 17 insertions(+), 22 deletions(-)

diff --git a/appholder/src/main/java/com/android/identity/wallet/fragment/TransferDocumentFragment.kt b/appholder/src/main/java/com/android/identity/wallet/fragment/TransferDocumentFragment.kt
index 886dba211..0751d82c3 100644
--- a/appholder/src/main/java/com/android/identity/wallet/fragment/TransferDocumentFragment.kt
+++ b/appholder/src/main/java/com/android/identity/wallet/fragment/TransferDocumentFragment.kt
@@ -163,13 +163,13 @@ class TransferDocumentFragment : Fragment() {
     private fun onTransferDisconnected() {
         log("Disconnected")
         hideButtons()
-        TransferManager.getInstance(requireContext()).disconnect()
+        TransferManager.getInstance(requireContext()).disconnect(true)
     }
 
     private fun onTransferError() {
         Toast.makeText(requireContext(), "An error occurred.", Toast.LENGTH_SHORT).show()
         hideButtons()
-        TransferManager.getInstance(requireContext()).disconnect()
+        TransferManager.getInstance(requireContext()).disconnect(true)
     }
 
     private fun hideButtons() {
diff --git a/appholder/src/main/java/com/android/identity/wallet/transfer/Communication.kt b/appholder/src/main/java/com/android/identity/wallet/transfer/Communication.kt
index 3b2d04891..e53de72b0 100644
--- a/appholder/src/main/java/com/android/identity/wallet/transfer/Communication.kt
+++ b/appholder/src/main/java/com/android/identity/wallet/transfer/Communication.kt
@@ -45,21 +45,16 @@ class Communication private constructor(
             if (progress == max) {
                 log("Completed...")
             }
+            if (progress == max && closeAfterSending) {
+                deviceRetrievalHelper?.disconnect()
+            }
         }
-        if (closeAfterSending) {
-            deviceRetrievalHelper?.sendDeviceResponse(
-                deviceResponse,
-                OptionalLong.of(Constants.SESSION_DATA_STATUS_SESSION_TERMINATION),
-                progressListener,
-                context.mainExecutor())
-            deviceRetrievalHelper?.disconnect()
-        } else {
-            deviceRetrievalHelper?.sendDeviceResponse(
-                deviceResponse,
-                OptionalLong.empty(),
-                progressListener,
-                context.mainExecutor())
-        }
+
+        deviceRetrievalHelper?.sendDeviceResponse(
+            deviceResponse,
+            OptionalLong.of(Constants.SESSION_DATA_STATUS_SESSION_TERMINATION),
+            progressListener,
+            context.mainExecutor())
     }
 
     fun stopPresentation(
diff --git a/appholder/src/main/java/com/android/identity/wallet/transfer/TransferManager.kt b/appholder/src/main/java/com/android/identity/wallet/transfer/TransferManager.kt
index f57797695..3376c2763 100644
--- a/appholder/src/main/java/com/android/identity/wallet/transfer/TransferManager.kt
+++ b/appholder/src/main/java/com/android/identity/wallet/transfer/TransferManager.kt
@@ -244,11 +244,13 @@ class TransferManager private constructor(private val context: Context) {
             sendSessionTerminationMessage,
             useTransportSpecificSessionTermination
         )
-        disconnect()
+        disconnect(true)
     }
 
-    fun disconnect() {
-        communication.disconnect()
+    fun disconnect(disconnectCommunication: Boolean) {
+        if (disconnectCommunication) {
+            communication.disconnect()
+        }
         qrCommunicationSetup?.close()
         transferStatusLd = MutableLiveData<TransferStatus>()
         destroy()
@@ -262,9 +264,7 @@ class TransferManager private constructor(private val context: Context) {
 
     fun sendResponse(deviceResponse: ByteArray, closeAfterSending: Boolean) {
         communication.sendResponse(deviceResponse, closeAfterSending)
-        if (closeAfterSending) {
-            disconnect()
-        }
+        disconnect(!closeAfterSending)
     }
 
     fun readDocumentEntries(documentName: String): DocumentElements {