Add support for compressing websockets

rustyx · rustyx · commit 6a719268cc00 · 2019-02-11T21:16:41.000+01:00
PR ipkn#329 fran6co
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -11,10 +11,15 @@ set(CMAKE_MODULE_PATH ${CMAKE_MODULE_PATH} "${CMAKE_SOURCE_DIR}/cmake/")
 find_package(Tcmalloc)
 find_package(Threads)
 find_package(OpenSSL)
+find_package(zlib)
 
 if(OPENSSL_FOUND)
 	include_directories(${OPENSSL_INCLUDE_DIR})
 endif()
+if(ZLIB_FOUND)
+	add_definitions(-DHAVE_ZLIB)
+	include_directories(${ZLIB_INCLUDE_DIR})
+endif()
 
 find_program(CCACHE_FOUND ccache)
 if(CCACHE_FOUND)
diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt
@@ -3,21 +3,30 @@ project (crow_examples)
 
 if (MSVC)
 	add_executable(example_vs example_vs.cpp)
-	target_link_libraries(example_vs ${Boost_LIBRARIES})
+	target_link_libraries(example_vs ${Boost_LIBRARIES} ${ZLIB_LIBRARIES})
 	target_link_libraries(example_vs ${CMAKE_THREAD_LIBS_INIT})
+	add_executable(example_websocket websocket/example_ws.cpp)
+	target_link_libraries(example_websocket ${Boost_LIBRARIES} ${ZLIB_LIBRARIES})
+	target_link_libraries(example_websocket ${CMAKE_THREAD_LIBS_INIT})
+	add_custom_command(OUTPUT ws.html
+	    COMMAND ${CMAKE_COMMAND} -E
+	    copy ${PROJECT_SOURCE_DIR}/websocket/templates/ws.html ${CMAKE_CURRENT_BINARY_DIR}/templates/ws.html
+	    DEPENDS ${PROJECT_SOURCE_DIR}/websocket/templates/ws.html
+	)
+	add_custom_target(example_ws_copy ALL DEPENDS ws.html)
 else ()
 	add_executable(helloworld helloworld.cpp)
-	target_link_libraries(helloworld ${Boost_LIBRARIES})
+	target_link_libraries(helloworld ${Boost_LIBRARIES} ${ZLIB_LIBRARIES})
 	target_link_libraries(helloworld ${CMAKE_THREAD_LIBS_INIT})
 
 if (OPENSSL_FOUND)
 	add_executable(example_ssl ssl/example_ssl.cpp)
-	target_link_libraries(example_ssl ${Boost_LIBRARIES})
+	target_link_libraries(example_ssl ${Boost_LIBRARIES} ${ZLIB_LIBRARIES})
 	target_link_libraries(example_ssl ${CMAKE_THREAD_LIBS_INIT} ${OPENSSL_LIBRARIES})
 endif()
 
 add_executable(example_websocket websocket/example_ws.cpp)
