From 834d2abb95e8339f7e1b73fbfa3cfc219791c822 Mon Sep 17 00:00:00 2001
From: Maicss <maicss@126.com>
Date: Sat, 20 Jul 2024 13:08:53 +0800
Subject: [PATCH] Add callback type parameters to the callback function to
 facilitate the user to receive Close and Pong messages

---
 easywsclient.cpp | 15 ++++++++++-----
 easywsclient.hpp |  9 +++++----
 2 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/easywsclient.cpp b/easywsclient.cpp
index 0c98098..bc7489a 100644
--- a/easywsclient.cpp
+++ b/easywsclient.cpp
@@ -259,9 +259,9 @@ class _RealWebSocket : public easywsclient::WebSocket
         {
             Callback_Imp& callable;
             CallbackAdapter(Callback_Imp& callable) : callable(callable) { }
-            void operator()(const std::vector<uint8_t>& message) {
+            void operator()(easywsclient::callbackType type, const std::vector<uint8_t>& message) {
                 std::string stringMessage(message.begin(), message.end());
-                callable(stringMessage);
+                callable(type, stringMessage);
             }
         };
         CallbackAdapter bytesCallback(callable);
@@ -347,7 +347,7 @@ class _RealWebSocket : public easywsclient::WebSocket
                 if (ws.mask) { for (size_t i = 0; i != ws.N; ++i) { rxbuf[i+ws.header_size] ^= ws.masking_key[i&0x3]; } }
                 receivedData.insert(receivedData.end(), rxbuf.begin()+ws.header_size, rxbuf.begin()+ws.header_size+(size_t)ws.N);// just feed
                 if (ws.fin) {
-                    callable((const std::vector<uint8_t>) receivedData);
+                    callable(easywsclient::callbackType::MESSAGE, (const std::vector<uint8_t>) receivedData);
                     receivedData.erase(receivedData.begin(), receivedData.end());
                     std::vector<uint8_t> ().swap(receivedData);// free memory
                 }
@@ -357,8 +357,13 @@ class _RealWebSocket : public easywsclient::WebSocket
                 std::string data(rxbuf.begin()+ws.header_size, rxbuf.begin()+ws.header_size+(size_t)ws.N);
                 sendData(wsheader_type::PONG, data.size(), data.begin(), data.end());
             }
-            else if (ws.opcode == wsheader_type::PONG) { }
-            else if (ws.opcode == wsheader_type::CLOSE) { close(); }
+            else if (ws.opcode == wsheader_type::PONG) { 
+                callable(easywsclient::callbackType::PONG, {});
+            }
+            else if (ws.opcode == wsheader_type::CLOSE) {
+                callable(easywsclient::callbackType::CLOSE, {});
+                close();
+            }
             else { fprintf(stderr, "ERROR: Got unexpected WebSocket message.\n"); close(); }
 
             rxbuf.erase(rxbuf.begin(), rxbuf.begin() + ws.header_size+(size_t)ws.N);
diff --git a/easywsclient.hpp b/easywsclient.hpp
index 08c4a7b..f862234 100644
--- a/easywsclient.hpp
+++ b/easywsclient.hpp
@@ -12,9 +12,10 @@
 #include <vector>
 
 namespace easywsclient {
+typedef enum callbackType { MESSAGE, PONG, CLOSE } callbackType;
 
-struct Callback_Imp { virtual void operator()(const std::string& message) = 0; };
-struct BytesCallback_Imp { virtual void operator()(const std::vector<uint8_t>& message) = 0; };
+struct Callback_Imp { virtual void operator()(callbackType type, const std::string& message) = 0; };
+struct BytesCallback_Imp { virtual void operator()(callbackType type, const std::vector<uint8_t>& message) = 0; };
 
 class WebSocket {
   public:
@@ -43,7 +44,7 @@ class WebSocket {
         struct _Callback : public Callback_Imp {
             Callable& callable;
             _Callback(Callable& callable) : callable(callable) { }
-            void operator()(const std::string& message) { callable(message); }
+            void operator()(callbackType type, const std::string& message) { callable(type, message); }
         };
         _Callback callback(callable);
         _dispatch(callback);
@@ -56,7 +57,7 @@ class WebSocket {
         struct _Callback : public BytesCallback_Imp {
             Callable& callable;
             _Callback(Callable& callable) : callable(callable) { }
-            void operator()(const std::vector<uint8_t>& message) { callable(message); }
+            void operator()(callbackType type, const std::vector<uint8_t>& message) { callable(type, message); }
         };
         _Callback callback(callable);
         _dispatchBinary(callback);