Fix of WebSocket close status frame

chronoxor · chronoxor · commit 826e690b01e3 · 2023-11-20T02:03:03.000+04:00
diff --git a/examples/ws_chat_server.cpp b/examples/ws_chat_server.cpp
@@ -42,7 +42,7 @@ class ChatSession : public CppServer::WS::WSSession
 
         // If the buffer starts with '!' the disconnect the current session
         if (message == "!")
-            Close(1000);
+            Close();
     }
 
     void onError(int error, const std::string& category, const std::string& message) override
diff --git a/examples/wss_chat_server.cpp b/examples/wss_chat_server.cpp
@@ -42,7 +42,7 @@ class ChatSession : public CppServer::WS::WSSSession
 
         // If the buffer starts with '!' the disconnect the current session
         if (message == "!")
-            Close(1000);
+            Close();
     }
 
     void onError(int error, const std::string& category, const std::string& message) override
diff --git a/include/server/version.h b/include/server/version.h
@@ -28,7 +28,7 @@ GitHub: https://github.com/chronoxor/CppServer
 namespace CppServer {
 
 //! Project version
-const char version[] = "1.0.4.0";
+const char version[] = "1.0.4.1";
 
 } // namespace CppServer
 
diff --git a/include/server/ws/ws.h b/include/server/ws/ws.h
@@ -142,8 +142,9 @@ class WebSocket
     /*!
         \param buffer - Received buffer
         \param size - Received buffer size
+        \param status - WebSocket status (defualt is 1000)
     */
