query.cpp

#include "query.h"
#include "worker.h"
#include "mainwindow.h"
#include <QDateTime>

const qint32 k_nAppIDTheShip = 2400;
const qint32 k_nAppIDKillingFloor = 1250;
const qint32 k_nAppIDKillingFloor2 = 232090;

int g_queryMaxRetries = QUERY_MAX_RETRIES_DEFAULT;

QString GetStringFromStream(QDataStream &stream)
{
    qint64 pos = stream.device()->pos();
    qint64 bytes = stream.device()->bytesAvailable();

    do
    {
        qint8 byte;
        stream >> byte;

        if(byte == 0x00)
        {
            qint64 size = bytes - stream.device()->bytesAvailable();

            char *ret = new char[size];

            stream.device()->reset();
            stream.skipRawData(pos);

            stream.readRawData(ret, size);

            QString res = QString(ret);
            delete [] ret;
            return res;
        }

    }while(stream.device()->bytesAvailable() > 0);

    stream.device()->reset();
    stream.skipRawData(pos);

    return QString();
}

QString GetRichUEStringFromStream(QDataStream &stream)
{
    qint64 pos = stream.device()->pos();
    qint64 bytes = stream.device()->bytesAvailable();

    do
    {
        qint8 byte;
        stream >> byte;

        if(byte == 0x00)
        {
            qint64 size = bytes - stream.device()->bytesAvailable();

            char *ret = new char[size];

            stream.device()->reset();
            stream.skipRawData(pos);

            stream.readRawData(ret, size);

            std::string res;
            bool inColor = false;
            for (size_t i = 0; i < size; ++i)
            {
                if (ret[i] == 0x1B && ((i + 3) < size))
                {
                    if (inColor)
                    {
                        res += "</font>";
                    }

                    unsigned char r = ret[i+1];
                    unsigned char g = ret[i+2];
                    unsigned char b = ret[i+3];
                    res += QString("<font color=\"%1\">").arg(
                        QColor::fromRgb(r, g, b).name()
                        ).toStdString();

                    i += 3;
                }
                else
                {
                    res += ret[i];
                }
            }

            if (inColor)
            {
                res += "</font>";
            }

            delete [] ret;

            return QString::fromLatin1(res.c_str());
        }

    }while(stream.device()->bytesAvailable() > 0);

    stream.device()->reset();
    stream.skipRawData(pos);

    return QString();
}

