114 lines
4.1 KiB
C++
114 lines
4.1 KiB
C++
#include <fstream>
|
|
#include <corvusoft/restbed/session.hpp>
|
|
#include <corvusoft/restbed/request.hpp>
|
|
#include <corvusoft/restbed/response.hpp>
|
|
#include <spdlog/spdlog.h>
|
|
#include <stb_image.h>
|
|
#include <stb_image_resize2.h>
|
|
#include <stb_image_write.h>
|
|
#include "server_internal.hxx"
|
|
|
|
static constexpr std::size_t chunk_size = 1024*1024, max_image_size = 1024*1024*50, preview_size=480;
|
|
static const std::set<std::string> image_extension = {".png", ".jpg", ".jpeg", ".tga", ".bmp", ".psd", ".gif", ".jfif", ".pjpeg", ".pjp"};
|
|
|
|
struct UploadInfo {
|
|
Server *server;
|
|
std::shared_lock<std::shared_mutex> node_lock;
|
|
std::size_t to_read;
|
|
std::filesystem::path path;
|
|
std::ofstream file;
|
|
std::shared_ptr<Node> node;
|
|
};
|
|
|
|
void make_preview(const std::shared_ptr<UploadInfo>& info) {
|
|
int x, y, channels;
|
|
auto img = std::unique_ptr<stbi_uc, decltype(&free)>
|
|
{stbi_load(info->path.c_str(), &x, &y, &channels, 0), &free};
|
|
if (!img)
|
|
return;
|
|
|
|
float x_ration = (float)preview_size / (float)x, y_ration = (float)preview_size / (float)y;
|
|
float ratio = std::min(x_ration, y_ration);
|
|
int new_x = (int)((float)(x)*ratio), new_y = (int)((float)(y)*ratio);
|
|
stbir_pixel_layout layout;
|
|
switch (channels) {
|
|
case 1: layout = STBIR_1CHANNEL; break;
|
|
case 2: layout = STBIR_2CHANNEL; break;
|
|
case 3: layout = STBIR_RGB; break;
|
|
case 4: layout = STBIR_RGBA; break;
|
|
default: return;
|
|
}
|
|
|
|
auto rimg = std::unique_ptr<unsigned char, decltype(&free)>
|
|
{stbir_resize_uint8_linear(img.get(), x, y, 0, nullptr, new_x, new_y, 0, layout), &free};
|
|
if (!rimg)
|
|
return;
|
|
|
|
auto png_path = info->path.replace_extension("png");
|
|
if (!stbi_write_png(png_path.c_str(), new_x, new_y, channels, rimg.get(), 0))
|
|
return;
|
|
|
|
info->node->preview = true;
|
|
}
|
|
|
|
void fetch_handler(const std::shared_ptr<restbed::Session> &s, const restbed::Bytes &bytes) {
|
|
std::shared_ptr<UploadInfo> info = s->get("upload");
|
|
|
|
std::size_t read = bytes.size();
|
|
info->to_read -= std::min(read, info->to_read);
|
|
info->file.write((char*)bytes.data(), bytes.size());
|
|
if (info->to_read > 0)
|
|
return s->fetch(std::min(info->to_read, chunk_size), fetch_handler);
|
|
info->file.close();
|
|
s->close(200);
|
|
|
|
std::size_t real_size = std::filesystem::file_size(info->path);
|
|
info->node->size = real_size;
|
|
auto ext = std::filesystem::path{info->node->name}.extension().string();
|
|
if (real_size < max_image_size && image_extension.contains(ext))
|
|
make_preview(info);
|
|
|
|
info->node_lock.unlock();
|
|
info->server->save();
|
|
}
|
|
|
|
void Server::upload(const std::shared_ptr<restbed::Session> &s) {
|
|
const auto req = s->get_request();
|
|
if (!req->has_header("X-Node"))
|
|
return s->close(400, "Missing node");
|
|
if (!req->has_header("X-Token"))
|
|
return s->close(400, "Missing token");
|
|
|
|
if (req->get_header("Transfer-Encoding") == "chunked") {
|
|
spdlog::error("Encountered a chunked upload!");
|
|
return s->close(500, "Sorry but your browser is not supported yet");
|
|
}
|
|
|
|
std::uint64_t node_id = req->get_header("X-Node", 0);
|
|
std::string token = req->get_header("X-Token");
|
|
|
|
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");
|
|
if (!node->file) return s->close(400, "Can't upload to a directory");
|
|
std::size_t to_read = req->get_header("Content-Length", 0);
|
|
auto path = user->user_dir / std::to_string(node->id);
|
|
if (node->preview) {
|
|
node->preview = false;
|
|
std::filesystem::remove(path.replace_extension("png"));
|
|
}
|
|
std::shared_ptr<UploadInfo> info{new UploadInfo{
|
|
.server = this,
|
|
.node_lock = std::shared_lock{user->node_lock},
|
|
.to_read = to_read,
|
|
.path = path,
|
|
.file = std::ofstream{path, std::ios_base::out|std::ios_base::trunc|std::ios_base::binary},
|
|
.node = node
|
|
}};
|
|
s->set("upload", info);
|
|
s->fetch(std::min(to_read, chunk_size), fetch_handler);
|
|
}
|
|
}
|