nonbsocket2.c

C编写的用来实现search engine的推荐功能

/***************************************
 * Functions to implement SOCKET operations
 * Author: Lubing Xu
 * Date:   Nov. 1999
 ***************************************/
#ifndef WIN32
#include <sys/time.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <signal.h>
#include <unistd.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <netdb.h>
#include <fcntl.h>
#include <errno.h>
#define closesocket close
typedef struct sockaddr_in SOCKADDR_IN;
#else
#include <windows.h>
#endif
#include <string.h>
#include <stdio.h>
#include <stdlib.h>

#include "NonbSocket2.h"
typedef void (*sighandler_t)(int);

#ifndef INADDR_NONE
#define INADDR_NONE -1
#endif

#ifdef WIN32
void StartWinsock()
{
	WSADATA wsaData;
	WSAStartup(MAKEWORD(1,1), &wsaData);
}
#endif

/***************************************
 * To create a TCP server socket listening
 * on the given port
 * Parameter:
 *  port: the port to listen at
 * Return:
 *  If the server socket is created
 * successfully, the socket handle
 * will be returned.
 *  INVALID_SOCKET otherwise.
 ***************************************/
int SocketServer(int port){
	SOCKADDR_IN sin;
	int socket_handle;
	int err;
	int opt = 1;/* 1: Enable,  0: Disable */
	
	socket_handle = socket(AF_INET,SOCK_STREAM,0);
	if(socket_handle == INVALID_SOCKET)
		return INVALID_SOCKET;

	setsockopt(socket_handle, SOL_SOCKET, SO_REUSEADDR, (const void*)&opt, sizeof(opt));

	memset(&sin, 0, sizeof(sin));
	sin.sin_family = AF_INET;
	sin.sin_port = htons((short)port);
	sin.sin_addr.s_addr = INADDR_ANY;
	
	err = bind(socket_handle, (struct sockaddr*)&sin, sizeof(sin));
	if(err==SOCKET_ERROR)
	{
		closesocket(socket_handle);
		return INVALID_SOCKET;
	}

	err = listen(socket_handle, 20);
	if(err==SOCKET_ERROR)
	{
		closesocket(socket_handle);
		return INVALID_SOCKET;
	}

	return socket_handle;
}

/**
 * accept a connection from one client
 * Parameter:
 *  serversock: the socket handle of the listening server
 * Return:
 *  the client socket handle or INVALID_SOCKET
 * NOTE:
 *  If you want to know the information of client IP & port,
 *  this function doesn't provide them. Please use function
 *  accept() directly.
 */
int Accept(int serversock)
{
	struct sockaddr addr;
	int addrlen;
	
	addrlen = sizeof(addr);
	return accept(serversock, &addr, &addrlen);
}

static unsigned long GetIPAddr(const char* host){
	unsigned long addr;
	
	if (host == NULL)
		return 0;
		
	addr = inet_addr(host);
	if (addr == INADDR_NONE){
		struct hostent* pHostent = gethostbyname(host);
		if (pHostent == NULL)
			return 0;
		else
			addr = *(unsigned long*)pHostent->h_addr_list[0];
	}

	return addr;
}

#define DEFAULT_TIMEOUT     300 /* timeout in seconds */
#define DEFAULT_NONB_TIMEOUT 30 /* non-block timeout */

#ifdef LINUX
void connect_alarm(int sig)
{
	return;
}
#endif

/********************************************
 * Create a non-block TCP client socket
 * Parameters:
 *  server: server name or IP
 *  port:   server port
 *  timeout:seconds to wait for connection
 * Return:
 *  INVALID_SOCKET: any error occurrs
 *  a socket handle: successful
 ********************************************/
