/*  Class for handling sockets in the Internet domain.
 *
 *  Copyright (c) 1994 Christopher J. Kane.
 *
 *  This software is subject to the terms of the MiscKit license
 *  agreement. Refer to the license document included with the
 *  MiscKit distribution for these terms.
 *
 *  Version 1.0 BETA (26 April 1994)
 */

#import <misckit/MiscINETSocket.h>
#import <objc/HashTable.h>
#import <libc.h>
#import <netdb.h>
#import <netinet/in_systm.h>
#import <netinet/ip.h>
#import <netinet/ip_icmp.h>
#import <netinet/tcp.h>
#import <sys/errno.h>
extern int errno;

static MiscINETSocket *handySocket = nil;

static unsigned short in_cksum(unsigned char *cp, int nbytes)
{
	unsigned short wa[(nbytes>>1)+1], *wp, result;
	long sum;
	wp = &wa[0];
	memset((char *)wp, 0, sizeof(wp));
	memcpy((char *)wp, cp, nbytes);
	for (sum = 0; 0 < nbytes; nbytes -= 2)
		sum += *wp++;
	while (sum >> 16)
		sum = (sum >> 16) + (sum & 0xffff);
	result = (unsigned short)~sum;
	if (result == 0x0)
		result = 0xffff;
	return result;
}

static int wait_for_socket_read(int sock, unsigned int ms)
{
	struct timeval timeout = {ms/1000, (ms%1000)*1000};
	fd_set rfds;
	int ret;
	FD_ZERO(&rfds);
	FD_SET(sock, &rfds);
	while ((ret = select(sock+1, &rfds, NULL, NULL, (ms==0 ? NULL : &timeout))) < 0) {
		if (errno == EINTR)
			continue;
		return -1;
	}
	return ret;
}

static int get_service_port(const char *service, int type)
{
	struct servent *sent;
	switch (type) {
	case MiscSOCK_DGRAM :
		sent = getservbyname((char *)service, "udp");
		if (sent == NULL)
			sent = getservbyport(atoi(service), "udp");
		break;
	case MiscSOCK_STREAM:
		sent = getservbyname((char *)service, "tcp");
		if (sent == NULL)
			sent = getservbyport(atoi(service), "tcp");
		break;
	default    :
		/* RAW connections to services are not supported. */
		errno = ESOCKTNOSUPPORT;
		return -1;
	}
	if (sent == NULL) {
		errno = ENOENT;
		return -1;
	}
	return sent->s_port;
}

@implementation MiscINETSocket

#define ECHOPKTLEN 32

+ (int)ping:(MiscINETAddress *)addr timeout:(unsigned int)ms useECHO:(BOOL)aBool
{
	if (addr == nil) {
		errno = EDESTADDRREQ;
		return -1;
	}
	if (aBool) {
		char out_pkt[ECHOPKTLEN], in_pkt[2*ECHOPKTLEN];
		int i, ret, in_len = sizeof(in_pkt);
		for (i = 0; i < ECHOPKTLEN; i++)
			out_pkt[i] = (char)(random()%256);
		ret = [MiscINETSocket sendDgram:out_pkt length:sizeof(out_pkt) to:addr service:"echo" timeout:ms retries:0 withReply:in_pkt length:&in_len];
		if (ret <= 0)
			return ret;
		if (in_len != ECHOPKTLEN)
			return -1;
		for (i = 0; i < ECHOPKTLEN; i++)
			if (out_pkt[i] != in_pkt[i])
				return -1;
		return 1;
	} else {
		struct {
			unsigned char type, code;
			unsigned short cksum, id, seq;
			unsigned char data[2*ICMP_MINLEN];
		} out_pkt = {ICMP_ECHO, 0, 0, getpid() & 0xffff, 1};
		struct sockaddr_in sockaddr;
		char buffer[sizeof(out_pkt)+64];
		int ret;
		struct icmp *icmpp;
		if (handySocket == nil)
			handySocket = [MiscINETSocket alloc];
		handySocket = [handySocket init];
		handySocket->sock = socket(PF_INET, SOCK_RAW, IPPROTO_ICMP);
		if (handySocket->sock < 0)
			return -1;
		handySocket->domain = PF_INET;
		handySocket->type = RAW;
		out_pkt.cksum = in_cksum((u_char *)&out_pkt, sizeof(out_pkt));
		sockaddr.sin_family = AF_INET;
		sockaddr.sin_addr = [addr address];
		sockaddr.sin_port = 0;
		if (sendto(handySocket->sock, &out_pkt, sizeof(out_pkt), 0, (struct sockaddr *)&sockaddr, sizeof(sockaddr)) < 0)
			return -1;
		ret = wait_for_socket_read(handySocket->sock, ms);
		if (ret <= 0)
			return ret;
		memset(buffer, '\0', sizeof(buffer));
		if (recv(handySocket->sock, buffer, sizeof(buffer), 0) < 0)
			return -1;
		icmpp = (struct icmp *)(buffer + 4*((struct ip *)buffer)->ip_hl);
		if (icmpp->icmp_type == ICMP_ECHOREPLY && icmpp->icmp_id == (getpid() & 0xffff))
			return 1;
		return -1;
	}
}

