/* net.c : PUBLIC DOMAIN - Jon Mayo - August 5, 2006
 * - You may remove any comments you wish, modify this code any way you wish,
 *   and distribute any way you wish.*/
/* generic networking routines */

#include <assert.h>
#include <arpa/inet.h>
#include <errno.h>
#include <fcntl.h>
#include <poll.h>
#include <netdb.h>
#include <netinet/in.h>
#include <stdarg.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <unistd.h>
#include "net.h"

static struct socket_info **socket_list;
static struct pollfd *sys_fds;
static unsigned sys_fds_max, sys_fds_cnt, nr_socket_list;

static unsigned grow_sys_fds(void) {
	unsigned new_max;
	if((sys_fds_cnt+1)>sys_fds_max) {
		new_max=sys_fds_cnt+64;
		sys_fds=realloc(sys_fds, sizeof *sys_fds * new_max);
		for(;sys_fds_max<new_max;sys_fds_max++) {
			memset(&sys_fds[sys_fds_max], 0, sizeof *sys_fds);
		}
	}
	return sys_fds_cnt++;
}

static void sys_fds_set(unsigned idx, int fd, unsigned events) {
	sys_fds[idx].fd=fd;
	sys_fds[idx].events=events;
	sys_fds[idx].revents=0;
}

static void grow_socket_list(unsigned fd) {
	unsigned new_max;
	if(nr_socket_list<=fd) {
		new_max=(fd|63)+1; /* round up to next 64 */
		socket_list=realloc(socket_list, sizeof *socket_list * new_max);

		memset(socket_list+nr_socket_list, 0, (new_max-nr_socket_list)*sizeof *socket_list);
		nr_socket_list=new_max;
	}
}

struct socket_info *socket_insert(int fd, short ev, void (*handler)(int fd, short ev, struct socket_info *si)) {
	struct socket_info *new_socket;
	new_socket=malloc(sizeof *new_socket);
	if(!new_socket) {
		perror("malloc()");
		return 0;
	}
	grow_socket_list(fd);
	new_socket->fd=fd;
	new_socket->pollent=grow_sys_fds();
	new_socket->nh.io_handler=handler;
	new_socket->nh.connect_handler=0;
	new_socket->output_len=0;
	new_socket->output_max=0;
	new_socket->output_buffer=NULL;
	new_socket->output_overflow=0;
	new_socket->userdata=0;
	new_socket->userdata_free=0;
	sys_fds_set(new_socket->pollent, fd, ev);
	fprintf(stderr, "inserting %p at %d\n", (void*)new_socket, fd);
	socket_list[fd]=new_socket;
	return new_socket;
}

static void dump_addrinfo(struct addrinfo *ai) {
	printf(
	"[addrinfo:%p]\n"
	"  ai_flags=%X\n"
	"  ai_family=%u\n"
	"  ai_socktype=%u\n"
	"  ai_protocol=%u\n"
	"  ai_addrlen=%u\n"
	"  ai_addr=%p\n"
	"  ai_canonname=%s\n"
	"  ai_next=%p\n",
	(void*)ai,
	ai->ai_flags,
	ai->ai_family,
	ai->ai_socktype,
	ai->ai_protocol,
	ai->ai_addrlen,
	(void*)ai->ai_addr,
	ai->ai_canonname,
	(void*)ai->ai_next);
}

