/*
 * Copyright (C) 2015 by the University of Southern California
 * $Id: client_proxy.cc 24217 2015-05-28 07:19:22Z 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 "client_proxy.hh"
#include <iostream>
#include <signal.h>
#include <assert.h>
#include <sys/socket.h>
#include <arpa/inet.h>

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

#include <ldns/ldns.h>

#include "global_var.h"
using namespace std;

#define MAX_BUF_SIZE    4096
#define MIN_BUF_SIZE    64

#define STARTTLS_RETRY 3

/*
  get the id in DNS header
*/
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);
}

/*
  fill in the address information
*/
void fill_addr(struct sockaddr_in *addr, const char *ip, int port)
{
  memset(addr, 0, sizeof(struct sockaddr_in));
  addr->sin_family = AF_INET;
  if (inet_aton(ip, &(addr->sin_addr)) == 0)
    err(1, "[error] ip[%s] is invalid, abort!", ip);
  addr->sin_port = htons(port);
}

/*
  create a tcp payload: add length of to the first two bytes
*/
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);
}

/*
  get a tcp socket
*/
static int 
tcp_conn_svr(struct sockaddr *addr, socklen_t addr_len)
{ 
  DBG("start\n");
  int fd = -1;
  if ((fd = socket(AF_INET, SOCK_STREAM, 0)) == -1)
    err(1, "socket");
  
  if (connect (fd, addr, addr_len) < 0)
    err(1, "connect");

  //nonblocking
  if (evutil_make_socket_nonblocking(fd) < 0) {
    evutil_closesocket(fd);
    errx(1, "[error] evutil_make_socket_nonblocking fails, abort!");
  }

  DBG("end\n");
  return fd;
}

client_proxy::client_proxy(const char *server_ip, int server_port,
			   const char *listen_ip, int listen_port,
			   int connection_type)
{
  svr_ip    = server_ip;
  svr_port  = server_port;
  lis_ip    = listen_ip;
  lis_port  = listen_port;
  conn_type = connection_type;
  svr_fd    = -1;
  my_udp_fd = -1;
  my_tcp_fd = -1;
  //set the server and my addresses
  fill_addr(&svr_addr, svr_ip.c_str(), svr_port);
  fill_addr(&my_addr, lis_ip.c_str(), lis_port);
  clt_map = new safe_map<uint16_t, client_t *>;
  starttls_success = false;
}

client_proxy::~client_proxy()
{
  client_t *tmp_clt = NULL;
  while (!clt_map->empty()) { //delete client event and client
    tmp_clt = clt_map->pop_second();
    assert(tmp_clt != NULL);
    delete tmp_clt;
    tmp_clt = NULL;
  }
  delete clt_map;
}

/*
  signal callback; 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 << "Interrupt signal (" << sig << ") received, exit in two seconds!\n";
  event_base_loopexit(base, &delay);
}

/*
  helper for starttls event callback
*/
void
client_proxy::starttls_event_cb_helper(struct bufferevent *bev, short events, void *ctx)
{
  (static_cast<client_proxy *>(ctx))->starttls_event_cb(bev, events);
}

/*
  the real starttls event callback
*/
void
client_proxy::starttls_event_cb(struct bufferevent *bev, short events)
{
  DBG("start\n");
  assert(bev);
  evutil_socket_t fd = bufferevent_getfd(bev);
  assert(fd != -1);
  struct event_base *base = bufferevent_get_base(bev);
  assert(base);
  if (events & BEV_EVENT_CONNECTED) {
    DBG("fd[%d] connected\n", (int)fd);
  } else if (events & (BEV_EVENT_EOF|BEV_EVENT_ERROR|BEV_EVENT_TIMEOUT)) {
    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()):"");
    starttls_success = false;
    event_base_loopexit(base, NULL);
  }
  DBG("end\n");
}

/*
  helper for starttls read callback
*/
void
client_proxy::starttls_read_cb_helper(struct bufferevent *bev, void *ctx)
{
  (static_cast<client_proxy *>(ctx))->starttls_read_cb(bev);
}

