From f40de822dde0dfe4273e7d6ee4306dbc4d363d44 Mon Sep 17 00:00:00 2001
From: Aaron Piotrowski <aaron@trowski.com>
Date: Wed, 13 Sep 2023 23:17:20 -0500
Subject: [PATCH] Fix onClose handler on socket disconnect

---
 src/Internal/Rfc6455FrameHandler.php | 15 +++++++++++++++
 src/Rfc6455Client.php                |  6 +-----
 2 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/src/Internal/Rfc6455FrameHandler.php b/src/Internal/Rfc6455FrameHandler.php
index 55b59e0..d2fcd11 100644
--- a/src/Internal/Rfc6455FrameHandler.php
+++ b/src/Internal/Rfc6455FrameHandler.php
@@ -7,6 +7,7 @@
 use Amp\DeferredFuture;
 use Amp\ForbidCloning;
 use Amp\ForbidSerialization;
+use Amp\Future;
 use Amp\Pipeline\ConcurrentIterator;
 use Amp\Pipeline\DisposedException;
 use Amp\Pipeline\Queue;
@@ -299,6 +300,20 @@ public function close(int $code, string $reason = '', bool $byPeer = false): voi
         $this->socket->close();
     }
 
+    /**
+     * @param \Closure(int, WebsocketCloseInfo):void $onClose
+     */
+    public function onClose(\Closure $onClose): void
+    {
+        $future = $this->closeDeferred?->getFuture() ?? Future::complete();
+
+        $metadata = $this->metadata;
+        $future->finally(static function () use ($onClose, $metadata): void {
+            \assert($metadata->closeInfo !== null, 'Client was not closed when onClose invoked');
+            $onClose($metadata->id, $metadata->closeInfo);
+        });
+    }
+
     private static function createMessage(WebsocketFrameType $frameType, ReadableStream|string $stream): WebsocketMessage
     {
         if ($frameType === WebsocketFrameType::Binary) {
diff --git a/src/Rfc6455Client.php b/src/Rfc6455Client.php
index d0568ad..92a3bcb 100644
--- a/src/Rfc6455Client.php
+++ b/src/Rfc6455Client.php
@@ -307,10 +307,6 @@ public function close(int $code = WebsocketCloseCode::NORMAL_CLOSE, string $reas
 
     public function onClose(\Closure $onClose): void
     {
-        $metadata = $this->metadata;
-        $this->socket->onClose(static function () use ($onClose, $metadata): void {
-            \assert($metadata->closeInfo !== null, 'Client was not closed when onClose invoked');
-            $onClose($metadata->id, $metadata->closeInfo);
-        });
+        $this->frameHandler->onClose($onClose);
     }
 }