diff --git a/Core/socket/Socket_Coro.h b/Core/socket/Socket_Coro.h new file mode 100644 index 0000000..7d9efde --- /dev/null +++ b/Core/socket/Socket_Coro.h @@ -0,0 +1,229 @@ +#pragma once + +#include "ASIO_Utils.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace Psc::asio_socket::coro { + +struct TCP_Read_Result { + std::vector data; +}; + +struct UDP_Read_Result { + Sockaddr_In remote; + std::vector data; +}; + +inline void throw_if_error(const asio::error_code& ec) +{ + if (ec) { + throw std::system_error(ec); + } +} + +inline ucoro::awaitable +tcp_resolve(asio::ip::tcp::resolver& resolver, + std::string host, + std::uint16_t port) +{ + struct Result { + asio::error_code ec; + asio::ip::tcp::resolver::results_type endpoints; + }; + + auto result = co_await ucoro::callback_awaitable( + [&resolver, host = std::move(host), port](auto done) mutable { + resolver.async_resolve( + host, + std::to_string(port), + [done = std::move(done)](const asio::error_code& ec, + asio::ip::tcp::resolver::results_type endpoints) mutable { + done(Result{ec, std::move(endpoints)}); + }); + }); + + throw_if_error(result.ec); + co_return std::move(result.endpoints); +} + +inline ucoro::awaitable +tcp_connect(asio::ip::tcp::socket& socket, + const asio::ip::tcp::resolver::results_type& endpoints) +{ + struct Result { + asio::error_code ec; + asio::ip::tcp::endpoint endpoint; + }; + + auto result = co_await ucoro::callback_awaitable( + [&socket, &endpoints](auto done) mutable { + asio::async_connect( + socket, + endpoints, + [done = std::move(done)](const asio::error_code& ec, + const asio::ip::tcp::endpoint& endpoint) mutable { + done(Result{ec, endpoint}); + }); + }); + + throw_if_error(result.ec); + co_return result.endpoint; +} + +inline ucoro::awaitable +tcp_connect(asio::ip::tcp::socket& socket, + asio::ip::tcp::resolver& resolver, + std::string host, + std::uint16_t port) +{ + auto endpoints = co_await tcp_resolve(resolver, std::move(host), port); + co_return co_await tcp_connect(socket, endpoints); +} + +inline ucoro::awaitable> +tcp_accept(asio::ip::tcp::acceptor& acceptor) +{ + auto socket = std::make_shared(acceptor.get_executor()); + + auto ec = co_await ucoro::callback_awaitable( + [&acceptor, socket](auto done) mutable { + acceptor.async_accept( + *socket, + [done = std::move(done)](const asio::error_code& ec) mutable { + done(ec); + }); + }); + + throw_if_error(ec); + co_return socket; +} + +inline ucoro::awaitable +tcp_read_some(asio::ip::tcp::socket& socket, + std::size_t max_size = 16 * 1024) +{ + struct Result { + asio::error_code ec; + std::size_t size{}; + }; + + auto buffer = std::make_shared>(max_size); + auto result = co_await ucoro::callback_awaitable( + [&socket, buffer](auto done) mutable { + socket.async_read_some( + asio::buffer(*buffer), + [done = std::move(done)](const asio::error_code& ec, + std::size_t size) mutable { + done(Result{ec, size}); + }); + }); + + throw_if_error(result.ec); + buffer->resize(result.size); + co_return TCP_Read_Result{std::move(*buffer)}; +} + +inline ucoro::awaitable +tcp_write(asio::ip::tcp::socket& socket, + std::string data) +{ + struct Result { + asio::error_code ec; + std::size_t size{}; + }; + + auto buffer = std::make_shared(std::move(data)); + auto result = co_await ucoro::callback_awaitable( + [&socket, buffer](auto done) mutable { + asio::async_write( + socket, + asio::buffer(*buffer), + [done = std::move(done)](const asio::error_code& ec, + std::size_t size) mutable { + done(Result{ec, size}); + }); + }); + + throw_if_error(result.ec); + co_return result.size; +} + +inline ucoro::awaitable +udp_receive_from(asio::ip::udp::socket& socket, + std::size_t max_size = 16 * 1024) +{ + struct Result { + asio::error_code ec; + std::size_t size{}; + asio::ip::udp::endpoint remote; + }; + + auto buffer = std::make_shared>(max_size); + auto remote = std::make_shared(); + auto result = co_await ucoro::callback_awaitable( + [&socket, buffer, remote](auto done) mutable { + socket.async_receive_from( + asio::buffer(*buffer), + *remote, + [done = std::move(done), remote](const asio::error_code& ec, + std::size_t size) mutable { + done(Result{ec, size, *remote}); + }); + }); + + throw_if_error(result.ec); + buffer->resize(result.size); + co_return UDP_Read_Result{endpoint_to_sockaddr(result.remote), std::move(*buffer)}; +} + +inline ucoro::awaitable +udp_send_to(asio::ip::udp::socket& socket, + std::string data, + asio::ip::udp::endpoint remote) +{ + struct Result { + asio::error_code ec; + std::size_t size{}; + }; + + auto buffer = std::make_shared(std::move(data)); + auto result = co_await ucoro::callback_awaitable( + [&socket, buffer, remote = std::move(remote)](auto done) mutable { + socket.async_send_to( + asio::buffer(*buffer), + remote, + [done = std::move(done)](const asio::error_code& ec, + std::size_t size) mutable { + done(Result{ec, size}); + }); + }); + + throw_if_error(result.ec); + co_return result.size; +} + +inline ucoro::awaitable +udp_send_to(asio::ip::udp::socket& socket, + std::string data, + const Sockaddr_In& remote) +{ + asio::error_code ec; + auto address = asio::ip::make_address(remote.ip, ec); + throw_if_error(ec); + co_return co_await udp_send_to( + socket, + std::move(data), + asio::ip::udp::endpoint(address, static_cast(remote.port))); +} + +} // namespace Psc::asio_socket::coro