/*
  the real starttls read callback
*/
void
client_proxy::starttls_read_cb(struct bufferevent *bev)
{
  DBG("start\n");
  assert(bev);
  struct evbuffer *input = bufferevent_get_input(bev);
  assert(input);
  size_t len = evbuffer_get_length(input);
  uint8_t *data = new uint8_t[len];
  assert(data);
  bufferevent_read(bev, data, len);
  if (has_to_bit(data+2, len-2)) {//tcp/ssl payload
    DBG("TO bit is set!\n");
    starttls_success = true;
  } else {
    DBG("TO bit is NOT set!\n");
    starttls_success = false;
  }
  delete[] data;
  struct event_base *base = bufferevent_get_base(bev);
  assert(base);
  event_base_loopexit(base, NULL);
  DBG("end\n");
}

/*
  do starttls
*/
void client_proxy::do_starttls()
{
  DBG("start\n");
  struct event_base *base = event_base_new();
  assert(base);
  struct bufferevent *bev = bufferevent_socket_new(base, svr_fd, 0);
  assert(bev);
  
  struct timeval read_to = {5, 0};  //5 sec timeout for waiting for response 
  bufferevent_set_timeouts(bev, &read_to, NULL);
  bufferevent_setcb(bev, &client_proxy::starttls_read_cb_helper, NULL, &client_proxy::starttls_event_cb_helper, this);
  bufferevent_enable(bev, EV_READ|EV_WRITE);

  //send a starttls query
  uint8_t *data;
  size_t len;
  query_t query = {"STARTTLS", "TXT", "CH"};
  format_dns_query(&data, &len, &query, EDNS_Z_DO | EDNS_Z_TO);
  uint8_t *tcp_data = new uint8_t[len + 2];
  memset(tcp_data, 0, len+2);
  tcp_payload(tcp_data, len+2, data, len);
  delete[] data;
  bufferevent_write(bev, tcp_data, len+2);

  //done, clean up
  event_base_dispatch(base);
  bufferevent_free(bev);
  event_base_free(base);
  delete[] tcp_data;
  DBG("end\n");
}

/*
  tcp connect to server, starttls as necessary
*/
void client_proxy::connect_server()
{
  DBG("start\n");
  svr_fd = tcp_conn_svr((struct sockaddr *)&svr_addr, sizeof(svr_addr));
  DBG("tcp socket connected\n");
  int i = 0;
  if (conn_type == CONN_STARTTLS) {
    for (i=0; i<STARTTLS_RETRY; i++) {//retry
      do_starttls();
      if (starttls_success) break;
      sleep(3);
    }
    if (!starttls_success)
      errx(1, "[error] server does not support TLS/STARTLS");
    else
      DBG("STARTTLS successful!\n");
  }
  DBG("end\n");
}

/*
  helper for ssl event callback
*/
void
client_proxy::ssl_event_cb_helper(struct bufferevent *bev, short events, void *ctx)
{
  (static_cast<client_proxy *>(ctx))->ssl_event_cb(bev, events);
}

/*
  the real ssl event callback
*/
void
client_proxy::ssl_event_cb(struct bufferevent *bev, short events)
{
  DBG("start\n");
  //mtx_bev.lock();
#ifdef DEBUG
  evutil_socket_t fd = bufferevent_getfd(bev);
#endif
  struct event_base *base = bufferevent_get_base(bev);
  assert(base);
  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)) {
    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()):"");
    mtx_svr_bev.lock();
    renew_ssl(bev);
    mtx_svr_bev.unlock();
  }
  //mtx_bev.unlock();
  DBG("end\n");
}

/*
  renew the ssl connection
*/
void 
client_proxy::renew_ssl(struct bufferevent *bev)
{
  DBG("start\n");
  assert(bev);
    
  //clean up ssl and socket
  //DO NOT free ssl object here, otherwise seg fault
  //bufferevent_free probably will does SSL_free()
  DBG("clean up ssl\n");
  SSL_set_shutdown(ssl, SSL_RECEIVED_SHUTDOWN);
  SSL_shutdown(ssl);

  //the following is probably not necessary; 
  //but just to be clean since it is old
  SSL_CTX_free(ssl_ctx);
  ssl_ctx = SSL_CTX_new(SSLv23_client_method());
  if (!ssl_ctx) ERR_print_errors_fp (stderr);

  //clean up old event
  DBG("stop old event\n");
  bufferevent_free(bev);

  //try to renew ssl connection
  connect_server();
  DBG("renew a connection successfully, fd[%d]\n", svr_fd);

  //create a new bufferevent
  DBG("new bufferevent\n");
  ssl = SSL_new(ssl_ctx);
  if(!ssl) {
    ERR_print_errors_fp (stderr);
    errx(1, "[error] cannot create ssl, abort!");
  }
  DBG("svr_fd[%d]\n", svr_fd);
  svr_bev = bufferevent_openssl_socket_new(base, svr_fd, ssl,
				       BUFFEREVENT_SSL_CONNECTING,
				       BEV_OPT_CLOSE_ON_FREE);
  bufferevent_setcb(svr_bev, 
		    &client_proxy::svr_read_cb_helper,
		    NULL, 
		    &client_proxy::ssl_event_cb_helper, this);
  bufferevent_enable(svr_bev, EV_READ|EV_WRITE);
}

