/* select_group.h
 * By Ron Bowes
 * Created August, 2008
 *
 * (See LICENSE.md)
 *
 * This module implements a simple interface to the select() function that
 * works across Windows, Linux, BSD, and Mac. Any (reasonable) number of
 * sockets can be added and when any one of them has data, callbacks are used
 * to notify the main program.
 *
 * This library is single-threaded (except on Windows.. I'll get to that). I've
 * compiled and tested it on Linux, FreeBSD, Mac, and Windows, and it works
 * beautifully on all of them. On any of those platforms it can select just
 * fine on stream sockets, datagram sockets, listeners, and pipes (including
 * stdin on Windows).
 *
 * Windows support for pipes is a special case. Because Windows can't select()
 * on a pipe or HANDLE, I had to implement some special code. Basically, it
 * polls -- instead of adding pipes to the select(), it times out the select
 * after a set amount of time (right now, it's 100ms). That means that every
 * 100ms, select() returns and checks if any input is waiting on the pipes. Not
 * the greatest solution, but it isn't the greatest OS for networking stuff.
 *
 * Even worse, stdin is a special case. stdin can be read through a pipe, so
 * the polling code works great -- except that Windows won't echo types
 * characters unless stdin is being actively read. Rather than introducing a
 * 100ms-delay to everything types, that would probably make me even more
 * crazy, I created the Windows-specific function get_stdin_handle(). The first
 * time it's called, it creates a thread that reads from stdin and writes to a
 * pipe. That pipe can be added to select() and everything else works the same.
 * It's an ugly hack, I know, but when writing Ncat (http://nmap.org/ncat)
 * David Fifield came up with the same solution. Apparently, it's the best
 * we've got.
 */


#ifndef __SELECT_GROUP_H__
#define __SELECT_GROUP_H__

/* Updates for dnscat2 */
#define SELECT_GROUP_VERSION "1.01"

#include <stdlib.h>

#ifdef WIN32
#include <winsock2.h>
#else
#endif

#include "types.h"

/* The maximum number of possible sockets (huge number, but I want to prevent overflows). Note that this is
 * sort of a range, because the number of sockets are doubled each time. So it's between 32768 and 65536. */
#define SOCKET_LIST_MAX_SOCKETS (65536/2)

/* The time, in milliseconds, between select() timing out and polling for pipe data (on Windows). */
#ifdef WIN32
#define TIMEOUT_INTERVAL 100
#endif

/* Different types of sockets, which will affect different aspects. */
typedef enum
{
  /* No special treatment (anything can technically use this one).  */
  SOCKET_TYPE_STREAM,

  /* Uses recvfrom() and passes along the socket address). */
  SOCKET_TYPE_DATAGRAM,

  /* Listening implies a stream. */
  SOCKET_TYPE_LISTEN,
#ifdef WIN32
  /* For use on Windows, only, is handled separately. */
  SOCKET_TYPE_PIPE
#endif
} SOCKET_TYPE_t;

/* Possible return values from callback functions. */
typedef enum
{
  SELECT_OK,           /* Everything went well. */
  SELECT_REMOVE,       /* Remove the socket from the list. */
  SELECT_CLOSE_REMOVE, /* Close the socket and remove it from the list. */
} SELECT_RESPONSE_t;

/* Define callback function types. I have to make the first parameter 'void*' because the struct hasn't been defined
 * yet, and the struct requires these typedefs to be in place. */
typedef SELECT_RESPONSE_t(select_ready)(void *group, int s, void *param);
/* 'addr' will only be filled in for datagram requests. */
typedef SELECT_RESPONSE_t(select_recv)(void *group, int s, uint8_t *data, size_t length, char *addr, uint16_t port, void *param);
typedef SELECT_RESPONSE_t(select_listen)(void *group, int s, void *param);
typedef SELECT_RESPONSE_t(select_error)(void *group, int s, int err, void *param);
typedef SELECT_RESPONSE_t(select_closed)(void *group, int s, void *param);
typedef SELECT_RESPONSE_t(select_timeout)(void *group, void *param);

