Skip to content

Commit 08a748d

Browse files
tml1024tml1024
committed
Add WebSocket::receiveFrame() that appends to a Poco::Buffer<char>
1 parent 4336528 commit 08a748d

6 files changed

Lines changed: 160 additions & 50 deletions

File tree

‎Net/include/Poco/Net/WebSocket.h‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
#include "Poco/Net/Net.h"
2424
#include "Poco/Net/StreamSocket.h"
2525
#include "Poco/Net/HTTPCredentials.h"
26+
#include "Poco/Buffer.h"
2627

2728

2829
namespace Poco {
@@ -221,6 +222,21 @@ class Net_API WebSocket: public StreamSocket
221222
/// The frame flags and opcode (FrameFlags and FrameOpcodes)
222223
/// is stored in flags.
223224

225+
int receiveFrame(Poco::Buffer<char>& buffer, int& flags);
226+
/// Receives a frame from the socket and stores it
227+
/// after any previous content in buffer.
228+
///
229+
/// Returns the number of bytes received.
230+
/// A return value of 0 means that the peer has
231+
/// shut down or closed the connection.
232+
///
233+
/// Throws a TimeoutException if a receive timeout has
234+
/// been set and nothing is received within that interval.
235+
/// Throws a NetException (or a subclass) in case of other errors.
236+
///
237+
/// The frame flags and opcode (FrameFlags and FrameOpcodes)
238+
/// is stored in flags.
239+
224240
Mode mode() const;
225241
/// Returns WS_SERVER if the WebSocket is a server-side
226242
/// WebSocket, or WS_CLIENT otherwise.

‎Net/include/Poco/Net/WebSocketImpl.h‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121

2222

2323
#include "Poco/Net/StreamSocketImpl.h"
24+
#include "Poco/Buffer.h"
2425
#include "Poco/Random.h"
2526

2627

@@ -43,6 +44,9 @@ class Net_API WebSocketImpl: public StreamSocketImpl
4344
virtual int receiveBytes(void* buffer, int length, int flags);
4445
/// Receives a WebSocket protocol frame.
4546

47+
virtual int receiveBytes(Poco::Buffer<char>& buffer, int flags);
48+
/// Receives a WebSocket protocol frame.
49+
4650
virtual SocketImpl* acceptConnection(SocketAddress& clientAddr);
4751
virtual void connect(const SocketAddress& address);
4852
virtual void connect(const SocketAddress& address, const Poco::Timespan& timeout);
@@ -77,6 +81,9 @@ class Net_API WebSocketImpl: public StreamSocketImpl
7781
MAX_HEADER_LENGTH = 14
7882
};
7983

84+
int receiveHeader(char mask[4], bool& useMask);
85+
int receivePayload(char *buffer, int payloadLength, char mask[4], bool useMask);
86+
8087
int receiveNBytes(void* buffer, int bytes);
8188
virtual ~WebSocketImpl();
8289

‎Net/src/WebSocket.cpp‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
#include "Poco/String.h"
3030
#include "Poco/Random.h"
3131
#include "Poco/StreamCopier.h"
32+
#include "Poco/Buffer.h"
3233
#include <sstream>
3334

3435

@@ -114,6 +115,14 @@ int WebSocket::receiveFrame(void* buffer, int length, int& flags)
114115
}
115116

116117

118+
int WebSocket::receiveFrame(Poco::Buffer<char>& buffer, int& flags)
119+
{
120+
int n = static_cast<WebSocketImpl*>(impl())->receiveBytes(buffer, 0);
121+
flags = static_cast<WebSocketImpl*>(impl())->frameFlags();
122+
return n;
123+
}
124+
125+
117126
WebSocket::Mode WebSocket::mode() const
118127
{
119128
return static_cast<WebSocketImpl*>(impl())->mustMaskPayload() ? WS_CLIENT : WS_SERVER;

‎Net/src/WebSocketImpl.cpp‎

Lines changed: 68 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -104,7 +104,7 @@ int WebSocketImpl::sendBytes(const void* buffer, int length, int flags)
104104
}
105105

106106

