Teaspeak-Server/file/local_server/LocalFileTransfer.cpp

//
// Created by WolverinDEV on 04/05/2020.
//

#include <cassert>
#include <event2/event.h>
#include <log/LogUtils.h>
#include <random>
#include "./LocalFileProvider.h"
#include "LocalFileProvider.h"

using namespace ts::server::file;
using namespace ts::server::file::transfer;

Buffer* transfer::allocate_buffer(size_t size) {
    auto total_size = sizeof(Buffer) + size;
    auto buffer = (Buffer*) malloc(total_size);
    new (buffer) Buffer{};
    buffer->capacity = size;
    return buffer;
}

void transfer::free_buffer(Buffer* buffer) {
    buffer->~Buffer();
    free(buffer);
}

FileClient::~FileClient() {
    auto head = this->buffer.buffer_head;
    while (head) {
        auto next = head->next;
        free_buffer(head);
        head = next;
    }

    assert(!this->file.file_descriptor);
    assert(!this->file.currently_processing);
    assert(!this->file.next_client);

    assert(!this->networking.event_read);
    assert(!this->networking.event_write);

    assert(this->state == STATE_DISCONNECTED);
}

LocalFileTransfer::LocalFileTransfer(filesystem::LocalFileSystem &fs) : file_system_{fs} {}
LocalFileTransfer::~LocalFileTransfer() = default;

bool LocalFileTransfer::start(const std::deque<std::shared_ptr<NetworkBinding>>& bindings, const std::shared_ptr<pipes::SSL::Options>& ssl_options) {
    assert(ssl_options);
    this->ssl_options_ = ssl_options;

    (void) this->start_client_worker();

    {
        auto start_result = this->start_disk_io();
        switch (start_result) {
            case DiskIOStartResult::SUCCESS:
                break;
            case DiskIOStartResult::OUT_OF_MEMORY:
                logError(LOG_FT, "Failed to start disk worker (Out of memory)");
                goto error_exit_disk;
            default:
                logError(LOG_FT, "Failed to start disk worker ({})", (int) start_result);
                goto error_exit_disk;
        }
    }


    {
        this->network.bindings = bindings;
        auto start_result = this->start_networking();
        switch (start_result) {
            case NetworkingStartResult::SUCCESS:
                break;
            case NetworkingStartResult::OUT_OF_MEMORY:
                logError(LOG_FT, "Failed to start networking (Out of memory)");
                goto error_exit_network;
            case NetworkingStartResult::NO_BINDINGS:
                logError(LOG_FT, "Failed to start networking (No address could be bound)");
                goto error_exit_network;
            default:
                logError(LOG_FT, "Failed to start networking ({})", (int) start_result);
                goto error_exit_network;
        }
    }

    return true;
    error_exit_network:
    this->shutdown_networking();

    error_exit_disk:
    this->shutdown_disk_io();
    this->shutdown_client_worker();
    return false;
}

void LocalFileTransfer::stop() {
    this->shutdown_networking();
    this->shutdown_disk_io();
    this->shutdown_client_worker();
}

std::shared_ptr<ExecuteResponse<TransferInitError, std::shared_ptr<Transfer>>> LocalFileTransfer::initialize_icon_transfer(Transfer::Direction direction, ServerId sid, const TransferInfo &info) {
    return this->initialize_transfer(direction, sid, 0, Transfer::TARGET_TYPE_ICON, info);
}

std::shared_ptr<ExecuteResponse<TransferInitError, std::shared_ptr<Transfer>>> LocalFileTransfer::initialize_avatar_transfer(Transfer::Direction direction, ServerId sid, const TransferInfo &info) {
    return this->initialize_transfer(direction, sid, 0, Transfer::TARGET_TYPE_AVATAR, info);
}

std::shared_ptr<ExecuteResponse<TransferInitError, std::shared_ptr<Transfer>>> LocalFileTransfer::initialize_channel_transfer(Transfer::Direction direction, ServerId sid, ChannelId cid, const TransferInfo &info) {
    return this->initialize_transfer(direction, sid, cid, Transfer::TARGET_TYPE_CHANNEL_FILE, info);
}

