Line data    Source code

       1              : #include "socket_pipe.hpp"
       2              : #include "../util/logger.hpp"
       3              : 
       4              : #include <boost/asio/dispatch.hpp>
       5              : #include <boost/asio/this_coro.hpp>
       6              : #include <vector>
       7              : 
       8              : namespace thinger::asio {
       9              : 
      10              : namespace {
      11              : 
      12              : // Close a socket on its own io_context, as sockets are not thread-safe: right away when called
      13              : // from a thread running it, otherwise once it gets there
      14          484 : void close_on_own_context(const std::shared_ptr<socket>& sock) {
      15          968 :     boost::asio::dispatch(sock->get_io_context(), [sock] { sock->close(); });
      16          484 : }
      17              : 
      18              : } // namespace
      19              : 
      20          135 : socket_pipe::socket_pipe(std::shared_ptr<socket> source, std::shared_ptr<socket> target)
      21          135 :     : source_(std::move(source)), target_(std::move(target)) {
      22          135 : }
      23              : 
      24          135 : socket_pipe::~socket_pipe() {
      25              :     try {
      26          135 :         if (on_end_) on_end_();
      27            0 :     } catch (...) {}
      28          135 : }
      29              : 
      30          132 : awaitable<void> socket_pipe::run() {
      31              :     auto self = shared_from_this();
      32              :     if (idle_timeout_.count() > 0) {
      33              :         // One timer for the whole run, which each forwarded write moves later without touching
      34              :         // it (see watchdog). A local of run(): it never calls back once run() has ended
      35           27 :         watchdog idle(co_await boost::asio::this_coro::executor, [this] {
      36           27 :             LOG_DEBUG("socket pipe idle for {} ms: closing it", idle_timeout_.count());
      37           27 :             cancel();
      38           27 :         });
      39              :         idle.expires_at(watchdog::clock::now() + idle_timeout_);
      40              :         co_await (
      41              :             forward(source_, target_, bytes_s2t_, &idle) ||
      42              :             forward(target_, source_, bytes_t2s_, &idle)
      43              :         );
      44              :     } else {
      45              :         co_await (
      46              :             forward(source_, target_, bytes_s2t_, nullptr) ||
      47              :             forward(target_, source_, bytes_t2s_, nullptr)
      48              :         );
      49              :     }
      50              :     cancel();
      51              :     // Closed before returning, also when cancel() was called from another thread and has not
      52              :     // closed them yet (a direction may see the flag and end first)
      53              :     close_on_own_context(source_);
      54              :     close_on_own_context(target_);
      55          264 : }
      56              : 
      57            6 : void socket_pipe::start() {
      58            6 :     auto self = shared_from_this();
      59           12 :     co_spawn(source_->get_io_context(), [self]() -> awaitable<void> {
      60              :         co_await self->run();
      61           12 :     }, detached);
      62            6 : }
      63              : 
      64          391 : void socket_pipe::cancel() {
      65          391 :     if (cancelled_.exchange(true)) return;
      66          132 :     close_on_own_context(source_);
      67          132 :     close_on_own_context(target_);
      68              : }
      69              : 
      70           15 : void socket_pipe::set_on_end(sync_callback<void()> listener) {
      71           15 :     on_end_ = std::move(listener);
      72           15 : }
      73              : 
      74           57 : void socket_pipe::set_idle_timeout(std::chrono::milliseconds timeout) {
      75           57 :     idle_timeout_ = timeout;
      76           57 : }
      77              : 
      78            3 : size_t socket_pipe::bytes_source_to_target() const {
      79            6 :     return bytes_s2t_.load();
      80              : }
      81              : 
      82            3 : size_t socket_pipe::bytes_target_to_source() const {
      83            6 :     return bytes_t2s_.load();
      84              : }
      85              : 
      86            3 : std::shared_ptr<socket> socket_pipe::get_source() const {
      87            3 :     return source_;
      88              : }
      89              : 
      90            3 : std::shared_ptr<socket> socket_pipe::get_target() const {
      91            3 :     return target_;
      92              : }
      93              : 
      94          264 : awaitable<void> socket_pipe::forward(
      95              :     std::shared_ptr<socket> from,
      96              :     std::shared_ptr<socket> to,
      97              :     std::atomic<size_t>& bytes_transferred,
      98              :     watchdog* idle)
      99              : {
     100              :     std::vector<uint8_t> buffer(BUFFER_SIZE);
     101              :     while (!cancelled_) {
     102              :         auto [read_ec, n] = co_await from->read_some(buffer.data(), BUFFER_SIZE);
     103              :         if (read_ec) break;
     104              :         auto [write_ec, written] = co_await to->write(buffer.data(), n);
     105              :         if (write_ec) break;
     106              :         bytes_transferred.fetch_add(n, std::memory_order_relaxed);
     107              :         if (idle) idle->expires_at(watchdog::clock::now() + idle_timeout_);
     108              :     }
     109              :     // Close both sockets to interrupt the other direction
     110              :     cancel();
     111          528 : }
     112              : 
     113              : } // namespace thinger::asio