file_manager.cpp

#include "file_manager.hpp"
#include "Polyweb/binary.hpp"
#include "json.hpp"
#include "util.hpp"
#include <FL/Fl_Native_File_Chooser.H>
#include <FL/fl_callback_macros.H>
#include <exception>
#include <inttypes.h>
#include <iomanip>
#include <iostream>
#include <sstream>
#include <utility>
#ifdef _WIN32
    #include "theme.hpp"
    #include <FL/x.H>
#endif

using nlohmann::json;

constexpr auto PROGRESS_UPDATE_INTERVAL = std::chrono::milliseconds(250);

ProgressWindow::ProgressWindow(const std::string& path, uint64_t value, uint64_t size, std::function<void()> cancel_cb):
    Fl_Double_Window(500, 120, "File Transfer") {
    auto path_box = new Fl_Box(10, 10, w() - 20, 30);
    path_box->copy_label(("Copying " + path + "...").c_str());

    progress = new Fl_Progress(10, 45, w() - 20, 30);
    progress->selection_color(FL_SELECTION_COLOR);
    progress->maximum(size);
    this->value(value);

    auto cancel_button = new Fl_Button(w() - 85, h() - 40, 75, 30, "Cancel");
    FL_INLINE_CALLBACK_2(cancel_button, Fl_Double_Window*, window, this, std::function<void()>, cancel_cb, cancel_cb, {
        window->hide();
        cancel_cb();
    });

    FL_INLINE_CALLBACK_2(this, Fl_Double_Window*, window, this, std::function<void()>, cancel_cb, cancel_cb, {
        window->hide();
        cancel_cb();
    });

    end();
}

void ProgressWindow::value(float value) {
    progress->value(value);

    std::ostringstream ss;
    ss << std::fixed << std::setprecision(1) << value / progress->maximum() * 100 << "% (" << std::setprecision(2) << value / 1e+6f << '/' << progress->maximum() / 1e+6f << " MB)";
    progress->copy_label(ss.str().c_str());
}

void FileManager::on_buffered_amount_low() {
    if (!buffered_amount_low_running) {
        buffered_amount_low_running = true;

        std::lock_guard<std::mutex> lock(mutex);
        while (channel->isOpen() && channel->bufferedAmount() <= 512 * 1024 && !outgoing_transfers.empty()) {
            for (auto transfer_it = outgoing_transfers.begin(); transfer_it != outgoing_transfers.end();) {
                if (transfer_it->second->started) {
                    rtc::binary message(chunk_size + 4);
                    if (transfer_it->second->file.read((char*) message.data() + 4, chunk_size).bad() && !transfer_it->second->file.eof()) {
                        uint32_t id = transfer_it->first;
                        transfer_it = outgoing_transfers.erase(transfer_it);
                        cancel_transfer(id);
                        awake([id]() {
                            fl_alert("Error: File transfer #%" PRIu32 " failed", id);
                        });
                        continue;
                    }
                    message.resize(transfer_it->second->file.gcount() + 4);
#if BYTE_ORDER == BIG_ENDIAN
                    memcpy(message.data(), &transfer_it->first, 4);
#else
                    pw::reverse_memcpy(message.data(), &transfer_it->first, 4);
#endif
                    if (!try_send(channel, std::move(message))) {
                        // The connection is going away, so there is nothing to
                        // report to the peer and no point alerting per transfer
                        transfer_it = outgoing_transfers.erase(transfer_it);
                        continue;
                    }
                    transfer_it->second->sent += transfer_it->second->file.gcount();

                    if (auto now = std::chrono::steady_clock::now(); now - transfer_it->second->last_progress_update >= PROGRESS_UPDATE_INTERVAL) {
                        awake([weak_transfer = std::weak_ptr<OutgoingTransfer>(transfer_it->second)]() {
                            if (auto transfer = weak_transfer.lock()) {
                                if (auto window = transfer->progress_window.load()) {
                                    window->value(transfer->sent);
                                }
                            }
                        });
                        transfer_it->second->last_progress_update = now;
                    }

                    if (transfer_it->second->file.eof() || transfer_it->second->sent >= transfer_it->second->size) {
                        transfer_it = outgoing_transfers.erase(transfer_it);
                        continue;
                    }
                }
                ++transfer_it;
            }
        }

        buffered_amount_low_running = false;
    }
}

