/*
 * Copyright (C) 2015 by the University of Southern California
 * $Id: server_proxy.cc 24241 2015-05-29 22:56:14Z liangzhu $
 *
 * This program is free software; you can redistribute it and/or
 * modify it under the terms of the GNU General Public License,
 * version 2, as published by the Free Software Foundation.
 *
 * 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.,
 * 59 Temple Place, Suite 330, Boston, MA 02111-1307, USA. 
 * 
 */

#include "server_proxy.hh"
#include <iostream>
#include <string.h>
#include <signal.h>
#include <assert.h>
#include <sys/types.h>
#include <unistd.h>

#include <linux/tcp.h> //fast open

#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>

//#include <openssl/ssl.h>
#include <openssl/err.h>
#include <openssl/rand.h>

#include <event2/listener.h>
#include <event2/bufferevent.h>
#include <event2/buffer.h>
#include <event2/bufferevent_ssl.h>

#include "global_var.h"
#include "socket_util.h"

#include <ldns/ldns.h>

using namespace std;

#define MAX_BUF_SIZE 	4096
#define MIN_BUF_SIZE 	64

#define DNS_MESSAGEEXTFLAG_DO  0x8000U
#define DNS_MESSAGEEXTFLAG_TO  0x4000U

//#define USE_PTHREAD 1

#ifdef USE_PTHREAD
#include <event2/thread.h>
#endif

void enq_msg(client_t *, uint8_t *, size_t);
void ssl_read_cb(struct bufferevent *, void *);

/*
  signal call back; exit event loop; clean up is done when the caller returns
*/
static void
signal_cb(evutil_socket_t sig, short what, void *arg)
{
  assert(what & EV_SIGNAL);
  struct event_base *base = (struct event_base *)arg;
  struct timeval delay = {2, 0};
  cerr << "[" << getpid() << "] Interrupt signal ("
       << sig << ") received, exit in two seconds!\n";
  event_base_loopexit(base, &delay);
}

/*
  clean up one client
*/
static void
cleanup_client(client_t *clt)
{
  //client socket is cleaned up by libevent since "BEV_OPT_CLOSE_ON_FREE"
  DBG("start\n");
  assert(clt != NULL);
  safe_map<uint16_t, uint16_t> *id_map = clt->id_map;
  uint16_t new_id = 0;
  if (clt->type == CONN_TLS) {
    DBG("clean up ssl\n");
    //clean up ssl: do lazy SSL shutdowns manually
    //http://www.wangafu.net/~nickm/libevent-book/Ref6a_advanced_bufferevents.html
    SSL_set_shutdown(clt->ctx, SSL_RECEIVED_SHUTDOWN);
    SSL_shutdown(clt->ctx);
  }
  while(!id_map->empty()){
    new_id = id_map->pop_first();
    if (new_id == 0)
      break;
    clt->map->erase(new_id);
    clt->id_q->push(new_id);
  }
  delete clt;
  DBG("end\n");
}

/*
  buffer event callback (non-read or non-write)
*/
static void
buf_ev_cb(struct bufferevent *bev, short events, void *ctx)
{
  DBG("start\n");
#ifdef DEBUG  
  evutil_socket_t fd = bufferevent_getfd(bev);
#endif
  client_t *clt = (client_t *)ctx;
  if (events & events & BEV_EVENT_CONNECTED) {
    DBG("fd[%d] connected\n", (int)fd);
  } else if (events & (BEV_EVENT_EOF|BEV_EVENT_ERROR|BEV_EVENT_TIMEOUT)) {
    if (events & BEV_EVENT_ERROR) {
      cerr << "error!!!!["<< EVUTIL_SOCKET_ERROR() <<"]"<<endl;
    }
    DBG("fd[%d] %s%s%s\n",
	(int)fd,
	(events & BEV_EVENT_TIMEOUT) ? "timeout":"",
	(events & BEV_EVENT_EOF) ? "got a close":"",
	(events & BEV_EVENT_ERROR) ? evutil_socket_error_to_string(EVUTIL_SOCKET_ERROR()):"");
    cleanup_client(clt);
    bufferevent_free(bev);
  }
  DBG("end\n");
}

/*
  print dns packet
*/
void print_dns_pkt(uint8_t *buf, int buf_sz)
{
  ldns_status status;
  ldns_pkt* pkt = NULL;
  //parse the reply packet
  status = ldns_wire2pkt(&pkt, buf, (size_t)buf_sz);
  if(status != LDNS_STATUS_OK) {
    ldns_get_errorstr_by_id(status);
    errx(1, "cannot parse packet when printing");
  }
  assert(pkt != NULL);
  //print the pkt
  ldns_pkt_print(stdout, pkt);
  ldns_pkt_free(pkt);
}

