Line data Source code
1 : #include "websocket_client.hpp"
2 : #include "../../asio/io_context_of.hpp"
3 : #include "../../util/logger.hpp"
4 :
5 : #include <boost/asio/use_future.hpp>
6 :
7 : #include <future>
8 : #include <stdexcept>
9 :
10 : namespace thinger::http {
11 :
12 : namespace {
13 :
14 : // A synchronous call: `operation` on the executor of the WebSocket, waiting for it to
15 : // complete. On an io_context owned by its client (http::client) the call waits running
16 : // it, until the operation completes (not until it runs out of work: the keepalive keeps it
17 : // busy). Otherwise something else runs it, and the call only waits. Refused, before
18 : // anything starts, on a thread running the io_context (a handler, callback or coroutine
19 : // of it, its strands included): it would wait there for a handler it holds up, or run
20 : // the handlers of other connections within its caller
21 : template<typename T>
22 558 : T run_blocking(asio::websocket& websocket, bool owned_io_context, const char* call, awaitable<T> operation) {
23 558 : auto& io_context = websocket.get_io_context();
24 1086 : if (io_context.get_executor().running_in_this_thread() ||
25 528 : asio::running_in_this_thread(websocket.get_executor())) {
26 90 : throw std::logic_error(std::string("websocket_client::") + call +
27 : "() called on a thread running the io_context of its WebSocket, where it "
28 : "cannot wait: use the coroutine API (send_text_async(), receive_async(), "
29 : "close_async()...), or stop() to close");
30 : }
31 528 : auto future = boost::asio::co_spawn(websocket.get_executor(), std::move(operation), boost::asio::use_future);
32 528 : if (owned_io_context) {
33 6690 : asio::run_until(io_context, [&] {
34 6174 : return future.wait_for(std::chrono::seconds(0)) == std::future_status::ready;
35 : });
36 : }
37 1056 : return future.get();
38 528 : }
39 :
40 24 : awaitable<void> apply_keepalive(std::shared_ptr<asio::websocket> websocket,
41 : std::chrono::milliseconds ping_interval, std::chrono::milliseconds idle_timeout) {
42 : websocket->set_keepalive(ping_interval, idle_timeout);
43 : co_return;
44 48 : }
45 :
46 : }
47 :
48 0 : websocket_client::websocket_client(std::shared_ptr<asio::websocket> ws)
49 0 : : websocket_(std::move(ws)) {}
50 :
51 262 : websocket_client::websocket_client(std::shared_ptr<asio::websocket> ws, bool owned_io_context)
52 262 : : websocket_(std::move(ws)), owned_io_context_(owned_io_context) {}
53 :
54 2444 : websocket_client::~websocket_client() {
55 2444 : stop();
56 2444 : }
57 :
58 2182 : websocket_client::websocket_client(websocket_client&& other) noexcept
59 2182 : : websocket_(std::move(other.websocket_)),
60 2182 : running_(other.running_.load()),
61 2182 : on_message_(std::move(other.on_message_)),
62 2182 : on_close_(std::move(other.on_close_)),
63 2182 : on_error_(std::move(other.on_error_)),
64 2182 : owned_io_context_(other.owned_io_context_) {
65 2182 : other.running_ = false;
66 2182 : }
67 :
68 0 : websocket_client& websocket_client::operator=(websocket_client&& other) noexcept {
69 0 : if (this != &other) {
70 0 : stop();
71 0 : websocket_ = std::move(other.websocket_);
72 0 : running_ = other.running_.load();
73 0 : on_message_ = std::move(other.on_message_);
74 0 : on_close_ = std::move(other.on_close_);
75 0 : on_error_ = std::move(other.on_error_);
76 0 : owned_io_context_ = other.owned_io_context_;
77 0 : other.running_ = false;
78 : }
79 0 : return *this;
80 : }
81 :
82 1050 : bool websocket_client::is_open() const {
83 1050 : return websocket_ && websocket_->is_open();
84 : }
85 :
86 36 : std::shared_ptr<asio::websocket> websocket_client::release_socket() {
87 36 : running_ = false;
88 36 : return std::move(websocket_);
89 : }
90 :
91 : // ============================================
92 : // Synchronous API
93 : // ============================================
94 :
95 179 : bool websocket_client::send_text(std::string_view message) {
96 179 : if (!is_open()) return false;
97 370 : return run_blocking(*websocket_, owned_io_context_, "send_text", send_text_async(std::string(message)));
98 : }
99 :
100 6 : bool websocket_client::send_binary(const void* data, size_t size) {
101 6 : return send_binary(std::string_view(static_cast<const char*>(data), size));
102 : }
103 :
104 51 : bool websocket_client::send_binary(std::string_view data) {
105 51 : if (!is_open()) return false;
106 108 : return run_blocking(*websocket_, owned_io_context_, "send_binary", send_binary_async(std::vector<uint8_t>(data.begin(), data.end())));
107 : }
108 :
109 179 : std::pair<std::string, bool> websocket_client::receive() {
110 179 : if (!is_open()) return {"", false};
111 179 : return run_blocking(*websocket_, owned_io_context_, "receive", receive_async());
112 : }
113 :
114 125 : void websocket_client::close() {
115 125 : if (!websocket_) return;
116 125 : run_blocking(*websocket_, owned_io_context_, "close", close_async());
117 : }
118 :
119 63 : void websocket_client::set_keepalive(std::chrono::milliseconds ping_interval, std::chrono::milliseconds idle_timeout) {
120 63 : if (!websocket_) return;
121 63 : if (asio::running_in_this_thread(websocket_->get_executor())) {
122 39 : websocket_->set_keepalive(ping_interval, idle_timeout);
123 39 : return;
124 : }
125 24 : run_blocking(*websocket_, owned_io_context_, "set_keepalive", apply_keepalive(websocket_, ping_interval, idle_timeout));
126 : }
127 :
128 21 : std::chrono::milliseconds websocket_client::get_ping_interval() const {
129 21 : return websocket_ ? websocket_->get_ping_interval() : std::chrono::milliseconds{0};
130 : }
131 :
132 18 : std::chrono::milliseconds websocket_client::get_idle_timeout() const {
133 18 : return websocket_ ? websocket_->get_idle_timeout() : std::chrono::milliseconds{0};
134 : }
135 :
136 : // ============================================
137 : // Async/Callback API
138 : // ============================================
139 :
140 12 : void websocket_client::run() {
141 12 : if (!is_open() || running_) return;
142 :
143 12 : running_ = true;
144 :
145 12 : boost::asio::co_spawn(
146 12 : websocket_->get_executor(),
147 36 : [this]() -> awaitable<void> {
148 : co_await message_loop();
149 24 : },
150 : boost::asio::detached
151 : );
152 : }
153 :
154 2459 : void websocket_client::stop() {
155 2459 : running_ = false;
156 2459 : if (websocket_) {
157 241 : websocket_->close();
158 : }
159 2459 : }
160 :
161 12 : awaitable<void> websocket_client::message_loop() {
162 : while (running_ && is_open()) {
163 : // An empty message is a message: the loop ends when reading fails (closed or lost)
164 : auto message = co_await read_message();
165 : if (!message) break;
166 :
167 : if (on_message_) {
168 : on_message_(message->first, message->second);
169 : }
170 : }
171 :
172 : running_ = false;
173 :
174 : if (on_close_) {
175 : on_close_();
176 : }
177 24 : }
178 :
179 : // ============================================
180 : // Coroutine API
181 : // ============================================
182 :
183 220 : awaitable<bool> websocket_client::send_text_async(std::string message) {
184 : if (!is_open()) co_return false;
185 :
186 : websocket_->set_binary(false);
187 : auto [ec, bytes] = co_await websocket_->write_message(message);
188 : co_return !ec;
189 440 : }
190 :
191 51 : awaitable<bool> websocket_client::send_binary_async(std::vector<uint8_t> data) {
192 : if (!is_open()) co_return false;
193 :
194 : websocket_->set_binary(true);
195 : auto [ec, bytes] = co_await websocket_->write_message(
196 : std::string_view(reinterpret_cast<const char*>(data.data()), data.size()));
197 : co_return !ec;
198 102 : }
199 :
200 226 : awaitable<std::pair<std::string, bool>> websocket_client::receive_async() {
201 : auto message = co_await read_message();
202 : if (!message) co_return std::make_pair(std::string{}, false);
203 : co_return std::move(*message);
204 452 : }
205 :
206 241 : awaitable<std::optional<std::pair<std::string, bool>>> websocket_client::read_message() {
207 : if (!is_open()) co_return std::nullopt;
208 :
209 : // A whole message: every frame of it, read in pieces (it may be empty)
210 : std::string message;
211 : std::array<uint8_t, 65536> buffer;
212 : for (;;) {
213 : auto [ec, bytes_read] = co_await websocket_->read_message_some(buffer.data(), buffer.size());
214 : if (ec) co_return std::nullopt;
215 : message.append(reinterpret_cast<const char*>(buffer.data()), bytes_read);
216 : if (message.size() > MAX_MESSAGE_SIZE) {
217 : LOG_ERROR("websocket message larger than {} bytes, closing", MAX_MESSAGE_SIZE);
218 : websocket_->close();
219 : co_return std::nullopt;
220 : }
221 : if (websocket_->remaining_in_frame() == 0 && websocket_->is_message_complete()) break;
222 : }
223 :
224 : bool is_binary = websocket_->is_binary();
225 : co_return std::make_pair(std::move(message), is_binary);
226 482 : }
227 :
228 154 : awaitable<void> websocket_client::close_async() {
229 : if (websocket_) {
230 : co_await websocket_->close_graceful();
231 : }
232 : running_ = false;
233 308 : }
234 :
235 : } // namespace thinger::http
|