int net_lookup(const char *host, const char *sport, const char *pname, int (*foreach)(void *p, struct addrinfo *ai, const char *addr, const char *serv), void *p) {
	struct protoent *pe;
	struct addrinfo *ai=0, hint, *curr;
	int proto;
	int res;
	assert(foreach != NULL);
	fprintf(stderr, "net_lookup(\"%s\", \"%s\", \"%s\")\n", host, sport, pname);
	if(pname && pname[0]) {
		pe=getprotobyname(pname);
		proto=pe ? pe->p_proto :  0;
	} else {
		proto=0;
	}
	if(host && !host[0]) host=0;
	if(sport && !sport[0]) sport=0;

	hint.ai_flags=AI_ADDRCONFIG|AI_PASSIVE;
	hint.ai_family=PF_UNSPEC;
	hint.ai_protocol=proto;
	hint.ai_socktype=0;
	hint.ai_addrlen=0;
	hint.ai_addr=0;
	hint.ai_canonname=0;
	hint.ai_next=0;
	
	if((res=getaddrinfo(host, sport, &hint, &ai))) {
		fprintf(stderr, "getaddrinfo(): %s\n", gai_strerror(res));
		if(res==EAI_SYSTEM) {
			perror("getaddrinfo()");
		}
		return 0;
	}
	curr=ai;
	if(!curr) {
		fprintf(stderr, "no entries!\n"); 
		return 0; 
	}

	for(; curr; curr=curr->ai_next) {
		char addr[NI_MAXHOST];
		char serv[NI_MAXSERV];
		int res;
		/* Mac OS X 10.4 getnameinfo() is buggy without NI_NUMERICSERV */
		if((res=getnameinfo(curr->ai_addr, curr->ai_addrlen, addr, sizeof addr, serv, sizeof serv, NI_NUMERICHOST|NI_NUMERICSERV))) {
			/* Failed */
			if(res==EAI_SYSTEM) {
				perror("getnameinfo()");
			} else {
				fprintf(stderr, "getnameinfo(): %s\n", gai_strerror(res));
			}
			continue;
		}
		if(!foreach(p, curr, addr, serv)) {
			freeaddrinfo(ai);
			return 0;
		}
	}
	freeaddrinfo(ai);
	return 1;
}

static void accept_handler(int fd, short ev, struct socket_info *si) {
	struct sockaddr_storage ss;
	socklen_t sslen;
	int connfd;
	int res;
	char addr[NI_MAXHOST], serv[NI_MAXSERV];
	fprintf(stderr, "%s()\n", __func__);
	sslen=sizeof ss;
	connfd=accept(fd, (struct sockaddr*)&ss, &sslen);
	if(connfd<0) {
		perror("accept()");
		return;
	}
	/* Mac OS X 10.4 getnameinfo() is buggy without NI_NUMERICSERV */
	if((res=getnameinfo((struct sockaddr*)&ss, sslen, addr, sizeof addr, serv, sizeof serv, NI_NUMERICHOST|NI_NUMERICSERV))) {
		fprintf(stderr, "getnameinfo(): %s\n", gai_strerror(res));
		if(res==EAI_SYSTEM) { 
			perror("getnameinfo()"); 
		} 
		return;
	}
	if(si->nh.connect_handler) {
		si->nh.connect_handler(si->nh.connect_p, connfd, addr, serv);
	}
}

void net_clear_event(struct socket_info *si, unsigned flag) {
	sys_fds[si->pollent].events&=~flag;
}


int bind_entry(void *p, struct addrinfo *ai, const char *addr, const char *serv) {
	struct net_handler *nh=p;
	struct socket_info *si;
	int s;

	fprintf(stderr, "binding: %s/%s\n", addr, serv);
	s=socket(ai->ai_family, ai->ai_socktype, ai->ai_protocol);
	if (s < 0) {
		perror("socket()");
		return 0; /* failure */
	}
	if(bind(s,(struct sockaddr*)ai->ai_addr, ai->ai_addrlen)) {
		perror("bind()");
		close(s);
		return 0; /* failure */
	}
	listen(s, LISTEN_QUEUE);
	fcntl(s, F_SETFL, O_NONBLOCK);
	if(!(si=socket_insert(s, NET_READ_READY, accept_handler))) {
		close(s);
		return 0; /* failure */
	}
	si->nh.connect_handler=nh->connect_handler;
	si->nh.connect_p=nh->connect_p;
	return 1; /* success */
}

