diff --git a/httplib.h b/httplib.h index 35b84e3b..fdfbe0f8 100644 --- a/httplib.h +++ b/httplib.h @@ -5766,12 +5766,17 @@ inline FrameRead read_websocket_frame(Stream &strm, Opcode &opcode, } } - // Read payload - payload.resize(static_cast(payload_len)); - if (payload_len > 0 && - read_exact(strm, &payload[0], static_cast(payload_len)) != - FrameRead::Ok) { - return FrameRead::Fail; + // Grow the payload only as it arrives, since a header alone can announce + // max_len bytes that never come. + payload.clear(); + auto len = static_cast(payload_len); + while (payload.size() < len) { + auto offset = payload.size(); + auto n = (std::min)(len - offset, CPPHTTPLIB_RECV_BUFSIZ); + payload.resize(offset + n); + if (read_exact(strm, &payload[offset], n) != FrameRead::Ok) { + return FrameRead::Fail; + } } // Unmask if needed diff --git a/test/test.cc b/test/test.cc index 029356fe..2cc821ae 100644 --- a/test/test.cc +++ b/test/test.cc @@ -25543,10 +25543,14 @@ private: }; TEST(WebSocketFrameTest, MultiByteFieldsSplitAcrossReads) { - // A 200-byte payload uses the 16-bit extended length, so the header, the - // length and the mask key are all multi-byte fields here. Each used to be - // read with a single read() call, which fails as soon as one is split. - std::string body(200, 'x'); + // The payload uses the 16-bit extended length, so the header, the length + // and the mask key are all multi-byte fields here. Each used to be read + // with a single read() call, which fails as soon as one is split. The + // payload is also a little over one read chunk, so it is read in two. + std::string body; + for (size_t i = 0; i < CPPHTTPLIB_RECV_BUFSIZ + 3; i++) { + body += static_cast(i % 251); + } const uint8_t mask[4] = {0x0a, 0x0b, 0x0c, 0x0d}; std::string frame; @@ -25566,13 +25570,37 @@ TEST(WebSocketFrameTest, MultiByteFieldsSplitAcrossReads) { std::string payload; bool fin = false; ASSERT_EQ(ws::impl::read_websocket_frame(strm, opcode, payload, fin, - /*expect_masked=*/true, 1024), + /*expect_masked=*/true, body.size()), ws::impl::FrameRead::Ok); EXPECT_TRUE(fin); EXPECT_EQ(opcode, ws::Opcode::Text); EXPECT_EQ(payload, body); } +TEST(WebSocketFrameTest, PayloadGrowsOnlyAsItArrives) { + // The header announces the largest payload allowed, and only 3 bytes of it + // ever arrive. + uint64_t len = CPPHTTPLIB_WEBSOCKET_MAX_PAYLOAD_LENGTH; + std::string frame; + frame += static_cast(0x82); // FIN + Binary + frame += static_cast(127); // 64-bit length follows + for (int shift = 56; shift >= 0; shift -= 8) { + frame += static_cast((len >> shift) & 0xff); + } + frame += "abc"; + + detail::BufferStream strm; + strm.write(frame.data(), frame.size()); + ws::Opcode opcode; + std::string payload; + bool fin; + EXPECT_EQ( + ws::impl::read_websocket_frame(strm, opcode, payload, fin, false, + CPPHTTPLIB_WEBSOCKET_MAX_PAYLOAD_LENGTH), + ws::impl::FrameRead::Fail); + EXPECT_LT(payload.capacity(), 2 * CPPHTTPLIB_RECV_BUFSIZ); +} + TEST(WebSocketTest, QueryStringInHandshake) { Server svr;