int NonbSocket(const char* server, int port, int timeout) {
#ifdef WIN32
	return Socket(server, port);
#else
	unsigned long ip;
	fd_set rset, wset;
	struct timeval tv;
	SOCKADDR_IN sin;
	int sock;
	int flags, len, n;
#ifdef LINUX
	sighandler_t alarmfunc;
#endif
	
	if (server == NULL)
		return INVALID_SOCKET;
	ip = GetIPAddr(server);
	if(ip == 0)
		return INVALID_SOCKET;

	sock = socket(AF_INET,SOCK_STREAM,0);
	if(sock == INVALID_SOCKET)
		return INVALID_SOCKET;

#ifdef LINUX
	alarmfunc = signal(SIGALRM, connect_alarm);
	if (timeout <= 0) {
		timeout = DEFAULT_NONB_TIMEOUT;
	}
	if (alarm(timeout) != 0) {
		signal(SIGALRM, alarmfunc);
		goto ERR;
	}
	/* connect to server */
	memset(&sin, 0, sizeof(sin));
	sin.sin_family = AF_INET;
	sin.sin_addr.s_addr = ip;
	sin.sin_port = htons((u_short)port);
	if (connect(sock, (struct sockaddr*)&sin, sizeof(sin)) == SOCKET_ERROR) {
		alarm(0);
		signal(SIGALRM, alarmfunc);
		goto ERR;
	}
	alarm(0);
	signal(SIGALRM, alarmfunc);
	return sock;
#else
	/* get status flags */
	if((flags = fcntl(sock, F_GETFL, 0)) < 0){
		goto ERR;
	}
	/* set non-block status flag */
	if(fcntl(sock, F_SETFL, flags | O_NONBLOCK) == -1){
		goto ERR;
	}
	/* connect to server */
	memset(&sin, 0, sizeof(sin));
	sin.sin_family = AF_INET;
	sin.sin_addr.s_addr = ip;
	sin.sin_port = htons((u_short)port);
	if(connect(sock, (struct sockaddr*)&sin, sizeof(sin)) == SOCKET_ERROR){
		if(errno != EINPROGRESS){
			goto ERR;
		}
	}else goto DONE; /* connect completed immediately */
	
	/* check whether the connect will complete after timeout */
	if (timeout <= 0) {
		timeout = DEFAULT_NONB_TIMEOUT;
	}
	FD_ZERO(&rset);
	FD_SET(sock, &rset);
	wset = rset;
	tv.tv_sec = timeout;
	tv.tv_usec = 0;
	
	if(select(sock+1, &rset, &wset, NULL, timeout ? &tv : NULL) == 0){/* timeout */
		goto ERR;
	}
	if(FD_ISSET(sock, &rset) || FD_ISSET(sock, &wset)){
		len = sizeof(n);
		if(getsockopt(sock, SOL_SOCKET, SO_ERROR, (char*)&n, &len) < 0 || (n!=0))
			goto ERR; /* Solaris pending error */
	}else goto ERR;
#endif
	
DONE:
	/* restore status flags */
	if(fcntl(sock, F_SETFL, flags) == -1)goto ERR;
	return sock;

ERR:
	closesocket(sock);
	return INVALID_SOCKET;
#endif
}

/********************************************
 * Create a TCP client socket
 * Parameters:
 *  server: server name or IP
 *  prot:   server port
 * Return:
 *  INVALID_SOCKET: any error occurrs
 *  a socket handle: successful
 ********************************************/
int Socket(const char *server, int port){
	unsigned long ip = GetIPAddr(server);
	SOCKADDR_IN sin;
	int sock;
	
	if(ip == 0) {
		return INVALID_SOCKET;
	}

	sock = socket(AF_INET,SOCK_STREAM,0);
	if(sock == INVALID_SOCKET)
		return INVALID_SOCKET;

	/* connect to server */
	memset(&sin, 0, sizeof(sin));
	sin.sin_family = AF_INET;
	sin.sin_addr.s_addr = ip;
	sin.sin_port = htons((u_short)port);
	if(connect(sock, (struct sockaddr*)&sin, sizeof(sin)) == SOCKET_ERROR){
		closesocket(sock);
		return INVALID_SOCKET;
	}

	return sock;
}

/***************************************
 * Check timeout when sending
 * Parameters:
 *  hSocket: socket handle to check
 *  time:    timeout in seconds
 * Return:
 *  >0: OK.
 *  <0: any error occurrs
 *   0: time expires
 ***************************************/
static int CheckSocketSendTimeOut(int hSocket,int time)
{
	fd_set set;
	struct timeval tv;
	int result;
	
	FD_ZERO(&set);
	FD_SET((unsigned)hSocket,&set);

	tv.tv_sec = time;
	tv.tv_usec = 0;

	result = select(hSocket+1,NULL,&set,NULL,&tv);
	return result;
}

/***************************************
 * Check timeout when receiving
 * Parameters:
 *  hSocket: socket handle to check
 *  time:    timeout in seconds
 * Return:
 *  >0: OK.
 *  <0: any error occurrs
 *   0: time expires
 ***************************************/