int connect_entry(void *p, struct addrinfo *ai, const char *addr, const char *serv) {
	struct net_handler *nh=p;
	struct socket_info *si;
	int s;
	fprintf(stderr, "connecting: %s/%s\n", addr, serv);
	assert(nh->io_handler != NULL);
	s=socket(ai->ai_family, ai->ai_socktype, ai->ai_protocol);
	if (s < 0) {
		perror("socket()");
		return 0; /* failure */
	}
	/* TODO : attach the event initially to wait for a connect, then connect and run the real event */
	/* fcntl(s, F_SETFL, O_NONBLOCK); */

	if(connect(s,(struct sockaddr*)ai->ai_addr, ai->ai_addrlen)) {
		perror("connect()");
		close(s);
		return 0; /* failure */
	}
	if(!(si=socket_insert(s, NET_READ_READY, nh->io_handler))) {
		close(s);
		return 0; /* failure */
	}
	printf("client connect_p = %p\n",  nh->connect_p);
	if(nh->connect_handler)
		nh->connect_handler(nh->connect_p, s, addr, serv);
	return 1; /* success */
}

int net_bind(const char *host, const char *sport, const char *pname, void (*handler)(void *p, int fd, const char *addr, const char *serv), void *p) {
	struct net_handler nh;
	nh.connect_handler=handler;
	nh.connect_p=p;
	return net_lookup(host, sport, pname, bind_entry, &nh);
}

int net_connect(const char *host, const char *sport, const char *pname, void (*connect_handler)(void *p, int fd, const char *addr, const char *serv), void *connect_p, void (*io_handler)(int fd, short ev, struct socket_info *si), void *io_p) {
	struct net_handler nh;
	nh.io_handler=io_handler;
	nh.io_p=io_p;
	nh.connect_handler=connect_handler;
	nh.connect_p=connect_p;
	fprintf(stderr, "net_connect(\"%s\", \"%s\", \"%s\")\n", host, sport, pname);
	return net_lookup(host, sport, pname, connect_entry, &nh);
}

/* call this to clean sys_fds entry when the sock_list entry is NULL */
static void scrub_fd(unsigned fd, unsigned pollent) {
	assert(pollent < sys_fds_cnt);
	assert(sys_fds_cnt > 0);
	fprintf(stderr, "fd=%u: scrubbing\n", fd);
	close(fd);
	/* move last entry into this entry */
	sys_fds[pollent]=sys_fds[--sys_fds_cnt];
	fd=sys_fds[pollent].fd;
	/* touch up the moved entry's socket_list */
	if(fd<nr_socket_list && socket_list[fd]) {
		socket_list[fd]->pollent=pollent;
	}
}

/* process connections. return 0 when there are no sockets left to poll() */
int net_poll(void) {
	int r, i;
	again:
	if(!sys_fds_cnt) return 0;
	assert(sys_fds_cnt>0);
	r=poll(sys_fds, sys_fds_cnt, -1);
	if(r<0) {
		perror("poll()");
		if(errno==EINTR) goto again;
		exit(EXIT_FAILURE);
	}
	for(i=0;r>0 && i<sys_fds_cnt;) {
		int fd;
		fd=sys_fds[i].fd;
		
		/* look for a dead connection */
		if(fd>nr_socket_list || !socket_list[fd]) {
			scrub_fd(fd, i);
			continue; /* try entry i again */
		}

		if(sys_fds[i].revents) {
			r--;
			/* TODO: socket had something to do */
			if(fd<nr_socket_list && socket_list[fd]) {
				fprintf(stderr, "fd=%d has something to do. revents=%X\n", fd, sys_fds[i].revents);
				if(socket_list[fd]->nh.io_handler) 
					socket_list[fd]->nh.io_handler(fd, sys_fds[i].revents, socket_list[fd]);
				/* TODO: should we clean the revents ? */
				if(fd>nr_socket_list || !socket_list[fd]) {
					/* we were deleted */
					scrub_fd(fd, i);
					continue; /* try entry i again */
				}
			} else {
				fprintf(stderr, "fd=%d activated, but nothing is registered to handle it.\n", fd);
			}
		}
		i++; /* we don't want this ran as part of continue */
	}
	return 1;
}

