dns_client.hpp

浏览该文件的文档.

#ifndef NEFORCE_NETWORK_DNS_DNS_CLIENT_HPP__
#define NEFORCE_NETWORK_DNS_DNS_CLIENT_HPP__

 
#include "NeForce/core/async/cancellation_slot.hpp"
#include "NeForce/core/async/future.hpp"
#include "NeForce/core/async/io_context.hpp"
#include "NeForce/core/async/promise.hpp"
#include "NeForce/core/async/shared_mutex.hpp"
#include "NeForce/core/async/use_awaitable.hpp"
#include "NeForce/core/container/unordered_map.hpp"
#include "NeForce/core/memory/weak_ptr.hpp"
#include "NeForce/core/time/clocks.hpp"
#include "NeForce/core/utility/optional.hpp"
#include "NeForce/network/dns/dns_message.hpp"
#include "NeForce/network/udp_socket.hpp"
#include "NeForce/network/util/ports.hpp"
NEFORCE_BEGIN_NAMESPACE__


class NEFORCE_API dns_client {
public:
    struct config {
        string server{"8.8.8.8"};   
        ports port{ports::DNS};     
        milliseconds timeout{5000}; 
    };
 
private:
    struct dns_query_op : enable_shared_from_this<dns_query_op> {
        dns_client* client;
        string domain;
        dns_record::raw type;
        dns_class qclass;
        function<void(error_code, dns_query_result)> handler;
        cancellation_slot* cancel_slot{nullptr};
        atomic<bool> fired_{false};
        uint16_t query_id_{0};
        byte_vector query_data_;
        size_t timer_id_{0};
        steady_clock::time_point start_time_;
 
        void start();
        void on_response(error_code ec, dns_query_result result);
        void on_timeout();
        void do_cancel();
    };
 
    struct pending_entry {
        weak_ptr<dns_query_op> op;                                
        byte_vector query_data;                                   
        steady_clock::time_point created_at{steady_clock::now()}; 
    };
 
    config config_{};                                                               
    unordered_map<string, pair<dns_query_result, steady_clock::time_point>> cache_; 
    mutable shared_mutex cache_mutex_;                                              
    seconds cache_ttl_{300};                                                        
    bool use_tcp_ = false;                                                          
 
    bool recursion_desired_ = true;                        
    uint16_t edns_udp_payload_{edns::DEFAULT_UDP_PAYLOAD}; 
    bool dnssec_ok_ = false;                               
 
    udp_socket shared_socket_;                               
    unordered_map<uint16_t, pending_entry> pending_queries_; 
    mutable mutex pending_mutex_;                            
    mutable mutex send_mutex_;                               
    io_context* ctx_{nullptr};                               
    atomic<bool> io_running_{false};                         
 
private:
    void ensure_io_started();
    void start_io();
    void stop_io();
    void register_shared_receive();
    void on_udp_readable(error_code ec);
    void process_udp_receive();
    void dispatch_to_op(uint16_t txid, byte_vector response);
    void unregister_query(uint16_t txid);
 
    void send_query(const byte_vector& query);
    byte_vector send_tcp_query(const byte_vector& query);
 
    optional<dns_query_result> check_cache(const string& key);
    void update_cache(const string& key, const dns_query_result& result);
 
public:
    dns_client() = default;

    explicit dns_client(config cfg, io_context& ctx, bool use_tcp = false);
 
    ~dns_client();
 
    dns_client(const dns_client&) = delete;
    dns_client& operator=(const dns_client&) = delete;
    dns_client(dns_client&&) = delete;
    dns_client& operator=(dns_client&&) = delete;

    void set_config(config cfg) noexcept { config_ = move(cfg); }

    void set_timeout(const milliseconds timeout) noexcept { config_.timeout = timeout; }

    void set_use_tcp(const bool use_tcp) noexcept { use_tcp_ = use_tcp; }

    void set_cache_ttl(const seconds ttl) noexcept { cache_ttl_ = ttl; }

    void set_recursion_desired(const bool rd) { recursion_desired_ = rd; }

    void set_edns_udp_payload(const uint16_t payload_size) noexcept { edns_udp_payload_ = payload_size; }

    void set_dnssec_ok(const bool ok) noexcept { dnssec_ok_ = ok; }

    void clear_cache() noexcept { cache_.clear(); }

    dns_query_result query(string_view domain, dns_record::raw type = dns_record::A,
                           dns_class qclass = dns_class::INTERNET);

    future<dns_query_result> query_async(const string& domain, dns_record::raw type = dns_record::A,
                                         dns_class qclass = dns_class::INTERNET);

    void async_query(string_view domain, dns_record::raw type, dns_class qclass,
                     function<void(error_code, dns_query_result)> handler);

    void async_query(string_view domain, dns_record::raw type, dns_class qclass, cancellation_slot& slot,
                     function<void(error_code, dns_query_result)> handler);

    template <typename Token,
              enable_if_t<!is_same_v<decay_t<Token>, function<void(error_code, dns_query_result)>>, int> = 0>
    void async_query(string_view domain, dns_record::raw type, dns_class qclass, Token&& token) {
        async_query(domain, type, qclass, function<void(error_code, dns_query_result)>(forward<Token>(token)));
    }

    auto async_query(string_view domain, dns_record::raw type, dns_class qclass, use_future_t);

    void async_query(string_view domain, dns_record::raw type, dns_class qclass, detached_t);
 
#ifdef NEFORCE_STANDARD_20
    auto async_query(string_view domain, dns_record::raw type, dns_class qclass, use_awaitable_t);
#endif

    vector<string> resolve_a(string_view domain);

    vector<string> resolve_aaaa(string_view domain);

    vector<string> resolve_cname(string_view domain);

    vector<string> resolve_mx(string_view domain);

    vector<string> resolve_txt(string_view domain);

    vector<dns_srv_record> resolve_srv(string_view domain);

    optional<dns_soa_record> resolve_soa(string_view domain);

    string reverse_query(string_view ip);

    vector<dns_query_result> batch_query(const vector<string>& domains, dns_record::raw type = dns_record::A);

    static byte_vector build_query(string_view domain, dns_record::raw type = dns_record::A,
                                   dns_class qclass = dns_class::INTERNET, bool rd = true, bool edns_enable = true,
                                   bool dnssec_ok = false, uint16_t edns_payload = edns::DEFAULT_UDP_PAYLOAD);

    static dns_query_result parse_response(const byte_vector& response, uint16_t expected_id = 0);
};
 
 
NEFORCE_BEGIN_INNER__
 
template <>
struct future_handler<error_code, dns_query_result> {
    shared_ptr<promise<dns_query_result>> promise_;
 
    void operator()(error_code ec, dns_query_result result) {
        if (ec) {
            promise_->set_exception(make_exception_ptr(system_exception(ec)));
        } else {
            promise_->set_value(move(result));
        }
    }
};
 
NEFORCE_END_INNER__
 
template <>
struct async_result<use_future_t, void(error_code, dns_query_result)> {
    using handler_type = inner::future_handler<error_code, dns_query_result>;
    using return_type = future<dns_query_result>;
    handler_type handler_;
    explicit async_result(use_future_t /*unused*/) { handler_.promise_ = make_shared<promise<dns_query_result>>(); }
    handler_type get_handler() { return handler_; }
    return_type get() { return handler_.promise_->get_future(); }
};
 // DNS
 
NEFORCE_END_NAMESPACE__
#endif // NEFORCE_NETWORK_DNS_DNS_CLIENT_HPP__