std::shared_ptr<ExecuteResponse<TransferInitError, std::shared_ptr<Transfer>>> LocalFileTransfer::initialize_transfer(
        Transfer::Direction direction, ServerId sid, ChannelId cid,
        Transfer::TargetType ttype,
        const TransferInfo &info) {
    auto response = this->create_execute_response<TransferInitError, std::shared_ptr<Transfer>>();

    /* TODO: test for a transfer limit */

    auto transfer = std::make_shared<Transfer>();
    transfer->server_transfer_id = ++this->current_transfer_id;
    transfer->server_id = sid;
    transfer->channel_id = cid;
    transfer->target_type = ttype;
    transfer->direction = direction;

    transfer->client_id = 0; /* must be provided externally */
    transfer->client_transfer_id = 0; /* must be provided externally */

    transfer->server_addresses.reserve(this->network.bindings.size());
    for(auto& binding : this->network.bindings) {
        if(!binding->file_descriptor) continue;

        transfer->server_addresses.emplace_back(Transfer::Address{binding->hostname, net::port(binding->address)});
    }

    transfer->target_file_path = info.file_path;
    transfer->file_offset = info.file_offset;
    transfer->expected_file_size = info.expected_file_size;
    transfer->max_bandwidth = info.max_bandwidth;

    constexpr static std::string_view kTokenCharacters{"abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"};
    for(auto& c : transfer->transfer_key)
        c = kTokenCharacters[transfer_random_token_generator() % kTokenCharacters.length()];
    transfer->transfer_key[0] = (char) 'r'; /* (114) */ /* a non valid SSL header type to indicate that we're using a file transfer key and not doing a SSL handshake */
    transfer->transfer_key[1] = (char) 'a'; /* ( 97) */
    transfer->transfer_key[2] = (char) 'w'; /* (119) */

    transfer->initialized_timestamp = std::chrono::system_clock::now();

    {
        std::lock_guard tlock{this->transfers_mutex};
        this->pending_transfers.push_back(transfer);
    }

    if(auto callback{this->callback_transfer_registered}; callback)
        callback(transfer);

    response->emplace_success(std::move(transfer));
    return response;
}

std::shared_ptr<ExecuteResponse<TransferActionError>> LocalFileTransfer::stop_transfer(transfer_id id, bool flush) {
    auto response = this->create_execute_response<TransferActionError>();

    std::shared_ptr<Transfer> transfer{};
    std::shared_ptr<FileClient> connected_transfer{};

    {
        std::lock_guard tlock{this->transfers_mutex};

        auto ct_it = std::find_if(this->transfers_.begin(), this->transfers_.end(), [&](const std::shared_ptr<FileClient>& t) {
            return t->transfer && t->transfer->server_transfer_id == id;
        });
        if(ct_it != this->transfers_.end())
            connected_transfer = *ct_it;
        else {
            auto t_it = std::find_if(this->pending_transfers.begin(), this->pending_transfers.end(), [&](const std::shared_ptr<Transfer>& t) {
                return t->server_transfer_id == id;
            });
            if(t_it != this->pending_transfers.end()) {
                transfer = *t_it;
                this->pending_transfers.erase(t_it);
            }
        }
    }

    if(!transfer) {
        if(connected_transfer)
            transfer = connected_transfer->transfer;
        else {
            response->emplace_fail(TransferActionError{TransferActionError::UNKNOWN_TRANSFER, ""});
            return response;
        }
    }

    if(connected_transfer) {
        logMessage(LOG_FT, "{} Stopping transfer due to an user request.", connected_transfer->log_prefix());

        std::unique_lock slock{connected_transfer->state_mutex};
        this->disconnect_client(connected_transfer, slock, flush);
    } else {
        logMessage(LOG_FT, "Removing pending file transfer for id {}", id);
    }

    response->emplace_success();
    return response;
}