static int CheckSocketRecvTimeOut(int hSocket,int time)
{
	fd_set set;
	struct timeval tv;
	int result;
	
	FD_ZERO(&set);
	FD_SET((unsigned)hSocket,&set);

	tv.tv_sec = time;
	tv.tv_usec = 0;

	result = select(hSocket+1,&set,NULL,NULL,&tv);
	return result;
}

/***************************************
 * Send specified length of bytes
 * Parameters:
 *  sock: the socket handle to send
 *  buffer: the bytes to send
 *  total: the total bytes to send
 *  timeout: timeout int senconds
 * Return:
 *  >=0: the number of bytes sent if successful,
 *  SOCKET_TIMEOUT: timeout for waiting
 *  SOCKET_ERROR: if any error occurrs.
 ***************************************/
int Send(int sock, const char *buffer, int total, int timeout)
{
	int retval;
	int sentSize = 0;
	
	if (buffer == NULL)
		return 0;
		
	if (timeout <= 0) {
		timeout = DEFAULT_TIMEOUT;
	}
	while(sentSize < total) {
		int res = CheckSocketSendTimeOut(sock, timeout);
		if(res < 0) {
			return SOCKET_ERROR;
		} else if (res == 0) {
			if (sentSize == 0) {
				return SOCKET_TIMEOUT;
			} else {
				return sentSize;
			}
		}
	
		retval = send(sock, buffer+sentSize, total-sentSize, 0);
		if (retval == SOCKET_ERROR) {
			return SOCKET_ERROR;
		}
		
		sentSize += retval;
	}
	
	return total;
}

/***************************************
 * Receive bytes from a socket
 * Parameters:
 *  sock: the socket handle to receive
 *  buffer: the buffer to store the data
 *  size: the size of the buffer
 *  timeout: timeout in seconds
 * Return:
 *  >0: the number of bytes received successfully
 *  SOCKET_CLOSED: the socket has been closed
 *  SOCKET_TIMEOUT: timeout for waiting.
 *  SOCKET_ERROR: if any error occurrs
 ***************************************/
int Receive(int sock, char *buffer, int size, int timeout)
{
	int receiveResult;
	int recvSize = 0;
	
	if (buffer == NULL)
		return 0;
	
	if (timeout <= 0) {
		timeout = DEFAULT_TIMEOUT;
	}

#ifdef OLD_NONBSOCKET2
	/* 以下代码不再维护 */
	receiveResult = CheckSocketRecvTimeOut(sock, timeout);
	if(receiveResult < 0) {
		return SOCKET_ERROR;
	} else if (receiveResult == 0) {
		return SOCKET_TIMEOUT;
	}
	
	receiveResult = recv(sock, buffer+recvSize, size-recvSize, MSG_DONTWAIT);

	/* if receiveResult ==0, the connection has been closed */
	if (receiveResult == 0) {
		return SOCKET_CLOSED;
	} else if (receiveResult < 0) {
		return SOCKET_ERROR;
	}
	
	return receiveResult;
#else
	while (recvSize < size)	{
		receiveResult = CheckSocketRecvTimeOut(sock, timeout);
		if(receiveResult < 0) {
			if (recvSize <= 0) {
				return SOCKET_ERROR;
			} else {
				return recvSize;
			}
		} else if (receiveResult == 0) {
			if (recvSize <= 0) {
				return SOCKET_TIMEOUT;
			} else {
				return recvSize;
			}
		}
#ifndef WIN32
		receiveResult = recv(sock, buffer+recvSize, size-recvSize, MSG_DONTWAIT);
#else
		receiveResult = recv(sock, buffer+recvSize, size-recvSize, 0);
#endif
	
		if (receiveResult == 0) {
			if (recvSize <= 0) {
				return SOCKET_CLOSED;
			} else {
				return recvSize;
			}
		} else if (receiveResult < 0) {
#ifndef WIN32
			if (EAGAIN == errno) {
				continue;
			}
#endif
			if (recvSize <= 0) {
				return SOCKET_ERROR;
			} else  {
				return recvSize;
			}
		}
		
		recvSize += receiveResult;
	}
	
	return size;
#endif
}

/***************************************
 * Send a string then send an extra new line
 * character
 * Parameters:
 *  sock: the socket handle to send
 *  buffer: a string ended with ASCIIZ
 *  timeout: timeout in seconds
 * Return:
 *  >0: if successfuly.
 *  SOCKET_ERROR: if any error occurrs.
 ***************************************/
