Blob Blame Raw

#include <cerrno>
#include <cassert>
#include <cstring>
#include <climits>
#include <cctype>

#include <fcntl.h>
#include <unistd.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/poll.h>
#include <netinet/in.h>
#include <netinet/tcp.h>
#include <netdb.h>

#include "tcp.h"


static const int pollTimeoutMs = 100;
static const int connTimeoutMs = 300000;


using namespace Tcp;


ErrorCode Address::resolve(const char *addrString) {
	*this = Address();
	
	if (!addrString)
		return ERR_TCP_BAD_ADDRESS;
	
	ErrorCode errorCode = ERR_NONE;
	
	const char *colon = addrString;
	while(*colon && *colon != ':') ++colon;
	
	char *hostBuf = nullptr;
	const char *host = nullptr;
	const char *port = nullptr;
	if (*colon) {
		hostBuf = new char[colon - addrString + 1];
		memcpy(hostBuf, addrString, colon - addrString);
		hostBuf[colon - addrString] = 0;
		host = hostBuf;
		port = colon + 1;
	} else {
		bool allDigits = true;
		for(const char *c = addrString; *c && allDigits; ++c)
			if (*c < '0' || *c > '9') allDigits = false;
		(allDigits ? port : host) = addrString;
	}
	
	if (host) {
		if (4 != sscanf(host, "%hhu.%hhu.%hhu.%hhu", &ip[0], &ip[1], &ip[2], &ip[3])) {
			memset(ip, 0, sizeof(ip));
			hostent *he;
			in_addr **addr_list;
			if ( (he = gethostbyname(host))
			  && (addr_list = (struct in_addr**)he->h_addr_list)
			  && (*addr_list) )
			{
				memcpy(ip, *addr_list, sizeof(ip));
			} else {
				errorCode = ERR_TCP_BAD_ADDRESS;
			}
		}
	}
	
	if (port) {
		if (1 != sscanf(port, "%hu", &this->port))
			errorCode = ERR_TCP_BAD_ADDRESS;
	}
	
	if (!host && !port)
		errorCode = ERR_TCP_BAD_ADDRESS;
	
	if (hostBuf) delete[] hostBuf;
	return errorCode;
}


Connection::Connection():
	status(STATUS_NONE),
	lastError(ERR_NONE),
	sockId(),
	stopping(false)
	{ }

Connection::Connection(SockId sockId, const Address &remoteAddr):
	status(STATUS_OPEN),
	lastError(ERR_NONE),
	sockId(sockId),
	remoteAddr(remoteAddr),
	stopping(false)
	{ }

Connection::~Connection()
	{ close(); }

bool Connection::open(const Address &remoteAddr) {
	close();
	
	this->remoteAddr = remoteAddr;
	if (!this->remoteAddr.isValid())
		{ close(ERR_TCP_BAD_ADDRESS); return false; }

	sockId = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
	status = STATUS_OPEN;
	
	struct sockaddr_in addr = {};
	addr.sin_family = AF_INET;
	memcpy(&addr.sin_addr, remoteAddr.ip, sizeof(remoteAddr.ip));
	addr.sin_port = htons(remoteAddr.port);
	
	if (0 != connect(sockId, (sockaddr*)&addr, sizeof(addr)))
		{ close(ERR_TCP_CONNECTION_FAILED); return false; }
	
	int tcp_nodelay = 1;
	fcntl(sockId, F_SETFL, fcntl(sockId, F_GETFL, 0) | O_NONBLOCK);
	setsockopt(sockId, IPPROTO_TCP, TCP_NODELAY, &tcp_nodelay, sizeof(int));
	
	return true;
}

void Connection::close(ErrorCode errorCode) {
	if (status == STATUS_OPEN)
		::close(sockId);
	sockId = 0;
	status = errorCode ? STATUS_LOST : STATUS_NONE;
	lastError = errorCode;
	stopping = false;
}

bool Connection::read(void *data, size_t size) {
	if (!data)
		return check() && size == 0;
	
	struct pollfd pfd = {};
	pfd.fd = sockId;
	pfd.events = POLLIN;
	
	int remainMs = connTimeoutMs;
	while(check() && size && remainMs >= 0) {
		int pr = poll(&pfd, 1, pollTimeoutMs);
		if (pr < 0)
			{ close(ERR_TCP_CONNECTION_LOST); return false; }
		if (pr > 0) {
			int s = size > INT_MAX ? INT_MAX : (int)size;
			int r = ::recv(sockId, data, s, MSG_DONTWAIT | MSG_NOSIGNAL);
			if (r < 0) {
				if (errno != EAGAIN)
					{ close(ERR_TCP_CONNECTION_LOST); return false; }
			} else {
				size -= r;
				data = (char*)data + r;
			}
		}
		remainMs -= pollTimeoutMs;
	}

	return size == 0;
}

