add ping and pong to message handler (#1143)

Kumamon38 · BillyONeal · commit 0f45af151345 · 2019-06-13T12:40:00.000-07:00
Add ping and pong to message handler, and optional pong timeout
diff --git a/Release/include/cpprest/ws_msg.h b/Release/include/cpprest/ws_msg.h
@@ -60,10 +60,23 @@ class websocket_outgoing_message
 public:
 #if !defined(__cplusplus_winrt)
     /// <summary>
-    /// Sets a the outgoing message to be an unsolicited pong message.
+    /// Sets the outgoing message to be a ping message.
     /// This is useful when the client side wants to check whether the server is alive.
     /// </summary>
-    void set_pong_message() { this->set_message_pong(); }
+    /// <param name="data">UTF-8 String containing the optional ping message.</param>
+    void set_ping_message(const std::string& data = {})
+    {
+        this->set_message_ping(concurrency::streams::container_buffer<std::string>(data));
+    }
+
+    /// <summary>
+    /// Sets the outgoing message to be an unsolicited pong message.
+    /// </summary>
+    /// <param name="data">UTF-8 String containing the optional pong message.</param>
+    void set_pong_message(const std::string& data = {})
+    { 
+        this->set_message_pong(concurrency::streams::container_buffer<std::string>(data));
+    }
 #endif
 
     /// <summary>
@@ -140,9 +153,14 @@ class websocket_outgoing_message
     const pplx::task_completion_event<void>& body_sent() const { return m_body_sent; }
 
 #if !defined(__cplusplus_winrt)
-    void set_message_pong()
+    void set_message_ping(const concurrency::streams::container_buffer<std::string>& buffer)
+    {
+        m_msg_type = websocket_message_type::ping;
+        m_length = static_cast<size_t>(buffer.size());
+        m_body = buffer;
+    }
+    void set_message_pong(const concurrency::streams::container_buffer<std::string>& buffer)
     {
-        concurrency::streams::container_buffer<std::string> buffer("");
         m_msg_type = websocket_message_type::pong;
         m_length = static_cast<size_t>(buffer.size());
         m_body = buffer;
diff --git a/Release/src/websockets/client/ws_client_wspp.cpp b/Release/src/websockets/client/ws_client_wspp.cpp
@@ -325,6 +325,35 @@ class wspp_callback_client : public websocket_client_callback_impl,
                 }
             });
 
+        client.set_ping_handler(
+            [this](websocketpp::connection_hdl, const std::string& msg) {
+                if (m_external_message_handler)
+                {
+                    _ASSERTE(m_state >= CONNECTED && m_state < CLOSED);
+                    websocket_incoming_message incoming_msg;
+
+                    incoming_msg.m_msg_type = websocket_message_type::ping;
+                    incoming_msg.m_body = concurrency::streams::container_buffer<std::string>(msg);
+
+                    m_external_message_handler(incoming_msg);
+                }
+                return true;
+            });
+
+        client.set_pong_handler(
+            [this](websocketpp::connection_hdl, const std::string& msg) {
+                if (m_external_message_handler)
+                {
+                    _ASSERTE(m_state >= CONNECTED && m_state < CLOSED);
+                    websocket_incoming_message incoming_msg;
+
+                    incoming_msg.m_msg_type = websocket_message_type::pong;
+                    incoming_msg.m_body = concurrency::streams::container_buffer<std::string>(msg);
+
+                    m_external_message_handler(incoming_msg);
+                }
+            });
+
         client.set_close_handler([this](websocketpp::connection_hdl con_hdl) {
             _ASSERTE(m_state != CLOSED);
             shutdown_wspp_impl<WebsocketConfigType>(con_hdl, false);
@@ -434,12 +463,13 @@ class wspp_callback_client : public websocket_client_callback_impl,
         {
             case websocket_message_type::text_message:
             case websocket_message_type::binary_message:
+            case websocket_message_type::ping:
             case websocket_message_type::pong: break;
             default: return pplx::task_from_exception<void>(websocket_exception("Message Type not supported."));
         }
 
         const auto length = msg.m_length;
-        if (length == 0 && msg.m_msg_type != websocket_message_type::pong)
+        if (length == 0 && msg.m_msg_type != websocket_message_type::ping && msg.m_msg_type != websocket_message_type::pong)
         {
             return pplx::task_from_exception<void>(websocket_exception("Cannot send empty message."));
         }
@@ -694,7 +724,18 @@ class wspp_callback_client : public websocket_client_callback_impl,
             case websocket_message_type::binary_message:
                 client.send(this_client->m_con, sp_allocated.get(), length, websocketpp::frame::opcode::binary, ec);
                 break;
-            case websocket_message_type::pong: client.pong(this_client->m_con, "", ec); break;
+            case websocket_message_type::ping:
+            {
+                std::string s(reinterpret_cast<char*>(sp_allocated.get()), length);
+                client.ping(this_client->m_con, s, ec);
+                break;
+            }
+            case websocket_message_type::pong:
+            {
+                std::string s(reinterpret_cast<char*>(sp_allocated.get()), length);
+                client.pong(this_client->m_con, s, ec);
+                break;
+            }
             default:
                 // This case should have already been filtered above.
                 std::abort();
diff --git a/Release/tests/functional/websockets/client/send_msg_tests.cpp b/Release/tests/functional/websockets/client/send_msg_tests.cpp
@@ -104,16 +104,32 @@ SUITE(send_msg_tests)
     }
 
     template<class SocketClientClass>
-    pplx::task<void> send_pong_msg_helper(SocketClientClass & client, web::uri uri, test_websocket_server & server)
+    pplx::task<void> send_ping_msg_helper(SocketClientClass & client, web::uri uri, test_websocket_server & server,
+                                          const std::string& body = "")
     {
         server.next_message(
-            [](test_websocket_msg msg) // Handler to verify the message sent by the client.
-            { websocket_asserts::assert_message_equals(msg, "", test_websocket_message_type::WEB_SOCKET_PONG_TYPE); });
+            [body](test_websocket_msg msg) // Handler to verify the message sent by the client.
+            { websocket_asserts::assert_message_equals(msg, body, test_websocket_message_type::WEB_SOCKET_PING_TYPE); });
 
         client.connect(uri).wait();
 
         websocket_outgoing_message msg;
-        msg.set_pong_message();
+        msg.set_ping_message(body);
+        return client.send(msg);
+    }
+
+    template<class SocketClientClass>
+    pplx::task<void> send_pong_msg_helper(SocketClientClass & client, web::uri uri, test_websocket_server & server,
+                                          const std::string& body = "")
+    {
+        server.next_message(
+            [body](test_websocket_msg msg) // Handler to verify the message sent by the client.
+            { websocket_asserts::assert_message_equals(msg, body, test_websocket_message_type::WEB_SOCKET_PONG_TYPE); });
+
+        client.connect(uri).wait();
+
+        websocket_outgoing_message msg;
+        msg.set_pong_message(body);
         return client.send(msg);
     }
 
