/* Socket I/O routines. */

#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <errno.h>
#include <fcntl.h>
#include <netinet/in.h>
#include <sys/socket.h>
#include <sys/time.h>

#include "socketio.h"

/*************************************************************************/

/* Array of socket info structures */
static Socket *sockets = NULL;
static int socketlen = 0;

/*************************************************************************/
/*************************************************************************/

/* Flush as much buffered write data as possible to the socket.  Return -1
 * if an error occurred.
 */

static int sock_flush(int sock)
{
    fd_set fds;
    struct timeval tv = {0,0};
    int n;
    Socket *s = &sockets[sock];

    if (!(s->status & SOCKSTAT_OPEN))
	return 0;
    FD_ZERO(&fds);
    FD_SET(sock, &fds);
    while (s->wbuflen > 0 && select(sock+1, NULL, &fds, NULL, &tv) > 0) {
	n = write(sock, s->wbuf, s->wbuflen);
	if (n <= 0 && errno != EAGAIN && errno != EINTR)
	    return -1;
	if (n > 0) {
	    if (n < s->wbuflen)
		memmove(s->wbuf, s->wbuf + n, s->wbuflen - n);
	    s->wbuflen -= n;
	}
    }
    return 0;
}

/*************************************************************************/

/* Try to establish a connection to a remote host.  Call the specified
 * callback routine when the connection either succeeds or fails.
 */

void sock_connect(int sock, struct sockaddr_in *sin, socket_callback callback)
{
    int res;
    Socket *s;

    /* This, by the way, also checks whether the socket number is valid, so
     * we don't accidentally go off and try to allocate a weird number of
     * bytes of memory. */
    if (fcntl(sock, F_SETFL, O_NONBLOCK) < 0) {
	callback(sock, 0);
	return;
    }

    if (socketlen <= sock) {
	sockets = realloc(sockets, sizeof(Socket) * (sock+1));
	memset(sockets+socketlen, 0, sizeof(Socket) * (sock-socketlen+1));
	socketlen = sock+1;
    }
    s = &sockets[sock];
    if (s->status & SOCKSTAT_OPEN) {
	callback(sock, 1);
	return;
    }

    if (!(s->status & SOCKSTAT_CONNECTING)) {
	s->status = SOCKSTAT_CONNECTING;
	s->rbuf = NULL;
	s->rbuflen = s->rbufsize = 0;
	s->connect_callback = callback;
    }

    res = connect(sock, (struct sockaddr *)sin, sizeof(*sin));
    if (res == 0) {
	s->status = SOCKSTAT_OPEN;
	sock_flush(sock);	/* Flush any accumulated writes */
	callback(sock, SOCKSTAT_OPEN);
    } else if (res < 0 && errno != EINPROGRESS && errno != EALREADY
		&& errno != EAGAIN && errno != EINTR) {
	s->status = 0;
	callback(sock, 0);
    }
}

/*************************************************************************/

/* Close the given socket.  This routine should be preferred over close()
 * because it frees any buffer memory associated with the socket.
 */

void sock_close(int sock)
{
    Socket *s = &sockets[sock];

    if (sock < 0 || sock >= socketlen ||
			!(s->status & (SOCKSTAT_OPEN | SOCKSTAT_CONNECTING)))
	return;
    if (s->rbufsize) {
	free(s->rbuf);
	s->rbuf = NULL;
	s->rbufsize = s->rbuflen = 0;
    }
    if (s->wbufsize) {
	free(s->wbuf);
	s->wbuf = NULL;
	s->wbufsize = s->wbuflen = 0;
    }
    s->status = 0;
    close(sock);
}

/*************************************************************************/

/* Read some data from the given socket.  Return 0 and errno==EAGAIN if the
 * requested amount of data is not yet available.
 */

int sock_read(int sock, char *buf, int len)
{
    Socket *s = &sockets[sock];
    fd_set fds;
    struct timeval tv = {0,0};
    int n;

    if (sock < 0 || sock >= socketlen ||
			!(s->status & (SOCKSTAT_OPEN | SOCKSTAT_CONNECTING))) {
	errno = EBADF;
	return -1;
    }
    if (len == 0) {
	errno = 0;
	return 0;
    }

    if (sock_flush(sock) < 0)
	return -1;

    if (len > s->rbufsize) {
	if (!(s->rbuf = realloc(s->rbuf, len))) {
	    s->rbuf = NULL;
	    s->rbufsize = s->rbuflen = 0;
	    errno = ENOMEM;
	    return -1;
	}
	s->rbufsize = len;
    }

    FD_ZERO(&fds);
    FD_SET(sock, &fds);
    while (s->rbuflen < len && select(sock+1,&fds,0,0,&tv) > 0) {
	errno = 0;
	n = read(sock, s->rbuf + s->rbuflen,
	               s->rbufsize - s->rbuflen);
	if (n < 0 || (n == 0 && errno != EAGAIN && errno != EINTR)) {
	    int errno_save = errno;
	    s->status = 0;
	    free(s->rbuf);
	    s->rbuf = NULL;
	    s->rbufsize = s->rbuflen = 0;
	    errno = errno_save;
	    return -1;
	}
	if (n == 0)
	    break;
	s->rbuflen += n;
    }
    if (s->rbuflen < len) {
	errno = EAGAIN;
	return 0;
    }
    memcpy(buf, s->rbuf, len);
    if (s->rbuflen > len)
	memmove(s->rbuf, s->rbuf + len, s->rbuflen - len);
    s->rbuflen -= len;
    return len;
}

/*************************************************************************/

/* Write some data to the requested socket.  Always succeeds, barring
 * out-of-memory conditions or other such failures.
 */

int sock_write(int sock, const char *buf, int len)
{
    Socket *s = &sockets[sock];

    if (sock < 0 || sock >= socketlen ||
			!(s->status & (SOCKSTAT_OPEN | SOCKSTAT_CONNECTING))) {
	errno = EBADF;
	return -1;
    }

    if (sock_flush(sock) < 0)
	return -1;
    if (s->wbuflen + len > s->wbufsize) {
	s->wbufsize = s->wbuflen + len;
	if (!(s->wbuf = realloc(s->wbuf, s->wbufsize))) {
	    s->wbuflen = s->wbufsize = 0;
	    errno = ENOMEM;
	    return -1;
	}
    }
    memcpy(s->wbuf + s->wbuflen, buf, len);
    s->wbuflen += len;
    if (sock_flush(sock) < 0)
	return -1;
    return len;
}

/*************************************************************************/