fileserver/src/server/download.cxx

#include <filesystem>
#include <string_view>
#include <ranges>
#include <charconv>
#include <deque>
#include "../util/miniz.hxx"
#include "server_internal.hxx"

net::awaitable<void> Server::download(tcp_stream &s, const http::request<http::string_body> &req) {
    auto body = req.body();
    if (body.empty()) { co_await send_error(s, req, "Empty body"); co_return; }

    std::string node_str, token;
    for (const auto part : std::views::split(body, '&')) {
        std::string_view part_view{part};
        auto equal_pos = part_view.find_first_of('=');
        auto key = part_view.substr(0, equal_pos);
        if (key == "node")
            node_str = part_view.substr(equal_pos+1);
        else if (key == "token")
            token = part_view.substr(equal_pos+1);
    }

    if (node_str.empty()) { co_await send_error(s, req, "Missing node"); co_return; }
    if (token.empty()) { co_await send_error(s, req, "Missing token"); co_return; }

    std::uint64_t node_id;
    auto fc_res = std::from_chars(node_str.data(), node_str.data() + node_str.size(), node_id);
    if (fc_res.ec != std::errc{}) { co_await send_error(s, req, "Invalid node"); co_return; }

    check_user() { co_await send_error(s, req, "Invalid user"); co_return; }

    std::shared_lock lock{user->node_lock};
    auto node = get_node(user, node_id);
    if (!node) { co_await send_error(s, req, "Invalid node"); co_return; }

    s.expires_never();

    beast::error_code ec;
    auto mime = get_mime_type(node->name);
    auto res = create_response<http::status::ok, http::file_body, http::string_body>(req);
    res.content_length(node->size);
    res.set(http::field::content_type, mime);
    res.set(http::field::content_disposition, "attachment; filename=\"" + node->name + "\"");
    res.body().open((user->user_dir / std::to_string(node->id)).c_str(), beast::file_mode::read, ec);
    co_await http::async_write(s, res, net::use_awaitable);

    s.expires_after(std::chrono::seconds(30));
}

struct Zip : public ZipArchive {
    tcp_stream *s;
    http::response<http::buffer_body> *body;
    http::response_serializer<http::buffer_body> *sr;

    Zip(tcp_stream *s, http::response<http::buffer_body> *body, http::response_serializer<http::buffer_body> *sr)
        : s(s), body(body), sr(sr) {}

protected:
    net::awaitable<void> write(const void *pVoid, size_t n) override {
        if (n == 0)
            co_return;

        body->body().data = (void*)pVoid;
        body->body().size = n;
        body->body().more = true;

        try {
            co_await http::async_write(*s, *sr, net::use_awaitable);
        } catch (beast::system_error &se) {
            if (se.code() != http::error::need_buffer)
                throw;
        }
    }
};

net::awaitable<void> Server::download_multi(tcp_stream &s, const http::request<http::string_body> &req) {
    auto body = req.body();
    if (body.empty()) { co_await send_error(s, req, "Empty body"); co_return; }

    std::string nodes_str, token;
    for (const auto part : std::views::split(body, '&')) {
        std::string_view part_view{part};
        auto equal_pos = part_view.find_first_of('=');
        auto key = part_view.substr(0, equal_pos);
        if (key == "nodes")
            nodes_str = part_view.substr(equal_pos+1);
        else if (key == "token")
            token = part_view.substr(equal_pos+1);
    }

    if (nodes_str.empty()) { co_await send_error(s, req, "Missing nodes"); co_return; }
    if (token.empty()) { co_await send_error(s, req, "Missing token"); co_return; }

    std::vector<std::uint64_t> node_ids;
    for (const auto part : std::views::split(nodes_str, '.')) {
        std::uint64_t node_id;
        auto res = std::from_chars(part.data(), part.data() + part.size(), node_id);
        if (res.ec != std::errc{}) { co_await send_error(s, req, "Invalid node " + std::string{std::string_view{part}}); co_return; }
        node_ids.push_back(node_id);
    }

    check_user() { co_await send_error(s, req, "Invalid user"); co_return; }

    std::shared_lock lock{user->node_lock};
    std::vector<std::shared_ptr<Node>> nodes;
    for (auto node_id : node_ids) {
        auto node = get_node(user, node_id);
        if (!node) { co_await send_error(s, req, "Invalid node " + std::to_string(node_id)); co_return; }
        nodes.push_back(node);
    }

    auto res = create_response<http::status::ok, http::buffer_body, http::string_body>(req);
    res.chunked(true);
    res.set(http::field::content_type, "application/zip");
    res.set(http::field::content_disposition, "attachment; filename=\"files.zip\"");
    res.body().data = nullptr;
    res.body().more = true;

    s.expires_never();

    http::response_serializer<http::buffer_body> sr{res};

    co_await http::async_write_header(s, sr, net::use_awaitable);

    Zip zip{&s, &res, &sr};

    std::deque<std::pair<std::shared_ptr<Node>, std::filesystem::path>> todo;
    for (const auto &node : nodes)
        todo.emplace_back(node, std::filesystem::path{});

    while (!todo.empty()) {
        const auto &node = todo.front();
        if (node.first->file) {
            auto path = node.second / node.first->name;
            auto real_path = user->user_dir / std::to_string(node.first->id);
            co_await zip.add_file(real_path, path);
        } else {
            auto path = node.second / node.first->name;
            auto dir_path = path.string() + "/";
            co_await zip.add_dir(dir_path);
            for (const auto &child : node.first->children) {
                auto p = std::make_pair(child, path);
                todo.push_back(p);
            }
        }
        todo.pop_front();
    }
    co_await zip.end();

    res.body().data = nullptr;
    res.body().more = false;
    co_await http::async_write(s, sr, net::use_awaitable);

    s.expires_after(std::chrono::seconds(30));
}