fix: ReceiveMessages premature client close (#134)

Fixes relaycorp/relaynet-gateway-android#411
relaycorp · Jul 26, 2021 · 389a981 · 389a981
1 parent 91846d2
commit 389a981
Showing 1 changed file with 14 additions and 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
@@ -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)
                 }
             }