From 389a98140b7c44893ded5a0d956855ae8facc2ea Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?S=C3=A9rgio=20Santos?= <me@sergiosantos.info>
Date: Mon, 26 Jul 2021 16:19:31 +0100
Subject: [PATCH] fix: ReceiveMessages premature client close (#134)

Fixes https://github.com/relaycorp/relaynet-gateway-android/issues/411
---
 .../awaladroid/messaging/ReceiveMessages.kt   | 24 +++++++++++--------
 1 file changed, 14 insertions(+), 10 deletions(-)

diff --git a/lib/src/main/java/tech/relaycorp/awaladroid/messaging/ReceiveMessages.kt b/lib/src/main/java/tech/relaycorp/awaladroid/messaging/ReceiveMessages.kt
index 8455906b..2ceb9469 100644
--- a/lib/src/main/java/tech/relaycorp/awaladroid/messaging/ReceiveMessages.kt
+++ b/lib/src/main/java/tech/relaycorp/awaladroid/messaging/ReceiveMessages.kt
@@ -5,6 +5,7 @@ import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.asFlow
 import kotlinx.coroutines.flow.flatMapLatest
 import kotlinx.coroutines.flow.mapNotNull
+import kotlinx.coroutines.flow.onCompletion
 import tech.relaycorp.awaladroid.Awala
 import tech.relaycorp.awaladroid.GatewayException
 import tech.relaycorp.awaladroid.GatewayProtocolException
@@ -36,16 +37,19 @@ internal class ReceiveMessages(
     fun receive(): Flow<IncomingMessage> =
         getNonceSigners()
             .flatMapLatest { nonceSigners ->
-                pdcClientBuilder().use {
-                    try {
-                        collectParcels(it, nonceSigners)
-                    } catch (exp: ServerException) {
-                        throw ReceiveMessageException("Server error", exp)
-                    } catch (exp: ClientBindingException) {
-                        throw GatewayProtocolException("Client error", exp)
-                    } catch (exp: NonceSignerException) {
-                        throw GatewayProtocolException("Client signing error", exp)
-                    }
+                val pdcClient = pdcClientBuilder()
+                try {
+                    collectParcels(pdcClient, nonceSigners)
+                        .onCompletion {
+                            @Suppress("BlockingMethodInNonBlockingContext")
+                            pdcClient.close()
+                        }
+                } catch (exp: ServerException) {
+                    throw ReceiveMessageException("Server error", exp)
+                } catch (exp: ClientBindingException) {
+                    throw GatewayProtocolException("Client error", exp)
+                } catch (exp: NonceSignerException) {
+                    throw GatewayProtocolException("Client signing error", exp)
                 }
             }