ReArchive/src/ReArchive.cpp

#include <fstream>
#include <unordered_set>
#include <ReArchive/ReArchive.h>
#include <ReArchive/types/Header.h>
#include <ReArchive/types/FileTable.h>
#include <ReArchive/types/FileEntry.h>
#include <cassert>

using ReArchive::Header;
using ReArchive::FileTable;
using ReArchive::FileEntry;

std::unordered_set<std::filesystem::path> locked {};

Header GetHeader(const unsigned char* archive) {
    ReArchive::Header h = ReArchive::Header::DeSerialize(archive);
    return h;
}

/// @param header our header.
/// @param in Our input stream to the file.
/// @note Does not close the input stream.
FileTable GetFileTable(const Header& header, std::ifstream& in) {
    in.seekg(0, std::ios::end);
    int64_t file_table_size = in.tellg() - header.FileTableOffset();
    in.seekg(header.FileTableOffset(), std::ios::beg);

    std::vector<unsigned char> buffer(file_table_size);
    in.read(reinterpret_cast<char *>(buffer.data()), buffer.size());

    unsigned char* ptr = buffer.data();
    int64_t file_table_entry_count = be64toh(*reinterpret_cast<int64_t*>(ptr));
    ptr += sizeof(int64_t);

    FileTable result;
    for (int64_t i = 0; i < file_table_entry_count; i++) {
        // Out of bounds.
        assert(ptr < (buffer.data() + buffer.size()));

        int64_t string_size = be64toh(*reinterpret_cast<const int64_t*>(ptr));
        ptr += sizeof(int64_t);

        std::string path(reinterpret_cast<const char*>(ptr), string_size);
        ptr += string_size;

        int64_t data_size = be64toh(*reinterpret_cast<const int64_t*>(ptr));
        ptr += sizeof(int64_t);

        int64_t data_offset = be64toh(*reinterpret_cast<const int64_t*>(ptr));
        ptr += sizeof(int64_t);

        result.Append({ data_size, data_offset, path });
    }
    return result;
}

std::pair<bool, FileTable> ReArchive::ReadFileTable(const std::filesystem::path& archive) {
    if (!std::filesystem::exists(archive))
        return {false, {}};

    // Busy-wait.
    while (locked.contains(archive)) {}
    locked.insert(archive);

    std::ifstream in(archive, std::ios::binary);
    if (!in)
        return {false, {}};

    in.seekg(0, std::ios::end);
    if (in.tellg() < Header::Size())
        return {false, {}};
    in.seekg(0, std::ios::beg);

    std::vector<unsigned char> buffer (Header::Size());
    in.read(reinterpret_cast<char *>(buffer.data()), (int64_t) buffer.size());
    if (buffer[0] != 'R' || buffer[1] != 'S' || buffer[2] != 'A')
        return {false, {}};

    auto header = GetHeader(buffer.data());
    auto file_table = GetFileTable(header, in);
    in.close();

    // Remove lock.
    auto position = locked.find(archive);
    if (position != locked.end())
        locked.erase(position);

    return {true, file_table};
}

bool ReArchive::CreateArchive(const std::filesystem::path& filesystem_path, bool use_compression, FileTable* current_file_table) {
    if (std::filesystem::exists(filesystem_path))
        return false;

    while (locked.contains(filesystem_path)) {}
    locked.insert(filesystem_path);

    std::ofstream file(filesystem_path, std::ios::binary);
    if (!file)
        return false;

    auto serialized_file_header = Header::Serialize(Header(use_compression, Header::Size()));
    auto file_table = FileTable();
    auto serialized_file_table = FileTable::Serialize(file_table);

    file.write(reinterpret_cast<const char*>(serialized_file_header.data()), (int64_t) serialized_file_header.size());
    file.write(reinterpret_cast<const char*>(serialized_file_table.data()), (int64_t) serialized_file_table.size());
    file.close();

    if (current_file_table)
        *current_file_table = file_table;

    // Remove lock.
    auto position = locked.find(filesystem_path);
    if (position != locked.end())
        locked.erase(position);

    return true;
}