107-
int WebSocketImpl::receiveBytes(void* buffer, int length, int)
107+
int WebSocketImpl::receiveHeader(char mask[4], bool& useMask)
108108
{
109109
char header[MAX_HEADER_LENGTH];
110110
int n = receiveNBytes(header, 2);
@@ -114,82 +114,100 @@ int WebSocketImpl::receiveBytes(void* buffer, int length, int)
114114
return n;
115115
}
116116
poco_assert (n == 2);
117+
Poco::UInt8 flags = static_cast<Poco::UInt8>(header[0]);
118+
_frameFlags = flags;
117119
Poco::UInt8 lengthByte = static_cast<Poco::UInt8>(header[1]);
118-
int maskOffset = 0;
119-
if (lengthByte & FRAME_FLAG_MASK) maskOffset += 4;
120+
useMask = ((lengthByte & FRAME_FLAG_MASK) != 0);
121+
int payloadLength;
120122
lengthByte &= 0x7f;
121-
if (lengthByte > 0 || maskOffset > 0)
123+
if (lengthByte == 127)
122124
{
123-
if (lengthByte + 2 + maskOffset < MAX_HEADER_LENGTH)
125+
n = receiveNBytes(header + 2, 8);
126+
if (n <= 0)
124127
{
125-
n = receiveNBytes(header + 2, lengthByte + maskOffset);
128+
_frameFlags = 0;
129+
return n;
126130
}
127-
else
128-
{
129-
n = receiveNBytes(header + 2, MAX_HEADER_LENGTH - 2);
130-
}
131-
if (n <= 0) throw WebSocketException("Incomplete header received", WebSocket::WS_ERR_INCOMPLETE_FRAME);
132-
n += 2;
133-
}
134-
Poco::MemoryInputStream istr(header, n);
135-
Poco::BinaryReader reader(istr, Poco::BinaryReader::NETWORK_BYTE_ORDER);
136-
Poco::UInt8 flags;
137-
char mask[4];
138-
reader >> flags >> lengthByte;
139-
_frameFlags = flags;
140-
int payloadLength = 0;
141-
int payloadOffset = 2;
142-
if ((lengthByte & 0x7f) == 127)
143-
{
131+
Poco::MemoryInputStream istr(header + 2, 8);
132+
Poco::BinaryReader reader(istr, Poco::BinaryReader::NETWORK_BYTE_ORDER);
144133
Poco::UInt64 l;
145134
reader >> l;
146-
if (l > length) throw WebSocketException(Poco::format("Insufficient buffer for payload size %Lu", l), WebSocket::WS_ERR_PAYLOAD_TOO_BIG);
147135
payloadLength = static_cast<int>(l);
148-
payloadOffset += 8;
149-
}
150-
else if ((lengthByte & 0x7f) == 126)
136+
} else if (lengthByte == 126)
151137
{
138+
n = receiveNBytes(header + 2, 2);
139+
if (n <= 0)
140+
{
141+
_frameFlags = 0;
142+
return n;
143+
}
144+
Poco::MemoryInputStream istr(header + 2, 2);
145+
Poco::BinaryReader reader(istr, Poco::BinaryReader::NETWORK_BYTE_ORDER);
152146
Poco::UInt16 l;
153147
reader >> l;
154-
if (l > length) throw WebSocketException(Poco::format("Insufficient buffer for payload size %hu", l), WebSocket::WS_ERR_PAYLOAD_TOO_BIG);
155148
payloadLength = static_cast<int>(l);
156-
payloadOffset += 2;
157149
}
158150
else
159151
{
160-
Poco::UInt8 l = lengthByte & 0x7f;
161-
if (l > length) throw WebSocketException(Poco::format("Insufficient buffer for payload size %u", unsigned(l)), WebSocket::WS_ERR_PAYLOAD_TOO_BIG);
162-
payloadLength = static_cast<int>(l);
163-
}
164-
if (lengthByte & FRAME_FLAG_MASK)
165-
{
166-
reader.readRaw(mask, 4);
167-
payloadOffset += 4;
152+
payloadLength = lengthByte;
168153
}
169-
int received = 0;
170-
if (payloadOffset < n)
171-
{
172-
std::memcpy(buffer, header + payloadOffset, n - payloadOffset);
173-
received = n - payloadOffset;
174-
}
175-
if (received < payloadLength)
154+
155+
if (useMask)
176156
{
177-
n = receiveNBytes(reinterpret_cast<char*>(buffer) + received, payloadLength - received);
178-
if (n <= 0) throw WebSocketException("Incomplete frame received", WebSocket::WS_ERR_INCOMPLETE_FRAME);
179-
received += n;
157+
n = receiveNBytes(mask, 4);
158+
if (n <= 0)
159+
{
160+
_frameFlags = 0;
161+
return n;
162+
}
180163
}
181-
if (lengthByte & FRAME_FLAG_MASK)
164+
165+
return payloadLength;
166+
}
167+
168+
169+
int WebSocketImpl::receivePayload(char *buffer, int payloadLength, char mask[4], bool useMask)
170+
{
171+
int received = receiveNBytes(reinterpret_cast<char*>(buffer), payloadLength);
172+
if (received <= 0) throw WebSocketException("Incomplete frame received", WebSocket::WS_ERR_INCOMPLETE_FRAME);
173+
174+
if (useMask)
182175
{
183-
char* p = reinterpret_cast<char*>(buffer);
184176
for (int i = 0; i < received; i++)
185177
{
186-
p[i] ^= mask[i % 4];
178+
buffer[i] ^= mask[i % 4];
187179
}
188180
}
189181
return received;
190182
}
191183

192184

