/*

  Copyright 2010 Nathan Phillip Brink

  This file is a part of DistRen.

  DistRen is free software: you can redistribute it and/or modify

  it under the terms of the GNU Affero General Public License as published by

  the Free Software Foundation, either version 3 of the License, or

  (at your option) any later version.

  DistRen 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 Affero General Public License for more details.

  You should have received a copy of the GNU Affero General Public License

  along with DistRen.  If not, see <http://www.gnu.org/licenses/>.

*/

#include "common/multiio.h"

#include <list.h>

#include <fcntl.h>

#include <malloc.h>

#include <poll.h>

#include <stdio.h>

#include <string.h>

struct multiio_socket_info

{

  /* the type of socket */

  multiio_socket_type_t socket_type;

  /* to be passed to the socket handler */

  void *socket_data;

};

struct multiio_socket_type_info

{

  /* type: struct multiio_socket_type_handler_info */

  list_t type_handlers;

};

/*

  The basic concept is that pollfds is an array. Thus, we have to

  have another array reflecting that array to hold socket-specific

  information.

 */

struct multiio_context

{

  /* the array passed to poll() */

  struct pollfd *pollfds;

  /* the number of entries pollfds could hold */

  nfds_t pollfds_len;

  /* the number of entries that pollfds does hold */

  nfds_t nfds;

  /*

    an array whose order mirrors pollfds containing information

    about the fd mentioned in pollfds.

  */

  struct multiio_socket_info *socket_infos;

  /* the number of socket types registered (equivilent to the type ID of the next registered type) */

  multiio_socket_type_t num_socket_types;

  /* the information about each individual socket type keyed by the type ID */

  struct multiio_socket_type_info *socket_types;

};

/*

  stores information about a handler associated with a socket type/class.

 */

struct multiio_socket_type_handler_info

{

  /* the poll() event to which this handler responds */

  short event;

  /* the handler function to call when the event is matched */

  multiio_event_handler_func_t handler;

  /* the pointer passed to multiio_event_handler_register() */

  void *handler_data;

};

/**

   Finds the index in the context->pollfds array of a particular socket.

   @param context the multiio context

   @param fd the socket to search for

   @param index a pointer which will be filled with the index of the socket if found

   @return 0 if the pollfds entry is found, 1 if the entry is not found

 */

int multiio_pollfd_index_by_fd(const multiio_context_t context, int fd, size_t *index);

multiio_context_t multiio_context_new()

{

  struct multiio_context *context;

  context = malloc(sizeof(struct multiio_context));

  if(!context)

    return NULL;

  context->pollfds = malloc(sizeof(struct pollfd) * 2);

  if(!context->pollfds)

    {

      free(context);

      return NULL;

    }

  context->pollfds_len = 2;

  context->nfds = 0;

  context->socket_infos = malloc(sizeof(struct multiio_socket_info) * 2);

  context->num_socket_types = 0;

  context->socket_types = NULL;

  if(!context->pollfds

     || !context->socket_infos)

    {

      free(context->pollfds);

      free(context->socket_infos);

      free(context);

      return NULL;

    }

  return context;

}

int multiio_context_free(multiio_context_t context)

{

  size_t counter;

  /*

    clean up pollfds and socket_infos

   */

  free(context->pollfds);

  free(context->socket_infos);

  /*

    wipe up the socket_types

   */

  for(counter = 0; counter < context->num_socket_types; counter ++)

    list_free(context->socket_types[counter].type_handlers, LIST_DEALLOC);

  free(context->socket_types);

  /*

    bye!

   */

  free(context);

  return 0;

}

struct multiio_poll_travinfo

{

  multiio_context_t multiio;

  struct multiio_socket_info *socket_info;

  struct pollfd *pollfd;

};

/**

   list traverser for multiio_poll()

 */

{

  short event;

  if(!event)

    return TRUE;

  handler_info->handler(travinfo->multiio,

			travinfo->pollfd->fd,

			event,

			handler_info->handler_data,

			travinfo->socket_info->socket_data);

  return TRUE;

}

int multiio_poll(multiio_context_t context)

{

  size_t counter;

  struct multiio_poll_travinfo travinfo;

  poll(context->pollfds, context->nfds, -1);

  for(counter = 0; counter < context->nfds; counter ++)

    if(context->pollfds[counter].revents)

      {

	travinfo.multiio = context;

	travinfo.socket_info = &context->socket_infos[counter];

	travinfo.pollfd = &context->pollfds[counter];

	list_traverse(context->socket_types[context->socket_infos[counter].socket_type].type_handlers,

		      &travinfo,

		      LIST_FRNT | LIST_SAVE);

      }

  return 0;

}

