client read ws timeout (#674)

qicosmos · Dec 16, 2024 · 8eb4847 · 8eb4847
1 parent b2adab6
commit 8eb4847
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 1 deletion.
diff --git a/include/cinatra/coro_http_client.hpp b/include/cinatra/coro_http_client.hpp
@@ -370,6 +370,8 @@ class coro_http_client : public std::enable_shared_from_this<coro_http_client> {
   }
 
   async_simple::coro::Lazy<resp_data> read_websocket() {
+    auto time_out_guard =
+        timer_guard(this, req_timeout_duration_, "websocket timer");
     co_return co_await async_read_ws();
   }
 
@@ -2141,6 +2143,9 @@ class coro_http_client : public std::enable_shared_from_this<coro_http_client> {
       if (auto [ec, _] = co_await async_read_ws(
               sock, read_buf, ws.left_header_len(), has_init_ssl);
           ec) {
+        if (socket_->is_timeout_) {
+          co_return resp_data{std::make_error_code(std::errc::timed_out), 404};
+        }
         data.net_err = ec;
         data.status = 404;
 

diff --git a/tests/test_cinatra_websocket.cpp b/tests/test_cinatra_websocket.cpp
@@ -106,6 +106,26 @@ TEST_CASE("test websocket") {
             break;
           }
 
+          auto ec = co_await req.get_conn()->write_websocket(result.data);
+          if (ec) {
+            break;
+          }
+        }
+      });
+  server.set_http_handler<cinatra::GET>(
+      "/test_client_timeout",
+      [](coro_http_request &req,
+         coro_http_response &resp) -> async_simple::coro::Lazy<void> {
+        CHECK(req.get_content_type() == content_type::websocket);
+        websocket_result result{};
+        while (true) {
+          result = co_await req.get_conn()->read_websocket();
+          if (result.ec) {
+            break;
+          }
+
+          std::this_thread::sleep_for(200ms);
+
           auto ec = co_await req.get_conn()->write_websocket(result.data);
           if (ec) {
             break;
@@ -114,7 +134,23 @@ TEST_CASE("test websocket") {
       });
   server.async_start();
 
-  std::this_thread::sleep_for(std::chrono::milliseconds(100));
+  auto client_timeout = []() -> async_simple::coro::Lazy<void> {
+    coro_http_client client{};
+    client.set_req_timeout(50ms);
+    client.set_ws_sec_key("s//GYHa/XO7Hd2F2eOGfyA==");
+
+    auto r = co_await client.connect("ws://localhost:8090/test_client_timeout");
+    if (r.net_err) {
+      co_return;
+    }
+
+    co_await client.write_websocket("hello websocket");
+    auto data = co_await client.read_websocket();
+    std::cout << data.net_err.message() << std::endl;
+    CHECK(data.net_err == std::errc::timed_out);
+  };
+
+  async_simple::coro::syncAwait(client_timeout());
 
   coro_http_client client{};
   client.set_ws_sec_key("s//GYHa/XO7Hd2F2eOGfyA==");