/* This struct is for internal use. */
typedef struct
{
  /* The socket. */
  int             s;
#ifdef WIN32
  /* A pipe (used for Windows' named pipes. */
  HANDLE          pipe;
#endif
  /* Datagram, stream, etc. */
  SOCKET_TYPE_t   type;

  /* The function to call when the socket is ready for data. */
  select_ready   *ready_callback;

  /* The function to call when data arrives. */
  select_recv    *recv_callback;

  /* The function to call when a connection arrives. */
  select_listen  *listen_callback;

  /* The function to call when there's an error. */
  select_error   *error_callback;

  /* The function to call when the connection is closed. */
  select_closed  *closed_callback;

  /* The number of bytes being waited on. If set to 0, will trigger on all
   * incoming data. */
  size_t          waiting_for;

  /* The buffer that holds the current bytes. */
  uint8_t        *buffer;

  /* The number of bytes currently stored in the buffer. */
  size_t          buffered;

  /* This is set after the socket has received the signal that it's ready to
   * receive data. */
  NBBOOL         ready;

  /* Set to 'false' when the socket is 'deleted'. It's easier than physically
   * removing it from the list, so until I implement something heavy weight
   * this will work. */
  NBBOOL         active;

  /* Stores a piece of arbitrary data that's sent to the callbacks. */
  void           *param;
} select_t;

/* This is the primary struct for this module. */
typedef struct
{
  /* A list of the select_t objects. */
  select_t **select_list;

  /* The current number of "select_t"s in the list. */
  size_t current_size;

  /* The maximum number of "select_t"s in the list before realloc() has to expand it. */
  size_t maximum_size;
#ifdef WIN32
  /* The number of milliseconds that have elapsed; used for timeouts. */
  uint32_t elapsed_time;
#endif

  /* The handle to the highest-numbered socket in the list (required for select() call). */
  int biggest_socket;

  /* The function to call when the timeout time expires. */
  select_timeout *timeout_callback;

  /* A parameter that is passed to the callback function. */
  void *timeout_param;
} select_group_t;

/* Allocate memory for a select group */
select_group_t *select_group_create();

/* Destroy and cleanup the group. */
void select_group_destroy(select_group_t *group);

/* Add a socket to the group. */
void select_group_add_socket(select_group_t *group, int s, SOCKET_TYPE_t type, void *param);

#ifdef WIN32
/* Add a pipe to the group. The 'identifier' is treated as a socket and is used in place of a socket
 * to look up the pipe. */
void select_group_add_pipe(select_group_t *group, int identifier, HANDLE pipe, void *param);
#endif

/* Set a callback that's called when the socket becomes ready to send data. */
select_ready   *select_set_ready(select_group_t *group, int s, select_ready *callback);

/* Set the recv() callback. This will return with as much data as comes in, or with the number of bytes set by
 * set_group_wait_for_bytes(), if that's set. Returns the old callback, if set. */
select_recv    *select_set_recv(select_group_t *group, int s, select_recv *callback);

/* Set the listen() callback for incoming connections. It's up to the callback to perform the accept() to
 * get the new socket. */
select_listen  *select_set_listen(select_group_t *group, int s, select_listen *callback);

/* Set the error callback, for socket errors. If SELECT_OK is returned, it assumes the error's been handled
 * and will continue to select() on the socket. In almost every case, SELECT_OK is the wrong thing to return.
 * If this isn't defined, the socket is automatically closed/removed. */
select_error   *select_set_error(select_group_t *group, int s, select_error *callback);

/* Set the closed callback. This is called when the connection is gracefully terminated. Like with errors,
 * SELECT_OK is probably not what you want. If this isn't handled, the socket is automatically removed from
 * the list. */
select_closed  *select_set_closed(select_group_t *group, int s, select_closed *callback);

/* Set the timeout callback, for when the time specified in select_group_do_select() elapses. */
select_timeout *select_set_timeout(select_group_t *group, select_timeout *callback, void *param);

/* Remove a socket from the group. Returns non-zero if successful. */
NBBOOL select_group_remove_socket(select_group_t *group, int s);

/* Remove a socket from the group, and close it. */
NBBOOL select_group_remove_and_close_socket(select_group_t *group, int s);

/* Perform the select() call across the various sockets. with the given timeout in milliseconds.
 * Note that the timeout (and therefore the timeout callback) only fires if _every_ socket is idle.
 * If timeout_ms < 0, it will block indefinitely (till data arrives on any socket). Because of polling,
 * on Windows, timeout_ms actually has a resolution defined by TIMEOUT_INTERVAL. */
void select_group_do_select(select_group_t *group, int timeout_ms);

/* Wait for the given number of bytes to arrive on the socket, rather than any number of bytes. This doesn't
 * work for datagram sockets.
 * Note: any data already queued up will be whacked. */
NBBOOL select_group_wait_for_bytes(select_group_t *group, int s, size_t bytes);

/* Check how many active sockets are left. */
size_t select_group_get_active_count(select_group_t *group);

#ifdef WIN32
/* Get a handle to stdin. This handle can be added to a select_group as a pipe. Behind the scenes,
 * it uses a thread. Don't ask. */
HANDLE get_stdin_handle();
#endif

#endif