#ifndef HYPERSWARM_PROTOMUX_H
#define HYPERSWARM_PROTOMUX_H
// Protomux — C++ port of holepunchto/protomux v3.10.1
// Multiplexes multiple message-oriented protocols over a single framed stream.
// Wire-compatible with the JS implementation.
//
// Threading: all methods run on the UDX thread. Callbacks fire inline.
#include "compact_encoding.h"
#include <cstdint>
#include <cstring>
#include <functional>
#include <memory>
#include <string>
#include <unordered_map>
#include <vector>
class Protomux;
class ProtomuxChannel;
// ---- ProtomuxMessage ----
struct ProtomuxMessage {
uint32_t type;
using RecvCallback = std::function<void(const uint8_t *, size_t, ProtomuxChannel *)>;
RecvCallback on_message;
};
// ---- ProtomuxChannel ----
class ProtomuxChannel {
public:
// Identity
std::string protocol;
std::vector<std::string> aliases;
std::vector<uint8_t> id;
// State
bool opened = false;
bool closed = false;
bool destroyed = false;
// Messages registered on this channel
std::vector<ProtomuxMessage> messages;
// Lifecycle callbacks (invoked on UDX thread)
std::function<void(const uint8_t *handshake, size_t len)> on_open;
std::function<void(bool is_remote)> on_close;
std::function<void()> on_destroy;
// Open the channel, optionally sending handshake data.
void open(const uint8_t *handshake = nullptr, size_t handshake_len = 0);
// Close the channel gracefully.
void close();
// Register a message type. Returns the type index.
uint32_t add_message(ProtomuxMessage::RecvCallback cb);
// Send a message of the given type. Returns false if channel not ready.
bool send(uint32_t type, const uint8_t *data, size_t len);
// Batching (delegates to Protomux).
void cork();
void uncork();
private:
friend class Protomux;
Protomux *_mux = nullptr;
uint32_t _local_id = 0; // 1-based, 0 = unassigned
uint32_t _remote_id = 0; // 1-based, 0 = unassigned
// Handshake received from remote (consumed in _fully_open).
std::vector<uint8_t> _remote_handshake;
void _fully_open();
void _close(bool is_remote);
void _destroy();
void _recv(uint32_t type, const uint8_t *data, size_t len);
};
// ---- Protomux ----
class Protomux {
public:
Protomux();
~Protomux();
// Create a channel for a protocol. Returns nullptr if duplicate (unique=true)
// or if the mux is shut down.
ProtomuxChannel *create_channel(
const std::string &protocol,
const std::vector<uint8_t> &id = {},
const std::vector<std::string> &aliases = {},
bool unique = true);
// Check if a channel is open for the given protocol/id.
bool opened(const std::string &protocol, const std::vector<uint8_t> &id = {}) const;
// Pairing: callback when remote opens a channel we don't have yet.
using PairCallback = std::function<void(const std::string &protocol,
const std::vector<uint8_t> &id)>;
void pair(const std::string &protocol, PairCallback cb);
void unpair(const std::string &protocol);
// Access local channels (for CMD_SEND to find godot-rpc channel).
const std::vector<ProtomuxChannel *> &local_channels() const { return _local; }
bool is_idle() const { return _local.size() == _free.size(); }
// Batching.
void cork();
void uncork();
// Feed a decrypted frame into protomux for dispatch.
void recv(const uint8_t *data, size_t len);
// Set the write callback (handles encryption + UDX write).
// Returns whether the carrier can take more after this write. FALSE does not
// mean the write failed — it means the socket's buffer is at its high
// watermark and producers should stop until on_drain fires. libudx reports
// this (udx_stream_write returns 0 there); discarding it here is what left
// every protocol above with no way to avoid burying the connection.
using WriteCallback = std::function<bool(const uint8_t *data, size_t len)>;
void set_write_callback(WriteCallback cb);
// The carrier drained. Invoked by the transport; forwarded to whoever
// registered interest, which is how a parked producer learns to resume.
using DrainCallback = std::function<void()>;
void set_drain_callback(DrainCallback cb);
void notify_drain();
// Shutdown: close all channels, clear state.
void shutdown();
private:
friend class ProtomuxChannel;
// Channel slots: _local[localId-1] = channel (null = freed).
std::vector<ProtomuxChannel *> _local;
std::vector<uint32_t> _free; // Free local slot indices (0-based).
// Ownership: all channels created by this mux, for cleanup.
std::vector<std::unique_ptr<ProtomuxChannel>> _owned;
// Remote channel tracking.
struct RemoteSlot {
ProtomuxChannel *session = nullptr;
std::vector<uint8_t> handshake; // Open message handshake data.
bool buffering = false; // true = buffer messages, false = dispatch.
struct PendingMsg {
uint32_t type;
std::vector<uint8_t> data;
};
std::vector<PendingMsg> pending;
};
std::vector<RemoteSlot *> _remote; // remoteId-1 → slot.
// Protocol info registry.
struct ProtocolInfo {
std::string key;
std::string protocol;
std::vector<uint8_t> id;
std::vector<std::string> alias_keys;
int opened = 0;
std::vector<uint32_t> incoming; // Remote IDs waiting for local createChannel.
std::vector<uint32_t> outgoing; // Local IDs waiting for remote Open.
ProtomuxChannel *last_channel = nullptr;
// Collected: out of _infos, and waiting to be freed. See _gc.
bool retired = false;
};
std::unordered_map<std::string, ProtocolInfo *> _infos;
// JS `_gc` only drops the map entries and lets the collector free the object
// whenever, so calling it twice on the same info is normal there — and the
// port kept the call sites that do (Session#_close, then _onrejectsession).
// We have to free, but not from inside a dispatch that is still walking the
// object. So _gc RETIRES an info here and recv() frees the batch once the
// whole dispatch has unwound. See test/test_protomux_gc.cpp.
std::vector<std::unique_ptr<ProtocolInfo>> _retired;
int _recv_depth = 0;
// Pairing callbacks.
std::unordered_map<std::string, PairCallback> _notify;
// Batching state.
int _corked = 0;
struct BatchEntry {
uint32_t local_id;
std::vector<uint8_t> data; // [type:varint][payload]
};
std::vector<BatchEntry> _batch;
size_t _batch_size = 0; // Accumulated batch size (tracked incrementally).
// Backpressure tracking.
size_t _buffered = 0;
int _remote_backlog = 0;
bool _shutdown_called = false;
static constexpr size_t MAX_BATCH = 60000; // UDX 2-byte BE frames max 65535.
// Output callback.
WriteCallback _write_cb;
DrainCallback _drain_cb;
// ---- Internal decode ----
void _decode(uint32_t remote_id, const uint8_t *data, size_t len);
void _on_control(uint32_t type, const uint8_t *data, size_t len);
void _on_open(const uint8_t *data, size_t len);
void _on_reject(const uint8_t *data, size_t len);
void _on_close(const uint8_t *data, size_t len);
void _on_batch(const uint8_t *data, size_t len);
// ---- Internal encode ----
void _write0(const uint8_t *data, size_t len);
void _send_batch();
void _push_batch(uint32_t local_id, const uint8_t *data, size_t len);
// ---- Helpers ----
ProtocolInfo *_get_info(const std::string &protocol, const std::vector<uint8_t> &id,
const std::vector<std::string> &aliases = {});
void _gc(ProtocolInfo *info);
void _reject_session(ProtocolInfo *info, uint32_t remote_id);
void _request_session(const std::string &protocol, const std::vector<uint8_t> &id,
ProtocolInfo *info);
static std::string _to_key(const std::string &protocol, const std::vector<uint8_t> &id);
};
#endif // HYPERSWARM_PROTOMUX_H