+ runServer:(id)sockets target:(id)target action:(SEL)action fork:(BOOL)fFlag loop:(BOOL)lFlag
{
	BOOL isCollection = NO;
	fd_set fds;
	int fd, nfd;
	if (![target respondsTo:action] || sockets == nil) {
		errno = EINVAL;
		return nil;
	}
	FD_ZERO(&fds);
	if ([sockets isKindOf:[HashTable class]]) {
		Class Socket = [MiscSocket class];
		for (fd = sizeof(fds)*8; fd--;) {
			id obj;
			if (![sockets isKey:(void *)fd])
				continue;
			obj = [sockets valueForKey:(void *)fd];
			if ([obj isKindOf:Socket] && ![obj isClosed])
				FD_SET([obj socket], &fds);
		}
		isCollection = YES;
	} else if ([sockets isKindOf:[MiscSocket class]] && ![sockets isClosed]) {
		FD_SET([sockets socket], &fds);
	} else {
		return nil;
	}
	for (fd = sizeof(fds)*8; fd--;)
		if (FD_ISSET(fd, &fds))
			break;
	nfd = fd + 1;
	if (nfd <= 0) {
		errno = EINVAL;
		return nil;
	}
	do {
		fd_set rfds;
		id sock_obj;
		memcpy(&rfds, &fds, sizeof(rfds));
		if (select(nfd, &rfds, NULL, NULL, NULL) < 0) {
			if (errno == EINTR)
				continue;
			return nil;
		}
		for (fd = 0; fd < nfd; fd++) {
			if (!FD_ISSET(fd, &rfds))
				continue;
			sock_obj = isCollection ? (id)[sockets valueForKey:(void *)fd] : sockets;
			if ([sock_obj type] == MiscSOCK_STREAM)
				if ([sock_obj acceptNewConnection:&sock_obj timeout:0] < 0)
					return nil;
			if (fFlag) {
				switch (fork()) {
				case -1:
					return nil;
				case 0 :
					[target perform:action with:sock_obj];
					exit(0);
				}
			} else
				[target perform:action with:sock_obj];
			if ([sock_obj type] == MiscSOCK_STREAM)
				[sock_obj free];
		}
	} while (lFlag);
	return self;
}

+ sendDgram:(void *)data length:(int)dlen to:(MiscINETAddress *)addr port:(int)portNum
{
	struct sockaddr_in sockaddr;
	if (addr == nil) {
		errno = EDESTADDRREQ;
		return nil;
	}
	if (handySocket == nil)
		handySocket = [MiscINETSocket alloc];
	handySocket = [handySocket initDomain:PF_INET type:MiscSOCK_DGRAM];
	if (handySocket == nil || portNum < 0)
		return nil;
	sockaddr.sin_family = AF_INET;
	sockaddr.sin_addr = [addr address];
	sockaddr.sin_port = htons((short)(portNum & 0xffff));
	if (sendto(handySocket->sock, data, dlen, 0, (struct sockaddr *)&sockaddr, sizeof(sockaddr)) < 0)
		return nil;
	return self;
}

+ sendDgram:(void *)data length:(int)dlen to:(MiscINETAddress *)addr service:(const char *)service
{
	return [self sendDgram:data length:dlen to:addr port:get_service_port(service, MiscSOCK_DGRAM)];
}