bool ReArchive::WriteFile(const std::filesystem::path& archive, const std::filesystem::path& file_path, const unsigned char* file_data, const int64_t& byte_count, FileTable* current_file_table) {
    if (!std::filesystem::exists(archive))
        return false;

    // Busy-wait.
    while (locked.contains(archive)) {}
    locked.insert(archive);

    std::ifstream in(archive, std::ios::binary);
    if (!in)
        return false;

    in.seekg(0, std::ios::end);
    if (in.tellg() < Header::Size())
        return false;
    in.seekg(0, std::ios::beg);

    std::vector<unsigned char> buffer (Header::Size());
    in.read(reinterpret_cast<char *>(buffer.data()), (int64_t) buffer.size());
    if (buffer[0] != 'R' || buffer[1] != 'S' || buffer[2] != 'A')
        return false;

    auto header = GetHeader(buffer.data());

    FileTable* file_table;
    if (current_file_table == nullptr)
        file_table = new FileTable(GetFileTable(header, in));
    else
        file_table = current_file_table;

    if (file_table->Contains(file_path))
        return false;

    in.close();

    std::ofstream out(archive, std::ios::binary | std::ios::out | std::ios::in);
    if (!out)
        return false;

    out.seekp(header.FileTableOffset(), std::ios::beg);
    out.write(reinterpret_cast<const char *>(file_data), byte_count);

    file_table->Append(FileEntry(byte_count, header.FileTableOffset(), file_path));
    header.FileTableOffset(out.tellp());

    auto new_file_table = FileTable::Serialize(*file_table);
    out.write(reinterpret_cast<const char *>(new_file_table.data()), (int64_t) new_file_table.size());

    auto new_header = Header::Serialize(header);
    out.seekp(0, std::ios::beg);
    out.write(reinterpret_cast<const char *>(new_header.data()), (int64_t) new_header.size());
    out.close();

    if (current_file_table)
        *current_file_table = *file_table;

    if (current_file_table == nullptr)
        delete file_table;

    // Remove lock.
    auto position = locked.find(archive);

    if (position != locked.end())
        locked.erase(position);

    return true;
}

bool ReArchive::OverwriteFile(const std::filesystem::path& archive, const std::filesystem::path& file_path, const unsigned char* file_data, const int64_t& byte_count, FileTable* current_file_table) {
    if (!std::filesystem::exists(archive))
        return false;

    // Busy-wait.
    while (locked.contains(archive)) {}
    locked.insert(archive);

    std::ifstream in(archive, std::ios::binary);
    if (!in)
        return false;

    in.seekg(0, std::ios::end);
    if (in.tellg() < Header::Size())
        return false;
    in.seekg(0, std::ios::beg);

    std::vector<unsigned char> buffer (Header::Size());
    in.read(reinterpret_cast<char *>(buffer.data()), (int64_t) buffer.size());
    if (buffer[0] != 'R' || buffer[1] != 'S' || buffer[2] != 'A')
        return false;

    auto header = GetHeader(buffer.data());
    FileTable* file_table;
    if (current_file_table == nullptr)
        file_table = new FileTable(GetFileTable(header, in));
    else
        file_table = current_file_table;

    const FileEntry* target = nullptr;
    auto file_entries = file_table->GetEntries();

    auto value = file_entries->find(file_path);
    if (value != file_entries->end())
        target = &value->second;

    if (!target)
        return false;

    if (byte_count != target->Size())
        return false;

    in.close();
    std::ofstream out(archive, std::ios::binary | std::ios::out | std::ios::in);
    if (!out)
        return false;

    out.seekp(target->Offset(), std::ios::beg);
    out.write(reinterpret_cast<const char *>(file_data), byte_count);
    out.close();

    if (current_file_table == nullptr)
        delete file_table;

    // Remove lock.
    auto position = locked.find(archive);
    if (position != locked.end())
        locked.erase(position);

    return true;
}

