// Copyright 2021 The ChromiumOS Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "dns-proxy/resolver.h"

#include <sys/socket.h>

#include <algorithm>
#include <cmath>
#include <memory>
#include <optional>
#include <set>
#include <utility>

#include <base/containers/contains.h>
#include <base/functional/bind.h>
#include <base/logging.h>
#include <base/memory/ref_counted.h>
#include <base/rand_util.h>
#include <base/strings/string_util.h>
#include <base/task/single_thread_task_runner.h>
#include <chromeos/patchpanel/dns/dns_protocol.h>
#include <chromeos/patchpanel/dns/dns_query.h>
#include <chromeos/patchpanel/dns/io_buffer.h>
#include <chromeos/patchpanel/net_util.h>
#include <net-base/ip_address.h>
#include <net-base/socket.h>

// Using directive is necessary to have the overloaded function for socket data
// structure available.
using patchpanel::operator<<;

namespace dns_proxy {
namespace {
constexpr uint32_t kMaxClientTcpConn = 16;
// Given multiple DNS and DoH servers, Resolver will query each servers
// concurrently. |kMaxConcurrentQueries| sets the maximum number of servers to
// query concurrently.
constexpr int kMaxConcurrentQueries = 3;
// Retry delays are reduced by at most |kRetryDelayJitterMultiplier| times to
// avoid coordinated spikes. Having the value >= 1 might introduce an undefined
// behavior.
constexpr float kRetryJitterMultiplier = 0.2;

constexpr base::TimeDelta kProbeInitialDelay = base::Seconds(1);
constexpr base::TimeDelta kProbeMaximumDelay = base::Hours(1);
constexpr float kProbeRetryMultiplier = 1.5;

// The size of TCP header representing the length of the DNS query.
constexpr int kDNSTCPHeaderLength = 2;

// DNS query for resolving "www.gstatic.com" in wire-format data used for
// probing. Transaction ID for the query is empty. This is safe because we
// don't care about the resolving result of the query.
constexpr char kDNSQueryGstatic[] =
    "\x00\x00\x01\x20\x00\x01\x00\x00\x00\x00\x00\x01\x03\x77\x77\x77"
    "\x07\x67\x73\x74\x61\x74\x69\x63\x03\x63\x6f\x6d\x00\x00\x01\x00"
    "\x01\x00\x00\x29\x10\x00\x00\x00\x00\x00\x00\x00";

// Get the time to wait until the next probe.
static base::TimeDelta GetTimeUntilProbe(int num_attempts) {
  base::TimeDelta delay = kProbeInitialDelay;
  delay *= pow(kProbeRetryMultiplier, num_attempts);
  delay -= base::RandDouble() * kRetryJitterMultiplier * delay;
  return std::min(delay, kProbeMaximumDelay);
}

Metrics::QueryError AresStatusMetric(int status) {
  switch (status) {
    case ARES_SUCCESS:
      return Metrics::QueryError::kNone;
    case ARES_ENODATA:
      return Metrics::QueryError::kNoData;
    case ARES_ENOTFOUND:
      return Metrics::QueryError::kDomainNotFound;
    case ARES_ENOTIMP:
      return Metrics::QueryError::kNotImplemented;
    case ARES_EREFUSED:
      return Metrics::QueryError::kQueryRefused;
    case ARES_EFORMERR:
    case ARES_EBADQUERY:
    case ARES_EBADNAME:
    case ARES_EBADFAMILY:
      return Metrics::QueryError::kBadQuery;
    case ARES_ESERVFAIL:
    case ARES_EBADRESP:
      return Metrics::QueryError::kOtherServerError;
    case ARES_ECONNREFUSED:
      return Metrics::QueryError::kConnectionRefused;
    case ARES_ETIMEOUT:
      return Metrics::QueryError::kQueryTimeout;
    default:
      return Metrics::QueryError::kOtherClientError;
  }
}

Metrics::QueryError CurlCodeMetric(int code) {
  switch (code) {
    case CURLE_OK:
      return Metrics::QueryError::kNone;
    case CURLE_UNSUPPORTED_PROTOCOL:
      return Metrics::QueryError::kUnsupportedProtocol;
    case CURLE_URL_MALFORMAT:
    case CURLE_BAD_CONTENT_ENCODING:
      return Metrics::QueryError::kBadQuery;
    case CURLE_COULDNT_RESOLVE_HOST:
    case CURLE_COULDNT_RESOLVE_PROXY:
      return Metrics::QueryError::kBadHost;
    case CURLE_COULDNT_CONNECT:
    case CURLE_SSL_CONNECT_ERROR:
    case CURLE_PEER_FAILED_VERIFICATION:
      return Metrics::QueryError::kConnectionFailed;
    case CURLE_REMOTE_ACCESS_DENIED:
    case CURLE_SSL_CLIENTCERT:
      return Metrics::QueryError::kConnectionRefused;
    case CURLE_OPERATION_TIMEDOUT:
      return Metrics::QueryError::kQueryTimeout;
    case CURLE_TOO_MANY_REDIRECTS:
      return Metrics::QueryError::kTooManyRedirects;
    case CURLE_GOT_NOTHING:
      return Metrics::QueryError::kNoData;
    case CURLE_SEND_ERROR:
    case CURLE_WRITE_ERROR:
    case CURLE_AGAIN:
      return Metrics::QueryError::kSendError;
    case CURLE_RECV_ERROR:
    case CURLE_READ_ERROR:
      return Metrics::QueryError::kReceiveError;
    case CURLE_WEIRD_SERVER_REPLY:
    case CURLE_RANGE_ERROR:
      return Metrics::QueryError::kOtherServerError;
    default:
      return Metrics::QueryError::kOtherClientError;
  }
}

// Return the next ID for SocketFds.
int NextId() {
  static int next_id = 1;
  return next_id++;
}

}  // namespace

std::ostream& operator<<(std::ostream& stream, const Resolver& resolver) {
  resolver.logger_.Run(stream);
  return stream;
}

Resolver::SocketFd::SocketFd(int type, int fd, size_t buf_size)
    : type(type),
      fd(fd),
      len(0),
      num_retries(0),
      num_active_queries(0),
      id(NextId()) {
  buf.resize(buf_size > kMaxDNSBufSize ? kMaxDNSBufSize : buf_size);
  msg = buf.data();
  if (type == SOCK_STREAM) {
    msg += kTCPBufferPaddingLength;
    socklen = 0;
    return;
  }
  socklen = sizeof(src);
}

size_t Resolver::SocketFd::try_resize() {
  size_t remaining_len = buf.size() - len;
  if (type == SOCK_STREAM) {
    remaining_len -= kTCPBufferPaddingLength;
  }

  // Resize only if buffer is full.
  if (remaining_len > 0) {
    return buf.size();
  }
  buf.resize(std::min(2 * buf.size(), kMaxDNSBufSize));

  // Reset |msg| pointer in case of re-allocation.
  msg = buf.data();
  if (type == SOCK_STREAM) {
    msg += kTCPBufferPaddingLength;
  }

  return buf.size();
}

const char* Resolver::SocketFd::get_message() const {
  if (type == SOCK_STREAM) {
    return msg + kDNSTCPHeaderLength;
  }
  return msg;
}

const ssize_t Resolver::SocketFd::get_length() const {
  if (type == SOCK_STREAM) {
    return len - kDNSTCPHeaderLength;
  }
  return len;
}

Resolver::TCPConnection::TCPConnection(
    std::unique_ptr<net_base::Socket> sock,
    const base::RepeatingCallback<void(int, int)>& callback)
    : sock(std::move(sock)) {
  watcher = base::FileDescriptorWatcher::WatchReadable(
      TCPConnection::sock->Get(),
      base::BindRepeating(callback, TCPConnection::sock->Get(), SOCK_STREAM));
}

Resolver::ProbeState::ProbeState(const std::string& target,
                                 bool doh,
                                 bool validated)
    : target(target), doh(doh), validated(validated), num_retries(0) {}

Resolver::Resolver(base::RepeatingCallback<void(std::ostream& stream)> logger,
                   base::TimeDelta timeout,
                   base::TimeDelta retry_delay,
                   int max_num_retries)
    : logger_(logger),
      always_on_doh_(false),
      doh_enabled_(false),
      retry_delay_(retry_delay),
      max_num_retries_(max_num_retries),
      metrics_(new Metrics) {
  ares_client_ = std::make_unique<AresClient>(timeout);
  curl_client_ = std::make_unique<DoHCurlClient>(timeout);
}

Resolver::Resolver(std::unique_ptr<AresClient> ares_client,
                   std::unique_ptr<DoHCurlClientInterface> curl_client,
                   std::unique_ptr<net_base::SocketFactory> socket_factory,
                   bool disable_probe,
                   std::unique_ptr<Metrics> metrics)
    : logger_(base::DoNothing()),
      socket_factory_(std::move(socket_factory)),
      always_on_doh_(false),
      doh_enabled_(false),
      disable_probe_(disable_probe),
      metrics_(std::move(metrics)),
      ares_client_(std::move(ares_client)),
      curl_client_(std::move(curl_client)) {}

bool Resolver::ListenTCP(struct sockaddr* addr) {
  std::unique_ptr<net_base::Socket> tcp_src =
      socket_factory_->Create(addr->sa_family, SOCK_STREAM | SOCK_NONBLOCK);
  if (!tcp_src) {
    PLOG(ERROR) << *this << " Failed to create TCP socket";
    return false;
  }

  socklen_t len =
      addr->sa_family == AF_INET ? sizeof(sockaddr_in) : sizeof(sockaddr_in6);
  if (!tcp_src->Bind(addr, len)) {
    PLOG(ERROR) << *this << " Cannot bind TCP listening socket to " << *addr;
    return false;
  }

  if (!tcp_src->Listen(kMaxClientTcpConn)) {
    PLOG(ERROR) << *this << " Cannot listen on " << *addr;
    return false;
  }

  // Run the accept loop.
  LOG(INFO) << *this << " Accepting TCP connections on " << *addr;
  tcp_src_watcher_ = base::FileDescriptorWatcher::WatchReadable(
      tcp_src->Get(), base::BindRepeating(&Resolver::OnTCPConnection,
                                          weak_factory_.GetWeakPtr()));
  tcp_src_ = std::move(tcp_src);

  return true;
}

bool Resolver::ListenUDP(struct sockaddr* addr) {
  std::unique_ptr<net_base::Socket> udp_src =
      socket_factory_->Create(addr->sa_family, SOCK_DGRAM | SOCK_NONBLOCK);
  if (!udp_src) {
    PLOG(ERROR) << *this << " Failed to create UDP socket";
    return false;
  }

  socklen_t len =
      addr->sa_family == AF_INET ? sizeof(sockaddr_in) : sizeof(sockaddr_in6);
  if (!udp_src->Bind(addr, len)) {
    PLOG(ERROR) << *this << " Cannot bind UDP socket to " << *addr;
    return false;
  }

  // Start listening.
  LOG(INFO) << *this << " Accepting UDP queries on " << *addr;
  udp_src_watcher_ = base::FileDescriptorWatcher::WatchReadable(
      udp_src->Get(),
      base::BindRepeating(&Resolver::OnDNSQuery, weak_factory_.GetWeakPtr(),
                          udp_src->Get(), SOCK_DGRAM));
  udp_src_ = std::move(udp_src);
  return true;
}

void Resolver::OnTCPConnection() {
  struct sockaddr_storage client_src = {};
  socklen_t sockaddr_len = sizeof(client_src);
  std::unique_ptr<net_base::Socket> client_conn =
      tcp_src_->Accept((struct sockaddr*)&client_src, &sockaddr_len);
  if (!client_conn) {
    PLOG(ERROR) << *this << " Failed to accept TCP connection";
    return;
  }
  tcp_connections_.emplace(
      client_conn->Get(),
      new TCPConnection(std::move(client_conn),
                        base::BindRepeating(&Resolver::OnDNSQuery,
                                            weak_factory_.GetWeakPtr())));
}

void Resolver::HandleAresResult(base::WeakPtr<SocketFd> sock_fd,
                                base::WeakPtr<ProbeState> probe_state,
                                int status,
                                unsigned char* msg,
                                size_t len) {
  // Query is already handled.
  if (!sock_fd) {
    return;
  }

  // Query failed, restart probing.
  // Errors that may be caused by its query's data are not considered as
  // failures:
  // - ARES_FORMERR means that the query data is incorrect.
  // - ARES_ENODATA means that the domain has no answers.
  // - ARES_ENOTIMP means that the operation requested is not implemented.
  // We don't treat this as an error as the user can create these packets
  // manually.
  static const std::set<int> query_success_statuses = {
      ARES_SUCCESS, ARES_EFORMERR, ARES_ENODATA, ARES_ENOTIMP};
  if (probe_state && probe_state->validated &&
      !base::Contains(query_success_statuses, status)) {
    auto target = probe_state->target;
    // |probe_state| will be invalidated by RestartProbe.
    RestartProbe(probe_state);
    int attempt = sock_fd->num_retries + 1;
    LOG(ERROR) << *this << " Do53 query to " << target << " failed after "
               << attempt << " attempt: " << ares_strerror(status) << ". "
               << validated_name_servers_.size() << "/" << name_servers_.size()
               << " validated name servers";
  }

  sock_fd->num_active_queries--;
  // Don't process failing result that is not the last result.
  if (status != ARES_SUCCESS && sock_fd->num_active_queries > 0)
    return;

  sock_fd->timer.StopResolve(status == ARES_SUCCESS);
  if (metrics_)
    metrics_->RecordQueryResult(Metrics::QueryType::kPlainText,
                                AresStatusMetric(status));

  if (status == ARES_SUCCESS) {
    ReplyDNS(sock_fd, msg, len);
    sock_fds_.erase(sock_fd->id);
    return;
  }

  // Process the last unsuccessful result.
  // Retry query upon failure.
  if (sock_fd->num_retries++ >= max_num_retries_) {
    LOG(ERROR) << *this
               << " Failed to do ares lookup: " << ares_strerror(status);
    sock_fds_.erase(sock_fd->id);
    return;
  }

  // Retry resolving the domain.
  base::SingleThreadTaskRunner::GetCurrentDefault()->PostTask(
      FROM_HERE, base::BindOnce(&Resolver::Resolve, weak_factory_.GetWeakPtr(),
                                sock_fd, false /* fallback */));
}

void Resolver::HandleCurlResult(base::WeakPtr<SocketFd> sock_fd,
                                base::WeakPtr<ProbeState> probe_state,
                                const DoHCurlClient::CurlResult& res,
                                unsigned char* msg,
                                size_t len) {
  // Query is already handled.
  if (!sock_fd) {
    return;
  }

  // Query failed, restart probing.
  if (probe_state && probe_state->validated && res.http_code != kHTTPOk) {
    auto target = probe_state->target;
    // |probe_state| will be invalidated by RestartProbe.
    RestartProbe(probe_state);
    int attempt = sock_fd->num_retries + 1;
    LOG(WARNING) << *this << " DoH query to " << target << " failed after "
                 << attempt << " attempt, http status code: " << res.http_code
                 << ". " << validated_doh_providers_.size() << "/"
                 << doh_providers_.size() << " validated DoH providers";
  }

  sock_fd->num_active_queries--;
  // Don't process failing result that is not the last result.
  if (res.http_code != kHTTPOk && sock_fd->num_active_queries > 0)
    return;

  sock_fd->timer.StopResolve(res.curl_code == CURLE_OK);
  if (metrics_)
    metrics_->RecordQueryResult(Metrics::QueryType::kDnsOverHttps,
                                CurlCodeMetric(res.curl_code), res.http_code);

  // Process result.
  if (res.curl_code != CURLE_OK) {
    LOG(ERROR) << *this << " DoH resolution failed: "
               << curl_easy_strerror(res.curl_code);
    if (always_on_doh_) {
      // TODO(jasongustaman): Send failure reply with RCODE.
      sock_fds_.erase(sock_fd->id);
      return;
    }
    base::SingleThreadTaskRunner::GetCurrentDefault()->PostTask(
        FROM_HERE,
        base::BindOnce(&Resolver::Resolve, weak_factory_.GetWeakPtr(), sock_fd,
                       true /* fallback */));
    return;
  }

  switch (res.http_code) {
    case kHTTPOk: {
      ReplyDNS(sock_fd, msg, len);
      sock_fds_.erase(sock_fd->id);
      return;
    }
    case kHTTPTooManyRequests: {
      if (sock_fd->num_retries >= max_num_retries_) {
        LOG(ERROR) << *this << " Failed to resolve hostname, retried "
                   << max_num_retries_ << " tries";
        sock_fds_.erase(sock_fd->id);
        return;
      }

      // Add jitter to avoid coordinated spikes of retries.
      base::TimeDelta retry_delay_jitter =
          (1 - (base::RandDouble() * kRetryJitterMultiplier)) * retry_delay_;

      // Retry resolving the domain.
      base::SingleThreadTaskRunner::GetCurrentDefault()->PostDelayedTask(
          FROM_HERE,
          base::BindOnce(&Resolver::Resolve, weak_factory_.GetWeakPtr(),
                         sock_fd, false /* fallback */),
          retry_delay_jitter);
      sock_fd->num_retries++;
      return;
    }
    default: {
      LOG(ERROR) << *this << " Failed to do curl lookup, HTTP status code: "
                 << res.http_code;
      if (always_on_doh_) {
        // TODO(jasongustaman): Send failure reply with RCODE.
        sock_fds_.erase(sock_fd->id);
        return;
      }
      base::SingleThreadTaskRunner::GetCurrentDefault()->PostTask(
          FROM_HERE,
          base::BindOnce(&Resolver::Resolve, weak_factory_.GetWeakPtr(),
                         sock_fd, true /* fallback */));
    }
  }
}

void Resolver::HandleDoHProbeResult(base::WeakPtr<ProbeState> probe_state,
                                    const ProbeData& probe_data,
                                    const DoHCurlClient::CurlResult& res,
                                    unsigned char* msg,
                                    size_t len) {
  if (!probe_state) {
    return;
  }

  int attempt = probe_data.num_retries + 1;
  auto now = base::Time::Now();
  auto attempt_latency = now - probe_data.start_time;

  if (res.curl_code != CURLE_OK) {
    LOG(INFO) << *this << " DoH probe attempt " << attempt << " to "
              << probe_state->target
              << " failed: " << curl_easy_strerror(res.curl_code) << " ("
              << attempt_latency << ")";
    return;
  }
  if (res.http_code != kHTTPOk) {
    LOG(INFO) << *this << " DoH probe attempt " << attempt << " to "
              << probe_state->target
              << " failed, HTTP status code: " << res.http_code << " ("
              << attempt_latency << ")";
    return;
  }

  validated_doh_providers_.push_back(probe_state->target);

  LOG(INFO) << *this << " DoH probe attempt " << attempt << " to "
            << probe_state->target << " succeeded (" << attempt_latency << "). "
            << validated_doh_providers_.size() << "/" << doh_providers_.size()
            << " validated DoH providers";

  // Clear the old probe state to stop running probes.
  // Entry must be valid as |probe_state| is still valid.
  const auto& doh_provider = doh_providers_.find(probe_state->target);
  doh_provider->second = std::make_unique<ProbeState>(
      doh_provider->first, probe_state->doh, /*validated=*/true);
}

void Resolver::HandleDo53ProbeResult(base::WeakPtr<ProbeState> probe_state,
                                     const ProbeData& probe_data,
                                     int status,
                                     unsigned char* msg,
                                     size_t len) {
  if (metrics_) {
    metrics_->RecordProbeResult(probe_data.family, probe_data.num_retries,
                                AresStatusMetric(status));
  }
  if (!probe_state) {
    return;
  }

  int attempt = probe_data.num_retries + 1;
  auto now = base::Time::Now();
  auto attempt_latency = now - probe_data.start_time;

  if (status != ARES_SUCCESS) {
    LOG(INFO) << *this << " Do53 probe attempt " << attempt << " to "
              << probe_state->target << " failed: " << ares_strerror(status)
              << " (" << attempt_latency << ")";
    return;
  }

  validated_name_servers_.push_back(probe_state->target);

  LOG(INFO) << *this << " Do53 probe attempt " << attempt << " to "
            << probe_state->target << " succeeded (" << attempt_latency << "). "
            << validated_name_servers_.size() << "/" << name_servers_.size()
            << " validated name servers";

  // Clear the old probe state to stop running probes.
  // Entry must be valid as |probe_state| is still valid.
  const auto& name_server = name_servers_.find(probe_state->target);
  name_server->second = std::make_unique<ProbeState>(
      name_server->first, name_server->second->doh, /*validated=*/true);
}

void Resolver::ReplyDNS(base::WeakPtr<SocketFd> sock_fd,
                        unsigned char* msg,
                        size_t len) {
  sock_fd->timer.StartReply();
  // For TCP, DNS messages have an additional 2-bytes header representing
  // the length of the query. Add the additional header for the reply.
  uint16_t dns_len = htons(len);
  struct iovec iov_out[2];
  iov_out[0].iov_base = &dns_len;
  iov_out[0].iov_len = 2;
  // For UDP, skip the additional header. By setting |iov_len| to 0, the
  // additional header |dns_len| will not be sent.
  if (sock_fd->type == SOCK_DGRAM) {
    iov_out[0].iov_len = 0;
  }
  iov_out[1].iov_base = static_cast<void*>(msg);
  iov_out[1].iov_len = len;
  struct msghdr hdr = {
      .msg_name = nullptr,
      .msg_namelen = 0,
      .msg_iov = iov_out,
      .msg_iovlen = 2,
      .msg_control = nullptr,
      .msg_controllen = 0,
  };
  if (sock_fd->type == SOCK_DGRAM) {
    hdr.msg_name = &sock_fd->src;
    hdr.msg_namelen = sock_fd->socklen;
  }
  const bool ok = sendmsg(sock_fd->fd, &hdr, 0) >= 0;
  sock_fd->timer.StopReply(ok);
  if (!ok) {
    PLOG(ERROR) << *this << " sendmsg() " << sock_fd->fd << " failed";
  }
}

void Resolver::SetNameServers(const std::vector<std::string>& name_servers) {
  SetServers(name_servers, /*doh=*/false);
}

void Resolver::SetDoHProviders(const std::vector<std::string>& doh_providers,
                               bool always_on_doh) {
  always_on_doh_ = always_on_doh;
  doh_enabled_ = !doh_providers.empty();

  SetServers(doh_providers, /*doh=*/true);
}

void Resolver::SetServers(const std::vector<std::string>& new_servers,
                          bool doh) {
  auto& servers = doh ? doh_providers_ : name_servers_;
  auto& validated_servers =
      doh ? validated_doh_providers_ : validated_name_servers_;
  const std::set<std::string> new_servers_set(new_servers.begin(),
                                              new_servers.end());
  bool servers_equal = true;

  // Remove any removed servers.
  for (auto it = servers.begin(); it != servers.end();) {
    if (base::Contains(new_servers_set, it->first)) {
      ++it;
      continue;
    }
    it = servers.erase(it);
    servers_equal = false;
  }

  // Remove any removed servers from validated servers.
  for (auto it = validated_servers.begin(); it != validated_servers.end();) {
    if (base::Contains(new_servers_set, *it)) {
      ++it;
      continue;
    }
    it = validated_servers.erase(it);
  }

  // Probe the new servers.
  for (const auto& new_server : new_servers_set) {
    if (base::Contains(servers, new_server)) {
      continue;
    }
    const auto& probe_state =
        servers
            .emplace(new_server, std::make_unique<ProbeState>(new_server, doh))
            .first;
    Probe(probe_state->second->weak_factory.GetWeakPtr());
    servers_equal = false;
  }

  if (servers_equal)
    return;

  if (doh) {
    LOG(INFO) << *this << " DoH providers are updated, "
              << validated_doh_providers_.size() << "/" << doh_providers_.size()
              << " validated DoH providers";
  } else {
    LOG(INFO) << *this << " Name servers are updated, "
              << validated_name_servers_.size() << "/" << name_servers_.size()
              << " validated name servers";
  }
}

std::unique_ptr<Resolver::SocketFd> Resolver::PopPendingSocketFd(int fd) {
  auto it = pending_sock_fds_.find(fd);
  if (it != pending_sock_fds_.end()) {
    auto sock_fd = std::move(it->second);
    pending_sock_fds_.erase(it);
    return sock_fd;
  }
  return {};
}

void Resolver::OnDNSQuery(int fd, int type) {
  auto sock_fd = PopPendingSocketFd(fd);
  if (!sock_fd) {
    // Initialize SocketFd to carry necessary data.
    sock_fd = std::make_unique<SocketFd>(type, fd);
    // Metrics will be recorded automatically when this object is deleted.
    sock_fd->timer.set_metrics(metrics_.get());
  }

  // If the current buffer is full, resize the container.
  ssize_t buf_size = sock_fd->try_resize();

  // For TCP, it is possible for the packets to be chunked. Move the buffer
  // to the last empty position and adjust the size.
  // For UDP, |sock_fd->len| is always initialized to 0.
  char* buf = sock_fd->buf.data() + sock_fd->len;
  buf_size -= sock_fd->len;

  // For TCP, DNS has an additional 2-bytes header representing the length
  // of the query. Move the receiving buffer, so it is 4-bytes aligned.
  if (sock_fd->type == SOCK_STREAM) {
    buf += kTCPBufferPaddingLength;
    buf_size -= kTCPBufferPaddingLength;
  }

  // Only the last recvfrom call is considered for receive metrics.
  sock_fd->timer.StartReceive();
  ssize_t len = Receive(fd, buf, buf_size,
                        reinterpret_cast<struct sockaddr*>(&sock_fd->src),
                        &sock_fd->socklen);
  // Assume success - on failure, the correct value will be recorded.
  sock_fd->timer.StopReceive(true);
  if (len < 0) {
    sock_fd->timer.StopReceive(false);
    PLOG(WARNING) << *this << " recvfrom failed";
    return;
  }
  // Handle TCP connection closed.
  if (len == 0) {
    sock_fd->timer.StopReceive(false);
    tcp_connections_.erase(fd);
    return;
  }
  sock_fd->len += len;

  HandleDNSQuery(std::move(sock_fd));
}

void Resolver::HandleDNSQuery(std::unique_ptr<SocketFd> sock_fd) {
  // Handle DNS query over UDP.
  if (sock_fd->type == SOCK_DGRAM) {
    const auto& sock_fd_it =
        sock_fds_.emplace(sock_fd->id, std::move(sock_fd)).first;
    Resolve(sock_fd_it->second->weak_factory.GetWeakPtr());
    return;
  }

  // Handle DNS query over TCP.
  while (sock_fd) {
    if (sock_fd->len < kDNSTCPHeaderLength) {
      pending_sock_fds_.emplace(sock_fd->fd, std::move(sock_fd));
      return;
    }

    // Check if the current buffer contains a complete DNS query noted by the
    // length taken from the TCP header.
    uint16_t dns_len = ntohs(*reinterpret_cast<uint16_t*>(sock_fd->msg));
    if (sock_fd->len < kDNSTCPHeaderLength + dns_len) {
      pending_sock_fds_.emplace(sock_fd->fd, std::move(sock_fd));
      return;
    }

    // Check if the current buffer contains extra data.
    // Move any remaining data to a new SocketFd.
    std::unique_ptr<SocketFd> tmp_sock_fd = {};
    ssize_t tcp_dns_len = kDNSTCPHeaderLength + dns_len;
    if (sock_fd->len > tcp_dns_len) {
      tmp_sock_fd = std::make_unique<SocketFd>(sock_fd->type, sock_fd->fd,
                                               sock_fd->buf.size());
      tmp_sock_fd->len = sock_fd->len - tcp_dns_len;
      // This copy is safe because the buffer size of |tmp_sock_fd| and
      // |sock_fd| is equal and only partial data of |sock_fd| is copied.
      memcpy(tmp_sock_fd->msg, sock_fd->msg + tcp_dns_len, tmp_sock_fd->len);
      sock_fd->len = tcp_dns_len;
    }

    // Start the DNS query for a complete DNS query data.
    const auto& sock_fd_it =
        sock_fds_.emplace(sock_fd->id, std::move(sock_fd)).first;
    Resolve(sock_fd_it->second->weak_factory.GetWeakPtr());

    // Process additional query data.
    sock_fd = std::move(tmp_sock_fd);
  }
}

bool Resolver::ResolveDNS(base::WeakPtr<SocketFd> sock_fd, bool doh) {
  if (!sock_fd) {
    LOG(ERROR) << *this
               << " Unexpected ResolveDNS() call with deleted SocketFd";
    return false;
  }

  const auto query_type =
      doh ? Metrics::QueryType::kDnsOverHttps : Metrics::QueryType::kPlainText;
  const auto& name_servers = GetActiveNameServers();
  if (name_servers.empty()) {
    LOG(ERROR) << *this << " Name server list must not be empty";
    if (metrics_) {
      metrics_->RecordQueryResult(query_type,
                                  Metrics::QueryError::kEmptyNameServers);
    }
    return false;
  }

  const auto& doh_providers = GetActiveDoHProviders();
  if (doh && doh_providers.empty()) {
    // No DoH providers are currently validated, fallback to Do53.
    if (!doh_providers_.empty()) {
      return false;
    }
    LOG(ERROR) << *this << " DoH provider list must not be empty";
    if (metrics_) {
      metrics_->RecordQueryResult(Metrics::QueryType::kDnsOverHttps,
                                  Metrics::QueryError::kEmptyDoHProviders);
    }
    return false;
  }

  // Start multiple concurrent queries.
  const auto& targets = doh ? doh_providers : name_servers;
  for (const auto& target : targets) {
    if (doh) {
      if (!curl_client_->Resolve(
              sock_fd->get_message(), sock_fd->get_length(),
              base::BindRepeating(
                  &Resolver::HandleCurlResult, weak_factory_.GetWeakPtr(),
                  sock_fd, doh_providers_[target]->weak_factory.GetWeakPtr()),
              name_servers, target)) {
        continue;
      }
    } else {
      if (!ares_client_->Resolve(
              reinterpret_cast<const unsigned char*>(sock_fd->get_message()),
              sock_fd->get_length(),
              base::BindRepeating(
                  &Resolver::HandleAresResult, weak_factory_.GetWeakPtr(),
                  sock_fd, name_servers_[target]->weak_factory.GetWeakPtr()),
              target, sock_fd->type)) {
        continue;
      }
    }
    if (++sock_fd->num_active_queries >= kMaxConcurrentQueries) {
      break;
    }
  }

  if (sock_fd->num_active_queries > 0)
    return true;

  LOG(ERROR) << *this << " No requests successfully started for query";
  if (metrics_) {
    metrics_->RecordQueryResult(
        query_type, Metrics::QueryError::kClientInitializationError);
  }
  return false;
}

std::vector<std::string> Resolver::GetActiveDoHProviders() {
  if (!always_on_doh_ || !validated_doh_providers_.empty())
    return validated_doh_providers_;

  std::vector<std::string> doh_providers;
  for (const auto& doh_provider : doh_providers_) {
    doh_providers.push_back(doh_provider.first);
  }
  return doh_providers;
}

std::vector<std::string> Resolver::GetActiveNameServers() {
  if (!validated_name_servers_.empty())
    return validated_name_servers_;

  std::vector<std::string> name_servers;
  for (const auto& name_server : name_servers_) {
    name_servers.push_back(name_server.first);
  }
  return name_servers;
}

void Resolver::RestartProbe(base::WeakPtr<ProbeState> probe_state) {
  if (!probe_state)
    return;

  auto& targets = probe_state->doh ? doh_providers_ : name_servers_;
  auto& validated_targets =
      probe_state->doh ? validated_doh_providers_ : validated_name_servers_;
  validated_targets.erase(
      std::remove(validated_targets.begin(), validated_targets.end(),
                  probe_state->target),
      validated_targets.end());

  const auto& target = targets.find(probe_state->target);
  target->second =
      std::make_unique<ProbeState>(target->first, probe_state->doh);
  Probe(target->second->weak_factory.GetWeakPtr());
}

void Resolver::Probe(base::WeakPtr<ProbeState> probe_state) {
  if (disable_probe_)
    return;

  if (!probe_state)
    return;

  // Schedule the next probe now as the probe may run for a long time.
  base::SingleThreadTaskRunner::GetCurrentDefault()->PostDelayedTask(
      FROM_HERE,
      base::BindOnce(&Resolver::Probe, weak_factory_.GetWeakPtr(), probe_state),
      GetTimeUntilProbe(probe_state->num_retries));

  // Run the probe.
  const auto target_ip =
      net_base::IPAddress::CreateFromString(probe_state->target);
  const sa_family_t target_family =
      target_ip ? net_base::ToSAFamily(target_ip->GetFamily()) : AF_UNSPEC;
  const ProbeData probe_data = {target_family, probe_state->num_retries,
                                base::Time::Now()};
  if (probe_state->doh) {
    curl_client_->Resolve(kDNSQueryGstatic, sizeof(kDNSQueryGstatic),
                          base::BindRepeating(&Resolver::HandleDoHProbeResult,
                                              weak_factory_.GetWeakPtr(),
                                              probe_state, probe_data),
                          GetActiveNameServers(), probe_state->target);
  } else {
    ares_client_->Resolve(
        reinterpret_cast<const unsigned char*>(kDNSQueryGstatic),
        sizeof(kDNSQueryGstatic),
        base::BindRepeating(&Resolver::HandleDo53ProbeResult,
                            weak_factory_.GetWeakPtr(), probe_state,
                            probe_data),
        probe_state->target);
  }
  probe_state->num_retries++;
}

void Resolver::Resolve(base::WeakPtr<SocketFd> sock_fd, bool fallback) {
  if (!sock_fd) {
    LOG(ERROR) << *this << " Unexpected Resolve() call with deleted SocketFd";
    return;
  }

  if (doh_enabled_ && !fallback) {
    sock_fd->timer.StartResolve(true);
    if (ResolveDNS(sock_fd, /*doh=*/true))
      return;

    sock_fd->timer.StopResolve(false);
  }
  if (!always_on_doh_) {
    sock_fd->timer.StartResolve();
    if (ResolveDNS(sock_fd, /*doh=*/false))
      return;

    sock_fd->timer.StopResolve(false);
  }

  // Construct and send a response indicating that there is a failure.
  patchpanel::DnsResponse response =
      ConstructServFailResponse(sock_fd->get_message(), sock_fd->get_length());
  ReplyDNS(sock_fd,
           reinterpret_cast<unsigned char*>(response.io_buffer()->data()),
           response.io_buffer_size());

  // Query is completed, remove SocketFd.
  sock_fds_.erase(sock_fd->id);
}

patchpanel::DnsResponse Resolver::ConstructServFailResponse(const char* msg,
                                                            int len) {
  // Construct a DNS query from the message buffer.
  std::optional<patchpanel::DnsQuery> query;
  if (len > 0 && len <= dns_proxy::kMaxDNSBufSize) {
    scoped_refptr<patchpanel::IOBufferWithSize> query_buf =
        base::MakeRefCounted<patchpanel::IOBufferWithSize>(len);
    memcpy(query_buf->data(), msg, len);
    query = patchpanel::DnsQuery(query_buf);
  }

  // Set the query id as 0 if the query is invalid.
  uint16_t query_id = 0;
  if (query.has_value() && query->Parse(len)) {
    query_id = query->id();
  } else {
    query.reset();
  }

  // Returns RCODE SERVFAIL response corresponding to the query.
  patchpanel::DnsResponse response(query_id, false /* is_authoritative */,
                                   {} /* answers */, {} /* authority_records */,
                                   {} /* additional_records */, query,
                                   patchpanel::dns_protocol::kRcodeSERVFAIL);
  return response;
}

void Resolver::SetProbingEnabled(bool enable_probe) {
  disable_probe_ = !enable_probe;
}

ssize_t Resolver::Receive(int fd,
                          char* buffer,
                          size_t buffer_size,
                          struct sockaddr* src_addr,
                          socklen_t* addrlen) {
  return recvfrom(fd, buffer, buffer_size, 0, src_addr, addrlen);
}
}  // namespace dns_proxy