hole_punch.h

Go to the documentation of this file.

#pragma once

 
#include "librats/util/rats_export.h"
#include "librats/core/address.h"
#include "librats/core/types.h"
#include "librats/node/nat_status.h"
#include "librats/node/peer_network.h"
#include "librats/peer/peer.h"
#include "librats/peer/peer_id.h"
#include "librats/subsystems/hole_punch_service.h"
#include "librats/subsystems/relay_service.h"
 
#include <atomic>
#include <chrono>
#include <condition_variable>
#include <cstdint>
#include <mutex>
#include <thread>
#include <unordered_map>
#include <vector>
 
namespace librats {
 
class DialService;
class ServiceRegistry;

class RATS_API HolePunch final : public Subsystem, public HolePunchService {
public:
    struct Config {
        size_t max_relays = 3;

        size_t max_addresses = 4;

        int attempts = 3;

        std::chrono::milliseconds round_timeout{6000};

        std::chrono::milliseconds cooldown{60000};

        size_t max_sessions = 32;

        bool enable_relay = true;

        size_t                    relay_budget = 12;
        std::chrono::milliseconds relay_window{1000};

        bool skip_when_endpoint_dependent = true;

        DialProfile profile = DialProfile::punch();

        std::chrono::milliseconds tick{250};

        bool relay_on_failure = true;
    };
 
    HolePunch();
    explicit HolePunch(Config config);
    ~HolePunch() override;
 
    void attach(NodeContext& ctx) override;
    void start() override;
    void stop() override;

    bool punch(const PeerId& target) override;

    size_t active_sessions() const;
    uint64_t punches_started() const noexcept { return punches_started_.load(); }
    uint64_t relayed() const noexcept { return relayed_.load(); }
 
private:
    enum class Phase : uint8_t {
        AwaitingPeerConnect,  
        AwaitingSync,         
        Punching,             
    };
 
    struct Session {
        bool                                  initiator = false;
        Phase                                 phase = Phase::AwaitingPeerConnect;
        int                                   attempt = 0;
        std::vector<Address>                  peer_addresses;
        std::chrono::steady_clock::time_point deadline{};
        PeerId                                via{};
        bool                                  have_via = false;
    };
 
    // — message handling (reactor threads) —
    void handle(const Peer& from, ByteView payload);
    void handle_relay_request(const Peer& from, const PeerId& dst, ByteView inner);
    void handle_relayed(const Peer& via, const PeerId& src, ByteView inner);
    void handle_connect(const Peer& via, const PeerId& src, bool opening,
                        std::vector<Address> addresses);
    void handle_sync(const PeerId& src);
 
    // — outgoing —
    size_t relay_to(const PeerId& dst, const Bytes& inner, const PeerId* via);
    bool   send_connect(const PeerId& target, bool opening, const PeerId* via);
    void   send_sync(const PeerId& target, const PeerId* via);
    void   fire_punch(const PeerId& target, const std::vector<Address>& addresses);
 
    // — worker —
    void loop();
    void service_sessions();
 
    // — helpers —
    void                 escalate_to_relay(const PeerId& target);
    std::vector<PeerId>  directly_connected() const;
    std::vector<Address> own_punch_addresses() const;
    bool                 relay_budget_ok(const PeerId& from);
    bool                 in_cooldown(const PeerId& target) const;   
    void                 begin_cooldown(const PeerId& target);      
 
    Config                 config_;
    PeerNetwork*           network_  = nullptr;
    DialService*           dialer_   = nullptr;
    ExternalAddressService* external_ = nullptr;
    ServiceRegistry*       services_ = nullptr;
    std::atomic<RelayService*> relay_{nullptr};
 
    std::atomic<bool>     running_{false};
    std::atomic<uint64_t> punches_started_{0};
    std::atomic<uint64_t> relayed_{0};
 
    mutable std::mutex      mutex_;
    std::condition_variable cv_;
    std::thread             worker_;
    std::unordered_map<PeerId, Session, PeerId::Hash> sessions_;
    std::unordered_map<PeerId, std::chrono::steady_clock::time_point, PeerId::Hash> cooldown_;
 
    struct RelayBudget {
        std::chrono::steady_clock::time_point window_started{};
        size_t                                spent = 0;
    };
    std::mutex                                              relay_mutex_;
    std::unordered_map<PeerId, RelayBudget, PeerId::Hash>    relay_budget_;
};
 
} // namespace librats