From 604c28af54e2d29cac01d77054b9891425ad35cd Mon Sep 17 00:00:00 2001
From: joffrey-bion <joffrey.bion@gmail.com>
Date: Mon, 17 Jan 2022 00:02:19 +0100
Subject: [PATCH] Prevent exceptions due to heartbeats

Resolves:
https://github.com/joffrey-bion/krossbow/issues/173
---
 .../org/hildan/krossbow/stomp/BaseStompSession.kt     | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/krossbow-stomp-core/src/commonMain/kotlin/org/hildan/krossbow/stomp/BaseStompSession.kt b/krossbow-stomp-core/src/commonMain/kotlin/org/hildan/krossbow/stomp/BaseStompSession.kt
index 25dec62f2..16d43e3c0 100644
--- a/krossbow-stomp-core/src/commonMain/kotlin/org/hildan/krossbow/stomp/BaseStompSession.kt
+++ b/krossbow-stomp-core/src/commonMain/kotlin/org/hildan/krossbow/stomp/BaseStompSession.kt
@@ -13,6 +13,7 @@ import org.hildan.krossbow.stomp.headers.*
 import org.hildan.krossbow.stomp.heartbeats.HeartBeater
 import org.hildan.krossbow.stomp.heartbeats.NO_HEART_BEATS
 import org.hildan.krossbow.utils.generateUuid
+import org.hildan.krossbow.websocket.WebSocketException
 import kotlin.coroutines.CoroutineContext
 import kotlin.coroutines.EmptyCoroutineContext
 import kotlin.coroutines.cancellation.CancellationException
@@ -35,7 +36,15 @@ internal class BaseStompSession(
         HeartBeater(
             heartBeat = heartBeat,
             tolerance = config.heartBeatTolerance,
-            sendHeartBeat = { stompSocket.sendHeartBeat() },
+            sendHeartBeat = {
+                // The web socket could have errored or be closed, and the heart beater's job not yet be cancelled.
+                // In this case, we don't want the heart beater to crash
+                try {
+                    stompSocket.sendHeartBeat()
+                } catch(e : WebSocketException) {
+                    shutdown("STOMP session failed: heart beat couldn't be sent", cause = e)
+                }
+            },
             onMissingHeartBeat = {
                 val cause = MissingHeartBeatException(heartBeat.expectedPeriod)
                 sharedStompEvents.emit(StompEvent.Error(cause))