/*
 * Copyright (c) 2004 Teodor Sigaev <teodor@sigaev.ru>
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 * 1. Redistributions of source code must retain the above copyright
 *        notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *        notice, this list of conditions and the following disclaimer in the
 *        documentation and/or other materials provided with the distribution.
 * 3. Neither the name of the author nor the names of any co-contributors
 *        may be used to endorse or promote products derived from this software
 *        without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY CONTRIBUTORS ``AS IS'' AND ANY EXPRESS
 * OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED. IN NO EVENT SHALL CONTRIBUTORS BE LIABLE FOR ANY
 * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE
 * GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER
 * IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR
 * OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN
 * IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */
#include <stdio.h>
#include <errno.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <fcntl.h>

#ifdef HAVE_POLL_H
#include <poll.h>
#else /* HAVE_POLL */
#ifdef HAVE_SYS_POLL_H
#include <sys/poll.h>
#else
#error Not defined HAVE_POLL_H or HAVE_SYS_POLL_H
#endif /* HAVE_SYS_POLL_H */
#endif /* HAVE_POLL */

#ifdef HAVE_HSTRERROR 
#include <netdb.h>
#endif


#include "connection.h"
#include "tlog.h"
#include "tmalloc.h"

static u_int32_t
setlinger( TC_Connection *cs ) {
	struct linger ling;
	int	val = 0;
	socklen_t size = sizeof(val); 

	if (getsockopt(cs->fd, SOL_SOCKET,SO_ERROR,&val,&size) == -1) {
		tlog(TL_ALARM,"getsockopt: %s:%d - %s(%d)",inet_ntoa(cs->serv_addr.sin_addr),
			ntohs(cs->serv_addr.sin_port), strerror(errno), errno);
		shutdown(cs->fd,SHUT_RDWR);
		close(cs->fd);
		cs->fd = 0;
		cs->state = CS_ERROR;
		return CS_ERROR;
	}

	if ( val ) {
		tlog(TL_ALARM,"getsockopt return: %s:%d - %s(%d)",inet_ntoa(cs->serv_addr.sin_addr),
			ntohs(cs->serv_addr.sin_port), strerror(val), val);
		shutdown(cs->fd,SHUT_RDWR);
		close(cs->fd);
		cs->fd = 0;
		cs->state = CS_ERROR;
		return CS_ERROR;
	}


	ling.l_onoff = ling.l_linger = 0;
	if (setsockopt(cs->fd, SOL_SOCKET,SO_LINGER,(char *)&ling,sizeof(ling))==-1) {
		tlog(TL_ALARM,"setsockopt: LINGER %s:%d - %s",inet_ntoa(cs->serv_addr.sin_addr), ntohs(cs->serv_addr.sin_port),
			strerror(errno));
		shutdown(cs->fd,SHUT_RDWR);
		close(cs->fd);
		cs->fd = 0;
		cs->state = CS_ERROR;
		return CS_ERROR;
	}
	cs->state = CS_CONNECTED;
	return CS_CONNECTED;
}

u_int32_t
TC_ClientInitConnection(TC_Connection *cs, char *name, u_int32_t port) {
	int flags, val=1;

	cs = TC_fillConnection(cs, name, port);

	cs->state = CS_OK;
	if ((cs->fd= socket(AF_INET, SOCK_STREAM, 0)) < 0)
		tlog(TL_CRIT|TL_EXIT,"socket4: %s:%d - %s",inet_ntoa(cs->serv_addr.sin_addr),
			ntohs(cs->serv_addr.sin_port),strerror(errno));

        if (setsockopt(cs->fd, SOL_SOCKET, SO_REUSEADDR, &val, sizeof(val)) < 0) {
                tlog(TL_CRIT|TL_EXIT, "socketsockopt failed: %s (%d)", strerror(errno), errno);
                return CS_ERROR;
        }

	if ((flags=fcntl(cs->fd,F_GETFL,0)) == -1)
		tlog(TL_ALARM,"fcntl F_GETFL - %s",strerror(errno));
	if (fcntl(cs->fd,F_SETFL,flags|O_NDELAY) < 0 )
		tlog(TL_ALARM,"fcntl O_NDELAY - %s",strerror(errno));

	if (bind(cs->fd, (struct sockaddr *) &(cs->serv_addr), sizeof(cs->serv_addr)) < 0)
		tlog(TL_CRIT|TL_EXIT, "cannot bind to %s address: %s",
			inet_ntoa(cs->serv_addr.sin_addr), strerror(errno));
	
	if (listen(cs->fd, 0) < 0)
		tlog(TL_CRIT|TL_EXIT, "cannot listen to %s address: %s",
			inet_ntoa(cs->serv_addr.sin_addr), strerror(errno));
	
	return CS_OK;
}