void FileManager::on_binary_message(rtc::binary message) {
    if (message.size() > 4) {
        uint32_t id;
#if BYTE_ORDER == BIG_ENDIAN
        memcpy(&id, message.data(), 4);
#else
        pw::reverse_memcpy(&id, message.data(), 4);
#endif

        std::lock_guard<std::mutex> lock(mutex);
        if (auto transfer_it = incoming_transfers.find(id); transfer_it != incoming_transfers.end()) {
            if (transfer_it->second->file.write((const char*) message.data() + 4, message.size() - 4).fail()) {
                uint32_t id = transfer_it->first;
                incoming_transfers.erase(transfer_it);
                cancel_transfer(id);
                awake([id]() {
                    fl_alert("Error: File transfer #%" PRIu32 " failed", id);
                });
                return;
            }
            transfer_it->second->received += message.size() - 4;

            if (transfer_it->second->progress_window) {
                if (auto now = std::chrono::steady_clock::now(); now - transfer_it->second->last_progress_update >= PROGRESS_UPDATE_INTERVAL) {
                    awake([weak_transfer = std::weak_ptr<IncomingTransfer>(transfer_it->second)]() {
                        if (auto transfer = weak_transfer.lock()) {
                            if (auto window = transfer->progress_window.load()) {
                                window->value(transfer->received);
                            }
                        }
                    });
                    transfer_it->second->last_progress_update = now;
                }
            }

            if (transfer_it->second->received >= transfer_it->second->size) {
                transfer_it = incoming_transfers.erase(transfer_it);
                return;
            }
        }
    }
}

void FileManager::on_string_message(rtc::string message) {
    std::lock_guard<std::mutex> lock(mutex);
    try {
        json message_json = json::parse(message);
        if (message_json["type"] == "transferready") {
            if (message_json.contains("size")) {
                if (auto transfer_it = incoming_transfers.find(message_json["id"]); transfer_it != incoming_transfers.end()) {
                    transfer_it->second->size = message_json["size"];
                    // Must not block: this runs on a libdatachannel thread, inside
                    // a callback that holds the channel's callback mutex and while
                    // holding our own. Waiting for the main thread here deadlocks
                    // against ~FileManager, which takes both from the other side
                    awake([this, id = transfer_it->first, weak_transfer = std::weak_ptr<IncomingTransfer>(transfer_it->second)]() {
                        if (auto transfer = weak_transfer.lock()) {
                            auto window = new ProgressWindow(transfer->path, transfer->received, transfer->size, [this, id]() {
                                std::lock_guard<std::mutex> lock(mutex);
                                incoming_transfers.erase(id);
                                cancel_transfer(id);
                            });
                            transfer->progress_window = window;
                            window->show();
#ifdef _WIN32
                            Fl::flush();
                            set_window_dark_mode(fl_xid(window));
#endif
                        }
                    });
                }
            } else {
                if (auto transfer_it = outgoing_transfers.find(message_json["id"]); transfer_it != outgoing_transfers.end()) {
                    // Gates on_buffered_amount_low(), which used to wait on the
                    // progress window existing. See the note above on why creating
                    // that window must not block this thread
                    transfer_it->second->started = true;
                    awake([this, id = transfer_it->first, weak_transfer = std::weak_ptr<OutgoingTransfer>(transfer_it->second)]() {
                        if (auto transfer = weak_transfer.lock()) {
                            auto window = new ProgressWindow(transfer->path, transfer->sent, transfer->size, [this, id]() {
                                std::lock_guard<std::mutex> lock(mutex);
                                outgoing_transfers.erase(id);
                                cancel_transfer(id);
                            });
                            transfer->progress_window = window;
                            window->show();
#ifdef _WIN32
                            Fl::flush();
                            set_window_dark_mode(fl_xid(window));
#endif
                        }
                    });

                    while (channel->isOpen() && channel->bufferedAmount() <= 512 * 1024) {
                        rtc::binary message(chunk_size + 4);
                        if (transfer_it->second->file.read((char*) message.data() + 4, chunk_size).bad() && !transfer_it->second->file.eof()) {
                            uint32_t id = transfer_it->first;
                            outgoing_transfers.erase(transfer_it);
                            cancel_transfer(id);
                            awake([id]() {
                                fl_alert("Error: File transfer #%" PRIu32 " failed", id);
                            });
                            break;
                        }
                        message.resize(transfer_it->second->file.gcount() + 4);
#if BYTE_ORDER == BIG_ENDIAN
                        memcpy(message.data(), &transfer_it->first, 4);
#else
                        pw::reverse_memcpy(message.data(), &transfer_it->first, 4);
#endif
                        if (!try_send(channel, std::move(message))) {
                            outgoing_transfers.erase(transfer_it);
                            break;
                        }
                        transfer_it->second->sent += transfer_it->second->file.gcount();

                        if (auto now = std::chrono::steady_clock::now(); now - transfer_it->second->last_progress_update >= PROGRESS_UPDATE_INTERVAL) {
                            awake([weak_transfer = std::weak_ptr<OutgoingTransfer>(transfer_it->second)]() {
                                if (auto transfer = weak_transfer.lock()) {
                                    if (auto window = transfer->progress_window.load()) {
                                        window->value(transfer->sent);
                                    }
                                }
                            });
                            transfer_it->second->last_progress_update = now;
                        }

                        if (transfer_it->second->file.eof() || transfer_it->second->sent >= transfer_it->second->size) {
                            outgoing_transfers.erase(transfer_it);
                            break;
                        }
                    }
                }
            }
        } else if (message_json["type"] == "canceltransfer") {
            // Report only transfers that were still running. An id below transfer_id
            // just means we allocated it at some point, so testing that alone claims
            // a transfer failed when a stray or duplicate cancel arrives for one that
            // already finished. Both erases run; ids are unique across the two maps
            uint32_t id = message_json["id"].get<uint32_t>();
            size_t erased = incoming_transfers.erase(id) + outgoing_transfers.erase(id);
            if (erased) {
                awake([id]() {
                    fl_alert("File transfer #%" PRIu32 " has been cancelled.", id);
                });
            }
        }
    } catch (const std::exception& e) {
        std::cerr << "Error parsing message: " << e.what() << std::endl;
    }
}