/*
  is TO bit set?
*/
bool has_to_bit(uint8_t *p, size_t l)
{
  ldns_status status;
  ldns_pkt* pkt = NULL;
  uint8_t *msg = new uint8_t[l];
  memcpy(msg, p, l);
  status = ldns_wire2pkt(&pkt, msg, l);
  if(status != LDNS_STATUS_OK) {
    ldns_get_errorstr_by_id(status);
    DBG("msg len=%d\n", (int)l);
    errx(1, "cannot parse the packet");
  }
  assert(pkt != NULL);
  //if(!ldns_pkt_edns(pkt) || !ldns_pkt_edns_data(pkt)) {
  if(!ldns_pkt_edns(pkt)) {
    DBG("NO edns in packet\n");
    ldns_pkt_free(pkt);
    return false; //no edns
  }

  bool found = false;
  ldns_rdf* qname_rdf = NULL;
  char* qname_str = NULL;
  ldns_rr *q_rr = NULL;
  //find STARTTLS query name
  for (int i = 0; i < ldns_pkt_qdcount(pkt); i++) {
    q_rr = ldns_rr_list_rr(ldns_pkt_question(pkt), i);
    qname_rdf = ldns_rr_owner(q_rr);
    qname_str = ldns_rdf2str(qname_rdf);
    DBG("qname=[%s]\n", qname_str);
    
    if (strcmp(qname_str, "STARTTLS") == 0 ||
	strcmp(qname_str, "STARTTLS.") == 0)
      found = true;

    //clean up
    if (qname_str) {
      LDNS_FREE(qname_str);
      qname_str = NULL;
    }
    /*  if (qname_rdf) {
      ldns_rdf_free(qname_rdf);
      qname_rdf = NULL;
    }
    if (q_rr) {
      ldns_rr_free(q_rr);
      q_rr = NULL;
    }*/
    if (found)
      break;
  }
  uint16_t zbit = ldns_pkt_edns_z(pkt); //get the z bit
  ldns_pkt_free(pkt);
  delete[] msg;
  return found && (zbit & DNS_MESSAGEEXTFLAG_TO);
}

/*
  create a socket to the real recursive server
*/
int get_svr_fd(const char *ip, int port)
{
  int fd = mk_conn (ip, port, false);
  if (fd < 0)
    errx(1, "[error] creating server socket fails, abort!");
  return fd;
}

/*
  get id of a DNS packet, either udp or tcp
*/
uint16_t get_dns_id(uint8_t* msg, int msg_len, bool tcp)
{
  uint16_t len_host = 0;
  uint16_t offset = tcp ? 2 : 0;
  memcpy(&len_host, msg + offset, 2);
  return ntohs(len_host);
}

/*
  write a dns id to a packet, either udp or tcp
*/
void put_dns_id(uint8_t* msg, uint16_t id, bool tcp)
{
  uint16_t id_net = htons(id);
  uint16_t offset = tcp ? 2 : 0;
  memcpy(msg + offset, &id_net, sizeof(id_net));
}

/*
  clean up a pointer
*/
void clean_pt(uint8_t *p, bool array)
{
  if (p != NULL) {
    if (array)
      delete[] p;
    else
      delete p;
  }
}

/*
  constructor
*/
server_proxy::server_proxy(const char *sip, int sp,
			   const char *lip, int lp, int t,
			   const char *key, const char *cert)
{
  svr_ip = sip;
  svr_port = sp;
  lis_ip = lip;
  lis_port = lp;
  conn_type = t;
  key_file = key;
  cert_file = cert;
  mypid = getpid();
  check_input();
  id_q = new safe_queue<int>(MAX_QSIZE);
  clt_map = new safe_map<int, client_t *>;
  for (int i=1; i<= MAX_QSIZE; i++)
    id_q->push(i);
  svr_fd = get_svr_fd(svr_ip, svr_port);

  //set the proxy address
  memset(&tcp_addr, 0, sizeof(struct sockaddr_in));
  tcp_addr.sin_family = AF_INET;
  if (strcmp(lis_ip, "ANY") == 0 ||
      strcmp(lis_ip, "any") == 0)
    tcp_addr.sin_addr.s_addr = htonl(INADDR_ANY);
  else if (inet_aton(lis_ip, &tcp_addr.sin_addr) == 0)
    err(1, "[%s:%d] [error] listening ip[%s] is invalid, abort!",
	conn_str, mypid, lis_ip);
  tcp_addr.sin_port = htons(lis_port);
}