std::vector<unsigned char> ReArchive::ReadFile(const std::filesystem::path& archive, const std::filesystem::path& file_path, FileTable* current_file_table) {
    if (!std::filesystem::exists(archive))
        return {};

    // Busy-wait.
    while (locked.contains(archive)) {}
    locked.insert(archive);

    std::ifstream in(archive, std::ios::binary);
    if (!in)
        return {};

    in.seekg(0, std::ios::end);
    if (in.tellg() < Header::Size())
        return {};
    in.seekg(0, std::ios::beg);

    std::vector<unsigned char> buffer (Header::Size());
    in.read(reinterpret_cast<char *>(buffer.data()), (int64_t) buffer.size());
    if (buffer[0] != 'R' || buffer[1] != 'S' || buffer[2] != 'A')
        return {};

    auto header = GetHeader(buffer.data());
    FileTable* file_table;
    if (current_file_table == nullptr)
        file_table = new FileTable(GetFileTable(header, in));
    else
        file_table = current_file_table;

    const FileEntry* target = nullptr;
    auto file_entries = file_table->GetEntries();

    auto value = file_entries->find(file_path);
    if (value != file_entries->end())
        target = &value->second;

    if (!target)
        return {};

    std::vector<unsigned char> result(target->Size());
    in.seekg(target->Offset(), std::ios::beg);
    in.read(reinterpret_cast<char*>(result.data()), (int64_t) result.size());
    in.close();

    // delete the one we allocated if one was not passed in.
    if (current_file_table == nullptr)
        delete file_table;

    // Remove lock.
    auto position = locked.find(archive);
    if (position != locked.end())
        locked.erase(position);
    return result;
}

// I tried to do this several different ways but this seems to be the best approach - Redacted.
bool ReArchive::EraseFile(const std::filesystem::path& archive, const std::filesystem::path& file_path, FileTable* running_tally) {

    if (!std::filesystem::exists(archive))
        return false;

    // Busy-wait.
    while (locked.contains(archive)) {}

    std::ifstream in(archive, std::ios::binary);
    if (!in)
        return false;

    in.seekg(0, std::ios::end);
    int64_t file_size = in.tellg();
    if (file_size < Header::Size())
        return false;
    in.seekg(0, std::ios::beg);

    std::vector<unsigned char> buffer (Header::Size());
    in.read(reinterpret_cast<char *>(buffer.data()), (int64_t) buffer.size());
    if (buffer[0] != 'R' || buffer[1] != 'S' || buffer[2] != 'A')
        return false;

    auto current_header = GetHeader(buffer.data());
    auto current_file_table = GetFileTable(current_header, in);

    // TODO randomize the name more than that.
    if (!CreateArchive(archive.string() + ".tmp", current_header.Compressed()))
        return false;

    auto file_entries = current_file_table.GetEntries();
    auto value = file_entries->find(file_path);
    if (value != file_entries->end())
        file_entries->erase(value);

    for (auto& e : *file_entries) {
        auto file_buffer = ReadFile(archive, e.first);
        if (!WriteFile(archive.string() + ".tmp", e.first, file_buffer.data(), (int64_t) file_buffer.size())) {
            std::filesystem::remove(archive.string() + ".tmp");
            return false;
        }
    }

    // Busy-wait.
    while (locked.contains(archive)) {}
    locked.insert(archive);

    std::filesystem::remove(archive);
    std::filesystem::rename(archive.string() + ".tmp", archive);

    // Remove lock.
    auto position = locked.find(archive);
    if (position != locked.end())
        locked.erase(position);

    if (running_tally) {
        auto result = ReadFileTable(archive);
        if (result.first)
            *running_tally = result.second;
    }
    return true;
}