myvtm/libtoolkit/sockutil.c

#include "precompile.h"
#include "sockutil.h"
#include "strutil.h"

//define replaced by cmake synatx
//#pragma comment(lib, "ws2_32.lib")

TOOLKIT_API int winsock_init()
{
	WSADATA wsaData;
	return WSAStartup(0x0202, &wsaData);
}

TOOLKIT_API int winsock_term()
{
	return WSACleanup();
}

TOOLKIT_API char* get_local_ip(char *buf, int n)
{
	char name[64];
	struct hostent * he;
	gethostname(name, sizeof(name));
	he = gethostbyname(name);
	if (he) {
		if (he->h_addr_list[0]) {
			return strncpy(buf, inet_ntoa(*(struct in_addr*)he->h_addr_list[0]), n);
		}
	}
	return NULL;
}

// < 0 : error
// < n : timeout
// ==n : ok
TOOLKIT_API int tsend_n(int fd, const char *buf, int n, long timeout/* = -1 */)
{
	struct timeval _tm;
	struct timeval *tm;
	int offset = 0;
	int left = n;
	DWORD dwLast;

	tm = (timeout >= 0) ? &_tm : NULL;
	if (tm)
		dwLast = GetTickCount();
	while (left > 0 && (timeout == -1 || timeout > 0)) {
		int t = send(fd, buf+offset, left, 0);
		if (t == 0) {
			return offset;
		} else if (t < 0) {
			if (WSAGetLastError() == WSAEWOULDBLOCK) {
				int rc;
				fd_set wr_fs;
				fd_set ex_fs;
				FD_ZERO(&wr_fs);
				FD_ZERO(&ex_fs);
				FD_SET(fd, &wr_fs);
				FD_SET(fd, &ex_fs);
				if (tm) {
					DWORD dwNow = GetTickCount();
					long diff = (long)(dwNow - dwLast);
					if (timeout >= diff) {
						timeout -= diff;
					} else {
						timeout = 0;
					}
					tm->tv_sec = timeout/1000;
					tm->tv_usec = timeout%1000;
					dwLast = dwNow;
				}
				rc = select(1, NULL, &wr_fs, &ex_fs, tm);
				if (rc < 0)
					return -1;
				if (rc == 0)
					return offset;
				if (FD_ISSET(fd, &ex_fs))
					return -1;
				assert(FD_ISSET(fd, &wr_fs));
			} else {
				return -1;
			}
		} else { // (t > 0)
			offset += t;
			left -= t;
		}
	}
	return offset;
}

TOOLKIT_API int trecv_n(int fd, char *buf, size_t n, long timeout/* = -1*/)
{
	struct timeval _tm;
	struct timeval *tm;
	int offset = 0;
	int left = n;
	DWORD dwLast;

	tm = (timeout >= 0) ? &_tm : NULL;
	if (tm)
		dwLast = GetTickCount();
	while (left > 0 && (timeout == -1 || timeout > 0)) {
		int t = recv(fd, buf+offset, left, 0);
		if (t == 0) {
			return offset;
		} else if (t < 0) {
			if (WSAGetLastError() == WSAEWOULDBLOCK) {
				int rc;
				fd_set rd_fs;
				fd_set ex_fs;
				FD_ZERO(&rd_fs);
				FD_ZERO(&ex_fs);
				FD_SET(fd, &rd_fs);
				FD_SET(fd, &ex_fs);
				if (tm) {
					DWORD dwNow = GetTickCount();
					long diff = (long)(dwNow - dwLast);
					if (timeout >= diff) {
						timeout -= diff;
					} else {
						timeout = 0;
					}
					tm->tv_sec = timeout/1000;
					tm->tv_usec = timeout%1000;
					dwLast = dwNow;
				}
				rc = select(1, &rd_fs, NULL, &ex_fs, tm);
				if (rc < 0)
					return -1;
				if (rc == 0)
					return offset;
				if (FD_ISSET(fd, &ex_fs))
					return -1;
				assert(FD_ISSET(fd, &rd_fs));
			} else {
				return -1;
			}
		} else {
			offset += t;
			left -= t;
		}
	}
	return offset;
}