TC_Connection*
TC_AcceptTcp(TC_Connection *cs) {
	TC_Connection *nc;
	struct sockaddr_in cli_addr;
	int ret, flags;
	socklen_t clilen = sizeof(cli_addr);

	cs->state = CS_READ;
	if ( (ret = accept(cs->fd,(struct sockaddr *)&cli_addr, &clilen)) < 0 ) {
		if ( errno == EAGAIN || errno == EWOULDBLOCK )
			return NULL;
		tlog(TL_ALARM,"TC_AcceptTcp: accept: %s", strerror(errno));
		return NULL;
	}
	nc = (TC_Connection*)t0malloc(sizeof(TC_Connection));

	nc->fd = ret;
	if ((flags=fcntl(nc->fd,F_GETFL,0)) == -1)
		tlog(TL_ALARM,"fcntl F_GETFL - %s",strerror(errno));
	if (fcntl(nc->fd,F_SETFL,flags|O_NDELAY) < 0 )
		tlog(TL_ALARM,"fcntl O_NDELAY - %s",strerror(errno));
	memcpy( &(nc->serv_addr), &cli_addr, clilen );
	nc->state = CS_CONNECTED;

	setlinger(nc);
	return nc;
}

TC_Connection *
TC_fillConnection(TC_Connection *sc, char *name, u_int32_t port) {
	if ( !sc ) 
		sc = (TC_Connection *)tmalloc(sizeof(TC_Connection));
	memset(sc, 0, sizeof(TC_Connection));
	sc->serv_addr.sin_family = AF_INET;
	sc->serv_addr.sin_addr.s_addr = (name && *name != '*' ) ? inet_addr(name) : htonl(INADDR_ANY);
	if ( sc->serv_addr.sin_addr.s_addr == INADDR_NONE ) {
		struct hostent *host;

		/*
		 * Can't parse address: it's a DNS Name
		 */
		host = gethostbyname(name);
		if ( host && host->h_addrtype == AF_INET ) {
			memcpy(&sc->serv_addr.sin_addr.s_addr, host->h_addr_list[0], 
				sizeof(&sc->serv_addr.sin_addr.s_addr));
		} else {
			tlog(TL_CRIT,"gethostbyname: %s - %s", name, hstrerror(h_errno));
			sc->state = CS_ERROR;
			return sc;
		}
	}
	
	sc->serv_addr.sin_port = htons(port);
	sc->state = CS_NOTINITED;
	return sc; 
}


u_int32_t
TC_ServerInitConnect( TC_Connection	*cs ) {
	int flags;

	if ( cs->state == CS_ERROR )
		return CS_ERROR;

	if ((cs->fd= socket(AF_INET, SOCK_STREAM, 0)) < 0) {
		tlog(TL_CRIT,"socket4: %s:%d - %s",inet_ntoa(cs->serv_addr.sin_addr),
			ntohs(cs->serv_addr.sin_port),strerror(errno));
		cs->state  = CS_ERROR;
		return	CS_ERROR;
	}

	if ((flags=fcntl(cs->fd,F_GETFL,0)) == -1)
		tlog(TL_ALARM,"fcntl F_GETFL - %s",strerror(errno));
	if (fcntl(cs->fd,F_SETFL,flags|O_NDELAY) < 0 )
		tlog(TL_ALARM,"fcntl O_NDELAY - %s",strerror(errno));

	flags=1;
	if (setsockopt(cs->fd, SOL_SOCKET, SO_REUSEADDR, &flags, sizeof(flags)) < 0)
		tlog(TL_ALARM, "socketsockopt failed: %s (%d)", strerror(errno), errno);

	if ( connect(cs->fd, (struct sockaddr *) &(cs->serv_addr),
		sizeof(struct sockaddr_in)) < 0 ) {
		if ( errno == EINPROGRESS || errno == EALREADY ) {
			cs->state = CS_INPROCESS;
			return CS_INPROCESS; 
		} else if (errno != EISCONN && errno != EALREADY &&
			errno != EWOULDBLOCK && errno != EAGAIN) {
			tlog(TL_DEBUG,"connect: %s:%d - %s",
				inet_ntoa(cs->serv_addr.sin_addr), ntohs(cs->serv_addr.sin_port),
				strerror(errno));
			shutdown(cs->fd,SHUT_RDWR);
			close(cs->fd);
			cs->fd = 0;
		} else {
			tlog(TL_DEBUG,"nonblock connect: %s:%d - %s [%d]",
				inet_ntoa(cs->serv_addr.sin_addr),
				ntohs(cs->serv_addr.sin_port),
				strerror(errno),errno);
		}
		cs->state = CS_ERROR;
		return CS_ERROR;
	}

	cs->state = CS_INPROCESS;
	return TC_ServerConnect( cs, 0 );
}
	