/*
  clean up
*/
server_proxy::~server_proxy()
{
  clean_fd(clt_fd);
  clean_fd(svr_fd);
  delete id_q;
  delete clt_map;
}

/*
  server write call back
*/
void svr_write_cb(evutil_socket_t fd, short what, void *ctx)
{
  DBG("starts\n");
  client_t *clt = (client_t *)ctx;
  net_msg *m = clt->recv_q->pop(); //this has been coverted to udp payload
  uint16_t old_id = get_dns_id(m->data, m->len, false);
  uint16_t new_id = clt->id_q->pop();
  clt->id_map->push(new_id, old_id);
  put_dns_id(m->data, new_id, false); //udp packet sent to server
  clt->map->push(new_id, clt);
  DBG("server fd[%d]\n", clt->svr_fd);
  if (send(clt->svr_fd, m->data, m->len, 0) == -1)
    err(1, "sendto fails");
  delete m; //query has been set, let's clean data
  DBG("a query sent to server\n");
  DBG("ends\n");
}

/*
  init ssl
*/
void init_ssl()
{
  SSL_load_error_strings();
  SSL_library_init();
}

/*
  upgrade to tls
*/
void upgrade_to_ssl(client_t *clt, uint8_t *d, size_t l)
{
  DBG("starts\n");
  struct bufferevent *old_tcp_bev = clt->bev;
  clt->type = CONN_TLS;
  assert(clt->svr_ctx != NULL);
  clt->ctx = SSL_new(clt->svr_ctx);
  struct event_base *base = clt->base;
  assert(base != NULL);
  //delete the original event
  struct bufferevent *bev = bufferevent_openssl_socket_new(base, clt->s, clt->ctx,
							   BUFFEREVENT_SSL_ACCEPTING,
							   BEV_OPT_CLOSE_ON_FREE);
  assert(bev != NULL);
  clt->bev = bev;
  struct timeval read_to = {READ_TIMEOUT, 0};
  bufferevent_set_timeouts(bev, &read_to, NULL);
  bufferevent_setcb(bev, ssl_read_cb, NULL, buf_ev_cb, (void *)clt);
  bufferevent_enable(bev, EV_READ);

  //change the query to response, just the qr bit for simplicity
  uint8_t *data = new uint8_t[l];
  memcpy(data, d, l);
  print_dns_pkt(data+2, l-2);
  uint16_t flag = 0;
  uint8_t offset = 4; //2 for length, 2 for ID
  memcpy(&flag, data + offset, sizeof(flag));
  flag = ntohs(flag);
  flag |= 0x8000U; //set qr = 1
  flag &= 0xfff0U; //set rcode = 0
  flag = htons(flag);
  memcpy(data + offset, &flag, sizeof(flag));
  //print_dns_pkt(data+2, l-2);
  if (bufferevent_write(old_tcp_bev, data, l) != 0)
    errx(1, "bufferevent_write fails!");
  clt->status = STAT_UPGRADE;
  
  //this old event will be freed when the response has been sent back to client
  bufferevent_disable(old_tcp_bev, EV_READ);

  delete[] data;
  DBG("ends\n");
}

/*
  process the client message received
*/
void process_clt_msg(client_t *clt)
{
  DBG("starts\n");
  net_msg *msg = clt->pending_msg;
  if (msg->len <= 2)
    return;

  uint8_t *cp = msg->data; // current position
  size_t r    = msg->len;  // the rest of length
  uint16_t l  = 0;   // length for each DNS payload
  struct event_base *base = clt->base;
  assert(base != NULL);
  struct event *svr_write_ev = NULL;
  while (r > 0) {
    l = 0;
    memcpy(&l, cp, sizeof(l));
    l = ntohs(l) + 2;
    if (l > r) {
      warnx("[warn] TCP payload is not completed or malformed DNS TCP packet!");
      break;
    }

    //check TO, upgrade to tls or not?
    print_dns_pkt(cp+2, l-2);
    if (has_to_bit(cp+2, l-2)) {
      DBG("TO bit is set!\n");
      if (clt->type != CONN_TLS) //upgrade to tls
	upgrade_to_ssl(clt, cp, l);
      cp += l;
      r -= l;
      break;
    } else {
      DBG("TO bit is NOT set!\n");
    }
    
    DBG("push one packet to client recv queue\n");
    clt->enq_recv(cp+2, l-2);
    
    //schedual a server write event
    svr_write_ev = event_new(base, clt->svr_fd, EV_WRITE, svr_write_cb, (void *)clt);
    assert(svr_write_ev != NULL);
    if (event_add(svr_write_ev, NULL) < 0)
      err(1, "[error] cannot add server write event, abort!");
    DBG("done schedual a server write event\n");
    
    cp += l;
    r -= l;
  }
  if (r == 0)
    msg->clean();
  else {
    uint8_t *d = new uint8_t[r];
    assert(d);
    memcpy(d, cp, r);
    msg->copy(d, r, false);
    delete[] d;
  }
  DBG("ends\n");
}