TOOLKIT_API int trecv_until(int fd, char *buf, int n, const char *delimer, int *header_len, long timeout)
{
	struct timeval _tm;
	struct timeval *tm;
	int offset = 0;
	int left = n;
	DWORD dwLast;

	assert(delimer);

	*header_len = 0;

	tm = (timeout >= 0) ? &_tm : NULL;
	if (tm)
		dwLast = GetTickCount();
	while (left > 0 && (timeout == -1 || timeout > 0)) {
		int t = recv(fd, buf+offset, left, 0);
		if (t == 0) {
			*header_len = 0;
			return offset;
		} else if (t < 0) {
			if (WSAGetLastError() == WSAEWOULDBLOCK) {
				int rc;
				fd_set rd_fs;
				fd_set ex_fs;
				FD_ZERO(&rd_fs);
				FD_ZERO(&ex_fs);
				FD_SET(fd, &rd_fs);
				FD_SET(fd, &ex_fs);
				if (tm) {
					DWORD dwNow = GetTickCount();
					long diff = (long)(dwNow - dwLast);
					if (timeout >= diff) {
						timeout -= diff;
					} else {
						timeout = 0;
					}
					tm->tv_sec = timeout/1000;
					tm->tv_usec = timeout%1000;
					dwLast = dwNow;
				}
				rc = select(1, &rd_fs, NULL, &ex_fs, tm);
				if (rc < 0)
					return -1;
				if (rc == 0)
					return offset;
				if (FD_ISSET(fd, &ex_fs))
					return -1;
				assert(FD_ISSET(fd, &rd_fs));
			} else {
				return -1;
			}
		} else {
			const char *p;
			offset += t;
			left -= t;
			p = memstr(buf, offset, delimer);
			if (p) {
				p += strlen(delimer);
				*header_len = p - buf;
				return offset;
			}
		}
	}

	return offset;
}

TOOLKIT_API int tsend_until(int fd, const char *buf, int n, const char *delimer, long timeout)
{
	const char *p;

	assert(delimer);

	p = memstr(buf, n, delimer);
	if (!p)
		return -1;
	p += strlen(delimer);

	return tsend_n(fd, buf, p - buf, timeout);
}


#ifndef SIO_UDP_CONNRESET
#define SIO_UDP_CONNRESET _WSAIOW(IOC_VENDOR,12)
#endif

TOOLKIT_API int disable_udp_connreset(SOCKET so)
{
	DWORD dwBytesReturned = 0;
	BOOL bNewBehavior = FALSE;
	return WSAIoctl(so, SIO_UDP_CONNRESET, &bNewBehavior, sizeof(bNewBehavior),
		NULL, 0, &dwBytesReturned, NULL, NULL);
}

TOOLKIT_API int tsend_n_v(int fd, LPWSABUF lpBuffers, DWORD dwBufferCount, long timeout)
{
	struct timeval _tm;
	struct timeval *tm;
	int sendedbufs = 0;
	int leftbufs = (int)dwBufferCount;
	int offset = 0;
	int sendedbytes = 0;
	DWORD dwLast;

	tm = (timeout >= 0) ? &_tm : NULL;
	if (tm)
		dwLast = GetTickCount();
	while (leftbufs > 0 && (timeout == -1 || timeout > 0)) {
		DWORD dwSent = 0;
		int rc = WSASend(fd, lpBuffers+sendedbufs, leftbufs, &dwSent, 0, NULL, NULL);
		if (rc < 0) {
			if (WSAGetLastError() == WSAEWOULDBLOCK) {
				int rc;
				fd_set wr_fs;
				fd_set ex_fs;
				FD_ZERO(&wr_fs);
				FD_ZERO(&ex_fs);
				FD_SET(fd, &wr_fs);
				FD_SET(fd, &ex_fs);
				if (tm) {
					DWORD dwNow = GetTickCount();
					long diff = (long)(dwNow - dwLast);
					if (timeout >= diff) {
						timeout -= diff;
					} else {
						timeout = 0;
					}
					tm->tv_sec = timeout/1000;
					tm->tv_usec = timeout%1000;
					dwLast = dwNow;
				}
				rc = select(1, NULL, &wr_fs, &ex_fs, tm);
				if (rc < 0)
					return -1;
				if (rc == 0)
					return sendedbytes;
				if (FD_ISSET(fd, &ex_fs))
					return -1;
				assert(FD_ISSET(fd, &wr_fs));
			} else {
				return -1;
			}
		} else { // (t > 0)
			if (dwSent == 0)
				return sendedbytes;
			sendedbytes += dwSent;
			while (dwSent > 0) {
				if (dwSent >= (lpBuffers[sendedbufs].len-offset)) {
					dwSent -= lpBuffers[sendedbufs].len-offset;
					sendedbufs++;
					leftbufs--;
					offset = 0;
				} else {
					lpBuffers[sendedbufs].buf += dwSent;
					lpBuffers[sendedbufs].len -= dwSent;
					dwSent = 0;
				}
			}
		}
	}
	return sendedbytes;
}

