BareGit
#include "socket_probe.h"

#include <algorithm>
#include <array>
#include <atomic>
#include <cerrno>
#include <climits>
#include <cstring>
#include <memory>
#include <utility>
#include <vector>

#include <ares.h>
#include <arpa/inet.h>
#include <netinet/icmp6.h>
#include <netinet/ip_icmp.h>
#include <poll.h>
#include <sys/socket.h>
#include <unistd.h>

namespace
{

using Clock = std::chrono::steady_clock;
using Addresses = std::unique_ptr<ares_addrinfo, decltype(&ares_freeaddrinfo)>;

enum class Protocol { TCP, UDP, ICMP };
enum class Phase { CONNECTING, WRITING, READING, FAILED };

struct Socket
{
    int fd;

    explicit Socket(int fd) : fd(fd) {}
    Socket(const Socket&) = delete;
    Socket& operator=(const Socket&) = delete;
    Socket(Socket&& other) noexcept : fd(std::exchange(other.fd, -1)) {}

    ~Socket()
    {
        if(fd >= 0)
        {
            close(fd);
        }
    }
};

struct Resolution
{
    int status = ARES_ETIMEOUT;
    Addresses addresses{nullptr, ares_freeaddrinfo};
};

void resolved(void* context, int status, [[maybe_unused]] int timeouts,
              ares_addrinfo* addresses)
{
    auto& result = *static_cast<Resolution*>(context);
    result.status = status;
    result.addresses.reset(addresses);
}

int remainingMilliseconds(Clock::time_point deadline)
{
    const auto remaining = std::chrono::ceil<std::chrono::milliseconds>(
        deadline - Clock::now()).count();
    return static_cast<int>(std::clamp<std::int64_t>(remaining, 0, INT_MAX));
}

mw::E<Addresses> resolve(const std::string& host, std::uint16_t port,
                         Clock::time_point deadline)
{
    Resolution result;
    ares_channel_t* raw = nullptr;
    ares_options options{};
    options.evsys = ARES_EVSYS_DEFAULT;
    const int initialized = ares_init_options(&raw, &options,
                                              ARES_OPT_EVENT_THREAD);
    if(initialized != ARES_SUCCESS)
    {
        return std::unexpected(mw::runtimeError(ares_strerror(initialized)));
    }
    std::unique_ptr<ares_channel_t, decltype(&ares_destroy)> channel(
        raw, ares_destroy);
    ares_addrinfo_hints hints{};
    hints.ai_family = AF_UNSPEC;
    hints.ai_socktype = SOCK_DGRAM;
    hints.ai_flags = ARES_AI_NUMERICSERV | ARES_AI_NOSORT;
    const auto service = std::to_string(port);
    ares_getaddrinfo(channel.get(), host.c_str(), service.c_str(), &hints,
                     resolved, &result);
    const auto waited = ares_queue_wait_empty(channel.get(),
                                              remainingMilliseconds(deadline));
    // Destruction joins the resolver thread before we inspect callback state.
    channel.reset();
    if(waited != ARES_SUCCESS || result.status != ARES_SUCCESS)
    {
        if(result.status == ARES_ENOMEM)
        {
            return std::unexpected(mw::runtimeError("DNS allocation failed"));
        }
        return Addresses(nullptr, ares_freeaddrinfo);
    }
    return std::move(result.addresses);
}

mw::Error socketError(const char* operation)
{
    return mw::runtimeError(std::string(operation) + ": " +
                            std::strerror(errno));
}

bool retryable(int code)
{
    return code == EAGAIN || code == EWOULDBLOCK || code == EINTR;
}

bool networkFailure(int code)
{
    return code == ECONNREFUSED || code == ECONNRESET || code == ETIMEDOUT ||
        code == EHOSTUNREACH || code == ENETUNREACH || code == ENETDOWN ||
        code == EHOSTDOWN || code == EPIPE;
}

struct Connection
{
    Socket socket;
    const ares_addrinfo_node* address;
    Phase phase;
    std::array<unsigned char, 16> echo{};
};

std::array<unsigned char, 16> echoRequest(int family)
{
    static std::atomic<std::uint64_t> next_nonce{1};
    const auto nonce = next_nonce.fetch_add(1);
    std::array<unsigned char, 16> packet{};
    packet[0] = family == AF_INET ? ICMP_ECHO : ICMP6_ECHO_REQUEST;
    packet[7] = 1;
    for(std::size_t i = 0; i < sizeof(nonce); ++i)
    {
        packet[8 + i] = static_cast<unsigned char>(nonce >> (i * 8));
    }
    if(family == AF_INET)
    {
        unsigned int sum = 0;
        for(std::size_t i = 0; i < packet.size(); i += 2)
        {
            sum += (packet[i] << 8) | packet[i + 1];
        }
        while(sum >> 16)
        {
            sum = (sum & 0xffff) + (sum >> 16);
        }
        const auto checksum = static_cast<std::uint16_t>(~sum);
        packet[2] = checksum >> 8;
        packet[3] = checksum & 0xff;
    }
    return packet;
}

bool sameHost(const sockaddr_storage& source, const sockaddr* expected)
{
    if(source.ss_family != expected->sa_family)
    {
        return false;
    }
    if(source.ss_family == AF_INET)
    {
        const auto& actual = reinterpret_cast<const sockaddr_in&>(source);
        const auto* wanted = reinterpret_cast<const sockaddr_in*>(expected);
        return actual.sin_addr.s_addr == wanted->sin_addr.s_addr;
    }
    const auto& actual = reinterpret_cast<const sockaddr_in6&>(source);
    const auto* wanted = reinterpret_cast<const sockaddr_in6*>(expected);
    return std::memcmp(&actual.sin6_addr, &wanted->sin6_addr,
                        sizeof(in6_addr)) == 0 &&
        (wanted->sin6_scope_id == 0 ||
         actual.sin6_scope_id == wanted->sin6_scope_id);
}

bool echoReply(const Connection& connection, const unsigned char* data,
               std::size_t size, const sockaddr_storage& source)
{
    const int type = connection.address->ai_family == AF_INET
        ? ICMP_ECHOREPLY : ICMP6_ECHO_REPLY;
    return size == connection.echo.size() && data[0] == type && data[1] == 0 &&
        data[6] == 0 && data[7] == 1 &&
        std::memcmp(data + 8, connection.echo.data() + 8, 8) == 0 &&
        sameHost(source, connection.address->ai_addr);
}

mw::E<ProbeStatus> check(const std::string& host, std::uint16_t port,
                         Protocol protocol, const std::string& payload,
                         Clock::time_point deadline)
{
    auto resolved = resolve(host, port, deadline);
    if(!resolved)
    {
        return std::unexpected(resolved.error());
    }
    if(!*resolved || Clock::now() >= deadline)
    {
        return ProbeStatus::BAD;
    }
    auto addresses = std::move(*resolved);
    std::vector<Connection> connections;
    std::optional<mw::Error> setup_error;
    bool attempted = false;
    for(auto* address = addresses->nodes; address; address = address->ai_next)
    {
        if(address->ai_family != AF_INET && address->ai_family != AF_INET6)
        {
            continue;
        }
        const int type = protocol == Protocol::TCP ? SOCK_STREAM : SOCK_DGRAM;
        const int ip_protocol = protocol == Protocol::ICMP
            ? (address->ai_family == AF_INET ? static_cast<int>(IPPROTO_ICMP)
                                           : static_cast<int>(IPPROTO_ICMPV6))
            : 0;
        Socket socket(::socket(address->ai_family,
                                type | SOCK_NONBLOCK | SOCK_CLOEXEC,
                                ip_protocol));
        if(socket.fd < 0)
        {
            setup_error = socketError("Create probe socket");
            continue;
        }
        attempted = true;
        Phase phase = Phase::WRITING;
        if(connect(socket.fd, address->ai_addr, address->ai_addrlen) < 0)
        {
            if(errno == EINPROGRESS || errno == EINTR)
            {
                phase = Phase::CONNECTING;
            }
            else if(networkFailure(errno))
            {
                continue;
            }
            else
            {
                return std::unexpected(socketError("Connect probe socket"));
            }
        }
        else if(protocol == Protocol::TCP)
        {
            return ProbeStatus::GOOD;
        }
        connections.push_back({std::move(socket), address, phase,
                               echoRequest(address->ai_family)});
    }
    if(!attempted && setup_error)
    {
        return std::unexpected(*setup_error);
    }
    std::vector<pollfd> descriptors(connections.size());
    while(Clock::now() < deadline)
    {
        bool active = false;
        for(std::size_t i = 0; i < connections.size(); ++i)
        {
            const auto& connection = connections[i];
            const bool failed = connection.phase == Phase::FAILED;
            descriptors[i] = {failed ? -1 : connection.socket.fd,
                static_cast<short>(connection.phase == Phase::READING
                    ? POLLIN : POLLOUT), 0};
            active |= !failed;
        }
        if(!active)
        {
            return ProbeStatus::BAD;
        }
        const int ready = poll(descriptors.data(), descriptors.size(),
                                remainingMilliseconds(deadline));
        if(ready < 0)
        {
            if(errno == EINTR)
            {
                continue;
            }
            return std::unexpected(socketError("Wait for probe socket"));
        }
        for(std::size_t i = 0; i < connections.size(); ++i)
        {
            if(!descriptors[i].revents)
            {
                continue;
            }
            auto& connection = connections[i];
            if(descriptors[i].revents & POLLNVAL)
            {
                return std::unexpected(mw::runtimeError(
                    "Invalid probe socket"));
            }
            if(connection.phase == Phase::CONNECTING)
            {
                int code = 0;
                socklen_t size = sizeof(code);
                if(getsockopt(connection.socket.fd, SOL_SOCKET, SO_ERROR,
                               &code, &size) < 0)
                {
                    return std::unexpected(socketError(
                        "Read connection error"));
                }
                if(code)
                {
                    if(!networkFailure(code))
                    {
                        return std::unexpected(mw::runtimeError(
                            std::string("Connect: ") + std::strerror(code)));
                    }
                    connection.phase = Phase::FAILED;
                    continue;
                }
                if(protocol == Protocol::TCP)
                {
                    return ProbeStatus::GOOD;
                }
                connection.phase = Phase::WRITING;
            }
            if(connection.phase == Phase::WRITING)
            {
                const void* data = protocol == Protocol::ICMP
                    ? static_cast<const void*>(connection.echo.data())
                    : static_cast<const void*>(payload.data());
                const auto size = protocol == Protocol::ICMP
                    ? connection.echo.size() : payload.size();
                const auto sent = send(connection.socket.fd, data, size,
                                        MSG_NOSIGNAL);
                if(sent < 0)
                {
                    if(retryable(errno))
                    {
                        continue;
                    }
                    if(!networkFailure(errno))
                    {
                        return std::unexpected(socketError("Send probe"));
                    }
                    connection.phase = Phase::FAILED;
                    continue;
                }
                connection.phase = Phase::READING;
            }
            if(connection.phase == Phase::READING)
            {
                std::array<unsigned char, 65536> buffer;
                sockaddr_storage source{};
                socklen_t source_size = sizeof(source);
                const auto received = recvfrom(connection.socket.fd,
                    buffer.data(), buffer.size(), 0,
                    reinterpret_cast<sockaddr*>(&source), &source_size);
                if(received < 0)
                {
                    if(retryable(errno))
                    {
                        continue;
                    }
                    if(!networkFailure(errno))
                    {
                        return std::unexpected(socketError("Receive probe"));
                    }
                    connection.phase = Phase::FAILED;
                    continue;
                }
                if(protocol == Protocol::UDP || echoReply(connection,
                    buffer.data(), static_cast<std::size_t>(received), source))
                {
                    return ProbeStatus::GOOD;
                }
            }
        }
    }
    if(protocol == Protocol::UDP)
    {
        for(const auto& connection : connections)
        {
            if(connection.phase == Phase::READING)
            {
                return ProbeStatus::OTHER;
            }
        }
    }
    return ProbeStatus::BAD;
}

}

namespace probe_internal
{

mw::E<ProbeStatus> probeSocket(const TcpEndpoint& endpoint,
                              Clock::time_point deadline)
{
    return check(endpoint.host, endpoint.port, Protocol::TCP, {}, deadline);
}

mw::E<ProbeStatus> probeSocket(const UdpEndpoint& endpoint,
                              Clock::time_point deadline)
{
    return check(endpoint.host, endpoint.port, Protocol::UDP, endpoint.payload,
                  deadline);
}

mw::E<ProbeStatus> probeSocket(const IcmpEndpoint& endpoint,
                              Clock::time_point deadline)
{
    return check(endpoint.host, 0, Protocol::ICMP, {}, deadline);
}

}