Line data    Source code

       1              : #ifndef THINGER_ASIO_SOCKET_PIPE_HPP
       2              : #define THINGER_ASIO_SOCKET_PIPE_HPP
       3              : 
       4              : #include <memory>
       5              : #include <functional>
       6              : #include <atomic>
       7              : #include <chrono>
       8              : 
       9              : #include "sockets/socket.hpp"
      10              : #include "watchdog.hpp"
      11              : #include "../util/types.hpp"
      12              : 
      13              : namespace thinger::asio {
      14              : 
      15              : /// Bidirectional coroutine-based pipe between two sockets: a tunnel between any two sockets
      16              : /// (TCP, TLS, Unix, a WebSocket, or one defined by the application, see sockets/socket.hpp).
      17              : /// Each direction writes what it read before reading more, so a slow reader slows its writer
      18              : /// down. Takes exclusive ownership of both sockets and closes them when the pipe ends, which
      19              : /// it does when either direction ends (the end of a stream, or an error): closing one end of
      20              : /// the tunnel closes the other.
      21              : /// Both sockets are used from the coroutine running run() (start() runs it on the io_context
      22              : /// of the source), so they must belong to that io_context.
      23              : /// set_on_end() and set_idle_timeout() must be called before run()/start(). With an idle
      24              : /// deadline, both directions move one timer, so run() must run where they never run at once:
      25              : /// an io_context run by one thread (as the library's workers are), or a strand.
      26              : /// cancel(), the transfer stats and the socket getters may be called from any thread.
      27              : class socket_pipe : public std::enable_shared_from_this<socket_pipe> {
      28              : public:
      29              :     static constexpr size_t BUFFER_SIZE = 8192;
      30              : 
      31              :     socket_pipe(std::shared_ptr<socket> source, std::shared_ptr<socket> target);
      32              :     ~socket_pipe();
      33              : 
      34              :     /// Run bidirectionally. Completes when either direction ends or errors.
      35              :     awaitable<void> run();
      36              : 
      37              :     /// Fire-and-forget: co_spawn(run()) on the source socket's io_context.
      38              :     void start();
      39              : 
      40              :     /// Close both sockets, each on its own io_context (right away when called from a thread
      41              :     /// running it), so run() completes. Thread-safe; only the first call closes them.
      42              :     void cancel();
      43              : 
      44              :     /// Completion callback (called from destructor).
      45              :     void set_on_end(sync_callback<void()> listener);
      46              : 
      47              :     /// End the pipe (both sockets closed, as cancel() does) once nothing was forwarded in
      48              :     /// either direction for `timeout`. A byte counts once written to the other socket, so a
      49              :     /// peer that does not read lets the pipe end too. Off by default (zero), as in 0.8: a
      50              :     /// generic tunnel, often long-lived and quiet with its own keepalive, has no deadline
      51              :     /// unless one is set (the reverse proxy sets one, see proxy_options::tunnel_idle_timeout).
      52              :     void set_idle_timeout(std::chrono::milliseconds timeout);
      53              : 
      54              :     /// Transfer stats (safe to read from any thread after run() completes).
      55              :     size_t bytes_source_to_target() const;
      56              :     size_t bytes_target_to_source() const;
      57              : 
      58              :     /// Socket access.
      59              :     std::shared_ptr<socket> get_source() const;
      60              :     std::shared_ptr<socket> get_target() const;
      61              : 
      62              : private:
      63              :     /// Forward data in one direction: from -> to, moving the idle deadline (if any) after
      64              :     /// each write.
      65              :     awaitable<void> forward(
      66              :         std::shared_ptr<socket> from,
      67              :         std::shared_ptr<socket> to,
      68              :         std::atomic<size_t>& bytes_transferred,
      69              :         watchdog* idle);
      70              : 
      71              :     std::shared_ptr<socket> source_;
      72              :     std::shared_ptr<socket> target_;
      73              :     sync_callback<void()> on_end_;
      74          135 :     std::chrono::milliseconds idle_timeout_{0};
      75              :     std::atomic<bool> cancelled_{false};
      76              :     std::atomic<size_t> bytes_s2t_{0};
      77              :     std::atomic<size_t> bytes_t2s_{0};
      78              : };
      79              : 
      80              : } // namespace thinger::asio
      81              : 
      82              : #endif