/*
  helper for server read (ssl or tcp) callback
*/
void
client_proxy::svr_read_cb_helper(struct bufferevent *bev, void *ctx)
{
  (static_cast<client_proxy *>(ctx))->svr_read_cb(bev);
}

/*
  the real server read callback
*/
void
client_proxy::svr_read_cb(struct bufferevent *bev)
{
  DBG("start\n");
  //mtx_bev.lock();
  assert(bev);
  struct evbuffer *input = bufferevent_get_input(bev);
  assert(input);
  size_t len = evbuffer_get_length(input);
  uint8_t *data = new uint8_t[len];
  assert(data);
  bufferevent_read(bev, data, len);
  //mtx_bev.unlock();

  //send back to client
  uint16_t dns_id = get_dns_id(data, len, true);
  if (!clt_map->find(dns_id)) {
    warnx("[warn] id[%u] is not in clt map", dns_id);
    return;
  }
  client_t *clt = clt_map->get(dns_id);
  clt_map->erase(dns_id);
  int send_len = sendto(my_udp_fd, data+2, len-2, 0,
			(struct sockaddr*)&(clt->addr), clt->addr_len);
  if (send_len == -1)
    err(1, "[error] sendto");

  //clean up
  delete[] data;
  delete clt;
  DBG("end\n");
}

/*
  helper for tcp event callback
*/
void
client_proxy::tcp_event_cb_helper(struct bufferevent *bev, short events, void *ctx)
{
  (static_cast<client_proxy *>(ctx))->tcp_event_cb(bev, events);
}

/*
  the real tcp event callback
*/
void
client_proxy::tcp_event_cb(struct bufferevent *bev, short events)
{
  DBG("start\n");
  //mtx_bev.lock();
#ifdef DEBUG
  evutil_socket_t fd = bufferevent_getfd(bev);
#endif
  struct event_base *base = bufferevent_get_base(bev);
  assert(base);
  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)) {
    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()):"");
    mtx_svr_bev.lock();
    bufferevent_free(svr_bev);
    if (!setup_tcp_event()) {
      event_base_free(base);
      errx(1, "[error] cannot setup tcp event, abort!");
    }
    mtx_svr_bev.unlock();
  }
  //mtx_bev.unlock();
  DBG("end\n");
}

/*
  setup tcp connection
*/
bool client_proxy::setup_tcp_event()
{
  svr_bev = bufferevent_socket_new(base, -1, BEV_OPT_CLOSE_ON_FREE);
  assert(svr_bev);
  //connect_server();
  
  if (bufferevent_socket_connect(svr_bev, (struct sockaddr *)&svr_addr, sizeof(svr_addr))<0) {
    bufferevent_free(svr_bev);
    cerr << "[error] bufferevent_socket_connect fails in run_tcp" << endl;
    return false;
  }
  
  //struct timeval read_to = {READ_TIMEOUT, 0};
  //bufferevent_set_timeouts(svr_bev, &read_to, NULL);
  bufferevent_setcb(svr_bev,
		    &client_proxy::svr_read_cb_helper, NULL,
		    &client_proxy::tcp_event_cb_helper, this);
  bufferevent_enable(svr_bev, EV_READ|EV_WRITE);
  return true;
}

/*
  start with tcp connection to recursive server
*/
void client_proxy::run_tcp()
{
  DBG("start\n");

  if (!setup_tcp_event()) {
    warnx("[error] fail to setup tcp server event!");
    return;
  }
  
  //set up event for recv queries from client
  clt_read_ev_init();

  event_base_dispatch(base);
  DBG("end\n");
}