-    virtual void onWSClose(const void* buffer, size_t size) {}
+    virtual void onWSClose(const void* buffer, size_t size, int status = 1000) {}
     //! Handle WebSocket ping notification
     /*!
         \param buffer - Received buffer
diff --git a/include/server/ws/ws_client.h b/include/server/ws/ws_client.h
@@ -40,8 +40,14 @@ class WSClient : public HTTP::HTTPClient, protected WebSocket
     bool Connect(const std::shared_ptr<Asio::TCPResolver>& resolver) override;
     bool ConnectAsync() override;
     bool ConnectAsync(const std::shared_ptr<Asio::TCPResolver>& resolver) override;
-    virtual bool Close(int status) { SendClose(status, nullptr, 0); HTTPClient::Disconnect(); return true; }
-    virtual bool CloseAsync(int status) { SendCloseAsync(status, nullptr, 0); HTTPClient::DisconnectAsync(); return true; }
+    virtual bool Close() { return Close(0, nullptr, 0); }
+    virtual bool Close(int status) { return Close(status, nullptr, 0); }
+    virtual bool Close(int status, const void* buffer, size_t size) { SendClose(status, buffer, size); HTTPClient::Disconnect(); return true; }
+    virtual bool Close(int status, std::string_view text) { SendClose(status, text); HTTPClient::Disconnect(); return true; }
+    virtual bool CloseAsync() { return CloseAsync(0, nullptr, 0); }
+    virtual bool CloseAsync(int status) { return CloseAsync(status, nullptr, 0); }
+    virtual bool CloseAsync(int status, const void* buffer, size_t size) { SendCloseAsync(status, buffer, size); HTTPClient::DisconnectAsync(); return true; }
+    virtual bool CloseAsync(int status, std::string_view text) { SendCloseAsync(status, text); HTTPClient::DisconnectAsync(); return true; }
 
     // WebSocket send text methods
     size_t SendText(const void* buffer, size_t size) { std::scoped_lock locker(_ws_send_lock); PrepareSendFrame(WS_FIN | WS_TEXT, true, buffer, size); return HTTPClient::Send(_ws_send_buffer.data(), _ws_send_buffer.size()); }
@@ -98,7 +104,7 @@ class WSClient : public HTTP::HTTPClient, protected WebSocket
     void onReceivedResponseError(const HTTP::HTTPResponse& response, const std::string& error) override;
 
     //! Handle WebSocket close notification
-    void onWSClose(const void* buffer, size_t size) override { CloseAsync(1000); }
+    void onWSClose(const void* buffer, size_t size, int status = 1000) override { CloseAsync(); }
     //! Handle WebSocket ping notification
     void onWSPing(const void* buffer, size_t size) override { SendPongAsync(buffer, size); }
     //! Handle WebSocket error notification
diff --git a/include/server/ws/ws_server.h b/include/server/ws/ws_server.h
@@ -38,7 +38,10 @@ class WSServer : public HTTP::HTTPServer, protected WebSocket
     WSServer& operator=(WSServer&&) = delete;
 
     // WebSocket connection methods
-    virtual bool CloseAll(int status);
+    virtual bool CloseAll() { return CloseAll(0, nullptr, 0); }
+    virtual bool CloseAll(int status) { return CloseAll(status, nullptr, 0); }
+    virtual bool CloseAll(int status, const void* buffer, size_t size);
+    virtual bool CloseAll(int status, std::string_view text);
 
     //! Multicast data to all connected WebSocket sessions
     bool Multicast(const void* buffer, size_t size) override;
diff --git a/include/server/ws/ws_session.h b/include/server/ws/ws_session.h
@@ -37,7 +37,10 @@ class WSSession : public HTTP::HTTPSession, protected WebSocket
     WSSession& operator=(WSSession&&) = delete;
 
     // WebSocket connection methods
-    virtual bool Close(int status) { SendCloseAsync(status, nullptr, 0); HTTPSession::Disconnect(); return true; }
+    virtual bool Close() { return Close(0, nullptr, 0); }
+    virtual bool Close(int status) { return Close(status, nullptr, 0); }
+    virtual bool Close(int status, const void* buffer, size_t size) { SendCloseAsync(status, buffer, size); HTTPSession::Disconnect(); return true; }
+    virtual bool Close(int status, std::string_view text) { SendCloseAsync(status, text); HTTPSession::Disconnect(); return true; }
 
     // WebSocket send text methods
     size_t SendText(const void* buffer, size_t size) { std::scoped_lock locker(_ws_send_lock); PrepareSendFrame(WS_FIN | WS_TEXT, false, buffer, size); return HTTPSession::Send(_ws_send_buffer.data(), _ws_send_buffer.size()); }
@@ -93,7 +96,7 @@ class WSSession : public HTTP::HTTPSession, protected WebSocket
     void onReceivedRequestError(const HTTP::HTTPRequest& request, const std::string& error) override;
 
     //! Handle WebSocket close notification
-    void onWSClose(const void* buffer, size_t size) override { Close(1000); }
+    void onWSClose(const void* buffer, size_t size, int status = 1000) override { Close(); }
     //! Handle WebSocket ping notification
     void onWSPing(const void* buffer, size_t size) override { SendPongAsync(buffer, size); }
     //! Handle WebSocket error notification
diff --git a/include/server/ws/wss_client.h b/include/server/ws/wss_client.h
@@ -40,8 +40,14 @@ class WSSClient : public HTTP::HTTPSClient, protected WebSocket
     bool Connect(const std::shared_ptr<Asio::TCPResolver>& resolver) override;
     bool ConnectAsync() override;
     bool ConnectAsync(const std::shared_ptr<Asio::TCPResolver>& resolver) override;
-    virtual bool Close(int status) { SendClose(status, nullptr, 0); HTTPSClient::Disconnect(); return true; }
-    virtual bool CloseAsync(int status) { SendCloseAsync(status, nullptr, 0); HTTPSClient::DisconnectAsync(); return true; }
+    virtual bool Close() { return Close(0, nullptr, 0); }
+    virtual bool Close(int status) { return Close(status, nullptr, 0); }
+    virtual bool Close(int status, const void* buffer, size_t size) { SendClose(status, buffer, size); HTTPSClient::Disconnect(); return true; }
+    virtual bool Close(int status, std::string_view text) { SendClose(status, text); HTTPSClient::Disconnect(); return true; }
+    virtual bool CloseAsync() { return CloseAsync(0, nullptr, 0); }
+    virtual bool CloseAsync(int status) { return CloseAsync(status, nullptr, 0); }
+    virtual bool CloseAsync(int status, const void* buffer, size_t size) { SendCloseAsync(status, buffer, size); HTTPSClient::DisconnectAsync(); return true; }
+    virtual bool CloseAsync(int status, std::string_view text) { SendCloseAsync(status, text); HTTPSClient::DisconnectAsync(); return true; }
 
     // WebSocket send text methods
     size_t SendText(const void* buffer, size_t size) { std::scoped_lock locker(_ws_send_lock); PrepareSendFrame(WS_FIN | WS_TEXT, true, buffer, size); return HTTPSClient::Send(_ws_send_buffer.data(), _ws_send_buffer.size()); }
@@ -98,7 +104,7 @@ class WSSClient : public HTTP::HTTPSClient, protected WebSocket
     void onReceivedResponseError(const HTTP::HTTPResponse& response, const std::string& error) override;
 
     //! Handle WebSocket close notification
-    void onWSClose(const void* buffer, size_t size) override { CloseAsync(1000); }
+    void onWSClose(const void* buffer, size_t size, int status = 1000) override { CloseAsync(); }
     //! Handle WebSocket ping notification
     void onWSPing(const void* buffer, size_t size) override { SendPongAsync(buffer, size); }
     //! Handle WebSocket error notification
diff --git a/include/server/ws/wss_server.h b/include/server/ws/wss_server.h
@@ -38,7 +38,10 @@ class WSSServer : public HTTP::HTTPSServer, protected WebSocket
     WSSServer& operator=(WSSServer&&) = delete;
 
     // WebSocket connection methods
-    virtual bool CloseAll(int status);
+    virtual bool CloseAll() { return CloseAll(0, nullptr, 0); }
+    virtual bool CloseAll(int status) { return CloseAll(status, nullptr, 0); }
+    virtual bool CloseAll(int status, const void* buffer, size_t size);
+    virtual bool CloseAll(int status, std::string_view text);
 
     //! Multicast data to all connected WebSocket sessions
     bool Multicast(const void* buffer, size_t size) override;
diff --git a/include/server/ws/wss_session.h b/include/server/ws/wss_session.h
@@ -37,7 +37,10 @@ class WSSSession : public HTTP::HTTPSSession, protected WebSocket
     WSSSession& operator=(WSSSession&&) = delete;
 
     // WebSocket connection methods
-    virtual bool Close(int status) { SendCloseAsync(status, nullptr, 0); HTTPSSession::Disconnect(); return true; }
+    virtual bool Close() { return Close(0, nullptr, 0); }
+    virtual bool Close(int status) { return Close(status, nullptr, 0); }
+    virtual bool Close(int status, const void* buffer, size_t size) { SendCloseAsync(status, buffer, size); HTTPSSession::Disconnect(); return true; }
+    virtual bool Close(int status, std::string_view text) { SendCloseAsync(status, text); HTTPSSession::Disconnect(); return true; }
 
     // WebSocket send text methods
     size_t SendText(const void* buffer, size_t size) { std::scoped_lock locker(_ws_send_lock); PrepareSendFrame(WS_FIN | WS_TEXT, false, buffer, size); return HTTPSSession::Send(_ws_send_buffer.data(), _ws_send_buffer.size()); }
@@ -93,7 +96,7 @@ class WSSSession : public HTTP::HTTPSSession, protected WebSocket
     void onReceivedRequestError(const HTTP::HTTPRequest& request, const std::string& error) override;
 
     //! Handle WebSocket close notification
-    void onWSClose(const void* buffer, size_t size) override { Close(1000); }
+    void onWSClose(const void* buffer, size_t size, int status = 1000) override { Close(); }
     //! Handle WebSocket ping notification
     void onWSPing(const void* buffer, size_t size) override { SendPongAsync(buffer, size); }
     //! Handle WebSocket error notification
diff --git a/source/server/ws/ws.cpp b/source/server/ws/ws.cpp
@@ -211,6 +211,11 @@ bool WebSocket::PerformServerUpgrade(const HTTP::HTTPRequest& request, HTTP::HTT
 
 void WebSocket::PrepareSendFrame(uint8_t opcode, bool mask, const void* buffer, size_t size, int status)
 {
+    // Check if we need to store additional 2 bytes of close status frame
+    bool store_status = ((opcode & WS_CLOSE) == WS_CLOSE) && ((size > 0) || (status != 0));
+    if (store_status)
+        size += 2;
+
     // Clear the previous WebSocket send buffer
     _ws_send_buffer.clear();
 
@@ -246,10 +251,23 @@ void WebSocket::PrepareSendFrame(uint8_t opcode, bool mask, const void* buffer,
     size_t offset = _ws_send_buffer.size();
     _ws_send_buffer.resize(offset + size);
 
-    // Mask WebSocket frame content
+    size_t index = 0;
     const uint8_t* data = (const uint8_t*)buffer;
-    for (size_t i = 0; i < size; ++i)
-        _ws_send_buffer[offset + i] = data[i] ^ _ws_send_mask[i % 4];
+
+    // Append WebSocket close status
+    // RFC 6455: If there is a body, the first two bytes of the body MUST
+    // be a 2-byte unsigned integer (in network byte order) representing
+    // a status code with value code.
+    if (store_status)
+    {
+        index += 2;
+        _ws_send_buffer[offset + 0] = ((status >> 8) & 0xFF) ^ _ws_send_mask[0];
+        _ws_send_buffer[offset + 1] = (status & 0xFF) ^ _ws_send_mask[1];
+    }
+
+    // Mask WebSocket frame content
+    for (size_t i = index; i < size; ++i)
+        _ws_send_buffer[offset + i] = data[i - index] ^ _ws_send_mask[i % 4];
 }
 
 void WebSocket::PrepareReceiveFrame(const void* buffer, size_t size)
@@ -410,8 +428,18 @@ void WebSocket::PrepareReceiveFrame(const void* buffer, size_t size)
                     }
                     case WS_CLOSE:
                     {
+                        size_t sindex = 0;
+                        int status = 1000;
+
+                        // Read WebSocket close status
+                        if (_ws_receive_final_buffer.size() >= 2)
+                        {
+                            sindex += 2;
+                            status = ((_ws_receive_final_buffer[0] << 8) | (_ws_receive_final_buffer[1] << 0));
+                        }
+
                         // Call the WebSocket close handler
-                        onWSClose(_ws_receive_final_buffer.data(), _ws_receive_final_buffer.size());
+                        onWSClose(_ws_receive_final_buffer.data() + sindex, _ws_receive_final_buffer.size() - sindex, status);
                         break;
                     }
                     case WS_BINARY:
diff --git a/source/server/ws/ws_server.cpp b/source/server/ws/ws_server.cpp
@@ -11,11 +11,22 @@
 namespace CppServer {
 namespace WS {
 
-bool WSServer::CloseAll(int status)
+bool WSServer::CloseAll(int status, const void* buffer, size_t size)
 {
     std::scoped_lock locker(_ws_send_lock);
 
-    PrepareSendFrame(WS_FIN | WS_CLOSE, false, nullptr, 0, status);
+    PrepareSendFrame(WS_FIN | WS_CLOSE, false, buffer, size, status);
+    if (!Multicast(_ws_send_buffer.data(), _ws_send_buffer.size()))
+        return false;
+
+    return HTTPServer::DisconnectAll();
+}
+
+bool WSServer::CloseAll(int status, std::string_view text)
+{
+    std::scoped_lock locker(_ws_send_lock);
+
+    PrepareSendFrame(WS_FIN | WS_CLOSE, false, text.data(), text.size(), status);
     if (!Multicast(_ws_send_buffer.data(), _ws_send_buffer.size()))
         return false;
 
diff --git a/source/server/ws/wss_server.cpp b/source/server/ws/wss_server.cpp
@@ -11,11 +11,22 @@
 namespace CppServer {
 namespace WS {
 
-bool WSSServer::CloseAll(int status)
+bool WSSServer::CloseAll(int status, const void* buffer, size_t size)
 {
     std::scoped_lock locker(_ws_send_lock);
 
-    PrepareSendFrame(WS_FIN | WS_CLOSE, false, nullptr, 0, status);
+    PrepareSendFrame(WS_FIN | WS_CLOSE, false, buffer, size, status);
+    if (!Multicast(_ws_send_buffer.data(), _ws_send_buffer.size()))
+        return false;
+
+    return HTTPSServer::DisconnectAll();
+}
+
+bool WSSServer::CloseAll(int status, std::string_view text)
+{
+    std::scoped_lock locker(_ws_send_lock);
+
+    PrepareSendFrame(WS_FIN | WS_CLOSE, false, text.data(), text.size(), status);
     if (!Multicast(_ws_send_buffer.data(), _ws_send_buffer.size()))
         return false;
 

Original file line number	Diff line number	Diff line change
`@@ -42,7 +42,7 @@ class ChatSession : public CppServer::WS::WSSession`
`42`	`42`
`43`	`43`	`// If the buffer starts with '!' the disconnect the current session`
`44`	`44`	`if (message == "!")`
`45`		`- Close(1000);`
	`45`	`+ Close();`
`46`	`46`	`}`
`47`	`47`
`48`	`48`	`void onError(int error, const std::string& category, const std::string& message) override`