/*
  something (queries) coming from client
*/
void clt_read_cb(struct bufferevent *bev, void *ctx)
{
  DBG("starts\n");
  client_t *clt = (client_t *)ctx;
#ifdef DEBUG
  char *ip = inet_ntoa(((sockaddr_in *)(clt->addr))->sin_addr);
  uint16_t port = ntohs(((sockaddr_in *)(clt->addr))->sin_port);
#endif
  DBG("data from client[%s:%u]:\n", ip, port);  
  struct evbuffer *input = bufferevent_get_input(bev);
  size_t len = evbuffer_get_length(input);
  
  //tcp payload length (2 bytes) and the real payload might be sent
  //in TWO different tcp segments
  if (len == 0) {
    warnx("buffer len == 0, do not process, return");
    return;
  }
  uint8_t *data = new uint8_t[len];
  bufferevent_read(bev, data, len);
  clt->pending_msg->append(data, len);
  clean_pt(data, true);  
  if (clt->pending_msg->len <= (size_t)2) { //only tcp length is read
    DBG("no real data, do not process\n");
    return;
  }
  DBG("process client message\n");
  process_clt_msg(clt);
  DBG("ends\n");
}

/*
  if buffer is flushed/empty AND ssl upgrade is done
  free the old tcp buffer event and clean up
*/
void clt_write_check_ssl_cb(struct bufferevent *bev, void *ctx)
{
  DBG("starts\n");  
  client_t *clt = (client_t *)ctx;
  struct evbuffer *output = bufferevent_get_output(bev);
  size_t len = evbuffer_get_length(output);
  if (len == 0 && clt->status == STAT_UPGRADE) {
    //next line is very important since we don't want close the real tcp socket
    //ssl still uses the tcp socket; note "BEV_OPT_CLOSE_ON_FREE"
    assert(bufferevent_setfd(bev, -1) == 0);
    bufferevent_free(bev);
    DBG("free bufferevent since this TLS upgarde is done\n");
  }
  DBG("ends\n");
}

/*
  call back for a new connection
*/
void accept_conn_cb(struct evconnlistener *listener, evutil_socket_t fd, 
		    struct sockaddr *addr, int len, void *ctx)
{
  DBG("starts\n");
#ifdef DEBUG
  char *ip = inet_ntoa(((sockaddr_in *)addr)->sin_addr); //ipv4 only for now
  uint16_t port = ntohs(((sockaddr_in *)addr)->sin_port);
#endif
  DBG("a new connection from[%s:%u] fd[%d]...\n", ip, port, fd);
  struct event_base *base = evconnlistener_get_base(listener);
  
  //note that "BEV_OPT_CLOSE_ON_FREE";
  //this requires reset the fd before free the bufferevent in the case of upgrading to ssl
#ifdef USE_PTHREAD    
  struct bufferevent *bev = bufferevent_socket_new(base, fd, BEV_OPT_CLOSE_ON_FREE|BEV_OPT_THREADSAFE);
  //struct bufferevent *bev = bufferevent_socket_new(base, fd, BEV_OPT_THREADSAFE);
#else
  struct bufferevent *bev = bufferevent_socket_new(base, fd, BEV_OPT_CLOSE_ON_FREE);
#endif

  clt_arg_t *arg = (clt_arg_t *)ctx;
  client_t *clt = new client_t(CONN_TCP, fd, addr, len, bev, base, 
			       arg->clt_map, arg->id_q, arg->svr_fd);
  //nonblocking
  if (evutil_make_socket_nonblocking(fd) < 0) {
    evutil_closesocket(fd);
    errx(1, "[error] evutil_make_socket_nonblocking fails, abort!");
  }
  clt->svr_ctx = arg->svr_ctx;
  struct timeval read_to = {READ_TIMEOUT, 0};
  bufferevent_set_timeouts(bev, &read_to, NULL);
  bufferevent_setcb(bev, clt_read_cb, clt_write_check_ssl_cb,
		    buf_ev_cb, (void *)clt);
  bufferevent_enable(bev, EV_READ|EV_WRITE);
  DBG("ends\n");
}