int multiio_socket_type_register(multiio_context_t context, multiio_socket_type_t *new_type)

{

  struct multiio_socket_type_info *new_socket_types;

  new_socket_types = malloc(sizeof(struct multiio_socket_type_info) * (context->num_socket_types + 1));

  if(!new_socket_types)

    return 1;

  *new_type = context->num_socket_types;

  new_socket_types[*new_type].type_handlers = list_init();

  if(!new_socket_types[*new_type].type_handlers)

    {

      free(new_socket_types);

      return 2;

    }

  if(context->num_socket_types)

    {

      memcpy(new_socket_types, context->socket_types, sizeof(struct multiio_socket_type_info) * context->num_socket_types);

      free(context->socket_types);

    }

  context->socket_types = new_socket_types;

  context->num_socket_types ++;

  return 0;

}

int multiio_event_handler_register(multiio_context_t context, multiio_socket_type_t socket_type, short event, multiio_event_handler_func_t handler_func, void *handler_data)

{

  struct multiio_socket_type_handler_info handler_info;

  if(socket_type >= context->num_socket_types)

    return 1;

  if(!event)

    return 2;

  handler_info.event = event;

  handler_info.handler = handler_func;

  handler_info.handler_data = handler_data;

  list_insert_after(context->socket_types[socket_type].type_handlers, &handler_info, sizeof(struct multiio_socket_type_handler_info));

  return 0;

}

int multiio_socket_add(multiio_context_t context, int fd, multiio_socket_type_t socket_type, void *socket_data, short events)

{

  struct pollfd *pollfds;

  struct multiio_socket_info *socket_infos;

  size_t new_pollfds_len;

  int fdstatus;

  int tmp;

  fdstatus = fcntl(fd, F_GETFL);

  fdstatus |= O_NONBLOCK;

  tmp = fcntl(fd, F_SETFL, fdstatus);

  if(tmp == -1)

    {

      perror("fcntl");

      return 1;

    }

  if(socket_type >= context->num_socket_types)

    return 1;

  /**

     extend sockfds array

   */

  if(context->nfds >= context->pollfds_len)

    {

      new_pollfds_len = (context->pollfds_len + 1) * 2;

      pollfds = malloc(sizeof(struct pollfd) * new_pollfds_len);

      socket_infos = malloc(sizeof(struct multiio_socket_info) * new_pollfds_len);

      if(!pollfds

	 || !socket_infos)

	return 1;

      memcpy(pollfds, context->pollfds, sizeof(struct pollfd) * context->pollfds_len);

      memcpy(socket_infos, context->socket_infos, sizeof(struct multiio_socket_info) * context->pollfds_len);

      free(context->pollfds);

      free(context->socket_infos);

      context->pollfds = pollfds;

      context->socket_infos = socket_infos;

      context->pollfds_len = new_pollfds_len;

    }

  pollfds = context->pollfds;

  socket_infos = context->socket_infos;

  pollfds[context->nfds].fd = fd;

  pollfds[context->nfds].events = events;

  socket_infos[context->nfds].socket_type = socket_type;

  socket_infos[context->nfds].socket_data = socket_data;

  context->nfds ++;

  return 0;

}

int multiio_socket_del(multiio_context_t context, int fd)

{

  size_t index;

  int tmp;

  tmp = multiio_pollfd_index_by_fd(context, fd, &index);

  if(tmp)

    return 1;

  /**

   * Special case: one left, if we tried filling in holes... segfault ;-)

   */

  if(context->nfds == 1)

    {

      context->nfds --;

      return 0;

    }

  /**

   * For now we'll unconditionally fill in holes.

   * In fact, it looks simple enough that this method

   * may be what we always use ;-).

   */

  memcpy(&context->pollfds[index], &context->pollfds[context->nfds - 1], sizeof(struct pollfd));

  memcpy(&context->socket_infos[index], &context->socket_infos[context->nfds - 1], sizeof(struct multiio_socket_info));

  context->nfds --;

  return 0;

}

int multiio_socket_event_enable(multiio_context_t context, int fd, short event)

{

  size_t index;

  int tmp;

  tmp = multiio_pollfd_index_by_fd(context, fd, &index);

  if(tmp)

    return 1;

  context->pollfds[index].events |= event;

  return 0;

}

int multiio_socket_event_disable(multiio_context_t context, int fd, short event)

{

  size_t index;

  int tmp;

  tmp = multiio_pollfd_index_by_fd(context, fd, &index);

  if(tmp)

    return 1;

  context->pollfds[index].events &= ~event;

  return 0;

}

