ws_connection.hpp

Go to the documentation of this file.

// SPDX-License-Identifier: Apache-2.0
 
#pragma once
#include <algorithm>
#include <array>
#include <atomic>
#include <chrono>
#include <core/async/Task.hpp>
#include <core/net/EventLoop.hpp>
#include <core/net/ISocket.hpp>
#include <cstddef>
#include <deque>
#include <functional>
#include <memory>
#include <optional>
#include <span>
#include <string>
#include <string_view>
#include <utility>
 
#include "../../core/connection_writer.hpp"
#include "../../core/profiler.hpp"
#include "ws_handshake.hpp"
 
/// @file
/// @brief One WebSocket connection on an I/O loop: its socket, the frames
///        waiting to be written, and the two flows that move bytes.
///
/// Shared by `SocketBackend` (client role) and `SocketServer` (server role).
/// Everything here runs on the loop's thread and is touched nowhere else; the
/// `LoopConnection` itself is held by `shared_ptr` from each flow that uses it,
/// so a flow that resumes after its owner has let go of the connection still
/// finds the socket object alive. The one exception is the backlog counter,
/// `pending`, which other threads add to and read.
 
namespace morph::net::detail {
 
/// @brief One whole frame waiting for the writer.
struct OutboundFrame {
    /// @brief The encoded WebSocket frame.
    std::string bytes;
    /// @brief What the frame added to `LoopConnection::pending`: the size of
    ///        the envelope it carries, or `0` for a frame that is not counted.
    std::size_t counted{0};
};
 
/// @brief A connected socket on the loop, plus its outgoing frame queue.
struct LoopConnection {
    /// @param eventLoop    The loop the socket is registered with.
    /// @param connected    The socket; owned from here on.
    /// @param writeTimeout Bound on writing one chunk of a frame; zero disables it.
    LoopConnection(::core::net::EventLoop& eventLoop, std::unique_ptr<::core::net::ISocket> connected,
                   std::chrono::milliseconds writeTimeout)
        : loop{eventLoop}, socket{std::move(connected)}, sendTimeout{writeTimeout} {}
 
    LoopConnection(const LoopConnection&) = delete;
    LoopConnection& operator=(const LoopConnection&) = delete;
    LoopConnection(LoopConnection&&) = delete;
    LoopConnection& operator=(LoopConnection&&) = delete;
    ~LoopConnection() = default;
 
