utun2/tools/chatgui/transport/chat_propagator.cpp

#include "chat_propagator.h"
#include "chat_proto.h"
#include "utun_node.h"
#include "crypto.h"
#include "db_manager.h"

#include <QDateTime>
#include <algorithm>
#include <random>

static uint64_t read64(const uint8_t* p) {
    uint64_t v; memcpy(&v, p, 8); return v;
}
static void write64(uint8_t* p, uint64_t v) { memcpy(p, &v, 8); }
static uint32_t read32(const uint8_t* p) {
    uint32_t v; memcpy(&v, p, 4); return v;
}
static void write32(uint8_t* p, uint32_t v) { memcpy(p, &v, 4); }
static uint16_t read16(const uint8_t* p) {
    uint16_t v; memcpy(&v, p, 2); return v;
}
static void write16(uint8_t* p, uint16_t v) { memcpy(p, &v, 2); }

ChatPropagator::ChatPropagator(UtunNode* node, DbManager* db, QObject* parent)
    : QObject(parent)
    , m_node(node)
    , m_db(db)
    , m_ackTimer(new QTimer(this))
{
    connect(m_node, &UtunNode::messageReceived,
            this, &ChatPropagator::onReceived);
    connect(m_ackTimer, &QTimer::timeout,
            this, &ChatPropagator::onAckTimeout);
}

void ChatPropagator::setKeys(const QByteArray& edPriv, const QByteArray& edPub, quint64 nodeId) {
    m_edPriv   = edPriv;
    m_edPub    = edPub;
    m_myNodeId = nodeId;
}

void ChatPropagator::start() {
    if (!m_edPriv.isEmpty() && !m_ackTimer->isActive())
        m_ackTimer->start(250);
    startCatchUp();
}

/* ======================================================================== */

void ChatPropagator::propagateNewMessage(const QString& channelId,
                                          const QByteArray& content,
                                          qint64 timestamp,
                                          const QByteArray& signature) {
    if (m_edPriv.isEmpty() || !m_myNodeId) return;

    QByteArray packet = buildNewMsgPacket(channelId, content, timestamp, signature);
    QList<quint64> exclude;
    exclude.append(m_myNodeId);
    fanoutGossip(channelId, packet, signature, ChatProp::kDefaultFanout, exclude);
}

QByteArray ChatPropagator::buildNewMsgPacket(const QString& channelId,
                                               const QByteArray& content,
                                               qint64 timestamp,
                                               const QByteArray& signature) {
    QByteArray cid    = channelId.toUtf8();
    uint8_t    cidLen = static_cast<uint8_t>(cid.size());
    size_t     total  = 1 + 1 + cidLen + 8 + 2 + content.size() + ChatProp::kSignSize;

    QByteArray pkt(static_cast<int>(total), Qt::Uninitialized);
    uint8_t* p = reinterpret_cast<uint8_t*>(pkt.data());
    *p++ = ChatProp::kNewMsg;
    *p++ = cidLen;
    memcpy(p, cid.constData(), cidLen);  p += cidLen;
    write64(p, static_cast<uint64_t>(timestamp)); p += 8;
    write16(p, static_cast<uint16_t>(content.size())); p += 2;
    memcpy(p, content.constData(), content.size());  p += content.size();
    memcpy(p, signature.constData(), ChatProp::kSignSize);
    return pkt;
}

QByteArray ChatPropagator::buildAckPacket(const QByteArray& sig) {
    QByteArray pkt(static_cast<int>(1 + ChatProp::kSignSize), Qt::Uninitialized);
    pkt[0] = static_cast<char>(ChatProp::kNewAck);
    memcpy(pkt.data() + 1, sig.constData(), ChatProp::kSignSize);
    return pkt;
}

/* ======================================================================== */

void ChatPropagator::fanoutGossip(const QString& channelId, const QByteArray& packet,
                                   const QByteArray& sig, int fanout,
                                   const QList<quint64>& exclude) {
    auto peers = pickRandomPeers(fanout, exclude);
    if (peers.isEmpty()) return;

    GossipEntry e;
    e.channelId = channelId;
    e.packet    = packet;
    e.sig       = sig;
    e.peers     = peers;
    e.sentAt    = QDateTime::currentMSecsSinceEpoch();
    m_pending[sig] = e;

    for (auto peer : peers)
        sendGossip(peer, packet);
}

void ChatPropagator::sendGossip(quint64 peer, const QByteArray& packet) {
    m_node->send(peer, packet);
}

/* ======================================================================== */

void ChatPropagator::onReceived(quint64 peerId, const QByteArray& data) {
    if (data.isEmpty()) return;
    uint8_t type = static_cast<uint8_t>(data[0]);
    switch (type) {
    case ChatProp::kNewMsg: handleNewMsg(peerId, data); break;
    case ChatProp::kNewAck: handleNewAck(data);         break;
    case ChatProp::kBucketsRsp: handleBucketsRsp(data); break;
    case ChatProp::kHashesRsp:  handleHashesRsp(data);  break;
    case ChatProp::kMsgsRsp:    handleMsgsRsp(data);    break;
    }
}

