/*

 * Copyright (C) Internet Systems Consortium, Inc. ("ISC")

 *

 * SPDX-License-Identifier: MPL-2.0

 *

 * This Source Code Form is subject to the terms of the Mozilla Public

 * License, v. 2.0. If a copy of the MPL was not distributed with this

 * file, you can obtain one at https://mozilla.org/MPL/2.0/.

 *

 * See the COPYRIGHT file distributed with this work for additional

 * information regarding copyright ownership.

 */

#include <stdbool.h>

#include <stdint.h>

#include <string.h>

#include <isc/ascii.h>

#include <isc/buffer.h>

#include <isc/hash.h>

#include <isc/mem.h>

#include <isc/util.h>

#include <dns/compress.h>

#include <dns/name.h>

#define HASH_INIT_DJB2 5381

#define CCTX_MAGIC    ISC_MAGIC('C', 'C', 'T', 'X')

#define CCTX_VALID(x) ISC_MAGIC_VALID(x, CCTX_MAGIC)

void

dns_compress_init(dns_compress_t *cctx, isc_mem_t *mctx,

      dns_compress_flags_t flags) {

  dns_compress_slot_t *set = NULL;

  uint16_t mask;

  REQUIRE(cctx != NULL);

  REQUIRE(mctx != NULL);

  if ((flags & DNS_COMPRESS_LARGE) != 0) {

    size_t count = (1 << DNS_COMPRESS_LARGEBITS);

    mask = count - 1;

    set = isc_mem_callocate(mctx, count, sizeof(*set));

  } else {

    mask = ARRAY_SIZE(cctx->smallset) - 1;

    set = cctx->smallset;

  }

  /*

   * The lifetime of this object is limited to the stack frame of the

   * caller, so we don't need to attach to the memory context.

   */

  *cctx = (dns_compress_t){

    .magic = CCTX_MAGIC,

    .flags = flags | DNS_COMPRESS_PERMITTED,

    .mctx = mctx,

    .mask = mask,

    .set = set,

    .coff = 0xffff,

  };

}

void

dns_compress_invalidate(dns_compress_t *cctx) {

  REQUIRE(CCTX_VALID(cctx));

  if (cctx->set != cctx->smallset) {

    isc_mem_free(cctx->mctx, cctx->set);

  }

  *cctx = (dns_compress_t){ 0 };

}

void

dns_compress_setmultiuse(dns_compress_t *cctx, bool multi) {

  REQUIRE(CCTX_VALID(cctx));

  if (multi) {

    cctx->flags |= DNS_COMPRESS_MULTIUSE;

  } else {

    cctx->flags &= ~DNS_COMPRESS_MULTIUSE;

  }

  cctx->coff = 0xffff;

}

bool

dns_compress_getmultiuse(dns_compress_t *cctx) {

  REQUIRE(CCTX_VALID(cctx));

  return (cctx->flags & DNS_COMPRESS_MULTIUSE) != 0;

}

void

dns_compress_setpermitted(dns_compress_t *cctx, bool permitted) {

  REQUIRE(CCTX_VALID(cctx));

  if (permitted) {

    cctx->flags |= DNS_COMPRESS_PERMITTED;

  } else {

    cctx->flags &= ~DNS_COMPRESS_PERMITTED;

  }

  dns_compress_setmultiuse(cctx, false);

}

bool

dns_compress_getpermitted(dns_compress_t *cctx) {

  REQUIRE(CCTX_VALID(cctx));

  return (cctx->flags & DNS_COMPRESS_PERMITTED) != 0;

}

/*

 * Our hash value needs to cover the entire suffix of a name, and we need

 * to calculate it one label at a time. So this function mixes a label into

 * an existing hash. (We don't use isc_hash32() because the djb2 hash is a

 * lot faster, and we limit the impact of collision attacks by restricting

 * the size and occupancy of the hash set.) The accumulator is 32 bits to

 * keep more of the fun mixing that happens in the upper bits.

 */

static uint16_t

hash_label(uint16_t init, uint8_t *ptr, bool sensitive) {

  unsigned int len = ptr[0] + 1;

  uint32_t hash = init;

  if (sensitive) {

    while (len-- > 0) {

      hash = hash * 33 + *ptr++;

    }

  } else {

    /* using the autovectorize-friendly tolower() */

    while (len-- > 0) {

      hash = hash * 33 + isc__ascii_tolower1(*ptr++);

    }

  }

  return isc_hash_bits32(hash, 16);

}

static bool

