finish short writes in write_websocket_frame and open_stream_ (#2642)

write_websocket_frame and the request-body write in open_stream_ called
Stream::write once and only failed on a negative return, so a short
write (one record per mbedtls_ssl_write call, or SO_SNDTIMEO expiring
after part of the data went out) dropped the rest of the frame or body.
Route both through write_data, as the other body writers do.

Co-authored-by: yhirose <yuji.hirose.bug@gmail.com>

httplib.h

@@ -3797,6 +3797,8 @@ template <typename T>
 bool check_and_write_headers(Stream &strm, Headers &headers, T header_writer,
                              Error &error);
 
+bool write_data(Stream &strm, const char *d, size_t l);
+
 std::string trim_copy(const std::string &s);
 
 void divide(
@@ -5815,32 +5817,40 @@ inline bool write_websocket_frame(Stream &strm, ws::Opcode opcode,
   if (len < 126) {
     header[1] = static_cast<uint8_t>(len);
     if (mask) { header[1] |= 0x80; }
-    if (strm.write(reinterpret_cast<char *>(header), 2) < 0) { return false; }
+    if (!write_data(strm, reinterpret_cast<char *>(header), 2)) {
+      return false;
+    }
   } else if (len <= 0xFFFF) {
     header[1] = 126;
     if (mask) { header[1] |= 0x80; }
-    if (strm.write(reinterpret_cast<char *>(header), 2) < 0) { return false; }
+    if (!write_data(strm, reinterpret_cast<char *>(header), 2)) {
+      return false;
+    }
     uint8_t ext[2];
     ext[0] = static_cast<uint8_t>((len >> 8) & 0xFF);
     ext[1] = static_cast<uint8_t>(len & 0xFF);
-    if (strm.write(reinterpret_cast<char *>(ext), 2) < 0) { return false; }
+    if (!write_data(strm, reinterpret_cast<char *>(ext), 2)) { return false; }
   } else {
     header[1] = 127;
     if (mask) { header[1] |= 0x80; }
-    if (strm.write(reinterpret_cast<char *>(header), 2) < 0) { return false; }
+    if (!write_data(strm, reinterpret_cast<char *>(header), 2)) {
+      return false;
+    }
     uint8_t ext[8];
     for (int i = 7; i >= 0; i--) {
       ext[7 - i] =
           static_cast<uint8_t>((static_cast<uint64_t>(len) >> (i * 8)) & 0xFF);
     }
-    if (strm.write(reinterpret_cast<char *>(ext), 8) < 0) { return false; }
+    if (!write_data(strm, reinterpret_cast<char *>(ext), 8)) { return false; }
   }
 
   if (mask) {
     // Generate random mask key
     uint8_t mask_key[4];
     fill_random_bytes(mask_key, sizeof(mask_key));
-    if (strm.write(reinterpret_cast<char *>(mask_key), 4) < 0) { return false; }
+    if (!write_data(strm, reinterpret_cast<char *>(mask_key), 4)) {
+      return false;
+    }
 
     // Write masked payload in chunks
     const size_t chunk_size = 4096;
@@ -5851,11 +5861,11 @@ inline bool write_websocket_frame(Stream &strm, ws::Opcode opcode,
         buf[i] =
             data[offset + i] ^ static_cast<char>(mask_key[(offset + i) % 4]);
       }
-      if (strm.write(buf.data(), n) < 0) { return false; }
+      if (!write_data(strm, buf.data(), n)) { return false; }
     }
   } else {
     if (len > 0) {
-      if (strm.write(data, len) < 0) { return false; }
+      if (!write_data(strm, data, len)) { return false; }
     }
   }
 
@@ -15969,7 +15979,7 @@ inline ClientImpl::StreamHandle ClientImpl::open_stream_(Request &req) {
   }
 
   if (!req.body.empty()) {
-    if (strm.write(req.body.data(), req.body.size()) < 0) {
+    if (!detail::write_data(strm, req.body.data(), req.body.size())) {
       handle.error = Error::Write;
       handle.response.reset();
       return handle;

test/test.cc

@@ -23859,6 +23859,55 @@ TEST(OpenStreamMalformedContentLength, OutOfRange) {
   server_thread.join();
 }
 
+#ifndef _WIN32
+// As above for the body of a streamed request: a send() cut short by the write
+// timeout must not leave the rest of the body unsent.
+TEST(OpenStreamPartialWriteTest, BodyIsWrittenWhole) {
+  signal(SIGPIPE, SIG_IGN);
+
+  Server svr;
+  svr.set_socket_options([](socket_t sock) {
+    int size = 8192;
+    setsockopt(sock, SOL_SOCKET, SO_RCVBUF,
+               reinterpret_cast<const char *>(&size), sizeof(size));
+  });
+  svr.Post("/size", [](const Request &, Response &res,
+                       const ContentReader &content_reader) {
+    // Outlast the client's write timeout before reading the body.
+    std::this_thread::sleep_for(std::chrono::milliseconds(750));
+    size_t received = 0;
+    content_reader([&](const char *, size_t len) {
+      received += len;
+      return true;
+    });
+    res.set_content(std::to_string(received), "text/plain");
+  });
+
+  auto port = svr.bind_to_any_port(HOST);
+  std::thread t([&]() { svr.listen_after_bind(); });
+  auto se = detail::scope_exit([&] {
+    svr.stop();
+    t.join();
+  });
+  svr.wait_until_ready();
+
+  Client cli(HOST, port);
+  cli.set_write_timeout(0, 500000);
+  cli.set_socket_options([](socket_t sock) {
+    int size = 8192;
+    setsockopt(sock, SOL_SOCKET, SO_SNDBUF,
+               reinterpret_cast<const char *>(&size), sizeof(size));
+  });
+
+  const std::string body(4 * 1024 * 1024, 'A');
+  auto handle = cli.open_stream("POST", "/size", {}, {}, body,
+                                "application/octet-stream");
+  ASSERT_TRUE(handle.is_valid()) << to_string(handle.error);
+  EXPECT_EQ(StatusCode::OK_200, handle.response->status);
+  EXPECT_EQ(std::to_string(body.size()), read_all(handle));
+}
+#endif
+
 // Serves `response` to the single request `fn` makes with a fresh client.
 template <typename Fn>
 static void with_single_response(const std::string &response, Fn fn) {
@@ -28195,6 +28244,64 @@ TEST_F(WebSocketIntegrationTest, ChronoTimeoutSetters) {
   EXPECT_LT(elapsed, 30000);
 }
 
+#ifndef _WIN32
+// A send() that times out after part of the data went out returns the count
+// written so far. Every byte of the frame must still be sent, or the peer reads
+// the next frame's header and payload as the rest of this one.
+TEST(WebSocketPartialWriteTest, FrameIsWrittenWhole) {
+  signal(SIGPIPE, SIG_IGN);
+
+  const std::string large(4 * 1024 * 1024, 'A');
+  std::atomic<bool> large_sent{false};
+  std::atomic<bool> small_sent{false};
+
+  Server svr;
+  svr.set_write_timeout(0, 500000);
+  svr.set_socket_options([](socket_t sock) {
+    int size = 8192;
+    setsockopt(sock, SOL_SOCKET, SO_SNDBUF,
+               reinterpret_cast<const char *>(&size), sizeof(size));
+  });
+  svr.WebSocket("/ws", [&](const Request &, ws::WebSocket &ws) {
+    large_sent = ws.send(large.data(), large.size());
+    small_sent = ws.send("after");
+    std::string msg;
+    while (ws.read(msg)) {}
+  });
+
+  auto port = svr.bind_to_any_port(HOST);
+  std::thread t([&]() { svr.listen_after_bind(); });
+  auto se = detail::scope_exit([&] {
+    svr.stop();
+    t.join();
+  });
+  svr.wait_until_ready();
+
+  ws::WebSocketClient client("ws://" + std::string(HOST) + ":" +
+                             std::to_string(port) + "/ws");
+  client.set_read_timeout(5);
+  client.set_socket_options([](socket_t sock) {
+    int size = 8192;
+    setsockopt(sock, SOL_SOCKET, SO_RCVBUF,
+               reinterpret_cast<const char *>(&size), sizeof(size));
+  });
+  ASSERT_TRUE(client.connect());
+
+  // Outlast the server's write timeout, then drain.
+  std::this_thread::sleep_for(std::chrono::milliseconds(750));
+
+  std::string msg;
+  EXPECT_EQ(ws::ReadResult::Binary, client.read(msg));
+  EXPECT_TRUE(msg == large);
+  EXPECT_EQ(ws::ReadResult::Text, client.read(msg));
+  EXPECT_EQ("after", msg);
+  EXPECT_TRUE(large_sent);
+  EXPECT_TRUE(small_sent);
+
+  client.close();
+}
+#endif
+
 TEST(WebSocketPreRoutingTest, RejectWithoutAuth) {
   Server svr;