From 894609e6381b03d592889b906314544c58fc2511 Mon Sep 17 00:00:00 2001
From: Alex Hultman <alexhultman@gmail.com>
Date: Sun, 4 Nov 2018 10:44:11 +0100
Subject: [PATCH] Pass opCode around, add Loop::setPostHandler

---
 src/App.h                  | 2 +-
 src/Loop.h                 | 9 +++++++++
 src/LoopData.h             | 4 ++++
 src/WebSocket.h            | 4 ++--
 src/WebSocketContext.h     | 2 +-
 src/WebSocketContextData.h | 4 +++-
 6 files changed, 20 insertions(+), 5 deletions(-)

diff --git a/src/App.h b/src/App.h
index c4edf20..614fbce 100644
--- a/src/App.h
+++ b/src/App.h
@@ -52,7 +52,7 @@ public:
     }
 
     // this method creates a new websocket context and attaches it to a path
-    TemplatedApp &ws(std::string pattern, std::function<void(void *, HttpRequest *)> connectHandler, std::function<void(uWS::WebSocket<SSL, true> *, std::string_view)> messageHandler) {
+    TemplatedApp &ws(std::string pattern, std::function<void(void *, HttpRequest *)> connectHandler, std::function<void(uWS::WebSocket<SSL, true> *, std::string_view, uWS::OpCode)> messageHandler) {
         // init the websocket context here!
         uWS::WebSocketContext<SSL> *webSocketContext = uWS::WebSocketContext<SSL>::create(uWS::Loop::defaultLoop(), (typename StaticDispatch<SSL>::SOCKET_CONTEXT_TYPE *) httpContext);
 
diff --git a/src/Loop.h b/src/Loop.h
index ff17ffa..9cb2711 100644
--- a/src/Loop.h
+++ b/src/Loop.h
@@ -53,7 +53,9 @@ private:
     }
 
     static void postCb(us_loop *loop) {
+        LoopData *loopData = (LoopData *) us_loop_ext(loop);
 
+        loopData->postHandler((Loop *) loop);
     }
 
     Loop() = delete;
@@ -95,6 +97,13 @@ public:
         us_loop_free((us_loop *) this);
     }
 
+    /* Set postCb callback */
+    void setPostHandler(std::function<void(Loop *)> handler) {
+        LoopData *loopData = (LoopData *) us_loop_ext((us_loop *) this);
+
+        loopData->postHandler = handler;
+    }
+
     /* Defer this callback on Loop's thread of execution */
     void defer(std::function<void()> cb) {
         LoopData *loopData = (LoopData *) us_loop_ext((us_loop *) this);
diff --git a/src/LoopData.h b/src/LoopData.h
index 108886c..443bc73 100644
--- a/src/LoopData.h
+++ b/src/LoopData.h
@@ -24,6 +24,8 @@
 
 namespace uWS {
 
+struct Loop;
+
 struct LoopData {
     friend struct Loop;
 private:
@@ -31,6 +33,8 @@ private:
     int currentDeferQueue = 0;
     std::vector<std::function<void()>> deferQueues[2];
 
+    std::function<void(Loop *)> postHandler;
+
 public:
     /* Good 16k for SSL perf. */
     static const int CORK_BUFFER_SIZE = 16 * 1024;
diff --git a/src/WebSocket.h b/src/WebSocket.h
index 434d292..3d97b8e 100644
--- a/src/WebSocket.h
+++ b/src/WebSocket.h
@@ -15,13 +15,13 @@ private:
     typedef AsyncSocket<SSL> Super;
 public:
 
-    void send(std::string_view message) {
+    void send(std::string_view message, uWS::OpCode opCode) {
 
         // if corkAllocate(size) then corkFree(unused)
 
         // format the response
         char buf[100];
-        int writeLength = WebSocketProtocol<isServer, WebSocket<SSL, isServer>>::formatMessage(buf, message.data(), message.length(), uWS::OpCode::TEXT, message.length(), false);
+        int writeLength = WebSocketProtocol<isServer, WebSocket<SSL, isServer>>::formatMessage(buf, message.data(), message.length(), opCode, message.length(), false);
 
 
 
diff --git a/src/WebSocketContext.h b/src/WebSocketContext.h
index e2ad05c..76ae041 100644
--- a/src/WebSocketContext.h
+++ b/src/WebSocketContext.h
@@ -48,7 +48,7 @@ private:
             WebSocketContextData<SSL> *webSocketContextData = (WebSocketContextData<SSL> *) us_socket_context_ext(us_socket_get_context((us_socket *) s));
 
             // emit message
-            webSocketContextData->messageHandler((WebSocket<SSL, true> *) s, std::string_view(data, length));
+            webSocketContextData->messageHandler((WebSocket<SSL, true> *) s, std::string_view(data, length), (uWS::OpCode) opCode);
 
             // the only thing here to check is probably closed
 
diff --git a/src/WebSocketContextData.h b/src/WebSocketContextData.h
index 45b8aaf..0c3634b 100644
--- a/src/WebSocketContextData.h
+++ b/src/WebSocketContextData.h
@@ -4,6 +4,8 @@
 #include <functional>
 #include <string_view>
 
+#include "WebSocketProtocol.h"
+
 namespace uWS {
 
 template <bool, bool> struct WebSocket;
@@ -11,7 +13,7 @@ template <bool, bool> struct WebSocket;
 template <bool SSL>
 struct WebSocketContextData {
 
-    std::function<void(WebSocket<SSL, true> *, std::string_view)> messageHandler;
+    std::function<void(WebSocket<SSL, true> *, std::string_view, uWS::OpCode)> messageHandler;
 };
 
 }