hewei.it пре 4 година
родитељ
комит
7a500fefca
2 измењених фајлова са 12 додато и 5 уклоњено
  1. 5 2
      http/Http1Parser.cpp
  2. 7 3
      http/WebSocketParser.cpp

+ 5 - 2
http/Http1Parser.cpp

@@ -1,5 +1,7 @@
 #include "Http1Parser.h"
 
+#define MAX_CONTENT_LENGTH  (1 << 24)   // 16M
+
 static int on_url(http_parser* parser, const char *at, size_t length);
 static int on_status(http_parser* parser, const char *at, size_t length);
 static int on_header_field(http_parser* parser, const char *at, size_t length);
@@ -112,8 +114,9 @@ int on_headers_complete(http_parser* parser) {
     if (iter != hp->parsed->headers.end()) {
         int content_length = atoi(iter->second.c_str());
         hp->parsed->content_length = content_length;
-        if ((!skip_body) && content_length > hp->parsed->body.capacity()) {
-            hp->parsed->body.reserve(content_length);
+        int reserve_length = MIN(content_length, MAX_CONTENT_LENGTH);
+        if ((!skip_body) && reserve_length > hp->parsed->body.capacity()) {
+            hp->parsed->body.reserve(reserve_length);
         }
     }
     hp->state = HP_HEADERS_COMPLETE;

+ 7 - 3
http/WebSocketParser.cpp

@@ -1,6 +1,9 @@
 #include "WebSocketParser.h"
 
 #include "websocket_parser.h"
+#include "hdef.h"
+
+#define MAX_PAYLOAD_LENGTH  (1 << 24)   // 16M
 
 static int on_frame_header(websocket_parser* parser) {
     WebSocketParser* wp = (WebSocketParser*)parser->data;
@@ -10,12 +13,13 @@ static int on_frame_header(websocket_parser* parser) {
         wp->opcode = opcode;
     }
     int length = parser->length;
-    if (length && length > wp->message.capacity()) {
-        wp->message.reserve(length);
+    int reserve_length = MIN(length, MAX_PAYLOAD_LENGTH);
+    if (reserve_length > wp->message.capacity()) {
+        wp->message.reserve(reserve_length);
     }
     if (wp->state == WS_FRAME_BEGIN ||
         wp->state == WS_FRAME_END) {
-        wp->message.resize(0);
+        wp->message.clear();
     }
     wp->state = WS_FRAME_HEADER;
     return 0;