u_int32_t
TC_ServerConnect( TC_Connection *cs, int timeout ) {
	struct pollfd	pfd;
	int ret;

	if ( cs->state != CS_INPROCESS )
		return cs->state;

	pfd.fd = cs->fd;
	pfd.events = POLLOUT;
	pfd.revents = 0;
	ret = poll( &pfd, 1, timeout );
	if ( ret<0 ) {
		tlog( TL_CRIT, "TC_ServerConnect: poll: %s",
			strerror(errno));
		cs->state = CS_ERROR;
		return CS_ERROR;
	} else if ( ret == 0 ) 
		return CS_INPROCESS;

	if ( (pfd.revents & (POLLHUP | POLLNVAL | POLLERR)) ) {
		tlog( TL_CRIT, "TC_ServerConnect: poll return connect error for %s:%d",
			inet_ntoa(cs->serv_addr.sin_addr), ntohs(cs->serv_addr.sin_port));
		cs->state = CS_ERROR;
		return CS_ERROR;
	}

	if ( ! (pfd.revents & POLLOUT) )
		return CS_INPROCESS;


	return setlinger( cs );
}

int
TC_ReadyIO( TC_Connection **cs, int number, int timeout ) {
	struct pollfd   *pfd;
	int ret,i, fdnum=0;

	if ( number==0 || cs ==NULL ) {
		if (timeout<0)
			timeout=1000;
		usleep( timeout * 1000.0 );
		return 0;
	}
	pfd = (struct pollfd*) tmalloc( sizeof(struct pollfd) * number );

	for(i=0; i<number;i++) {
		if ( cs[i]->fd>0 && (cs[i]->state == CS_READ || cs[i]->state == CS_SEND) ) {
			pfd[fdnum].fd = cs[i]->fd;
			pfd[fdnum].events = ( cs[i]->state == CS_READ ) ? POLLIN : POLLOUT;
			pfd[fdnum].revents = 0;
			fdnum++;
		}
		cs[i]->readyio=0;
	}

	if ( fdnum==0 ) {
		tfree(pfd);
		if (timeout<0)
			timeout=1000;
		usleep( timeout * 1000.0 );
		return 0;
	}	
	ret = poll( pfd, fdnum, timeout );
	if ( ret<0 ) {
		tlog( TL_CRIT, "TC_ReadyIO: poll: %s",
			strerror(errno));
		tfree(pfd);
		return 0;
	}

	if ( ret == 0 ) {
		tfree(pfd);
		return 0;
	}

	fdnum=0; ret=0;
	for(i=0; i<number;i++) {
		if ( cs[i]->fd>0 && (cs[i]->state == CS_READ || cs[i]->state == CS_SEND) ) {
			if ( pfd[fdnum].revents & (POLLHUP | POLLNVAL | POLLERR) ) { 
				tlog( TL_ALARM, "TC_ReadyIO: poll return error for %s:%d",
					inet_ntoa(cs[i]->serv_addr.sin_addr), 
					ntohs(cs[i]->serv_addr.sin_port));
				cs[i]->state = CS_ERROR;
				ret = 1;
			} else if ( pfd[fdnum].revents & ( ( cs[i]->state == CS_READ ) ? POLLIN : POLLOUT ) ) {
				cs[i]->readyio=1;
				ret = 1;
			}
			fdnum++;
		}
	}

	tfree(pfd);
	return ret;
}

u_int32_t
TC_Send( TC_Connection *cs ) {
	int sz;
	
	if ( cs->state == CS_ERROR )
		return CS_ERROR;

	if ( cs->state != CS_SEND || cs->ptr == NULL ) {
		cs->state = CS_SEND;
		cs->ptr = (char*)cs->buf;
		cs->len = cs->buf->len;

		/* convert fields to network byteorder */
		cs->buf->len = htonl(cs->buf->len);
		cs->buf->type = htonl(cs->buf->type);
	}

	if ( cs->ptr - (char*)cs->buf >= cs->len ) {
		cs->state = CS_FINISHSEND;
		return CS_FINISHSEND;
	}

	if ((sz=write(cs->fd, cs->ptr, cs->len - (cs->ptr - (char*)cs->buf)))==0 ||
		(sz < 0 && (errno == EWOULDBLOCK || errno == EAGAIN))) {

		/* SunOS 4.1.x, are broken and select() says that
		 * O_NDELAY sockets are always writable even when
		 * they're actually not.
     		 */
		cs->state = CS_SEND;
		return CS_SEND;
	}
	if ( sz<0 ) {
		if (errno != EPIPE && errno != EINVAL)
			tlog(TL_ALARM, "write[%s:%d] - %s",
				inet_ntoa(cs->serv_addr.sin_addr),
				ntohs(cs->serv_addr.sin_port), 
				strerror(errno));
		cs->state = CS_ERROR;
		return CS_ERROR;
	}

	cs->ptr += sz;

	if ( cs->ptr - (char*)cs->buf >= cs->len ) {
		cs->state = CS_FINISHSEND;
		/* revert byteorder conversion */
		cs->buf->len = ntohl(cs->buf->len);
		cs->buf->type = ntohl(cs->buf->type);
		return CS_FINISHSEND;
	}
	
	return cs->state;
}

