Add optional keep-alive mechanism to WS connections

include/crow/routing.h

@@ -512,6 +512,14 @@ namespace crow // NOTE: Already documented in "crow/app.h"
             return *this;
         }
 
+        template<typename Func>
+        self_t& ontimeout(Func f, uint64_t timeout_in_seconds = 5)
+        {
+            timeout_handler_.first = f;
+            timeout_handler_.second = timeout_in_seconds;
+            return *this;
+        }
+
         template<typename Func>
         self_t& onaccept(Func f)
         {
@@ -525,6 +533,7 @@ namespace crow // NOTE: Already documented in "crow/app.h"
         std::function<void(crow::websocket::connection&, const std::string&, bool)> message_handler_;
         std::function<void(crow::websocket::connection&, const std::string&, uint16_t)> close_handler_;
         std::function<void(crow::websocket::connection&, const std::string&)> error_handler_;
+        std::pair<std::function<void(crow::websocket::connection&, const std::string&)>, uint64_t> timeout_handler_;
         std::function<bool(const crow::request&, void**)> accept_handler_;
         uint64_t max_payload_;
         bool max_payload_override_ = false;

include/crow/websocket.h

@@ -116,6 +116,7 @@ namespace crow // NOTE: Already documented in "crow/app.h"
                        std::function<void(crow::websocket::connection&, const std::string&, bool)> message_handler,
                        std::function<void(crow::websocket::connection&, const std::string&, uint16_t)> close_handler,
                        std::function<void(crow::websocket::connection&, const std::string&)> error_handler,
+                       std::pair<std::function<void(crow::websocket::connection&, const std::string&)>, uint64_t> receiver_timeout_handler,
                        std::function<bool(const crow::request&, void**)> accept_handler):
               adaptor_(std::move(adaptor)),
               handler_(handler),
@@ -124,7 +125,9 @@ namespace crow // NOTE: Already documented in "crow/app.h"
               message_handler_(std::move(message_handler)),
               close_handler_(std::move(close_handler)),
               error_handler_(std::move(error_handler)),
-              accept_handler_(std::move(accept_handler))
+              timeout_handler_(std::move(receiver_timeout_handler)),
+              accept_handler_(std::move(accept_handler)),
+              task_timer_(adaptor_.get_io_service())
             {
                 if (!utility::string_equals(req.get_header_value("upgrade"), "websocket"))
                 {
@@ -339,6 +342,26 @@ namespace crow // NOTE: Already documented in "crow/app.h"
                 do_read();
             }
 
+            void start_deadline(/*int timeout = 5*/)
+            {
+                cancel_deadline_timer();
+
+                if (close_connection_ || !timeout_handler_.first) return;
+
+                task_timer_.set_default_timeout(timeout_handler_.second);
+                task_id_ = task_timer_.schedule([this] {
+                    timeout_handler_.first(*this, "timeout");
+                });
+                CROW_LOG_DEBUG << this << " websocket timer added: " << &task_timer_ << ' ' << task_id_;
+            }
+
+            void cancel_deadline_timer()
+            {
+                if (!timeout_handler_.first) return;
+                CROW_LOG_DEBUG << this << " websocket timer cancelled: " << &task_timer_ << ' ' << task_id_;
+                task_timer_.cancel(task_id_);
+            }
+
             /// Read a websocket message.
 
             ///
@@ -357,6 +380,8 @@ namespace crow // NOTE: Already documented in "crow/app.h"
                     return;
                 }
 
+                start_deadline();
+
                 is_reading = true;
                 switch (state_)
                 {
@@ -822,7 +847,11 @@ namespace crow // NOTE: Already documented in "crow/app.h"
             std::function<void(crow::websocket::connection&, const std::string&, bool)> message_handler_;
             std::function<void(crow::websocket::connection&, const std::string&, uint16_t status_code)> close_handler_;
             std::function<void(crow::websocket::connection&, const std::string&)> error_handler_;
+            std::pair<std::function<void(crow::websocket::connection&, const std::string&)>, uint64_t> timeout_handler_;
             std::function<bool(const crow::request&, void**)> accept_handler_;
+
+            detail::task_timer task_timer_;
+            detail::task_timer::identifier_type task_id_;
         };
     } // namespace websocket
 } // namespace crow

tests/unittest.cpp

@@ -2741,6 +2741,11 @@ TEST_CASE("websocket")
           else if (isbin && message == "Hello bin")
               conn.send_binary("Hello back bin");
       })
+      .ontimeout([&](websocket::connection& conn, const std::string&) {
+          CROW_LOG_INFO << "Websocket Time Out";
+          conn.send_text("TimeOut");
+      },
+                 2 /* seconds */)
       .onclose([&](websocket::connection&, const std::string&, uint16_t) {
           CROW_LOG_INFO << "Closing websocket";
       });
@@ -2875,6 +2880,17 @@ TEST_CASE("websocket")
         std::string checkstring(std::string(buf).substr(0, 12));
         CHECK(checkstring == "\x81\x0AHello back");
     }
+
+    //----------TimeOut----------
+    {
+        std::fill_n(buf, 2048, 0);
+        CROW_LOG_INFO << "Waiting Time Out";
+        c.receive(asio::buffer(buf, 2048));
+        std::this_thread::sleep_for(std::chrono::milliseconds(5));
+        std::string checkstring(std::string(buf).substr(0, 10));
+        CHECK(checkstring == "\x81\x07TimeOut");
+    }
+
     //----------Close----------
     {
         std::fill_n(buf, 2048, 0);