Server/SourceFiles/FileHandler/FileHandler.cpp - master - Shymer123/Messenger

Server/SourceFiles/FileHandler/FileHandler.cpp
395 строк12 KB
Shymer123
changed sending files to the client
25 мар 2025, 17:09
25 мар 2025, 17:09b6b912b
#include "FileHandler.h"
#include "Logger/Logger.h"
#include "MainServer/Server.h"
#include "Utils/Utils.h"
#include <filesystem>
#include <fstream>
#include <iostream>
#include <netinet/in.h>
#include <sys/socket.h>
#include <thread>

namespace fs = std::filesystem;
using json = nlohmann::json;

std::optional<nlohmann::json> FileHandler::receiveFile(int clientSocket)
{
    auto filesCountOpt = receiveFileCount(clientSocket);
    if(!filesCountOpt)
    {
        return std::nullopt;
    }

    uint32_t fileCount = *filesCountOpt;
    json filesMetadata = json::array();

    for(uint32_t i = 0; i < fileCount; ++i)
    {
        auto fileMetadata = receiveSingleFile(clientSocket);
        if(!fileMetadata)
        {
            continue;
        }

        filesMetadata.push_back(*fileMetadata);
    }

    json content;
    content["content"] = filesMetadata;

    return content;
}

std::optional<uint32_t> FileHandler::receiveFileCount(int clientSocket)
{
    uint32_t fileCount;
    size_t bytesReadTotal = 0;

    while(bytesReadTotal < sizeof(fileCount))
    {
        int bytesRead = recv(clientSocket, reinterpret_cast<char*>(&fileCount) + bytesReadTotal, sizeof(fileCount) - bytesReadTotal, 0);

        if(bytesRead == 0)
        {
            Logger::getInstance().logError("Client disconnected while reading files count", Logger::logLevel::Error);
            return std::nullopt;
        }
        else if(bytesRead < 0)
        {
            if(errno == EAGAIN || errno == EWOULDBLOCK)
            {
                std::this_thread::sleep_for(std::chrono::milliseconds(10));
                continue;
            }
            Logger::getInstance().logError("Error reading from socket: " + std::string(strerror(errno)), Logger::logLevel::Error);
            return std::nullopt;
        }

        bytesReadTotal += bytesRead;
    }

    return ntohl(fileCount);
}


std::optional<json> FileHandler::receiveSingleFile(int clientSocket)
{
    auto fileNameSizeOpt = receiveFileNameSize(clientSocket);
    if(!fileNameSizeOpt)
    {
        return std::nullopt;
    }

    auto fileNameOpt = receiveFileName(clientSocket, *fileNameSizeOpt);
    if(!fileNameOpt)
    {
        return std::nullopt;
    }

    auto fileSizeOpt = receiveFileSize(clientSocket);
    if(!fileSizeOpt)
    {
        return std::nullopt;
    }

    auto chunkCountOpt = receiveChunkCount(clientSocket);
    if(!chunkCountOpt)
    {
        return std::nullopt;
    }

    return saveFile(clientSocket, *fileNameOpt, *fileSizeOpt, *chunkCountOpt);
}

std::optional<uint32_t> FileHandler::receiveFileNameSize(int clientSocket)
{
    uint32_t nameSize;
    size_t bytesReadTotal = 0;

    while(bytesReadTotal < sizeof(nameSize))
    {
        int bytesRead = recv(clientSocket, reinterpret_cast<char*>(&nameSize) + bytesReadTotal, sizeof(nameSize) - bytesReadTotal, 0);

        if(bytesRead == 0)
        {
            Logger::getInstance().logError("Client disconnected while reading files count", Logger::logLevel::Error);
            return std::nullopt;
        }
        else if(bytesRead < 0)
        {
            if(errno == EAGAIN || errno == EWOULDBLOCK)
            {
                std::this_thread::sleep_for(std::chrono::milliseconds(10));
                continue;
            }
            Logger::getInstance().logError("Error reading from socket: " + std::string(strerror(errno)), Logger::logLevel::Error);
            return std::nullopt;
        }

        bytesReadTotal += bytesRead;
    }

    return ntohl(nameSize);
}