+ (int)sendDgram:(void *)data length:(int)dlen to:(MiscINETAddress *)addr port:(int)portNum timeout:(unsigned int)ms retries:(int)retries withReply:(void *)repl length:(int *)rlen
{
	struct sockaddr_in to_addr;
	if (addr == nil) {
		errno = EDESTADDRREQ;
		return -1;
	}
	if (handySocket == nil)
		handySocket = [MiscINETSocket alloc];
	handySocket = [handySocket initDomain:PF_INET type:MiscSOCK_DGRAM];
	if (handySocket == nil || portNum < 0)
		return -1;
	to_addr.sin_family = AF_INET;
	to_addr.sin_addr = [addr address];
	to_addr.sin_port = htons((short)(portNum & 0xffff));
	do {
		int ret;
		if (sendto(handySocket->sock, data, dlen, 0, (struct sockaddr *)&to_addr, sizeof(to_addr)) < 0)
			return -1;
		ret = wait_for_socket_read(handySocket->sock, ms);
		if (ret < 0)
			return -1;
		if (0 < ret) {
			struct sockaddr_in from_addr;
			int len = sizeof(from_addr);
			ret = recvfrom(handySocket->sock, repl, *rlen, 0, (struct sockaddr *)&from_addr, &len);
			if (ret < 0)
				return -1;
			if (!memcmp(&to_addr, &from_addr, len-sizeof(to_addr.sin_zero))) {
				*rlen = ret;
				return 1;
			}
		}
	} while (retries-- > 0);
	return 0;
}

+ (int)sendDgram:(void *)data length:(int)dlen to:(MiscINETAddress *)addr service:(const char *)service timeout:(unsigned int)ms retries:(int)retries withReply:(void *)repl length:(int *)rlen
{
	return [self sendDgram:data length:dlen to:addr port:get_service_port(service, MiscSOCK_DGRAM) timeout:ms retries:retries withReply:repl length:rlen];
}

- init
{
	[super init];
	if (localAddress != nil)
		localAddress = [localAddress free];
	if (remoteAddress != nil)
		remoteAddress = [remoteAddress free];
	localPortNum = remotePortNum = -1;
	return self;
}

- copyFromZone:(NXZone *)zone
{
	MiscINETSocket *copy = [super copyFromZone:zone];
	copy->localAddress = [localAddress copyFromZone:zone];
	copy->remoteAddress = [remoteAddress copyFromZone:zone];
	return copy;
}

- read:(NXTypedStream *)stream;
{
	[super read:stream];
	localAddress = NXReadObject(stream);
	NXReadType(stream, "i", &localPortNum);
	remoteAddress = NXReadObject(stream);
	NXReadType(stream, "i", &remotePortNum);
	return self;
}

- write:(NXTypedStream *)stream
{
	[super write:stream];
	NXWriteRootObject(stream, localAddress);
	NXWriteType(stream, "i", &localPortNum);
	NXWriteRootObject(stream, remoteAddress);
	NXWriteType(stream, "i", &remotePortNum);
	return self;
}

- (int)acceptNewConnection:(MiscINETSocket **)newSocket timeout:(unsigned int)ms
{
	int ret;
	*newSocket = nil;
	if (type != MiscSOCK_STREAM) {
		errno = EINVAL;
		return -1;
	}
	ret = wait_for_socket_read(sock, ms);
	if (0 < ret) {
		struct sockaddr_in sockaddr;
		int len = sizeof(sockaddr);
		ret = accept(sock, (struct sockaddr *)&sockaddr, &len);
		if (0 <= ret) {
			*newSocket = [[MiscINETSocket alloc] init];
			(*newSocket)->sock = ret;
			(*newSocket)->domain = PF_INET;
			(*newSocket)->type = MiscSOCK_STREAM;
			ret = 1;
		}
	}
	return ret;
}

- connectTo:(MiscINETAddress *)addr port:(int)portNum type:(int)aType
{
	struct sockaddr_in sockaddr;
	if (addr == nil) {
		errno = EDESTADDRREQ;
		return [self free];
	}
	self = [self initDomain:PF_INET type:aType];
	if (self == nil || portNum < 0)
		return [self free];
	sockaddr.sin_family = AF_INET;
	sockaddr.sin_addr = [addr address];
	sockaddr.sin_port = htons((short)(portNum & 0xffff));
	if (connect(sock, (struct sockaddr *)&sockaddr, sizeof(sockaddr)) < 0)
		return [self free];
	return self;
}