185+
int WebSocketImpl::receiveBytes(void* buffer, int length, int)
186+
{
187+
char mask[4];
188+
bool useMask;
189+
int payloadLength = receiveHeader(mask, useMask);
190+
if (payloadLength <= 0)
191+
return payloadLength;
192+
if (payloadLength > length)
193+
throw WebSocketException(Poco::format("Insufficient buffer for payload size %hu", payloadLength), WebSocket::WS_ERR_PAYLOAD_TOO_BIG);
194+
return receivePayload(reinterpret_cast<char*>(buffer), payloadLength, mask, useMask);
195+
}
196+
197+
198+
int WebSocketImpl::receiveBytes(Poco::Buffer<char>& buffer, int)
199+
{
200+
char mask[4];
201+
bool useMask;
202+
int payloadLength = receiveHeader(mask, useMask);
203+
if (payloadLength <= 0)
204+
return payloadLength;
205+
int oldSize = buffer.size();
206+
buffer.resize(oldSize + payloadLength);
207+
return receivePayload(buffer.begin() + oldSize, payloadLength, mask, useMask);
208+
}
209+
210+
193211
int WebSocketImpl::receiveNBytes(void* buffer, int bytes)
194212
{
195213
int received = _pStreamSocketImpl->receiveBytes(reinterpret_cast<char*>(buffer), bytes);

‎Net/testsuite/src/WebSocketTest.cpp‎

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -141,6 +141,13 @@ void WebSocketTest::testWebSocket()
141141
assert (n == payload.size());
142142
assert (payload.compare(0, payload.size(), buffer, 0, n) == 0);
143143
assert (flags == WebSocket::FRAME_TEXT);
144+
145+
ws.sendFrame(payload.data(), (int) payload.size());
146+
Poco::Buffer<char> pocobuffer(0);
147+
n = ws.receiveFrame(pocobuffer, flags);
148+
assert (n == payload.size());
149+
assert (payload.compare(0, payload.size(), pocobuffer.begin(), 0, n) == 0);
150+
assert (flags == WebSocket::FRAME_TEXT);
144151
}
145152

146153
for (int i = 125; i < 129; i++)
@@ -151,6 +158,13 @@ void WebSocketTest::testWebSocket()
151158
assert (n == payload.size());
152159
assert (payload.compare(0, payload.size(), buffer, 0, n) == 0);
153160
assert (flags == WebSocket::FRAME_TEXT);
161+
162+
ws.sendFrame(payload.data(), (int) payload.size());
163+
Poco::Buffer<char> pocobuffer(0);
164+
n = ws.receiveFrame(pocobuffer, flags);
165+
assert (n == payload.size());
166+
assert (payload.compare(0, payload.size(), pocobuffer.begin(), 0, n) == 0);
167+
assert (flags == WebSocket::FRAME_TEXT);
154168
}
155169

156170
payload = "Hello, world!";
@@ -210,6 +224,49 @@ void WebSocketTest::testWebSocketLarge()
210224
}
211225

212226

227+
void WebSocketTest::testOneLargeFrame(int msgSize)
228+
{
229+
Poco::Net::ServerSocket ss(0);
230+
Poco::Net::HTTPServer server(new WebSocketRequestHandlerFactory(msgSize), ss, new Poco::Net::HTTPServerParams);
231+
server.start();
232+
233+
Poco::Thread::sleep(200);
234+
235+
HTTPClientSession cs("localhost", ss.address().port());
236+
HTTPRequest request(HTTPRequest::HTTP_GET, "/ws");
237+
HTTPResponse response;
238+
WebSocket ws(cs, request, response);
239+
ws.setSendBufferSize(msgSize);
240+
ws.setReceiveBufferSize(msgSize);
241+
std::string payload(msgSize, 'x');
242+
243+
ws.sendFrame(payload.data(), msgSize);
244+
245+
char buffer[msgSize];
246+
int flags;
247+
int n;
248+
249+
n = ws.receiveFrame(buffer, sizeof(buffer), flags);
250+
assert (n == payload.size());
251+
assert (payload.compare(0, payload.size(), buffer, 0, n) == 0);
252+
253+
ws.sendFrame(payload.data(), msgSize);
254+
255+
Poco::Buffer<char> pocobuffer(0);
256+
257+
n = ws.receiveFrame(pocobuffer, flags);
258+
assert (n == payload.size());
259+
assert (payload.compare(0, payload.size(), pocobuffer.begin(), 0, n) == 0);
260+
}
261+
262+
263+
void WebSocketTest::testWebSocketLargeInOneFrame()
264+
{
265+
testOneLargeFrame(64000);
266+
testOneLargeFrame(70000);
267+
}
268+
269+
213270
void WebSocketTest::setUp()
214271
{
215272
}
@@ -226,6 +283,7 @@ CppUnit::Test* WebSocketTest::suite()
226283

227284
CppUnit_addTest(pSuite, WebSocketTest, testWebSocket);
228285
CppUnit_addTest(pSuite, WebSocketTest, testWebSocketLarge);
286+
CppUnit_addTest(pSuite, WebSocketTest, testWebSocketLargeInOneFrame);
229287

230288
return pSuite;
231289
}

‎Net/testsuite/src/WebSocketTest.h‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,13 +28,15 @@ class WebSocketTest: public CppUnit::TestCase
2828

2929
void testWebSocket();
3030
void testWebSocketLarge();
31+
void testWebSocketLargeInOneFrame();
3132

3233
void setUp();
3334
void tearDown();
3435

3536
static CppUnit::Test* suite();
3637

3738
private:
39+
void testOneLargeFrame(int msgSize);
3840
};
3941

4042

0 commit comments

Comments
 (0)