/* Copyright Joyent, Inc. and other Node contributors. All rights reserved.

 *

 * 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 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.

 */

#include "uv.h"

#include "task.h"

#include <stdio.h>

#include <stdlib.h>

#include <string.h>

typedef struct {

  uv_write_t req;

  uv_buf_t buf;

} write_req_t;

/* used to track multiple DNS requests received */

typedef struct {

  char* prevbuf_ptr;

  int prevbuf_pos;

  int prevbuf_rem;

} dnsstate;

/* modify handle to append dnsstate */

typedef struct {

  uv_tcp_t handle;

  dnsstate state;

} dnshandle;

static uv_loop_t* loop;

static uv_tcp_t server;

static void after_write(uv_write_t* req, int status);

static void after_read(uv_stream_t*, ssize_t nread, const uv_buf_t* buf);

static void on_close(uv_handle_t* peer);

static void on_connection(uv_stream_t*, int status);

#define WRITE_BUF_LEN   (64*1024)

#define DNSREC_LEN      (4)

#define LEN_OFFSET 0

#define QUERYID_OFFSET 2

static unsigned char DNSRsp[] = {

  0, 43, 0, 0, 0x81, 0x80, 0, 1, 0, 1, 0, 0, 0, 0

};

static unsigned char qrecord[] = {

  5, 'e', 'c', 'h', 'o', 's', 3, 's', 'r', 'v', 0, 0, 1, 0, 1

};

static unsigned char arecord[] = {

  0xc0, 0x0c, 0, 1, 0, 1, 0, 0, 5, 0xbd, 0, 4, 10, 0, 1, 1

};

static void after_write(uv_write_t* req, int status) {

  write_req_t* wr;

  if (status) {

    fprintf(stderr, "uv_write error: %s\n", uv_strerror(status));

    ASSERT(0);

  }

  wr = (write_req_t*) req;

  /* Free the read/write buffer and the request */

  free(wr->buf.base);

  free(wr);

}

static void after_shutdown(uv_shutdown_t* req, int status) {

  uv_close((uv_handle_t*) req->handle, on_close);

  free(req);

}

static void addrsp(write_req_t* wr, char* hdr) {

  char * dnsrsp;

  short int rsplen;

  short int* reclen;

  rsplen = sizeof(DNSRsp) + sizeof(qrecord) + sizeof(arecord);

  ASSERT (rsplen + wr->buf.len < WRITE_BUF_LEN);

  dnsrsp = wr->buf.base + wr->buf.len;

  /* copy stock response */

  memcpy(dnsrsp, DNSRsp, sizeof(DNSRsp));

  memcpy(dnsrsp + sizeof(DNSRsp), qrecord, sizeof(qrecord));

  memcpy(dnsrsp + sizeof(DNSRsp) + sizeof(qrecord), arecord, sizeof(arecord));

  /* overwrite with network order length and id from request header */

  reclen = (short int*)dnsrsp;

  *reclen = htons(rsplen-2);

  dnsrsp[QUERYID_OFFSET] = hdr[QUERYID_OFFSET];

  dnsrsp[QUERYID_OFFSET+1] = hdr[QUERYID_OFFSET+1];

  wr->buf.len += rsplen;

}