QByteArray SendUDPQuery(QByteArray query, QHostAddress host, quint16 port, bool goldsrcSplits = false)
{
    qint8 tryCount = 0;

    //Use a while loop for rust and other servers that use gamserverport+1 for query
    do
    {
        QUdpSocket socket;
        socket.connectToHost(host, port+tryCount);

        if(socket.isValid())
        {
            int retries = 0;
            while(retries <= g_queryMaxRetries)
            {
            if(socket.write(query) != -1 && socket.waitForReadyRead(QUERY_TIMEOUT))
            {
                QByteArray reply;
                reply.resize(socket.pendingDatagramSize());
                socket.readDatagram(reply.data(), reply.size());

                QDataStream response(reply);
                response.setByteOrder(QDataStream::LittleEndian);

                qint32 header;
                response >> header;

                // Handle challenge-response on the SAME socket
                // Challenge is tied to source IP:port, so we must reuse the socket
                if(header == -1)
                {
                    qint8 check;
                    response >> check;

                    if(check == 0x41 && reply.size() >= 9) // Challenge response
                    {
                        qint32 challenge;
                        response >> challenge;

                        // Build challenge query: replace trailing -1 with challenge,
                        // or append challenge if query doesn't end with -1
                        QByteArray challengeQuery = query;
                        if(challengeQuery.size() >= 4)
                        {
                            // Check if last 4 bytes are -1 (0xFFFFFFFF) = player/rules style
                            QDataStream tailCheck(challengeQuery.right(4));
                            tailCheck.setByteOrder(QDataStream::LittleEndian);
                            qint32 tail;
                            tailCheck >> tail;
                            if(tail == -1)
                            {
                                // Replace the -1 with the challenge value
                                challengeQuery.chop(4);
                            }
                        }
                        QDataStream challengeData(&challengeQuery, QIODevice::WriteOnly | QIODevice::Append);
                        challengeData.setByteOrder(QDataStream::LittleEndian);
                        challengeData << challenge;

                        if(socket.write(challengeQuery) != -1 && socket.waitForReadyRead(QUERY_TIMEOUT))
                        {
                            reply.resize(socket.pendingDatagramSize());
                            socket.readDatagram(reply.data(), reply.size());

                            QDataStream response2(reply);
                            response2.setByteOrder(QDataStream::LittleEndian);
                            response2 >> header;
                        }
                        else
                        {
                            retries++;
                            continue;
                        }
                    }
                    else
                    {
                        // Not a challenge, rewind and return the full reply
                    }

                    socket.close();
                    return reply;
                }
                else if(header == -2)
                {
                    //We have multiple split packets
                    qint32 id;
                    qint8 skip = 0;
                    response >> id;

                    if(goldsrcSplits)
                    {
                        //GoldSrc split format: single byte encodes (packetNum << 4) | numPackets
                        quint8 packetByte;
                        response >> packetByte;
                        qint8 total = packetByte & 0x0F;
                        qint8 packetNum = (packetByte & 0xF0) >> 4;

                        reply.clear();
                        reply.append(response.device()->read(response.device()->size()));

                        while(packetNum < total-1)
                        {
                            if(!socket.isValid())
                                break;
                            if(socket.write(query) == -1 || !socket.waitForReadyRead(QUERY_TIMEOUT))
                                break;

                            QByteArray replyTemp;
                            replyTemp.resize(socket.pendingDatagramSize());
                            socket.readDatagram(replyTemp.data(), replyTemp.size());

                            QDataStream tempStream(replyTemp);
                            tempStream.setByteOrder(QDataStream::LittleEndian);

                            tempStream >> header;

                            qint32 tempId;
                            tempStream >> tempId;

                            if(header != -2 || id != tempId)
                                break;

                            tempStream >> packetByte;
                            packetNum = (packetByte & 0xF0) >> 4;

                            reply.append(tempStream.device()->read(tempStream.device()->size()));
                        }

                        socket.close();

                        //Strip the inner header (0xFFFFFFFF) from assembled payload
                        if(reply.size() >= 4)
                        {
                            QDataStream assembled(reply);
                            assembled.setByteOrder(QDataStream::LittleEndian);
                            assembled >> header;

                            if(header == -1)
                            {
                                QByteArray result;
                                result.append(reply);
                                return result;
                            }
                        }
                        return reply;
                    }

                    bool compressed;
                    response.device()->seek(response.device()->pos()-1);

                    response >> compressed;

                    // TODO:Support compressed packs
                    if(!compressed)
                    {
                        qint8 total;
                        response >> total;
                        qint8 packetNum;
                        response >> packetNum;

                        response >> header;

                        if(header != -1)
                        {
                            response.device()->seek(response.device()->pos()-sizeof(qint32)+sizeof(qint16));
                            skip = sizeof(qint16);
                            response >> header;
                        }

                        if(header == -1)
                        {
                            reply.clear();

                            response.device()->seek(response.device()->pos()-sizeof(qint32));

                            reply.append(response.device()->read(response.device()->size()));

                            do
                            {
                                if(socket.isValid())
                                {
                                    if(socket.write(query) != -1 && socket.waitForReadyRead(QUERY_TIMEOUT))
                                    {
                                        QByteArray replyTemp;
                                        replyTemp.resize(socket.pendingDatagramSize());
                                        socket.readDatagram(replyTemp.data(), replyTemp.size());

                                        QDataStream tempStream(replyTemp);
                                        tempStream.setByteOrder(QDataStream::LittleEndian);

                                        tempStream >> header;

                                        qint32 tempId;
                                        tempStream >> tempId;

                                        if(header != -2 || id != tempId)
                                            break;

                                        tempStream.skipRawData(sizeof(qint8));

                                        tempStream >> packetNum;
                                        tempStream.skipRawData(skip);

                                        reply.append(tempStream.device()->read(tempStream.device()->size()));
                                    }
                                }

                            }while(packetNum != total-1);

                            socket.close();
                            return reply;
                        }
                    }
                }
            }
            retries++;
            } // retry loop
            socket.close();
        }

        tryCount ++;

    }while(tryCount != 2);

    return QByteArray();
}