@@ -493,6 +509,24 @@ SUITE(send_msg_tests)
     }
 
 #if !defined(__cplusplus_winrt)
+    // Send a ping message to the server
+    TEST_FIXTURE(uri_address, send_ping_msg)
+    {
+        test_websocket_server server;
+        websocket_client client;
+        send_ping_msg_helper(client, m_uri, server).wait();
+        client.close().wait();
+    }
+
+    // Send a ping message to the server with a body
+    TEST_FIXTURE(uri_address, send_ping_msg_body)
+    {
+        test_websocket_server server;
+        websocket_client client;
+        send_ping_msg_helper(client, m_uri, server, "abcdefghijklmnopqrstuvwxyz").wait();
+        client.close().wait();
+    }
+
     // Send an unsolicited pong message to the server
     TEST_FIXTURE(uri_address, send_pong_msg)
     {
@@ -502,6 +536,15 @@ SUITE(send_msg_tests)
         client.close().wait();
     }
 
+    // Send an unsolicited pong message to the server with a body
+    TEST_FIXTURE(uri_address, send_pong_msg_body)
+    {
+        test_websocket_server server;
+        websocket_client client;
+        send_pong_msg_helper(client, m_uri, server, "abcdefghijklmnopqrstuvwxyz").wait();
+        client.close().wait();
+    }
+
     // Send an unsolicited pong message to the server with websocket_callback_client
     TEST_FIXTURE(uri_address, send_pong_msg_callback_client)
     {
diff --git a/Release/tests/functional/websockets/utilities/test_websocket_server.cpp b/Release/tests/functional/websockets/utilities/test_websocket_server.cpp
@@ -122,6 +122,20 @@ class _test_websocket_server
             m_server_connected.set_exception(std::runtime_error("Connection attempt failed."));
         });
 
+        m_srv.set_ping_handler([this](websocketpp::connection_hdl hdl, std::string input) {
+            auto fn = m_test_srv->get_next_message_handler();
+            assert(fn);
+
+            test_websocket_msg wsmsg;
+
+            wsmsg.set_data(std::vector<uint8_t>(input.begin(), input.end()));
+
+            wsmsg.set_msg_type(WEB_SOCKET_PING_TYPE);
+            fn(wsmsg);
+
+            return true;
+        });
+
         m_srv.set_pong_handler([this](websocketpp::connection_hdl hdl, std::string input) {
             auto fn = m_test_srv->get_next_message_handler();
             assert(fn);
diff --git a/Release/tests/functional/websockets/utilities/test_websocket_server.h b/Release/tests/functional/websockets/utilities/test_websocket_server.h
@@ -46,6 +46,7 @@ enum test_websocket_message_type
     WEB_SOCKET_UTF8_MESSAGE_TYPE,
     WEB_SOCKET_UTF8_FRAGMENT_TYPE,
     WEB_SOCKET_CLOSE_TYPE,
+    WEB_SOCKET_PING_TYPE,
     WEB_SOCKET_PONG_TYPE
 };
 
