/*
 * This file is part of PowerDNS or dnsdist.
 * Copyright -- PowerDNS.COM B.V. and its contributors
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of version 2 of the GNU General Public License as
 * published by the Free Software Foundation.
 *
 * In addition, for the avoidance of any doubt, permission is granted to
 * link this program with OpenSSL and to (re)distribute the binaries
 * produced as the result of such linking.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
 */
#pragma once
#include "ednsoptions.hh"
#include "misc.hh"
#include "iputils.hh"

class PacketCache : public boost::noncopyable
{
public:

  /* hash the packet from the provided position, which should point right after tje qname. This skips:
     - the query ID ;
     - EDNS Cookie options, if any ;
     - Any given option code present in optionsToSkip
  */
  static uint32_t hashAfterQname(const std::string_view& packet, uint32_t currentHash, size_t pos, const std::unordered_set<uint16_t>& optionsToSkip = {EDNSOptionCode::COOKIE}, const std::vector<uint16_t>& payloadRanks = {})
  {
    const size_t packetSize = packet.size();
    assert(packetSize >= sizeof(dnsheader));

    /* we need at least 2 (QTYPE) + 2 (QCLASS)

       + OPT root label (1), type (2), class (2) and ttl (4)
       + the OPT RR rdlen (2)
       = 15
    */
    const dnsheader_aligned dnsheaderdata(packet.data());
    const struct dnsheader *dh = dnsheaderdata.get();
    if (ntohs(dh->qdcount) != 1 || ntohs(dh->ancount) != 0 || ntohs(dh->nscount) != 0 || ntohs(dh->arcount) != 1 || (pos + 15) > packetSize) {
      if (packetSize > pos) {
        currentHash = burtle(reinterpret_cast<const unsigned char*>(&packet.at(pos)), packetSize - pos, currentHash);
      }
      return currentHash;
    }

    if (payloadRanks.empty()) {
      currentHash = burtle(reinterpret_cast<const unsigned char*>(&packet.at(pos)), 15, currentHash);
    }
    else {
      std::vector<unsigned char> optrr(packet.begin() + pos, packet.begin() + pos + 15);
      uint16_t bufSize = optrr.at(7) * 256 + optrr.at(8);
      auto it = std::upper_bound(payloadRanks.begin(), payloadRanks.end(), bufSize);
      if (it != payloadRanks.begin()) {
        it--;
        optrr[7] = (*it) >> 8;
        optrr[8] = (*it) & 0xff;
      }
      currentHash = burtle(reinterpret_cast<const unsigned char*>(&optrr.at(0)), 15, currentHash);
    }
    if ( (pos + 15) == packetSize ) {
      return currentHash;
    }

    /* skip the qtype (2), qclass (2) */
    /* root label (1), type (2), class (2) and ttl (4) */
    /* already hashed above */
    pos += 13;

    const uint16_t rdLen = ((static_cast<uint16_t>(packet.at(pos)) * 256) + static_cast<uint16_t>(packet.at(pos + 1)));
    /* skip the rd length */
    /* already hashed above */
    pos += 2;

    if (rdLen > (packetSize - pos)) {
      if (pos < packetSize) {
        currentHash = burtle(reinterpret_cast<const unsigned char*>(&packet.at(pos)), packetSize - pos, currentHash);
      }
      return currentHash;
    }

    uint16_t rdataRead = 0;
    uint16_t optionCode;
    uint16_t optionLen;

    while (pos < packetSize && rdataRead < rdLen && getNextEDNSOption(&packet.at(pos), rdLen - rdataRead, optionCode, optionLen)) {
      if (optionLen > (rdLen - rdataRead - 4)) {
        if (packetSize > pos) {
          currentHash = burtle(reinterpret_cast<const unsigned char*>(&packet.at(pos)), packetSize - pos, currentHash);
        }
        return currentHash;
      }

      if (optionsToSkip.count(optionCode) == 0) {
        /* hash the option code, length and content */
        currentHash = burtle(reinterpret_cast<const unsigned char*>(&packet.at(pos)), 4 + optionLen, currentHash);
      }
      else {
        /* skip option: hash only its code and length */
        currentHash = burtle(reinterpret_cast<const unsigned char*>(&packet.at(pos)), 4, currentHash);
      }

      pos += 4 + optionLen;
      rdataRead += 4 + optionLen;
    }

    if (pos < packetSize) {
      currentHash = burtle(reinterpret_cast<const unsigned char*>(&packet.at(pos)), packetSize - pos, currentHash);
    }

    return currentHash;
  }