    /// @brief The loop the socket belongs to.
    ::core::net::EventLoop& loop;
    /// @brief The socket. Closed, never replaced, for the connection's life.
    std::unique_ptr<::core::net::ISocket> socket;
    /// @brief Bound on one chunk's write; see `writerFlow`.
    std::chrono::milliseconds sendTimeout;
    /// @brief Whole frames waiting for the writer, oldest first.
    std::deque<OutboundFrame> outbox;
    /// @brief Bytes of encoded envelopes handed to this connection and not yet
    ///        written: its backlog, as `RemoteServer` reads it.
    ///
    /// Added to by whoever hands over a frame, on any thread, before the frame
    /// reaches the loop, so the count covers frames still on their way to the
    /// outbox. Subtracted on the loop by the same amount once the frame's last
    /// chunk is written, or once the frame is dropped. Counting envelopes, not
    /// framed bytes, makes the add and the subtract the same number, so the
    /// count cannot drift. Held by `shared_ptr` so a thread off the loop can
    /// keep the counter without keeping the connection, whose socket must be
    /// released on the loop.
    std::shared_ptr<std::atomic<std::size_t>> const pending{std::make_shared<std::atomic<std::size_t>>(0)};
    /// @brief Called on the loop when a written frame takes `pending` from at
    ///        least `backend::kEventBacklogLow` to below it. Not called once
    ///        the connection is closed or closing. Empty means nobody is told.
    std::function<void()> onDrained;
    /// @brief Whether a writer flow is running.
    bool writing{false};
    /// @brief Whether the socket has been closed.
    bool closed{false};
    /// @brief Close once the outbox has drained (a Close echo, say).
    bool closeWhenFlushed{false};
};
 
/// @brief Takes @p bytes off @p conn's backlog, and calls `onDrained` if
///        that takes it below `backend::kEventBacklogLow`.
///
/// The crossing is judged from the value this subtraction replaced, so of
/// several releases that pass the bound together exactly one reports it.
/// Frames handed over meanwhile on other threads may lift the count again;
/// the server's drain handling is idempotent, so a second report is harmless.
/// @param conn  The connection.
/// @param bytes What the frame added when it was handed over.
inline void releaseBacklog(LoopConnection& conn, std::size_t bytes) {
    if (bytes == 0) {
        return;
    }
    std::size_t const before = conn.pending->fetch_sub(bytes);
    bool const crossed =
        before >= ::morph::backend::kEventBacklogLow && before - bytes < ::morph::backend::kEventBacklogLow;
    if (crossed && !conn.closed && !conn.closeWhenFlushed && conn.onDrained) {
        conn.onDrained();
    }
}
 
/// @brief Closes @p conn's socket now, dropping any unwritten frames.
///
/// A read or write parked on it resumes with `NetErrorCode::Cancelled` in a
/// later drain step of the loop, which is how the flows learn of it. The
/// dropped frames leave the backlog; a frame the writer holds leaves it when
/// the writer lets go of it.
/// @param conn The connection; idempotent.
inline void closeConnection(LoopConnection& conn) noexcept {
    if (conn.closed) {
        return;
    }
    conn.closed = true;
    std::size_t dropped = 0;
    for (auto const& frame : conn.outbox) {
        dropped += frame.counted;
    }
    conn.outbox.clear();
    conn.pending->fetch_sub(dropped);
    conn.socket->close();
}
 
/// @brief The largest single write the writer issues, so that `sendTimeout`
///        bounds a write making no progress rather than a large frame on a
///        slow link.
inline constexpr std::size_t kWriteChunk = std::size_t{64} * 1024;
 
/// @brief Timer callback: a write, or a handshake read, overran its bound.
/// @param connPtr The `LoopConnection`.
inline void closeOnDeadline(void* connPtr) { closeConnection(*static_cast<LoopConnection*>(connPtr)); }
 
/// @brief Writes every queued frame, in order, one chunk at a time.
///
/// Each chunk's write is bounded by `sendTimeout`: a loop timer that closes
/// the socket if the write has not finished, so a peer that stops reading
/// retires the connection instead of growing its queue forever. A frame
/// leaves the backlog once its last chunk is written, or once the connection
/// closes under it.
/// @param conn The connection; the flow's own share keeps its socket alive.
/// @return The flow.
inline ::core::async::Task<void> writerFlow(std::shared_ptr<LoopConnection> conn) {
    while (!conn->closed && !conn->outbox.empty()) {
        OutboundFrame const next = std::move(conn->outbox.front());
        conn->outbox.pop_front();
        std::string const& frame = next.bytes;
        std::size_t offset = 0;
        while (offset < frame.size() && !conn->closed) {
            std::size_t const length = std::min(kWriteChunk, frame.size() - offset);
            ::core::net::TimerId timer{};
            if (conn->sendTimeout.count() > 0) {
                timer =
                    conn->loop.addTimer(conn->loop.clock().now() + conn->sendTimeout, &closeOnDeadline, conn.get());
            }
            auto const written =
                co_await conn->socket->write(std::as_bytes(std::span<char const>{frame}.subspan(offset, length)));
            static_cast<void>(conn->loop.cancelTimer(timer));
            if (!written || *written == 0) {
                closeConnection(*conn);
                break;
            }
            offset += *written;
        }
        releaseBacklog(*conn, next.counted);
    }
    conn->writing = false;
    if (conn->closeWhenFlushed) {
        closeConnection(*conn);
    }
}
 
/// @brief Queues @p frame on @p conn and starts the writer if it is idle.
/// @param conn    The connection.
/// @param frame   One complete, encoded WebSocket frame.
/// @param counted What the caller added to `conn->pending` for this frame;
///                `0` for a frame that is not counted.
/// @return `false` if the connection is closed and the frame was dropped, its
///         count released with it.
inline bool enqueueFrame(std::shared_ptr<LoopConnection> const& conn, std::string frame, std::size_t counted = 0) {
    // The send side's zone is here and not in `writerFlow`: a zone must not
    // stay open across a `co_await`, where the loop runs other work on the
    // same thread before the writer resumes.
    MORPH_ZONE("ws::enqueueFrame");
    if (conn->closed || conn->closeWhenFlushed) {
        releaseBacklog(*conn, counted);
        return false;
    }
    conn->outbox.push_back(OutboundFrame{.bytes = std::move(frame), .counted = counted});
    if (!conn->writing) {
        conn->writing = true;
        conn->loop.spawn(writerFlow(conn));
    }
    return true;
}
 
/// @brief Closes @p conn once every queued frame is written, or now if none is.
/// @param conn The connection.
inline void closeAfterFlush(std::shared_ptr<LoopConnection> const& conn) {
    if (conn->writing) {
        conn->closeWhenFlushed = true;
        return;
    }
    closeConnection(*conn);
}
 
/// @brief Reads an HTTP header block — up to and including `\r\n\r\n` — off
///        @p conn, as the WebSocket handshake needs.
///
/// @p timeout bounds the whole read, from the first byte to the terminator:
/// a loop timer closes the socket when it runs out, so a peer that dribbles
/// one byte at a time cannot stretch the handshake past it. The header is
/// capped at 64 KiB, checked before each read, like `readHttpHeaderBlock`.
/// @param conn    The connection; must have no other read outstanding.
/// @param timeout Bound on the whole read; zero disables it.
/// @return The header and any bytes read past it, or `std::nullopt` if the
///         peer closed, the read failed, the cap was reached or the time ran out.
inline ::core::async::Task<std::optional<HandshakeReadResult>> readHeaderBlockAsync(
    std::shared_ptr<LoopConnection> conn, std::chrono::milliseconds timeout) {
    ::core::net::TimerId timer{};
    if (timeout.count() > 0) {
        timer = conn->loop.addTimer(conn->loop.clock().now() + timeout, &closeOnDeadline, conn.get());
    }
    std::string buf;
    std::array<std::byte, 4096> chunk{};
    std::optional<HandshakeReadResult> result;
    for (;;) {
        if (auto const pos = buf.find("\r\n\r\n"); pos != std::string::npos) {
            result = HandshakeReadResult{.header = buf.substr(0, pos), .leftover = buf.substr(pos + 4)};
            break;
        }
        if (buf.size() > std::size_t{64} * 1024 || conn->closed) {
            break;
        }
        auto const got = co_await conn->socket->read(chunk);
        if (!got || *got == 0) {
            break;
        }
        // NOLINTNEXTLINE(cppcoreguidelines-pro-type-reinterpret-cast): the socket yields bytes, the buffer holds text; same object representation.
        buf.append(reinterpret_cast<char const*>(chunk.data()), *got);
    }
    static_cast<void>(conn->loop.cancelTimer(timer));
    co_return result;
}
 
}  // namespace morph::net::detail