/* assumes fd points to a valid socket_list entry */
void sockkill(struct socket_info *si) {
	int fd=si->fd;
	assert(fd<nr_socket_list);
	assert(socket_list[fd]);
	/* try and improve the odds that the second pass on poll will find the
	 * dead connection. */
	sys_fds[si->pollent].events|=POLLIN|POLLOUT|POLLHUP|POLLERR;
	if(socket_list[fd]->userdata_free) 
		socket_list[fd]->userdata_free(socket_list[fd]->userdata);
	free(socket_list[fd]->output_buffer);
	free(socket_list[fd]);
	socket_list[fd]=0;
	printf("fd=%d: close\n", fd);
}

/* writes to a socket buffer, don't truncate */
int sockwrite(struct socket_info *si, const void *buf, unsigned len) {
	int left;

	assert(si != NULL);
	left=(int)si->output_max-(int)si->output_len;
	if(left < len) {
		fprintf(stderr, "fd=%u: output overflow!\n", si->fd);
		if(si->output_overflow) {
			si->output_overflow(si);
		} else {
			/* if no overflow handler blow it away */
			sockkill(si);
		}
		return 0;
	}
	memcpy(si->output_buffer+si->output_len, buf, len);
	si->output_len+=len;
	sys_fds[si->pollent].events|=POLLOUT; /* mark as write ready */
	return len;
}

int sockvprintf(struct socket_info *si, const char *fmt, va_list ap) {
	int len;
	assert(si != NULL);

	/* overflow check */
	if(si->output_len >= si->output_max) {
		fprintf(stderr, "fd=%u: output overflow!\n", si->fd);
		if(si->output_overflow) {
			si->output_overflow(si);
		} else {
			/* if no overflow handler blow it away */
			sockkill(si);
		}
		return 0;
	}
	len=vsnprintf(si->output_buffer+si->output_len, si->output_max-si->output_len, fmt, ap);
	fprintf(stderr, "fd=%u: sending \"%.*s\"\n", si->fd, len, si->output_buffer+si->output_len);
	si->output_len+=len;
	sys_fds[si->pollent].events|=POLLOUT; /* mark as write ready */
	return len;
}

int sockprintf(struct socket_info *si, const char *fmt, ...) {
	int len;
	va_list ap;
	va_start(ap, fmt);
	len=sockvprintf(si, fmt, ap);
	va_end(ap);
	return len;
}

struct socket_info *socket_info_fd(int fd) {
	assert(fd >= 0);
	return (fd<nr_socket_list) ? socket_list[fd] : 0;
}

void socket_info_setuserdata(struct socket_info *si, void *userdata, void (*userdata_free)(void *)) {
	si->userdata=userdata;
	si->userdata_free=userdata_free;
}

int net_get_nr_sockets(void) {
	return nr_socket_list;
}

int socket_output_buffer_size(struct socket_info *si, unsigned newsize) {
	char *tmp;

	tmp=realloc(si->output_buffer, newsize);
	if(!tmp) return 0; /* failure */
	si->output_buffer=tmp;
	si->output_max=newsize;
	if(si->output_len>si->output_max) {
		/* truncated buffer */
		si->output_len=si->output_max;
	}
	return 1; /* success */
}

#ifdef STAND_ALONE
static int dump_entry(void *p, struct addrinfo *ai, const char *addr, const char *serv) {
	fprintf(stderr, "family: %u\tsocktype: %u\tprotocol: %u\taddr: %s/%s\n", ai->ai_family, ai->ai_socktype, ai->ai_protocol, addr, serv);
	return 1;
}