void FileManager::cancel_transfer(uint32_t id) {
    json message = {
        {"type", "canceltransfer"},
        {"id", id},
    };
    try_send(channel, message.dump());
}

FileManager::~FileManager() {
    *alive = false;

    // Both of these wait for a callback already in flight to return, which is
    // only safe because neither of them blocks on the main thread anymore
    channel->onMessage(nullptr, nullptr);
    channel->onBufferedAmountLow(nullptr);

    std::unique_lock<std::mutex> lock(mutex);
    for (auto& transfer : incoming_transfers) {
        if (auto window = transfer.second->progress_window.exchange(nullptr)) {
            awake([window]() { // By value, so the capture cannot outlive this loop
                delete window;
            },
                true);
        }
        cancel_transfer(transfer.first);
    }
    for (auto& transfer : outgoing_transfers) {
        if (auto window = transfer.second->progress_window.exchange(nullptr)) {
            awake([window]() {
                delete window;
            },
                true);
        }
        cancel_transfer(transfer.first);
    }
    incoming_transfers.clear();
    outgoing_transfers.clear();
    lock.unlock();
}

void FileManager::upload() {
    // The chooser pumps the event loop (the GTK backend runs Fl::check() while its
    // own dialog is up), so the connection can be torn down and this FileManager
    // destroyed before it returns. Nothing below may touch a member until checked
    auto alive_copy = alive;

    Fl_Native_File_Chooser chooser(Fl_Native_File_Chooser::BROWSE_FILE);
    chooser.title("Choose File");
    int result = chooser.show();
    if (!*alive_copy) {
        return;
    }
    if (result == -1) {
        fl_alert("Error: Failed to open file chooser: %s", chooser.errmsg());
        return;
    } else if (result != 0 || !chooser.filename() || !*chooser.filename()) {
        return; // Cancelled
    }
    const char* filename = chooser.filename();

    auto transfer = std::make_shared<OutgoingTransfer>();
    transfer->file.open((transfer->path = filename), std::ios::binary | std::ios::ate);
    if (!transfer->file.is_open()) {
        fl_alert("Error: Failed to open file");
        return;
    }

    uint64_t size = transfer->size = transfer->file.tellg();
    transfer->file.seekg(0, std::ios::beg);

    mutex.lock();
    uint32_t id = transfer_id++;
    outgoing_transfers[id] = std::move(transfer);
    mutex.unlock();

    json message = {
        {"type", "requesttransfer"},
        {"id", id},
        {"size", size},
    };
    if (!try_send(channel, message.dump())) {
        {
            std::lock_guard<std::mutex> lock(mutex);
            outgoing_transfers.erase(id);
        }
        // Outside the lock: fl_alert() pumps the event loop, which can re-enter
        // us through a progress window's cancel callback
        fl_alert("Error: There is no active connection");
    }
}

void FileManager::download() {
    auto alive_copy = alive;

    Fl_Native_File_Chooser chooser(Fl_Native_File_Chooser::BROWSE_SAVE_FILE);
    chooser.title("Save File");
    chooser.options(Fl_Native_File_Chooser::SAVEAS_CONFIRM | Fl_Native_File_Chooser::NEW_FOLDER);
    int result = chooser.show();
    if (!*alive_copy) {
        return;
    }
    if (result == -1) {
        fl_alert("Error: Failed to open file chooser: %s", chooser.errmsg());
        return;
    } else if (result != 0 || !chooser.filename() || !*chooser.filename()) {
        return; // Cancelled
    }
    const char* filename = chooser.filename();

    auto transfer = std::make_shared<IncomingTransfer>();
    transfer->file.open((transfer->path = filename), std::ios::binary);
    if (!transfer->file.is_open()) {
        fl_alert("Error: Failed to open file");
        return;
    }

    mutex.lock();
    uint32_t id = transfer_id++;
    incoming_transfers[id] = std::move(transfer);
    mutex.unlock();

    json message = {
        {"type", "requesttransfer"},
        {"id", id},
    };
    if (!try_send(channel, message.dump())) {
        {
            std::lock_guard<std::mutex> lock(mutex);
            incoming_transfers.erase(id);
        }
        fl_alert("Error: There is no active connection");
    }
}