Add websocket feature

1 files changed, 166 insertions, 1 deletions

diff --git a/include/routing.h b/include/routing.h
index 418209c..4fc2de8 100644
--- a/include/routing.h
+++ b/include/routing.h
@@ -13,6 +13,7 @@
 #include "http_request.h"
 #include "utility.h"
 #include "logging.h"
+#include "websocket.h"
 
 namespace crow
 {
@@ -29,8 +30,26 @@ namespace crow
         }
         
         virtual void validate() = 0;
+        std::unique_ptr<BaseRule> upgrade()
+		{
+			if (rule_to_upgrade_)
+				return std::move(rule_to_upgrade_);
+			return {};
+		}
 
         virtual void handle(const request&, response&, const routing_params&) = 0;
+        virtual void handle_upgrade(const request&, response& res, SocketAdaptor&&) 
+		{
+			res = response(404);
+			res.end();
+		}
+#ifdef CROW_ENABLE_SSL
+        virtual void handle_upgrade(const request&, response& res, SSLAdaptor&&) 
+		{
+			res = response(404);
+			res.end();
+		}
+#endif
 
         uint32_t get_methods()
         {
@@ -42,6 +61,9 @@ namespace crow
 
         std::string rule_;
         std::string name_;
+
+		std::unique_ptr<BaseRule> rule_to_upgrade_;
+
         friend class Router;
         template <typename T>
         friend struct RuleParameterTraits;
@@ -233,10 +255,82 @@ namespace crow
         }
     }
 
+	class WebSocketRule : public BaseRule
+	{
+        using self_t = WebSocketRule;
+	public:
+        WebSocketRule(std::string rule)
+            : BaseRule(std::move(rule))
+        {
+        }
+
+		void validate() override
+		{
+		}
+
+		void handle(const request&, response& res, const routing_params&) override
+		{
+			res = response(404);
+			res.end();
+		}
+
+        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_);
+		}
+#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_);
+		}
+#endif
+
+		template <typename Func>
+		self_t& onopen(Func f)
+		{
+			open_handler_ = f;
+			return *this;
+		}
+
+		template <typename Func>
+		self_t& onmessage(Func f)
+		{
+			message_handler_ = f;
+			return *this;
+		}
+
+		template <typename Func>
+		self_t& onclose(Func f)
+		{
+			close_handler_ = f;
+			return *this;
+		}
+
+		template <typename Func>
+		self_t& onerror(Func f)
+		{
+			error_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_;
+	};
+
     template <typename T>
     struct RuleParameterTraits
     {
         using self_t = T;
+		WebSocketRule& websocket() 
+		{
+			auto p =new WebSocketRule(((self_t*)this)->rule_);
+            ((self_t*)this)->rule_to_upgrade_.reset(p);
+			return *p;
+		}
+
         self_t& name(std::string name) noexcept
         {
             ((self_t*)this)->name_ = std::move(name);
@@ -256,6 +350,7 @@ namespace crow
             ((self_t*)this)->methods_ |= 1 << (int)method;
             return (self_t&)*this;
         }
+
     };
 
     class DynamicRule : public BaseRule, public RuleParameterTraits<DynamicRule>
@@ -343,7 +438,7 @@ namespace crow
         {
         }
 
-        void validate()
+        void validate() override
         {
             if (!handler_)
             {
@@ -809,10 +904,80 @@ public:
             for(auto& rule:rules_)
             {
                 if (rule)
+				{
+					auto upgraded = rule->upgrade();
+					if (upgraded)
+						rule = std::move(upgraded);
                     rule->validate();
+				}
             }
         }
 
+		template <typename Adaptor> 
+		void handle_upgrade(const request& req, response& res, Adaptor&& adaptor)
+		{
+            auto found = trie_.find(req.url);
+            unsigned rule_index = found.first;
+            if (!rule_index)
+            {
+                CROW_LOG_DEBUG << "Cannot match rules " << req.url;
+                res = response(404);
+                res.end();
+                return;
+            }
+
+            if (rule_index >= rules_.size())
+                throw std::runtime_error("Trie internal structure corrupted!");
+
+            if (rule_index == RULE_SPECIAL_REDIRECT_SLASH)
+            {
+                CROW_LOG_INFO << "Redirecting to a url with trailing slash: " << req.url;
+                res = response(301);
+
+                // TODO absolute url building
+                if (req.get_header_value("Host").empty())
+                {
+                    res.add_header("Location", req.url + "/");
+                }
+                else
+                {
+                    res.add_header("Location", "http://" + req.get_header_value("Host") + req.url + "/");
+                }
+                res.end();
+                return;
+            }
+
+            if ((rules_[rule_index]->get_methods() & (1<<(uint32_t)req.method)) == 0)
+            {
+                CROW_LOG_DEBUG << "Rule found but method mismatch: " << req.url << " with " << method_name(req.method) << "(" << (uint32_t)req.method << ") / " << rules_[rule_index]->get_methods();
+                res = response(404);
+                res.end();
+                return;
+            }
+
+            CROW_LOG_DEBUG << "Matched rule (upgrade) '" << rules_[rule_index]->rule_ << "' " << (uint32_t)req.method << " / " << rules_[rule_index]->get_methods();
+
+            // any uncaught exceptions become 500s
+            try
+            {
+                rules_[rule_index]->handle_upgrade(req, res, std::move(adaptor));
+            }
+            catch(std::exception& e)
+            {
+                CROW_LOG_ERROR << "An uncaught exception occurred: " << e.what();
+                res = response(500);
+                res.end();
+                return;   
+            }
+            catch(...)
+            {
+                CROW_LOG_ERROR << "An uncaught exception occurred. The type was unknown so no information was available.";
+                res = response(500);
+                res.end();
+                return;   
+            }
+		}
+
         void handle(const request& req, response& res)
         {
             auto found = trie_.find(req.url);