static void data_handler(int fd, short ev, struct socket_info *si) {
	size_t len;
	if(ev&NET_READ_READY) {
		char buf[128];
		len=read(fd, buf, sizeof buf);
		if(len<0) {
			/* close on an error */
			perror("read()");
			sockkill(si);
			return;
		} 
		if(len==0) {
			sockkill(si);
			return;
		}
		printf("fd=%d: %u bytes read\n", fd, len);
	}
	if(ev&NET_WRITE_READY) {
		len=write(fd, si->output_buffer, si->output_len);
		if(len<0) {
			/* close on an error */
			perror("write()");
			sockkill(si);
			return;
		}
		printf("fd=%d: %u bytes written\n", fd, len);

		si->output_len-=len;
		if(si->output_len>0)
			memmove(si->output_buffer, si->output_buffer+len, si->output_len);
		else /* clear the write bit */
			sys_fds[si->pollent].events&=~POLLOUT;
	}
}

static void client_handler(void *p, int fd, const char *addr, const char *serv) {
	struct socket_info *si;
	si=socket_info_fd(fd);
	assert(si != NULL);
	fprintf(stderr, "client: %s %s\n", addr, serv);
	if(p) sockprintf(si, "%s\r\n", p);
}

static void connect_handler(void *p, int fd, const char *addr, const char *serv) {
	struct socket_info *si;
	fprintf(stderr, "connect: %s %s\n", addr, serv);
	/* TODO: do something */
	si=socket_insert(fd, NET_READ_READY, data_handler);
	if(!si) {
		close(fd);
		/* failure */
	}
	socket_output_buffer_size(si, 4072);
}

static void usage(const char *argv0) {
	printf("usage: %s\n"
	"    -P <proto>  sets the protocol\n"
	"    -t <secs>   timeout after s seconds\n"
	"    -c          use client mode\n"
	"    -n <cnt>    number of connections for the client to perform\n"
	"    -s          use server mode\n"
	"    -m <msg>    sent message on connect\n"
	, argv0);
	exit(EXIT_FAILURE);
}

int main(int argc, char **argv) {
	int c;
	int dump_mode=0, client_mode=0, server_mode=0, timeout=0, client_connect=10;
	const char *message_str=0;
	const char *sport="5800", *proto="tcp";

	while((c=getopt(argc, argv, "dcshP:p:n:m:t:"))!=-1) {
		switch(c) {
			case 'P': proto = optarg; break;
			case 'p': sport = optarg; break;
			case 'd': dump_mode=1; break;
			case 'c': client_mode=1; break;
			case 's': server_mode=1; break;
			case 't': timeout=strtoul(optarg, 0, 0); break;
			case 'n': client_connect=strtoul(optarg, 0, 0); break;
			case 'm': message_str=optarg; break;
			case 'h':
			default:
				usage(argv[0]);
		}
	}
	
	/* if two or more are set */
	if((client_mode+server_mode+dump_mode) > 1) usage(argv[0]);	
	if(!client_mode && !server_mode && !dump_mode) server_mode=1;
	
	if(timeout>0) {
		fprintf(stderr, "Timeout in %d seconds\n", timeout);
		alarm(timeout);
	}
	if(dump_mode) {
		int i;
		if(optind==argc) {
			net_lookup(NULL, sport, proto, dump_entry, 0);
		} else {
			for(i=optind;i<argc;i++) {
				net_lookup(argv[i], sport, proto, dump_entry, 0);
			}
		}
		return 0;
	} else if(client_mode) {
		/* client_mode */
		int i,j;
		for(j=0;j<client_connect;j++) {
			if(optind==argc) {
				net_connect("localhost", sport, proto, client_handler, message_str, data_handler, 0);
			} else {
				for(i=optind;i<argc;i++) {
					net_connect(argv[i], sport, proto, client_handler, message_str, data_handler, 0);
				}
			}
		}
	} else {
		/* server_mode */
		int i;
		if(optind==argc) {
			net_bind(NULL, sport, proto, connect_handler, 0);
		} else {
			for(i=optind;i<argc;i++) {
				net_bind(argv[i], sport, proto, connect_handler, 0);
			}
		}
	}
	
	while(1) net_poll();
	return 0;
}
#endif