InfoReply::InfoReply(QByteArray response, qint64 ping)
{
    this->success = false;
    this->appId = -1;
    this->rawServerId = 0;
    this->goldsrc = false;
    this->ping = ping;

    if(response.size() > 5)
    {
        QDataStream data(response);

        data.setFloatingPointPrecision(QDataStream::SinglePrecision);
        data.setByteOrder(QDataStream::LittleEndian);

        qint32 header;
        qint8 check;

        data >> header;
        data >> check;

        if(header == -1 && check == A2S_INFO_CHECK)
        {
            this->success = true;
            data >> this->protocol;
            qint64 hostnamePos = data.device()->pos();
            this->hostnameRich = GetStringFromStream(data);
            this->map = GetStringFromStream(data);
            this->mod = GetStringFromStream(data);
            this->gamedesc = GetStringFromStream(data);
            qint16 id;
            data >> id;
            this->appId = id;
            data >> this->players;
            data >> this->maxplayers;
            data >> this->bots;
            data.device()->getChar(&(this->type));
            data.device()->getChar(&(this->os));
            data >> this->visibility;
            data >> this->vac;

            //The ship stuff...
            if(this->appId == k_nAppIDTheShip)
            {
                data.skipRawData(sizeof(qint8)*3);
            }

            this->version = GetStringFromStream(data);//Version

            qint8 edf;
            data >> edf;

            if(edf & 0x80)
            {
                data.skipRawData(sizeof(qint16));
            }
            if(edf & 0x10)
            {
                data >> this->rawServerId;

                quint32 accountID = (this->rawServerId & 0xFFFFFFFF);
                quint64 accountInst = (this->rawServerId >> 32) & 0xFFFFF;
                quint64 accountType = (this->rawServerId >> 52) & 0xF;
                quint8 accountUni = (this->rawServerId >> 56) & 0xFF;

                if(accountType == 4 || accountType == 3)
                {
                    if(accountType == 4)
                    {
                        this->serverID = QString("[A:%1:%2:%3]").arg(QString::number(accountUni), QString::number(accountID), QString::number(accountInst));
                    }
                    else
                    {
                        this->serverID = QString("[G:%1:%2]").arg(QString::number(accountUni), QString::number(accountID));
                    }
                }

            }
            if(edf & 0x40)
            {
                data.skipRawData(sizeof(qint16));
                GetStringFromStream(data);
            }
            if(edf & 0x20)
            {
                this->tags = GetStringFromStream(data);
            }
            if(edf & 0x01)
            {
                qint64 temp;
                data >> temp;

                this->appId = temp & ((1 << 24) - 1);
            }

            //GoldSrc servers may respond with Source format but protocol 48
            if(this->protocol == 48)
            {
                this->goldsrc = true;
            }

            // TODO: move these out to a config and add more games.
            // Unreal Engine uses Extended ASCII instead of UTF8, and also supports colors in hostname.
            if (this->appId == k_nAppIDKillingFloor || this->appId == k_nAppIDKillingFloor2)
            {
                data.device()->seek(hostnamePos);
                this->hostnameRich = GetRichUEStringFromStream(data);
            }
        }
        else if(header == -1 && check == A2S_INFO_GOLDSRC_CHECK)
        {
            this->success = true;
            this->goldsrc = true;

            GetStringFromStream(data);//Address (ip:port), not used

            this->hostnameRich = GetStringFromStream(data);
            this->map = GetStringFromStream(data);
            this->mod = GetStringFromStream(data);
            this->gamedesc = GetStringFromStream(data);

            data >> this->players;
            data >> this->maxplayers;
            data >> this->protocol;

            data.device()->getChar(&(this->type));
            data.device()->getChar(&(this->os));

            data >> this->visibility;

            qint8 isMod;
            data >> isMod;

            if(isMod)
            {
                GetStringFromStream(data);//Mod URL info
                GetStringFromStream(data);//Mod URL download
                data.skipRawData(sizeof(qint8));//Null byte
                data.skipRawData(sizeof(qint32));//Mod version
                data.skipRawData(sizeof(qint32));//Mod size
                data.skipRawData(sizeof(qint8));//Server-side only
                data.skipRawData(sizeof(qint8));//Custom client DLL
            }

            data >> this->vac;
            data >> this->bots;

            // GoldSrc A2S_INFO doesn't include AppID, map mod folder to AppID
            static const QMap<QString, int> goldsrcModMap = {
                {"cstrike", 10},       // Counter-Strike 1.6
                {"czero", 80},         // Counter-Strike: Condition Zero
                {"dod", 30},           // Day of Defeat
                {"dmc", 40},           // Deathmatch Classic
                {"tfc", 20},           // Team Fortress Classic
                {"valve", 70},         // Half-Life
                {"gearbox", 50},       // Opposing Force
                {"ricochet", 60},      // Ricochet
                {"bshift", 130},       // Blue Shift
                {"svencoop", 225840},  // Sven Co-op
            };
            if (goldsrcModMap.contains(this->mod))
                this->appId = goldsrcModMap.value(this->mod);
        }
    }
}

InfoQuery::InfoQuery(MainWindow *main)
{
    Worker *worker = new Worker;
    worker->moveToThread(&workerThread);
    connect(&workerThread, &QThread::finished, worker, &Worker::deleteLater);
    connect(&workerThread, &QThread::finished, this, &InfoQuery::deleteLater);
    connect(this, &InfoQuery::query, worker, &Worker::getServerInfo);
    connect(worker, &Worker::serverInfoReady, main, &MainWindow::ServerInfoReady);

    workerThread.start();
}

