Fix infinite loop

Author: an-tao

lib/src/WebSocketConnectionImpl.cc

@@ -17,6 +17,7 @@
 #include <json/value.h>
 #include <json/writer.h>
 #include <thread>
+#include <limits>
 
 using namespace drogon;
 
@@ -327,8 +328,13 @@ bool WebSocketMessageParser::parse(trantor::MsgBuffer *buffer)
         {
             indexFirstMask = 10;
         }
-        if (indexFirstMask > 2 && buffer->readableBytes() >= indexFirstMask)
+        if (indexFirstMask > 2)
         {
+            if (buffer->readableBytes() < indexFirstMask)
+            {
+                // Not enough data yet, wait for more.
+                return true;
+            }
             if (isControlFrame)
             {
                 // rfc6455-5.5
@@ -344,14 +350,17 @@ bool WebSocketMessageParser::parse(trantor::MsgBuffer *buffer)
             }
             else if (indexFirstMask == 10)
             {
-                length = (unsigned char)(*buffer)[2];
-                length = (length << 8) + (unsigned char)(*buffer)[3];
-                length = (length << 8) + (unsigned char)(*buffer)[4];
-                length = (length << 8) + (unsigned char)(*buffer)[5];
-                length = (length << 8) + (unsigned char)(*buffer)[6];
-                length = (length << 8) + (unsigned char)(*buffer)[7];
-                length = (length << 8) + (unsigned char)(*buffer)[8];
-                length = (length << 8) + (unsigned char)(*buffer)[9];
+                length = 0;
+                for (int i = 2; i <= 9; ++i)
+                {
+                    if (length > ((std::numeric_limits<size_t>::max)() >> 8))
+                    {
+                        LOG_ERROR
+                            << "Payload length too large to handle safely";
+                        return false;
+                    }
+                    length = (length << 8) + (unsigned char)(*buffer)[i];
+                }
             }
             else
             {
@@ -387,6 +396,11 @@ bool WebSocketMessageParser::parse(trantor::MsgBuffer *buffer)
                     return true;
                 }
             }
+            else
+            {
+                // Not enough data yet, wait for more.
+                return true;
+            }
         }
         else
         {
@@ -401,6 +415,11 @@ bool WebSocketMessageParser::parse(trantor::MsgBuffer *buffer)
                     return true;
                 }
             }
+            else
+            {
+                // Not enough data yet, wait for more.
+                return true;
+            }
         }
     }
     return true;