Improve WebSocket mask handling (#875)

CMakeLists.txt

@@ -64,7 +64,7 @@ if (BUILD_DROGON_SHARED)
         SOVERSION ${DROGON_MAJOR_VERSION})
     target_link_libraries(${PROJECT_NAME} PUBLIC Threads::Threads)
     if (WIN32)
-        target_link_libraries(${PROJECT_NAME} PUBLIC Rpcrt4 ws2_32)
+        target_link_libraries(${PROJECT_NAME} PUBLIC Rpcrt4 ws2_32 crypt32 Advapi32)
         if (CMAKE_CXX_COMPILER_ID MATCHES MSVC)
             # Ignore MSVC C4251 and C4275 warning of exporting std objects with no dll export
             # We export class to facilitate maintenance, thus if you compile

lib/inc/drogon/utils/Utilities.h

@@ -163,5 +163,18 @@ DROGON_EXPORT void replaceAll(std::string &s,
                               const std::string &from,
                               const std::string &to);
 
+/**
+ * @brief Generates cryptographically secure random bytes.
+ *
+ * @param ptr the pointer which the random bytes are stored to
+ * @param size number of bytes to generate
+ *
+ * @return true if generation is successfull. False otherwise
+ *
+ * @note DO NOT abuse this function. Especially if Drogon is built without
+ * OpenSSL. Entropy running low is a real issue.
+ */
+DROGON_EXPORT bool secureRandomBytes(void *ptr, size_t size);
+
 }  // namespace utils
 }  // namespace drogon

lib/src/Utilities.cc

@@ -17,6 +17,7 @@
 #include <drogon/config.h>
 #ifdef OpenSSL_FOUND
 #include <openssl/md5.h>
+#include <openssl/rand.h>
 #else
 #include "ssl_funcs/Md5.h"
 #endif
@@ -28,8 +29,10 @@
 #include <Rpc.h>
 #include <direct.h>
 #include <io.h>
+#include <ntsecapi.h>
 #else
 #include <uuid.h>
+#include <unistd.h>
 #endif
 #include <zlib.h>
 #include <iomanip>
@@ -44,9 +47,6 @@
 #include <cstdlib>
 #include <stdio.h>
 #include <string.h>
-#ifndef _WIN32
-#include <unistd.h>
-#endif
 #include <sys/stat.h>
 #include <fcntl.h>
 #include <stdarg.h>
@@ -1199,5 +1199,48 @@ void replaceAll(std::string &s, const std::string &from, const std::string &to)
     }
 }
 
+/**
+ * @brief Generates `size` random bytes from the systems random source and
+ * stores them into `ptr`.
+ */
+static bool systemRandomBytes(void *ptr, size_t size)
+{
+#if defined(__BSD__) || defined(__APPLE__)
+    arc4random_buf(ptr, size);
+    return true;
+#elif defined(__linux__) && \
+    ((defined(__GLIBC__) && \
+      (__GLIBC__ > 2 || (__GLIBC__ == 2 && __GLIBC_MINOR__ >= 25))))
+    return getentropy(ptr, size) != -1;
+#elif defined(_WIN32)    // Windows
+    return RtlGenRandom(ptr, size);
+#elif defined(__unix__)  // fallback to /dev/urandom for other/old UNIX
+    static std::unique_ptr<FILE, std::function<void(FILE *)> > fptr(
+        fopen("/dev/urandom", "rb"), [](FILE *ptr) {
+            if (ptr != nullptr)
+                fclose(ptr);
+        });
+    if (fptr == nullptr)
+    {
+        LOG_FATAL << "Failed to open /dev/urandom for randomness";
+        abort();
+    }
+    if (fread(ptr, 1, size, fptr.get()) != 0)
+        return true;
+#endif
+    return false;
+}
+
+bool secureRandomBytes(void *ptr, size_t size)
+{
+#ifdef OpenSSL_FOUND
+    if (RAND_bytes((unsigned char *)ptr, size) == 0)
+        return true;
+#endif
+    if (systemRandomBytes(ptr, size))
+        return true;
+    return false;
+}
+
 }  // namespace utils
 }  // namespace drogon