/*
  this function is copied from 
  http://www.wangafu.net/~nickm/libevent-book/Ref8_listener.html
*/
static void accept_error_cb(struct evconnlistener *listener, void *ctx)
{
  struct event_base *base = evconnlistener_get_base(listener);
  int err = EVUTIL_SOCKET_ERROR();
  fprintf(stderr, "Got an error %d (%s) on the listener. "
	  "Shutting down.\n", err, evutil_socket_error_to_string(err));
  event_base_loopexit(base, NULL);
}

/*
  start tcp
*/
void server_proxy::start_tcp()
{
  DBG("[%s:%d] starts\n", conn_str, mypid);
  struct evconnlistener *listener;
  int yes = 1, tcpfo = 1, rv = -1, on = 1;
#ifdef USE_PTHREAD
  evthread_use_pthreads(); //otherwise, Segmentation fault
  if (evthread_make_base_notifiable(base) < 0) //this line may be commented out since there is no thread
    err(1, "[%s:%d] [error] couldn't make base notifiable, abort!", conn_str, mypid);
#endif
  
  //client fd
  if ((clt_fd = socket(AF_INET, SOCK_STREAM, 0)) == -1)
    err(1, "[%s:%d] [error] creating tcp socket fails, abort!", conn_str, mypid);
  DBG("clt_fd:%d", clt_fd);
  //set a few options
  if (setsockopt(clt_fd, SOL_SOCKET, SO_REUSEADDR, &yes, sizeof(int)) == -1)
    err(1, "[%s:%d] [error] setsockopt", conn_str, mypid);

  //enable TCP fast open
  rv = setsockopt(clt_fd,          /* socket affected */
		  IPPROTO_TCP,     /* set option at TCP level */
		  TCP_FASTOPEN,    /* name of option */
		  (char *) &tcpfo,  /* the cast is historical cruft */
		  sizeof(int));    /* length of option value */
  if(rv < 0)
    warnx("[%s:%d] [warn] enable TCP fastopen failed!", conn_str, mypid);

  //nonblocking
  if (evutil_make_socket_nonblocking(clt_fd) < 0) {
    evutil_closesocket(clt_fd);
    errx(1, "[%s:%d] [error] evutil_make_socket_nonblocking fails, abort!", conn_str, mypid);
  }
  //keeplive, see evconnlistener_new_bind of listen.c in libevent source code
  if (setsockopt(clt_fd, SOL_SOCKET, SO_KEEPALIVE, (void*)&on, sizeof(on))<0) {
    evutil_closesocket(clt_fd);
    err(1, "[%s:%d] [error] evutil_make_socket_nonblocking fails, abort!", conn_str, mypid);
  }

  //bind address  
  DBG("[%s:%d] binding port[%d], ip[%s]\n", conn_str, mypid, lis_port, lis_ip);
  if (bind(clt_fd, (struct sockaddr *)&tcp_addr, sizeof(tcp_addr)) == -1)
    err(1, "[%s:%d] [error] bind tcp socket fails, abort!", conn_str, mypid);

  //create a libevent listener
  int backlog = 1024; //this may be set to 1024*2 in kernel

  //client argument to pass to the accept call back
  //used to create a client object
  clt_arg_t *clt_arg = NULL;
  clt_arg = new clt_arg_t;
  assert(clt_arg != NULL);
  clt_arg->clt_map = clt_map;
  clt_arg->id_q    = id_q;
  clt_arg->svr_fd  = svr_fd;
  clt_arg->svr_ctx = ctx;
#ifdef USE_PTHREAD
  listener = evconnlistener_new(base, accept_conn_cb, (void *)clt_arg,
				LEV_OPT_CLOSE_ON_FREE|LEV_OPT_REUSEABLE|LEV_OPT_THREADSAFE,
				backlog, clt_fd);
#else
  listener = evconnlistener_new(base, accept_conn_cb, (void *)clt_arg,
				LEV_OPT_CLOSE_ON_FREE|LEV_OPT_REUSEABLE,
				backlog, clt_fd);
#endif
  if (!listener) {
    evutil_closesocket(clt_fd);
    err(1, "[%s:%d] [error] evconnlistener_new fails, abort!", conn_str, mypid);
  }
  evconnlistener_set_error_cb(listener, accept_error_cb);

  //start event loop
  DBG("[%s:%d] start tcp event loop...\n", conn_str, mypid);
  event_base_dispatch(base);

  evconnlistener_free(listener);
  delete clt_arg;
  DBG("[%s:%d] ends\n", conn_str, mypid);
}