-target_link_libraries(example_websocket ${Boost_LIBRARIES})
+target_link_libraries(example_websocket ${Boost_LIBRARIES} ${ZLIB_LIBRARIES})
 target_link_libraries(example_websocket ${CMAKE_THREAD_LIBS_INIT} ${OPENSSL_LIBRARIES})
 add_custom_command(OUTPUT ws.html
     COMMAND ${CMAKE_COMMAND} -E
@@ -28,7 +37,7 @@ add_custom_target(example_ws_copy ALL DEPENDS ws.html)
 
 add_executable(example example.cpp)
 #target_link_libraries(example crow)
-target_link_libraries(example ${Boost_LIBRARIES})
+target_link_libraries(example ${Boost_LIBRARIES} ${ZLIB_LIBRARIES})
 target_link_libraries(example ${CMAKE_THREAD_LIBS_INIT})
 
 if (Tcmalloc_FOUND)
@@ -38,7 +47,7 @@ endif(Tcmalloc_FOUND)
 add_executable(example_with_all example_with_all.cpp)
 add_dependencies(example_with_all amalgamation)
 #target_link_libraries(example crow)
-target_link_libraries(example_with_all ${Boost_LIBRARIES})
+target_link_libraries(example_with_all ${Boost_LIBRARIES} ${ZLIB_LIBRARIES})
 target_link_libraries(example_with_all ${CMAKE_THREAD_LIBS_INIT})
 
 add_custom_command(OUTPUT example_test.py
@@ -49,7 +58,7 @@ add_custom_command(OUTPUT example_test.py
 add_custom_target(example_copy ALL DEPENDS example_test.py)
 
 add_executable(example_chat example_chat.cpp)
-target_link_libraries(example_chat ${Boost_LIBRARIES})
+target_link_libraries(example_chat ${Boost_LIBRARIES} ${ZLIB_LIBRARIES})
 target_link_libraries(example_chat ${CMAKE_THREAD_LIBS_INIT})
 add_custom_command(OUTPUT example_chat.html
 	COMMAND ${CMAKE_COMMAND} -E
diff --git a/include/crow/settings.h b/include/crow/settings.h
@@ -11,6 +11,11 @@
 /* #ifdef - enables ssl */
 //#define CROW_ENABLE_SSL
 
+/* #ifdef - enables websocket compression */
+#ifdef HAVE_ZLIB
+#define CROW_ENABLE_WEBSOCKET_COMPRESSION
+#endif
+
 /* #define - specifies log level */
 /*
     Debug       = 0
diff --git a/include/crow/socket_adaptors.h b/include/crow/socket_adaptors.h
@@ -90,11 +90,11 @@ namespace crow
 
         void close()
         {
-            if (ssl_socket_)
-            {
-                boost::system::error_code ec;
-                raw_socket().close(ec);
-            }
+            if (ssl_socket_)
+            {
+                boost::system::error_code ec;
+                raw_socket().close(ec);
+            }
         }
 
         boost::asio::io_service& get_io_service()
diff --git a/include/crow/websocket.h b/include/crow/websocket.h
@@ -1,9 +1,11 @@
 #pragma once
 #include <boost/algorithm/string/predicate.hpp>
+#include <boost/algorithm/string/split.hpp>
 #include <boost/array.hpp>
 #include "crow/socket_adaptors.h"
 #include "crow/http_request.h"
 #include "crow/TinySHA1.hpp"
+#include "crow/zlib.hpp"
 
 namespace crow
 {
@@ -61,8 +63,19 @@ namespace crow
 							return;
 						}
 					}
-
-					// Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==
+#ifdef CROW_ENABLE_WEBSOCKET_COMPRESSION
+                    std::string extensionsHeader = req.get_header_value("Sec-WebSocket-Extensions");
+                    std::vector<std::string> extensions;
+                    boost::split(extensions, extensionsHeader, boost::is_any_of(";"));
+                    if (std::find(extensions.begin(), extensions.end(), "permessage-deflate") != extensions.end())
+                    {
+                        bool reset_compressor_on_send_ = std::find(extensions.begin(), extensions.end(), "server_no_context_takeover") != extensions.end();
+                        compressor_.reset(new zlib_compressor(reset_compressor_on_send_, true, 15, Z_BEST_COMPRESSION, 8, Z_DEFAULT_STRATEGY));
+                        bool reset_decompressor_on_send_ = std::find(extensions.begin(), extensions.end(), "client_no_context_takeover") != extensions.end();
+                        decompressor_.reset(new zlib_decompressor(reset_decompressor_on_send_, true, 15));
+                    }
+#endif
+                    // Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==
 					// Sec-WebSocket-Version: 13
                     std::string magic = req.get_header_value("Sec-WebSocket-Key") +  "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
                     sha1::SHA1 s;
@@ -98,19 +111,35 @@ namespace crow
                 void send_binary(const std::string& msg) override
                 {
                     dispatch([this, msg]{
-                        auto header = build_header(2, msg.size());
-                        write_buffers_.emplace_back(std::move(header));
-                        write_buffers_.emplace_back(msg);
+                        if (compressor_) {
+                            std::string msg_ = compressor_->compress(msg);
+                            auto header = build_header(0x42, msg_.size());
+                            write_buffers_.emplace_back(std::move(header));
+                            write_buffers_.emplace_back(std::move(msg_));
+                        }
+                        else {
+                            auto header = build_header(2, msg.size());
+                            write_buffers_.emplace_back(std::move(header));
+                            write_buffers_.emplace_back(msg);
+                        }
                         do_write();
                     });
                 }
 
                 void send_text(const std::string& msg) override
                 {
                     dispatch([this, msg]{
-                        auto header = build_header(1, msg.size());
-                        write_buffers_.emplace_back(std::move(header));
-                        write_buffers_.emplace_back(msg);
+                        if (compressor_) {
+                            std::string msg_ = compressor_->compress(msg);
+                            auto header = build_header(0x41, msg_.size());
+                            write_buffers_.emplace_back(std::move(header));
+                            write_buffers_.emplace_back(std::move(msg_));
+                        }
+                        else {
+                            auto header = build_header(1, msg.size());
+                            write_buffers_.emplace_back(std::move(header));
+                            write_buffers_.emplace_back(msg);
+                        }
                         do_write();
                     });
                 }
@@ -167,6 +196,16 @@ namespace crow
                     write_buffers_.emplace_back(header);
                     write_buffers_.emplace_back(std::move(hello));
                     write_buffers_.emplace_back(crlf);
+                    if (compressor_ && decompressor_) {
+                        write_buffers_.emplace_back(
+                            "Sec-WebSocket-Extensions: permessage-deflate"
+                            "; server_max_window_bits=" + std::to_string(decompressor_->window_bits) +
+                            "; client_max_window_bits=" + std::to_string(compressor_->window_bits) +
+                            (compressor_->reset_before_compress ? "; server_no_context_takeover" : "") +
+                            (decompressor_->reset_before_decompress ? "; client_no_context_takeover" : "")
+                        );
+                        write_buffers_.emplace_back(crlf);
+                    }
                     write_buffers_.emplace_back(crlf);
                     do_write();
                     if (open_handler_)
@@ -368,6 +407,11 @@ namespace crow
                     return mini_header_ & 0x8000;
                 }
 
+                bool is_compressed()
+                {
+                    return mini_header_ & 0x4000;
+                }
+
                 int opcode()
                 {
                     return (mini_header_ & 0x0f00) >> 8;
@@ -387,7 +431,7 @@ namespace crow
                                 if (is_FIN())
                                 {
                                     if (message_handler_)
-                                        message_handler_(*this, message_, is_binary_);
+                                        message_handler_(*this, (is_compressed() && decompressor_) ? decompressor_->decompress(message_) : message_, is_binary_);
                                     message_.clear();
                                 }
                             }
@@ -398,7 +442,7 @@ namespace crow
                                 if (is_FIN())
                                 {
                                     if (message_handler_)
-                                        message_handler_(*this, message_, is_binary_);
+                                        message_handler_(*this, (is_compressed() && decompressor_) ? decompressor_->decompress(message_) : message_, is_binary_);
                                     message_.clear();
                                 }
                             }
@@ -410,7 +454,7 @@ namespace crow
                                 if (is_FIN())
                                 {
                                     if (message_handler_)
-                                        message_handler_(*this, message_, is_binary_);
+                                        message_handler_(*this, (is_compressed() && decompressor_) ? decompressor_->decompress(message_) : message_, is_binary_);
                                     message_.clear();
                                 }
                             }
@@ -514,7 +558,12 @@ namespace crow
                 bool pong_received_{false};
                 bool is_close_handler_called_{false};
 
-				std::function<void(crow::websocket::connection&)> open_handler_;
+                bool reset_compressor_on_send_{ false };
+                bool reset_decompressor_on_send_{ false };
+                std::unique_ptr<zlib_compressor> compressor_;
+                std::unique_ptr<zlib_decompressor> decompressor_;
+
+                std::function<void(crow::websocket::connection&)> open_handler_;
 				std::function<void(crow::websocket::connection&, const std::string&, bool)> message_handler_;
 				std::function<void(crow::websocket::connection&, const std::string&)> close_handler_;
 				std::function<void(crow::websocket::connection&)> error_handler_;
diff --git a/include/crow/zlib.hpp b/include/crow/zlib.hpp
@@ -0,0 +1,132 @@
+#pragma once
+#include "crow/settings.h"
+#include <boost/asio/streambuf.hpp>
+#include <string>
+#ifdef HAVE_ZLIB
+#include <zlib.h>
+
+namespace crow {
+    class zlib_compressor {
+    public:
+        zlib_compressor(bool reset_before_compress, bool noheader, int window_bits, int level, int mem_level, int strategy)
+            : reset_before_compress(reset_before_compress)
+            , window_bits(window_bits) {
+            stream = std::make_unique<z_stream>();
+            stream->zalloc = 0;
+            stream->zfree = 0;
+            stream->opaque = 0;
+
+            ::deflateInit2(stream.get(),
+                level,
+                Z_DEFLATED,
+                (noheader ? -window_bits : window_bits),
+                mem_level,
+                strategy
+            );
+        }
+
+        ~zlib_compressor() {
+            ::deflateEnd(stream.get());
+        }
+
+        std::string compress(const std::string & src) {
+            if (reset_before_compress)
+                ::deflateReset(stream.get());
+
+            stream->next_in = reinterpret_cast<uint8_t*>(const_cast<char*>(src.c_str()));
+            stream->avail_in = src.size();
+
+            const uint64_t bufferSize = 256;
+            boost::asio::streambuf buffer;
+            do {
+                boost::asio::streambuf::mutable_buffers_type chunk = buffer.prepare(bufferSize);
+
+                uint8_t* next_out = boost::asio::buffer_cast<uint8_t*>(chunk);
+
+                stream->next_out = next_out;
+                stream->avail_out = bufferSize;
+
+                ::deflate(stream.get(), reset_before_compress ? Z_FINISH : Z_SYNC_FLUSH);
+
+                uint64_t outputSize = stream->next_out - next_out;
+                buffer.commit(outputSize);
+            } while (stream->avail_out == 0);
+
+            uint64_t buffer_size = buffer.size();
+            if (!reset_before_compress) buffer_size -= 4;
+
+            return std::string(boost::asio::buffer_cast<const char*>(buffer.data()), buffer_size);
+        }
+
+        std::unique_ptr<z_stream> stream;
+
+        bool reset_before_compress;
+        int window_bits;
+    };
+
+    class zlib_decompressor {
+    public:
+        zlib_decompressor(bool reset_before_decompress, bool noheader, int window_bits)
+            : reset_before_decompress(reset_before_decompress)
+            , window_bits(window_bits) {
+            stream = std::make_unique<z_stream>();
+            stream->zalloc = 0;
+            stream->zfree = 0;
+            stream->opaque = 0;
+
+            ::inflateInit2(stream.get(), (noheader ? -window_bits : window_bits));
+        }
+
+        ~zlib_decompressor() {
+            inflateEnd(stream.get());
+        }
+
+        std::string decompress(std::string src) {
+            if (reset_before_decompress)
+                inflateReset(stream.get());
+
+            src.push_back('\x00');
+            src.push_back('\x00');
+            src.push_back('\xff');
+            src.push_back('\xff');
+
+            stream->next_in = reinterpret_cast<uint8_t*>(const_cast<char*>(src.c_str()));
+            stream->avail_in = src.size();
+
+            const uint64_t bufferSize = 256;
+            boost::asio::streambuf buffer;
+            do {
+                boost::asio::streambuf::mutable_buffers_type chunk = buffer.prepare(bufferSize);
+
+                uint8_t* next_out = boost::asio::buffer_cast<uint8_t*>(chunk);
+
+                stream->next_out = next_out;
+                stream->avail_out = bufferSize;
+
+                ::inflate(stream.get(), reset_before_decompress ? Z_FINISH : Z_SYNC_FLUSH);
+                buffer.commit(stream->next_out - next_out);
+            } while (stream->avail_out == 0);
+
+            return std::string(boost::asio::buffer_cast<const char*>(buffer.data()), buffer.size());
+        }
+
+        std::unique_ptr<z_stream> stream;
+
+        bool reset_before_decompress;
+        int window_bits;
+    };
+}
+#else
+namespace crow {
+    struct zlib_compressor {
+        std::string compress(const std::string& src) { return src; }
+        bool reset_before_compress;
+        int window_bits;
+    };
+    struct zlib_decompressor {
+        std::string decompress(const std::string& src) { return src; }
+        bool reset_before_decompress;
+        int window_bits;
+    };
+}
+#endif