match_wirename(uint8_t *a, uint8_t *b, unsigned int len, bool sensitive) {

  if (sensitive) {

    return memcmp(a, b, len) == 0;

  } else {

    /* label lengths are < 'A' so unaffected by tolower() */

    return isc_ascii_lowerequal(a, b, len);

  }

}

/*

 * We have found a hash set entry whose hash value matches the current

 * suffix of our name, which is passed to this function via `sptr` and

 * `slen`. We need to verify that the suffix in the message (referred to

 * by `new_coff`) actually matches, in case of hash collisions.

 *

 * We know that the previous suffix of this name (after the first label)

 * occurs in the message at `old_coff`, and all the compression offsets in

 * the hash set and in the message refer to the first occurrence of a

 * particular name or suffix.

 *

 * First, we need to match the label that was just added to our suffix,

 * and second, verify that it is followed by the previous suffix.

 *

 * There are a few ways to match the previous suffix:

 *

 * When the first occurrence of this suffix is also the first occurrence

 * of the previous suffix, `old_coff` points just after the new label.

 *

 * Otherwise, if this suffix occurs in a compressed name, it will be

 * followed by a compression pointer that refers to the previous suffix,

 * which must be equal to `old_coff`.

 *

 * The final possibility is that this suffix occurs in an uncompressed

 * name, so we have to compare the rest of the suffix in full.

 *

 * A special case is when this suffix is a TLD. That can be handled by

 * the case for uncompressed names, but it is common enough that it is

 * worth taking a short cut. (In the TLD case, the `old_coff` will be

 * zero, and the quick checks for the previous suffix will fail.)

 */

static bool

match_suffix(isc_buffer_t *buffer, unsigned int new_coff, uint8_t *sptr,

       unsigned int slen, unsigned int old_coff, bool sensitive) {

  uint8_t pptr[] = { 0xC0 | (old_coff >> 8), old_coff & 0xff };

  uint8_t *bptr = isc_buffer_base(buffer);

  unsigned int blen = isc_buffer_usedlength(buffer);

  unsigned int llen = sptr[0] + 1;

  INSIST(llen <= 64 && llen < slen);

  if (blen < new_coff + llen) {

    return false;

  }

  blen -= new_coff;

  bptr += new_coff;

  /* does the first label of the suffix appear here? */

  if (!match_wirename(bptr, sptr, llen, sensitive)) {

    return false;

  }

  /* is this label followed by the previously matched suffix? */

  if (old_coff == new_coff + llen) {

    return true;

  }

  blen -= llen;

  bptr += llen;

  slen -= llen;

  sptr += llen;

  /* are both labels followed by the root label? */

  if (blen >= 1 && slen == 1 && bptr[0] == 0 && sptr[0] == 0) {

    return true;

  }

  /* is this label followed by a pointer to the previous match? */

  if (blen >= 2 && bptr[0] == pptr[0] && bptr[1] == pptr[1]) {

    return true;

  }

  /* is this label followed by a copy of the rest of the suffix? */

  return blen >= slen && match_wirename(bptr, sptr, slen, sensitive);

}

/*

 * Robin Hood hashing aims to minimize probe distance when inserting a

 * new element by ensuring that the new element does not have a worse

 * probe distance than any other element in its probe sequence. During

 * insertion, if an existing element is encountered with a shorter

 * probe distance, it is swapped with the new element, and insertion

 * continues with the displaced element.

 */

static unsigned int

probe_distance(dns_compress_t *cctx, unsigned int slot) {

  return (slot - cctx->set[slot].hash) & cctx->mask;

}

static unsigned int

slot_index(dns_compress_t *cctx, unsigned int hash, unsigned int probe) {

  return (hash + probe) & cctx->mask;

}

static bool

insert_label(dns_compress_t *cctx, isc_buffer_t *buffer,

       const dns_offsets_t offsets, unsigned int label, uint16_t hash,

       unsigned int probe) {

  /*

   * hash set entries must have valid compression offsets

   * and the hash set must not get too full (75% load)

   */

  unsigned int prefix_len = offsets[label];

  unsigned int coff = isc_buffer_usedlength(buffer) + prefix_len;

  if (coff >= 0x4000 || cctx->count > cctx->mask * 3 / 4) {

    return false;

  }

  for (;;) {

    unsigned int slot = slot_index(cctx, hash, probe);

    /* we can stop when we find an empty slot */

    if (cctx->set[slot].coff == 0) {

      cctx->set[slot].hash = hash;

      cctx->set[slot].coff = coff;

      cctx->count++;

      return true;

    }

    /* he steals from the rich and gives to the poor */

    if (probe > probe_distance(cctx, slot)) {

      probe = probe_distance(cctx, slot);

      ISC_SWAP(cctx->set[slot].hash, hash);

      ISC_SWAP(cctx->set[slot].coff, coff);

    }

    probe++;

  }

}

