BareGit
#include "multipart_reader.h"

#include <atomic>
#include <chrono>
#include <cstddef>
#include <cstdint>
#include <filesystem>
#include <fstream>
#include <optional>
#include <string>
#include <string_view>
#include <system_error>
#include <unordered_set>
#include <utility>

#include "form_limits.h"

namespace
{

inline constexpr std::size_t MAX_IMAGE_SIZE = 32 * 1024 * 1024;
inline constexpr std::size_t MAX_TOTAL_IMAGE_SIZE = 72 * 1024 * 1024;

std::atomic<std::uint64_t> NEXT_STAGING_ID = 0;

bool isTextField(std::string_view name)
{
    if(name.starts_with("game.") && name.size() > 5)
    {
        return true;
    }
    static const std::unordered_set<std::string> names = {
        "game",
        "name",
        "rarity",
        "short_description",
        "long_description",
        "source_mode",
        "front_url",
        "foil_url",
        "front_action",
        "foil_action",
        "revision",
        "game_revision",
        "series_id",
        "csrf_token",
    };
    return names.contains(std::string(name));
}

std::optional<std::string> stagedFilename(std::string_view name)
{
    if(name == "front")
    {
        return "upload_front";
    }
    if(name == "foil")
    {
        return "upload_foil";
    }
    if(name == "thumbnail")
    {
        return "upload_thumbnail";
    }
    return std::nullopt;
}

mw::E<std::filesystem::path> createStagingDirectory(
    const std::filesystem::path& card_storage_root)
{
    const std::filesystem::path staging_root =
        card_storage_root / ".staging";
    std::error_code filesystem_error;
    std::filesystem::create_directories(staging_root, filesystem_error);
    if(filesystem_error)
    {
        return std::unexpected(mw::runtimeError(
            "Failed to create the upload staging root: " +
            filesystem_error.message()));
    }

    for(int attempt = 0; attempt < 32; ++attempt)
    {
        const auto ticks = std::chrono::steady_clock::now()
                               .time_since_epoch()
                               .count();
        const std::uint64_t sequence = NEXT_STAGING_ID.fetch_add(1);
        const std::filesystem::path directory = staging_root /
            (std::to_string(ticks) + "-" + std::to_string(sequence));
        filesystem_error.clear();
        if(std::filesystem::create_directory(directory, filesystem_error))
        {
            return directory;
        }
        if(filesystem_error &&
           filesystem_error != std::errc::file_exists)
        {
            return std::unexpected(mw::runtimeError(
                "Failed to create an upload staging directory: " +
                filesystem_error.message()));
        }
    }
    return std::unexpected(mw::runtimeError(
        "Failed to allocate an upload staging directory"));
}

} // namespace

MultipartReader::MultipartReader(std::filesystem::path card_storage_root)
        : card_storage_root_(std::move(card_storage_root))
{}

MultipartReader::~MultipartReader()
{
    if(staging_directory_.empty())
    {
        return;
    }
    std::error_code filesystem_error;
    std::filesystem::remove_all(staging_directory_, filesystem_error);
}

mw::E<CardUpload> MultipartReader::read(
    const httplib::ContentReader& content_reader)
{
    auto staging = createStagingDirectory(card_storage_root_);
    if(!staging)
    {
        return std::unexpected(std::move(staging.error()));
    }
    staging_directory_ = *staging;

    CardUpload upload;
    upload.staging_directory = staging_directory_;
    std::ofstream binary_stream;
    std::string* text_value = nullptr;
    bool ignore_part = false;
    std::size_t part_size = 0;
    std::size_t total_image_size = 0;
    std::size_t total_text_size = 0;
    std::string error_message;
    int error_status = 400;

    const bool read = content_reader(
        [&](const httplib::FormData& part)
        {
            binary_stream.close();
            text_value = nullptr;
            ignore_part = false;
            part_size = 0;

            if(isTextField(part.name))
            {
                if(part.name == "series_id")
                {
                    upload.series_ids.emplace_back();
                    text_value = &upload.series_ids.back();
                    return true;
                }
                auto [position, inserted] =
                    upload.fields.try_emplace(part.name);
                if(!inserted)
                {
                    error_message =
                        "The form field '" + part.name +
                        "' was supplied more than once";
                    return false;
                }
                text_value = &position->second;
                return true;
            }

            const auto filename = stagedFilename(part.name);
            if(!filename)
            {
                error_message =
                    "The form contains an unknown field: " + part.name;
                return false;
            }
            if(part.filename.empty())
            {
                ignore_part = true;
                return true;
            }
            std::optional<std::filesystem::path>* target = nullptr;
            if(part.name == "front")
            {
                target = &upload.front;
            }
            else if(part.name == "foil")
            {
                target = &upload.foil;
            }
            else
            {
                target = &upload.thumbnail;
            }
            if(*target)
            {
                error_message =
                    "The image field '" + part.name +
                    "' was supplied more than once";
                return false;
            }
            *target = staging_directory_ / *filename;
            binary_stream.open(**target, std::ios::binary);
            if(!binary_stream)
            {
                error_message = "Failed to stage an uploaded image";
                error_status = 500;
                return false;
            }
            return true;
        },
        [&](const char* data, std::size_t size)
        {
            part_size += size;
            if(ignore_part)
            {
                total_image_size += size;
                if(part_size > MAX_IMAGE_SIZE ||
                   total_image_size > MAX_TOTAL_IMAGE_SIZE)
                {
                    error_message = "An uploaded image is too large";
                    error_status = 413;
                    return false;
                }
                return true;
            }
            if(text_value != nullptr)
            {
                total_text_size += size;
                if(part_size > MAX_FORM_TEXT_FIELD_SIZE)
                {
                    error_message = "A form field is too large";
                    error_status = 413;
                    return false;
                }
                if(total_text_size > MAX_FORM_TOTAL_TEXT_SIZE)
                {
                    error_message = "Card metadata is too large";
                    error_status = 413;
                    return false;
                }
                text_value->append(data, size);
                return true;
            }
            if(!binary_stream)
            {
                error_message = "The multipart request is malformed";
                return false;
            }
            total_image_size += size;
            if(part_size > MAX_IMAGE_SIZE ||
               total_image_size > MAX_TOTAL_IMAGE_SIZE)
            {
                error_message = "An uploaded image is too large";
                error_status = 413;
                return false;
            }
            binary_stream.write(
                data, static_cast<std::streamsize>(size));
            if(!binary_stream)
            {
                error_message = "Failed to stage an uploaded image";
                error_status = 500;
                return false;
            }
            return true;
        });
    binary_stream.close();

    if(!read)
    {
        if(error_message.empty())
        {
            error_message = "Failed to read the multipart request";
        }
        return std::unexpected(mw::httpError(
            error_status, error_message));
    }
    return upload;
}