InfoReply *GetInfoReply(QHostAddress host, quint16 port)
{
    QByteArray query;
    QDataStream data(&query, QIODevice::ReadWrite);
    data.setByteOrder(QDataStream::LittleEndian);

    data << A2S_HEADER << (qint8)A2S_INFO;
    data.writeRawData(A2S_INFO_STRING, sizeof(A2S_INFO_STRING));

    qint64 ping = QDateTime::currentMSecsSinceEpoch();
    QByteArray byteResponse = SendUDPQuery(query, host, port);
    ping = QDateTime::currentMSecsSinceEpoch()-ping;
    if(ping > 2000)
    {
        ping = ping - 2000;
    }
    if(ping > 2000)
    {
        ping = 2000;
    }

    // Challenge-response is now handled inside SendUDPQuery on the same socket

    return new InfoReply(byteResponse, ping);
}

PlayerQuery::PlayerQuery(MainWindow *main)
{
    pMain = main;
    this->worker = new Worker;
    this->worker->moveToThread(&workerThread);
    connect(&workerThread, &QThread::finished, this, &PlayerQuery::deleteLater);
    connect(&workerThread, &QThread::finished, this->worker, &Worker::deleteLater);
    connect(this, &PlayerQuery::query, this->worker, &Worker::getPlayerInfo);
    connect(this->worker, &Worker::playerInfoReady, main, &MainWindow::PlayerInfoReady);

    workerThread.start();
}
PlayerQuery::~PlayerQuery()
{
    workerThread.quit();
    workerThread.wait();

    if(this == pMain->pPlayerQuery)
        pMain->pPlayerQuery = nullptr;
}

QList<PlayerInfo> *GetPlayerReply(QHostAddress host, quint16 port, bool goldsrcSplits)
{
    QList<PlayerInfo> *list = new QList<PlayerInfo>();

    QByteArray query;
    QDataStream data(&query, QIODevice::ReadWrite);
    data.setByteOrder(QDataStream::LittleEndian);
    data << A2S_HEADER << (qint8)A2S_PLAYER << qint32(-1);

    QByteArray byteResponse = SendUDPQuery(query, host, port, goldsrcSplits);
    QDataStream response(byteResponse);
    response.setByteOrder(QDataStream::LittleEndian);

    qint32 header;
    qint8 byteCheck;

    // Challenge-response is now handled inside SendUDPQuery on the same socket
    response.setFloatingPointPrecision(QDataStream::SinglePrecision);
    response >> header;
    response >> byteCheck;

    if(header == -1 && byteCheck == A2S_PLAYER_CHECK)
    {
        quint8 count;
        response >> count;

        for(int i = 0; i < count; i++)
        {
            PlayerInfo playerInfo;
            response >> byteCheck;
            playerInfo.name = GetStringFromStream(response);
            response >> playerInfo.score;
            response >> playerInfo.time;
            list->append(playerInfo);
        }
    }
    return list;
}

RulesQuery::RulesQuery(MainWindow *main)
{
    pMain = main;
    this->worker = new Worker;
    this->worker->moveToThread(&workerThread);
    connect(&workerThread, &QThread::finished, this, &RulesQuery::deleteLater);
    connect(&workerThread, &QThread::finished, this->worker, &Worker::deleteLater);
    connect(this, &RulesQuery::query, this->worker, &Worker::getRulesInfo);
    connect(this->worker, &Worker::rulesInfoReady, main, &MainWindow::RulesInfoReady);

    workerThread.start();
}
RulesQuery::~RulesQuery()
{
    workerThread.quit();
    workerThread.wait();

    if(this == pMain->pRulesQuery)
        pMain->pRulesQuery = nullptr;
}

QList<RulesInfo> *GetRulesReply(QHostAddress host, quint16 port, bool goldsrcSplits)
{
    QList<RulesInfo> *list = new QList<RulesInfo>();

    QByteArray query;
    QDataStream data(&query, QIODevice::ReadWrite);
    data.setByteOrder(QDataStream::LittleEndian);
    data << A2S_HEADER << (qint8)A2S_RULES << qint32(-1);

    QByteArray byteResponse = SendUDPQuery(query, host, port, goldsrcSplits);
    QDataStream response(byteResponse);
    response.setByteOrder(QDataStream::LittleEndian);

    qint32 header;
    qint8 byteCheck;

    // Challenge-response is now handled inside SendUDPQuery on the same socket
    response >> header;
    response >> byteCheck;

    if(header == -1 && byteCheck == A2S_RULES_CHECK)
    {
        quint16 count;
        response >> count;

        QString name;
        QString value;
        for(int i = 0; i < count; i++)
        {
            name = GetStringFromStream(response);
            value = GetStringFromStream(response);
            list->append(RulesInfo(name, value));
            name = "";
            value = "";
        }
    }
    return list;
}