1 files changed, 21 insertions, 3 deletions

diff --git a/http.cpp b/http.cpp
index 4738ad8..469f1aa 100644
--- a/http.cpp
+++ b/http.cpp
@@ -4,6 +4,7 @@
 
 #include "server.h"
 #include "response.h"
+#include "websocket.h"
 
 #include <boost/beast/version.hpp>
 #include <boost/beast/core.hpp>
@@ -27,6 +28,7 @@ namespace beast = boost::beast;         // from <boost/beast.hpp>
 namespace http = beast::http;           // from <boost/beast/http.hpp>
 namespace net = boost::asio;            // from <boost/asio.hpp>
 using tcp = boost::asio::ip::tcp;       // from <boost/asio/ip/tcp.hpp>
+namespace websocket = beast::websocket;
 
 namespace {
 
@@ -41,6 +43,7 @@ void fail(beast::error_code ec, char const* what)
 // Handles an HTTP server connection
 class session : public std::enable_shared_from_this<session>
 {
+ boost::asio::io_context& ioc_;
  beast::tcp_stream stream_;
  beast::flat_buffer buffer_;
  Server& m_server;
@@ -65,13 +68,21 @@ class session : public std::enable_shared_from_this<session>
           sp->need_eof()));
  }
 
+ void handle_websocket()
+ {
+  beast::get_lowest_layer(stream_).expires_never();
+  make_websocket_session(ioc_, std::move(stream_), response::get_websocket_address(req_, m_server), parser_->release());
+ }
+
 public:
     // Take ownership of the stream
     session(
+        boost::asio::io_context& ioc,
         tcp::socket&& socket,
-        Server& server)
-        : stream_(std::move(socket))
-        , m_server(server)
+        Server& server):
+         ioc_(ioc),
+         stream_(std::move(socket)),
+         m_server(server)
     {
     }
 
@@ -127,6 +138,12 @@ public:
 
         req_ = parser_->get();
         
+        if (websocket::is_upgrade(req_))
+        {
+         handle_websocket();
+         return;
+        }
+        
         // Send the response
         handle_request(m_server, std::move(req_));
     }
@@ -251,6 +268,7 @@ private:
         {
             // Create the session and run it
             std::make_shared<session>(
+                ioc_,
                 std::move(socket),
                 m_server)->run();
         }