int multiio_pollfd_index_by_fd(const multiio_context_t context, int fd, size_t *index)

{

  size_t counter;

  for(counter = 0; counter < context->nfds; counter ++)

    if(context->pollfds[counter].fd == fd)

      {

	*index = counter;

	return 0;

      }

  return 1;

}

       * shifting seems the simplest solution.

       */

      packet->len -= written_amount;

      memmove(packet->data, packet->data + written_amount, packet->len);

    }

  return 0;

}

int remoteio_close(struct remoteio *rem)

{

  int rtn;

  /**

   * See careful_free's and _remoteio_handle_read()'s docs.  If

   * careful_free is nonzero, then we shouldn't free it here because

   * such a free would cause a segfault. However, whoever set

   * rem->careful_free to nonzero will handle resetting

   * rem->careful_free to zero and calling remoteio_close() if

   * necessary.

   */

  if(rem->careful_free)

    {

      rem->careful_free = 2;

      return 0;

    }

  /* cleanup multiiio stuff */

  multiio_socket_del(rem->opts->multiio, rem->sock);

  /* backend-specific cleanup */

  rtn = funcmap[rem->method].close_func(rem);

  /* this part is normal ;-) */

  free(rem->inbuf.data);

  q_free(rem->outmsgs, (list_dealloc_func_t)remoteio_packet_free);

  free(rem);

  return rtn;

}

/**

 * Frees an entire packet, including the passed pointer. If you just

 * want the contents of the packet free()ed, just do

 * free(packet.data);

 */

void remoteio_packet_free(struct remoteio_packet *packet)

{

  free(packet->data);

  free(packet); 

}

int _remoteio_getserver_traverse(char *servername, struct remoteio_server *aserver)

{

  if(!strcmp(aserver->name, servername))

    return FALSE; /* stop traversal */

  return TRUE;

}

struct remoteio_server *remoteio_getserver(const struct remoteio_opts *opts, const char *servername)

{

  int traversal_result;

  char *dispensible_servername;

  dispensible_servername = strdup(servername); /* for the sake of constness... */

  traversal_result = list_traverse(opts->servers, dispensible_servername, (list_traverse_func_t)&_remoteio_getserver_traverse, LIST_FRNT | LIST_ALTR);

  free(dispensible_servername);

  if(traversal_result == LIST_OK)

    return (struct remoteio_server *)list_curr(opts->servers);

  return (struct remoteio_server *)NULL;

}

/**

   different remoteio methods' implementations:

 */

/*

  SSH, via execio

*/

int _remoteio_ssh_open(struct remoteio *rem, struct remoteio_server *server)

{

  char *userhost;

  char *sshargs[] = {rem->opts->ssh_command, NULL /* userhost */, "distrend", "-d", (char *)NULL};

  int rtn;

  if(server->username)

    _distren_asprintf(&userhost, "%s@%s", server->username, server->hostname);

  else

    userhost = strdup(server->hostname);

  sshargs[1] = userhost;

  rtn = execio_open( &rem->execio, "ssh", sshargs);

  if(rtn)

    {

      fprintf(stderr, "error opening remoteio channel to ssh userhost ``%s''\n" , userhost);

      free(userhost);

      return 1;

    }

  free(userhost);

  return 0;

}

int _remoteio_ssh_read(struct remoteio *rem, void *buf, size_t len, size_t *bytesread)

{

  return execio_read(rem->execio, buf, len, bytesread);

}

int _remoteio_ssh_write(struct remoteio *rem, void *buf, size_t len, size_t *byteswritten)

{

  return execio_write(rem->execio, buf, len, byteswritten);

}

int _remoteio_ssh_close(struct remoteio *rem)

{

  int rtn;

  rtn = execio_close(rem->execio);

  if(rtn)

    fprintf(stderr, "%s:%d: error closing execio\n", __FILE__, __LINE__);

  return rtn;

}

#ifndef _WIN32

/*

  local sockets implementation (``named pipes''), unix-only

 */

int _remoteio_sock_open(struct remoteio *rem, struct remoteio_server *server)

