Source

`stdlib/websocket/websocket.cpp`

2// Original work; see ACKNOWLEDGMENTS.md for the open-source ideas we build upon.

3// websocket.cpp — a fast, from-scratch RFC 6455 WebSocket client over the

4// cheatah tls 1.3 client + socket. See websocket.hpp for the contract and the

5// performance design (reused read buffer, unmasked-recv fast path, word-XOR

6// masking on send, in-place header parsing, large TLS drains).

8#include "websocket.hpp"

9#include "websocket_lowlevel.hpp" // the C++-only raw handle API this module implements (+ Client uses)

10#include "tls_lowlevel.hpp" // websocket rides tls's low-level (C++-only) session API

12#include <cstdint>

13#include <cstring>

14#include <stdexcept>

16#include "os.hpp"

17#include "socket.hpp"

18#include "tls.hpp"

20namespace cheatah::websocket {

22namespace {

24// A session is heap-allocated and its address IS the handle — so send/recv do a

25// single pointer cast, never a map lookup or lock, on the hot path (cheatah is

26// single-trust; an invalid handle is a caller bug, like a bad pointer in C).

27struct Session {

28 long long fd = -1; // the TCP fd (we own it; close it ourselves)

29 long long tls = -1; // the TLS session riding that fd

30 std::string buf; // read buffer, REUSED across every frame

31 std::size_t pos = 0; // parse offset into buf (consumed prefix is buf[0,pos))

32 bool closed = false;

33 std::uint64_t max_frame = 0; // per-frame payload cap (0 -> the kMaxFramePayload default)

34 std::uint64_t max_message = 0; // reassembled-message cap (0 -> the kMaxMessageBytes default)

35};

37// Hard caps a client MUST impose itself — RFC 6455 sets no upper bound on a frame or a

38// reassembled message, so a malicious server could otherwise (a) overflow `header + len`

39// in size_t math and drive an out-of-bounds unmask, or (b) stream an unbounded body to

40// exhaust memory. A single frame's payload and the reassembled message are each capped;

41// anything larger fails the connection. Control frames are bounded to 125 bytes (§5.5).

42constexpr std::uint64_t kMaxFramePayload = 64ull << 20; // 64 MiB per frame

43constexpr std::uint64_t kMaxMessageBytes = 64ull << 20; // 64 MiB reassembled total

44constexpr std::uint64_t kMaxControlPayload = 125; // RFC 6455 §5.5

46[[nodiscard]] Session* as_session(long long h) {

47 return reinterpret_cast<Session*>(static_cast<std::uintptr_t>(h));

48}

49[[nodiscard]] long long as_handle(Session* s) {

50 return static_cast<long long>(reinterpret_cast<std::uintptr_t>(s));

51}

53// Standard base64 (for the Sec-WebSocket-Key). Tiny inputs (16/raw bytes).

54[[nodiscard]] std::string base64(const std::string& in) {

55 static const char* T = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";

56 std::string out;

57 out.reserve((in.size() + 2) / 3 * 4);

58 std::size_t i = 0;

59 for (; i + 3 <= in.size(); i += 3) {

60 const std::uint32_t n = (static_cast<unsigned char>(in[i]) << 16) |

61 (static_cast<unsigned char>(in[i + 1]) << 8) |

62 static_cast<unsigned char>(in[i + 2]);

63 out.push_back(T[(n >> 18) & 63]);

64 out.push_back(T[(n >> 12) & 63]);

65 out.push_back(T[(n >> 6) & 63]);

66 out.push_back(T[n & 63]);

67 }

68 if (i < in.size()) {

69 std::uint32_t n = static_cast<unsigned char>(in[i]) << 16;

70 if (i + 1 < in.size()) n |= static_cast<unsigned char>(in[i + 1]) << 8;

71 out.push_back(T[(n >> 18) & 63]);

72 out.push_back(T[(n >> 12) & 63]);

73 out.push_back(i + 1 < in.size() ? T[(n >> 6) & 63] : '=');

74 out.push_back('=');

75 }

76 return out;

77}

79void destroy(Session* s) {

80 if (s == nullptr) return;

81 if (s->tls >= 0) tls::close(s->tls);

82 if (s->fd >= 0) socket::close(s->fd);

83 delete s;

84}

86// Ensure the buffer holds at least @p need unconsumed bytes (from pos), draining

87// the TLS stream in big chunks. Compacts the consumed prefix only when needed,

88// so steady-state framing does no front-erase per frame.

89// The transport, chosen per session. A Session carries `tls = -1` when it is PLAINTEXT — a

90// sentinel the struct has always declared and destroy() has always honoured. Everything above

91// the transport (masking, framing, reassembly) is identical either way, so these two are the

92// only places the difference exists.

93//

94// Pointer for the session, const reference for the payload: that is this file's convention

95// (see send_frame), because a session arrives as a handle cast straight to Session* and never

96// through a lookup, while a payload is an ordinary parameter.

97long long sess_send(Session* s, const std::string& data) {

98 return s->tls >= 0 ? tls::send(s->tls, data) : socket::send(s->fd, data);

99}

100

101std::string sess_recv(Session* s, long long n) {

102 return s->tls >= 0 ? tls::recv(s->tls, n) : socket::recv(s->fd, n);

103}

104

105std::string sess_error(Session* s) {

106 return s->tls >= 0 ? tls::last_error() : socket::last_error();

107}

108

109void ensure(Session* s, std::size_t need) {

110 while (s->buf.size() - s->pos < need) {

111 if (s->pos > 0 && (s->pos == s->buf.size() || s->pos >= (1u << 16))) {

112 s->buf.erase(0, s->pos);

113 s->pos = 0;

114 }

115 std::string chunk = sess_recv(s, 1 << 16); // 64 KiB drains

116 if (chunk.empty()) throw std::runtime_error("websocket: connection closed by peer");

117 s->buf.append(chunk);

118 }

119}

120

121// Build one frame header for a client (always-masked) text/control send.

122void put_header(std::string& frame, unsigned char opcode, std::size_t n) {

123 frame.push_back(static_cast<char>(0x80 | opcode)); // FIN | opcode

124 if (n < 126) {

125 frame.push_back(static_cast<char>(0x80 | n)); // MASK | len

126 } else if (n <= 0xFFFF) {

127 frame.push_back(static_cast<char>(0x80 | 126));

128 frame.push_back(static_cast<char>((n >> 8) & 0xFF));

129 frame.push_back(static_cast<char>(n & 0xFF));

130 } else {

131 frame.push_back(static_cast<char>(0x80 | 127));

132 for (int sh = 56; sh >= 0; sh -= 8) frame.push_back(static_cast<char>((n >> sh) & 0xFF));

133 }

134}

135

136// XOR-mask @p n bytes of @p src into @p dst with the 4-byte @p key, eight bytes

137// per step (the key tiled into a 64-bit word) — the fast masking path.

138void mask_into(char* dst, const char* src, std::size_t n, const unsigned char key[4]) {

139 std::uint64_t k64;

140 unsigned char tile[8] = {key[0], key[1], key[2], key[3], key[0], key[1], key[2], key[3]};

141 std::memcpy(&k64, tile, 8);

142 std::size_t i = 0;

143 for (; i + 8 <= n; i += 8) {

144 std::uint64_t w;

145 std::memcpy(&w, src + i, 8);

146 w ^= k64;

147 std::memcpy(dst + i, &w, 8);

148 }

149 for (; i < n; ++i) dst[i] = static_cast<char>(src[i] ^ key[i & 3]);

150}

151

152long long send_frame(Session* s, unsigned char opcode, const std::string& payload) {

153 const std::size_t n = payload.size();

154 std::string frame;

155 frame.reserve(n + 14);

156 put_header(frame, opcode, n);

157 const std::string key = os::urandom(4);

158 const unsigned char k[4] = {static_cast<unsigned char>(key[0]), static_cast<unsigned char>(key[1]),

159 static_cast<unsigned char>(key[2]), static_cast<unsigned char>(key[3])};

160 frame.append(reinterpret_cast<const char*>(k), 4);

161 const std::size_t off = frame.size();

162 frame.resize(off + n);

163 mask_into(frame.data() + off, payload.data(), n, k);

164 if (sess_send(s, frame) < 0)

165 throw std::runtime_error("websocket: send failed: " + sess_error(s));

166 return static_cast<long long>(n);

167}

168

169} // namespace

170

171/// @cond INTERNAL — the C++-only low-level session API (websocket_lowlevel.hpp); cheatah uses the Client guard

172// Whether @p host names this machine. Plaintext is permitted ONLY here: a cleartext WebSocket

173// that cannot leave the loopback interface is a local control plane (Chrome's DevTools port is

174// the motivating one), not a network protocol.

175bool is_loopback(const std::string& host) {

176 return host == "127.0.0.1" || host == "::1" || host == "[::1]" || host == "localhost";

177}

178

179// Send the upgrade request, or destroy the session and report why.

180//

181// Extracted so the failure branch is REACHABLE from a test. It cannot be forced through

182// connect(): it needs the peer to have reset the connection between tcp_connect returning and

183// this very send, which is a race, and a test that only sometimes covers a line is a test that

184// only sometimes passes. testonly::send_upgrade_on_closed_fd drives THIS function with an

185// invalid fd, where socket::send fails with EBADF every time.

186void send_upgrade(Session* s, const std::string& req) {

187 if (sess_send(s, req) < 0) {

188 const std::string err = sess_error(s);

189 destroy(s);

190 throw std::runtime_error("websocket: upgrade request failed: " + err);

191 }

192}

193

194long long connect(const std::string& host, long long port, const std::string& path,

195 const std::string& server_name, bool insecure, const std::string& ca_file,

196 bool secure) {

197 // TLS is the default and the norm. Plaintext must be asked for explicitly AND can only ever

198 // reach this machine — so no amount of configuration turns this into a cleartext socket to

199 // the internet.

200 if (!secure && !is_loopback(host))

201 throw std::runtime_error(

202 "websocket: plaintext ws:// is allowed only to a loopback host (127.0.0.1, ::1, "

203 "localhost); refusing " + host + " — use wss://");

204

205 const long long fd = socket::tcp_connect(host, port);

206 if (fd < 0) throw std::runtime_error("websocket: TCP connect to " + host + " failed");

207 long long tlss = -1;

208 if (secure) {

209 tlss = tls::client_connect(fd, server_name, insecure, ca_file);

210 if (tlss < 0) {

211 socket::close(fd);

212 throw std::runtime_error("websocket: TLS handshake failed: " + tls::last_error());

213 }

214 }

215 Session* s = new Session();

216 s->fd = fd;

217 s->tls = tlss; // -1 => plaintext; sess_send/sess_recv branch on it

218

219 const std::string key = base64(os::urandom(16));

220 const std::string req = "GET " + path + " HTTP/1.1\r\n" + "Host: " + server_name + "\r\n" +

221 "Upgrade: websocket\r\n" + "Connection: Upgrade\r\n" +

222 "Sec-WebSocket-Key: " + key + "\r\n" + "Sec-WebSocket-Version: 13\r\n\r\n";

223 send_upgrade(s, req);

224

225 // Read the response header block (up to and including the blank line). Any

226 // bytes after it are the start of the WebSocket frame stream and stay in buf.

227 std::size_t hdr_end = std::string::npos;

228 for (;;) {

229 const std::string chunk = sess_recv(s, 1 << 12);

230 if (chunk.empty()) {

231 destroy(s);

232 throw std::runtime_error("websocket: connection closed during upgrade");

233 }

234 s->buf.append(chunk);

235 hdr_end = s->buf.find("\r\n\r\n");

236 if (hdr_end != std::string::npos) break;

237 if (s->buf.size() > (1u << 16)) {

238 destroy(s);

239 throw std::runtime_error("websocket: upgrade response too large / not a WebSocket server");

240 }

241 }

242 const std::string status = s->buf.substr(0, s->buf.find("\r\n"));

243 if (status.find(" 101") == std::string::npos) {

244 destroy(s);

245 throw std::runtime_error("websocket: server did not switch protocols: " + status);

246 }

247 s->pos = hdr_end + 4; // frame stream begins after the blank line

248 return as_handle(s);

249}

250

251long long connect_url(const std::string& url, bool insecure, const std::string& ca_file) {

252 // wss:// is the default and behaves exactly as it always has. ws:// exists for a local

253 // control plane and has to be spelled out; connect() then refuses any non-loopback host.

254 bool secure = true;

255 std::string scheme = "wss://";

256 if (url.compare(0, scheme.size(), scheme) != 0) {

257 scheme = "ws://";

258 if (url.compare(0, scheme.size(), scheme) != 0)

259 throw std::runtime_error("websocket: only wss:// and ws:// URLs are supported: " + url);

260 secure = false;

261 }

262 const std::size_t host_start = scheme.size();

263 const std::size_t slash = url.find('/', host_start);

264 const std::string authority =

265 url.substr(host_start, slash == std::string::npos ? std::string::npos : slash - host_start);

266 const std::string path = slash == std::string::npos ? "/" : url.substr(slash);

267 const std::size_t colon = authority.find(':');

268 const std::string host = authority.substr(0, colon);

269 const long long port = colon == std::string::npos

270 ? (secure ? 443 : 80)

271 : static_cast<long long>(std::stol(authority.substr(colon + 1)));

272 return connect(host, port, path, host, insecure, ca_file, secure);

273}

274

275long long send_text(long long session, const std::string& message) {

276 return send_frame(as_session(session), 0x1, message);

277}

278

279std::string recv(long long session) {

280 Session* s = as_session(session);

281 if (s->closed) return std::string();

282 const std::uint64_t max_frame = s->max_frame != 0 ? s->max_frame : kMaxFramePayload;

283 const std::uint64_t max_message = s->max_message != 0 ? s->max_message : kMaxMessageBytes;

284 std::string message; // reassembly buffer for fragmented messages

285 bool fragmenting = false; // mid multi-frame message

286 for (;;) {

287 ensure(s, 2);

288 const unsigned char b0 = static_cast<unsigned char>(s->buf[s->pos]);

289 const unsigned char b1 = static_cast<unsigned char>(s->buf[s->pos + 1]);

290 const bool fin = (b0 & 0x80) != 0;

291 const int opcode = b0 & 0x0F;

292 const bool masked = (b1 & 0x80) != 0; // server frames are unmasked

293 std::uint64_t len = b1 & 0x7F;

294 std::size_t header = 2;

295 if (len == 126) {

296 ensure(s, 4);

297 len = (static_cast<std::uint64_t>(static_cast<unsigned char>(s->buf[s->pos + 2])) << 8) |

298 static_cast<unsigned char>(s->buf[s->pos + 3]);

299 header = 4;

300 } else if (len == 127) {

301 ensure(s, 10);

302 len = 0;

303 for (int i = 0; i < 8; ++i)

304 len = (len << 8) | static_cast<unsigned char>(s->buf[s->pos + 2 + i]);

305 header = 10;

306 }

307 // Validate BEFORE trusting `len` in size math or a copy: reserved bits must be

308 // clear (no extension negotiated), a control frame must be <=125 bytes and not

309 // fragmented, and no frame may exceed the cap. Capping `len` here is what makes

310 // `header + len` below unable to overflow size_t.

311 if ((b0 & 0x70) != 0) throw std::runtime_error("websocket: reserved bits set");

312 if ((opcode & 0x08) != 0 && (len > kMaxControlPayload || !fin))

313 throw std::runtime_error("websocket: invalid control frame");

314 if (len > max_frame) throw std::runtime_error("websocket: frame too large");

315

316 std::size_t mask_off = 0;

317 if (masked) {

318 mask_off = header;

319 header += 4;

320 }

321 ensure(s, header + len);

322 char* payload = s->buf.data() + s->pos + header;

323 if (masked) { // protocol-irregular for a server, but unmask defensively

324 const unsigned char k[4] = {static_cast<unsigned char>(s->buf[s->pos + mask_off]),

325 static_cast<unsigned char>(s->buf[s->pos + mask_off + 1]),

326 static_cast<unsigned char>(s->buf[s->pos + mask_off + 2]),

327 static_cast<unsigned char>(s->buf[s->pos + mask_off + 3])};

328 mask_into(payload, payload, len, k);

329 }

330

331 switch (opcode) {

332 // No braces on the control cases: neither declares anything needing the scope, and

333 // a `}` after a continue/return is unreachable by construction — it shows up

334 // forever as an uncovered line that no test can ever reach.

335 case 0x9: // ping -> pong with the same payload (control, off hot path)

336 send_frame(s, 0xA, std::string(payload, len));

337 s->pos += header + len;

338 continue;

339 case 0xA: // pong -> ignore

340 s->pos += header + len;

341 continue;

342 case 0x8: // close -> echo close, mark done, signal EOF

343 send_frame(s, 0x8, std::string());

344 s->pos += header + len;

345 s->closed = true;

346 return std::string();

347 case 0x1: // text

348 case 0x2: // binary

349 if (fragmenting)

350 throw std::runtime_error("websocket: new data frame during a fragmented message");

351 if (fin) { // single-frame message (the common case) — no reassembly buffer

352 std::string out(payload, len);

353 s->pos += header + len;

354 return out;

355 }

356 message.append(payload, len); // first fragment (frame cap bounds it; total capped below)

357 s->pos += header + len;

358 fragmenting = true;

359 continue;

360 case 0x0: // continuation

361 if (!fragmenting)

362 throw std::runtime_error("websocket: continuation frame with no message in progress");

363 if (message.size() + len > max_message)

364 throw std::runtime_error("websocket: message too large");

365 message.append(payload, len);

366 s->pos += header + len;

367 if (fin) return message;

368 continue;

369 default: // 0x3-0x7, 0xB-0xF are reserved/undefined — fail the connection

370 throw std::runtime_error("websocket: unknown opcode");

371 }

372 }

373}

374

375#ifdef CHEATAH_WEBSOCKET_TESTING

376namespace testonly {

377// A white-box seam (test builds only): build a session pre-loaded with raw frame bytes and

378// overridable caps, so recv()'s frame parser can be driven with crafted/hostile server frames

379// without any TLS/socket. The frames must be self-contained (recv never has to read more).

380// Free it with close(). Compiled ONLY into cheatah_tests; absent from the shipped library.

381long long session_from_bytes(const std::string& frames, std::uint64_t max_frame,

382 std::uint64_t max_message) {

383 Session* s = new Session();

384 s->buf = frames;

385 s->max_frame = max_frame;

386 s->max_message = max_message;

387 return as_handle(s);

388}

389

390// Drive send_upgrade()'s failure branch deterministically: a plaintext session on an invalid

391// descriptor, so socket::send returns EBADF with no peer, no timing and no network.

392void send_upgrade_on_closed_fd() {

393 Session* s = new Session();

394 s->fd = -1; // invalid on purpose

395 s->tls = -1; // plaintext, so sess_send takes the socket:: branch

396 send_upgrade(s, std::string("GET / HTTP/1.1\r\n\r\n"));

397 destroy(s); // unreachable: send_upgrade throws, having destroyed it

398}

399} // namespace testonly

400#endif // CHEATAH_WEBSOCKET_TESTING

401

402long long shutdown(long long session) {

403 // Wake a reader blocked in recv() WITHOUT freeing the session: half-close the

404 // underlying socket so recv returns "". The owner still calls close() after it

405 // has joined the reader. Safe to call from another thread concurrently with the

406 // reader's recv (that is exactly what ::shutdown is for).

407 Session* s = as_session(session);

408 if (s == nullptr || s->tls < 0) return -1;

409 if (s->tls < 0) return socket::shutdown(s->fd);

410 return tls::shutdown(s->tls);

411}

412

413long long close(long long session) {

414 Session* s = as_session(session);

415 if (s == nullptr) return 0;

416 if (!s->closed && s->tls >= 0) {

417 try {

418 send_frame(s, 0x8, std::string());

419 } catch (...) {

420 }

421 }

422 destroy(s);

423 return 0;

424}

425/// @endcond

426

427// ---- owning RAII client ----

428// Each method forwards to the handle-based free function above; the guard adds deterministic

429// close() (close frame + TLS + socket teardown) on scope exit, so a `with` block cannot leak.

430

431Client& Client::operator=(Client&& other) noexcept {

432 if (this != &other) {

433 if (session_ != 0) cheatah::websocket::close(session_);

434 session_ = other.session_;

435 other.session_ = 0;

436 }

437 return *this;

438}

439Client::~Client() {

440 if (session_ != 0) cheatah::websocket::close(session_);

441}

442long long Client::send_text(const std::string& message) {

443 return cheatah::websocket::send_text(session_, message);

444}

445std::string Client::recv() { return cheatah::websocket::recv(session_); }

446long long Client::shutdown() { return cheatah::websocket::shutdown(session_); }

447long long Client::close() {

448 if (session_ == 0) return -1;

449 const long long rc = cheatah::websocket::close(session_);

450 session_ = 0;

451 return rc;

452}

453Client open(const std::string& host, long long port, const std::string& path,

454 const std::string& server_name, bool insecure, const std::string& ca_file,

455 bool secure) {

456 return Client(connect(host, port, path, server_name, insecure, ca_file, secure));

457}

458Client open_url(const std::string& url, bool insecure, const std::string& ca_file) {

459 return Client(connect_url(url, insecure, ca_file));

460}

461

462} // namespace cheatah::websocket