Line data    Source code

       1              : // Protocol upgrades (WebSocket...) through the reverse proxy: a tunnel between the client and
       2              : // the target
       3              : #include "proxy_exchange.hpp"
       4              : #include "../connection_tracker.hpp"
       5              : #include "../../../asio/socket_pipe.hpp"
       6              : #include "../../../asio/sockets/websocket.hpp"
       7              : #include "../../../asio/with_timeout.hpp"
       8              : 
       9              : namespace thinger::http::proxy_detail {
      10              : 
      11              : // A quiet WebSocket behind the proxy, with the server's default keepalive, pings before the
      12              : // default tunnel deadline cuts it (see websocket_connection, upgrade.hpp's ping_interval)
      13              : static_assert(proxy_options::DEFAULT_TUNNEL_IDLE_TIMEOUT > asio::websocket::CONNECTION_TIMEOUT_SECONDS,
      14              :               "the default tunnel idle timeout must exceed the default WebSocket ping interval");
      15              : 
      16              : namespace {
      17              : 
      18              : // Write what a side sent right after the handshake, within the idle deadline, if any: a
      19              : // peer that does not read it ends the tunnel as the pipe would
      20          108 : awaitable<bool> write_early(asio::socket& to, const std::string& data, std::chrono::milliseconds idle_timeout) {
      21              :     if (data.empty()) co_return true;
      22              :     if (idle_timeout.count() <= 0) {
      23              :         auto [ec, bytes] = co_await to.write(data);
      24              :         co_return !ec;
      25              :     }
      26              :     auto result = co_await asio::with_timeout(to.write(data), idle_timeout);
      27              :     co_return result && !std::get<0>(*result);
      28          216 : }
      29              : 
      30              : }
      31              : 
      32           54 : awaitable<void> run_tunnel(std::shared_ptr<asio::socket> client, std::shared_ptr<asio::socket> target,
      33              :                            std::string to_client, std::string to_target,
      34              :                            std::chrono::milliseconds idle_timeout) {
      35              :     // Bytes each side sent right after the handshake go first
      36              :     if (!co_await write_early(*client, to_client, idle_timeout) ||
      37              :         !co_await write_early(*target, to_target, idle_timeout)) {
      38              :         client->close();
      39              :         target->close();
      40              :         co_return;
      41              :     }
      42              :     auto pipe = std::make_shared<asio::socket_pipe>(client, target);
      43              :     pipe->set_idle_timeout(idle_timeout);
      44              :     co_await pipe->run();
      45          108 : }
      46              : 
      47           60 : void proxy_exchange::upgrade() {
      48           60 :     const auto& protocol = head_->get_header(header::upgrade);
      49           60 :     if (!upgrade_requested_ || protocol.empty()) {
      50           18 :         fail(http_response::status::bad_gateway, "Invalid response from the target", "unexpected protocol switch");
      51            6 :         return;
      52              :     }
      53              : 
      54           54 :     http_response out;
      55           54 :     out.set_status(http_response::status::switching_protocols);
      56           54 :     copy_end_to_end(*head_, out);
      57           54 :     if (options_.on_response) options_.on_response(req_, out);
      58           54 :     remove_framing_headers(out);
      59           54 :     if (!has_valid_headers(out)) {
      60            0 :         fail(http_response::status::bad_gateway, "Invalid response from the target", "invalid response header");
      61            0 :         return;
      62              :     }
      63              : 
      64           54 :     res_.status(http_response::status::switching_protocols);
      65           60 :     for (const auto& [name, value] : out.get_headers()) res_.header(name, value);
      66           54 :     res_.header(header::upgrade, protocol);
      67          108 :     res_.header(header::connection, "Upgrade");
      68              : 
      69              :     // Once the 101 is written, the client connection and the target become a tunnel,
      70              :     // starting with what each one sent past its head
      71           54 :     auto target = std::exchange(target_socket_, nullptr);
      72          108 :     res_.take_over([target, early = std::exchange(pending_, {}),
      73           54 :                     idle_timeout = options_.tunnel_idle_timeout](upgraded_connection client) {
      74           54 :         client.scope.set_type(connection_type::tunnel);
      75           54 :         auto& io_context = client.socket->get_io_context();
      76           54 :         co_spawn(io_context, run_tunnel(std::move(client.socket), target, early, std::move(client.buffered),
      77              :                                         idle_timeout), detached);
      78           54 :     });
      79           54 :     result_.status = 101;
      80           54 :     result_.from_target = true;
      81           54 :     result_.upgraded = true;
      82           54 :     result_.complete = true;
      83           54 : }
      84              : 
      85              : }