void ChatPropagator::handleNewMsg(quint64 fromPeer, const QByteArray& payload) {
    QString    channelId;
    QByteArray content;
    qint64     timestamp;
    QByteArray sig;
    if (!parseNewMsg(payload, channelId, content, timestamp, sig))
        return;

    /* verify sender's signature */
    QByteArray peerEdPub = m_db->getNodeEdPub(fromPeer);
    if (!peerEdPub.isEmpty() && !Crypto::verify(peerEdPub, payload.left(payload.size() - 64), sig))
        return;

    /* insert into DB */
    m_db->insertMessage(channelId, fromPeer, content, timestamp, 0);

    emit messageReceived(channelId, fromPeer, content, timestamp);

    /* re-gossip */
    QList<quint64> exclude;
    exclude.append(m_myNodeId);
    exclude.append(fromPeer);
    fanoutGossip(channelId, payload, sig, ChatProp::kDefaultFanout, exclude);

    /* ack to sender */
    m_node->send(fromPeer, buildAckPacket(sig));
}

void ChatPropagator::handleNewAck(const QByteArray& payload) {
    if (payload.size() < static_cast<int>(1 + ChatProp::kSignSize)) return;
    QByteArray sig = payload.mid(1, static_cast<int>(ChatProp::kSignSize));

    auto it = m_pending.find(sig);
    if (it == m_pending.end()) return;

    /* mark all peers as acked (simple version: ACK means message delivered) */
    m_pending.erase(it);
}

void ChatPropagator::onAckTimeout() {
    qint64 now = QDateTime::currentMSecsSinceEpoch();
    QList<QByteArray> toRemove;

    for (auto it = m_pending.begin(); it != m_pending.end(); ++it) {
        auto& e = it.value();
        if (now - e.sentAt < ChatProp::kDefaultAckTimeoutMs) continue;

        if (e.retries >= 3) {
            toRemove.append(it.key());
            continue;
        }

        /* retry: pick new peers for each missing ACK */
        QList<quint64> unacked;
        for (auto p : e.peers)
            if (!e.acked.contains(p)) unacked.append(p);

        if (!unacked.isEmpty()) {
            auto newPeers = pickRandomPeers(1, e.peers + e.acked +
                                            QList<quint64>{m_myNodeId});
            if (!newPeers.isEmpty()) {
                sendGossip(newPeers.first(), e.packet);
                e.peers.append(newPeers.first());
            }
        }

        e.sentAt  = now;
        e.retries++;
    }

    for (auto& k : toRemove) m_pending.remove(k);
}

/* ========================================================================
 * Catch-up
 * ======================================================================== */

void ChatPropagator::startCatchUp() {
    if (!m_db) return;
    auto channels = m_db->getChannels();
    if (channels.isEmpty()) return;

    auto peers = pickRandomPeers(ChatProp::kDefaultFanout, {m_myNodeId});
    if (peers.isEmpty()) return;

    for (auto& ch : channels) {
        QByteArray cid = ch.channelId.toUtf8();
        QByteArray pkt(static_cast<int>(2 + cid.size()), Qt::Uninitialized);
        pkt[0] = static_cast<char>(ChatProp::kBucketsReq);
        pkt[1] = static_cast<char>(cid.size());
        memcpy(pkt.data() + 2, cid.constData(), cid.size());

        for (auto peer : peers)
            sendGossip(peer, pkt);
    }
}

void ChatPropagator::handleBucketsRsp(const QByteArray& payload) {
    if (payload.size() < 3) return;
    const uint8_t* p = reinterpret_cast<const uint8_t*>(payload.constData());
    if (*p != ChatProp::kBucketsRsp) return; p++;
    uint8_t cidLen = *p++;
    if (p + cidLen + 2 > reinterpret_cast<const uint8_t*>(payload.constData()) + payload.size()) return;
    QString channelId = QString::fromUtf8(reinterpret_cast<const char*>(p), cidLen); p += cidLen;
    uint16_t bCount = read16(p); p += 2;

    for (int i = 0; i < bCount; i++) {
        if (p + 4 + 1 + 4 + 32 > reinterpret_cast<const uint8_t*>(payload.constData()) + payload.size()) break;
        qint64 day   = static_cast<qint64>(read32(p)); p += 4;
        int    hour  = *p++;
        int    rCount = static_cast<int>(read32(p)); p += 4;
        QByteArray rHash(reinterpret_cast<const char*>(p), 32); p += 32;

        auto local = m_db->getBucketMeta(channelId);
        bool found = false;
        for (auto& b : local) {
            if (b.day == day && b.hour == hour && b.count == rCount && b.hash == rHash) {
                found = true; break;
            }
        }
        if (found) continue;

        /* mismatch — request signatures */
        QByteArray cid = channelId.toUtf8();
        QByteArray req(static_cast<int>(7 + cid.size()), Qt::Uninitialized);
        uint8_t* rp = reinterpret_cast<uint8_t*>(req.data());
        *rp++ = ChatProp::kHashesReq;
        *rp++ = static_cast<uint8_t>(cid.size());
        memcpy(rp, cid.constData(), cid.size()); rp += cid.size();
        write32(rp, static_cast<uint32_t>(day)); rp += 4;
        *rp = static_cast<uint8_t>(hour);

        auto peers = pickRandomPeers(1, {m_myNodeId});
        if (!peers.isEmpty())
            sendGossip(peers.first(), req);
    }
}