int SendLine(int sock, const char *buffer, int timeout) {
	int len;
	int retval;
	int sentSize = 0;
#ifndef OLD_NONBSOCKET2
	char buf[1024];
	char* dbuf = NULL;
#endif

	if (buffer == NULL) {
		return 1;
	}
	
	/* send the total buffer */
	len = strlen(buffer);
#ifndef OLD_NONBSOCKET2
	if (len < 1024) {
		/* <1024个字节，复制到临时缓冲区 */
		memcpy(buf, buffer, len);
		dbuf = buf;
	} else {
		/* >=1024个字节，复制到动态缓冲区 */
		dbuf = strdup(buffer);
	}
	if (dbuf != NULL) {
		dbuf[len++] = '\n'; /* 加上换行符一次发送 */
		
		while (sentSize < len) {
			retval = Send(sock, dbuf+sentSize, len-sentSize, timeout);
			switch (retval) {
			case SOCKET_ERROR:
				if (dbuf != buf) {
					free(dbuf); /* 释放动态内存 */
				}
				return SOCKET_ERROR;
			case SOCKET_TIMEOUT:
				break;
			default:
				sentSize += retval;
				break;
			}
		}
		
		if (dbuf != buf) {
			free(dbuf); /* 释放动态内存 */
		}
		return sentSize;
	} else {
		/* 动态内存申请失败，仍采用原来的发送方式 */
#endif
		while (sentSize < len) {
			retval = Send(sock, buffer+sentSize, len-sentSize, timeout);
			switch (retval) {
			case SOCKET_ERROR:
				return SOCKET_ERROR;
			case SOCKET_TIMEOUT:
				break;
			default:
				sentSize += retval;
				break;
			}
		}
	
		/* send an extra new line */
		while (1) {
			switch (Send(sock, "\n", 1, timeout)) {
			case SOCKET_ERROR:
				return SOCKET_ERROR;
			case SOCKET_TIMEOUT:
				break;
			default:
				return 1;
			}
		}

		return 1;
#ifndef OLD_NONBSOCKET2
	}
#endif
}

/***************************************
 * Read a line ended with new line char
 * Parameters:
 *  sock: the socket handle to read
 *  buffer: a buffer to store data
 *  size: the size of the buffer
 *  timeout: timeout in seconds
 * Return:
 *  >0: the length of line
 *  SOCKET_CLOSED: the socket has been closed
 *  SOCKET_TIMEOUT: timeout for waiting.
 *  SOCKET_ERROR: if any error occurrs
 * NOTE:
 *  any '\r' is ignored
 ***************************************/
int ReadLine(int sock, char *buffer, int size, int timeout){
	int readByte = 0;
	int receiveResult;

	if (buffer == NULL) {
		return 0;
	}
	
	if (timeout <= 0) {
		timeout = DEFAULT_TIMEOUT;
	}
	
	while(readByte < size){
		int res = CheckSocketRecvTimeOut(sock, timeout);
		if(res < 0) {
			return SOCKET_ERROR;
		} else if (res == 0) {
			if (readByte == 0) {
				return SOCKET_TIMEOUT;
			} else {
				buffer[readByte] = '\0';
				return readByte;
			}
		}
		
#ifndef WIN32
		receiveResult = recv(sock, buffer+readByte, 1, MSG_DONTWAIT);
#else
		receiveResult = recv(sock, buffer+readByte, 1, 0);
#endif
		/* if receiveResult ==0, the connection has been closed */
		if(receiveResult == 0) {
			if (readByte == 0) {
				return SOCKET_CLOSED;
			} else {
				buffer[readByte] = '\0';
				return readByte;
			}
		} else if(receiveResult < 0) {
#ifndef WIN32
			if (EAGAIN == errno) {
				continue;
			}
#endif
			return SOCKET_ERROR;
		}
		
		switch (buffer[readByte]) {
		case '\r':
			continue;
			
		case '\n':
			buffer[readByte] = '\0';
			/* run case '0' */
		case '\0':
			return readByte;
		
		default:
			readByte++;
			break;
		}
	}

	/* the buffer is too small */
	return size;
}

/********************************************
 * Close a socket
 * Parameter:
 *  the socket handle to close
 ********************************************/
void CloseSocket(int sock)
{
	closesocket(sock);
}