/* MIT License
 *
 * Copyright (c) 2023 Brad House
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice (including the next
 * paragraph) shall be included in all copies or substantial portions of the
 * Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 *
 * SPDX-License-Identifier: MIT
 */
#include "ares_private.h"
struct ares_qcache {
  ares_htable_strvp_t *cache;
  ares_slist_t        *expire;
  unsigned int         max_ttl;
};
typedef struct {
  char              *key;
  ares_dns_record_t *dnsrec;
  time_t             expire_ts;
  time_t             insert_ts;
} ares_qcache_entry_t;
static char *ares_qcache_calc_key(const ares_dns_record_t *dnsrec)
{
  ares_buf_t      *buf = ares_buf_create();
  size_t           i;
  ares_status_t    status;
  ares_dns_flags_t flags;
  if (dnsrec == NULL || buf == NULL) {
    return NULL; /* LCOV_EXCL_LINE: DefensiveCoding */
  }
  /* Format is OPCODE|FLAGS[|QTYPE1|QCLASS1|QNAME1]... */
  status = ares_buf_append_str(
    buf, ares_dns_opcode_tostr(ares_dns_record_get_opcode(dnsrec)));
  if (status != ARES_SUCCESS) {
    goto fail; /* LCOV_EXCL_LINE: OutOfMemory */
  }
  status = ares_buf_append_byte(buf, '|');
  if (status != ARES_SUCCESS) {
    goto fail; /* LCOV_EXCL_LINE: OutOfMemory */
  }
  flags = ares_dns_record_get_flags(dnsrec);
  /* Only care about RD and CD */
  if (flags & ARES_FLAG_RD) {
    status = ares_buf_append_str(buf, "rd");
    if (status != ARES_SUCCESS) {
      goto fail; /* LCOV_EXCL_LINE: OutOfMemory */
    }
  }
  if (flags & ARES_FLAG_CD) {
    status = ares_buf_append_str(buf, "cd");
    if (status != ARES_SUCCESS) {
      goto fail; /* LCOV_EXCL_LINE: OutOfMemory */
    }
  }
  for (i = 0; i < ares_dns_record_query_cnt(dnsrec); i++) {
    const char         *name;
    size_t              name_len;
    ares_dns_rec_type_t qtype;
    ares_dns_class_t    qclass;
    status = ares_dns_record_query_get(dnsrec, i, &name, &qtype, &qclass);
    if (status != ARES_SUCCESS) {
      goto fail; /* LCOV_EXCL_LINE: DefensiveCoding */
    }
    status = ares_buf_append_byte(buf, '|');
    if (status != ARES_SUCCESS) {
      goto fail; /* LCOV_EXCL_LINE: OutOfMemory */
    }
    status = ares_buf_append_str(buf, ares_dns_rec_type_tostr(qtype));
    if (status != ARES_SUCCESS) {
      goto fail; /* LCOV_EXCL_LINE: OutOfMemory */
    }
    status = ares_buf_append_byte(buf, '|');
    if (status != ARES_SUCCESS) {
      goto fail; /* LCOV_EXCL_LINE: OutOfMemory */
    }
    status = ares_buf_append_str(buf, ares_dns_class_tostr(qclass));
    if (status != ARES_SUCCESS) {
      goto fail; /* LCOV_EXCL_LINE: OutOfMemory */
    }
    status = ares_buf_append_byte(buf, '|');
    if (status != ARES_SUCCESS) {
      goto fail; /* LCOV_EXCL_LINE: OutOfMemory */
    }
    /* On queries, a '.' may be appended to the name to indicate an explicit
     * name lookup without performing a search.  Strip this since its not part
     * of a cached response. */
    name_len = ares_strlen(name);
    if (name_len && name[name_len - 1] == '.') {
      name_len--;
    }
    if (name_len > 0) {
      status = ares_buf_append(buf, (const unsigned char *)name, name_len);
      if (status != ARES_SUCCESS) {
        goto fail; /* LCOV_EXCL_LINE: OutOfMemory */
      }
    }
  }
  return ares_buf_finish_str(buf, NULL);
/* LCOV_EXCL_START: OutOfMemory */
fail:
  ares_buf_destroy(buf);
  return NULL;
  /* LCOV_EXCL_STOP */
}
static void ares_qcache_expire(ares_qcache_t *cache, const ares_timeval_t *now)
{
  ares_slist_node_t *node;
  if (cache == NULL) {
    return;
  }
  while ((node = ares_slist_node_first(cache->expire)) != NULL) {
    const ares_qcache_entry_t *entry = ares_slist_node_val(node);
    /* If now is NULL, we're flushing everything, so don't break */
    if (now != NULL && entry->expire_ts > now->sec) {
      break;
    }
    ares_htable_strvp_remove(cache->cache, entry->key);
    ares_slist_node_destroy(node);
  }
}
void ares_qcache_flush(ares_qcache_t *cache)
{
  ares_qcache_expire(cache, NULL /* flush all */);
}
void ares_qcache_destroy(ares_qcache_t *cache)
{
  if (cache == NULL) {
    return;
  }
  ares_htable_strvp_destroy(cache->cache);
  ares_slist_destroy(cache->expire);
  ares_free(cache);
}
static int ares_qcache_entry_sort_cb(const void *arg1, const void *arg2)
{
  const ares_qcache_entry_t *entry1 = arg1;
  const ares_qcache_entry_t *entry2 = arg2;
  if (entry1->expire_ts > entry2->expire_ts) {
    return 1;
  }
  if (entry1->expire_ts < entry2->expire_ts) {
    return -1;
  }
  return 0;
}
static void ares_qcache_entry_destroy_cb(void *arg)
{
  ares_qcache_entry_t *entry = arg;
  if (entry == NULL) {
    return; /* LCOV_EXCL_LINE: DefensiveCoding */
  }
  ares_free(entry->key);
  ares_dns_record_destroy(entry->dnsrec);
  ares_free(entry);
}
ares_status_t ares_qcache_create(ares_rand_state *rand_state,
                                 unsigned int     max_ttl,
                                 ares_qcache_t  **cache_out)
{
  ares_status_t  status = ARES_SUCCESS;
  ares_qcache_t *cache;
  cache = ares_malloc_zero(sizeof(*cache));
  if (cache == NULL) {
    status = ARES_ENOMEM; /* LCOV_EXCL_LINE: OutOfMemory */
    goto done;            /* LCOV_EXCL_LINE: OutOfMemory */
  }
  cache->cache = ares_htable_strvp_create(NULL);
  if (cache->cache == NULL) {
    status = ARES_ENOMEM; /* LCOV_EXCL_LINE: OutOfMemory */
    goto done;            /* LCOV_EXCL_LINE: OutOfMemory */
  }
  cache->expire = ares_slist_create(rand_state, ares_qcache_entry_sort_cb,
                                    ares_qcache_entry_destroy_cb);
  if (cache->expire == NULL) {
    status = ARES_ENOMEM; /* LCOV_EXCL_LINE: OutOfMemory */
    goto done;            /* LCOV_EXCL_LINE: OutOfMemory */
  }
  cache->max_ttl = max_ttl;
done:
  if (status != ARES_SUCCESS) {
    *cache_out = NULL;
    ares_qcache_destroy(cache);
    return status;
  }
  *cache_out = cache;
  return status;
}
static unsigned int ares_qcache_calc_minttl(ares_dns_record_t *dnsrec)
{
  unsigned int minttl = 0xFFFFFFFF;
  size_t       sect;
  for (sect = ARES_SECTION_ANSWER; sect <= ARES_SECTION_ADDITIONAL; sect++) {
    size_t i;
    for (i = 0; i < ares_dns_record_rr_cnt(dnsrec, (ares_dns_section_t)sect);
         i++) {
      const ares_dns_rr_t *rr =
        ares_dns_record_rr_get(dnsrec, (ares_dns_section_t)sect, i);
      ares_dns_rec_type_t type = ares_dns_rr_get_type(rr);
      unsigned int        ttl  = ares_dns_rr_get_ttl(rr);
      /* TTL is meaningless on these record types */
      if (type == ARES_REC_TYPE_OPT || type == ARES_REC_TYPE_SOA ||
          type == ARES_REC_TYPE_SIG) {
        continue;
      }
      if (ttl < minttl) {
        minttl = ttl;
      }
    }
  }
  return minttl;
}
static unsigned int ares_qcache_soa_minimum(ares_dns_record_t *dnsrec)
{
  size_t i;
  /* RFC 2308 Section 5 says its the minimum of MINIMUM and the TTL of the
   * record. */
  for (i = 0; i < ares_dns_record_rr_cnt(dnsrec, ARES_SECTION_AUTHORITY); i++) {
    const ares_dns_rr_t *rr =
      ares_dns_record_rr_get(dnsrec, ARES_SECTION_AUTHORITY, i);
    ares_dns_rec_type_t type = ares_dns_rr_get_type(rr);
    unsigned int        ttl;
    unsigned int        minimum;
    if (type != ARES_REC_TYPE_SOA) {
      continue;
    }
    minimum = ares_dns_rr_get_u32(rr, ARES_RR_SOA_MINIMUM);
    ttl     = ares_dns_rr_get_ttl(rr);
    if (ttl > minimum) {
      return minimum;
    }
    return ttl;
  }
  return 0;
}
/* On success, takes ownership of dnsrec */
static ares_status_t ares_qcache_insert_int(ares_qcache_t           *qcache,
                                            ares_dns_record_t       *qresp,
                                            const ares_dns_record_t *qreq,
                                            const ares_timeval_t    *now)
{
  ares_qcache_entry_t *entry;
  unsigned int         ttl;
  ares_dns_rcode_t     rcode = ares_dns_record_get_rcode(qresp);
  ares_dns_flags_t     flags = ares_dns_record_get_flags(qresp);
  if (qcache == NULL || qresp == NULL) {
    return ARES_EFORMERR;
  }
  /* Only save NOERROR or NXDOMAIN */
  if (rcode != ARES_RCODE_NOERROR && rcode != ARES_RCODE_NXDOMAIN) {
    return ARES_ENOTIMP;
  }
  /* Don't save truncated queries */
  if (flags & ARES_FLAG_TC) {
    return ARES_ENOTIMP;
  }
  /* Look at SOA for NXDOMAIN for minimum */
  if (rcode == ARES_RCODE_NXDOMAIN) {
    ttl = ares_qcache_soa_minimum(qresp);
  } else {
    ttl = ares_qcache_calc_minttl(qresp);
  }
  if (ttl > qcache->max_ttl) {
    ttl = qcache->max_ttl;
  }
  /* Don't cache something that is already expired */
  if (ttl == 0) {
    return ARES_EREFUSED;
  }
  entry = ares_malloc_zero(sizeof(*entry));
  if (entry == NULL) {
    goto fail; /* LCOV_EXCL_LINE: OutOfMemory */
  }
  entry->dnsrec    = qresp;
  entry->expire_ts = (time_t)now->sec + (time_t)ttl;
  entry->insert_ts = (time_t)now->sec;
  /* We can't guarantee the server responded with the same flags as the
   * request had, so we have to re-parse the request in order to generate the
   * key for caching, but we'll only do this once we know for sure we really
   * want to cache it */
  entry->key = ares_qcache_calc_key(qreq);
  if (entry->key == NULL) {
    goto fail; /* LCOV_EXCL_LINE: OutOfMemory */
  }
  if (!ares_htable_strvp_insert(qcache->cache, entry->key, entry)) {
    goto fail; /* LCOV_EXCL_LINE: OutOfMemory */
  }
  if (ares_slist_insert(qcache->expire, entry) == NULL) {
    goto fail; /* LCOV_EXCL_LINE: OutOfMemory */
  }
  return ARES_SUCCESS;
/* LCOV_EXCL_START: OutOfMemory */
fail:
  if (entry != NULL && entry->key != NULL) {
    ares_htable_strvp_remove(qcache->cache, entry->key);
    ares_free(entry->key);
    ares_free(entry);
  }
  return ARES_ENOMEM;
  /* LCOV_EXCL_STOP */
}
ares_status_t ares_qcache_fetch(ares_channel_t           *channel,
                                const ares_timeval_t     *now,
                                const ares_dns_record_t  *dnsrec,
                                const ares_dns_record_t **dnsrec_resp)
{
  char                *key = NULL;
  ares_qcache_entry_t *entry;
  ares_status_t        status = ARES_SUCCESS;
  if (channel == NULL || dnsrec == NULL || dnsrec_resp == NULL) {
    return ARES_EFORMERR;
  }
  if (channel->qcache == NULL) {
    return ARES_ENOTFOUND;
  }
  ares_qcache_expire(channel->qcache, now);
  key = ares_qcache_calc_key(dnsrec);
  if (key == NULL) {
    status = ARES_ENOMEM; /* LCOV_EXCL_LINE: OutOfMemory */
    goto done;            /* LCOV_EXCL_LINE: OutOfMemory */
  }
  entry = ares_htable_strvp_get_direct(channel->qcache->cache, key);
  if (entry == NULL) {
    status = ARES_ENOTFOUND;
    goto done;
  }
  ares_dns_record_ttl_decrement(entry->dnsrec,
                                (unsigned int)(now->sec - entry->insert_ts));
  *dnsrec_resp = entry->dnsrec;
done:
  ares_free(key);
  return status;
}
ares_status_t ares_qcache_insert(ares_channel_t       *channel,
                                 const ares_timeval_t *now,
                                 const ares_query_t   *query,
                                 ares_dns_record_t    *dnsrec)
{
  return ares_qcache_insert_int(channel->qcache, dnsrec, query->query, now);
}