From 515217652a0969f1b02526df44a19b69163097f7 Mon Sep 17 00:00:00 2001
From: Ako Tulu <ako@koodinurk.ee>
Date: Tue, 3 Sep 2024 02:00:31 +0300
Subject: [PATCH] Implemented WebSocket connected callback.

---
 src/Connection/TcpConnection.php |  9 ++++++++-
 src/Protocols/Websocket.php      | 16 ++++++++++++++--
 src/Worker.php                   |  9 ++++++++-
 3 files changed, 30 insertions(+), 4 deletions(-)

diff --git a/src/Connection/TcpConnection.php b/src/Connection/TcpConnection.php
index 35654099f..c76046b0d 100644
--- a/src/Connection/TcpConnection.php
+++ b/src/Connection/TcpConnection.php
@@ -110,12 +110,19 @@ class TcpConnection extends ConnectionInterface implements JsonSerializable
     public $onConnect = null;
 
     /**
-     * Emitted when websocket handshake completed (Only work when protocol is ws).
+     * Emitted before websocket handshake (Only works when protocol is ws).
      *
      * @var ?callable
      */
     public $onWebSocketConnect = null;
 
+    /**
+     * Emitted after websocket handshake (Only works when protocol is ws).
+     *
+     * @var ?callable
+     */
+    public $onWebSocketConnected = null;
+
     /**
      * Emitted when data is received.
      *
diff --git a/src/Protocols/Websocket.php b/src/Protocols/Websocket.php
index eb72e4a3f..6255bde11 100644
--- a/src/Protocols/Websocket.php
+++ b/src/Protocols/Websocket.php
@@ -382,12 +382,14 @@ public static function dealHandshake(string $buffer, TcpConnection $connection):
             $connection->context->websocketCurrentFrameBuffer = '';
             // Consume handshake data.
             $connection->consumeRecvBuffer($headerLength);
+            // Request from buffer
+            $request = new Request($buffer);
 
             // Try to emit onWebSocketConnect callback.
             $onWebsocketConnect = $connection->onWebSocketConnect ?? $connection->worker->onWebSocketConnect ?? false;
             if ($onWebsocketConnect) {
                 try {
-                    $onWebsocketConnect($connection, new Request($buffer));
+                    $onWebsocketConnect($connection, $request);
                 } catch (Throwable $e) {
                     Worker::stopAll(250, $e);
                 }
@@ -413,10 +415,20 @@ public static function dealHandshake(string $buffer, TcpConnection $connection):
             }
             $handshakeMessage .= "\r\n";
             // Send handshake response.
-            $connection->send($handshakeMessage, true);
+            $status = $connection->send($handshakeMessage, true);
             // Mark handshake complete.
             $connection->context->websocketHandshake = true;
 
+            // Try to emit onWebSocketConnected callback.
+            $onWebsocketConnected = $connection->onWebSocketConnected ?? $connection->worker->onWebSocketConnected ?? false;
+            if ($status && $onWebsocketConnected) {
+                try {
+                    $onWebsocketConnected($connection, $request);
+                } catch (Throwable $e) {
+                    Worker::stopAll(250, $e);
+                }
+            }
+
             // There are data waiting to be sent.
             if (!empty($connection->context->tmpWebsocketData)) {
                 $connection->send($connection->context->tmpWebsocketData, true);
diff --git a/src/Worker.php b/src/Worker.php
index 8ed426fda..a610609ad 100644
--- a/src/Worker.php
+++ b/src/Worker.php
@@ -159,12 +159,19 @@ class Worker
     public $onConnect = null;
 
     /**
-     * Emitted when websocket handshake did complete (Only called when protocol is ws).
+     * Emitted before websocket handshake (Only works when protocol is ws).
      *
      * @var ?callable
      */
     public $onWebSocketConnect = null;
 
+    /**
+     * Emitted after websocket handshake (Only works when protocol is ws).
+     *
+     * @var ?callable
+     */
+    public $onWebSocketConnected = null;
+
     /**
      * Emitted when data is received.
      *