/*
  start with ssl connection to recursive server
*/
void client_proxy::run_ssl()
{
  DBG("start\n");
  SSL_load_error_strings();
  SSL_library_init();
  connect_server();
  DBG("successful server connection\n");

  //set up a ssl buffer event as client
  ssl_ctx = SSL_CTX_new(SSLv23_client_method());
  if (!ssl_ctx) ERR_print_errors_fp (stderr);
  ssl = SSL_new(ssl_ctx);
  if(!ssl) {
    ERR_print_errors_fp (stderr);
    errx(1, "[error] cannot create ssl, abort!");
  }

  //http://stackoverflow.com/questions/7784573/callback-to-member-function-using-libevent
  //mtx_bev.lock();
  svr_bev = bufferevent_openssl_socket_new(base, svr_fd, ssl,
					   BUFFEREVENT_SSL_CONNECTING,
					   BEV_OPT_CLOSE_ON_FREE);
  bufferevent_setcb(svr_bev,
		    &client_proxy::svr_read_cb_helper, 
		    NULL, 
		    &client_proxy::ssl_event_cb_helper, this);
  bufferevent_enable(svr_bev, EV_READ|EV_WRITE);
  //mtx_bev.unlock();
  
  //set up event for recv queries from client
  clt_read_ev_init();

  event_base_dispatch(base);
  DBG("end\n");
}

/*
  get started
*/
void client_proxy::start()
{
  //create event base
  base = event_base_new();
  if (!base)
    err(1, "[error] couldn't open event base, abort!");

  //create udp socket for receiving queries
  if ((my_udp_fd = socket(AF_INET, SOCK_DGRAM, 0)) == -1)
    err(1, "[error] create udp socket fails, abort!");

  //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, "[error] cannot add signal event, abort!");
  
  DBG("start\n");
  if (conn_type == CONN_TCP)
    run_tcp();
  else if (conn_type == CONN_SSL ||
	   conn_type == CONN_STARTTLS)
    run_ssl();
  else
    errx(1, "[error] unknown connection type[%d], abort!", conn_type);
  
  event_free(signal_event);
  event_base_free(base);
  DBG("end\n");
}

/*
  helper for client read callback
*/
void
client_proxy::clt_read_cb_helper(evutil_socket_t fd, short events, void *ctx)
{
  (static_cast<client_proxy *>(ctx))->clt_read_cb(fd, events);
}

/*
  the real client read callback
*/
void
client_proxy::clt_read_cb(evutil_socket_t fd, short events)
{
  DBG("data coming from local, fd[%d]...\n", (int)fd);
  assert(events & EV_READ);

  //create a new client buffer
  client_t *clt = new client_t;
  clt->addr_len = sizeof(clt->addr);
  assert(clt);

  //recv buffer
  uint8_t *recv_buf = new uint8_t[MAX_BUF_SIZE];
  assert(recv_buf);
  memset(recv_buf, 0, MAX_BUF_SIZE);

  //udp recv from
  int recv_len = recvfrom(fd, recv_buf, MAX_BUF_SIZE, 0, (struct sockaddr *)&(clt->addr), &(clt->addr_len));
  if (recv_len == -1)
    err(1, "[error] recvfrom return -1 in recv_udp_cb");
  DBG("recv udp, len[%d]...\n", recv_len);

  //record dns id to client mapping
  uint16_t dns_id = get_dns_id(recv_buf, recv_len, false);
  clt_map->push(dns_id, clt);
  
  //copy udp query as tcp payload
  size_t msg_len = recv_len + 2;
  uint8_t *msg = new uint8_t[msg_len];
  assert(msg);
  memset(msg, 0, msg_len);
  tcp_payload(msg, msg_len, recv_buf, recv_len);
  //print_dns_pkt(msg+2, msg_len-2);

  //bufferevent write
  DBG("bufferevent write to server, len[%lu]...\n", msg_len);
  mtx_svr_bev.lock();
  bufferevent_write(svr_bev, msg, msg_len);
  mtx_svr_bev.unlock();
  delete[] recv_buf;
  delete[] msg;
  DBG("ends\n");
}

/*
  init client read event to receive queries from local application
*/
void client_proxy::clt_read_ev_init()
{
  DBG("start\n");
  if (bind(my_udp_fd, (struct sockaddr *)&my_addr, sizeof(my_addr)) != 0)
    err(1, "[error] bind udp socket fails, abort!");

  assert(base != NULL);
  struct event *clt_read_ev = NULL;
  clt_read_ev = event_new(base, my_udp_fd, EV_READ|EV_PERSIST, clt_read_cb_helper, this);
  assert(clt_read_ev != NULL);
  if (event_add(clt_read_ev, NULL) < 0)
    err(1, "[error] cannot add receive query event, abort!");
  event_list.push_back(clt_read_ev);
  DBG("end\n");
}