- connectTo:(MiscINETAddress *)addr service:(const char *)service type:(int)aType
{
	return [self connectTo:addr port:get_service_port(service, aType) type:aType];
}

- openServerPort:(int *)portNum type:(int)aType
{
	struct sockaddr_in sockaddr;
	int on = 1, old_errno = errno;
	self = [self initDomain:PF_INET type:aType];
	if (self == nil || *portNum < 0)
		return [self free];
	setsockopt(sock, SOL_SOCKET, SO_REUSEADDR, (char *)&on, sizeof(on));
	errno = old_errno;
	sockaddr.sin_family = AF_INET;
	sockaddr.sin_port = htons((short)(*portNum & 0xffff));
	sockaddr.sin_addr.s_addr = INADDR_ANY;
	if (bind(sock, (struct sockaddr *)&sockaddr, sizeof(sockaddr)) < 0)
		return [self free];
	if (aType == MiscSOCK_STREAM)
		listen(sock, 2);
	return self;
}

- openServerService:(const char *)service type:(int)aType
{
	int port = get_service_port(service, aType);
	if (port < 0)
		return [self free];
	return [self openServerPort:&port type:aType];
}

- (int)receiveData:(void *)data length:(int *)dlen
{
	return [self receiveData:data length:dlen timeout:0 from:nil port:NULL];
}

- (int)receiveData:(void *)data length:(int *)dlen timeout:(unsigned int)ms from:(MiscINETAddress *)addr port:(int *)port
{
	int ret = wait_for_socket_read(sock, ms);
	if (0 <= ret) {
		struct sockaddr_in sockaddr;
		int len = sizeof(sockaddr);
		ret = recvfrom(sock, data, *dlen, 0, (struct sockaddr *)&sockaddr, &len);
		if (0 == ret) {
			[self close];
			errno = EPIPE;
			ret = -1;
		} else if (0 < ret) {
			[addr initTo:sockaddr.sin_addr];
			*dlen = ret;
			if (port != NULL)
				*port = (int)ntohs(sockaddr.sin_port);
			ret = 1;
		}
	}
	return ret;
}

- (int)receiveData:(void *)data length:(int *)dlen timeout:(unsigned int)ms toNext:(unsigned char)sentinel
{
	struct timeval t1, t2;
	int i, ret;
	if (type != MiscSOCK_STREAM) {
		errno = EINVAL;
		return -1;
	}
	gettimeofday(&t1, NULL);
	for (i = *dlen, *dlen = 0; i--;) {
		if (0 < ms) {
			int ms_left;
			gettimeofday(&t2, NULL);
			ms_left = ms - 1000*(t2.tv_sec-t1.tv_sec) + (t2.tv_usec-t1.tv_usec)/1000;
			if (ms_left < 1)
				return 0;
			ret = wait_for_socket_read(sock, ms_left);
			if (ret <= 0)
				return ret;
		}
		ret = read(sock, data, 1);
		if (ret < 0)
			return -1;
		else if (0 == ret) {
			[self close];
			errno = EPIPE;
			return -1;
		} 
		(*dlen)++;
		if (*(unsigned char *)data == sentinel)
			return 1;
		data++;
	}
	return 1;
}

- sendData:(void *)data length:(int)dlen
{
	while (0 < dlen) {
		int ret = write(sock, data, dlen);
		if (ret <= 0) {
			if (errno == EINTR)
				continue;
			else if (errno == EPIPE || errno == 0) {
				[self close];
				return nil;
			}
			return nil;
		}
		dlen -= ret;
		data += ret;
	}
	return self;
}

- shutdownLocalEnd
{
	if (localAddress != nil)
		localAddress = [localAddress free];
	localPortNum = -1;
	return (shutdown(sock, 1) < 0 ? nil : self);
}

- shutdownRemoteEnd
{
	if (remoteAddress != nil)
		remoteAddress = [remoteAddress free];
	remotePortNum = -1;
	return (shutdown(sock, 0) < 0 ? nil : self);
}

- enableDebug:(BOOL)aBool
{
	int on = aBool;
	if (![self isClosed]) {
		int size = sizeof(on), old_errno = errno;
		setsockopt(sock, SOL_SOCKET, SO_DEBUG, (char *)&on, size);
		errno = old_errno;
	}
	return self;
}