static void 
resizeCS( TC_Connection *cs, int sz ) {
	int diff = cs->ptr - (char*)cs->buf;

	if ( cs->len >= sz )
		return; 
	cs->len = sz;
	cs->buf = (TCMsg*)trealloc( (void*)cs->buf, cs->len );
	cs->ptr = ((char*)cs->buf) + diff;
}

u_int32_t
TC_Read( TC_Connection *cs, size_t maxsize ) {
	int sz, totalread = -1, toread=0, alreadyread;

	if ( cs->state == CS_ERROR )
		return CS_ERROR;

	if (cs->state != CS_READ || cs->ptr == NULL ) {
		cs->state = CS_READ;
		cs->ptr = (char*)cs->buf;
		cs->len = 0;
	}

	alreadyread = cs->ptr - (char*)cs->buf;
	if ( alreadyread < TCMSGHDRSZ ) {
		toread = TCMSGHDRSZ - alreadyread;
		resizeCS(cs, TCMSGHDRSZ);
	} else {
		totalread = ntohl(cs->buf->len);
		if ( maxsize > 0 && totalread > maxsize )
		{
			tlog(TL_ALARM,"TC_Read: message size (%d b) is greater than max allowed (%d b)", totalread, maxsize);
			cs->state = CS_ERROR;
			return CS_ERROR;
		}
		toread = totalread - alreadyread;
		if ( toread == 0 ) {
			cs->state = CS_FINISHREAD;
			return CS_FINISHREAD;
		}
		resizeCS(cs, totalread);
	}

	if ((sz=read( cs->fd, cs->ptr, toread))<0) {
		if (errno == EAGAIN || errno == EINTR) {
			cs->state = CS_READ;
			return CS_READ;
		}
		tlog(TL_ALARM,"read: finish - %s",strerror(errno));
		cs->state = CS_ERROR;
		return CS_ERROR;
	}
	
	if ( alreadyread < TCMSGHDRSZ  && alreadyread + sz >= TCMSGHDRSZ ) {
		/* 
		 * we just read header - we can get totalread value. 
		 */
		totalread = ntohl(cs->buf->len);
	}
	
	cs->ptr += sz;
	alreadyread += sz;
	if ( sz == 0 && alreadyread != totalread ) {
		tlog(TL_ALARM,"read: disconnecting");
		cs->state = CS_ERROR;
		return CS_ERROR;
	}

	if ( alreadyread == totalread ) {
		cs->buf->len = ntohl(cs->buf->len);
		cs->buf->type = ntohl(cs->buf->type);
		cs->state = CS_FINISHREAD;
	}

	return cs->state;
}

void
TC_FreeConnection( TC_Connection *cs ) {
	if ( cs->state == CS_CLOSED )
		return;
	if ( cs->buf ) {
		tfree(cs->buf);
		cs->buf = NULL;
	}
	if ( cs->fd && cs->state != CS_NOTINITED ) {
		shutdown(cs->fd,SHUT_RDWR);
		close(cs->fd);
	}
	cs->fd = 0;
	cs->state = CS_CLOSED;
}

u_int32_t 
TC_Talk( TC_Connection *cs, size_t maxsize  ) {
	if ( cs->state==CS_NOTINITED ) 
		TC_ServerInitConnect( cs );

	while( cs->state == CS_INPROCESS ) 
		TC_ServerConnect(cs, 100);

	if ( cs->state != CS_CONNECTED )
		return cs->state;
	
	cs->state = CS_SEND;
	cs->ptr = NULL;
	while( cs->state != CS_FINISHSEND ) {
		while( !TC_ReadyIO( &cs, 1, 100) );
		if ( TC_Send(cs) == CS_ERROR ) return CS_ERROR;
	}

	cs->state = CS_READ;
	cs->ptr = NULL;
	while( cs->state != CS_FINISHREAD ) {
		while( !TC_ReadyIO( &cs, 1, 100) );
		if ( TC_Read(cs, maxsize) == CS_ERROR ) return CS_ERROR;
	}

	return CS_OK; 
}