260 lines
7.3 KiB
C++
260 lines
7.3 KiB
C++
#ifndef RING_BUFFER_H
|
|
#define RING_BUFFER_H
|
|
|
|
#include <cstdint>
|
|
#include <cstring>
|
|
#include <atomic>
|
|
|
|
// we use a ring buffer with unmasked indices with power-of-2 size
|
|
|
|
class RingBuffer {
|
|
public:
|
|
explicit RingBuffer(size_t capacity = 65536) {
|
|
size_t pow2 = 1;
|
|
while (pow2 < capacity) pow2 <<= 1;
|
|
|
|
capacity_ = pow2;
|
|
mask_ = pow2 - 1;
|
|
buffer_ = new uint8_t[capacity_];
|
|
read_idx_ = 0;
|
|
write_idx_ = 0;
|
|
}
|
|
|
|
~RingBuffer() {
|
|
delete[] buffer_;
|
|
}
|
|
|
|
// no copy and reset buffer to empty state
|
|
RingBuffer(const RingBuffer&) = delete;
|
|
RingBuffer& operator=(const RingBuffer&) = delete;
|
|
|
|
void reset() {
|
|
read_idx_ = 0;
|
|
write_idx_ = 0;
|
|
}
|
|
|
|
size_t writeAvailable() const {
|
|
return capacity_ - size();
|
|
}
|
|
|
|
size_t size() const {
|
|
return write_idx_ - read_idx_;
|
|
}
|
|
|
|
bool empty() const {
|
|
return read_idx_ == write_idx_;
|
|
}
|
|
|
|
bool full() const {
|
|
return size() == capacity_;
|
|
}
|
|
|
|
size_t capacity() const {
|
|
return capacity_;
|
|
}
|
|
|
|
size_t write(const uint8_t* data, size_t len) {
|
|
size_t available = writeAvailable();
|
|
if (len > available) len = available;
|
|
if (len == 0) return 0;
|
|
|
|
size_t write_pos = write_idx_ & mask_;
|
|
size_t first_chunk = capacity_ - write_pos;
|
|
|
|
if (first_chunk >= len) {
|
|
// single contiguous write
|
|
memcpy(buffer_ + write_pos, data, len);
|
|
} else {
|
|
// wrap around - two writes
|
|
memcpy(buffer_ + write_pos, data, first_chunk);
|
|
memcpy(buffer_, data + first_chunk, len - first_chunk);
|
|
}
|
|
|
|
write_idx_ += len;
|
|
return len;
|
|
}
|
|
|
|
// returns bytes available at offset
|
|
size_t peek(uint8_t* dest, size_t len, size_t offset = 0) const {
|
|
size_t available = size();
|
|
if (offset >= available) return 0;
|
|
available -= offset;
|
|
if (len > available) len = available;
|
|
if (len == 0) return 0;
|
|
|
|
size_t read_pos = (read_idx_ + offset) & mask_;
|
|
size_t first_chunk = capacity_ - read_pos;
|
|
|
|
if (first_chunk >= len) {
|
|
memcpy(dest, buffer_ + read_pos, len);
|
|
} else {
|
|
memcpy(dest, buffer_ + read_pos, first_chunk);
|
|
memcpy(dest + first_chunk, buffer_, len - first_chunk);
|
|
}
|
|
|
|
return len;
|
|
}
|
|
|
|
// read the single byte at offset for parsing the header frame
|
|
uint8_t peekByte(size_t offset) const {
|
|
if (offset >= size()) return 0;
|
|
return buffer_[(read_idx_ + offset) & mask_];
|
|
}
|
|
|
|
void consume(size_t len) {
|
|
size_t available = size();
|
|
if (len > available) len = available;
|
|
read_idx_ += len;
|
|
}
|
|
|
|
size_t read(uint8_t* dest, size_t len) {
|
|
size_t bytes_read = peek(dest, len);
|
|
consume(bytes_read);
|
|
return bytes_read;
|
|
}
|
|
|
|
const uint8_t* getReadPtr(size_t* contiguous_len) const {
|
|
size_t read_pos = read_idx_ & mask_;
|
|
size_t available = size();
|
|
size_t to_end = capacity_ - read_pos;
|
|
*contiguous_len = (available < to_end) ? available : to_end;
|
|
return buffer_ + read_pos;
|
|
}
|
|
|
|
uint8_t* getWritePtr(size_t* contiguous_len) {
|
|
size_t write_pos = write_idx_ & mask_;
|
|
size_t available = writeAvailable();
|
|
size_t to_end = capacity_ - write_pos;
|
|
*contiguous_len = (available < to_end) ? available : to_end;
|
|
return buffer_ + write_pos;
|
|
}
|
|
|
|
void commitWrite(size_t len) {
|
|
write_idx_ += len;
|
|
}
|
|
|
|
private:
|
|
uint8_t* buffer_;
|
|
size_t capacity_;
|
|
size_t mask_;
|
|
size_t read_idx_;
|
|
size_t write_idx_;
|
|
};
|
|
|
|
enum class WsFrameResult {
|
|
WS_COMPLETE,
|
|
WS_INCOMPLETE, // needs more data
|
|
WS_ERROR
|
|
};
|
|
struct WsFrameInfo {
|
|
bool fin;
|
|
uint8_t opcode;
|
|
bool masked;
|
|
uint64_t payload_length;
|
|
uint32_t mask_key;
|
|
size_t header_size; // total size 2-14 bytes
|
|
size_t total_size; // header_size + payload_length
|
|
};
|
|
|
|
// frame header from ring buffer returns result and fills frame_info on success
|
|
inline WsFrameResult parseWsFrameHeader(const RingBuffer& buf, WsFrameInfo& info) {
|
|
size_t available = buf.size();
|
|
if (available < 2) return WsFrameResult::WS_INCOMPLETE;
|
|
|
|
uint8_t byte1 = buf.peekByte(0);
|
|
uint8_t byte2 = buf.peekByte(1);
|
|
|
|
info.fin = (byte1 & 0x80) != 0;
|
|
uint8_t rsv = (byte1 >> 4) & 0x07;
|
|
info.opcode = byte1 & 0x0F;
|
|
info.masked = (byte2 & 0x80) != 0;
|
|
info.payload_length = byte2 & 0x7F;
|
|
|
|
// RSV bits must be 0 unless extension negotiated
|
|
if (rsv != 0) {
|
|
return WsFrameResult::WS_ERROR;
|
|
}
|
|
|
|
size_t header_size = 2;
|
|
|
|
if (info.payload_length == 126) {
|
|
if (available < 4) return WsFrameResult::WS_INCOMPLETE;
|
|
info.payload_length = ((uint64_t)buf.peekByte(2) << 8) | buf.peekByte(3);
|
|
header_size = 4;
|
|
} else if (info.payload_length == 127) {
|
|
if (available < 10) return WsFrameResult::WS_INCOMPLETE;
|
|
info.payload_length = 0;
|
|
for (int i = 0; i < 8; i++) {
|
|
info.payload_length = (info.payload_length << 8) | buf.peekByte(2 + i);
|
|
}
|
|
header_size = 10;
|
|
|
|
// max 16MB payload
|
|
if (info.payload_length > 16 * 1024 * 1024) {
|
|
return WsFrameResult::WS_ERROR;
|
|
}
|
|
}
|
|
|
|
// mask key only for client-to-server
|
|
if (info.masked) {
|
|
if (available < header_size + 4) return WsFrameResult::WS_INCOMPLETE;
|
|
info.mask_key = ((uint32_t)buf.peekByte(header_size) << 24) |
|
|
((uint32_t)buf.peekByte(header_size + 1) << 16) |
|
|
((uint32_t)buf.peekByte(header_size + 2) << 8) |
|
|
buf.peekByte(header_size + 3);
|
|
header_size += 4;
|
|
} else {
|
|
info.mask_key = 0;
|
|
}
|
|
|
|
info.header_size = header_size;
|
|
info.total_size = header_size + info.payload_length;
|
|
|
|
if (available < info.total_size) {
|
|
return WsFrameResult::WS_INCOMPLETE;
|
|
}
|
|
|
|
return WsFrameResult::WS_COMPLETE;
|
|
}
|
|
|
|
// extract and unmask payload from ring buffer following uWebSockets patterns: fast path for unmasked, 4-byte XOR for masked
|
|
inline void extractWsPayload(RingBuffer& buf, const WsFrameInfo& info, uint8_t* dest) {
|
|
size_t payload_start = info.header_size;
|
|
size_t len = (size_t)info.payload_length;
|
|
|
|
// skip header bytes first
|
|
buf.consume(payload_start);
|
|
|
|
// unmasked data for server-to-client messages which are never masked
|
|
if (!info.masked) {
|
|
buf.read(dest, len);
|
|
return;
|
|
}
|
|
|
|
// otherwise read then unmask
|
|
buf.read(dest, len);
|
|
|
|
// extracted mask bytes in network byte order
|
|
uint8_t mask[4] = {
|
|
(uint8_t)((info.mask_key >> 24) & 0xFF),
|
|
(uint8_t)((info.mask_key >> 16) & 0xFF),
|
|
(uint8_t)((info.mask_key >> 8) & 0xFF),
|
|
(uint8_t)(info.mask_key & 0xFF)
|
|
};
|
|
|
|
// use 4-byte operations like uWebSockets
|
|
size_t i = 0;
|
|
size_t chunks = len >> 2; // len / 4
|
|
for (size_t c = 0; c < chunks; c++, i += 4) {
|
|
dest[i] ^= mask[0];
|
|
dest[i + 1] ^= mask[1];
|
|
dest[i + 2] ^= mask[2];
|
|
dest[i + 3] ^= mask[3];
|
|
}
|
|
for (; i < len; i++) {
|
|
dest[i] ^= mask[i % 4];
|
|
}
|
|
}
|
|
|
|
#endif
|