#include "netstack/dns_resolver.h"

#include <algorithm>
#include <cctype>
#include <cstring>

#include "common/error.h"
#include "common/logging.h"
#include "common/metrics.h"

namespace ovg::netstack {
namespace {

constexpr const char *kMod = "dns";

constexpr uint16_t kPortDns = 53;
constexpr size_t kHeaderLen = 12;
constexpr uint16_t kTypeA = 1;
constexpr uint16_t kClassIn = 1;

// UDP answers larger than this are truncated by definition (no EDNS0 means a
// 512-byte payload limit), but tunnels have been seen to deliver more, so the
// buffer is a full MTU rather than 512.
constexpr size_t kRxBuf = 1500;

// Floor on a single server's share of the total budget. Below this, a server on
// a high-latency VPN node is being written off before it has had a chance to
// answer, and the failover is pure loss.
constexpr Millis kMinPerServer{700};

metrics::Counter *queries_total() {
  static auto *c = metrics::counter("ovg_dns_queries_total",
                                    "DNS queries sent through a tunnel");
  return c;
}
metrics::Counter *cache_hits_total() {
  static auto *c =
      metrics::counter("ovg_dns_cache_hits_total", "DNS lookups served from cache");
  return c;
}
metrics::Counter *failures_total() {
  static auto *c = metrics::counter("ovg_dns_failures_total",
                                    "DNS lookups that produced no address");
  return c;
}

std::string lowercase(const std::string &s) {
  std::string r = s;
  std::transform(r.begin(), r.end(), r.begin(), [](unsigned char c) {
    return static_cast<char>(std::tolower(c));
  });
  return r;
}

inline uint16_t rd16(const uint8_t *p) {
  return static_cast<uint16_t>((p[0] << 8) | p[1]);
}
inline uint32_t rd32(const uint8_t *p) {
  return (static_cast<uint32_t>(p[0]) << 24) |
         (static_cast<uint32_t>(p[1]) << 16) |
         (static_cast<uint32_t>(p[2]) << 8) | static_cast<uint32_t>(p[3]);
}

// Advances *pos past one wire-format name. Compression pointers are not
// followed -- a pointer terminates the name, and every name in a response we
// care about is one we are only skipping over.
bool skip_name(const uint8_t *d, size_t len, size_t *pos) {
  size_t p = *pos;
  for (int labels = 0; labels < 128; ++labels) {
    if (p >= len) return false;
    const uint8_t l = d[p];
    if ((l & 0xC0) == 0xC0) {
      if (p + 1 >= len) return false;
      *pos = p + 2;
      return true;
    }
    if ((l & 0xC0) != 0) return false;  // reserved label type
    if (l == 0) {
      *pos = p + 1;
      return true;
    }
    p += 1u + l;
  }
  return false;
}

}  // namespace

namespace dns {

bool build_query(const std::string &name, uint16_t id,
                 std::vector<uint8_t> *out) {
  if (out == nullptr) return false;

  // A single trailing dot is the root label and is implied by the encoding.
  std::string n = name;
  if (!n.empty() && n.back() == '.') n.pop_back();
  if (n.empty() || n.size() > 253) return false;

  std::vector<uint8_t> qname;
  qname.reserve(n.size() + 2);
  size_t start = 0;
  while (start <= n.size()) {
    const size_t dot = n.find('.', start);
    const size_t end = (dot == std::string::npos) ? n.size() : dot;
    const size_t label_len = end - start;
    if (label_len == 0 || label_len > 63) return false;
    qname.push_back(static_cast<uint8_t>(label_len));
    qname.insert(qname.end(), n.begin() + static_cast<long>(start),
                 n.begin() + static_cast<long>(end));
    if (dot == std::string::npos) break;
    start = dot + 1;
  }
  qname.push_back(0);
  if (qname.size() > 255) return false;

  out->clear();
  out->reserve(kHeaderLen + qname.size() + 4);
  out->push_back(static_cast<uint8_t>(id >> 8));
  out->push_back(static_cast<uint8_t>(id & 0xFF));
  out->push_back(0x01);  // RD
  out->push_back(0x00);
  out->push_back(0x00);
  out->push_back(0x01);  // QDCOUNT = 1
  for (int i = 0; i < 6; ++i) out->push_back(0x00);  // AN/NS/AR = 0
  out->insert(out->end(), qname.begin(), qname.end());
  out->push_back(0x00);
  out->push_back(static_cast<uint8_t>(kTypeA));
  out->push_back(0x00);
  out->push_back(static_cast<uint8_t>(kClassIn));
  return true;
}

bool parse_response(const uint8_t *data, size_t len, ParseResult *out) {
  if (data == nullptr || out == nullptr || len < kHeaderLen) return false;

  const uint16_t flags = rd16(data + 2);
  if ((flags & 0x8000) == 0) return false;  // a query, not a response

  out->id = rd16(data);
  out->rcode = flags & 0x000F;
  out->truncated = (flags & 0x0200) != 0;
  out->addrs.clear();
  out->min_ttl = 0;

  const uint16_t qdcount = rd16(data + 4);
  const uint16_t ancount = rd16(data + 6);

  size_t pos = kHeaderLen;
  for (uint16_t i = 0; i < qdcount; ++i) {
    if (!skip_name(data, len, &pos)) return false;
    if (pos + 4 > len) return false;
    pos += 4;  // QTYPE + QCLASS
  }

  bool have_ttl = false;
  for (uint16_t i = 0; i < ancount; ++i) {
    if (!skip_name(data, len, &pos)) return false;
    if (pos + 10 > len) return false;
    const uint16_t rtype = rd16(data + pos);
    const uint16_t rclass = rd16(data + pos + 2);
    const uint32_t ttl = rd32(data + pos + 4);
    const uint16_t rdlen = rd16(data + pos + 8);
    pos += 10;
    if (pos + rdlen > len) return false;

    if (rtype == kTypeA && rclass == kClassIn && rdlen == 4) {
      out->addrs.push_back(IpAddress::from_bytes_v4(data + pos));
      // The shortest TTL in the set governs the whole set: caching an address
      // past its own TTL because a sibling record lived longer is how stale
      // entries outlive a failover.
      if (!have_ttl || ttl < out->min_ttl) {
        out->min_ttl = ttl;
        have_ttl = true;
      }
    }
    // CNAMEs are skipped rather than chased: a resolver that answers a CNAME
    // without the A record it points at is broken, and every real one inlines
    // the whole chain.
    pos += rdlen;
  }
  return true;
}

}  // namespace dns

// ---------------------------------------------------------------------------
// DnsResolver
// ---------------------------------------------------------------------------

std::shared_ptr<DnsResolver> DnsResolver::create(std::shared_ptr<Netif> netif,
                                                 DnsConfig cfg) {
  return std::shared_ptr<DnsResolver>(
      new DnsResolver(std::move(netif), std::move(cfg)));
}

DnsResolver::DnsResolver(std::shared_ptr<Netif> netif, DnsConfig cfg)
    : netif_(std::move(netif)),
      cfg_(std::move(cfg)),
      strand_(asio::make_strand(netif_->stack().io())),
      rng_(std::random_device{}()) {
  rx_buf_.resize(kRxBuf);

  auto add = [this](const IpAddress &a) {
    if (!a.valid() || !a.is_v4()) return;
    if (std::find(servers_.begin(), servers_.end(), a) != servers_.end()) return;
    servers_.push_back(a);
  };
  // Pushed servers first: they are inside the tunnel's own network and are the
  // only ones guaranteed to be reachable from it.
  for (const auto &a : netif_->dns_servers()) add(a);
  for (const auto &s : cfg_.fallback_servers) {
    if (auto a = IpAddress::parse(s)) add(*a);
  }

  if (servers_.empty()) {
    LOG_WARN(kMod,
             "{}: no usable DNS server (node pushed none and no fallback "
             "parsed); every lookup on this tunnel will fail",
             netif_->label());
  } else {
    std::string list;
    for (const auto &a : servers_) {
      if (!list.empty()) list += ", ";
      list += a.to_string();
    }
    LOG_DEBUG(kMod, "{}: resolver servers: {}", netif_->label(), list);
  }
}

DnsResolver::~DnsResolver() {
  // The last reference is gone, so nothing else can be touching this object --
  // which is what makes it safe to answer the stragglers inline rather than
  // posting to a strand that may never run again.
  const auto ec = make_error_code(Error::Cancelled);
  for (auto &[id, q] : by_id_) {
    (void)id;
    if (q->done) continue;
    q->done = true;
    for (auto &w : q->waiters) w(ec, {});
  }
  by_id_.clear();
  by_name_.clear();
}

void DnsResolver::async_resolve(const std::string &host, Handler h) {
  auto self = shared_from_this();
  asio::post(strand_, [self, host, h = std::move(h)]() mutable {
    self->start(host, std::move(h));
  });
}

void DnsResolver::start(const std::string &host, Handler h) {
  // A literal needs no server, no socket and no cache entry. SOCKS5 clients
  // send these constantly (anything that resolved on its own), so short-
  // circuiting here is not a micro-optimization.
  if (auto lit = IpAddress::parse(host)) {
    std::vector<IpAddress> one{*lit};
    asio::post(strand_, [h = std::move(h), one]() mutable {
      h(std::error_code{}, std::move(one));
    });
    return;
  }

  const std::string name = lowercase(host);

  std::vector<IpAddress> cached;
  if (cache_get(name, &cached)) {
    {
      std::lock_guard<std::mutex> lk(stats_mu_);
      stats_.cache_hits++;
    }
    cache_hits_total()->inc();
    asio::post(strand_, [h = std::move(h), cached]() mutable {
      h(std::error_code{}, std::move(cached));
    });
    return;
  }

  // Two SOCKS5 sessions opening the same site at once should cost one query,
  // not two -- and with 1000 concurrent connections that ratio matters.
  if (auto it = by_name_.find(name); it != by_name_.end()) {
    it->second->waiters.push_back(std::move(h));
    std::lock_guard<std::mutex> lk(stats_mu_);
    stats_.coalesced++;
    return;
  }

  auto q = std::make_shared<Query>(strand_);
  q->name = name;
  q->id = allocate_id();
  q->deadline = Clock::now() + cfg_.timeout;
  q->waiters.push_back(std::move(h));

  if (!dns::build_query(name, q->id, &q->wire)) {
    LOG_DEBUG(kMod, "{}: refusing to look up a malformed name: {}",
              netif_->label(), host);
    auto waiters = std::move(q->waiters);
    for (auto &w : waiters) {
      asio::post(strand_, [w = std::move(w)]() mutable {
        w(make_error_code(Error::ProtocolError), {});
      });
    }
    return;
  }

  by_id_[q->id] = q;
  by_name_[name] = q;
  {
    std::lock_guard<std::mutex> lk(stats_mu_);
    stats_.queries++;
  }
  queries_total()->inc();

  if (servers_.empty()) {
    finish(q, make_error_code(Error::ResolveFailed), {}, 0);
    return;
  }
  if (sock_) {
    send_query(q);
    return;
  }
  pending_open_.push_back(q);
  open_socket();
}

void DnsResolver::open_socket() {
  if (opening_ || sock_) return;
  opening_ = true;

  auto self = shared_from_this();
  netif_->async_open_udp(
      strand_, [self](const std::error_code &ec, UdpSocketPtr sock) {
        self->opening_ = false;
        auto pending = std::move(self->pending_open_);
        self->pending_open_.clear();

        if (ec) {
          LOG_WARN(kMod, "{}: cannot open a DNS socket: {}",
                   self->netif_->label(), ec.message());
          for (auto &q : pending) self->finish(q, ec, {}, 0);
          return;
        }
        self->sock_ = std::move(sock);
        LOG_DEBUG(kMod, "{}: resolver socket {}", self->netif_->label(),
                  self->sock_->local_endpoint().to_string());
        self->arm_receive();
        for (auto &q : pending) self->send_query(q);
      });
}

void DnsResolver::arm_receive() {
  if (!sock_ || receiving_) return;
  receiving_ = true;
  auto self = shared_from_this();
  sock_->async_receive_from(
      asio::buffer(rx_buf_),
      [self](const std::error_code &ec, size_t n, const Endpoint &from) {
        self->on_datagram(ec, n, from);
      });
}

void DnsResolver::send_query(const std::shared_ptr<Query> &q) {
  if (q->done) return;
  if (!sock_) {
    finish(q, make_error_code(Error::EgressGone), {}, 0);
    return;
  }
  const auto now = Clock::now();
  if (q->server_idx >= servers_.size() || now >= q->deadline) {
    // Every server either timed out or refused. Report the last RCODE if we got
    // one -- "no such host" and "nothing answered" are different problems for
    // whoever reads the log.
    const auto ec = make_error_code(Error::ResolveFailed);
    LOG_DEBUG(kMod, "{}: {} unresolved after {} server(s) (last rcode {})",
              netif_->label(), q->name, servers_.size(), q->last_rcode);
    finish(q, ec, {}, 0);
    return;
  }

  const Endpoint server(servers_[q->server_idx], kPortDns);

  // Split the total budget across the servers so that a dead first resolver
  // cannot consume it all, but never below the floor.
  Millis per = cfg_.timeout / static_cast<int>(std::max<size_t>(1, servers_.size()));
  if (per < kMinPerServer) per = kMinPerServer;
  const auto remaining =
      std::chrono::duration_cast<Millis>(q->deadline - now);
  q->timer.expires_after(std::min(per, remaining));

  auto self = shared_from_this();
  q->timer.async_wait([self, q](const std::error_code &ec) {
    if (ec) return;  // cancelled: the answer arrived
    self->on_query_timeout(q);
  });

  LOG_TRACE(kMod, "{}: query {} A {} -> {}", netif_->label(), q->id, q->name,
            server.to_string());

  sock_->async_send_to(asio::buffer(q->wire), server,
                       [self, q](const std::error_code &ec, size_t) {
                         if (!ec || q->done) return;
                         // A send that fails outright is a dead server; do not
                         // spend the timeout waiting to learn that.
                         q->timer.cancel();
                         q->server_idx++;
                         self->send_query(q);
                       });
}

void DnsResolver::on_query_timeout(const std::shared_ptr<Query> &q) {
  if (q->done) return;
  {
    std::lock_guard<std::mutex> lk(stats_mu_);
    stats_.timeouts++;
  }
  LOG_TRACE(kMod, "{}: query {} for {} timed out on server {}",
            netif_->label(), q->id, q->name, q->server_idx);
  q->server_idx++;
  send_query(q);
}

void DnsResolver::on_datagram(const std::error_code &ec, size_t n,
                              const Endpoint &from) {
  receiving_ = false;

  if (ec) {
    if (ec != asio::error::operation_aborted) {
      LOG_WARN(kMod, "{}: resolver socket failed: {}", netif_->label(),
               ec.message());
    }
    // Drop the socket so the next lookup opens a fresh one; the tunnel itself
    // may still be perfectly healthy.
    sock_.reset();
    fail_all(ec);
    return;
  }

  dns::ParseResult r;
  if (!dns::parse_response(rx_buf_.data(), n, &r)) {
    LOG_TRACE(kMod, "{}: discarded a {}-byte malformed response from {}",
              netif_->label(), n, from.to_string());
    arm_receive();
    return;
  }

  auto it = by_id_.find(r.id);
  if (it == by_id_.end()) {
    // Late answer to a query we already gave up on, or an unsolicited packet.
    arm_receive();
    return;
  }
  auto q = it->second;

  // Cheap anti-spoofing: only accept from an address we actually asked. Inside
  // a tunnel this is close to redundant, but the tunnel is a hostile network by
  // assumption -- it belongs to a stranger who volunteered a VPN node.
  const bool known_server =
      !from.is_domain() &&
      std::find(servers_.begin(), servers_.end(), from.address()) !=
          servers_.end();
  if (!known_server) {
    LOG_DEBUG(kMod, "{}: ignoring a response for query {} from {}, which is "
                    "not one of our servers",
              netif_->label(), r.id, from.to_string());
    arm_receive();
    return;
  }

  if (r.rcode != 0 || r.addrs.empty()) {
    // Fail over rather than trust it. VPNGate nodes push whatever resolver
    // their operator happened to have, and answering NXDOMAIN for names that
    // plainly exist is common enough that treating one negative answer as
    // authoritative would break the proxy on those nodes. The cost is that a
    // genuinely nonexistent name is asked of every server before it fails.
    q->last_rcode = r.rcode;
    LOG_TRACE(kMod, "{}: server {} answered rcode={} with {} address(es) for {}",
              netif_->label(), q->server_idx, r.rcode, r.addrs.size(), q->name);
    q->timer.cancel();
    q->server_idx++;
    send_query(q);
    arm_receive();
    return;
  }

  if (r.truncated) {
    LOG_TRACE(kMod, "{}: truncated answer for {}, keeping the {} address(es) "
                    "that fit",
              netif_->label(), q->name, r.addrs.size());
  }
  finish(q, {}, std::move(r.addrs), r.min_ttl);
  arm_receive();
}

void DnsResolver::finish(const std::shared_ptr<Query> &q,
                         const std::error_code &ec, std::vector<IpAddress> addrs,
                         uint32_t ttl_seconds) {
  if (q->done) return;
  q->done = true;
  q->timer.cancel();

  by_id_.erase(q->id);
  if (auto it = by_name_.find(q->name); it != by_name_.end() && it->second == q)
    by_name_.erase(it);

  if (!ec && !addrs.empty()) {
    cache_put(q->name, addrs, ttl_seconds);
    LOG_TRACE(kMod, "{}: {} -> {} (+{} more), ttl {}s", netif_->label(), q->name,
              addrs.front().to_string(), addrs.size() - 1, ttl_seconds);
  } else {
    {
      std::lock_guard<std::mutex> lk(stats_mu_);
      stats_.failures++;
    }
    failures_total()->inc();
  }

  auto waiters = std::move(q->waiters);
  q->waiters.clear();
  for (auto &w : waiters) {
    // Posted, not called: a handler that starts another lookup would otherwise
    // re-enter start() from inside finish(), while by_name_ is mid-erase.
    asio::post(strand_, [w = std::move(w), ec, addrs]() mutable {
      w(ec, addrs);
    });
  }
}

void DnsResolver::fail_all(const std::error_code &ec) {
  std::vector<std::shared_ptr<Query>> all;
  all.reserve(by_id_.size());
  for (auto &[id, q] : by_id_) {
    (void)id;
    all.push_back(q);
  }
  for (auto &q : all) finish(q, ec, {}, 0);
}

uint16_t DnsResolver::allocate_id() {
  // Random, not sequential: a predictable transaction ID is the other half of
  // the spoofing check above.
  for (int i = 0; i < 64; ++i) {
    const uint16_t id = static_cast<uint16_t>(rng_() & 0xFFFF);
    if (by_id_.find(id) == by_id_.end()) return id;
  }
  // 64 collisions means the table is saturated; any id will do at that point,
  // and the loser is answered by whichever query completes first.
  return static_cast<uint16_t>(rng_() & 0xFFFF);
}

void DnsResolver::cache_put(const std::string &name,
                            const std::vector<IpAddress> &addrs,
                            uint32_t ttl_seconds) {
  if (addrs.empty() || cfg_.cache_entries == 0) return;

  // Clamping both ends: a 30-second TTL would have us re-querying constantly on
  // a high-latency tunnel, and a 7-day one would outlive several node switches.
  Millis ttl(static_cast<int64_t>(ttl_seconds) * 1000);
  ttl = std::clamp(ttl, cfg_.min_ttl, cfg_.max_ttl);

  const auto expires = Clock::now() + ttl;
  auto [it, inserted] = cache_.insert_or_assign(name, CacheEntry{addrs, expires});
  if (inserted) cache_order_.push_back(name);

  while (cache_.size() > cfg_.cache_entries && !cache_order_.empty()) {
    // FIFO rather than LRU: keeping a per-entry access timestamp costs more
    // than it saves at this size, and the TTL clamp already bounds staleness.
    const std::string victim = cache_order_.front();
    cache_order_.pop_front();
    cache_.erase(victim);
  }
}

bool DnsResolver::cache_get(const std::string &name,
                            std::vector<IpAddress> *out) {
  auto it = cache_.find(name);
  if (it == cache_.end()) return false;
  if (Clock::now() >= it->second.expires) {
    // Deliberately left in place. cache_order_ mirrors cache_ one entry per
    // insertion; erasing here without touching the deque would let it drift and
    // grow without bound as names expire and are re-inserted. The stale entry
    // is either overwritten by the next answer or evicted with the rest.
    return false;
  }
  *out = it->second.addrs;
  return true;
}

void DnsResolver::clear_cache() {
  auto self = shared_from_this();
  asio::post(strand_, [self] {
    const size_t n = self->cache_.size();
    self->cache_.clear();
    self->cache_order_.clear();
    LOG_DEBUG(kMod, "{}: dropped {} cached name(s)", self->netif_->label(), n);
  });
}

DnsResolver::Stats DnsResolver::stats() const {
  std::lock_guard<std::mutex> lk(stats_mu_);
  Stats s = stats_;
  // cache_ and by_id_ are strand-owned; reading their size off-strand is a
  // benign race on a number that is only ever displayed.
  s.cached = cache_.size();
  s.in_flight = by_id_.size();
  return s;
}

}  // namespace ovg::netstack