// Copyright (C) Microsoft Corporation. All rights reserved.

#include <arpa/inet.h>

#include <sys/epoll.h>

#include <netinet/in.h>

#include <sys/socket.h>

#include "DnsServer.h"

#include "RuntimeErrorWithSourceLocation.h"

#include "Syscall.h"

#include "util.h"

// Port used by DNS server

constexpr int c_dnsServerPort = 53;

// Max number of events to be returned by epoll_wait()

constexpr int c_epollWaitMaxEvents = 100;

// Maximum size of DNS over UDP requests is 4096 bytes (max size is reached for EDNS UDP requests)

constexpr int c_maxUdpDnsBufferSize = 4096;

// Max number of pending connections in the TCP listen queue

constexpr int c_maxListenBacklog = 1000;

DnsServer::DnsServer(DnsTunnelingCallback&& tunnelDnsRequest) : m_tunnelDnsRequest(std::move(tunnelDnsRequest))

{

}

DnsServer::~DnsServer() noexcept

{

Stop();

}

void DnsServer::Start(const std::string& ipAddress) noexcept

try

{

// Create epoll handler fd. 0 represents default flags

m_epollFd = Syscall(epoll_create1, 0);

StartUdpDnsServer(ipAddress);

StartTcpDnsServer(ipAddress);

// Create and register the shutdown pipe with epoll

m_shutdownServerLoopPipe = wil::unique_pipe::create(0);

epoll_event event{};

event.events = EPOLLIN;

event.data.fd = m_shutdownServerLoopPipe.read().get();

Syscall(epoll_ctl, m_epollFd.get(), EPOLL_CTL_ADD, m_shutdownServerLoopPipe.read().get(), &event);

// Start server loop

m_serverThread = std::thread([this]() { ServerLoop(); });

}

CATCH_LOG()

void DnsServer::StartUdpDnsServer(const std::string& ipAddress) noexcept

try

{

sockaddr_in serverAddr{};

serverAddr.sin_family = AF_INET;

Syscall(inet_pton, AF_INET, ipAddress.c_str(), &serverAddr.sin_addr);

serverAddr.sin_port = htons(c_dnsServerPort);

// Create IPv4 UDP socket

m_udpSocket = Syscall(socket, AF_INET, SOCK_DGRAM | SOCK_NONBLOCK, 0);

// Bind socket

Syscall(bind, m_udpSocket.get(), reinterpret_cast<sockaddr*>(&serverAddr), sizeof(serverAddr));

// Configure epoll to track the UDP socket. EPOLLIN is used to get epoll notifications

// whenever there is data available to be read from the socket

epoll_event event{};

event.events = EPOLLIN;

event.data.fd = m_udpSocket.get();

Syscall(epoll_ctl, m_epollFd.get(), EPOLL_CTL_ADD, m_udpSocket.get(), &event);

GNS_LOG_INFO("Successfully started UDP server on IP {}", ipAddress.c_str());

}

CATCH_LOG()

void DnsServer::StartTcpDnsServer(const std::string& ipAddress) noexcept

try

{

sockaddr_in serverAddr{};

serverAddr.sin_family = AF_INET;

Syscall(inet_pton, AF_INET, ipAddress.c_str(), &serverAddr.sin_addr);

serverAddr.sin_port = htons(c_dnsServerPort);

// Create IPv4 TCP socket

m_tcpListenSocket = Syscall(socket, AF_INET, SOCK_STREAM | SOCK_NONBLOCK, 0);

// Bind socket

Syscall(bind, m_tcpListenSocket.get(), reinterpret_cast<sockaddr*>(&serverAddr), sizeof(serverAddr));

// Listen for incoming connections

Syscall(listen, m_tcpListenSocket.get(), c_maxListenBacklog);

// Configure epoll to track the TCP listening socket. EPOLLIN is used to get epoll notifications

// whenever there is a new incoming TCP connection.

epoll_event event{};

event.events = EPOLLIN;

event.data.fd = m_tcpListenSocket.get();

Syscall(epoll_ctl, m_epollFd.get(), EPOLL_CTL_ADD, m_tcpListenSocket.get(), &event);

GNS_LOG_INFO("Successfully started TCP server on IP {}", ipAddress.c_str());

}