- (BOOL)debugEnabled
{
	int on = NO;
	if (![self isClosed]) {
		int size = sizeof(on), old_errno = errno;
		getsockopt(sock, SOL_SOCKET, SO_DEBUG, (char *)&on, &size);
		errno = old_errno;
	}
	return (BOOL)on;
}

- enableDelay:(BOOL)aBool
{
	int on = aBool;
	if (![self isClosed]) {
		int size = sizeof(on), old_errno = errno;
		setsockopt(sock, IPPROTO_TCP, TCP_NODELAY, (char *)&on, size);
		errno = old_errno;
	}
	return self;
}

- (BOOL)delayEnabled
{
	int on = YES;
	if (![self isClosed]) {
		int size = sizeof(on), old_errno = errno;
		getsockopt(sock, IPPROTO_TCP, TCP_NODELAY, (char *)&on, &size);
		errno = old_errno;
	}
	return (BOOL)on;
}

- enableKeepAlive:(BOOL)aBool
{
	int on = aBool;
	if (![self isClosed]) {
		int size = sizeof(on), old_errno = errno;
		setsockopt(sock, SOL_SOCKET, SO_KEEPALIVE, (char *)&on, size);
		errno = old_errno;
	}
	return self;
}

- (BOOL)keepAliveEnabled
{
	int on = NO;
	if (![self isClosed]) {
		int size = sizeof(on), old_errno = errno;
		getsockopt(sock, SOL_SOCKET, SO_KEEPALIVE, (char *)&on, &size);
		errno = old_errno;
	}
	return (BOOL)on;
}

- setLingerTime:(int)secs
{
	if (![self isClosed]) {
		struct linger ling;
		int size = sizeof(ling), old_errno = errno;
		ling.l_onoff = 0 < secs ? 1 : 0;
		ling.l_linger = 0 < secs ? secs : 0;
		setsockopt(sock, SOL_SOCKET, SO_LINGER, (char *)&ling, size);
		errno = old_errno;
	}
	return self;
}

- (int)lingerTime
{
	int time = 0;
	if (![self isClosed]) {
		struct linger ling;
		int size = sizeof(ling), old_errno = errno;
		getsockopt(sock, SOL_SOCKET, SO_LINGER, (char *)&ling, &size);
		errno = old_errno;
		if (ling.l_onoff)
			time = ling.l_linger;
	}
	return time;
}

- (BOOL)dataAvailable
{
	int old_errno = errno;
	BOOL avail = NO;
	if (![self isClosed]) {
		struct timeval timeout = {0, 0};
		fd_set rfds;
		FD_ZERO(&rfds);
		FD_SET(sock, &rfds);
		if (0 < select(sock+1, &rfds, NULL, NULL, &timeout))
			avail = YES;
	}
	errno = old_errno;
	return avail;
}

- (MiscINETAddress *)localAddress
{
	if (localAddress == nil) {
		struct sockaddr_in sockaddr;
		int len = sizeof(sockaddr);
		if (getsockname(sock, (struct sockaddr *)&sockaddr, &len) < 0)
			return nil;
		localAddress = [[MiscINETAddress allocFromZone:[self zone]] initTo:sockaddr.sin_addr];
		if (localAddress != nil)
			localPortNum = ntohs(sockaddr.sin_port & 0xffff);
	}
	return localAddress;
}

- (int)localPortNum
{
	if (localPortNum == -1)
		[self localAddress];
	return localPortNum;
}

- (MiscINETAddress *)remoteAddress
{
	if (remoteAddress == nil) {
		struct sockaddr_in sockaddr;
		int len = sizeof(sockaddr);
		if (getpeername(sock, (struct sockaddr *)&sockaddr, &len) < 0)
			return nil;
		remoteAddress = [[MiscINETAddress allocFromZone:[self zone]] initTo:sockaddr.sin_addr];
		if (remoteAddress != nil)
			remotePortNum = ntohs(sockaddr.sin_port & 0xffff);
	}
	return remoteAddress;
}

- (int)remotePortNum
{
	if (remotePortNum == -1)
		[self remoteAddress];
	return remotePortNum;
}

- (int)socketError
{
	int err = 0;
	if (![self isClosed]) {
		int size = sizeof(err), old_errno = errno;
		getsockopt(sock, SOL_SOCKET, SO_ERROR, (char *)&err, &size);
		errno = old_errno;
	}
	return err;
}

@end
