Line data Source code
1 : #include <iostream>
2 : #include "tcp_socket.hpp"
3 :
4 : namespace thinger::asio {
5 :
6 769068 : tcp_socket::tcp_socket(const std::string& context, boost::asio::io_context& io_context)
7 769068 : : socket(context, io_context), socket_(io_context) {
8 769068 : }
9 :
10 0 : tcp_socket::tcp_socket(const std::string& context, const std::shared_ptr<tcp_socket>& tcp_socket)
11 0 : : socket(context, tcp_socket->get_io_context()), socket_(std::move(tcp_socket->get_socket())) {
12 0 : }
13 :
14 15 : tcp_socket::tcp_socket(const std::string& context, boost::asio::ip::tcp::socket&& sock)
15 15 : : socket(context, static_cast<boost::asio::io_context&>(sock.get_executor().context())), socket_(std::move(sock)) {
16 15 : }
17 :
18 769083 : tcp_socket::~tcp_socket() {
19 769083 : LOG_TRACE("releasing tcp connection");
20 769083 : close();
21 769083 : }
22 :
23 770119 : void tcp_socket::close() {
24 770119 : boost::system::error_code ec;
25 770119 : if (socket_.is_open()) {
26 1660 : socket_.shutdown(boost::asio::ip::tcp::socket::shutdown_both, ec);
27 : }
28 770119 : socket_.close(ec);
29 770119 : LOG_TRACE("closing tcp socket result: {}", ec.message());
30 770119 : }
31 :
32 83 : void tcp_socket::cancel() {
33 83 : socket_.cancel();
34 79 : }
35 :
36 826 : awaitable<boost::system::error_code> tcp_socket::connect(
37 : const std::string& host,
38 : const std::string& port,
39 : std::chrono::seconds timeout)
40 : {
41 : close();
42 :
43 : // Resolve host
44 : boost::asio::ip::tcp::resolver resolver(io_context_);
45 : auto [ec_resolve, endpoints] = co_await resolver.async_resolve(
46 : host, port, use_nothrow_awaitable);
47 :
48 : if (ec_resolve) {
49 : co_return ec_resolve;
50 : }
51 :
52 : // Setup timeout timer
53 : boost::asio::steady_timer timer(io_context_);
54 : timer.expires_after(timeout);
55 :
56 : // Race between connect and timeout
57 : bool timed_out = false;
58 : boost::system::error_code connect_ec;
59 :
60 : // Start timeout
61 817 : auto timeout_coro = [&]() -> awaitable<void> {
62 : auto [ec] = co_await timer.async_wait(use_nothrow_awaitable);
63 : if (!ec) {
64 : timed_out = true;
65 : socket_.cancel();
66 : }
67 1634 : };
68 :
69 : // Start connect
70 817 : auto connect_coro = [&]() -> awaitable<void> {
71 : auto [ec, endpoint] = co_await boost::asio::async_connect(
72 : socket_, endpoints, use_nothrow_awaitable);
73 : timer.cancel();
74 : if (ec) {
75 : connect_ec = timed_out ? boost::asio::error::timed_out : ec;
76 : }
77 1634 : };
78 :
79 : // Wait for connect (timeout will cancel if needed)
80 : co_spawn(io_context_, timeout_coro(), detached);
81 : co_await connect_coro();
82 :
83 : if (connect_ec) {
84 : co_return connect_ec;
85 : }
86 :
87 : // Run handshake if required (for SSL sockets)
88 : if (requires_handshake()) {
89 : auto hs_ec = co_await handshake(host);
90 : if (hs_ec) {
91 : close();
92 : co_return hs_ec;
93 : }
94 : }
95 :
96 : co_return boost::system::error_code{};
97 1652 : }
98 :
99 1421 : boost::asio::ip::tcp::socket& tcp_socket::get_socket() {
100 1421 : return socket_;
101 : }
102 :
103 831 : std::string tcp_socket::get_remote_ip() const {
104 831 : boost::system::error_code ec;
105 831 : auto remote_ep = socket_.remote_endpoint(ec);
106 831 : if (!ec) {
107 831 : return remote_ep.address().to_string();
108 : }
109 0 : return "0.0.0.0";
110 : }
111 :
112 834 : std::string tcp_socket::get_local_port() const {
113 834 : boost::system::error_code ec;
114 834 : auto local_ep = socket_.local_endpoint(ec);
115 834 : if (!ec) {
116 834 : return std::to_string(local_ep.port());
117 : }
118 0 : return "0";
119 : }
120 :
121 3 : std::string tcp_socket::get_remote_port() const {
122 3 : boost::system::error_code ec;
123 3 : auto remote_ep = socket_.remote_endpoint(ec);
124 3 : if (!ec) {
125 3 : return std::to_string(remote_ep.port());
126 : }
127 0 : return "0";
128 : }
129 :
130 4417 : awaitable<size_t> tcp_socket::read_some(uint8_t* buffer, size_t max_size) {
131 : auto [ec, bytes] = co_await socket_.async_read_some(
132 : boost::asio::buffer(buffer, max_size),
133 : use_nothrow_awaitable);
134 : if (ec) co_return 0;
135 : co_return bytes;
136 8834 : }
137 :
138 559 : awaitable<size_t> tcp_socket::read(uint8_t* buffer, size_t size) {
139 : auto [ec, bytes] = co_await boost::asio::async_read(
140 : socket_,
141 : boost::asio::buffer(buffer, size),
142 : boost::asio::transfer_exactly(size),
143 : use_nothrow_awaitable);
144 : if (ec) co_return 0;
145 : co_return bytes;
146 1118 : }
147 :
148 3 : awaitable<size_t> tcp_socket::read(boost::asio::streambuf& buffer, size_t size) {
149 : auto [ec, bytes] = co_await boost::asio::async_read(
150 : socket_,
151 : buffer,
152 : boost::asio::transfer_exactly(size),
153 : use_nothrow_awaitable);
154 : if (ec) co_return 0;
155 : co_return bytes;
156 6 : }
157 :
158 3 : awaitable<size_t> tcp_socket::read_until(boost::asio::streambuf& buffer, std::string_view delim) {
159 : auto [ec, bytes] = co_await boost::asio::async_read_until(
160 : socket_,
161 : buffer,
162 : std::string(delim),
163 : use_nothrow_awaitable);
164 : if (ec) co_return 0;
165 : co_return bytes;
166 6 : }
167 :
168 846 : awaitable<size_t> tcp_socket::write(const uint8_t* buffer, size_t size) {
169 : auto [ec, bytes] = co_await boost::asio::async_write(
170 : socket_,
171 : boost::asio::buffer(buffer, size),
172 : use_nothrow_awaitable);
173 : if (ec) co_return 0;
174 : co_return bytes;
175 1692 : }
176 :
177 9 : awaitable<size_t> tcp_socket::write(std::string_view str) {
178 : auto [ec, bytes] = co_await boost::asio::async_write(
179 : socket_,
180 : boost::asio::buffer(str.data(), str.size()),
181 : use_nothrow_awaitable);
182 : if (ec) co_return 0;
183 : co_return bytes;
184 18 : }
185 :
186 1687 : awaitable<size_t> tcp_socket::write(const std::vector<boost::asio::const_buffer>& buffers) {
187 : auto [ec, bytes] = co_await boost::asio::async_write(
188 : socket_,
189 : buffers,
190 : use_nothrow_awaitable);
191 : if (ec) co_return 0;
192 : co_return bytes;
193 3374 : }
194 :
195 6 : awaitable<boost::system::error_code> tcp_socket::wait(boost::asio::socket_base::wait_type type) {
196 : auto [ec] = co_await socket_.async_wait(type, use_nothrow_awaitable);
197 : co_return ec;
198 12 : }
199 :
200 833 : void tcp_socket::enable_tcp_no_delay() {
201 833 : socket_.set_option(boost::asio::ip::tcp::no_delay(true));
202 831 : }
203 :
204 5 : void tcp_socket::disable_tcp_no_delay() {
205 5 : socket_.set_option(boost::asio::ip::tcp::no_delay(false));
206 3 : }
207 :
208 5405 : bool tcp_socket::is_open() const {
209 5405 : return socket_.is_open();
210 : }
211 :
212 1583 : bool tcp_socket::is_secure() const {
213 1583 : return false;
214 : }
215 :
216 7 : size_t tcp_socket::available() const {
217 7 : boost::system::error_code ec;
218 7 : auto size = socket_.available(ec);
219 7 : if (ec) {
220 4 : LOG_ERROR("error while getting socket available bytes ({}): {}", size, ec.message());
221 : }
222 7 : return size;
223 : }
224 :
225 : }
|