void ChatPropagator::handleHashesRsp(const QByteArray& payload) {
    if (payload.size() < 3) return;
    const uint8_t* p = reinterpret_cast<const uint8_t*>(payload.constData());
    if (*p != ChatProp::kHashesRsp) return; p++;
    uint16_t sigCount = read16(p); p += 2;

    QList<QByteArray> remoteSigs;
    for (int i = 0; i < sigCount; i++) {
        if (p + 64 > reinterpret_cast<const uint8_t*>(payload.constData()) + payload.size()) break;
        remoteSigs.append(QByteArray(reinterpret_cast<const char*>(p), 64));
        p += 64;
    }
    if (remoteSigs.isEmpty()) return;

    /* find missing */
    QList<QByteArray> missing;
    QList<QByteArray> remote = remoteSigs;
    for (auto& s : remote) {
        /* TODO: efficient lookup — query DB for this sig */
        auto msgs = m_db->getMessagesBySigs({s});
        if (msgs.isEmpty()) missing.append(s);
    }
    if (missing.isEmpty()) return;

    /* request missing messages */
    int mCount = missing.size();
    QByteArray req(static_cast<int>(3 + mCount * 64), Qt::Uninitialized);
    uint8_t* rp = reinterpret_cast<uint8_t*>(req.data());
    *rp++ = ChatProp::kMsgsReq;
    write16(rp, static_cast<uint16_t>(mCount)); rp += 2;
    for (auto& s : missing) { memcpy(rp, s.constData(), 64); rp += 64; }

    auto peers = pickRandomPeers(1, {m_myNodeId});
    if (!peers.isEmpty())
        sendGossip(peers.first(), req);
}

void ChatPropagator::handleMsgsRsp(const QByteArray& payload) {
    if (payload.size() < 4) return;
    const uint8_t* p  = reinterpret_cast<const uint8_t*>(payload.constData());
    const uint8_t* end = p + payload.size();
    if (*p != ChatProp::kMsgsRsp) return; p++;
    uint8_t cidLen = *p++;
    if (p + cidLen + 2 > end) return;
    QString channelId = QString::fromUtf8(reinterpret_cast<const char*>(p), cidLen); p += cidLen;
    uint16_t msgCount = read16(p); p += 2;

    for (int i = 0; i < msgCount; i++) {
        if (p + 64 + 8 + 8 + 2 > end) break;
        QByteArray sig(reinterpret_cast<const char*>(p), 64); p += 64;
        quint64 author = read64(p); p += 8;
        qint64 ts = static_cast<qint64>(read64(p)); p += 8;
        uint16_t cLen = read16(p); p += 2;
        if (p + cLen > end) break;
        QByteArray content(reinterpret_cast<const char*>(p), cLen); p += cLen;

        auto existing = m_db->getMessagesBySigs({sig});
        if (!existing.isEmpty()) continue;

        m_db->insertMessage(channelId, author, content, ts, 0);
        emit messageReceived(channelId, author, content, ts);
    }
}

/* ======================================================================== */

bool ChatPropagator::parseNewMsg(const QByteArray& payload, QString& channelId,
                                  QByteArray& content, qint64& timestamp,
                                  QByteArray& sig) {
    if (payload.size() < 14) return false;  /* minimum: 1+1+1+8+2+1 */

    const uint8_t* p = reinterpret_cast<const uint8_t*>(payload.constData());
    if (*p != ChatProp::kNewMsg) return false;
    p++;

    uint8_t cidLen = *p++;
    if (p + cidLen + 8 + 2 > reinterpret_cast<const uint8_t*>(payload.constData()) + payload.size())
        return false;
    channelId = QString::fromUtf8(reinterpret_cast<const char*>(p), cidLen); p += cidLen;

    timestamp = static_cast<qint64>(read64(p)); p += 8;

    uint16_t contentLen = read16(p); p += 2;
    if (p + contentLen + ChatProp::kSignSize >
        reinterpret_cast<const uint8_t*>(payload.constData()) + payload.size())
        return false;

    content = QByteArray(reinterpret_cast<const char*>(p), contentLen); p += contentLen;
    sig     = QByteArray(reinterpret_cast<const char*>(p), static_cast<int>(ChatProp::kSignSize));
    return true;
}

QList<quint64> ChatPropagator::pickRandomPeers(int count, const QList<quint64>& exclude) {
    QList<quint64> result;
    if (!m_db) return result;

    auto accs = m_db->getAccounts(true);
    QList<quint64> pool;
    for (auto& a : accs) {
        if (!exclude.contains(a.nodeId))
            pool.append(a.nodeId);
    }
    if (pool.isEmpty()) return result;

    std::random_device rd;
    std::mt19937 g(rd());
    std::shuffle(pool.begin(), pool.end(), g);

    for (int i = 0; i < count && i < pool.size(); i++)
        result.append(pool[i]);
    return result;
}