file_transfer.h

Go to the documentation of this file.

#pragma once
 
#include "util/rats_export.h"
#include "node/peer_network.h"
#include "peer/peer.h"
#include "core/bytes.h"
#include "peer/peer_id.h"
#include "util/fs.h"
 
extern "C" {
#include "sha256.h"
}
 
#include <atomic>
#include <chrono>
#include <condition_variable>
#include <cstdint>
#include <functional>
#include <memory>
#include <mutex>
#include <queue>
#include <string>
#include <thread>
#include <unordered_map>
#include <vector>
 
namespace librats {
 
class RATS_API FileTransfer final : public Subsystem {
public:
    struct Config {
        uint32_t    chunk_size           = 64 * 1024;        
        uint32_t    window_bytes         = 4 * 1024 * 1024;  
        uint32_t    progress_interval    = 256 * 1024;       
        uint32_t    transfer_timeout_secs = 60;              
        uint32_t    worker_threads       = 4;                
        uint32_t    disk_threads         = 4;                
        bool        verify_integrity     = true;             
        std::string temp_directory       = ".";              
    };
 
    enum class Status   { Pending, Active, Paused, Completed, Failed, Cancelled };
    enum class Direction { Sending, Receiving };
 
    struct FileEntry {
        std::string relative_path;  
        uint64_t    size = 0;
    };
 
    struct Offer {
        PeerId                 from;
        uint64_t               id = 0;
        std::string            name;             
        uint64_t               size = 0;         
        bool                   is_directory = false;
        std::vector<FileEntry> files;            
    };
 
    struct Progress {
        uint64_t   id = 0;
        PeerId     peer;
        Direction  direction = Direction::Sending;
        Status     status = Status::Pending;
        uint64_t   bytes_transferred = 0;
        uint64_t   total_bytes = 0;
        uint32_t   files_completed = 0;
        uint32_t   total_files = 0;
 
        double transfer_rate_bps = 0.0;  
        double average_rate_bps  = 0.0;  
        std::chrono::milliseconds elapsed{0};                  
        std::chrono::milliseconds estimated_time_remaining{0}; 
 
        double percent() const {
            if (total_bytes == 0) return status == Status::Completed ? 100.0 : 0.0;
            return static_cast<double>(bytes_transferred) / static_cast<double>(total_bytes) * 100.0;
        }
    };
 
    struct Stats {
        uint64_t bytes_sent = 0, bytes_received = 0;
        uint64_t completed = 0, failed = 0;
    };
 
    using OfferHandler    = std::function<void(const Offer&)>;
    using ProgressHandler = std::function<void(const Progress&)>;
    using CompleteHandler = std::function<void(uint64_t id, bool success, const std::string& path)>;
 
    explicit FileTransfer(std::string temp_dir = ".");
    explicit FileTransfer(Config config);
    ~FileTransfer() override;
 
    void on_offer(OfferHandler handler)       { offer_handler_ = std::move(handler); }
    void on_progress(ProgressHandler handler) { progress_handler_ = std::move(handler); }
    void on_complete(CompleteHandler handler) { complete_handler_ = std::move(handler); }
 
    uint64_t send_file(const PeerId& to, const std::string& path);
    uint64_t send_directory(const PeerId& to, const std::string& dir_path);
 
    void accept(const PeerId& from, uint64_t id, const std::string& dest_path);
    void reject(const PeerId& from, uint64_t id);
 
    bool cancel(const PeerId& peer, uint64_t id);
    bool pause(const PeerId& peer, uint64_t id);
    bool resume(const PeerId& peer, uint64_t id);
 
    Stats stats() const;
 
    static bool is_safe_relative_path(const std::string& p);
 
    void attach(NodeContext& ctx) override;
    void start() override;
    void stop() override;
 
private:
    // Smoothed throughput + elapsed/ETA tracking for one transfer. Accessed only
    // under the owning transfer's mutex. The clock starts lazily on the first
    // sample (first byte activity); a long gap (pause/idle) is treated as a
    // discontinuity so it doesn't register as a throughput dip.
    struct RateTracker {
        using clock = std::chrono::steady_clock;
        clock::time_point start{};      
        clock::time_point mark{};       
        uint64_t          mark_bytes = 0;
        double            rate_bps   = 0.0;  
 
