HTTPParser and WebSocket fix.

- HTTPParser: fix leftover data after HTTP upgrade - WebSocket: Now correctly handle leftover data upon handshake
nidium · Oct 15, 2016 · 472e805 · 472e805
1 parent 457725e
commit 472e805
Show file tree

Hide file tree

Showing 4 changed files with 31 additions and 8 deletions.
diff --git a/src/Net/HTTPParser.cpp b/src/Net/HTTPParser.cpp
@@ -6,6 +6,7 @@
 #include "Net/HTTPParser.h"
 
 #include <string.h>
+#include <stdio.h>
 
 namespace Nidium {
 namespace Net {
@@ -73,7 +74,7 @@ static int body_cb(http_parser *p, const char *buf, size_t len)
                            reinterpret_cast<const unsigned char *>(buf), len);
     }
 
-    nhttp->HTTPOnData(nhttp->m_Data->used - len, len);
+    nhttp->HTTPOnData(buf, len);
 
     return 0;
 }
@@ -171,6 +172,14 @@ bool HTTPParser::HTTPParse(const char *data, size_t len)
     nparsed = http_parser_execute(&m_Parser, &settings,
                                   static_cast<const char *>(data), len);
 
+    if (m_Parser.upgrade) {
+        this->HTTPOnUpgrade();
+
+        if (nparsed < len) {
+            this->HTTPOnData(data + nparsed, len - nparsed);
+        }
+    }
+
     return nparsed != 0;
 }
 

diff --git a/src/Net/HTTPParser.h b/src/Net/HTTPParser.h
@@ -21,9 +21,10 @@ class HTTPParser
     void HTTPClearState();
     const char *HTTPGetHeader(const char *key);
 
+    virtual void HTTPOnUpgrade(){};
     virtual void HTTPHeaderEnded() = 0;
     virtual void HTTPRequestEnded() = 0;
-    virtual void HTTPOnData(size_t offset, size_t len) = 0;
+    virtual void HTTPOnData(const char *data, size_t len) = 0;
 
     enum PrevState
     {

diff --git a/src/Net/WebSocketClient.cpp b/src/Net/WebSocketClient.cpp
@@ -45,6 +45,7 @@ static void nidium_on_ws_client_frame(websocket_state *state,
                                       int binary)
 {
     ape_socket *sock = state->socket;
+
     if (sock == NULL) {
         return;
     }
@@ -108,19 +109,25 @@ void WebSocketClient::HTTPHeaderEnded()
     const char *swa = this->HTTPGetHeader("Sec-WebSocket-Accept");
 
     /* Check handshake key integrity */
-    if (swa == NULL || strcmp(swa, m_ComputedKey) != 0) {
+    if (swa == NULL || strcmp(swa, m_ComputedKey) != 0 || !m_Parser.upgrade) {
         APE_socket_shutdown_now(m_Socket);
         return;
     }
+
+    m_HandShakeDone = true;
 }
 
 void WebSocketClient::HTTPRequestEnded()
 {
-    m_Socket->callbacks.on_read = nidium_ws_read_ws;
-
     Args args;
     args[0].set(this);
 
+    if (!m_HandShakeDone) {
+        return;
+    }
+
+    m_Socket->callbacks.on_read = nidium_ws_read_ws;
+
     this->fireEvent<WebSocketClient>(WebSocketClient::kEvents_ClientConnect,
                                      args);
 }
@@ -146,8 +153,13 @@ void WebSocketClient::ping()
     ape_ws_ping(&m_WSState);
 }
 
-void WebSocketClient::HTTPOnData(size_t offset, size_t len)
+void WebSocketClient::HTTPOnData(const char *data, size_t len)
 {
+    if (!m_HandShakeDone) {
+        return;
+    }
+
+    this->onDataWS((const uint8_t *)data, len);
 }
 
 
@@ -213,7 +225,7 @@ void WebSocketClient::onConnected()
     */
     APE_socket_write(m_Socket, m_HandShakeKey, strlen(m_HandShakeKey),
                      APE_DATA_STATIC);
-    // TODO: new style cast
+
     APE_socket_write(m_Socket, (unsigned char *)CONST_STR_LEN("\r\n\r\n"),
                      APE_DATA_STATIC);
 

diff --git a/src/Net/WebSocketClient.h b/src/Net/WebSocketClient.h
@@ -47,7 +47,7 @@ class WebSocketClient : public Nidium::Core::Events, public HTTPParser
 
     virtual void HTTPHeaderEnded();
     virtual void HTTPRequestEnded();
-    virtual void HTTPOnData(size_t offset, size_t len);
+    virtual void HTTPOnData(const char *data, size_t len);
 
 private:
     char m_HandShakeKey[32];
@@ -57,6 +57,7 @@ class WebSocketClient : public Nidium::Core::Events, public HTTPParser
     char *m_URL;
     uint16_t m_Port;
     bool m_SSL;
+    bool m_HandShakeDone = false;
 
     websocket_state m_WSState;
 };