pex

pex.c

/* This file is part of ToaruOS and is released under the terms
 * of the NCSA / University of Illinois License - see LICENSE.md
 * Copyright (C) 2014 Kevin Lange
 */
#include <sys/socket.h>
#include <sys/un.h>

#include <alloca.h>
#include <assert.h>
#include <errno.h>
#include <fcntl.h>
#include <ioleast.h>
#include <poll.h>
#include <stdbool.h>
#include <stddef.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#include "pex.h"

#undef PACKET_SIZE

#define MAX_CONNS 256

struct pex_server_connection
{
	int fd;
	unsigned char buffer[MAX_PACKET_SIZE];
	size_t received;
};

struct pex_server
{
	int fd;
	struct pex_server_connection conns[MAX_CONNS];
	size_t num_connections;
};

static int open_local_client_socket(const char* path, int flags)
{
	size_t path_length = strlen(path);
	size_t addr_size = offsetof(struct sockaddr_un, sun_path) + path_length + 1;
	struct sockaddr_un* sockaddr = (struct sockaddr_un*) malloc(addr_size);
	if ( !sockaddr )
		return -1;
	sockaddr->sun_family = AF_LOCAL;
	strcpy(sockaddr->sun_path, path);
	int fd = socket(AF_LOCAL, SOCK_STREAM | flags, 0);
	if ( fd < 0 )
		return free(sockaddr), -1;
	if ( connect(fd, (const struct sockaddr*) sockaddr, addr_size) < 0 )
		return close(fd), free(sockaddr), -1;
	free(sockaddr);
	return fd;
}

static int open_local_server_socket(const char* path, int flags)
{
	size_t path_length = strlen(path);
	size_t addr_size = offsetof(struct sockaddr_un, sun_path) + path_length + 1;
	struct sockaddr_un* sockaddr = (struct sockaddr_un*) malloc(addr_size);
	if ( !sockaddr )
		return -1;
	sockaddr->sun_family = AF_LOCAL;
	strcpy(sockaddr->sun_path, path);
	int fd = socket(AF_LOCAL, SOCK_STREAM | flags, 0);
	if ( fd < 0 )
		return free(sockaddr), -1;
	if ( bind(fd, (const struct sockaddr*) sockaddr, addr_size) < 0 )
		return close(fd), free(sockaddr), -1;
	if ( listen(fd, 5) < 0 )
		return close(fd), free(sockaddr), -1;
	free(sockaddr);
	return fd;
}

size_t pex_send(FILE * sock, unsigned int rcpt, size_t size, char * blob) {
	struct pex_server* server = (struct pex_server*) sock;
	assert(rcpt != 0);
	size_t i = rcpt - 1;
	char packet[MAX_PACKET_SIZE];
	assert(size <= sizeof(packet));
	memcpy(packet, blob, size);
	memset(packet + size, 0, sizeof(packet) - size);
	int fd = server->conns[i].fd;
	int fl = fcntl(fd, F_GETFL);
	fcntl(fd, F_SETFL, fl & ~O_NONBLOCK);
	size_t amount = writeall(fd, packet, sizeof(packet));
	assert(amount == sizeof(packet));
	fcntl(fd, F_SETFL, fl);
	return size;
}

size_t pex_broadcast(FILE * sock, size_t size, char * blob) {
	struct pex_server* server = (struct pex_server*) sock;
	for ( size_t i = 0; i < MAX_CONNS; i++ )
	{
		if ( 0 <= server->conns[i].fd )
			pex_send(sock, 1 + i, size, blob);
	}
	return size;
}