CATCH_LOG();

void DnsServer::HandleUdpDnsResponse(const gsl::span<gsl::byte> dnsBuffer, const LX_GNS_DNS_CLIENT_IDENTIFIER& dnsClientIdentifier) noexcept

try

{

GNS_LOG_INFO("New UDP DNS response DNS buffer size: {}, UDP request id: {}", dnsBuffer.size(), dnsClientIdentifier.DnsClientId);

std::scoped_lock<std::mutex> lock{m_udpLock};

auto it = m_udpRequests.find(dnsClientIdentifier.DnsClientId);

if (it == m_udpRequests.end())

{

GNS_LOG_ERROR("Received a response for a UDP request that is not tracked, UDP request id: {}", dnsClientIdentifier.DnsClientId);

return;

}

// Stop tracking the request, irrespective of the DNS response being succesfully sent

const auto removeDnsRequest = wil::scope_exit([&] { m_udpRequests.erase(dnsClientIdentifier.DnsClientId); });

sockaddr_in& remoteAddr = it->second;

// Send DNS response buffer back to the Linux DNS client

int bufferSize = dnsBuffer.size();

int totalBytesSent = 0;

while (totalBytesSent < bufferSize)

{

int bytesSent = Syscall(

sendto, m_udpSocket.get(), dnsBuffer.data() + totalBytesSent, bufferSize - totalBytesSent, 0, reinterpret_cast<sockaddr*>(&remoteAddr), sizeof(remoteAddr));

totalBytesSent += bytesSent;

}

}

CATCH_LOG()

void DnsServer::HandleTcpDnsResponse(const gsl::span<gsl::byte> dnsBuffer, const LX_GNS_DNS_CLIENT_IDENTIFIER& dnsClientIdentifier) noexcept

try

{

GNS_LOG_INFO(

"New TCP DNS response "

"DNS buffer size: {}, TCP connection id: {}",

dnsBuffer.size(),

dnsClientIdentifier.DnsClientId);

std::scoped_lock<std::mutex> lock{m_tcpLock};

auto it = m_tcpConnectionContexts.find(dnsClientIdentifier.DnsClientId);

if (it == m_tcpConnectionContexts.end())

{

GNS_LOG_ERROR("Received a response for an untracked TCP connection id: {}", dnsClientIdentifier.DnsClientId);

return;

}

auto tcpConnection = it->second->m_tcpConnection.get();

// Send DNS response buffer back to the Linux DNS client.

//

// Note: there might be more DNS requests sent on the same TCP connection. The DNS protocol allows sending the responses in a

// different order than the order of the corresponding DNS requests.

int bufferSize = dnsBuffer.size();

int totalBytesSent = 0;

while (totalBytesSent < bufferSize)

{

int bytesSent = Syscall(write, tcpConnection, dnsBuffer.data() + totalBytesSent, bufferSize - totalBytesSent);

totalBytesSent += bytesSent;

}

}

CATCH_LOG()

void DnsServer::HandleNewTcpConnection() noexcept

try

{

std::scoped_lock<std::mutex> lock{m_tcpLock};

// Accept new connection. Mark connection socket as non-blocking

wil::unique_fd connectionFd = Syscall(accept4, m_tcpListenSocket.get(), nullptr, nullptr, SOCK_NONBLOCK);

// Get next connection id. If value reaches UINT_MAX + 1 it will be automatically reset to 0

const auto connectionId = m_currentTcpConnectionId++;

// Track the new connection

auto [it, _] = m_tcpConnectionContexts.emplace(

connectionId, std::make_unique<DnsServer::TcpConnectionContext>(connectionId, std::move(connectionFd)));

auto& localContext = it->second;

auto removeContextOnError = wil::scope_exit([&] { m_tcpConnectionContexts.erase(connectionId); });

// Register the new connection with epoll. EPOLLIN is used to get epoll notifications

// whenever there is new data on the TCP connection.

epoll_event event{};

event.events = EPOLLIN;

event.data.fd = localContext->m_tcpConnection.get();

event.data.ptr = localContext.get();

Syscall(epoll_ctl, m_epollFd.get(), EPOLL_CTL_ADD, localContext->m_tcpConnection.get(), &event);

removeContextOnError.release();

}

