dht_discovery.h

Go to the documentation of this file.

#pragma once

 
#include "librats/util/rats_export.h"
#include "librats/node/peer_network.h"
#include "librats/subsystems/dht_service.h"
#include "librats/dht/dht.h"
 
#include <atomic>
#include <chrono>
#include <condition_variable>
#include <cstdint>
#include <memory>
#include <mutex>
#include <optional>
#include <string>
#include <thread>
#include <unordered_map>
#include <vector>
 
namespace librats {
 
class DialService;
class ReflexiveAddressService;
 
class RATS_API DhtDiscovery final : public Subsystem, public DhtService {
public:
    struct Config {
        uint16_t                  dht_port = 0;          
        std::string               bind_address = "";
        std::string               data_dir = "";         
        bool                      enable_ipv4 = true;    
        bool                      enable_ipv6 = true;    
        std::string               discovery_key = "";  
        std::vector<std::string>  keys;
        std::vector<HostEndpoint>    bootstrap_nodes;       
        std::chrono::milliseconds search_interval{30000}; 
        std::chrono::milliseconds announce_interval{600000};  // 10 min

        bool                      discover_external_ip = true;
        std::vector<HostEndpoint> stun_servers;          
        std::chrono::milliseconds stun_timeout{3000};
        std::chrono::milliseconds stun_every{20000};

        std::chrono::milliseconds redial_every{5000};
        std::chrono::milliseconds redial_for{120000};
        std::chrono::milliseconds redial_slow{25000};
    };

    class RATS_API Redials {
    public:
        enum class Dial : uint8_t {
            None,   
            First,  
            Again,  
        };
        using Clock = std::chrono::steady_clock;

        Redials(std::chrono::milliseconds every, std::chrono::milliseconds dial_for,
                std::chrono::milliseconds slow)
            : every_(every), for_(dial_for), slow_(slow) {}

        Dial found(const Address& address, Clock::time_point now);
        std::vector<Address> due(Clock::time_point now);
        std::optional<Clock::time_point> next_due() const;
        void connected(const Address& address) { found_.erase(address); }
        void forget_older_than(Clock::time_point before);
        size_t size() const { return found_.size(); }
 
    private:
        struct Dialed {
            Clock::time_point first;
            Clock::time_point last;
            Clock::time_point seen;
        };
        std::chrono::milliseconds every_;
        std::chrono::milliseconds for_;
        std::chrono::milliseconds slow_;
        std::unordered_map<Address, Dialed> found_;
    };
 
    explicit DhtDiscovery(Config config);
    ~DhtDiscovery() override;
 
    void attach(NodeContext& ctx) override;
    void start() override;
    void stop() override;
 
    bool     is_running() const;

    DhtClient* dht_client() override { return dht_ ? dht_.get() : dht6_.get(); }
 
    uint16_t dht_port() const;       
    uint16_t dht_port_v6() const;    
    InfoHash discovery_hash() const;
    std::vector<InfoHash> discovery_hashes() const;

    enum class Use : uint8_t { Announce = 1, Search = 2, Both = 3 };
    struct Key {
        std::string key;
        Use use = Use::Both;
    };

    void set_keys(const std::vector<Key>& keys);
    void set_keys(const std::vector<std::string>& keys);

    std::string external_address() const;

    bool is_own(const Address& address) const;

    static InfoHash hash_for_key(const std::string& key);
 
private:
    void loop();
    bool probe_external_ip();
    void on_peers(const std::vector<Address>& peers, const InfoHash& info_hash);
    bool searches(const InfoHash& hash) const;
    std::vector<Address> connected_addresses() const;
    void punch_due(std::chrono::steady_clock::time_point now);
    uint16_t announced_port(bool v6 = false) const;

    std::unique_ptr<DhtClient> make_client(AddressFamily family);

    void wake();

    void refresh_own_ips();

    template <typename Fn>
    void for_each_client(Fn fn) {
        if (dht_)  fn(*dht_);
        if (dht6_) fn(*dht6_);
    }

    struct Announced {
        InfoHash hash{};
        Use use = Use::Both;
        std::optional<std::chrono::steady_clock::time_point> announced[2];
        bool searched = false;  
    };
 
    Config                     config_;
    mutable std::mutex         keys_mutex_;
    std::vector<Announced>     hashes_;            
    bool                       keys_set_ = false;  
    PeerNetwork*               network_ = nullptr;
    DialService*               dial_ = nullptr;  
    ReflexiveAddressService*   reflexive_ = nullptr;  
    std::unique_ptr<DhtClient> dht_;    
    std::unique_ptr<DhtClient> dht6_;   
 
    std::thread             thread_;
    std::atomic<bool>       running_{false};
    std::atomic<bool>       recover_pending_{false};  
    std::atomic<bool>       announce_pending_{false}; 
    std::atomic<bool>       keys_pending_{false};     
    std::atomic<bool>       punch_pending_{false};    
    std::atomic<uint16_t>   mapped_port_{0};          
    std::atomic<uint16_t>   reflexive_port_{0};       
    bool                    mapping_checked_ = false; 
    std::optional<Address>  last_seen_;               
 
    mutable std::mutex      own_mutex_;
    std::vector<IpAddress>  own_local_;   
    std::optional<IpAddress> own_seen_;   
    std::optional<IpAddress> own_mapped_; 
    std::mutex              wait_mutex_;
    std::condition_variable wake_;
 
    std::mutex dialed_mutex_;
    Redials    dialed_;  
};
 
} // namespace librats