/*
  add length to the payload
*/
void tcp_payload(uint8_t *buf, size_t buf_len, uint8_t *d, uint16_t l)
{
  memset(buf, 0, buf_len);
  uint16_t len = htons(l);
  memcpy(buf, &len, sizeof(len));
  memcpy(buf + 2, d, l);  
}

/*
  send data to client
*/
void clt_write_cb(evutil_socket_t fd, short what, void *ctx)
{
  assert(what & EV_WRITE);
  client_t *clt = (client_t *)ctx;
  net_msg *m = clt->send_q->pop();
  uint16_t buf_len = m->len + 2;
  uint8_t buf[buf_len];
  tcp_payload(buf, buf_len, m->data, m->len);
  if (clt->type != CONN_TCP && clt->type != CONN_TLS)
    errx(1, "client is not TCP or TLS!"); 
  DBG("%s send back to client\n", clt->type == CONN_TCP?"TCP":"TLS");
  assert(clt->bev != NULL);
  if (bufferevent_write(clt->bev, buf, buf_len) != 0)
    errx(1, "bufferevent_write fails!");
}

/*
  call recvfrom to receive the responses from server
*/
void recv_from_server(void *ctx)
{
  DBG("starts...\n");
  svr_arg_t *arg = (svr_arg_t *)ctx;
  safe_map<int, client_t *> *clt_map = arg->clt_map; 
  safe_queue<int> *id_q = arg->id_q;
  int svr_fd = arg->svr_fd;

  uint8_t *recv_buf = NULL;
  recv_buf = new uint8_t[MAX_BUF_SIZE];
  assert(recv_buf != NULL);
  memset(recv_buf, 0, MAX_BUF_SIZE);

  sockaddr addr;
  size_t addr_len = sizeof(addr);
  memset(&addr, 0, addr_len);

  //recv from server
  int recv_len = udp_dns_pkt(svr_fd, recv_buf, MAX_BUF_SIZE, &addr, &addr_len);
  DBG("recv a resp from server, len=%d\n", recv_len);
  if (recv_len == 0)  //remote connection is closed, probably recursive is down
    errx(1, "recursive server is down"); //do we need retry here?
#ifdef DEBUG
  print_dns_pkt(recv_buf, recv_len);
#endif
  
  //restore dns id
  uint16_t dns_id = get_dns_id(recv_buf, recv_len, false); //get dns id
  if (!clt_map->find(dns_id)) {
    warnx("dns_id[%u] is not found in client map", dns_id);
    clean_pt(recv_buf, true);
    return;
  }
  client_t *clt = (client_t *)(clt_map->get(dns_id));      //find the client
  DBG("find client[%u] from client map...\n", dns_id);

  uint16_t old_id = clt->id_map->get(dns_id);
  put_dns_id(recv_buf, old_id, false);   //put the original dns id in packet
  DBG("restore dns id from new[%u] to old[%u]\n", dns_id, old_id);
  
  clt->enq_send(recv_buf, recv_len);
  DBG("push to clt sending queue\n");
  
  //clean up client
  DBG("delete client[%u] from client map...\n", dns_id);
  clt_map->erase(dns_id);
  id_q->push(dns_id);
  clt->id_map->erase(dns_id);

  //schedule a write event
  DBG("schedule a client write event...\n");
  struct event_base *base = clt->base;
  struct event *clt_write_ev = NULL;
  clt_write_ev = event_new(base, clt->s, EV_WRITE, clt_write_cb, (void *)clt);
  assert(clt_write_ev != NULL);
  if (event_add(clt_write_ev, NULL) < 0)
    err(1, "[error] cannot add client write event, abort!");

  clean_pt(recv_buf, true);
  DBG("ends...\n");
}

/*
  something (response) is coming from server
  if threadpool is used, do NOT use thread pool here, because:
  this event keeps triggered before thread reads the socket
*/
void svr_read_cb(evutil_socket_t fd, short what, void *ctx)
{
  DBG("data coming from server...\n");
  assert(what & EV_READ);
  svr_arg_t *arg = (svr_arg_t *)ctx;
  safe_map<int, client_t *> *clt_map = arg->clt_map;
  if (clt_map->empty()) {
    warnx("[warn] client_map is empty; server has data coming;return!");
    return;
  }
  recv_from_server(ctx);
  DBG("ends\n");
}