        void sample(uint64_t bytes, clock::time_point now) {
            constexpr int64_t kMinMs = 250, kMaxMs = 2000;
            constexpr double  kAlpha = 0.4;
            if (start == clock::time_point{}) { start = mark = now; mark_bytes = bytes; return; }
            const int64_t dt = std::chrono::duration_cast<std::chrono::milliseconds>(now - mark).count();
            if (dt < kMinMs) return;                              // too soon: hold the prior rate
            if (dt > kMaxMs) { mark = now; mark_bytes = bytes; return; }  // gap: reset, don't pollute
            const double inst = static_cast<double>(bytes - mark_bytes) * 1000.0 / static_cast<double>(dt);
            rate_bps = rate_bps <= 0.0 ? inst : kAlpha * inst + (1.0 - kAlpha) * rate_bps;
            mark = now; mark_bytes = bytes;
        }
 
        void fill(Progress& p, uint64_t bytes, uint64_t total, clock::time_point now) const {
            if (start == clock::time_point{}) return;            // not live yet
            const auto ms = std::chrono::duration_cast<std::chrono::milliseconds>(now - start);
            p.elapsed = ms;
            p.transfer_rate_bps = rate_bps;
            if (ms.count() > 0)
                p.average_rate_bps = static_cast<double>(bytes) * 1000.0 / static_cast<double>(ms.count());
            if (rate_bps > 1.0 && total > bytes)
                p.estimated_time_remaining = std::chrono::milliseconds(
                    static_cast<int64_t>(static_cast<double>(total - bytes) / rate_bps * 1000.0));
        }
    };
 
    // ── Per-transfer state (held via shared_ptr so workers/handlers keep it
    //    alive past map removal) ──────────────────────────────────────────────
    struct Outgoing {
        uint64_t                 id = 0;
        PeerId                   peer;
        std::string              name;
        std::string              root;          
        bool                     is_directory = false;
        std::vector<FileEntry>   files;         
        std::vector<std::string> sources;       
        uint64_t                 total_bytes = 0;
 
        std::mutex               mtx;
        std::condition_variable  cv;
        size_t                   cur_file = 0;
        uint64_t                 cur_offset = 0;
        uint64_t                 bytes_done = 0;
        uint64_t                 acked = 0;
        uint32_t                 files_done = 0;
        Status                   status = Status::Pending;
        bool                     worker_active = false;
        bool                     finished = false;
        sha256_context_t         hash{};
        std::chrono::steady_clock::time_point last_activity{};
        RateTracker              rate;
    };
 
    struct IncomingFile {
        std::string relative_path;
        uint64_t    size = 0;
        std::string final_path;
        std::string temp_path;
        uint64_t    enqueued = 0;   
        uint64_t    received = 0;   
        bool        temp_created = false;
        bool        finalized = false;
    };
 
    // A unit of disk work queued by the reactor for the writer pool. A pure data
    // job carries chunk bytes; a file-end job carries the sender's SHA-256 so the
    // writer can verify + finalize once all of the file's data is on disk.
    struct WriteJob {
        uint32_t fidx = 0;
        uint64_t offset = 0;
        Bytes    data;                   
        bool     is_file_end = false;
        uint8_t  sha[SHA256_HASH_SIZE]{};
    };
 
    struct Incoming {
        uint64_t                   id = 0;
        PeerId                     peer;
        std::string                name;
        bool                       is_directory = false;
        std::string                dest_root;
        std::vector<IncomingFile>  files;
 
        std::mutex                 mtx;
        size_t                     recv_file = 0;   
        uint64_t                   bytes_done = 0;  
        uint64_t                   last_ack = 0;
        uint32_t                   files_done = 0;
        Status                     status = Status::Pending;
        bool                       finished = false;
        std::chrono::steady_clock::time_point last_activity{};
        RateTracker                rate;
 
        // ── async disk writer (all fields below the queue are writer-thread only) ─
        std::queue<WriteJob>       wq;              
        uint64_t                   queued_bytes = 0;
        bool                       scheduled = false;
        FileStream                 out;             
        size_t                     out_idx = SIZE_MAX;
        int                        hashing_file = -1;
        sha256_context_t           hash{};
 