std::optional<std::string> FileHandler::receiveFileName(int clientSocket, uint32_t nameSize)
{
    std::vector<char> nameBuffer(nameSize);
    size_t bytesReadTotal = 0;

    while(bytesReadTotal < nameSize)
    {
        int bytesRead = recv(clientSocket, nameBuffer.data() + bytesReadTotal, nameSize - bytesReadTotal, 0);

        if(bytesRead == 0)
        {
            Logger::getInstance().logError("Client disconnected while reading files count", Logger::logLevel::Error);
            return std::nullopt;
        }
        else if(bytesRead < 0)
        {
            if(errno == EAGAIN || errno == EWOULDBLOCK)
            {
                std::this_thread::sleep_for(std::chrono::milliseconds(10));
                continue;
            }
            Logger::getInstance().logError("Error reading from socket: " + std::string(strerror(errno)), Logger::logLevel::Error);
            return std::nullopt;
        }

        bytesReadTotal += bytesRead;
    }

    return std::string(nameBuffer.begin(), nameBuffer.end());
}

std::optional<uint64_t> FileHandler::receiveFileSize(int clientSocket)
{
    uint64_t fileSize;
    size_t bytesReadTotal = 0;

    while(bytesReadTotal < sizeof(fileSize))
    {
        int bytesRead = recv(clientSocket, reinterpret_cast<char*>(&fileSize) + bytesReadTotal, sizeof(fileSize) - bytesReadTotal, 0);

        if(bytesRead == 0)
        {
            Logger::getInstance().logError("Client disconnected while reading files count", Logger::logLevel::Error);
            return std::nullopt;
        }
        else if(bytesRead < 0)
        {
            if(errno == EAGAIN || errno == EWOULDBLOCK)
            {
                std::this_thread::sleep_for(std::chrono::milliseconds(10));
                continue;
            }
            Logger::getInstance().logError("Error reading from socket: " + std::string(strerror(errno)), Logger::logLevel::Error);
            return std::nullopt;
        }

        bytesReadTotal += bytesRead;
    }

    return be64toh(fileSize);
}

std::optional<uint32_t> FileHandler::receiveChunkCount(int clientSocket)
{
    uint32_t chunkCount;
    size_t bytesReadTotal = 0;

    while(bytesReadTotal < sizeof(chunkCount))
    {
        int bytesRead = recv(clientSocket, reinterpret_cast<char*>(&chunkCount) + bytesReadTotal, sizeof(chunkCount) - bytesReadTotal, 0);

        if(bytesRead == 0)
        {
            Logger::getInstance().logError("Client disconnected while reading files count", Logger::logLevel::Error);
            return std::nullopt;
        }
        else if(bytesRead < 0)
        {
            if(errno == EAGAIN || errno == EWOULDBLOCK)
            {
                std::this_thread::sleep_for(std::chrono::milliseconds(10));
                continue;
            }
            Logger::getInstance().logError("Error reading from socket: " + std::string(strerror(errno)), Logger::logLevel::Error);
            return std::nullopt;
        }

        bytesReadTotal += bytesRead;
    }

    return ntohl(chunkCount);
}

std::optional<json> FileHandler::saveFile(int clientSocket, const std::string& fileName,
                                          uint64_t fileSize, uint32_t chunkCount)
{
    std::string fileUUID = Utils::generateUniqueID();
    std::string savedFileName = fileUUID + "_" + fileName;
    fs::create_directory("uploads");

    std::ofstream file("uploads/" + savedFileName, std::ios::binary);
    if(!file.is_open())
    {
        Logger::getInstance().logError("Failed to create file: " + savedFileName, Logger::logLevel::Error);
        return std::nullopt;
    }

    for(uint32_t currentChunk = 0; currentChunk < chunkCount; ++currentChunk)
    {
        if(!receiveAndWriteChunk(clientSocket, file, fileName))
        {
            file.close();
            fs::remove("uploads/" + savedFileName);
            return std::nullopt;
        }
    }

    file.close();
    return json{{"fileName", fileName}, {"fileSize", fileSize}, {"fileUUID", fileUUID}};
}

