#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);
}
}