resolver.hpp

#pragma once
 
#include <exception>
#include <unordered_map>
 
#include <coroio/promises.hpp>
#include <coroio/socket.hpp>
#include <coroio/corochain.hpp>
 
namespace NNet {
 
class TResolvConf {
public:
    TResolvConf(const std::string& fn = "/etc/resolv.conf");
    TResolvConf(std::istream& input);
    std::vector<TAddress> Nameservers;
 
private:
    void Load(std::istream& input);
};
 
enum class EDNSType {
    DEFAULT = 0, 
    A = 1, 
    AAAA = 28, 
};
 
struct TTypelessSocket {
    TTypelessSocket() = default;
 
    template<typename TSocket>
    TTypelessSocket(TSocket&& sock) {
        auto socket = std::make_shared<TSocket>(std::move(sock));
        Connector = [socket](const TAddress& addr, TTime deadline) -> TFuture<void> {
            co_await socket->Connect(addr, deadline);
        };
        Reader = [socket](void* buf, size_t size) -> TFuture<ssize_t> {
            co_return co_await socket->ReadSome(buf, size);
        };
        Writer = [socket](const void* buf, size_t size) -> TFuture<ssize_t> {
            co_return co_await socket->WriteSome(buf, size);
        };
    }
 
    TFuture<void> Connect(const TAddress& addr, TTime deadline = TTime::max()) {
        co_await Connector(addr, deadline);
    }
 
    TFuture<ssize_t> ReadSome(void* buf, size_t size) {
        co_return co_await Reader(buf, size);
    }
 
    TFuture<ssize_t> WriteSome(const void* buf, size_t size) {
        co_return co_await Writer(buf, size);
    }
 
private:
    std::function<TFuture<void>(const TAddress& addr, TTime deadline)> Connector;
    std::function<TFuture<ssize_t>(void* buf, size_t size)> Reader;
    std::function<TFuture<ssize_t>(const void* buf, size_t size)> Writer;
};
 
class TResolver {
public:
    template<typename TPoller>
    TResolver(TPoller& poller, EDNSType defaultType = EDNSType::A);
    template<typename TPoller>
    TResolver(const TResolvConf& conf, TPoller& poller, EDNSType defaultType = EDNSType::A);
    template<typename TPoller>
    TResolver(TAddress dnsAddr, TPoller& poller, EDNSType defaultType = EDNSType::A);
    ~TResolver();
 
    TFuture<std::vector<TAddress>> Resolve(const std::string& hostname, EDNSType type = EDNSType::DEFAULT);
 
private:
    TFuture<void> SenderTask();
    TFuture<void> ReceiverTask();
    TFuture<void> TimeoutsTask(std::function<TFuture<void>(TTime until)> timeout);
 
    void ResumeSender();
 
    TFuture<void> Sender;
    TFuture<void> Receiver;
    TFuture<void> Timeouts;
    std::coroutine_handle<> SenderSuspended;
 
    TAddress DnsAddr;
    TTypelessSocket Socket;
    EDNSType DefaultType;
 
    struct TResolveRequest {
        std::string Name;
        EDNSType Type;
 
        bool operator==(const TResolveRequest& other) const {
            return Name == other.Name && Type == other.Type;
        }
 
        size_t hash() const {
            std::size_t result = 0;
            result ^= std::hash<std::string>()(Name) + 0x9e3779b9 + (result << 6) + (result >> 2);
            result ^= std::hash<EDNSType>()(Type) + 0x9e3779b9 + (result << 6) + (result >> 2);
            return result;
        }
    };
 
    struct TResolveRequestHash
    {
        std::size_t operator()(const TResolveRequest& c) const
        {
            return c.hash();
        }
    };
 
    std::queue<TResolveRequest> AddResolveQueue;
    std::queue<std::pair<TTime, TResolveRequest>> TimeoutsQueue;
 
    struct TResolveResult {
        std::vector<TAddress> Addresses = {};
        std::exception_ptr Exception = nullptr;
        int Retries = 0;
    };
 
    void ResumeWaiters(TResolveResult&& result, const TResolveRequest& req);
 
    std::unordered_map<TResolveRequest, TResolveResult, TResolveRequestHash> Results;
    std::unordered_map<TResolveRequest, std::vector<std::coroutine_handle<>>, TResolveRequestHash> WaitingAddrs;
    std::unordered_map<uint64_t, TResolveRequest> Inflight;
 
    uint16_t Xid = 1;
};
 
template<typename TPoller>
TResolver::TResolver(TPoller& poller, EDNSType defaultType)
    : TResolver(TResolvConf(), poller, defaultType)
{ }
 
template<typename TPoller>
TResolver::TResolver(const TResolvConf& conf, TPoller& poller, EDNSType defaultType)
    : TResolver(conf.Nameservers[0], poller, defaultType)
{ }
 
template<typename TPoller>
TResolver::TResolver(TAddress dnsAddr, TPoller& poller, EDNSType defaultType)
    : DnsAddr(std::move(dnsAddr))
    , Socket(typename TPoller::TSocket(poller, DnsAddr.Domain(), SOCK_DGRAM))
    , DefaultType(defaultType)
{
    // Start tasks after fields initialization
    Sender = SenderTask();
    Receiver = ReceiverTask();
    Timeouts = TimeoutsTask([&poller](TTime until) -> TFuture<void> {
        co_await poller.Sleep(until);
    });
}
 
class THostPort {
public:
    THostPort(const std::string& hostPort);
    THostPort(const std::string& host, int port);
 
    TFuture<TAddress> Resolve(TResolver& resolver) {
        char buf[16];
        if (inet_pton(AF_INET, Host.c_str(), buf) == 1 || inet_pton(AF_INET6, Host.c_str(), buf)) {
            co_return TAddress{Host, Port};
        }
 
        auto addresses = co_await resolver.Resolve(Host);
        if (addresses.empty()) {
            throw std::runtime_error("Empty address");
        }
 
        co_return addresses.front().WithPort(Port);
    }
 
    const std::string ToString() const {
        return Host + ":" + std::to_string(Port);
    }
 
    const std::string& GetHost() const {
        return Host;
    }
 
    int GetPort() const {
        return Port;
    }
 
private:
    std::string Host;
    int Port;
};
 
} // namespace NNet