TOOLKIT_API int trecv_n_v(int fd, LPWSABUF lpBuffers, DWORD dwBufferCount, long timeout)
{
	struct timeval _tm;
	struct timeval *tm;
	int recvedbufs = 0;
	int leftbufs = (int)dwBufferCount;
	int offset = 0;
	int recvedbytes = 0;
	DWORD dwLast;

	tm = (timeout >= 0) ? &_tm : NULL;
	if (tm)
		dwLast = GetTickCount();
	while (leftbufs > 0 && (timeout == -1 || timeout > 0)) {
		DWORD dwRecv = 0;
		int rc = WSARecv(fd, lpBuffers+recvedbufs, leftbufs, &dwRecv, 0, NULL, NULL);
		if (rc < 0) {
			if (WSAGetLastError() == WSAEWOULDBLOCK) {
				int rc;
				fd_set rd_fs;
				fd_set ex_fs;
				FD_ZERO(&rd_fs);
				FD_ZERO(&ex_fs);
				FD_SET(fd, &rd_fs);
				FD_SET(fd, &ex_fs);
				if (tm) {
					DWORD dwNow = GetTickCount();
					long diff = (long)(dwNow - dwLast);
					if (timeout >= diff) {
						timeout -= diff;
					} else {
						timeout = 0;
					}
					tm->tv_sec = timeout/1000;
					tm->tv_usec = timeout%1000;
					dwLast = dwNow;
				}
				rc = select(1, &rd_fs, NULL, &ex_fs, tm);
				if (rc < 0)
					return -1;
				if (rc == 0)
					return recvedbytes;
				if (FD_ISSET(fd, &ex_fs))
					return -1;
				assert(FD_ISSET(fd, &rd_fs));
			} else {
				return -1;
			}
		} else { // (t > 0)
			if (dwRecv == 0)
				return recvedbytes;
			recvedbytes += dwRecv;
			while (dwRecv > 0) {
				if (dwRecv >= (lpBuffers[recvedbufs].len-offset)) {
					dwRecv -= lpBuffers[recvedbufs].len-offset;
					recvedbufs++;
					leftbufs--;
					offset = 0;
				} else {
					lpBuffers[recvedbufs].buf += dwRecv;
					lpBuffers[recvedbufs].len -= dwRecv;
					dwRecv = 0;
				}
			}
		}
	}
	return recvedbytes;
}

TOOLKIT_API int tsendto(int fd, const char *buf, int n, const struct sockaddr* to, int tolen, long timeout)
{
	struct timeval _tm;
	struct timeval *tm;
	DWORD dwLast;

	tm = (timeout >= 0) ? &_tm : NULL;
	if (tm)
		dwLast = GetTickCount();

	do {
		int rc = sendto(fd, buf, n, 0, to, tolen);
		if (rc >= 0)
			return rc;
		if (WSAGetLastError() == WSAEWOULDBLOCK) {
			int rc;
			fd_set wr_fs;
			fd_set ex_fs;
			FD_ZERO(&wr_fs);
			FD_ZERO(&ex_fs);
			FD_SET(fd, &wr_fs);
			FD_SET(fd, &ex_fs);
			if (tm) {
				DWORD dwNow = GetTickCount();
				long diff = (long)(dwNow - dwLast);
				if (timeout >= diff) {
					timeout -= diff;
				} else {
					timeout = 0;
				}
				tm->tv_sec = timeout/1000;
				tm->tv_usec = timeout%1000;
				dwLast = dwNow;
			}
			rc = select(1, NULL, &wr_fs, &ex_fs, tm);
			if (rc <= 0)
				break;
			if (FD_ISSET(fd, &ex_fs))
				break;
			assert(FD_ISSET(fd, &wr_fs));
		} else {
			break;
		}
	} while (timeout == -1 || timeout > 0);

	return -1;
}

TOOLKIT_API int trecvfrom(int fd, char *buf, int n, struct sockaddr* fromaddr, int *fromlen, long timeout)
{
	struct timeval _tm;
	struct timeval *tm;
	DWORD dwLast;

	tm = (timeout >= 0) ? &_tm : NULL;
	if (tm)
		dwLast = GetTickCount();

	do {
		int rc = recvfrom(fd, buf, n, 0, fromaddr, fromlen);
		if (rc >= 0)
			return rc;
		if (WSAGetLastError() == WSAEWOULDBLOCK) {
			int rc;
			fd_set rd_fs;
			fd_set ex_fs;
			FD_ZERO(&rd_fs);
			FD_ZERO(&ex_fs);
			FD_SET(fd, &rd_fs);
			FD_SET(fd, &ex_fs);
			if (tm) {
				DWORD dwNow = GetTickCount();
				long diff = (long)(dwNow - dwLast);
				if (timeout >= diff) {
					timeout -= diff;
				} else {
					timeout = 0;
				}
				tm->tv_sec = timeout/1000;
				tm->tv_usec = timeout%1000;
				dwLast = dwNow;
			}
			rc = select(1, &rd_fs, NULL, &ex_fs, tm);
			if (rc <= 0)
				break;
			if (FD_ISSET(fd, &ex_fs))
				break;
			assert(FD_ISSET(fd, &rd_fs));
		} else {
			break;
		}
	} while (timeout == -1 || timeout > 0);

	return -1;
}