Skip to content

Commit 010d52c

Browse files
committed
fix continus message bug with presise mask
1 parent 4caa84e commit 010d52c

File tree

1 file changed

+28
-6
lines changed

1 file changed

+28
-6
lines changed

src/WebSocketProtocol.h

Lines changed: 28 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -402,6 +402,28 @@ struct WebSocketProtocol {
402402
}
403403
}
404404

405+
static inline void unmaskPreciseInplace(char *src, size_t length, char *mask) {
406+
uint64_t maskInt;
407+
memcpy(&maskInt, mask, 4);
408+
memcpy(((char *)&maskInt) + 4, mask, 4);
409+
410+
size_t lengthu64 = length/8;
411+
uint64_t* u64I = (uint64_t*)src;
412+
uint64_t* u64O = (uint64_t*)src;
413+
for(size_t m=0; m<lengthu64; m++){
414+
u64O[m] = u64I[m] ^ maskInt;
415+
}
416+
size_t remain = length % 8;
417+
if (remain > 0) {
418+
size_t roffset = length - remain;
419+
char* rI = (src + roffset);
420+
char* rO = (src + roffset);
421+
for (size_t i=0; i<remain; i++) {
422+
rO[i] = rI[i] ^ mask[i % 4];
423+
}
424+
}
425+
}
426+
405427
template <unsigned int MESSAGE_HEADER, typename T>
406428
static inline bool consumeMessage(T payLength, char *&src, unsigned int &length, WebSocketState<isServer> *wState, void *user) {
407429
if (getOpCode(src)) {
@@ -450,10 +472,11 @@ struct WebSocketProtocol {
450472
bool fin = isFin(src);
451473
if constexpr (isServer) {
452474
memcpy(wState->mask, src + MESSAGE_HEADER - 4, 4);
453-
uint64_t mask;
454-
memcpy(&mask, src + MESSAGE_HEADER - 4, 4);
455-
memcpy(((char *)&mask) + 4, src + MESSAGE_HEADER - 4, 4);
456-
unmaskImprecise8<0>(src + MESSAGE_HEADER, mask, length);
475+
uint64_t maskInt;
476+
memcpy(&maskInt, src + MESSAGE_HEADER - 4, 4);
477+
memcpy(((char *)&maskInt) + 4, src + MESSAGE_HEADER - 4, 4);
478+
char* mask = src + MESSAGE_HEADER - 4;
479+
unmaskPrecise8(src + MESSAGE_HEADER, maskInt, mask, length - MESSAGE_HEADER);
457480
rotateMask(4 - (length - MESSAGE_HEADER) % 4, wState->mask);
458481
}
459482
Impl::handleFragment(src + MESSAGE_HEADER, length - MESSAGE_HEADER, wState->remainingBytes, wState->state.opCode[wState->state.opStack], fin, wState, user);
@@ -498,8 +521,7 @@ struct WebSocketProtocol {
498521
if /*constexpr*/ (LIBUS_RECV_BUFFER_LENGTH == length) {
499522
unmaskAll(src, wState->mask);
500523
} else {
501-
// Slow path
502-
unmaskInplace(src, src + ((length >> 2) + 1) * 4, wState->mask);
524+
unmaskPreciseInplace(src, length, wState->mask);
503525
}
504526
}
505527
}

0 commit comments

Comments
 (0)