bool Connection::write(const void *data, size_t size) {
	if (!data)
		return check() && size == 0;
	
	struct pollfd pfd = {};
	pfd.fd = sockId;
	pfd.events = POLLOUT;
	
	int remainMs = connTimeoutMs;
	while(check() && size && remainMs >= 0) {
		int pr = poll(&pfd, 1, pollTimeoutMs);
		if (pr < 0)
			{ close(ERR_TCP_CONNECTION_LOST); return false; }
		if (pr > 0) {
			int s = size > INT_MAX ? INT_MAX : (int)size;
			int r = ::send(sockId, data, s, MSG_DONTWAIT | MSG_NOSIGNAL);
			if (r < 0) {
				if (errno != EAGAIN)
					{ close(ERR_TCP_CONNECTION_LOST); return false; }
			} else {
				size -= r;
				data = (const char*)data + r;
			}
		}
		remainMs -= pollTimeoutMs;
	}

	return size == 0;
}


Handler::~Handler() { }


Server::Server():
	status(STATUS_NONE),
	lastError(ERR_NONE),
	sockId(),
	localAddr(),
	handler(),
	thread(),
	stopping(false)
	{ }

Server::~Server()
	{ stop(); }

void Server::handleConnection(ConnList::iterator iter) {
	handler->handleTcpConnection(*iter->connection);
	{
		std::lock_guard<std::mutex> lock(connectionsMutex);
		connectionsFinished.push_back(*iter);
		connections.erase(iter);
	}
}

void Server::cleanConnections() {
	while(!connectionsFinished.empty()) {
		ConnDesc &desc = connectionsFinished.back();
		desc.thread->join();
		delete desc.thread;
		delete desc.connection;
		connectionsFinished.pop_back();
	}
}

void Server::listen() {
	fcntl(sockId, F_SETFL, fcntl(sockId, F_GETFL, 0) | O_NONBLOCK);

	struct sockaddr_in addr = {};
	Address remoteAddr;
	
	struct pollfd pfd = {};
	pfd.fd = sockId;
	pfd.events = POLLIN;
	
	status = STATUS_OPEN;
	while(!stopping) {
		int pr = poll(&pfd, 1, pollTimeoutMs);
		if (pr < 0)
			{ lastError = ERR_TCP_LISTEN_LOST; break; }
			
		if (pr > 0 && (pr & POLLIN)) {
			socklen_t addrlen = sizeof(addr);
			int sid = ::accept(sockId, (sockaddr*)&addr, &addrlen);
			if (sid < 0) {
				if (errno != EAGAIN)
					{ lastError = ERR_TCP_LISTEN_LOST; break; }
				continue;
			}
			
			memcpy(remoteAddr.ip, &addr.sin_addr, sizeof(remoteAddr.ip));
			remoteAddr.port = ntohs(addr.sin_port);
			
			ConnList::iterator iter;
			{
				std::lock_guard<std::mutex> lock(connectionsMutex);
				cleanConnections();
				iter = connections.emplace(connections.end());
			}
			iter->connection = new Connection(sid, remoteAddr);
			iter->thread = new std::thread(&Server::handleConnection, this, iter);
		} else {
			std::lock_guard<std::mutex> lock(connectionsMutex);
			cleanConnections();
		}
	}
	status = lastError ? STATUS_LOST : STATUS_CLOSED;
	
	connectionsMutex.lock();
	for(ConnList::iterator i = connections.begin(); i != connections.end(); ++i)
		i->connection->stop();
	while(!connections.empty()) {
		std::thread *thread = connections.front().thread;
		connectionsMutex.unlock();
		thread->join();
		connectionsMutex.lock();
	}
	cleanConnections();
	connectionsMutex.unlock();
	
	stopping = false;
}

bool Server::start(const Address &localAddr, Handler *handler) {
	stop();
	
	this->localAddr = localAddr;
	this->handler = handler;
	
	if (!this->localAddr.isValidPort())
		{ status = STATUS_LOST; lastError = ERR_TCP_BAD_ADDRESS; return false; }
	if (!this->handler)
		{ status = STATUS_LOST; lastError = ERR_TCP_BAD_HANDLER; return false; }
	
	sockId = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
	
	struct sockaddr_in addr = {};
	addr.sin_family = AF_INET;
	memcpy(&addr.sin_addr, localAddr.ip, sizeof(localAddr.ip));
	addr.sin_port = htons(localAddr.port);
	
	if ( 0 != ::bind(sockId, (sockaddr*)&addr, sizeof(addr))
	  || 0 != ::listen(sockId, 32) )
	{
		status = STATUS_LOST;
		lastError = ERR_TCP_LISTEN_FAILED;
		this->handler = nullptr;
		return false;
	}
	
	int tcp_nodelay = 1;
	fcntl(sockId, F_SETFL, fcntl(sockId, F_GETFL, 0) | O_NONBLOCK);
	setsockopt(sockId, IPPROTO_TCP, TCP_NODELAY, &tcp_nodelay, sizeof(int));
	
	thread = new std::thread(&Server::listen, this);
	return true;
}

void Server::join() {
	if (!thread)
		return;
	thread->join();
	delete thread;
	thread = nullptr;
	handler = nullptr;
}

void Server::stop() {
	if (!thread)
		return;
	stopping = true;
	join();
}