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