CATCH_LOG();

void DnsServer::HandleNewTcpData(TcpConnectionContext* context) noexcept

try

{

std::vector<gsl::byte> dnsRequest;

uint32_t tcpConnectionId{};

// Scoped m_tcpLock

{

std::scoped_lock<std::mutex> lock{m_tcpLock};

// In case of any failure reading data, close the connection and stop tracking it.

// Note: Closing the connection automatically unregisters it from epoll.

auto removeConnectionOnError = wil::scope_exit([&] { m_tcpConnectionContexts.erase(context->m_connectionId); });

// Read the remaining bytes of the current DNS request

int bytesReceived = Syscall(

recv,

context->m_tcpConnection.get(),

context->m_currentDnsRequest.data() + context->m_currentRequestOffset,

context->m_currentDnsRequest.size() - context->m_currentRequestOffset,

0);

// 0 bytes received indicates connection was closed by the TCP client

if (bytesReceived == 0)

{

return;

}

context->m_currentRequestOffset += bytesReceived;

if (context->m_currentRequestOffset == context->m_currentDnsRequest.size())

{

// We read the 2 bytes that represent the DNS request length

// Resize buffer to fit the entire DNS request (2 bytes storing the request length + the actual DNS request)

if (context->m_currentDnsRequest.size() == c_byteCountTcpRequestLength)

{

uint16_t dnsRequestLength = 0;

memcpy(&dnsRequestLength, context->m_currentDnsRequest.data(), c_byteCountTcpRequestLength);

// The request length is stored in network byte order

dnsRequestLength = ntohs(dnsRequestLength);

context->m_currentDnsRequest.resize(c_byteCountTcpRequestLength + dnsRequestLength);

}

// We read a full DNS request

else

{

// Move request to a local variable

dnsRequest = std::move(context->m_currentDnsRequest);

tcpConnectionId = context->m_connectionId;

// Reset state to prepare for the next DNS request on the connection (if any)

context->m_currentRequestOffset = 0;

context->m_currentDnsRequest.resize(c_byteCountTcpRequestLength);

}

}

removeConnectionOnError.release();

}

if (!dnsRequest.empty())

{

// Tunnel request to Windows

LX_GNS_DNS_CLIENT_IDENTIFIER dnsClientIdentifier{};

dnsClientIdentifier.DnsClientId = tcpConnectionId;

dnsClientIdentifier.Protocol = IPPROTO_TCP;

GNS_LOG_INFO("New TCP DNS request DNS buffer size: {}, TCP connection id: {}", dnsRequest.size(), dnsClientIdentifier.DnsClientId);

m_tunnelDnsRequest(gsl::make_span(dnsRequest), dnsClientIdentifier);

}

}

CATCH_LOG();

void DnsServer::HandleDnsResponse(const gsl::span<gsl::byte> dnsBuffer, const LX_GNS_DNS_CLIENT_IDENTIFIER& dnsClientIdentifier) noexcept

try

{

switch (dnsClientIdentifier.Protocol)

{

case IPPROTO_UDP:

{

HandleUdpDnsResponse(dnsBuffer, dnsClientIdentifier);

break;

}

case IPPROTO_TCP:

{

HandleTcpDnsResponse(dnsBuffer, dnsClientIdentifier);

break;

}

default:

{

GNS_LOG_ERROR("Unexpected DNS protocol {}", dnsClientIdentifier.Protocol);

break;

}

}

}

CATCH_LOG()

void DnsServer::ServerLoop() noexcept

