diff --git a/include/crow/routing.h b/include/crow/routing.h index 5d887e4f2..cdfa48015 100644 --- a/include/crow/routing.h +++ b/include/crow/routing.h @@ -276,12 +276,12 @@ namespace crow void handle_upgrade(const request& req, response&, SocketAdaptor&& adaptor) override { - new crow::websocket::Connection(req, std::move(adaptor), open_handler_, message_handler_, close_handler_, error_handler_); + new crow::websocket::Connection(req, std::move(adaptor), open_handler_, message_handler_, close_handler_, error_handler_, accept_handler_); } #ifdef CROW_ENABLE_SSL void handle_upgrade(const request& req, response&, SSLAdaptor&& adaptor) override { - new crow::websocket::Connection(req, std::move(adaptor), open_handler_, message_handler_, close_handler_, error_handler_); + new crow::websocket::Connection(req, std::move(adaptor), open_handler_, message_handler_, close_handler_, error_handler_, accept_handler_); } #endif @@ -313,11 +313,19 @@ namespace crow return *this; } + template + self_t& onaccept(Func f) + { + accept_handler_ = f; + return *this; + } + protected: std::function open_handler_; std::function message_handler_; std::function close_handler_; std::function error_handler_; + std::function accept_handler_; }; template diff --git a/include/crow/websocket.h b/include/crow/websocket.h index c03547cd0..efcf834dc 100644 --- a/include/crow/websocket.h +++ b/include/crow/websocket.h @@ -40,8 +40,10 @@ namespace crow std::function open_handler, std::function message_handler, std::function close_handler, - std::function error_handler) + std::function error_handler, + std::function accept_handler) : adaptor_(std::move(adaptor)), open_handler_(std::move(open_handler)), message_handler_(std::move(message_handler)), close_handler_(std::move(close_handler)), error_handler_(std::move(error_handler)) + , accept_handler_(std::move(accept_handler)) { if (!boost::iequals(req.get_header_value("upgrade"), "websocket")) { @@ -49,6 +51,17 @@ namespace crow delete this; return; } + + if (accept_handler_) + { + if (!accept_handler_(req)) + { + adaptor.close(); + delete this; + return; + } + } + // Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ== // Sec-WebSocket-Version: 13 std::string magic = req.get_header_value("Sec-WebSocket-Key") + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"; @@ -485,6 +498,7 @@ namespace crow std::function message_handler_; std::function close_handler_; std::function error_handler_; + std::function accept_handler_; }; } }