/*
  start server read event
*/
struct event *
server_proxy::start_server_event(void *ctx)
{
  svr_arg_t *arg = (svr_arg_t *)ctx;
  struct event_base *base = arg->base;
  assert(base != NULL);

  //create read event
  struct event *svr_ev = NULL;
  svr_ev = event_new(base, svr_fd, EV_READ|EV_PERSIST, svr_read_cb, ctx);
  assert(svr_ev != NULL);
  if (event_add(svr_ev, NULL) < 0)
    err(1, "[error] cannot add server event, abort!");
  return svr_ev;
}
/*
  this function is copied from
  http://www.wangafu.net/~nickm/libevent-book/Ref6a_advanced_bufferevents.html
*/
static SSL_CTX *
evssl_init(const char *key_file, const char *cert_file)
{
  /* We MUST have entropy, or else there's no point to crypto. */
  if (!RAND_poll())
    return NULL;
  SSL_CTX  *server_ctx = NULL;
  server_ctx = SSL_CTX_new(SSLv23_server_method());
  assert(server_ctx != NULL);  
  if (! SSL_CTX_use_certificate_chain_file(server_ctx, cert_file) ||
      ! SSL_CTX_use_PrivateKey_file(server_ctx, key_file, SSL_FILETYPE_PEM)) {
    err(1, "[error] cannot read cert_file[%s] or key_file[%s]", cert_file, key_file);
    //return NULL;
  }
  SSL_CTX_set_options(server_ctx, SSL_OP_NO_SSLv2);
  return server_ctx;
}

/*
  start: create socket (UDP or TCP) or TLS session and threads
*/
void server_proxy::start()
{
  DBG("[%s:%d] starts\n", conn_str, mypid);

  //create event base
  base = event_base_new();
  if (!base)
    err(1, "[%s:%d] [error] couldn't open event base, abort!", conn_str, mypid);

  //initialize ssl
  init_ssl();
  ctx = evssl_init(key_file, cert_file);
  assert(ctx != NULL);

  //server argument to give to the server event callback
  svr_arg_t *svr_arg = NULL;
  svr_arg = new svr_arg_t;
  assert(svr_arg != NULL);
  svr_arg->svr_fd  = svr_fd;
  svr_arg->clt_map = clt_map;
  svr_arg->id_q    = id_q;
  svr_arg->ctx     = ctx;
  svr_arg->base    = base;

  //create server read event
  struct event * svr_ev = start_server_event((void*)svr_arg);
  
  //set up signal event
  struct event *signal_event = evsignal_new(base, SIGINT, signal_cb, (void *)base);
  assert(signal_event != NULL);
  if (event_add(signal_event, NULL) < 0)
    err(1, "[%s:%d] [error] cannot add signal event, abort!", conn_str, mypid);

  if (conn_type == CONN_UDP) {
    errx(1, "[warn] UDP should not be used, abort!");
  } else if (conn_type == CONN_TCP) {
    start_tcp();
  } else if (conn_type == CONN_TLS) {
    start_ssl();
  }

  //clean up
  DBG("[%s:%d] clear up!\n", conn_str, mypid);
  client_t *tmp_clt = NULL;
  while (!clt_map->empty()) { //delete client event and client
    tmp_clt = clt_map->pop_second();
    assert(tmp_clt != NULL);
    SSL_shutdown(tmp_clt->ctx);
    SSL_free(tmp_clt->ctx);
    bufferevent_free(tmp_clt->bev);
    delete tmp_clt;
    tmp_clt = NULL;
  }
  SSL_CTX_free(ctx);//clean up server
  event_free(signal_event);
  event_free(svr_ev);
  event_base_free(base);
  DBG("[%s:%d] ends\n", conn_str, mypid);
}

/*
  check input error
*/
void server_proxy::check_input()
{
  //check ip address
  if (!ipv4_addr(svr_ip))
    errx(1, "[%s:%d] [error] server ip must be ipv4 address, abort!", conn_str, mypid);
  if (strcmp(lis_ip, "ANY") != 0 && !ipv4_addr(lis_ip))
    errx(1, "[%s:%d] [error] listen ip must be ipv4 address, abort!", conn_str, mypid);

  //check connection type and set string
  if (conn_type == CONN_TCP)
    conn_str = "TCP";
  else if (conn_type == CONN_TLS)
    conn_str = "TLS";
  else
    errx(1, "[%s:%d] [error] unknow connection type[%d], abort!", conn_str, mypid, conn_type);
}