bool FileHandler::receiveAndWriteChunk(int clientSocket, std::ofstream& file, const std::string& fileName)
{
    uint32_t chunkSize;
    size_t bytesReadTotal = 0;

    while(bytesReadTotal < sizeof(chunkSize))
    {
        int bytesRead = recv(clientSocket, reinterpret_cast<char*>(&chunkSize) + bytesReadTotal, sizeof(chunkSize) - bytesReadTotal, 0);

        if(bytesRead == 0)
        {
            Logger::getInstance().logError("Client disconnected while reading files count", Logger::logLevel::Error);
            return false;
        }
        else if(bytesRead < 0)
        {
            if(errno == EAGAIN || errno == EWOULDBLOCK)
            {
                std::this_thread::sleep_for(std::chrono::milliseconds(10));
                continue;
            }
            Logger::getInstance().logError("Error reading from socket: " + std::string(strerror(errno)), Logger::logLevel::Error);
            return false;
        }

        bytesReadTotal += bytesRead;
    }
    chunkSize = ntohl(chunkSize);

    std::vector<char> compressedChunk(chunkSize);
    bytesReadTotal = 0;
    while(bytesReadTotal < chunkSize)
    {
        int bytesRead = recv(clientSocket, compressedChunk.data() + bytesReadTotal, chunkSize - bytesReadTotal, 0);

        if(bytesRead == 0)
        {
            Logger::getInstance().logError("Client disconnected while reading files count", Logger::logLevel::Error);
            return false;
        }
        else if(bytesRead < 0)
        {
            if(errno == EAGAIN || errno == EWOULDBLOCK)
            {
                std::this_thread::sleep_for(std::chrono::milliseconds(10));
                continue;
            }
            Logger::getInstance().logError("Error reading from socket: " + std::string(strerror(errno)), Logger::logLevel::Error);
            return false;
        }

        bytesReadTotal += bytesRead;
    }

    std::vector<char> fileData;
    if(shouldDecompress(fileName))
    {
        fileData = Utils::decompressData(compressedChunk);
    }
    else
    {
        fileData = std::move(compressedChunk);
    }

    file.write(fileData.data(), fileData.size());
    return true;
}

bool FileHandler::shouldDecompress(const std::string& fileName)
{
    std::vector<std::string> nonCompressExtensions = {".jpg", ".jpeg", ".mp4", ".zip"};

    for(const std::string& extension : nonCompressExtensions)
    {
        if(fileName.size() >= extension.size()
            && fileName.compare(fileName.size() - extension.size(), extension.size(), extension) == 0)
        {
            return false;
        }
    }

    return true;
}

void FileHandler::sendFile(Server* server, int clientSocket, uint32_t fileID, const std::string& fileName)
{
    std::ifstream inFile("uploads/" + fileName, std::ios::binary);
    if(!inFile.is_open())
    {
        Logger::getInstance().logError("File extract error", Logger::logLevel::Error);
        return;
    }

    inFile.seekg(0, std::ios::end);
    uint64_t fileSize = inFile.tellg();
    inFile.seekg(0, std::ios::beg);

    const int32_t CHUNK_SIZE = 8192;

    uint32_t chunkCount = static_cast<uint32_t>((fileSize + CHUNK_SIZE - 1) / CHUNK_SIZE);
    bool shouldCompress = shouldDecompress(fileName);

    char buffer[CHUNK_SIZE];

    while (inFile.read(buffer, sizeof(buffer)) || inFile.gcount() > 0)
    {
        std::vector<char> packet;
        std::vector<char> chunkData(buffer, buffer + inFile.gcount());
        std::vector<char> compressedData;

        if(shouldCompress)
        {
            compressedData = Utils::compressData(chunkData);
        }

        uint32_t chunkSize = static_cast<uint32_t>(shouldCompress ? compressedData.size()
                                                                  : chunkData.size());

        uint32_t beFileID = htonl(fileID);
        uint32_t beChunkCount = htonl(chunkCount);
        uint32_t beChunkSize = htonl(chunkSize);

        packet.reserve(packet.size() + sizeof(beFileID) + sizeof(beChunkCount) + sizeof(beChunkSize) + chunkSize);

        packet.insert(packet.end(), reinterpret_cast<char*>(&beFileID), reinterpret_cast<char*>(&beFileID) + sizeof(beFileID));
        packet.insert(packet.end(), reinterpret_cast<char*>(&beChunkCount), reinterpret_cast<char*>(&beChunkCount) + sizeof(beChunkCount));
        packet.insert(packet.end(), reinterpret_cast<char*>(&beChunkSize), reinterpret_cast<char*>(&beChunkSize) + sizeof(beChunkSize));

        if(shouldCompress)
        {
            packet.insert(packet.end(), compressedData.begin(), compressedData.end());
        }
        else
        {
            packet.insert(packet.end(), chunkData.begin(), chunkData.end());
        }

        server->safeSend(clientSocket, 0x02, packet);
    }
}