{

  int sock;

  struct sockaddr_un sockaddr;

  /*

    The POSIX docs pretty much say that I can't depend on sockpath being able to be longer than 

    some proprietary length. So, if the compiler specifies a long path for RUNSTATEDIR, it could

    cause a buffer overflow.

   */

  char *sockpath = RUNSTATEDIR "/distrend.sock";

  unsigned int sockaddr_len;

  sock = socket(AF_UNIX, SOCK_STREAM, 0);

  if(sock == -1)

    {

      perror("socket");

      return 1;

    }

  sockaddr.sun_family = AF_UNIX;

  /*

    The terminating NULL should not be included in what's copied to sun_path,

    although it won't hurt as long as strlen(sockpath) < max socket length

   */

  for(sockaddr_len = 0; sockpath[sockaddr_len]; sockaddr_len ++)

    sockaddr.sun_path[sockaddr_len] = sockpath[sockaddr_len];

  if(connect(sock, (struct sockaddr *)&sockaddr, sockaddr_len) == -1)

    {

      perror("connect");

      close(sock);

      return 1;

    }

  rem->sock = sock;

  return 0;

}

#endif /* _WIN32 */

int _remoteio_sock_close(struct remoteio *rem)

{

  close(rem->sock);

  return 0;

}

int _remoteio_sock_read(struct remoteio *rem, void *buf, size_t len, size_t *bytesread)

{

  ssize_t readrtn;

  readrtn = read(rem->sock, buf, len);

  /*

    The following is valid for blocking sockets:

   */

  if(readrtn == -1)

    {

      /*

	in this case, we may have been interrupted by a signal and errno == EINTR

	or the connection was reset and errno = ECONNRESET

	Some of these are not error conditions:

       */

      perror("read");

      if(errno != EINTR)

	{

	  fprintf(stderr, "error reading socket in remoteio_sock_read\n");

	  return 1;

	}

      return 0;

    }

  *bytesread = readrtn;

  if(!readrtn)

    {

      /*

	means EOF except when FD is in ``message-nondiscard'' or ``message-discard''

	modes.

       */

      return 1;

    }

  return 0;

}

int _remoteio_sock_write(struct remoteio *rem, void *buf, size_t len, size_t *byteswritten)

{

  ssize_t writertn;

  writertn = write(rem->sock, buf, len);

  if(writertn == -1)

    {

      perror("write");

      if(errno != EINTR)

	{

	  fprintf(stderr, "error writing to socket in remoteio_sock_write()\n");

	  return 1;

	}

    }

  *byteswritten = writertn;

  if(!writertn)

    {

      /*

	should we consider this an error? I'm pretty

	sure we should :-/

       */

      fprintf(stderr, "write() returned 0 in remoteio_sock_write()\n");

      return 1;

    }

  return 0;

}

/**

   TCP, taking advantage of the two generic socks read/write functions:

 */

int _remoteio_tcp_open(struct remoteio *rem, struct remoteio_server *server)

{

  int tmp;

  int tmp2;

  int sock;

  char *hostname;

  char *port;

  struct addrinfo addrinfo_hints;

  struct addrinfo *addrinfo_res;

  /**

     only hostname should be free()-ed, not port,

     because both are from the same block of malloc()-ed

     memory

   */

  hostname = strdup(server->hostname);

  for(port = hostname;

      *port && *port != ':';

      port ++)

    ;

  if(*port)

    {

      *port = '\0';

      port ++;

    }

  else

    port = REMOTEIO_DEFAULT_PORT;

  memset(&addrinfo_hints, '\0', sizeof(struct addrinfo));

  addrinfo_hints.ai_family = AF_UNSPEC;

#ifdef _WIN32

  /* windows lacks stuff documented in POSIX, I guess :-( */

  addrinfo_hints.ai_flags = 0;

#else

  addrinfo_hints.ai_flags = AI_V4MAPPED | AI_ADDRCONFIG;

#endif

  addrinfo_hints.ai_socktype = SOCK_STREAM;

  tmp = getaddrinfo(server->hostname, port, &addrinfo_hints, &addrinfo_res);

  if(tmp)

    fprintf(stderr, "error resolving %s:%s: %s\n", server->hostname, port, gai_strerror(tmp));

  fprintf(stderr, "connecting to %s[%s]:%s\n", server->hostname, addrinfo_res->ai_canonname, port);

  free(hostname);

  sock = socket(addrinfo_res->ai_family, SOCK_STREAM, addrinfo_res->ai_protocol);

  if(sock == -1)

    {

      perror("socket");

      freeaddrinfo(addrinfo_res);

      return 1;

    }

  tmp = connect(sock, addrinfo_res->ai_addr, addrinfo_res->ai_addrlen);

  tmp2 = errno;

  freeaddrinfo(addrinfo_res);

  errno = tmp2;

  if(tmp == -1)

    {

      perror("connect");

      close(sock);

      return 1;

    }

  rem->sock = sock;

  return 0;

}

int _remoteio_tcp_close(struct remoteio *rem)

{

  close(rem->sock);

  return 0;

}