#include "protomux.h"
#include "debug_log.h"
#include <algorithm>
#include <cstdio>
// ---- String encoding helpers (wire-identical to compact-encoding c.string) ----
static void string_preencode(EncodeState &s, const std::string &str) {
cenc::uint_preencode(s, str.size());
s.end += str.size();
}
static void string_encode(EncodeState &s, const std::string &str) {
cenc::uint_encode(s, str.size());
if (!str.empty()) {
memcpy(s.buffer + s.start, str.data(), str.size());
s.start += str.size();
}
}
static std::string string_decode(EncodeState &s) {
uint64_t len = cenc::uint_decode(s);
if (len == 0 || s.start + len > s.end) return "";
std::string result((const char *)(s.buffer + s.start), (size_t)len);
s.start += (size_t)len;
return result;
}
// Hex encoding for _to_key.
static std::string to_hex(const std::vector<uint8_t> &data) {
if (data.empty()) return "";
std::string hex;
hex.reserve(data.size() * 2);
static const char digits[] = "0123456789abcdef";
for (uint8_t b : data) {
hex.push_back(digits[b >> 4]);
hex.push_back(digits[b & 0x0f]);
}
return hex;
}
// ============================================================================
// ProtomuxChannel
// ============================================================================
void ProtomuxChannel::open(const uint8_t *handshake, size_t handshake_len) {
DHT_LOG("[Protomux] Channel::open('%s') _local_id=%u _remote_id=%u\n",
protocol.c_str(), _local_id, _remote_id);
if (!_mux || _local_id != 0) return; // Already opened or no mux.
// Allocate local ID from free list or grow.
uint32_t idx;
if (!_mux->_free.empty()) {
idx = _mux->_free.back();
_mux->_free.pop_back();
} else {
idx = (uint32_t)_mux->_local.size();
_mux->_local.push_back(nullptr);
}
_local_id = idx + 1; // 1-based.
_mux->_local[idx] = this;
// Get protocol info and track this channel.
auto *info = _mux->_get_info(protocol, id, aliases);
info->opened++;
info->last_channel = this;
if (_remote_id == 0) {
// Remote hasn't opened yet — queue for later pairing.
info->outgoing.push_back(_local_id);
}
// Encode Open control message: [0][1][localId][protocol][id][handshake?]
EncodeState pre = {0, 2, nullptr}; // 2 bytes for [0][1].
cenc::uint_preencode(pre, _local_id);
string_preencode(pre, protocol);
cenc::buffer_preencode(pre, id.data(), id.size());
if (handshake && handshake_len > 0) {
pre.end += handshake_len;
}
auto buf = cenc::state_alloc(pre);
EncodeState enc = {0, pre.end, buf.data()};
enc.buffer[enc.start++] = 0; // remoteId = 0 (control).
enc.buffer[enc.start++] = 1; // type = 1 (open).
cenc::uint_encode(enc, _local_id);
string_encode(enc, protocol);
cenc::buffer_encode(enc, id.data(), id.size());
if (handshake && handshake_len > 0) {
memcpy(enc.buffer + enc.start, handshake, handshake_len);
enc.start += handshake_len;
}
DHT_LOG("[Protomux] Channel::open('%s') SENT Open msg: local_id=%u, buf_size=%zu\n",
protocol.c_str(), _local_id, buf.size());
_mux->_write0(buf.data(), buf.size());
}
void ProtomuxChannel::close() {
if (closed) return;
// Encode Close control message: [0][3][localId]
EncodeState pre = {0, 2, nullptr};
cenc::uint_preencode(pre, _local_id);
auto buf = cenc::state_alloc(pre);
EncodeState enc = {0, pre.end, buf.data()};
enc.buffer[enc.start++] = 0; // remoteId = 0.
enc.buffer[enc.start++] = 3; // type = 3 (close).
cenc::uint_encode(enc, _local_id);
_close(false);
_mux->_write0(buf.data(), buf.size());
}
uint32_t ProtomuxChannel::add_message(ProtomuxMessage::RecvCallback cb) {
uint32_t type = (uint32_t)messages.size();
messages.push_back({type, std::move(cb)});
return type;
}
bool ProtomuxChannel::send(uint32_t type, const uint8_t *data, size_t len) {
if (!opened || closed) return false;
if (_mux->_corked > 0) {
// Batched: encode [type:varint][payload].
EncodeState pre = {0, 0, nullptr};
cenc::uint_preencode(pre, type);
pre.end += len;
auto buf = cenc::state_alloc(pre);
EncodeState enc = {0, pre.end, buf.data()};
cenc::uint_encode(enc, type);
if (len > 0) {
memcpy(enc.buffer + enc.start, data, len);
enc.start += len;
}
_mux->_push_batch(_local_id, buf.data(), buf.size());
// Corked: the frame is batched locally and goes out at uncork, so there is
// nothing to report about the socket yet.
return true;
}
// Non-batched: encode [_local_id:varint][type:varint][payload] and write directly.
// The receiver uses our _local_id to index their _remote[] array.
EncodeState pre = {0, 0, nullptr};
cenc::uint_preencode(pre, _local_id);
cenc::uint_preencode(pre, type);
pre.end += len;
auto buf = cenc::state_alloc(pre);
EncodeState enc = {0, pre.end, buf.data()};
cenc::uint_encode(enc, _local_id);
cenc::uint_encode(enc, type);
if (len > 0) {
memcpy(enc.buffer + enc.start, data, len);
enc.start += len;
}
if (_mux->_write_cb) {
// Forward the carrier's answer instead of asserting success. This is the
// only path by which a producer above learns the buffer is full.
return _mux->_write_cb(buf.data(), buf.size());
}
return true;
}
void ProtomuxChannel::cork() {
if (_mux) _mux->cork();
}
void ProtomuxChannel::uncork() {
if (_mux) _mux->uncork();
}
void ProtomuxChannel::_fully_open() {
DHT_LOG("[Protomux] Channel::_fully_open('%s') local_id=%u remote_id=%u opened=%d closed=%d\n",
protocol.c_str(), _local_id, _remote_id, opened, closed);
if (opened || closed) return;
opened = true;
// Deliver handshake to callback.
if (on_open) {
DHT_LOG("[Protomux] Channel::_fully_open('%s') firing on_open callback\n", protocol.c_str());
on_open(_remote_handshake.data(), _remote_handshake.size());
}
// Clear stored handshake.
_remote_handshake.clear();
_remote_handshake.shrink_to_fit();
// Set session on remote slot and drain any buffered messages.
if (_remote_id > 0 && (_remote_id - 1) < _mux->_remote.size()) {
auto *slot = _mux->_remote[_remote_id - 1];
if (slot) {
slot->session = this;
slot->handshake.clear();
if (slot->buffering) {
// Drain pending messages.
for (auto &pm : slot->pending) {
_mux->_buffered -= pm.data.size();
_recv(pm.type, pm.data.data(), pm.data.size());
}
slot->pending.clear();
slot->buffering = false;
}
}
}
}
void ProtomuxChannel::_close(bool is_remote) {
if (closed) return;
closed = true;
if (!_mux) return;
// Update protocol info.
auto *info = _mux->_get_info(protocol, id);
info->opened--;
if (info->last_channel == this) info->last_channel = nullptr;
// Free remote slot.
if (_remote_id > 0) {
uint32_t ridx = _remote_id - 1;
if (ridx < _mux->_remote.size()) {
delete _mux->_remote[ridx];
_mux->_remote[ridx] = nullptr;
}
_remote_id = 0;
// Free local slot for reuse.
_mux->_free.push_back(_local_id - 1);
}
// Null out local slot.
if (_local_id > 0 && (_local_id - 1) < _mux->_local.size()) {
_mux->_local[_local_id - 1] = nullptr;
}
_local_id = 0;
_mux->_gc(info);
if (on_close) on_close(is_remote);
_destroy();
}
void ProtomuxChannel::_destroy() {
if (destroyed) return;
destroyed = true;
if (on_destroy) on_destroy();
}
void ProtomuxChannel::_recv(uint32_t type, const uint8_t *data, size_t len) {
if (type < messages.size() && messages[type].on_message) {
messages[type].on_message(data, len, this);
}
}
// ============================================================================
// Protomux
// ============================================================================
Protomux::Protomux() {}
Protomux::~Protomux() {
// Clean up remote slots.
for (auto *slot : _remote) {
delete slot;
}
_remote.clear();
// Clean up protocol infos.
// Collect unique infos (aliases point to the same info).
std::vector<ProtocolInfo *> unique_infos;
for (auto &pair : _infos) {
if (std::find(unique_infos.begin(), unique_infos.end(), pair.second) == unique_infos.end()) {
unique_infos.push_back(pair.second);
}
}
for (auto *info : unique_infos) {
delete info;
}
_infos.clear();
// _owned holds unique_ptrs — clears automatically.
// _local has raw pointers (non-owning) — just clear.
_local.clear();
}
ProtomuxChannel *Protomux::create_channel(
const std::string &protocol,
const std::vector<uint8_t> &id,
const std::vector<std::string> &aliases,
bool unique) {
if (_shutdown_called) return nullptr;
auto *info = _get_info(protocol, id, aliases);
if (unique && info->opened > 0) return nullptr;
auto ch = std::make_unique<ProtomuxChannel>();
ch->protocol = protocol;
ch->aliases = aliases;
ch->id = id;
ch->_mux = this;
ProtomuxChannel *raw = ch.get();
if (info->incoming.empty()) {
// No pending remote Open — channel waits for remote.
_owned.push_back(std::move(ch));
return raw;
}
// Remote already sent Open for this protocol — pair immediately.
_remote_backlog--;
uint32_t remote_id = info->incoming.front();
info->incoming.erase(info->incoming.begin());
if (remote_id == 0 || (remote_id - 1) >= _remote.size()) {
_owned.push_back(std::move(ch));
return raw;
}
auto *slot = _remote[remote_id - 1];
if (!slot) {
_owned.push_back(std::move(ch));
return raw;
}
raw->_remote_id = remote_id;
raw->_remote_handshake = std::move(slot->handshake);
_owned.push_back(std::move(ch));
// Fully open immediately — _fully_open sets slot->session and drains pending.
raw->_fully_open();
return raw;
}
bool Protomux::opened(const std::string &protocol, const std::vector<uint8_t> &id) const {
std::string key = _to_key(protocol, id);
auto it = _infos.find(key);
return it != _infos.end() && it->second->opened > 0;
}
void Protomux::pair(const std::string &protocol, PairCallback cb) {
// Use protocol with empty id for wildcard pairing.
_notify[_to_key(protocol, {})] = std::move(cb);
}
void Protomux::unpair(const std::string &protocol) {
_notify.erase(_to_key(protocol, {}));
}
void Protomux::cork() {
if (++_corked == 1) {
_batch.clear();
_batch_size = 2; // Reserve for [0][0] batch header.
}
}
void Protomux::uncork() {
if (--_corked == 0) {
_send_batch();
_batch.clear();
_batch_size = 0;
}
}
void Protomux::recv(const uint8_t *data, size_t len) {
if (len == 0) return;
DHT_LOG("[Protomux] recv: %zu bytes, first_byte=0x%02x\n", len, data[0]);
// Decode the first varint: remoteId (= sender's localId).
EncodeState s = {0, len, const_cast<uint8_t *>(data)};
uint32_t remote_id = (uint32_t)cenc::uint_decode(s);
DHT_LOG("[Protomux] recv: decoded remote_id=%u, remaining=%zu bytes\n",
remote_id, s.end - s.start);
// Depth-counted, not a plain flag: a handler may write, and a write may be
// looped straight back into this mux. Only the outermost frame may free, and
// by then no frame below it is still holding a retired info.
_recv_depth++;
_decode(remote_id, s.buffer + s.start, s.end - s.start);
if (--_recv_depth == 0) _retired.clear();
}
void Protomux::set_drain_callback(DrainCallback cb) {
_drain_cb = std::move(cb);
}
void Protomux::notify_drain() {
if (_drain_cb) _drain_cb();
}
void Protomux::set_write_callback(WriteCallback cb) {
_write_cb = std::move(cb);
}
void Protomux::shutdown() {
_shutdown_called = true;
// Close all open channels.
for (auto *ch : _local) {
if (ch && !ch->closed) {
ch->_close(true);
}
}
}
// ---- Internal Decode ----
void Protomux::_decode(uint32_t remote_id, const uint8_t *data, size_t len) {
// Decode message type varint.
EncodeState s = {0, len, const_cast<uint8_t *>(data)};
uint32_t type = (uint32_t)cenc::uint_decode(s);
if (remote_id == 0) {
// Control message.
_on_control(type, s.buffer + s.start, s.end - s.start);
return;
}
// Data message — look up remote slot.
if (remote_id > _remote.size()) return;
auto *slot = _remote[remote_id - 1];
if (!slot) return;
const uint8_t *payload = s.buffer + s.start;
size_t payload_len = s.end - s.start;
if (slot->buffering) {
// Channel not fully open yet — buffer the message.
slot->pending.push_back({type, std::vector<uint8_t>(payload, payload + payload_len)});
_buffered += payload_len;
return;
}
if (slot->session) {
slot->session->_recv(type, payload, payload_len);
}
}
void Protomux::_on_control(uint32_t type, const uint8_t *data, size_t len) {
switch (type) {
case 0:
_on_batch(data, len);
break;
case 1:
_on_open(data, len);
break;
case 2:
_on_reject(data, len);
break;
case 3:
_on_close(data, len);
break;
default:
break;
}
}
void Protomux::_on_open(const uint8_t *data, size_t len) {
DHT_LOG("[Protomux] _on_open: %zu bytes\n", len);
EncodeState s = {0, len, const_cast<uint8_t *>(data)};
uint32_t remote_id = (uint32_t)cenc::uint_decode(s);
std::string protocol = string_decode(s);
BufView id_view = cenc::buffer_decode(s);
DHT_LOG("[Protomux] _on_open: remote_id=%u protocol='%s' id_len=%zu\n",
remote_id, protocol.c_str(), id_view.len);
// Remaining bytes = handshake data.
const uint8_t *hs_data = s.buffer + s.start;
size_t hs_len = (s.start < s.end) ? (s.end - s.start) : 0;
// Reject remoteId = 0 (control session can't be opened).
if (remote_id == 0) {
_reject_session(nullptr, 0);
return;
}
uint32_t ridx = remote_id - 1;
std::vector<uint8_t> id_vec;
if (id_view.data && id_view.len > 0) {
id_vec.assign(id_view.data, id_view.data + id_view.len);
}
auto *info = _get_info(protocol, id_vec);
// Grow _remote array if needed (allow +1).
if (_remote.size() == ridx) {
_remote.push_back(nullptr);
}
if (ridx >= _remote.size() || _remote[ridx] != nullptr) {
fprintf(stderr, "[Protomux] Invalid open message: remoteId=%u\n", remote_id);
return;
}
DHT_LOG("[Protomux] _on_open: info->outgoing.size()=%zu info->incoming.size()=%zu\n",
info->outgoing.size(), info->incoming.size());
if (!info->outgoing.empty()) {
// We already sent Open for this protocol — pair immediately.
uint32_t local_id = info->outgoing.front();
info->outgoing.erase(info->outgoing.begin());
DHT_LOG("[Protomux] _on_open: PAIRING with local_id=%u\n", local_id);
if (local_id == 0 || (local_id - 1) >= _local.size()) {
_free.push_back(local_id - 1);
return;
}
auto *session = _local[local_id - 1];
if (!session) {
// Already closed — free the slot.
_free.push_back(local_id - 1);
return;
}
// Create remote slot (no buffering — paired immediately).
auto *slot = new RemoteSlot();
if (hs_len > 0) {
slot->handshake.assign(hs_data, hs_data + hs_len);
}
slot->buffering = false;
_remote[ridx] = slot;
session->_remote_id = remote_id;
session->_remote_handshake = std::move(slot->handshake);
session->_fully_open();
return;
}
// No local channel yet — buffer and wait for createChannel().
auto *slot = new RemoteSlot();
if (hs_len > 0) {
slot->handshake.assign(hs_data, hs_data + hs_len);
}
slot->buffering = true; // Buffer incoming data messages.
_remote[ridx] = slot;
_remote_backlog++;
info->incoming.push_back(remote_id);
// Notify pair callbacks.
_request_session(protocol, id_vec, info);
}
void Protomux::_on_reject(const uint8_t *data, size_t len) {
EncodeState s = {0, len, const_cast<uint8_t *>(data)};
uint32_t local_id = (uint32_t)cenc::uint_decode(s);
// Find which protocol info has this localId in outgoing.
for (auto &pair : _infos) {
auto *info = pair.second;
auto it = std::find(info->outgoing.begin(), info->outgoing.end(), local_id);
if (it == info->outgoing.end()) continue;
info->outgoing.erase(it);
if (local_id > 0 && (local_id - 1) < _local.size()) {
auto *session = _local[local_id - 1];
_free.push_back(local_id - 1);
if (session) session->_close(true);
}
_gc(info);
return;
}
fprintf(stderr, "[Protomux] Invalid reject message: localId=%u\n", local_id);
}
void Protomux::_on_close(const uint8_t *data, size_t len) {
EncodeState s = {0, len, const_cast<uint8_t *>(data)};
uint32_t remote_id = (uint32_t)cenc::uint_decode(s);
if (remote_id == 0) return;
uint32_t ridx = remote_id - 1;
if (ridx >= _remote.size()) return;
auto *slot = _remote[ridx];
if (!slot) return;
if (slot->session) {
slot->session->_close(true);
}
}
void Protomux::_on_batch(const uint8_t *data, size_t len) {
EncodeState s = {0, len, const_cast<uint8_t *>(data)};
uint32_t remote_id = (uint32_t)cenc::uint_decode(s);
// Cork so replies are batched.
cork();
while (s.start < s.end) {
uint64_t msg_len = cenc::uint_decode(s);
if (msg_len == 0) {
// Channel switch sentinel.
remote_id = (uint32_t)cenc::uint_decode(s);
continue;
}
if (s.start + msg_len > s.end) break; // Truncated.
_decode(remote_id, s.buffer + s.start, (size_t)msg_len);
s.start += (size_t)msg_len;
}
uncork();
}
// ---- Internal Encode ----
void Protomux::_write0(const uint8_t *data, size_t len) {
if (len < 1) return;
if (_corked > 0) {
// Strip the leading [0] (remoteId=0) and push as control (localId=0).
_push_batch(0, data + 1, len - 1);
return;
}
if (_write_cb) {
_write_cb(data, len);
}
}
void Protomux::_push_batch(uint32_t local_id, const uint8_t *data, size_t len) {
// Check if batch exceeded MAX_BATCH.
if (_batch_size >= MAX_BATCH) {
_send_batch();
_batch.clear();
_batch_size = 2; // Reset for new batch header.
}
// Track channel-switch overhead.
if (_batch.empty() || _batch.back().local_id != local_id) {
_batch_size += 1; // Sentinel byte (or initial ID slot).
EncodeState pre = {0, 0, nullptr};
cenc::uint_preencode(pre, local_id);
_batch_size += pre.end;
}
// Track buffer encoding overhead: [len:varint][data].
EncodeState pre = {0, 0, nullptr};
cenc::buffer_preencode(pre, data, len);
_batch_size += pre.end;
_batch.push_back({local_id, std::vector<uint8_t>(data, data + len)});
}
void Protomux::_send_batch() {
if (_batch.empty()) return;
// Pre-encode to calculate buffer size.
EncodeState pre = {0, 2, nullptr}; // [0][0] header.
uint32_t prev_id = _batch[0].local_id;
cenc::uint_preencode(pre, prev_id);
for (size_t i = 0; i < _batch.size(); i++) {
auto &entry = _batch[i];
if (entry.local_id != prev_id) {
pre.end += 1; // Sentinel byte (varint 0).
cenc::uint_preencode(pre, entry.local_id);
prev_id = entry.local_id;
}
// Length-prefixed buffer: [len:varint][data].
cenc::buffer_preencode(pre, entry.data.data(), entry.data.size());
}
// Allocate and encode.
auto buf = cenc::state_alloc(pre);
EncodeState enc = {0, pre.end, buf.data()};
enc.buffer[enc.start++] = 0; // remoteId = 0 (control).
enc.buffer[enc.start++] = 0; // type = 0 (batch).
prev_id = _batch[0].local_id;
cenc::uint_encode(enc, prev_id);
for (size_t i = 0; i < _batch.size(); i++) {
auto &entry = _batch[i];
if (entry.local_id != prev_id) {
enc.buffer[enc.start++] = 0; // Sentinel.
cenc::uint_encode(enc, entry.local_id);
prev_id = entry.local_id;
}
cenc::buffer_encode(enc, entry.data.data(), entry.data.size());
}
if (_write_cb) {
_write_cb(buf.data(), buf.size());
}
}
// ---- Helpers ----
Protomux::ProtocolInfo *Protomux::_get_info(
const std::string &protocol,
const std::vector<uint8_t> &id,
const std::vector<std::string> &aliases) {
std::string key = _to_key(protocol, id);
auto it = _infos.find(key);
if (it != _infos.end()) return it->second;
auto *info = new ProtocolInfo();
info->key = key;
info->protocol = protocol;
info->id = id;
_infos[key] = info;
// Register aliases.
for (auto &alias : aliases) {
std::string alias_key = _to_key(alias, id);
info->alias_keys.push_back(alias_key);
_infos[alias_key] = info;
}
return info;
}
void Protomux::_gc(ProtocolInfo *info) {
// Idempotent, because JS's is and both of its callers rely on it: a reject
// runs Session#_close (which gcs) and then gcs again itself. Freeing here
// instead made the second call read and delete a dead object.
if (info->retired) return;
if (info->opened == 0 && info->outgoing.empty() && info->incoming.empty()) {
_infos.erase(info->key);
for (auto &alias : info->alias_keys) {
_infos.erase(alias);
}
// Out of the registry, so nothing can find it again; freed by recv() once
// the dispatch that retired it has returned.
info->retired = true;
_retired.emplace_back(info);
}
}
void Protomux::_reject_session(ProtocolInfo *info, uint32_t remote_id) {
// Clean up remote slot.
if (remote_id > 0 && (remote_id - 1) < _remote.size()) {
auto *slot = _remote[remote_id - 1];
if (slot) {
if (slot->buffering) {
for (auto &pm : slot->pending) {
_buffered -= pm.data.size();
}
}
delete slot;
_remote[remote_id - 1] = nullptr;
}
}
// Send Reject: [0][2][remoteId].
EncodeState pre = {0, 2, nullptr};
cenc::uint_preencode(pre, remote_id);
auto buf = cenc::state_alloc(pre);
EncodeState enc = {0, pre.end, buf.data()};
enc.buffer[enc.start++] = 0; // remoteId = 0.
enc.buffer[enc.start++] = 2; // type = 2 (reject).
cenc::uint_encode(enc, remote_id);
_write0(buf.data(), buf.size());
}
void Protomux::_request_session(const std::string &protocol,
const std::vector<uint8_t> &id,
ProtocolInfo *info) {
// Check for pair callback (exact match or wildcard).
std::string key = _to_key(protocol, id);
auto it = _notify.find(key);
if (it == _notify.end()) {
// Try wildcard (protocol with empty id).
key = _to_key(protocol, {});
it = _notify.find(key);
}
if (it != _notify.end()) {
it->second(protocol, id);
return;
}
// No pair callback — reject all pending incoming for this info.
while (!info->incoming.empty()) {
uint32_t rid = info->incoming.front();
info->incoming.erase(info->incoming.begin());
_reject_session(info, rid);
}
_gc(info);
}
std::string Protomux::_to_key(const std::string &protocol, const std::vector<uint8_t> &id) {
return protocol + "##" + to_hex(id);
}