squawk/core/storage/archive.cpp

/*
 * Squawk messenger. 
 * Copyright (C) 2019  Yury Gubich <blue@macaw.me>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

#include "archive.h"
#include <sys/stat.h>
#include <sys/types.h>
#include <QStandardPaths>
#include <QDebug>
#include <QDataStream>
#include <QDir>

Core::Archive::Archive(const QString& account, const QString& p_jid, QObject* parent):
    QObject(parent),
    jid(p_jid),
    account(account),
    opened(false),
    db(account + "/" + jid),
    messages(db.addStorage<QString, Shared::Message>("messages")),
    order(db.addStorage<uint64_t, QString>("order")),
    stats(db.addStorage<QString, QVariant>("stats")),
    avatars(db.addStorage<QString, AvatarInfo>("avatars")),
    stanzaIdToId(db.addStorage<QString, QString>("stanzaIdToId")),
    cursor(order->createCursor())
{}

Core::Archive::~Archive() {
    close();
}

void Core::Archive::open() {
    db.open();
    LMDBAL::WriteTransaction txn = db.beginTransaction();

    AvatarInfo info;
    bool hasAvatar = false;
    try {
        avatars->getRecord(jid, info, txn);
        hasAvatar = true;
    } catch (const LMDBAL::NotFound& e) {}

    if (!hasAvatar)
        return;

    QFile ava(db.getPath() + "/" + jid + "." + info.type);
    if (ava.exists())
        return;

    try {
        avatars->removeRecord(jid, txn);
        txn.commit();
    } catch (const std::exception& e) {
        qDebug() << e.what();
        qDebug() << "error opening archive" << jid << "for account" << account
        << ". There is supposed to be avatar but the file doesn't exist, couldn't even drop it, it surely will lead to an error";
    }
}

void Core::Archive::close() {
    db.close();
}

bool Core::Archive::addElement(const Shared::Message& message) {
    QString id = message.getId();
    qDebug() << "Adding message with id " << id;

    try {
        LMDBAL::WriteTransaction txn = db.beginTransaction();
        messages->addRecord(id, message, txn);
        order->addRecord(message.getTime().toMSecsSinceEpoch(), id, txn);
        QString stanzaId = message.getStanzaId();
        if (!stanzaId.isEmpty())
            stanzaIdToId->addRecord(stanzaId, id, txn);

        txn.commit();
        return true;
    } catch (const std::exception& e) {
        qDebug() << "Could not add message with id " + id;
        qDebug() << e.what();
    }

    return false;
}

void Core::Archive::clear() {
    db.drop();
}

Shared::Message Core::Archive::getElement(const QString& id) const {
    return messages->getRecord(id);
}

bool Core::Archive::hasElement(const QString& id) const {
    return messages->checkRecord(id);
}

void Core::Archive::changeMessage(const QString& id, const QMap<QString, QVariant>& data) {
    LMDBAL::WriteTransaction txn = db.beginTransaction();
    Shared::Message msg = messages->getRecord(id, txn);

    bool hadStanzaId = !msg.getStanzaId().isEmpty();
    QDateTime oTime = msg.getTime();
    bool idChange = msg.change(data);
    QString newId = msg.getId();
    QDateTime nTime = msg.getTime();

    bool orderChange = oTime != nTime;
    if (idChange || orderChange) {
        if (idChange)
            messages->removeRecord(id, txn);

        if (orderChange)
            order->removeRecord(oTime.toMSecsSinceEpoch(), txn);

        order->forceRecord(nTime.toMSecsSinceEpoch(), newId, txn);
    }

    QString sid = msg.getStanzaId();
    if (!sid.isEmpty() && (idChange || !hadStanzaId))
        stanzaIdToId->forceRecord(sid, newId, txn);

    messages->forceRecord(newId, msg, txn);
    txn.commit();
}

Shared::Message Core::Archive::newest() const {
    LMDBAL::Transaction txn = db.beginReadOnlyTransaction();

    try {
        cursor.open(txn);
        while (true) {
            std::pair<uint64_t, QString> pair = cursor.prev();
            Shared::Message msg = messages->getRecord(pair.second, txn);
            if (msg.serverStored()) {
                cursor.close();
                return msg;
            }
        }
    } catch (...) {
        cursor.close();
        throw;
    }
}

QString Core::Archive::newestId() const {
    Shared::Message msg = newest();
    return msg.getId();
}

QString Core::Archive::oldestId() const {
    Shared::Message msg = oldest();
    return msg.getId();
}

Shared::Message Core::Archive::oldest() const {
    LMDBAL::Transaction txn = db.beginReadOnlyTransaction();

    try {
        cursor.open(txn);
        while (true) {
            std::pair<uint64_t, QString> pair = cursor.next();
            Shared::Message msg = messages->getRecord(pair.second, txn);
            if (msg.serverStored()) {
                cursor.close();
                return msg;
            }
        }
    } catch (...) {
        cursor.close();
        throw;
    }
}

unsigned int Core::Archive::addElements(const std::list<Shared::Message>& messages) {
    unsigned int success = 0;
    LMDBAL::WriteTransaction txn = db.beginTransaction();
    for (const Shared::Message& message : messages) {
        QString id = message.getId();
        bool added = false;
        try {
            Core::Archive::messages->addRecord(id, message, txn);
            added = true;
        } catch (const LMDBAL::Exist& e) {}

        if (!added)
            continue;

        order->addRecord(message.getTime().toMSecsSinceEpoch(), id, txn);

        QString sid = message.getStanzaId();
        if (!sid.isEmpty())
            stanzaIdToId->addRecord(sid, id, txn);

        ++success;
    }
    txn.commit();
    
    return success;
}

long unsigned int Core::Archive::size() const {
    return order->count();
}

std::list<Shared::Message> Core::Archive::getBefore(unsigned int count, const QString& id) {
    LMDBAL::Transaction txn = db.beginReadOnlyTransaction();
    std::list<Shared::Message> res;
    try {
        cursor.open(txn);
        if (!id.isEmpty()) {
            Shared::Message reference = messages->getRecord(id, txn);
            uint64_t stamp = reference.getTime().toMSecsSinceEpoch();
            cursor.set(stamp);
        }

        for (unsigned int i = 0; i < count; ++i) {
            std::pair<uint64_t, QString> pair;
            cursor.prev(pair.first, pair.second);

            res.emplace_back();
            Shared::Message& msg = res.back();
            messages->getRecord(pair.second, msg, txn);
        }
        cursor.close();

        return res;
    } catch (const LMDBAL::NotFound& e) {
        cursor.close();
        if (res.empty())
            throw e;
        else
            return res;
    } catch (...) {
        cursor.close();
        throw;
    }
}

bool Core::Archive::isFromTheBeginning() const {
    try {
        return stats->getRecord("fromTheBeginning").toBool();
    } catch (const LMDBAL::NotFound& e) {
        return false;
    }
}

void Core::Archive::setFromTheBeginning(bool is) {
    stats->forceRecord("fromTheBeginning", is);
}

bool Core::Archive::isEncryptionEnabled() const {
    try {
        return stats->getRecord("isEncryptionEnabled").toBool();
    } catch (const LMDBAL::NotFound& e) {
        return false;
    }
}

bool Core::Archive::setEncryptionEnabled(bool is) {
    LMDBAL::WriteTransaction txn = db.beginTransaction();
    bool current = false;
    try {
        current = stats->getRecord("isEncryptionEnabled", txn).toBool();
    } catch (const LMDBAL::NotFound& e) {}

    if (is != current) {
        stats->forceRecord("isEncryptionEnabled", is, txn);
        txn.commit();
        return true;
    }

    return false;
}

QString Core::Archive::idByStanzaId(const QString& stanzaId) const {
    return stanzaIdToId->getRecord(stanzaId);
}

QString Core::Archive::stanzaIdById(const QString& id) const {
    try {
        Shared::Message msg = getElement(id);
        return msg.getStanzaId();
    } catch (const LMDBAL::NotFound& e) {
        return QString();
    }
}

bool Core::Archive::setAvatar(const QByteArray& data, AvatarInfo& newInfo, bool generated, const QString& resource) {
    LMDBAL::WriteTransaction txn = db.beginTransaction();
    AvatarInfo oldInfo;
    bool haveAvatar = false;
    QString res = resource.isEmpty() ? jid : resource;
    try {
        avatars->getRecord(res, oldInfo, txn);
        haveAvatar = true;
    } catch (const LMDBAL::NotFound& e) {}
    
    if (data.size() == 0) {
        if (!haveAvatar)
            return false;

        avatars->removeRecord(res, txn);
        txn.commit();
        return true;
    } else {
        QString currentPath = db.getPath();
        bool needToRemoveOld = false;
        QCryptographicHash hash(QCryptographicHash::Sha1);
        hash.addData(data);
        QByteArray newHash(hash.result());
        if (haveAvatar) {
            if (!generated && !oldInfo.autogenerated && oldInfo.hash == newHash)
                return false;

            QFile oldAvatar(currentPath + "/" + res + "." + oldInfo.type);
            if (oldAvatar.exists()) {
                if (oldAvatar.rename(currentPath + "/" + res + "." + oldInfo.type + ".bak")) {
                    needToRemoveOld = true;
                } else {
                    qDebug() << "Can't change avatar: couldn't get rid of the old avatar" << oldAvatar.fileName();
                    return false;
                }
            }
        }
        QMimeDatabase mimedb;
        QMimeType type = mimedb.mimeTypeForData(data);
        QString ext = type.preferredSuffix();
        QFile newAvatar(currentPath + "/" + res + "." + ext);
        if (newAvatar.open(QFile::WriteOnly)) {
            newAvatar.write(data);
            newAvatar.close();
            
            newInfo.type = ext;
            newInfo.hash = newHash;
            newInfo.autogenerated = generated;
            try {
                avatars->forceRecord(res, newInfo, txn);
                txn.commit();
            } catch (...) {
                qDebug() << "Can't change avatar: couldn't store changes to database for" << newAvatar.fileName() << "rolling back to the previous state";
                if (needToRemoveOld) {
                    QFile oldAvatar(currentPath + "/" + res + "." + oldInfo.type + ".bak");
                    oldAvatar.rename(currentPath + "/" + res + "." + oldInfo.type);
                }
                return false;
            }
            
            if (needToRemoveOld) {
                QFile oldAvatar(currentPath + "/" + res + "." + oldInfo.type + ".bak");
                oldAvatar.remove();
            }
            return true;
        } else {
            qDebug() << "Can't change avatar: cant open file to write" << newAvatar.fileName() << "rolling back to the previous state";
            if (needToRemoveOld) {
                QFile oldAvatar(currentPath + "/" + res + "." + oldInfo.type + ".bak");
                oldAvatar.rename(currentPath + "/" + res + "." + oldInfo.type);
            }
            return false;
        }
    }
}

bool Core::Archive::readAvatarInfo(Core::Archive::AvatarInfo& target, const QString& resource) const {
    try {
        avatars->getRecord(resource.isEmpty() ? jid : resource, target);
        return true;
    } catch (const LMDBAL::NotFound& e) {
        return false;
    }
}

void Core::Archive::readAllResourcesAvatars(std::map<QString, AvatarInfo>& data) const {
    avatars->readAll(data);
}

Core::Archive::AvatarInfo Core::Archive::getAvatarInfo(const QString& resource) const {
    return avatars->getRecord(resource);
}

Core::Archive::AvatarInfo::AvatarInfo():
    type(),
    hash(),
    autogenerated(false)
{}

Core::Archive::AvatarInfo::AvatarInfo(const QString& p_type, const QByteArray& p_hash, bool p_autogenerated):
    type(p_type),
    hash(p_hash),
    autogenerated(p_autogenerated)
{}

QDataStream & operator<<(QDataStream& out, const Core::Archive::AvatarInfo& info) {
    out << info.type;
    out << info.hash;
    out << info.autogenerated;

    return out;
}

QDataStream & operator>>(QDataStream& in, Core::Archive::AvatarInfo& info) {
    in >> info.type;
    in >> info.hash;
    in >> info.autogenerated;

    return in;
}