From bea1ba3797fbf8bbab209edbbf765d0e6d5960d1 Mon Sep 17 00:00:00 2001
From: Marcos Medeiros <marcos@engemap.com.br>
Date: Wed, 28 Dec 2016 11:06:56 -0200
Subject: [PATCH] Add onaccept handler to websocket rule

---
 include/crow/routing.h   | 12 ++++++++++--
 include/crow/websocket.h | 16 +++++++++++++++-
 2 files changed, 25 insertions(+), 3 deletions(-)

diff --git a/include/crow/routing.h b/include/crow/routing.h
index a8aa233c9..95fe3541b 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<SocketAdaptor>(req, std::move(adaptor), open_handler_, message_handler_, close_handler_, error_handler_);
+			new crow::websocket::Connection<SocketAdaptor>(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<SSLAdaptor>(req, std::move(adaptor), open_handler_, message_handler_, close_handler_, error_handler_);
+			new crow::websocket::Connection<SSLAdaptor>(req, std::move(adaptor), open_handler_, message_handler_, close_handler_, error_handler_, accept_handler_);
 		}
 #endif
 
@@ -313,11 +313,19 @@ namespace crow
 			return *this;
 		}
 
+		template <typename Func>
+		self_t& onaccept(Func f)
+		{
+		    accept_handler_ = f;
+		    return *this;
+		}
+
 	protected:
 		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_;
+		std::function<bool(const crow::request&)> accept_handler_;
 	};
 
     template <typename T>
diff --git a/include/crow/websocket.h b/include/crow/websocket.h
index a1e8e8f75..fa45d1b25 100644
--- a/include/crow/websocket.h
+++ b/include/crow/websocket.h
@@ -39,8 +39,10 @@ namespace crow
 						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)
+						std::function<void(crow::websocket::connection&)> error_handler,
+						std::function<bool(const crow::request&)> 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"))
 					{
@@ -48,6 +50,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";
@@ -484,6 +497,7 @@ namespace crow
 				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_;
+				std::function<bool(const crow::request&)> accept_handler_;
         };
     }
 }