Merge https://gitlab.com/eidheim/Simple-WebSocket-Server

3 files changed, 53 insertions, 37 deletions

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 5a9429d..9d2d4f7 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -5,15 +5,6 @@ project (Simple-WebSocket-Server)
 option(USE_STANDALONE_ASIO "set ON to use standalone Asio instead of Boost.Asio" OFF)
 option(BUILD_TESTING "set ON to build library tests" OFF)
 
-if(NOT MSVC)
-    add_compile_options(-std=c++11 -Wall -Wextra -Wsign-conversion)
-    if (CMAKE_CXX_COMPILER_ID MATCHES "Clang")
-        add_compile_options(-Wthread-safety)
-    endif()
-else()
-    add_compile_options(/W1)
-endif()
-
 add_library(simple-websocket-server INTERFACE)
 
 target_include_directories(simple-websocket-server INTERFACE ${CMAKE_CURRENT_SOURCE_DIR})
@@ -55,6 +46,15 @@ target_include_directories(simple-websocket-server INTERFACE ${OPENSSL_INCLUDE_D
 
 # If Simple-WebSocket-Server is not a sub-project:
 if(CMAKE_SOURCE_DIR STREQUAL "${CMAKE_CURRENT_SOURCE_DIR}")
+    if(NOT MSVC)
+        add_compile_options(-std=c++11 -Wall -Wextra -Wsign-conversion)
+        if (CMAKE_CXX_COMPILER_ID MATCHES "Clang")
+            add_compile_options(-Wthread-safety)
+        endif()
+    else()
+        add_compile_options(/W1)
+    endif()
+
     add_executable(ws_examples ws_examples.cpp)
     target_link_libraries(ws_examples simple-websocket-server)
     add_executable(wss_examples wss_examples.cpp)
diff --git a/client_ws.hpp b/client_ws.hpp
index b17ff5e..b5ce2d7 100644
--- a/client_ws.hpp
+++ b/client_ws.hpp
@@ -460,7 +460,7 @@ namespace SimpleWeb {
 
     void read_message(const std::shared_ptr<Connection> &connection, std::size_t num_additional_bytes) {
       connection->set_timeout();
-      asio::async_read(*connection->socket, connection->in_message->streambuf, asio::transfer_exactly(num_additional_bytes > 2 ? 0 : 2 - num_additional_bytes), [this, connection](const error_code &ec, std::size_t bytes_transferred) {
+      asio::async_read(*connection->socket, connection->in_message->streambuf, asio::transfer_exactly(num_additional_bytes > 2 ? 0 : 2 - num_additional_bytes), [this, connection, num_additional_bytes](const error_code &ec, std::size_t bytes_transferred) {
         connection->cancel_timeout();
         auto lock = connection->handler_runner->continue_lock();
         if(!lock)
@@ -470,7 +470,7 @@ namespace SimpleWeb {
             this->read_message(connection, 0);
             return;
           }
-          std::size_t num_additional_bytes = connection->in_message->streambuf.size() - bytes_transferred;
+          auto updated_num_additional_bytes = num_additional_bytes > 2 ? num_additional_bytes - 2 : 0;
 
           std::array<unsigned char, 2> first_bytes;
           connection->in_message->read(reinterpret_cast<char *>(&first_bytes[0]), 2);
@@ -490,14 +490,12 @@ namespace SimpleWeb {
           if(length == 126) {
             // 2 next bytes is the size of content
             connection->set_timeout();
-            asio::async_read(*connection->socket, connection->in_message->streambuf, asio::transfer_exactly(num_additional_bytes > 2 ? 0 : 2 - num_additional_bytes), [this, connection](const error_code &ec, std::size_t bytes_transferred) {
+            asio::async_read(*connection->socket, connection->in_message->streambuf, asio::transfer_exactly(updated_num_additional_bytes > 2 ? 0 : 2 - updated_num_additional_bytes), [this, connection, updated_num_additional_bytes](const error_code &ec, std::size_t /*bytes_transferred*/) {
               connection->cancel_timeout();
               auto lock = connection->handler_runner->continue_lock();
               if(!lock)
                 return;
               if(!ec) {
-                std::size_t num_additional_bytes = connection->in_message->streambuf.size() - bytes_transferred;
-
                 std::array<unsigned char, 2> length_bytes;
                 connection->in_message->read(reinterpret_cast<char *>(&length_bytes[0]), 2);
 
@@ -507,7 +505,7 @@ namespace SimpleWeb {
                   length += static_cast<std::size_t>(length_bytes[c]) << (8 * (num_bytes - 1 - c));
 
                 connection->in_message->length = length;
-                this->read_message_content(connection, num_additional_bytes);
+                this->read_message_content(connection, updated_num_additional_bytes > 2 ? updated_num_additional_bytes - 2 : 0);
               }
               else
                 this->connection_error(connection, ec);
@@ -516,14 +514,12 @@ namespace SimpleWeb {
           else if(length == 127) {
             // 8 next bytes is the size of content
             connection->set_timeout();
-            asio::async_read(*connection->socket, connection->in_message->streambuf, asio::transfer_exactly(num_additional_bytes > 8 ? 0 : 8 - num_additional_bytes), [this, connection](const error_code &ec, std::size_t bytes_transferred) {
+            asio::async_read(*connection->socket, connection->in_message->streambuf, asio::transfer_exactly(updated_num_additional_bytes > 8 ? 0 : 8 - updated_num_additional_bytes), [this, connection, updated_num_additional_bytes](const error_code &ec, std::size_t /*bytes_transferred*/) {
               connection->cancel_timeout();
               auto lock = connection->handler_runner->continue_lock();
               if(!lock)
                 return;
               if(!ec) {
-                std::size_t num_additional_bytes = connection->in_message->streambuf.size() - bytes_transferred;
-
                 std::array<unsigned char, 8> length_bytes;
                 connection->in_message->read(reinterpret_cast<char *>(&length_bytes[0]), 8);
 
@@ -533,7 +529,7 @@ namespace SimpleWeb {
                   length += static_cast<std::size_t>(length_bytes[c]) << (8 * (num_bytes - 1 - c));
 
                 connection->in_message->length = length;
-                this->read_message_content(connection, num_additional_bytes);
+                this->read_message_content(connection, updated_num_additional_bytes > 8 ? updated_num_additional_bytes - 8 : 0);
               }
               else
                 this->connection_error(connection, ec);
@@ -541,7 +537,7 @@ namespace SimpleWeb {
           }
           else {
             connection->in_message->length = length;
-            this->read_message_content(connection, num_additional_bytes);
+            this->read_message_content(connection, updated_num_additional_bytes);
           }
         }
         else
@@ -559,15 +555,15 @@ namespace SimpleWeb {
         return;
       }
       connection->set_timeout();
-      asio::async_read(*connection->socket, connection->in_message->streambuf, asio::transfer_exactly(num_additional_bytes > connection->in_message->length ? 0 : connection->in_message->length - num_additional_bytes), [this, connection](const error_code &ec, std::size_t bytes_transferred) {
+      asio::async_read(*connection->socket, connection->in_message->streambuf, asio::transfer_exactly(num_additional_bytes > connection->in_message->length ? 0 : connection->in_message->length - num_additional_bytes), [this, connection, num_additional_bytes](const error_code &ec, std::size_t /*bytes_transferred*/) {
         connection->cancel_timeout();
         auto lock = connection->handler_runner->continue_lock();
         if(!lock)
           return;
         if(!ec) {
-          std::size_t num_additional_bytes = connection->in_message->streambuf.size() - bytes_transferred;
+          auto updated_num_additional_bytes = num_additional_bytes > connection->in_message->length ? num_additional_bytes - connection->in_message->length : 0;
           std::shared_ptr<InMessage> next_in_message;
-          if(num_additional_bytes > 0) { // Extract bytes that are not extra bytes in buffer (only happen when several messages are sent in upgrade response)
+          if(updated_num_additional_bytes > 0) { // Extract bytes that are not extra bytes in buffer (only happen when several messages are sent in upgrade response)
             next_in_message = connection->in_message;
             connection->in_message = std::shared_ptr<InMessage>(new InMessage(next_in_message->fin_rsv_opcode, next_in_message->length));
 
@@ -605,7 +601,7 @@ namespace SimpleWeb {
 
             // Next message
             connection->in_message = next_in_message;
-            this->read_message(connection, num_additional_bytes);
+            this->read_message(connection, updated_num_additional_bytes);
           }
           // If pong
           else if((connection->in_message->fin_rsv_opcode & 0x0f) == 10) {
@@ -614,7 +610,7 @@ namespace SimpleWeb {
 
             // Next message
             connection->in_message = next_in_message;
-            this->read_message(connection, num_additional_bytes);
+            this->read_message(connection, updated_num_additional_bytes);
           }
           // If fragmented message and not final fragment
           else if((connection->in_message->fin_rsv_opcode & 0x80) == 0) {
@@ -633,7 +629,7 @@ namespace SimpleWeb {
 
             // Next message
             connection->in_message = next_in_message;
-            this->read_message(connection, num_additional_bytes);
+            this->read_message(connection, updated_num_additional_bytes);
           }
           else {
             if(this->on_message) {
@@ -655,7 +651,7 @@ namespace SimpleWeb {
             connection->in_message = next_in_message;
             // Only reset fragmented_message for non-control frames (control frames can be in between a fragmented message)
             connection->fragmented_in_message = nullptr;
-            this->read_message(connection, num_additional_bytes);
+            this->read_message(connection, updated_num_additional_bytes);
           }
         }
         else
diff --git a/server_ws.hpp b/server_ws.hpp
index e91a479..158813d 100644
--- a/server_ws.hpp
+++ b/server_ws.hpp
@@ -91,10 +91,20 @@ namespace SimpleWeb {
         return endpoint;
       }
 
+      asio::ip::tcp::endpoint local_endpoint() const noexcept {
+        try {
+          if(auto connection = this->connection.lock())
+            return connection->socket->lowest_layer().local_endpoint();
+        }
+        catch(...) {
+        }
+        return asio::ip::tcp::endpoint();
+      }
+
       /// Deprecated, please use remote_endpoint().address().to_string() instead.
       DEPRECATED std::string remote_endpoint_address() const noexcept {
         try {
-          return socket->lowest_layer().remote_endpoint().address().to_string();
+          return endpoint.address().to_string();
         }
         catch(...) {
         }
@@ -104,7 +114,7 @@ namespace SimpleWeb {
       /// Deprecated, please use remote_endpoint().port() instead.
       DEPRECATED unsigned short remote_endpoint_port() const noexcept {
         try {
-          return socket->lowest_layer().remote_endpoint().port();
+          return endpoint.port();
         }
         catch(...) {
         }
@@ -363,7 +373,7 @@ namespace SimpleWeb {
       std::unique_lock<std::mutex> lock(start_stop_mutex);
 
       asio::ip::tcp::endpoint endpoint;
-      if(config.address.size() > 0)
+      if(!config.address.empty())
         endpoint = asio::ip::tcp::endpoint(make_address(config.address), config.port);
       else
         endpoint = asio::ip::tcp::endpoint(asio::ip::tcp::v6(), config.port);
@@ -375,7 +385,17 @@ namespace SimpleWeb {
 
       if(!acceptor)
         acceptor = std::unique_ptr<asio::ip::tcp::acceptor>(new asio::ip::tcp::acceptor(*io_service));
-      acceptor->open(endpoint.protocol());
+      try {
+        acceptor->open(endpoint.protocol());
+      }
+      catch(const system_error &error) {
+        if(error.code() == asio::error::address_family_not_supported && config.address.empty()) {
+          endpoint = asio::ip::tcp::endpoint(asio::ip::tcp::v4(), config.port);
+          acceptor->open(endpoint.protocol());
+        }
+        else
+          throw;
+      }
       acceptor->set_option(asio::socket_base::reuse_address(config.reuse_address));
       if(config.fast_open) {
 #if defined(__linux__) && defined(TCP_FASTOPEN)
@@ -538,6 +558,12 @@ namespace SimpleWeb {
             auto sha1 = Crypto::sha1(key_it->second + ws_magic_string);
             response_header.emplace("Sec-WebSocket-Accept", Crypto::Base64::encode(sha1));
 
+            try {
+              connection->endpoint = connection->socket->lowest_layer().remote_endpoint();
+            }
+            catch(...) {
+            }
+
             if(regex_endpoint.second.on_handshake)
               status_code = regex_endpoint.second.on_handshake(connection, response_header);
 
@@ -561,12 +587,6 @@ namespace SimpleWeb {
             if(status_code != StatusCode::information_switching_protocols)
               return;
 
-            try {
-              connection->endpoint = connection->socket->lowest_layer().remote_endpoint();
-            }
-            catch(...) {
-            }
-
             if(!ec) {
               connection_open(connection, regex_endpoint.second);
               read_message(connection, regex_endpoint.second);