184 lines
7.2 KiB
C++
184 lines
7.2 KiB
C++
|
#include <filesystem>
|
||
|
|
#include <string_view>
|
||
|
|
#include <ranges>
|
||
|
|
#include <fstream>
|
||
|
|
#include <deque>
|
||
|
|
#include <charconv>
|
||
|
|
#include <miniz.h>
|
||
|
|
#include <corvusoft/restbed/session.hpp>
|
||
|
|
#include <corvusoft/restbed/request.hpp>
|
||
|
|
#include <corvusoft/restbed/response.hpp>
|
||
|
|
#include "server_internal.hxx"
|
||
|
|
|
||
|
|
void Server::download(const std::shared_ptr<restbed::Session> &s) {
|
||
|
|
const auto req = s->get_request();
|
||
|
|
std::size_t body_len = req->get_header("Content-Length", 0);
|
||
|
|
s->fetch(body_len, [this](const std::shared_ptr<restbed::Session> &s, const restbed::Bytes &b){
|
||
|
|
std::string body{b.cbegin(), b.cend()};
|
||
|
|
if (body.empty())
|
||
|
|
return s->close(400, "empty body");
|
||
|
|
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())
|
||
|
|
return s->close(400, "Missing node");
|
||
|
|
if (token.empty())
|
||
|
|
return s->close(400, "Missing token");
|
||
|
|
|
||
|
|
std::uint64_t node_id;
|
||
|
|
auto res = std::from_chars(node_str.data(), node_str.data() + node_str.size(), node_id);
|
||
|
|
if (res.ec != std::errc{})
|
||
|
|
return s->close(400, "Invalid node");
|
||
|
|
|
||
|
|
check_user() return s->close(400, "Invalid user");
|
||
|
|
{
|
||
|
|
std::shared_lock lock{user->node_lock};
|
||
|
|
auto node = get_node(user, node_id);
|
||
|
|
if (!node) return s->close(400, "Invalid node");
|
||
|
|
|
||
|
|
auto mime = get_mime_type(node->name);
|
||
|
|
|
||
|
|
s->yield(
|
||
|
|
200,
|
||
|
|
"",
|
||
|
|
std::multimap<std::string, std::string>{
|
||
|
|
{"Content-Type", mime},
|
||
|
|
{"Content-Length", std::to_string(node->size)},
|
||
|
|
{"Content-Disposition", "attachment; filename=\"" + node->name + "\""}
|
||
|
|
},
|
||
|
|
[&](const std::shared_ptr<restbed::Session>& s) {
|
||
|
|
std::shared_lock lock{user->node_lock};
|
||
|
|
restbed::Bytes buf(1024*1024*4, 0);
|
||
|
|
std::ifstream f{user->user_dir / std::to_string(node->id)};
|
||
|
|
while (!f.eof()) {
|
||
|
|
f.read((char*)buf.data(), buf.size());
|
||
|
|
buf.resize(f.gcount());
|
||
|
|
s->yield(buf);
|
||
|
|
}
|
||
|
|
s->close();
|
||
|
|
}
|
||
|
|
);
|
||
|
|
}
|
||
|
|
});
|
||
|
|
}
|
||
|
|
|
||
|
|
size_t zip_write_func(void *pOpaque, mz_uint64 _file_ofs, const void *pBuf, size_t n) {
|
||
|
|
auto s = (restbed::Session*)pOpaque;
|
||
|
|
if (n > 0) {
|
||
|
|
restbed::Bytes buf(n, 0);
|
||
|
|
std::memcpy(buf.data(), pBuf, n);
|
||
|
|
std::stringstream ss;
|
||
|
|
ss << std::hex << n;
|
||
|
|
s->yield(ss.str() + "\r\n");
|
||
|
|
s->yield(buf);
|
||
|
|
s->yield("\r\n");
|
||
|
|
}
|
||
|
|
return n;
|
||
|
|
}
|
||
|
|
|
||
|
|
void Server::download_multi(const std::shared_ptr<restbed::Session> &s) {
|
||
|
|
const auto req = s->get_request();
|
||
|
|
const auto body_len = req->get_header("Content-Length", 0);
|
||
|
|
s->fetch(body_len, [this](const std::shared_ptr<restbed::Session> &s, const restbed::Bytes &b){
|
||
|
|
std::string body{b.cbegin(), b.cend()};
|
||
|
|
if (body.empty())
|
||
|
|
return s->close(400, "empty body");
|
||
|
|
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())
|
||
|
|
return s->close(400, "Missing nodes");
|
||
|
|
if (token.empty())
|
||
|
|
return s->close(400, "Missing token");
|
||
|
|
|
||
|
|
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{})
|
||
|
|
return s->close(400, "Invalid node " + std::string{std::string_view{part}});
|
||
|
|
node_ids.push_back(node_id);
|
||
|
|
}
|
||
|
|
|
||
|
|
check_user() return s->close(400, "Invalid user");
|
||
|
|
{
|
||
|
|
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) return s->close(400, "Invalid node " + std::to_string(node_id));
|
||
|
|
nodes.push_back(node);
|
||
|
|
}
|
||
|
|
|
||
|
|
s->yield(
|
||
|
|
200,
|
||
|
|
"",
|
||
|
|
std::multimap<std::string, std::string>{
|
||
|
|
{"Content-Type", "application/zip"},
|
||
|
|
{"Content-Disposition", "attachment; filename=\"files.zip\""},
|
||
|
|
{"Transfer-Encoding", "chunked"}
|
||
|
|
}
|
||
|
|
);
|
||
|
|
std::thread zip_thread{[nodes = nodes, user = user, s = s] {
|
||
|
|
std::shared_lock lock{user->node_lock};
|
||
|
|
|
||
|
|
mz_zip_archive archive;
|
||
|
|
mz_zip_zero_struct(&archive);
|
||
|
|
archive.m_pWrite = zip_write_func;
|
||
|
|
archive.m_pIO_opaque = s.get();
|
||
|
|
|
||
|
|
mz_zip_writer_init_v2(&archive, 0, MZ_ZIP_FLAG_WRITE_ZIP64);
|
||
|
|
|
||
|
|
std::deque<std::pair<std::shared_ptr<Node>, std::filesystem::path>> todo;
|
||
|
|
for (const auto &node : nodes)
|
||
|
|
todo.emplace_back(node, std::filesystem::path{});
|
||
|
|
|
||
|
|
auto handle_file = [&user, &archive](const std::pair<std::shared_ptr<Node>, std::filesystem::path> &i) {
|
||
|
|
auto path = i.second / i.first->name;
|
||
|
|
auto real_path = user->user_dir / std::to_string(i.first->id);
|
||
|
|
mz_zip_writer_add_file(&archive, path.c_str(), real_path.c_str(), nullptr, 0, MZ_DEFAULT_COMPRESSION);
|
||
|
|
};
|
||
|
|
|
||
|
|
while (!todo.empty()) {
|
||
|
|
const auto &node = todo.front();
|
||
|
|
if (node.first->file) {
|
||
|
|
handle_file(node);
|
||
|
|
} else {
|
||
|
|
auto path = node.second / node.first->name;
|
||
|
|
auto dir_path = path.string() + "/";
|
||
|
|
mz_zip_writer_add_mem(&archive, dir_path.c_str(), nullptr, 0, 0);
|
||
|
|
for (const auto &child : node.first->children) {
|
||
|
|
auto p = std::make_pair(child, path);
|
||
|
|
if (child->file)
|
||
|
|
handle_file(p);
|
||
|
|
else
|
||
|
|
todo.push_back(p);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
todo.pop_front();
|
||
|
|
}
|
||
|
|
|
||
|
|
mz_zip_writer_finalize_archive(&archive);
|
||
|
|
mz_zip_writer_end(&archive);
|
||
|
|
|
||
|
|
s->close("0\r\n\r\n");
|
||
|
|
}};
|
||
|
|
zip_thread.detach();
|
||
|
|
}
|
||
|
|
});
|
||
|
|
}
|