  static uint32_t hashHeaderAndQName(const std::string& packet, size_t& pos)
  {
    const size_t packetSize = packet.size();
    assert(packetSize >= sizeof(dnsheader));
    // Quite some bits in the header are actually irrelevant for
    // incoming queries.  If we ever change that and ignore them for
    // hashing, don't forget to also adapt the `queryHeaderMatches`
    // code, as it should be consistent with the hash function.
    uint32_t currentHash = burtle(reinterpret_cast<const unsigned char*>(&packet.at(2)), sizeof(dnsheader) - 2, 0); // rest of dnsheader, skip id

    for (pos = sizeof(dnsheader); pos < packetSize; ) {
      const unsigned char labelLen = static_cast<unsigned char>(packet.at(pos));
      ++pos;
      if (labelLen == 0) {
        break;
      }
      pos = std::min(pos + labelLen, packetSize);
    }
    return burtleCI(reinterpret_cast<const unsigned char*>(&packet.at(sizeof(dnsheader))), pos - sizeof(dnsheader), currentHash);
  }

  /* hash the packet from the beginning, including the qname. This skips:
     - the query ID ;
     - EDNS Cookie options, if any ;
     - Any given option code present in optionsToSkip
  */
  static uint32_t canHashPacket(const std::string& packet, const std::unordered_set<uint16_t>& optionsToSkip = {EDNSOptionCode::COOKIE})
  {
    size_t pos = 0;
    uint32_t currentHash = hashHeaderAndQName(packet, pos);

    if (pos >= packet.size()) {
      return currentHash;
    }

    return hashAfterQname(packet, currentHash, pos, optionsToSkip);
  }

  static bool queryHeaderMatches(const std::string& cachedQuery, const std::string& query)
  {
    if (cachedQuery.size() != query.size()) {
      return false;
    }

    return (cachedQuery.compare(/* skip the ID */ 2, sizeof(dnsheader) - 2, query, 2, sizeof(dnsheader) - 2) == 0);
  }

  static bool queryMatches(const std::string& cachedQuery, const std::string& query, const DNSName& qname, const std::unordered_set<uint16_t>& optionsToIgnore)
  {
    const size_t querySize = query.size();
    const size_t cachedQuerySize = cachedQuery.size();
    if (querySize != cachedQuerySize) {
      return false;
    }

    if (!queryHeaderMatches(cachedQuery, query)) {
      return false;
    }

    size_t pos = sizeof(dnsheader) + qname.wirelength();

    /* we need at least 2 (QTYPE) + 2 (QCLASS)
       + OPT root label (1), type (2), class (2) and ttl (4)
       + the OPT RR rdlen (2)
       = 15
    */
    const dnsheader_aligned dnsheaderdata(query.data());
    const struct dnsheader* dh = dnsheaderdata.get();
    if (ntohs(dh->qdcount) != 1 || ntohs(dh->ancount) != 0 || ntohs(dh->nscount) != 0 || ntohs(dh->arcount) != 1 || (pos + 15) >= querySize || optionsToIgnore.empty()) {
      return cachedQuery.compare(pos, cachedQuerySize - pos, query, pos, querySize - pos) == 0;
    }

    /* compare up to the first option, if any */
    if (cachedQuery.compare(pos, 15, query, pos, 15) != 0) {
      return false;
    }

    /* skip the qtype (2), qclass (2) */
    /* root label (1), type (2), class (2) and ttl (4) */
    /* already compared above */
    pos += 13;

    const uint16_t rdLen = ((static_cast<unsigned char>(query.at(pos)) * 256) + static_cast<unsigned char>(query.at(pos + 1)));
    /* skip the rd length */
    /* already compared above */
    pos += sizeof(uint16_t);

    if (rdLen > (querySize - pos)) {
      /* something is wrong, let's just compare everything */
      return cachedQuery.compare(pos, cachedQuerySize - pos, query, pos, querySize - pos) == 0;
    }

    uint16_t rdataRead = 0;
    uint16_t optionCode;
    uint16_t optionLen;

    while (pos < querySize && rdataRead < rdLen && getNextEDNSOption(&query.at(pos), rdLen - rdataRead, optionCode, optionLen)) {
      if (optionLen > (rdLen - rdataRead)) {
        return cachedQuery.compare(pos, cachedQuerySize - pos, query, pos, querySize - pos) == 0;
      }

      /* compare the option code and length */
      if (cachedQuery.compare(pos, 4, query, pos, 4) != 0) {
        return false;
      }
      pos += 4;
      rdataRead += 4;

      if (optionLen > 0 && optionsToIgnore.count(optionCode) == 0) {
        if (cachedQuery.compare(pos, optionLen, query, pos, optionLen) != 0) {
          return false;
        }
      }
      pos += optionLen;
      rdataRead += optionLen;
    }

    if (pos >= querySize) {
        return true;
    }

    return cachedQuery.compare(pos, cachedQuerySize - pos, query, pos, querySize - pos) == 0;
  }

};