Line data Source code
1 : #include "websocket.hpp"
2 : #include "../io_context_of.hpp"
3 : #include "../../util/logger.hpp"
4 :
5 : #include <algorithm>
6 : #include <cstring>
7 : #include <random>
8 : #ifdef THINGER_HTTP_SSL_ENABLED
9 : #include <openssl/rand.h>
10 : #endif
11 :
12 : namespace thinger::asio {
13 :
14 : namespace {
15 :
16 : // Bytes of the input buffer; larger payloads are read straight into the reader's buffer
17 : constexpr size_t INPUT_BUFFER_SIZE = 4096;
18 : constexpr size_t DIRECT_READ_SIZE = 1024;
19 :
20 : constexpr uint8_t OPCODE_CONTINUATION = 0x0;
21 : constexpr uint8_t OPCODE_TEXT = 0x1;
22 : constexpr uint8_t OPCODE_BINARY = 0x2;
23 : constexpr uint8_t OPCODE_CLOSE = 0x8;
24 : constexpr uint8_t OPCODE_PING = 0x9;
25 : constexpr uint8_t OPCODE_PONG = 0xA;
26 :
27 : constexpr size_t MAX_CONTROL_PAYLOAD = 125;
28 :
29 : // Mask keys must be unpredictable (RFC 6455, 5.3): taken from the CSPRNG of OpenSSL (a few
30 : // bytes per frame sent by a client), or from std::random_device without it
31 7821 : void random_mask(uint8_t mask[websocket::MASK_SIZE_BYTES]) {
32 : #ifdef THINGER_HTTP_SSL_ENABLED
33 7821 : if (RAND_bytes(mask, websocket::MASK_SIZE_BYTES) == 1) return;
34 : #endif
35 0 : thread_local std::random_device device;
36 0 : uint32_t value = device();
37 0 : std::memcpy(mask, &value, websocket::MASK_SIZE_BYTES);
38 : }
39 :
40 : // The code to send in a close frame: 1000 for one that may not be sent
41 3113 : uint16_t sendable_close_code(uint16_t code) {
42 3113 : if (websocket::valid_close_code(code)) return code;
43 0 : LOG_WARNING("websocket close code {} may not be sent, closing with 1000", code);
44 0 : return websocket::CLOSE_NORMAL;
45 : }
46 :
47 20 : boost::system::error_code protocol_error() {
48 20 : return boost::system::errc::make_error_code(boost::system::errc::protocol_error);
49 : }
50 :
51 : }
52 :
53 : std::atomic<unsigned long> websocket::connections{0};
54 :
55 4216 : bool websocket::valid_close_code(uint16_t code) {
56 4216 : if (code >= 3000 && code <= 4999) return true;
57 4060 : return code >= 1000 && code <= 1014 && code != 1004 && code != 1005 && code != 1006;
58 : }
59 :
60 3919 : websocket::websocket(std::shared_ptr<socket> sock, bool binary, bool server, boost::asio::any_io_executor executor)
61 : : socket("websocket", sock->get_io_context())
62 3919 : , socket_(std::move(sock))
63 3919 : , executor_(executor ? std::move(executor) : boost::asio::any_io_executor(io_context_.get_executor()))
64 3919 : , binary_(binary)
65 3919 : , server_role_(server)
66 7838 : , in_(INPUT_BUFFER_SIZE)
67 3919 : , close_timer_(std::make_shared<boost::asio::steady_timer>(executor_))
68 3919 : , on_close_(executor_)
69 : // A member, called only while this lives: not kept alive while waiting
70 3919 : , keepalive_(executor_, [this] {
71 3571 : if (auto self = weak_from_this().lock()) self->check_keepalive();
72 3571 : })
73 23514 : , last_received_(std::chrono::steady_clock::now()) {
74 3919 : ++connections;
75 3919 : LOG_DEBUG("websocket created");
76 3919 : }
77 :
78 3919 : websocket::~websocket() {
79 3919 : --connections;
80 3919 : LOG_DEBUG("releasing websocket");
81 : // Destroyed open: never told closed
82 3919 : on_close_.release();
83 : // A websocket released open on its executor says goodbye to its peer. Anywhere else no
84 : // I/O is started, and nothing shared with the executor is touched: the socket closes the
85 : // connection once released
86 3919 : if (!closed_ && !writing_ && running_in_this_thread(executor_)) {
87 240 : close(CLOSE_GOING_AWAY);
88 : }
89 3919 : }
90 :
91 336 : void websocket::set_read_ahead(std::string data) {
92 336 : if (data.size() > in_.size()) in_.resize(data.size());
93 336 : std::memcpy(in_.data(), data.data(), data.size());
94 336 : in_begin_ = 0;
95 336 : in_end_ = data.size();
96 336 : }
97 :
98 : // ---------------------------------------------------------------------------
99 : // Frames
100 : // ---------------------------------------------------------------------------
101 :
102 105831 : bool websocket::parse_header(const uint8_t* data, size_t size, frame_header& header) {
103 105831 : header.size = 2;
104 105831 : if (size < 2) return false;
105 61396 : uint8_t length = data[1] & 0x7F;
106 61396 : header.masked = data[1] & 0x80;
107 61396 : header.size += (length == 126 ? 2 : length == 127 ? 8 : 0) + (header.masked ? MASK_SIZE_BYTES : 0);
108 61396 : if (size < header.size) return false;
109 :
110 61360 : header.fin = data[0] & 0x80;
111 61360 : header.rsv = data[0] & 0x70;
112 61360 : header.opcode = data[0] & 0x0F;
113 61360 : size_t offset = 2;
114 61360 : if (length == 126) {
115 18462 : header.payload_size = (uint64_t(data[2]) << 8) | data[3];
116 18462 : offset += 2;
117 42898 : } else if (length == 127) {
118 2325 : header.payload_size = 0;
119 20925 : for (int i = 0; i < 8; ++i) header.payload_size = (header.payload_size << 8) | data[2 + i];
120 2325 : offset += 8;
121 : } else {
122 40573 : header.payload_size = length;
123 : }
124 61360 : if (header.masked) std::memcpy(header.mask, data + offset, MASK_SIZE_BYTES);
125 61360 : return true;
126 : }
127 :
128 52775 : size_t websocket::build_header(uint8_t opcode, uint64_t payload_size, uint8_t header[14],
129 : uint8_t mask[MASK_SIZE_BYTES]) const {
130 52775 : size_t size = 2;
131 52775 : header[0] = 0x80 | opcode;
132 52775 : if (payload_size <= 125) {
133 18962 : header[1] = static_cast<uint8_t>(payload_size);
134 33813 : } else if (payload_size <= 65535) {
135 25287 : header[1] = 126;
136 25287 : header[2] = (payload_size >> 8) & 0xFF;
137 25287 : header[3] = payload_size & 0xFF;
138 25287 : size += 2;
139 : } else {
140 8526 : header[1] = 127;
141 76734 : for (int i = 0; i < 8; ++i) header[2 + i] = (payload_size >> ((7 - i) * 8)) & 0xFF;
142 8526 : size += 8;
143 : }
144 : // Clients mask what they send (RFC 6455, 5.3)
145 52775 : if (!server_role_) {
146 7821 : header[1] |= 0x80;
147 7821 : random_mask(mask);
148 7821 : std::memcpy(header + size, mask, MASK_SIZE_BYTES);
149 7821 : size += MASK_SIZE_BYTES;
150 : }
151 52775 : return size;
152 : }
153 :
154 124382 : void websocket::apply_mask(uint8_t* data, size_t size, const uint8_t mask[MASK_SIZE_BYTES], size_t offset) {
155 : // The key from `offset`, repeated over a word
156 : uint8_t key[8];
157 1119438 : for (size_t i = 0; i < sizeof(key); ++i) key[i] = mask[(offset + i) % MASK_SIZE_BYTES];
158 : uint64_t word_key;
159 124382 : std::memcpy(&word_key, key, sizeof(word_key));
160 124382 : size_t i = 0;
161 377455080 : for (; i + sizeof(word_key) <= size; i += sizeof(word_key)) {
162 : uint64_t word;
163 377330698 : std::memcpy(&word, data + i, sizeof(word));
164 377330698 : word ^= word_key;
165 377330698 : std::memcpy(data + i, &word, sizeof(word));
166 : }
167 186614 : for (; i < size; ++i) data[i] ^= key[i % sizeof(key)];
168 124382 : }
169 :
170 115915 : void websocket::unmask(uint8_t buffer[], size_t size) {
171 115915 : apply_mask(buffer, size, mask_, mask_offset_);
172 115915 : mask_offset_ = (mask_offset_ + size) % MASK_SIZE_BYTES;
173 115915 : }
174 :
175 : // ---------------------------------------------------------------------------
176 : // Read side
177 : // ---------------------------------------------------------------------------
178 :
179 109173 : awaitable<boost::system::error_code> websocket::fill(size_t size) {
180 : while (buffered() < size) {
181 : if (buffered() == 0) {
182 : in_begin_ = in_end_ = 0;
183 : } else if (in_.size() - in_begin_ < size) {
184 : std::memmove(in_.data(), in_.data() + in_begin_, buffered());
185 : in_end_ -= in_begin_;
186 : in_begin_ = 0;
187 : }
188 : if (read_cancelled_) co_return boost::asio::error::operation_aborted;
189 : auto [ec, bytes] = co_await socket_->read_some(in_.data() + in_end_, in_.size() - in_end_);
190 : // closed meanwhile: the socket is not this websocket's anymore
191 : if (closed_) co_return boost::asio::error::operation_aborted;
192 : if (bytes == 0 && !ec) ec = boost::asio::error::eof; // a socket breaking its guarantees
193 : if (ec) {
194 : check_lost(ec);
195 : co_return ec;
196 : }
197 : in_end_ += bytes;
198 : last_received_ = std::chrono::steady_clock::now();
199 : }
200 : co_return boost::system::error_code{};
201 218346 : }
202 :
203 135985 : awaitable<io_result> websocket::read_payload(uint8_t buffer[], size_t max_size) {
204 : auto wanted = static_cast<size_t>(std::min<uint64_t>(max_size, frame_remaining_));
205 : size_t bytes;
206 : if (buffered() == 0 && wanted >= DIRECT_READ_SIZE) {
207 : // Large reads go straight to the reader's buffer, never past this frame
208 : if (read_cancelled_) co_return io_result{boost::asio::error::operation_aborted, 0};
209 : auto [ec, read] = co_await socket_->read_some(buffer, wanted);
210 : if (closed_) co_return io_result{boost::asio::error::operation_aborted, 0};
211 : if (read == 0 && !ec) ec = boost::asio::error::eof; // a socket breaking its guarantees
212 : if (ec) {
213 : check_lost(ec);
214 : co_return io_result{ec, 0};
215 : }
216 : last_received_ = std::chrono::steady_clock::now();
217 : bytes = read;
218 : } else {
219 : if (auto ec = co_await fill(1)) co_return io_result{ec, 0};
220 : bytes = std::min(wanted, buffered());
221 : std::memcpy(buffer, in_.data() + in_begin_, bytes);
222 : in_begin_ += bytes;
223 : }
224 : if (frame_masked_) unmask(buffer, bytes);
225 : frame_remaining_ -= bytes;
226 : if (frame_remaining_ == 0 && frame_fin_) new_message_ = true;
227 : co_return io_result{boost::system::error_code{}, bytes};
228 271970 : }
229 :
230 20 : boost::system::error_code websocket::fail(const char* reason) {
231 20 : LOG_DEBUG("websocket protocol error: {}", reason);
232 20 : close(CLOSE_PROTOCOL_ERROR);
233 20 : return protocol_error();
234 : }
235 :
236 4093 : websocket_close websocket::close_status() const {
237 4093 : return websocket_close{get_close_code(), close_reason_, closed_by_peer_};
238 : }
239 :
240 52086 : void websocket::check_lost(const boost::system::error_code& ec) {
241 52086 : if (!ec || closed_) return;
242 1435 : if (ec == boost::asio::error::operation_aborted && socket_->is_open()) return;
243 1398 : on_close_.notify(close_status());
244 : }
245 :
246 102 : void websocket::on_close(sync_callback<void(const websocket_close&)> callback) {
247 102 : on_close_.set(std::move(callback));
248 102 : }
249 :
250 3249 : awaitable<boost::system::error_code> websocket::handle_control(const frame_header& header) {
251 : uint8_t payload[MAX_CONTROL_PAYLOAD];
252 : size_t size = header.payload_size;
253 : std::memcpy(payload, in_.data() + in_begin_ + header.size, size);
254 : if (header.masked) apply_mask(payload, size, header.mask, 0);
255 : in_begin_ += header.size + size;
256 :
257 : switch (header.opcode) {
258 : case OPCODE_PING: {
259 : LOG_DEBUG("received ping frame");
260 : if (close_sent_) co_return boost::system::error_code{};
261 : if (weak_from_this().expired()) {
262 : // Not owned by a std::shared_ptr: answered before reading on
263 : co_await send_frame(OPCODE_PONG, {boost::asio::buffer(payload, size)}, true);
264 : if (closed_) co_return boost::asio::error::operation_aborted;
265 : co_return boost::system::error_code{};
266 : }
267 : // Answered as soon as no frame is being written, while reading goes on (a pong may
268 : // answer only the last ping)
269 : pending_pong_.emplace(reinterpret_cast<const char*>(payload), size);
270 : queue_control();
271 : co_return boost::system::error_code{};
272 : }
273 : case OPCODE_PONG:
274 : LOG_DEBUG("received pong frame");
275 : co_return boost::system::error_code{};
276 : case OPCODE_CLOSE: {
277 : LOG_DEBUG("received close frame");
278 : if (size == 1) co_return fail("invalid close frame payload");
279 : close_received_ = true;
280 : closed_by_peer_ = !close_sent_;
281 : uint16_t code = size >= 2 ? uint16_t((payload[0] << 8) | payload[1]) : CLOSE_NORMAL;
282 : if (!valid_close_code(code)) code = CLOSE_PROTOCOL_ERROR;
283 : received_close_code_ = size >= 2 ? code : CLOSE_NO_STATUS;
284 : if (size > 2) close_reason_.assign(reinterpret_cast<const char*>(payload) + 2, size - 2);
285 : // Answered (unless this end closed first), and the connection closed: the server
286 : // closes it first, and a client does not need to wait for it
287 : close(code, false);
288 : co_return boost::asio::error::eof;
289 : }
290 : default:
291 : co_return fail("unknown control frame opcode");
292 : }
293 6498 : }
294 :
295 138393 : awaitable<io_result> websocket::read_frame_data(uint8_t buffer[], size_t max_size, bool skip_empty) {
296 : for (;;) {
297 : if (close_received_) co_return io_result{boost::asio::error::eof, 0};
298 : if (closed_) co_return io_result{boost::asio::error::bad_descriptor, 0};
299 :
300 : // Payload of the data frame being received
301 : if (frame_remaining_ > 0) {
302 : if (max_size == 0) co_return io_result{boost::system::error_code{}, 0};
303 : co_return co_await read_payload(buffer, max_size);
304 : }
305 :
306 : // Next frame header
307 : frame_header header;
308 : while (!parse_header(in_.data() + in_begin_, buffered(), header)) {
309 : if (auto ec = co_await fill(header.size)) co_return io_result{ec, 0};
310 : }
311 :
312 : if (header.rsv) co_return io_result{fail("reserved bits set without an extension"), 0};
313 : if (server_role_ && !header.masked) co_return io_result{fail("unmasked frame from a client"), 0};
314 : if (!server_role_ && header.masked) co_return io_result{fail("masked frame from a server"), 0};
315 : if (header.payload_size >> 63) co_return io_result{fail("invalid payload length"), 0};
316 :
317 : // Control frames, which may come between the frames of a message
318 : if (header.opcode >= 0x8) {
319 : if (!header.fin) co_return io_result{fail("fragmented control frame"), 0};
320 : if (header.payload_size > MAX_CONTROL_PAYLOAD) co_return io_result{fail("control frame too large"), 0};
321 : if (auto ec = co_await fill(header.size + header.payload_size)) co_return io_result{ec, 0};
322 : if (auto ec = co_await handle_control(header)) co_return io_result{ec, 0};
323 : continue;
324 : }
325 :
326 : // Data frames: a text or binary frame starts a message, continuation frames go on
327 : switch (header.opcode) {
328 : case OPCODE_CONTINUATION:
329 : if (new_message_) co_return io_result{fail("continuation frame without a message"), 0};
330 : break;
331 : case OPCODE_TEXT:
332 : case OPCODE_BINARY:
333 : if (!new_message_) co_return io_result{fail("new message before the previous one ended"), 0};
334 : message_opcode_ = header.opcode;
335 : break;
336 : default:
337 : co_return io_result{fail("unknown data frame opcode"), 0};
338 : }
339 :
340 : in_begin_ += header.size;
341 : frame_remaining_ = header.payload_size;
342 : frame_fin_ = header.fin;
343 : frame_masked_ = header.masked;
344 : std::memcpy(mask_, header.mask, MASK_SIZE_BYTES);
345 : mask_offset_ = 0;
346 : new_message_ = false;
347 :
348 : if (frame_remaining_ == 0) {
349 : if (frame_fin_) new_message_ = true;
350 : if (!skip_empty) co_return io_result{boost::system::error_code{}, 0};
351 : }
352 : }
353 276786 : }
354 :
355 68370 : awaitable<io_result> websocket::read_some(uint8_t buffer[], size_t max_size) {
356 : if (max_size == 0) co_return io_result{closed_ ? boost::asio::error::bad_descriptor : boost::system::error_code{}, 0};
357 : reading_ = true;
358 : auto result = co_await read_frame_data(buffer, max_size, true);
359 : reading_ = false;
360 : read_cancelled_ = false;
361 : co_return result;
362 136740 : }
363 :
364 70023 : awaitable<io_result> websocket::read_message_some(uint8_t buffer[], size_t max_size) {
365 : reading_ = true;
366 : auto result = co_await read_frame_data(buffer, max_size, false);
367 : reading_ = false;
368 : read_cancelled_ = false;
369 : co_return result;
370 140046 : }
371 :
372 : // ---------------------------------------------------------------------------
373 : // Write side
374 : // ---------------------------------------------------------------------------
375 :
376 48057 : awaitable<bool> websocket::acquire_write(bool control) {
377 : if (closed_) co_return false;
378 : if (!writing_) {
379 : writing_ = true;
380 : co_return true;
381 : }
382 : // Wait for the turn: control frames go before the data waiting, after other control frames
383 : auto waiter = std::make_shared<write_waiter>(executor_);
384 : if (control) {
385 : auto position = std::find_if(write_waiters_.begin(), write_waiters_.end(),
386 0 : [](const auto& other) { return !other->control; });
387 : waiter->control = true;
388 : write_waiters_.insert(position, waiter);
389 : } else {
390 : write_waiters_.push_back(waiter);
391 : }
392 : co_await waiter->turn.wait();
393 : if (waiter->granted) co_return true;
394 : // closed, cancelled, or the wait itself cancelled
395 : auto it = std::find(write_waiters_.begin(), write_waiters_.end(), waiter);
396 : if (it != write_waiters_.end()) write_waiters_.erase(it);
397 : co_return false;
398 96114 : }
399 :
400 53238 : void websocket::release_write() {
401 53238 : if (closed_) {
402 : // the closing handshake waited for the write in progress
403 138 : writing_ = false;
404 138 : if (linger_pending_) start_linger();
405 138 : return;
406 : }
407 53100 : if (pending_pong_ || pending_ping_) {
408 : // the turn goes to the control frames queued: writing_ stays set
409 2677 : if (auto self = weak_from_this().lock()) {
410 5354 : co_spawn(executor_, [self]() -> awaitable<void> { co_await self->flush_control(); }, detached);
411 2677 : return;
412 2677 : }
413 0 : pending_pong_.reset();
414 0 : pending_ping_ = false;
415 : }
416 50423 : if (!write_waiters_.empty()) {
417 : // the turn passes to the next writer: writing_ stays set
418 796 : auto next = std::move(write_waiters_.front());
419 796 : write_waiters_.pop_front();
420 796 : next->granted = true;
421 796 : next->turn.set();
422 796 : return;
423 796 : }
424 49627 : writing_ = false;
425 : }
426 :
427 2753 : void websocket::wake_writers() {
428 2753 : auto waiters = std::move(write_waiters_);
429 2753 : write_waiters_.clear();
430 2825 : for (auto& waiter : waiters) waiter->turn.set();
431 2753 : }
432 :
433 2738 : void websocket::queue_control() {
434 2738 : if (writing_) return; // sent when the turn is released
435 2598 : writing_ = true;
436 2598 : release_write();
437 : }
438 :
439 2677 : awaitable<void> websocket::flush_control() {
440 : while (!closed_ && (pending_pong_ || pending_ping_)) {
441 : boost::system::error_code ec;
442 : if (pending_pong_) {
443 : auto payload = std::move(*pending_pong_);
444 : pending_pong_.reset();
445 : ec = co_await write_frame(OPCODE_PONG, {boost::asio::buffer(payload)});
446 : } else {
447 : pending_ping_ = false;
448 : ping_in_flight_ = true;
449 : ec = co_await write_frame(OPCODE_PING, {});
450 : ping_in_flight_ = false;
451 : // a ping that cannot be sent leaves nothing to keep alive (cancel() aside)
452 : if (ec && ec != boost::asio::error::operation_aborted && !closed_) {
453 : LOG_DEBUG("websocket keepalive ping failed: {}", ec.message());
454 : close(CLOSE_GOING_AWAY);
455 : }
456 : }
457 : if (ec) break;
458 : }
459 : release_write();
460 5354 : }
461 :
462 50675 : awaitable<boost::system::error_code> websocket::write_frame(uint8_t opcode,
463 : const std::vector<boost::asio::const_buffer>& payload,
464 : const sync_callback<void()>* progress) {
465 : // The rest of an interrupted frame goes first
466 : if (!tail_.empty()) {
467 : auto [ec, bytes] = co_await socket_->write(tail_.data(), tail_.size());
468 : tail_.erase(tail_.begin(), tail_.begin() + static_cast<std::ptrdiff_t>(std::min(bytes, tail_.size())));
469 : if (ec) {
470 : if (ec != boost::asio::error::operation_aborted) write_error_ = ec;
471 : check_lost(ec);
472 : co_return ec;
473 : }
474 : }
475 : if (opcode == OPCODE_CLOSE) {
476 : close_sent_ = true;
477 : close_unwritten_ = true;
478 : }
479 :
480 : size_t size = boost::asio::buffer_size(payload);
481 : uint8_t header[14];
482 : uint8_t mask[MASK_SIZE_BYTES];
483 : size_t header_size = build_header(opcode, size, header, mask);
484 :
485 : std::vector<boost::asio::const_buffer> output;
486 : output.reserve(payload.size() + 1);
487 : output.emplace_back(header, header_size);
488 : std::vector<uint8_t> masked;
489 : if (server_role_) {
490 : output.insert(output.end(), payload.begin(), payload.end());
491 : } else if (size > 0) {
492 : masked.resize(size);
493 : boost::asio::buffer_copy(boost::asio::buffer(masked), payload);
494 : apply_mask(masked.data(), size, mask, 0);
495 : output.emplace_back(masked.data(), masked.size());
496 : }
497 :
498 : bool sliced = progress && header_size + size > WRITE_SLICE_SIZE;
499 : auto [ec, bytes] = sliced ? co_await socket_->write_slices(output, *progress) : co_await socket_->write(output);
500 : check_lost(ec);
501 : if (ec) {
502 : if (ec != boost::asio::error::operation_aborted) {
503 : write_error_ = ec;
504 : } else if (bytes > 0) {
505 : // Interrupted in the middle of the frame: its rest is sent before anything else
506 : // (write_slices() leaves it in `output`)
507 : if (!sliced) consume(output, bytes);
508 : tail_.resize(boost::asio::buffer_size(output));
509 : boost::asio::buffer_copy(boost::asio::buffer(tail_), output);
510 : }
511 : } else if (opcode == OPCODE_CLOSE) {
512 : close_unwritten_ = false;
513 : }
514 : co_return ec;
515 101350 : }
516 :
517 48057 : awaitable<io_result> websocket::send_frame(uint8_t opcode, const std::vector<boost::asio::const_buffer>& payload,
518 : bool control, const sync_callback<void()>* progress) {
519 : if (!co_await acquire_write(control)) {
520 : co_return io_result{closed_ ? boost::asio::error::bad_descriptor : boost::asio::error::operation_aborted, 0};
521 : }
522 :
523 : boost::system::error_code error;
524 : if (closed_) {
525 : error = boost::asio::error::bad_descriptor;
526 : } else if (write_error_) {
527 : error = write_error_;
528 : } else if (close_sent_ || close_received_) {
529 : // nothing is sent after a close frame
530 : error = boost::asio::error::broken_pipe;
531 : }
532 : if (error) {
533 : release_write();
534 : co_return io_result{error, 0};
535 : }
536 :
537 : auto ec = co_await write_frame(opcode, payload, progress);
538 : release_write();
539 : co_return io_result{ec, ec ? 0 : boost::asio::buffer_size(payload)};
540 96114 : }
541 :
542 511 : awaitable<boost::system::error_code> websocket::send_close() {
543 : LOG_DEBUG("sending close frame: {}", close_code_);
544 : auto payload = close_payload();
545 : auto [ec, bytes] = co_await send_frame(OPCODE_CLOSE, {boost::asio::buffer(payload)}, true);
546 : co_return ec;
547 1022 : }
548 :
549 2611 : std::vector<uint8_t> websocket::close_payload() const {
550 5222 : std::vector<uint8_t> payload{uint8_t(close_code_ >> 8), uint8_t(close_code_ & 0xFF)};
551 2611 : payload.insert(payload.end(), close_reason_sent_.begin(), close_reason_sent_.end());
552 2611 : return payload;
553 0 : }
554 :
555 0 : awaitable<void> websocket::send_ping(uint8_t buffer[], size_t size) {
556 : co_await send_frame(OPCODE_PING, {boost::asio::buffer(buffer, std::min(size, MAX_CONTROL_PAYLOAD))}, true);
557 0 : }
558 :
559 0 : awaitable<void> websocket::send_pong(uint8_t buffer[], size_t size) {
560 : co_await send_frame(OPCODE_PONG, {boost::asio::buffer(buffer, std::min(size, MAX_CONTROL_PAYLOAD))}, true);
561 0 : }
562 :
563 13707 : awaitable<io_result> websocket::write(const uint8_t buffer[], size_t size) {
564 : co_return co_await write(std::vector<boost::asio::const_buffer>{boost::asio::buffer(buffer, size)});
565 27414 : }
566 :
567 6335 : awaitable<io_result> websocket::write(std::string_view str) {
568 : co_return co_await write(std::vector<boost::asio::const_buffer>{boost::asio::buffer(str.data(), str.size())});
569 12670 : }
570 :
571 20123 : awaitable<io_result> websocket::write(const std::vector<boost::asio::const_buffer>& buffers) {
572 : // As a stream, writing nothing sends nothing (not an empty message)
573 : if (boost::asio::buffer_size(buffers) == 0) {
574 : co_return io_result{closed_ ? boost::asio::error::bad_descriptor : boost::system::error_code{}, 0};
575 : }
576 : co_return co_await send_frame(binary_ ? OPCODE_BINARY : OPCODE_TEXT, buffers, false);
577 40246 : }
578 :
579 27234 : awaitable<io_result> websocket::write_message(std::string_view data) {
580 : // An empty message is a valid message (RFC 6455, 5.6): a frame without payload
581 : co_return co_await send_frame(binary_ ? OPCODE_BINARY : OPCODE_TEXT,
582 : {boost::asio::buffer(data.data(), data.size())}, false);
583 54468 : }
584 :
585 198 : awaitable<io_result> websocket::write_message(std::string_view data, const sync_callback<void()>& progress) {
586 : co_return co_await send_frame(binary_ ? OPCODE_BINARY : OPCODE_TEXT,
587 : {boost::asio::buffer(data.data(), data.size())}, false, &progress);
588 396 : }
589 :
590 : // ---------------------------------------------------------------------------
591 : // Closing
592 : // ---------------------------------------------------------------------------
593 :
594 1722 : void websocket::close() {
595 1722 : close(CLOSE_NORMAL);
596 1722 : }
597 :
598 3384 : void websocket::close(uint16_t code, bool expect_reply) {
599 3384 : if (closed_) return;
600 2599 : code = sendable_close_code(code);
601 2599 : mark_closed();
602 :
603 : // The closing handshake runs on the executor: called anywhere else (e.g. by a synchronous
604 : // client between its calls, while nothing runs its io_context) the connection is closed now
605 2599 : if (close_sent_ || write_error_ || !socket_->is_open() || !running_in_this_thread(executor_)) {
606 498 : socket_->close();
607 498 : on_close_.notify(close_status());
608 498 : return;
609 : }
610 2101 : close_sent_ = true;
611 : // (a close frame asked for and waiting for its turn is the one sent)
612 2101 : if (!close_requested_) close_code_ = code;
613 2101 : linger_expect_reply_ = expect_reply;
614 :
615 : // Pending operations end now. A write in progress leaves the rest of its frame, sent
616 : // before the close frame once it returns the turn
617 2101 : socket_->cancel();
618 2101 : if (writing_) {
619 64 : linger_pending_ = true;
620 : } else {
621 2037 : start_linger();
622 : }
623 2101 : on_close_.notify(close_status());
624 : }
625 :
626 2695 : void websocket::mark_closed() {
627 2695 : closed_ = true;
628 2695 : keepalive_.cancel();
629 2695 : close_timer_->cancel();
630 2695 : wake_writers();
631 2695 : pending_pong_.reset();
632 2695 : pending_ping_ = false;
633 2695 : }
634 :
635 511 : void websocket::start_close_deadline() {
636 511 : close_timer_->expires_after(CLOSE_TIMEOUT);
637 : // (a websocket not owned by a std::shared_ptr is closed by close_graceful(), once the socket
638 : // ends what it waits for)
639 511 : close_timer_->async_wait([weak = weak_from_this(), socket = socket_](const boost::system::error_code& ec) {
640 499 : if (ec) return;
641 96 : if (auto self = weak.lock()) self->close_expired();
642 96 : else socket->close();
643 : });
644 511 : }
645 :
646 96 : void websocket::close_expired() {
647 96 : if (closed_) return;
648 96 : LOG_DEBUG("websocket closing handshake timed out");
649 96 : mark_closed();
650 96 : socket_->close();
651 96 : on_close_.notify(close_status());
652 : }
653 :
654 2100 : void websocket::start_linger() {
655 2100 : linger_pending_ = false;
656 2100 : std::vector<uint8_t> output = std::move(tail_);
657 2100 : tail_.clear();
658 : uint8_t header[14];
659 : uint8_t mask[MASK_SIZE_BYTES];
660 2100 : auto payload = close_payload();
661 2100 : size_t header_size = build_header(OPCODE_CLOSE, payload.size(), header, mask);
662 2100 : output.insert(output.end(), header, header + header_size);
663 2100 : if (!server_role_) apply_mask(payload.data(), payload.size(), mask, 0);
664 2100 : output.insert(output.end(), payload.begin(), payload.end());
665 :
666 : // The closing handshake goes on without this websocket, on its socket, with what was
667 : // received but not read yet
668 2100 : std::string input(reinterpret_cast<const char*>(in_.data() + in_begin_), buffered());
669 : // (asked for by close_graceful(), the time it has left)
670 2100 : auto deadline = close_requested_ ? close_timer_->expiry() : std::chrono::steady_clock::now() + CLOSE_TIMEOUT;
671 2100 : co_spawn(executor_, linger(socket_, std::move(output), std::move(input), frame_remaining_, linger_expect_reply_, deadline),
672 : detached);
673 2100 : }
674 :
675 2100 : awaitable<void> websocket::linger(std::shared_ptr<socket> socket, std::vector<uint8_t> output,
676 : std::string input, uint64_t skip, bool expect_reply,
677 : std::chrono::steady_clock::time_point deadline) {
678 : // Close the socket at the deadline, whatever is pending on it: on the executor of the
679 : // websocket, as the operations it ends
680 : auto timer = std::make_shared<boost::asio::steady_timer>(co_await boost::asio::this_coro::executor, deadline);
681 1143 : timer->async_wait([socket](const boost::system::error_code& ec) {
682 1143 : if (!ec) socket->close();
683 1143 : });
684 :
685 : auto [write_ec, written] = co_await socket->write(output.data(), output.size());
686 : if (!write_ec && expect_reply) {
687 : // Wait for the close frame of the peer, skipping anything else
688 : uint8_t chunk[1024];
689 : for (;;) {
690 : auto skipped = std::min<uint64_t>(skip, input.size());
691 : input.erase(0, skipped);
692 : skip -= skipped;
693 : if (skip == 0) {
694 : frame_header header;
695 : if (parse_header(reinterpret_cast<const uint8_t*>(input.data()), input.size(), header)) {
696 : if (header.opcode == OPCODE_CLOSE) break;
697 : input.erase(0, header.size);
698 : skip = header.payload_size;
699 : continue;
700 : }
701 : }
702 : auto [ec, bytes] = co_await socket->read_some(chunk, sizeof(chunk));
703 : if (ec || bytes == 0) break;
704 : input.append(reinterpret_cast<const char*>(chunk), bytes);
705 : }
706 : }
707 : timer->cancel();
708 : socket->close();
709 4200 : }
710 :
711 532 : awaitable<void> websocket::close_graceful(uint16_t code, std::string reason) {
712 : if (closed_) co_return;
713 : code = sendable_close_code(code);
714 : // Nothing else is sent after the close frame
715 : keepalive_.cancel();
716 : pending_ping_ = false;
717 : if (!close_requested_) {
718 : if (reason.size() > MAX_CLOSE_REASON_SIZE) {
719 : LOG_WARNING("websocket close reason of {} bytes does not fit in a close frame, closing without it", reason.size());
720 : reason.clear();
721 : }
722 : close_requested_ = true;
723 : close_code_ = code;
724 : close_reason_sent_ = std::move(reason);
725 : // One deadline for the whole handshake, its turn to write included
726 : start_close_deadline();
727 : if (co_await send_close()) {
728 : close(code);
729 : co_return;
730 : }
731 : }
732 : if (closed_) co_return;
733 : if (!close_received_) {
734 : if (reading_) {
735 : // The reader gets the close frame of the peer, and closes (which ends this wait)
736 : co_await close_timer_->async_wait(use_nothrow_awaitable);
737 : } else {
738 : // Read until the close frame of the peer, discarding anything else
739 178 : auto discard = [this]() -> awaitable<void> {
740 : uint8_t buffer[512];
741 : for (;;) {
742 : auto [ec, bytes] = co_await read_message_some(buffer, sizeof(buffer));
743 : if (ec) co_return;
744 : }
745 356 : };
746 : co_await (discard() || close_timer_->async_wait(use_nothrow_awaitable));
747 : }
748 : }
749 : close(code);
750 1064 : }
751 :
752 58 : void websocket::cancel() {
753 58 : if (reading_) read_cancelled_ = true;
754 58 : socket_->cancel();
755 58 : wake_writers();
756 58 : }
757 :
758 : // ---------------------------------------------------------------------------
759 : // Keepalive
760 : // ---------------------------------------------------------------------------
761 :
762 4018 : void websocket::set_keepalive(std::chrono::milliseconds ping_interval, std::chrono::milliseconds idle_timeout) {
763 4018 : ping_interval_ = ping_interval;
764 4018 : idle_timeout_ = idle_timeout;
765 4018 : keepalive_.cancel();
766 4018 : if (closed_ || close_requested_ || (ping_interval.count() <= 0 && idle_timeout.count() <= 0)) return;
767 3514 : if (weak_from_this().expired()) {
768 0 : LOG_ERROR("websocket keepalive requires the websocket to be owned by a std::shared_ptr");
769 0 : return;
770 : }
771 3514 : last_ping_ = first_ping_ = {};
772 3514 : last_received_ = std::chrono::steady_clock::now();
773 3514 : check_keepalive();
774 : }
775 :
776 : namespace {
777 :
778 : // When the idle timeout of a websocket closes it, counted from the last byte `received`, and
779 : // whether a ping goes out first (see websocket::check_keepalive())
780 : struct idle_deadline {
781 : std::chrono::steady_clock::time_point close_at;
782 : bool ping_now = false;
783 : };
784 :
785 : // In the client role, with pings before the idle timeout, it closes only once the first ping
786 : // sent since the last byte received (`first_ping`) went unanswered for the rest of the idle
787 : // timeout: a check run late (nothing ran the io_context, as a synchronous client between its
788 : // calls) pings first. On time, that ping went out at the ping interval: the same instant
789 6923 : idle_deadline idle_deadline_of(std::chrono::steady_clock::time_point now,
790 : std::chrono::steady_clock::time_point received,
791 : std::chrono::steady_clock::time_point first_ping,
792 : std::chrono::milliseconds idle_timeout, std::chrono::milliseconds ping_interval,
793 : bool server_role) {
794 6923 : idle_deadline deadline{received + idle_timeout};
795 6923 : bool pings_first = !server_role && ping_interval.count() > 0 && ping_interval < idle_timeout;
796 6923 : if (pings_first && now >= deadline.close_at) {
797 74 : if (first_ping < received) {
798 9 : deadline.ping_now = true;
799 9 : deadline.close_at = now + (idle_timeout - ping_interval);
800 : } else {
801 65 : deadline.close_at = std::max(deadline.close_at, first_ping + (idle_timeout - ping_interval));
802 : }
803 : }
804 6923 : return deadline;
805 : }
806 :
807 : }
808 :
809 7085 : void websocket::check_keepalive() {
810 : using clock = std::chrono::steady_clock;
811 7151 : if (closed_) return;
812 7085 : auto now = clock::now();
813 : // In the client role, bytes that arrived since the previous check and nothing read count
814 : // as received (e.g. the pongs of a client that only sends). A server always reads: its
815 : // timing does not change
816 7085 : if (!server_role_) {
817 2075 : auto arrived = socket_->is_open() ? available() : 0;
818 2075 : if (arrived > available_seen_) last_received_ = now;
819 2075 : available_seen_ = arrived;
820 : }
821 7085 : auto received = last_received_;
822 : // At most an interval from now: the deadlines below are never later
823 7085 : auto next = now + std::max(idle_timeout_, ping_interval_);
824 7085 : bool ping_now = false;
825 7085 : if (idle_timeout_.count() > 0) {
826 6923 : auto deadline = idle_deadline_of(now, received, first_ping_, idle_timeout_, ping_interval_, server_role_);
827 6923 : if (now >= deadline.close_at) {
828 66 : LOG_DEBUG("websocket idle for too long, closing");
829 66 : close(CLOSE_GOING_AWAY);
830 66 : return;
831 : }
832 6857 : ping_now = deadline.ping_now;
833 6857 : next = deadline.close_at;
834 : }
835 7019 : if (ping_interval_.count() > 0) {
836 7001 : auto since = std::max(received, last_ping_);
837 7001 : if (ping_now || now - since >= ping_interval_) {
838 : // Queued, not waited for: a ping stuck behind a write that does not end does not
839 : // hold the idle timeout. One ping at a time
840 1640 : last_ping_ = now;
841 1640 : if (first_ping_ < received) first_ping_ = now;
842 1640 : if (!pending_ping_ && !ping_in_flight_) {
843 1496 : pending_ping_ = true;
844 1496 : queue_control();
845 : }
846 1640 : since = now;
847 : }
848 7001 : next = std::min(next, since + ping_interval_);
849 : }
850 7019 : keepalive_.expires_at(next);
851 : }
852 :
853 262 : void websocket::start_timeout() {
854 262 : set_keepalive(CONNECTION_TIMEOUT_SECONDS, 2 * CONNECTION_TIMEOUT_SECONDS);
855 262 : }
856 :
857 : // ---------------------------------------------------------------------------
858 : // Delegated to the socket
859 : // ---------------------------------------------------------------------------
860 :
861 0 : awaitable<boost::system::error_code> websocket::wait(boost::asio::socket_base::wait_type type) {
862 : if (type == boost::asio::socket_base::wait_read && buffered() > 0) {
863 : co_return boost::system::error_code{};
864 : }
865 : co_return co_await socket_->wait(type);
866 0 : }
867 :
868 6 : awaitable<boost::system::error_code> websocket::connect(
869 : const std::string& host,
870 : const std::string& port,
871 : std::chrono::seconds timeout) {
872 : co_return co_await socket_->connect(host, port, timeout);
873 12 : }
874 :
875 0 : bool websocket::requires_handshake() const {
876 0 : return socket_->requires_handshake();
877 : }
878 :
879 0 : awaitable<boost::system::error_code> websocket::handshake(const std::string& host) {
880 : co_return co_await socket_->handshake(host);
881 0 : }
882 :
883 6041 : uint16_t websocket::get_close_code() const {
884 6041 : if (close_received_) return received_close_code_;
885 4937 : if (close_sent_ && !close_unwritten_) return close_code_;
886 2835 : return CLOSE_ABNORMAL;
887 : }
888 :
889 2534857 : bool websocket::is_open() const {
890 2534857 : return !closed_ && socket_->is_open();
891 : }
892 :
893 0 : bool websocket::is_secure() const {
894 0 : return socket_->is_secure();
895 : }
896 :
897 3 : std::optional<tls_session> websocket::tls() const {
898 3 : return socket_->tls();
899 : }
900 :
901 2075 : size_t websocket::available() const {
902 : // received, framing included
903 2075 : if (closed_) return 0;
904 2075 : return buffered() + socket_->available();
905 : }
906 :
907 98 : size_t websocket::send_backlog() const {
908 98 : return closed_ ? 0 : socket_->send_backlog();
909 : }
910 :
911 0 : std::string websocket::get_remote_ip() const {
912 0 : return socket_->get_remote_ip();
913 : }
914 :
915 0 : std::string websocket::get_local_port() const {
916 0 : return socket_->get_local_port();
917 : }
918 :
919 0 : std::string websocket::get_remote_port() const {
920 0 : return socket_->get_remote_port();
921 : }
922 :
923 66446 : uint64_t websocket::remaining_in_frame() const {
924 66446 : return frame_remaining_;
925 : }
926 :
927 3636 : bool websocket::is_message_complete() const {
928 3636 : return new_message_;
929 : }
930 :
931 : }
|