/*
  mutiple queries are sent in one TCP/TLS packet?
*/
void enq_msg(client_t *clt, uint8_t *msg, size_t len)
{
  DBG("starts\n");
  uint8_t *cp = msg; // current position
  size_t r    = len; // the rest of length
  uint16_t l  = 0;   // length for each DNS payload
  struct event_base *base = clt->base;
  assert(base != NULL);
  struct event *svr_write_ev = NULL;
  while (r > 0) {
    l = 0;
    memcpy(&l, cp, sizeof(l));
    l = ntohs(l) + 2;
    if (l > r) {
      warnx("[warn] malformed DNS TCP packet, ignored!");
      break;
    }
    DBG("push one packet to client recv queue\n");
    clt->enq_recv(cp+2, l-2);
    
    //schedual a server write event
    svr_write_ev = event_new(base, clt->svr_fd, EV_WRITE, svr_write_cb, (void *)clt);
    assert(svr_write_ev != NULL);
    if (event_add(svr_write_ev, NULL) < 0)
      err(1, "[error] cannot add server write event, abort!");
    DBG("done schedual a server write event\n");
    
    cp += l;
    r -= l;
  }
  DBG("ends\n");
}


/*
  read from ssl connection
*/
void ssl_read_cb(struct bufferevent *bev, void *arg)
{
  DBG("starts\n");
  client_t *clt = (client_t *)arg;
#ifdef DEBUG
  char *ip = inet_ntoa(((sockaddr_in *)(clt->addr))->sin_addr);
  uint16_t port = ntohs(((sockaddr_in *)(clt->addr))->sin_port);
#endif
  DBG("data from client[%s:%u]:\n", ip, port);
  struct evbuffer *input = bufferevent_get_input(bev);
  size_t len = evbuffer_get_length(input);
  if (len == 0) {
    warnx("buffer len == 0, do not process, return");
    return;
  }
  //tcp payload length (2 bytes) and the real payload might be sent
  //in TWO different tcp segments  
  uint8_t *data = new uint8_t[len];
  //evbuffer_copyout(input, clt->data, len);
  bufferevent_read(bev, data, len);
  clt->pending_msg->append(data, len);
  clean_pt(data, true);  
  if (clt->pending_msg->len <= (size_t)2) { //only tcp length is read
    DBG("no real data, do not process\n");
    return;
  }
  DBG("process client message\n");
  process_clt_msg(clt);
  DBG("ends\n");
}

/*
  new ssl connection
*/
void ssl_accept_cb(struct evconnlistener *listener, int fd,
		  struct sockaddr *addr, int addr_len, void *arg)
{
#ifdef DEBUG
  char *ip = inet_ntoa(((sockaddr_in *)addr)->sin_addr); //ipv4 only for now
  uint16_t port = ntohs(((sockaddr_in *)addr)->sin_port);
#endif
  DBG("a new SSL connection from[%s:%u] fd[%d]...\n", ip, port, fd);
  clt_arg_t *clt_arg = (clt_arg_t *)arg;
  SSL_CTX *svr_ctx = clt_arg->svr_ctx;
  SSL *clt_ctx = SSL_new(svr_ctx);
  struct event_base *base = evconnlistener_get_base(listener);
  assert(base != NULL);
  struct bufferevent *bev = bufferevent_openssl_socket_new(base, fd, clt_ctx,
							   BUFFEREVENT_SSL_ACCEPTING,
							   BEV_OPT_CLOSE_ON_FREE);
  assert(bev != NULL);
  client_t *clt = new client_t(CONN_TLS, fd, addr, addr_len, bev, base,
			       clt_arg->clt_map, clt_arg->id_q, clt_arg->svr_fd);
  assert(clt != NULL);
  clt->svr_ctx = svr_ctx;
  clt->ctx = clt_ctx;
  struct timeval read_to = {READ_TIMEOUT, 0};
  bufferevent_set_timeouts(bev, &read_to, NULL); //set timeout
  bufferevent_setcb(bev, ssl_read_cb, NULL, buf_ev_cb, (void *)clt);
  bufferevent_enable(bev, EV_READ);
}

/*
  start ssl
*/
void server_proxy::start_ssl()
{
  clt_arg_t *clt_arg = NULL;
  clt_arg = new clt_arg_t;
  assert(clt_arg != NULL);
  clt_arg->clt_map = clt_map;
  clt_arg->id_q    = id_q;
  clt_arg->svr_fd  = svr_fd;
  clt_arg->svr_ctx = ctx;
  
  struct evconnlistener *listener = NULL;
  listener = evconnlistener_new_bind(base, ssl_accept_cb, (void *)clt_arg,
				     LEV_OPT_CLOSE_ON_FREE | LEV_OPT_REUSEABLE, 1024,
				     (struct sockaddr *)&tcp_addr, sizeof(tcp_addr));
  assert(listener != NULL);
  event_base_loop(base, 0);

  evconnlistener_free(listener);
  delete clt_arg;
}