{

UtilSetThreadName("DnsServer");

epoll_event events[c_epollWaitMaxEvents];

memset(events, 0, sizeof(events));

for (;;)

{

try

{

// A fixed number of events is requested from epoll_wait (c_epollWaitMaxEvents). In case the number of ready events is

// greater than c_epollWaitMaxEvents, epoll will round-robin through the ready events until we get a notification for all of them.

size_t numReadyEvents = Syscall(epoll_wait, m_epollFd.get(), events, c_epollWaitMaxEvents, -1);

// No event

if (numReadyEvents == 0)

{

continue;

}

for (size_t index = 0; index < numReadyEvents; index++)

{

// Notification for the shutdown pipe == the server needs to exit

if (events[index].data.fd == m_shutdownServerLoopPipe.read().get())

{

return;

}

// Notification for the listen socket == a new incoming TCP connection

else if (events[index].data.fd == m_tcpListenSocket.get())

{

HandleNewTcpConnection();

}

// Notification for the UDP socket == There is data to be read from the UDP socket, indicating a new DNS request was received

else if (events[index].data.fd == m_udpSocket.get())

{

HandleUdpDnsRequest();

}

// Other notifications == new data was received on one of the active TCP connections

else

{

HandleNewTcpData(static_cast<TcpConnectionContext*>(events[index].data.ptr));

}

}

}

CATCH_LOG()

}

}

void DnsServer::HandleUdpDnsRequest() noexcept

try

{

static std::array<gsl::byte, c_maxUdpDnsBufferSize> s_dnsBuffer;

gsl::span<gsl::byte> dnsRequest;

uint32_t udpRequestId{};

// Scoped m_udpLock

{

std::scoped_lock<std::mutex> lock{m_udpLock};

// Since we only configure an IPv4 DNS server in Linux, we expect all Linux DNS clients to use IPv4 addresses

sockaddr_in remoteAddr{};

socklen_t remoteAddrLen = sizeof(remoteAddr);

// Read the DNS request

int bytesReceived = Syscall(

recvfrom, m_udpSocket.get(), s_dnsBuffer.data(), c_maxUdpDnsBufferSize, 0, reinterpret_cast<sockaddr*>(&remoteAddr), &remoteAddrLen);

if (bytesReceived == 0)

{

GNS_LOG_ERROR("recvfrom returned 0 bytes");

return;

}

// Get next request id. If value reaches UINT_MAX + 1 it will be automatically reset to 0

const auto requestId = m_currentUdpRequestId++;

GNS_LOG_INFO(

"New UDP DNS request DNS client IP: {}, DNS client port {}, DNS buffer size: {}, UDP request id: {}",

Address::FromBinary(AF_INET, 0, &remoteAddr.sin_addr).Addr().c_str(),

ntohs(remoteAddr.sin_port),

bytesReceived,

requestId);

// Move request to a local variable

dnsRequest = std::move(gsl::make_span(s_dnsBuffer).subspan(0, bytesReceived));

udpRequestId = requestId;

// Track the request

m_udpRequests.emplace(requestId, remoteAddr);

}

if (!dnsRequest.empty())

{

auto removeRequestOnError = wil::scope_exit([&] {

std::scoped_lock<std::mutex> lock{m_udpLock};

m_udpRequests.erase(udpRequestId);

});

// Tunnel request to Windows

LX_GNS_DNS_CLIENT_IDENTIFIER dnsClientIdentifier{};

dnsClientIdentifier.Protocol = IPPROTO_UDP;

dnsClientIdentifier.DnsClientId = udpRequestId;

m_tunnelDnsRequest(dnsRequest, dnsClientIdentifier);

removeRequestOnError.release();

}

}

CATCH_LOG()

void DnsServer::Stop() noexcept

try

{

GNS_LOG_INFO("stopping DNS server");

// Signal the server loop to stop by closing the write fd of the pipe

m_shutdownServerLoopPipe.write().reset();

if (m_serverThread.joinable())

{

m_serverThread.join();

}

}

CATCH_LOG()