lib/src/WebSocketConnectionImpl.cc

@@ -23,7 +23,8 @@ WebSocketConnectionImpl::WebSocketConnectionImpl(
     : tcpConnectionPtr_(conn),
       localAddr_(conn->localAddr()),
       peerAddr_(conn->peerAddr()),
-      isServer_(isServer)
+      isServer_(isServer),
+      usingMask_(false)
 {
 }
 WebSocketConnectionImpl::~WebSocketConnectionImpl()
@@ -105,16 +106,42 @@ void WebSocketConnectionImpl::sendWsData(const char *msg,
     }
     if (!isServer_)
     {
-        // Add masking key;
-        static std::once_flag once;
-        std::call_once(once, []() {
-            std::srand(static_cast<unsigned int>(time(nullptr)));
-        });
-        int random = std::rand();
+        int random;
+        // Use the cached randomness if no one else is also using it. Otherwise
+        // generate one from scratch.
+        if (!usingMask_.exchange(true))
+        {
+            if (masks_.empty())
+            {
+                masks_.resize(16);
+                bool status =
+                    utils::secureRandomBytes(masks_.data(),
+                                             masks_.size() * sizeof(uint32_t));
+                if (status == false)
+                {
+                    LOG_ERROR << "Failed to generate random numbers for "
+                                 "WebSocket mask";
+                    abort();
+                }
+            }
+            random = masks_.back();
+            masks_.pop_back();
+            usingMask_ = false;
+        }
+        else
+        {
+            bool status = utils::secureRandomBytes(&random, sizeof(random));
+            if (status == false)
+            {
+                LOG_ERROR
+                    << "Failed to generate random numbers for WebSocket mask";
+                abort();
+            }
+        }
 
         bytesFormatted[1] = (bytesFormatted[1] | 0x80);
         bytesFormatted.resize(indexStartRawData + 4 + len);
-        *((int *)&bytesFormatted[indexStartRawData]) = random;
+        memcpy(&bytesFormatted[indexStartRawData], &random, sizeof(random));
         for (size_t i = 0; i < len; ++i)
         {
             bytesFormatted[indexStartRawData + 4 + i] =

lib/src/WebSocketConnectionImpl.h

@@ -108,6 +108,8 @@ class WebSocketConnectionImpl final
     bool isServer_{true};
     WebSocketMessageParser parser_;
     trantor::TimerId pingTimerId_{trantor::InvalidTimerId};
+    std::vector<uint32_t> masks_;
+    std::atomic<bool> usingMask_;
 
     std::function<void(std::string &&,
                        const WebSocketConnectionImplPtr &,

orm_lib/src/sqlite3_impl/Sqlite3Connection.cc

@@ -308,8 +308,14 @@ void Sqlite3Connection::disconnect()
     std::promise<int> pro;
     auto f = pro.get_future();
     auto thisPtr = shared_from_this();
-    loopThread_.getLoop()->runInLoop([thisPtr, &pro]() {
-        thisPtr->connectionPtr_.reset();
+    std::weak_ptr<Sqlite3Connection> weakPtr = thisPtr;
+    loopThread_.getLoop()->runInLoop([weakPtr, &pro]() {
+        {
+            auto thisPtr = weakPtr.lock();
+            if (!thisPtr)
+                return;
+            thisPtr->connectionPtr_.reset();
+        }
         pro.set_value(1);
     });
     f.get();

orm_lib/tests/db_test.cc

@@ -1928,20 +1928,6 @@ int main(int argc, char **argv)
 #if USE_SQLITE3
     sqlite3Client = DbClient::newSqlite3Client("filename=:memory:", 1);
 #endif
-
     int testStatus = test::run(argc, argv);
-    std::this_thread::sleep_for(0.008s);
-
-    // Destruct the clients before event loop shutdown
-#if USE_MYSQL
-    mysqlClient.reset();
-#endif
-#if USE_POSTGRESQL
-    postgreClient.reset();
-#endif
-#if USE_SQLITE3
-    sqlite3Client.reset();
-#endif
-
     return testStatus;
 }