diff --git a/include/crow/routing.h b/include/crow/routing.h index 5477d0fac..550f5c394 100644 --- a/include/crow/routing.h +++ b/include/crow/routing.h @@ -512,6 +512,14 @@ namespace crow // NOTE: Already documented in "crow/app.h" return *this; } + template + self_t& ontimeout(Func f, uint64_t timeout_in_seconds = 5) + { + timeout_handler_.first = f; + timeout_handler_.second = timeout_in_seconds; + return *this; + } + template self_t& onaccept(Func f) { @@ -525,6 +533,7 @@ namespace crow // NOTE: Already documented in "crow/app.h" std::function message_handler_; std::function close_handler_; std::function error_handler_; + std::pair, uint64_t> timeout_handler_; std::function accept_handler_; uint64_t max_payload_; bool max_payload_override_ = false; diff --git a/include/crow/websocket.h b/include/crow/websocket.h index 8e5bd8731..94ddc74ea 100644 --- a/include/crow/websocket.h +++ b/include/crow/websocket.h @@ -116,6 +116,7 @@ namespace crow // NOTE: Already documented in "crow/app.h" std::function message_handler, std::function close_handler, std::function error_handler, + std::pair, uint64_t> receiver_timeout_handler, std::function 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 message_handler_; std::function close_handler_; std::function error_handler_; + std::pair, uint64_t> timeout_handler_; std::function accept_handler_; + + detail::task_timer task_timer_; + detail::task_timer::identifier_type task_id_; }; } // namespace websocket } // namespace crow diff --git a/tests/unittest.cpp b/tests/unittest.cpp index fe48576f5..282452c3b 100644 --- a/tests/unittest.cpp +++ b/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);