1/* Copyright (c) 2013-2014 Jeffrey Pfau
  2 *
  3 * This Source Code Form is subject to the terms of the Mozilla Public
  4 * License, v. 2.0. If a copy of the MPL was not distributed with this
  5 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
  6#ifndef SOCKET_H
  7#define SOCKET_H
  8
  9#include "util/common.h"
 10
 11#ifdef __cplusplus
 12#define restrict __restrict__
 13#endif
 14
 15#ifdef _WIN32
 16#include <winsock2.h>
 17#include <ws2tcpip.h>
 18
 19#define SOCKET_FAILED(s) ((s) == INVALID_SOCKET)
 20typedef SOCKET Socket;
 21#else
 22#include <errno.h>
 23#include <fcntl.h>
 24#include <netinet/in.h>
 25#include <netinet/tcp.h>
 26#include <sys/socket.h>
 27
 28#define INVALID_SOCKET (-1)
 29#define SOCKET_FAILED(s) ((s) < 0)
 30typedef int Socket;
 31#endif
 32
 33enum IP {
 34	IPV4,
 35	IPV6
 36};
 37
 38struct Address {
 39	enum IP version;
 40	union {
 41		uint32_t ipv4;
 42		uint8_t ipv6[16];
 43	};
 44};
 45
 46static inline void SocketSubsystemInit() {
 47#ifdef _WIN32
 48	WSADATA data;
 49	WSAStartup(MAKEWORD(2, 2), &data);
 50#endif
 51}
 52
 53static inline int SocketError() {
 54#ifdef _WIN32
 55	return WSAGetLastError();
 56#else
 57	return errno;
 58#endif
 59}
 60
 61static inline bool SocketWouldBlock() {
 62#ifdef _WIN32
 63	return SocketError() == WSAEWOULDBLOCK;
 64#else
 65	return SocketError() == EWOULDBLOCK || SocketError() == EAGAIN;
 66#endif
 67}
 68
 69static inline ssize_t SocketSend(Socket socket, const void* buffer, size_t size) {
 70#ifdef _WIN32
 71	return send(socket, (const char*) buffer, size, 0);
 72#else
 73	return write(socket, buffer, size);
 74#endif
 75}
 76
 77static inline ssize_t SocketRecv(Socket socket, void* buffer, size_t size) {
 78#ifdef _WIN32
 79	return recv(socket, (char*) buffer, size, 0);
 80#else
 81	return read(socket, buffer, size);
 82#endif
 83}
 84
 85static inline Socket SocketOpenTCP(int port, const struct Address* bindAddress) {
 86	Socket sock = socket(PF_INET, SOCK_STREAM, IPPROTO_TCP);
 87	if (SOCKET_FAILED(sock)) {
 88		return sock;
 89	}
 90
 91	int err;
 92	if (!bindAddress) {
 93		struct sockaddr_in bindInfo;
 94		memset(&bindInfo, 0, sizeof(bindInfo));
 95		bindInfo.sin_family = AF_INET;
 96		bindInfo.sin_port = htons(port);
 97		err = bind(sock, (const struct sockaddr*) &bindInfo, sizeof(bindInfo));
 98	} else if (bindAddress->version == IPV4) {
 99		struct sockaddr_in bindInfo;
100		memset(&bindInfo, 0, sizeof(bindInfo));
101		bindInfo.sin_family = AF_INET;
102		bindInfo.sin_port = htons(port);
103		bindInfo.sin_addr.s_addr = bindAddress->ipv4;
104		err = bind(sock, (const struct sockaddr*) &bindInfo, sizeof(bindInfo));
105	} else {
106		struct sockaddr_in6 bindInfo;
107		memset(&bindInfo, 0, sizeof(bindInfo));
108		bindInfo.sin6_family = AF_INET6;
109		bindInfo.sin6_port = htons(port);
110		memcpy(bindInfo.sin6_addr.s6_addr, bindAddress->ipv6, sizeof(bindInfo.sin6_addr.s6_addr));
111		err = bind(sock, (const struct sockaddr*) &bindInfo, sizeof(bindInfo));
112
113	}
114	if (err) {
115		close(sock);
116		return INVALID_SOCKET;
117	}
118	return sock;
119}
120
121static inline Socket SocketConnectTCP(int port, const struct Address* destinationAddress) {
122	Socket sock = socket(PF_INET, SOCK_STREAM, IPPROTO_TCP);
123	if (SOCKET_FAILED(sock)) {
124		return sock;
125	}
126
127	int err;
128	if (!destinationAddress) {
129		struct sockaddr_in bindInfo;
130		memset(&bindInfo, 0, sizeof(bindInfo));
131		bindInfo.sin_family = AF_INET;
132		bindInfo.sin_port = htons(port);
133		err = connect(sock, (const struct sockaddr*) &bindInfo, sizeof(bindInfo));
134	} else if (destinationAddress->version == IPV4) {
135		struct sockaddr_in bindInfo;
136		memset(&bindInfo, 0, sizeof(bindInfo));
137		bindInfo.sin_family = AF_INET;
138		bindInfo.sin_port = htons(port);
139		bindInfo.sin_addr.s_addr = destinationAddress->ipv4;
140		err = connect(sock, (const struct sockaddr*) &bindInfo, sizeof(bindInfo));
141	} else {
142		struct sockaddr_in6 bindInfo;
143		memset(&bindInfo, 0, sizeof(bindInfo));
144		bindInfo.sin6_family = AF_INET6;
145		bindInfo.sin6_port = htons(port);
146		memcpy(bindInfo.sin6_addr.s6_addr, destinationAddress->ipv6, sizeof(bindInfo.sin6_addr.s6_addr));
147		err = connect(sock, (const struct sockaddr*) &bindInfo, sizeof(bindInfo));
148	}
149
150	if (err) {
151		close(sock);
152		return INVALID_SOCKET;
153	}
154	return sock;
155}
156
157static inline Socket SocketListen(Socket socket, int queueLength) {
158	return listen(socket, queueLength);
159}
160
161static inline Socket SocketAccept(Socket socket, struct Address* address) {
162	if (!address) {
163		return accept(socket, 0, 0);
164	}
165	if (address->version == IPV4) {
166		struct sockaddr_in addrInfo;
167		memset(&addrInfo, 0, sizeof(addrInfo));
168		addrInfo.sin_family = AF_INET;
169		addrInfo.sin_addr.s_addr = address->ipv4;
170		socklen_t len = sizeof(addrInfo);
171		return accept(socket, (struct sockaddr*) &addrInfo, &len);
172	} else {
173		struct sockaddr_in6 addrInfo;
174		memset(&addrInfo, 0, sizeof(addrInfo));
175		addrInfo.sin6_family = AF_INET6;
176		memcpy(addrInfo.sin6_addr.s6_addr, address->ipv6, sizeof(addrInfo.sin6_addr.s6_addr));
177		socklen_t len = sizeof(addrInfo);
178		return accept(socket, (struct sockaddr*) &addrInfo, &len);
179	}
180}
181
182static inline int SocketClose(Socket socket) {
183	return close(socket) >= 0;
184}
185
186static inline int SocketSetBlocking(Socket socket, bool blocking) {
187#ifdef _WIN32
188	u_long unblocking = !blocking;
189	return ioctlsocket(socket, FIONBIO, &unblocking) == NO_ERROR;
190#else
191	int flags = fcntl(socket, F_GETFL);
192	if (flags == -1) {
193		return 0;
194	}
195	if (blocking) {
196		flags &= ~O_NONBLOCK;
197	} else {
198		flags |= O_NONBLOCK;
199	}
200	return fcntl(socket, F_SETFL, flags) >= 0;
201#endif
202}
203
204static inline int SocketSetTCPPush(Socket socket, int push) {
205	return setsockopt(socket, IPPROTO_TCP, TCP_NODELAY, (char*) &push, sizeof(int)) >= 0;
206}
207
208static inline int SocketPoll(size_t nSockets, Socket* reads, Socket* writes, Socket* errors, int64_t timeoutMillis) {
209	fd_set rset;
210	fd_set wset;
211	fd_set eset;
212	FD_ZERO(&rset);
213	FD_ZERO(&wset);
214	FD_ZERO(&eset);
215	size_t i;
216	Socket maxFd = 0;
217	if (reads) {
218		for (i = 0; i < nSockets; ++i) {
219			if (SOCKET_FAILED(reads[i])) {
220				break;
221			}
222			if (reads[i] > maxFd) {
223				maxFd = reads[i];
224			}
225			FD_SET(reads[i], &rset);
226			reads[i] = INVALID_SOCKET;
227		}
228	}
229	if (writes) {
230		for (i = 0; i < nSockets; ++i) {
231			if (SOCKET_FAILED(writes[i])) {
232				break;
233			}
234			if (writes[i] > maxFd) {
235				maxFd = writes[i];
236			}
237			FD_SET(writes[i], &wset);
238			writes[i] = INVALID_SOCKET;
239		}
240	}
241	if (errors) {
242		for (i = 0; i < nSockets; ++i) {
243			if (SOCKET_FAILED(errors[i])) {
244				break;
245			}
246			if (errors[i] > maxFd) {
247				maxFd = errors[i];
248			}
249			FD_SET(errors[i], &eset);
250			errors[i] = INVALID_SOCKET;
251		}
252	}
253	struct timeval tv;
254	tv.tv_sec = timeoutMillis / 1000;
255	tv.tv_usec = (timeoutMillis % 1000) * 1000;
256	int result = select(maxFd + 1, &rset, &wset, &eset, timeoutMillis < 0 ? 0 : &tv);
257	int r = 0;
258	int w = 0;
259	int e = 0;
260	Socket j;
261	for (j = 0; j < maxFd; ++j) {
262		if (reads && FD_ISSET(j, &rset)) {
263			reads[r] = j;
264			++r;
265		}
266		if (writes && FD_ISSET(j, &wset)) {
267			writes[w] = j;
268			++w;
269		}
270		if (errors && FD_ISSET(j, &eset)) {
271			errors[e] = j;
272			++e;
273		}
274	}
275	return result;
276}
277
278#endif