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	return send(socket, buffer, size, 0);
 71}
 72
 73static inline ssize_t SocketRecv(Socket socket, void* buffer, size_t size) {
 74	return recv(socket, buffer, size, 0);
 75}
 76
 77static inline Socket SocketOpenTCP(int port, const struct Address* bindAddress) {
 78	Socket sock = socket(PF_INET, SOCK_STREAM, IPPROTO_TCP);
 79	if (SOCKET_FAILED(sock)) {
 80		return sock;
 81	}
 82
 83	int err;
 84	if (!bindAddress) {
 85		struct sockaddr_in bindInfo;
 86		memset(&bindInfo, 0, sizeof(bindInfo));
 87		bindInfo.sin_family = AF_INET;
 88		bindInfo.sin_port = htons(port);
 89		err = bind(sock, (const struct sockaddr*) &bindInfo, sizeof(bindInfo));
 90	} else if (bindAddress->version == IPV4) {
 91		struct sockaddr_in bindInfo;
 92		memset(&bindInfo, 0, sizeof(bindInfo));
 93		bindInfo.sin_family = AF_INET;
 94		bindInfo.sin_port = htons(port);
 95		bindInfo.sin_addr.s_addr = bindAddress->ipv4;
 96		err = bind(sock, (const struct sockaddr*) &bindInfo, sizeof(bindInfo));
 97	} else {
 98		struct sockaddr_in6 bindInfo;
 99		memset(&bindInfo, 0, sizeof(bindInfo));
100		bindInfo.sin6_family = AF_INET6;
101		bindInfo.sin6_port = htons(port);
102		memcpy(bindInfo.sin6_addr.s6_addr, bindAddress->ipv6, sizeof(bindInfo.sin6_addr.s6_addr));
103		err = bind(sock, (const struct sockaddr*) &bindInfo, sizeof(bindInfo));
104
105	}
106	if (err) {
107		close(sock);
108		return INVALID_SOCKET;
109	}
110	return sock;
111}
112
113static inline Socket SocketConnectTCP(int port, const struct Address* destinationAddress) {
114	Socket sock = socket(PF_INET, SOCK_STREAM, IPPROTO_TCP);
115	if (SOCKET_FAILED(sock)) {
116		return sock;
117	}
118
119	int err;
120	if (!destinationAddress) {
121		struct sockaddr_in bindInfo;
122		memset(&bindInfo, 0, sizeof(bindInfo));
123		bindInfo.sin_family = AF_INET;
124		bindInfo.sin_port = htons(port);
125		err = connect(sock, (const struct sockaddr*) &bindInfo, sizeof(bindInfo));
126	} else if (destinationAddress->version == IPV4) {
127		struct sockaddr_in bindInfo;
128		memset(&bindInfo, 0, sizeof(bindInfo));
129		bindInfo.sin_family = AF_INET;
130		bindInfo.sin_port = htons(port);
131		bindInfo.sin_addr.s_addr = destinationAddress->ipv4;
132		err = connect(sock, (const struct sockaddr*) &bindInfo, sizeof(bindInfo));
133	} else {
134		struct sockaddr_in6 bindInfo;
135		memset(&bindInfo, 0, sizeof(bindInfo));
136		bindInfo.sin6_family = AF_INET6;
137		bindInfo.sin6_port = htons(port);
138		memcpy(bindInfo.sin6_addr.s6_addr, destinationAddress->ipv6, sizeof(bindInfo.sin6_addr.s6_addr));
139		err = connect(sock, (const struct sockaddr*) &bindInfo, sizeof(bindInfo));
140	}
141
142	if (err) {
143		close(sock);
144		return INVALID_SOCKET;
145	}
146	return sock;
147}
148
149static inline Socket SocketListen(Socket socket, int queueLength) {
150	return listen(socket, queueLength);
151}
152
153static inline Socket SocketAccept(Socket socket, struct Address* address) {
154	if (!address) {
155		return accept(socket, 0, 0);
156	}
157	if (address->version == IPV4) {
158		struct sockaddr_in addrInfo;
159		memset(&addrInfo, 0, sizeof(addrInfo));
160		addrInfo.sin_family = AF_INET;
161		addrInfo.sin_addr.s_addr = address->ipv4;
162		socklen_t len = sizeof(addrInfo);
163		return accept(socket, (struct sockaddr*) &addrInfo, &len);
164	} else {
165		struct sockaddr_in6 addrInfo;
166		memset(&addrInfo, 0, sizeof(addrInfo));
167		addrInfo.sin6_family = AF_INET6;
168		memcpy(addrInfo.sin6_addr.s6_addr, address->ipv6, sizeof(addrInfo.sin6_addr.s6_addr));
169		socklen_t len = sizeof(addrInfo);
170		return accept(socket, (struct sockaddr*) &addrInfo, &len);
171	}
172}
173
174static inline int SocketClose(Socket socket) {
175	return close(socket) >= 0;
176}
177
178static inline int SocketSetBlocking(Socket socket, bool blocking) {
179#ifdef _WIN32
180	u_long unblocking = !blocking;
181	return ioctlsocket(socket, FIONBIO, &unblocking) == NO_ERROR;
182#else
183	int flags = fcntl(socket, F_GETFL);
184	if (flags == -1) {
185		return 0;
186	}
187	if (blocking) {
188		flags &= ~O_NONBLOCK;
189	} else {
190		flags |= O_NONBLOCK;
191	}
192	return fcntl(socket, F_SETFL, flags) >= 0;
193#endif
194}
195
196static inline int SocketSetTCPPush(Socket socket, int push) {
197	return setsockopt(socket, IPPROTO_TCP, TCP_NODELAY, (char*) &push, sizeof(int)) >= 0;
198}
199
200static inline int SocketPoll(size_t nSockets, Socket* reads, Socket* writes, Socket* errors, int64_t timeoutMillis) {
201	fd_set rset;
202	fd_set wset;
203	fd_set eset;
204	FD_ZERO(&rset);
205	FD_ZERO(&wset);
206	FD_ZERO(&eset);
207	size_t i;
208	Socket maxFd = 0;
209	if (reads) {
210		for (i = 0; i < nSockets; ++i) {
211			if (SOCKET_FAILED(reads[i])) {
212				break;
213			}
214			if (reads[i] > maxFd) {
215				maxFd = reads[i];
216			}
217			FD_SET(reads[i], &rset);
218			reads[i] = INVALID_SOCKET;
219		}
220	}
221	if (writes) {
222		for (i = 0; i < nSockets; ++i) {
223			if (SOCKET_FAILED(writes[i])) {
224				break;
225			}
226			if (writes[i] > maxFd) {
227				maxFd = writes[i];
228			}
229			FD_SET(writes[i], &wset);
230			writes[i] = INVALID_SOCKET;
231		}
232	}
233	if (errors) {
234		for (i = 0; i < nSockets; ++i) {
235			if (SOCKET_FAILED(errors[i])) {
236				break;
237			}
238			if (errors[i] > maxFd) {
239				maxFd = errors[i];
240			}
241			FD_SET(errors[i], &eset);
242			errors[i] = INVALID_SOCKET;
243		}
244	}
245	struct timeval tv;
246	tv.tv_sec = timeoutMillis / 1000;
247	tv.tv_usec = (timeoutMillis % 1000) * 1000;
248	int result = select(maxFd, &rset, &wset, &eset, timeoutMillis < 0 ? 0 : &tv);
249	int r = 0;
250	int w = 0;
251	int e = 0;
252	Socket j;
253	for (j = 0; j < maxFd; ++j) {
254		if (reads && FD_ISSET(j, &rset)) {
255			reads[r] = j;
256			++r;
257		}
258		if (writes && FD_ISSET(j, &wset)) {
259			writes[w] = j;
260			++w;
261		}
262		if (errors && FD_ISSET(j, &eset)) {
263			errors[e] = j;
264			++e;
265		}
266	}
267	return result;
268}
269
270#endif