static void process_req(uv_stream_t* handle,

                        ssize_t nread,

                        const uv_buf_t* buf) {

  write_req_t* wr;

  dnshandle* dns = (dnshandle*)handle;

  char hdrbuf[DNSREC_LEN];

  int hdrbuf_remaining = DNSREC_LEN;

  int rec_remaining = 0;

  int readbuf_remaining;

  char* dnsreq;

  char* hdrstart;

  int usingprev = 0;

  wr = (write_req_t*) malloc(sizeof *wr);

  wr->buf.base = (char*)malloc(WRITE_BUF_LEN);

  wr->buf.len = 0;

  if (dns->state.prevbuf_ptr != NULL) {

    dnsreq = dns->state.prevbuf_ptr + dns->state.prevbuf_pos;

    readbuf_remaining = dns->state.prevbuf_rem;

    usingprev = 1;

  } else {

    dnsreq = buf->base;

    readbuf_remaining = nread;

  }

  hdrstart = dnsreq;

  while (dnsreq != NULL) {

    /* something to process */

    while (readbuf_remaining > 0) {

      /* something to process in current buffer */

      if (hdrbuf_remaining > 0) {

        /* process len and id */

        if (readbuf_remaining < hdrbuf_remaining) {

          /* too little to get request header. save for next buffer */

          memcpy(&hdrbuf[DNSREC_LEN - hdrbuf_remaining],

                 dnsreq,

                 readbuf_remaining);

          hdrbuf_remaining = DNSREC_LEN - readbuf_remaining;

          break;

        } else {

          /* save header */

          memcpy(&hdrbuf[DNSREC_LEN - hdrbuf_remaining],

                 dnsreq,

                 hdrbuf_remaining);

          dnsreq += hdrbuf_remaining;

          readbuf_remaining -= hdrbuf_remaining;

          hdrbuf_remaining = 0;

          /* get record length */

          rec_remaining = (unsigned) hdrbuf[0] * 256 + (unsigned) hdrbuf[1];

          rec_remaining -= (DNSREC_LEN - 2);

        }

      }

      if (rec_remaining <= readbuf_remaining) {

        /* prepare reply */

        addrsp(wr, hdrbuf);

        /* move to next record */

        dnsreq += rec_remaining;

        hdrstart = dnsreq;

        readbuf_remaining -= rec_remaining;

        rec_remaining = 0;

        hdrbuf_remaining = DNSREC_LEN;

      } else {

        /* otherwise this buffer is done. */

        rec_remaining -= readbuf_remaining;

        break;

      }

    }

    /* If we had to use bytes from prev buffer, start processing the current

     * one.

     */

    if (usingprev == 1) {

      /* free previous buffer */

      free(dns->state.prevbuf_ptr);

      dnsreq = buf->base;

      readbuf_remaining = nread;

      usingprev = 0;

    } else {

      dnsreq = NULL;

    }

  }

  /* send write buffer */

  if (wr->buf.len > 0) {

    if (uv_write((uv_write_t*) &wr->req, handle, &wr->buf, 1, after_write)) {

      FATAL("uv_write failed");

    }

  }

  if (readbuf_remaining > 0) {

    /* save start of record position, so we can continue on next read */

    dns->state.prevbuf_ptr = buf->base;

    dns->state.prevbuf_pos = hdrstart - buf->base;

    dns->state.prevbuf_rem = nread - dns->state.prevbuf_pos;

  } else {

    /* nothing left in this buffer */

    dns->state.prevbuf_ptr = NULL;

    dns->state.prevbuf_pos = 0;

    dns->state.prevbuf_rem = 0;

    free(buf->base);

  }

}

static void after_read(uv_stream_t* handle,

                       ssize_t nread,

                       const uv_buf_t* buf) {

  uv_shutdown_t* req;

  if (nread < 0) {

    /* Error or EOF */

    ASSERT(nread == UV_EOF);

    if (buf->base) {

      free(buf->base);

    }

    req = malloc(sizeof *req);

    uv_shutdown(req, handle, after_shutdown);

    return;

  }

  if (nread == 0) {

    /* Everything OK, but nothing read. */

    free(buf->base);

    return;

  }

  /* process requests and send responses */

  process_req(handle, nread, buf);

}

static void on_close(uv_handle_t* peer) {

  free(peer);

}

static void buf_alloc(uv_handle_t* handle,

                      size_t suggested_size,

                      uv_buf_t* buf) {

  buf->base = malloc(suggested_size);

  buf->len = suggested_size;

}

static void on_connection(uv_stream_t* server, int status) {

  dnshandle* handle;

  int r;

  ASSERT(status == 0);

  handle = (dnshandle*) malloc(sizeof *handle);

  ASSERT(handle != NULL);

  /* initialize read buffer state */

  handle->state.prevbuf_ptr = 0;

  handle->state.prevbuf_pos = 0;

  handle->state.prevbuf_rem = 0;

  r = uv_tcp_init(loop, (uv_tcp_t*)handle);

  ASSERT(r == 0);

  r = uv_accept(server, (uv_stream_t*)handle);

  ASSERT(r == 0);

  r = uv_read_start((uv_stream_t*)handle, buf_alloc, after_read);

  ASSERT(r == 0);

}

static int dns_start(int port) {

  struct sockaddr_in addr;

  int r;

  ASSERT(0 == uv_ip4_addr("0.0.0.0", port, &addr));

  r = uv_tcp_init(loop, &server);

  if (r) {

    /* TODO: Error codes */

    fprintf(stderr, "Socket creation error\n");

    return 1;

  }

  r = uv_tcp_bind(&server, (const struct sockaddr*) &addr, 0);

  if (r) {

    /* TODO: Error codes */

    fprintf(stderr, "Bind error\n");

    return 1;

  }

  r = uv_listen((uv_stream_t*)&server, 128, on_connection);

  if (r) {

    /* TODO: Error codes */

    fprintf(stderr, "Listen error\n");

    return 1;

  }

  return 0;

}

HELPER_IMPL(dns_server) {

  loop = uv_default_loop();

  if (dns_start(TEST_PORT_2))

    return 1;

  uv_run(loop, UV_RUN_DEFAULT);

  return 0;

}