/*

 * Add the unmatched prefix of the name to the hash set.

 */

static void

insert(dns_compress_t *cctx, isc_buffer_t *buffer, const dns_name_t *name,

       const dns_offsets_t offsets, unsigned int label, uint16_t hash,

       unsigned int probe) {

  bool sensitive = (cctx->flags & DNS_COMPRESS_CASE) != 0;

  /*

   * this insertion loop continues from the search loop inside

   * dns_compress_name() below, iterating over the remaining labels

   * of the name and accumulating the hash in the same manner

   */

  while (insert_label(cctx, buffer, offsets, label, hash, probe) &&

         label-- > 0)

  {

    unsigned int prefix_len = offsets[label];

    uint8_t *suffix_ptr = name->ndata + prefix_len;

    hash = hash_label(hash, suffix_ptr, sensitive);

    probe = 0;

  }

}

void

dns_compress_name(dns_compress_t *cctx, isc_buffer_t *buffer,

      const dns_name_t *name, unsigned int *return_prefix,

      unsigned int *return_coff) {

  REQUIRE(CCTX_VALID(cctx));

  REQUIRE(ISC_BUFFER_VALID(buffer));

  REQUIRE(dns_name_isabsolute(name));

  REQUIRE(return_prefix != NULL);

  REQUIRE(return_coff != NULL);

  REQUIRE(*return_coff == 0);

  if ((cctx->flags & DNS_COMPRESS_DISABLED) != 0) {

    return;

  }

  dns_offsets_t offsets;

  size_t labels = dns_name_offsets(name, offsets);

  INSIST(labels > 0);

  bool sensitive = (cctx->flags & DNS_COMPRESS_CASE) != 0;

  uint16_t hash = HASH_INIT_DJB2;

  size_t label = labels - 1; /* skip the root label */

  /*

   * find out how much of the name's suffix is in the hash set,

   * stepping backwards from the end one label at a time

   */

  while (label-- > 0) {

    unsigned int prefix_len = offsets[label];

    unsigned int suffix_len = name->length - prefix_len;

    uint8_t *suffix_ptr = name->ndata + prefix_len;

    hash = hash_label(hash, suffix_ptr, sensitive);

    for (unsigned int probe = 0; true; probe++) {

      unsigned int slot = slot_index(cctx, hash, probe);

      unsigned int coff = cctx->set[slot].coff;

      /*

       * if we would have inserted this entry here (as in

       * insert_label() above), our suffix cannot be in the

       * hash set, so stop searching and switch to inserting

       * the rest of the name (its prefix) into the set

       */

      if (coff == 0 || probe > probe_distance(cctx, slot)) {

        insert(cctx, buffer, name, offsets, label, hash,

               probe);

        return;

      }

      /*

       * this slot matches, so provisionally set the

       * return values and continue with the next label

       */

      if (hash == cctx->set[slot].hash &&

          match_suffix(buffer, coff, suffix_ptr, suffix_len,

           *return_coff, sensitive))

      {

        *return_coff = coff;

        *return_prefix = prefix_len;

        break;

      }

    }

  }

}

void

dns_compress_rollback(dns_compress_t *cctx, unsigned int coff) {

  REQUIRE(CCTX_VALID(cctx));

  for (unsigned int slot = 0; slot <= cctx->mask; slot++) {

    if (cctx->set[slot].coff < coff) {

      continue;

    }

    /*

     * The next few elements might be part of the deleted element's

     * probe sequence, so we slide them down to overwrite the entry

     * we are deleting and preserve the probe sequence. Moving an

     * element to the previous slot reduces its probe distance, so

     * we stop when we find an element whose probe distance is zero.

     */

    unsigned int prev = slot;

    unsigned int next = slot_index(cctx, prev, 1);

    while (cctx->set[next].coff != 0 &&

           probe_distance(cctx, next) != 0)

    {

      cctx->set[prev] = cctx->set[next];

      prev = next;

      next = slot_index(cctx, prev, 1);

    }

    cctx->set[prev].coff = 0;

    cctx->set[prev].hash = 0;

    cctx->count--;

  }

}