size_t pex_listen(FILE * sock, pex_packet_t * packet) {
	struct pex_server* server = (struct pex_server*) sock;
	while ( true )
	{
		for ( size_t i = 0; i < MAX_CONNS; i++ )
		{
			if ( server->conns[i].received == MAX_PACKET_SIZE )
			{
				packet->source = 1 + i;
				packet->size = 1;
				memcpy(packet->data, server->conns[i].buffer, MAX_PACKET_SIZE);
				server->conns[i].received = 0;
				return 1;
			}
		}

		struct pollfd pfds[1 + MAX_CONNS];
		memset(pfds, 0, sizeof(pfds));
		pfds[0].fd = server->fd;
		if ( server->num_connections < MAX_CONNS )
			pfds[0].events |= POLLIN;
		for ( size_t i = 0; i < MAX_CONNS; i++ )
		{
			if ( (pfds[1 + i].fd = server->conns[i].fd) < 0 )
				continue;
			if ( server->conns[i].received < MAX_PACKET_SIZE )
				pfds[1 + i].events |= POLLIN;
		}

		poll(pfds, 1 + MAX_CONNS, -1);

		if ( pfds[0].revents & POLLIN )
		{
			int client_fd = accept4(server->fd, NULL, NULL, SOCK_NONBLOCK);
			size_t i = 0;
			while ( 0 <= server->conns[i].fd )
				i++;
			memset(&server->conns[i], 0, sizeof(server->conns[i]));
			server->conns[i].fd = client_fd;
			server->num_connections++;
		}
		for ( size_t i = 0; i < MAX_CONNS; i++ )
		{
			if ( pfds[1 + i].revents & (POLLERR | POLLHUP | POLLNVAL) )
			{
				close(server->conns[i].fd);
				memset(&server->conns[i], 0, sizeof(server->conns[i]));
				server->conns[i].fd = -1;
				server->num_connections--;
				packet->source = 1 + i;
				packet->size = 0;
				return 0;
			}

			if ( pfds[1 + i].revents & POLLIN )
			{
				int fd = server->conns[i].fd;
				unsigned char* buffer = server->conns[i].buffer;
				size_t left = MAX_PACKET_SIZE - server->conns[i].received;
				ssize_t amount = read(fd, buffer, left);

				if ( amount < 0 && (errno == EAGAIN || errno == EWOULDBLOCK) )
					continue;
				if ( amount < 0 || amount == 0 )
					continue;
				server->conns[i].received += amount;
			}
		}
	}
}

size_t pex_reply(FILE * sock, size_t size, char * blob) {
	char packet[MAX_PACKET_SIZE];
	assert(size <= sizeof(packet));
	memcpy(packet, blob, size);
	memset(packet + size, 0, sizeof(packet) - size);
	size_t amount = writeall(fileno(sock), packet, sizeof(packet));
	assert(amount == sizeof(packet));
	return amount;
}

size_t pex_recv(FILE * sock, char * blob) {
	memset(blob, 0, MAX_PACKET_SIZE);
	size_t amount = readall(fileno(sock), blob, MAX_PACKET_SIZE);
	assert(amount == MAX_PACKET_SIZE);
	return amount;
}

FILE * pex_connect(char * target) {
	char* path;
	if ( asprintf(&path, "/run/%s", target) < 0 )
		return NULL;
	int fd = open_local_client_socket(path, SOCK_CLOEXEC);
	free(path);
	if ( fd < 0 )
		return NULL;
	FILE* fp = fdopen(fd, "r+");
	if ( !fp )
		return close(fd), (FILE*) NULL;
	return fp;
}

FILE * pex_bind(char * target) {
	char* path;
	if ( asprintf(&path, "/run/%s", target) < 0 )
		return NULL;
	int fd = open_local_server_socket(path, SOCK_NONBLOCK | SOCK_CLOEXEC);
	free(path);
	if ( fd < 0 )
		return NULL;
	struct pex_server* server = calloc(sizeof(struct pex_server), 1);
	if ( !server )
		return close(fd), (FILE*) NULL;
	server->fd = fd;
	for ( size_t i = 0; i < MAX_CONNS; i++ )
		server->conns[i].fd = -1;
	return (FILE*) server;
}

size_t pex_query(FILE * sock) {
	struct pollfd pfd = { 0 };
	pfd.fd = fileno(sock);
	pfd.events = POLLIN;
	poll(&pfd, 1, 0);
	return pfd.revents & POLLIN ? 1 : 0;
}