        ~Incoming();  
    };
 
    // ── message handling (reactor thread) ─────────────────────────────────────
    void on_message(const Peer& peer, ByteView payload);
    void handle_offer(const PeerId& from, uint64_t id, bool is_dir, uint64_t total,
                      std::string name, std::vector<FileEntry> files);
    void handle_chunk(const PeerId& from, uint64_t id, uint32_t fidx, uint64_t offset,
                      ByteView data);
    void handle_file_end(const PeerId& from, uint64_t id, uint32_t fidx, const uint8_t* sha);
 
    // ── sending ──────────────────────────────────────────────────────────────
    uint64_t start_send(std::shared_ptr<Outgoing> t);
    void     queue_send(uint64_t id);
    void     worker_loop();
    void     run_send(const std::shared_ptr<Outgoing>& t);
 
    // ── receiving ────────────────────────────────────────────────────────────
    // The reactor thread only validates + copies a chunk into the transfer's write
    // queue; a disk-writer thread does the blocking write, hashing and finalize off
    // the reactor. `schedule_writer` hands a transfer to the pool (call with the
    // transfer's mutex NOT held — pool lock is always taken after the transfer's).
    void schedule_writer(const std::shared_ptr<Incoming>& t);
    void disk_worker_loop();
    void drain_writes(const std::shared_ptr<Incoming>& t);
    void process_data(const std::shared_ptr<Incoming>& t, WriteJob& job);
    void process_file_end(const std::shared_ptr<Incoming>& t, WriteJob& job);
 
    // ── lifecycle helpers ─────────────────────────────────────────────────────
    void maintenance_loop();
    void finish_outgoing(const std::shared_ptr<Outgoing>& t, bool success);
    void finish_incoming(const std::shared_ptr<Incoming>& t, bool success, const std::string& error);
    void emit_progress(const std::shared_ptr<Outgoing>& t);
    void emit_progress(const std::shared_ptr<Incoming>& t);
 
    std::shared_ptr<Outgoing> find_outgoing(uint64_t id) const;
    std::shared_ptr<Incoming> find_incoming(const PeerId& peer, uint64_t id) const;
 
    void send_to(const PeerId& peer, const Bytes& msg) {
        if (network_) network_->send(peer, MessageType::FileChunk, ByteView(msg));
    }
    void send_simple(const PeerId& peer, uint8_t op, uint64_t id);
    void send_complete(const PeerId& peer, uint64_t id, bool ok);
 
    PeerNetwork*          network_ = nullptr;
    Config                config_;
    std::atomic<uint64_t> next_id_{1};
    std::atomic<bool>     running_{false};
 
    OfferHandler    offer_handler_;
    ProgressHandler progress_handler_;
    CompleteHandler complete_handler_;
 
    mutable std::mutex mutex_;  
    std::unordered_map<uint64_t, std::shared_ptr<Outgoing>> outgoing_;
    std::unordered_map<PeerId, std::unordered_map<uint64_t, std::shared_ptr<Incoming>>,
                       PeerId::Hash> incoming_;
 
    // send-side worker pool + send queue
    std::vector<std::thread>  workers_;
    std::mutex                queue_mutex_;
    std::condition_variable   queue_cv_;
    std::queue<uint64_t>      send_queue_;
 
    // receive-side disk-writer pool + ready queue. A transfer is pushed here when it
    // has pending write jobs; its `scheduled` flag keeps it single-owner so exactly
    // one worker drains a given transfer at a time (preserving chunk order).
    std::vector<std::thread>              disk_workers_;
    std::mutex                            disk_mutex_;
    std::condition_variable               disk_cv_;
    std::queue<std::shared_ptr<Incoming>> disk_ready_;
 
    // maintenance (idle timeout / purge)
    std::thread               maintenance_thread_;
    std::mutex                maintenance_mutex_;
    std::condition_variable   maintenance_cv_;
 
    mutable std::mutex stats_mutex_;
    Stats              stats_;
};
 
} // namespace librats