From 50de2cea89f227e087ed04966e4106bcd25a386d Mon Sep 17 00:00:00 2001
From: Alex Hultman <alexhultman@gmail.com>
Date: Sat, 29 Dec 2018 13:32:49 +0100
Subject: [PATCH] Hook up maxPayloadLength, not for Zlib

---
 src/App.h                  |  3 +++
 src/WebSocketContext.h     | 13 ++++++-------
 src/WebSocketContextData.h |  3 +++
 src/WebSocketProtocol.h    |  2 +-
 4 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/src/App.h b/src/App.h
index ea959bc..b15efbb 100644
--- a/src/App.h
+++ b/src/App.h
@@ -95,6 +95,9 @@ public:
         webSocketContext->getExt()->drainHandler = behavior.drain;
         webSocketContext->getExt()->closeHandler = behavior.close;
 
+        /* Copy settings */
+        webSocketContext->getExt()->maxPayloadLength = behavior.maxPayloadLength;
+
         return get(pattern, [webSocketContext, this, behavior](auto *res, auto *req) {
             /* If we have this header set, it's a websocket */
             std::string_view secWebSocketKey = req->getHeader("sec-websocket-key");
diff --git a/src/WebSocketContext.h b/src/WebSocketContext.h
index 596ecbf..37232ad 100644
--- a/src/WebSocketContext.h
+++ b/src/WebSocketContext.h
@@ -221,14 +221,13 @@ private:
         return false;
     }
 
-    // bug: todo
-    static bool refusePayloadLength(uint64_t length, uWS::WebSocketState<isServer> *wState) {
+    static bool refusePayloadLength(uint64_t length, uWS::WebSocketState<isServer> *wState, void *s) {
+        WebSocketContextData<SSL> *webSocketContextData = (WebSocketContextData<SSL> *) static_dispatch(us_ssl_socket_context_ext, us_socket_context_ext)(
+            static_dispatch(us_ssl_socket_get_context, us_socket_get_context)((SOCKET_TYPE *)s)
+            );
 
-        //std::cout << "refuse payload length" << std::endl;
-
-        /* We check if we want to accept such a frame based on size */
-        // for now, accept anything
-        return false;
+        /* Return true for refuse, false for accept */
+        return webSocketContextData->maxPayloadLength < length;
     }
 
     WebSocketContext<SSL, isServer> *init() {
diff --git a/src/WebSocketContextData.h b/src/WebSocketContextData.h
index a390bc1..45c4be2 100644
--- a/src/WebSocketContextData.h
+++ b/src/WebSocketContextData.h
@@ -32,6 +32,9 @@ struct WebSocketContextData {
     std::function<void(WebSocket<SSL, true> *, std::string_view, uWS::OpCode)> messageHandler = nullptr;
     std::function<void(uWS::WebSocket<SSL, true> *)> drainHandler = nullptr;
     std::function<void(uWS::WebSocket<SSL, true> *, int, std::string_view)> closeHandler = nullptr;
+
+    /* Settings for this context */
+    size_t maxPayloadLength = 0;
 };
 
 }
diff --git a/src/WebSocketProtocol.h b/src/WebSocketProtocol.h
index 5f0d370..30f0901 100644
--- a/src/WebSocketProtocol.h
+++ b/src/WebSocketProtocol.h
@@ -288,7 +288,7 @@ protected:
         }
         wState->state.lastFin = isFin(src);
 
-        if (Impl::refusePayloadLength(payLength, wState)) {
+        if (Impl::refusePayloadLength(payLength, wState, user)) {
             Impl::forceClose(wState, user);
             return true;
         }