rmps/src/client_pool.c

/*
 * client_pool.c
 *
 * Copyright (C) 2018 by Bogomil Vasilev <b.vasilev@smirky.net>
 *
 * This file is part of Remote Management and Provisioning System (RMPS).
 *
 * RMPS is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 2 of the License, or
 * (at your option) any later version.
 *
 * RMPS is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with RMPS.  If not, see <http://www.gnu.org/licenses/>.
 */

#include "client_pool.h"
#include "log.h"
#include "job_queue.h"
#include "protocol.h"
#include <pthread.h>
#include <unistd.h>
#include <string.h>

#define MAXJOBS 10

struct client_args {
	int busy;
	SSL *ssl;
	int sd;
	char ip[16]; /* IPv4 */
};

static void *servlet(void *args);
static void send_reject_msg(SSL *ssl);

static void *servlet(void *args)  /* Serve the connection -- threadable */
{
	struct msg_t buf;
	int bytes, ret;
	//unsigned short job[MAXJOBS] = { 0 };
	struct client_args *client = (struct client_args *)args;

	SSL_load_error_strings();
	ret = SSL_accept(client->ssl);
	/* We check for unclean (ret < 0) and clean (ret == 0) failures */
	if (ret <= 0) {
		log(WARNING, "SSL_accept() failed. Reason below:");
		log_ssl();
	} else {
		int queue_id = start_msg_queue();

		if (queue_id == FAIL)
			goto exit;
		do {
			//buf.meta.type = GET_MEMORY;
			//sleep(1);
			//SSL_write(client->ssl, &buf, sizeof(buf));
			bytes = SSL_read(client->ssl, &buf, sizeof(buf));
			if (bytes > 0) {
				if (bytes != sizeof(struct msg_t)) {
					log(WARNING,
						"Client [%s] sent non-standard data!",
						client->ip);
					continue;
				}

				log(VERBOSE, "Client msg: \"%s\"", buf.chunk.data);
				/* TODO: Insert msg handler here */
				add_msg_to_queue(queue_id, buf);

				continue;
			}
			if (SSL_get_shutdown(client->ssl) == SSL_RECEIVED_SHUTDOWN)
				log(VERBOSE, "SSL_RECEIVED_SHUTDOWN from client [%s]", client->ip);
			else {
				log(VERBOSE, "Client didn't send data!");
				//log_ssl(); /* We actually don't have anything to log from SSL */
				sprintf((char *)buf.chunk.data, "%s", "Where's the data, m8?");
				SSL_write(client->ssl, &buf, sizeof(struct msg_t));
			}
			log(INFO, "Client [%s] disconnected.", client->ip);
		} while (bytes);
	}
exit:
	SSL_free(client->ssl);          /* release SSL state */
	close(client->sd);              /* close connection */
	client->busy = 0;
	return 0;
}

static void send_reject_msg(SSL *ssl)
{
	char *reply = "FAILURE - The connection queue is full!\n";

	SSL_write(ssl, reply, strlen(reply));
}

void *client_pool(void *args)
{
	struct pool_data *pool = args;
	pthread_mutex_t mutex;
	pthread_attr_t attr;
	pthread_t *client_thread = malloc(pool->size * sizeof(*client_thread));
	struct client_args *client_struct =
		malloc(pool->size * sizeof(*client_struct));
	int i;

	memset(client_thread, 0, sizeof(pthread_t) * pool->size);
	memset(client_struct, 0, sizeof(struct client_args) * pool->size);

	pthread_mutex_init(&mutex, NULL);
	pthread_attr_init(&attr);
	pthread_attr_setdetachstate(&attr, PTHREAD_CREATE_JOINABLE);

	while (1) {
		struct sockaddr_in addr;
		char address[INET6_ADDRSTRLEN];
		socklen_t len = sizeof(addr);
		SSL *ssl;
		int agent = accept(pool->srv, (struct sockaddr *)&addr, &len);

		log(INFO,
			"Connection: %s:%d",
			inet_ntop(AF_INET, &addr.sin_addr, address, sizeof(address)),
			ntohs(addr.sin_port));

		for (i = 0; i < pool->size; i++) {
			if (!client_struct[i].busy) {
				client_struct[i].busy = 1;
				client_struct[i].ssl = SSL_new(pool->ctx);
				client_struct[i].sd = agent;
				memcpy(client_struct[i].ip,
					inet_ntop(AF_INET, &addr.sin_addr, address, sizeof(address)),
					sizeof(client_struct[i].ip));
				SSL_set_fd(client_struct[i].ssl, client_struct[i].sd);
				pthread_create(&client_thread[i],
						&attr,
						servlet,
						&client_struct[i]);
				break;
			}
		}
		if (i == pool->size) {
			log(WARNING,
				   "Agent [%s] dropped. Poolsize limit reached.",
				   inet_ntop(AF_INET, &addr.sin_addr, address, sizeof(address))
			);
			ssl = SSL_new(pool->ctx);
			SSL_set_fd(ssl, agent);
			if (SSL_accept(ssl) == FAIL) {
				SSL_free(ssl);
				close(agent);
				continue;
			}
			send_reject_msg(ssl);
			SSL_free(ssl);
			close(agent);
		}
	}
	pthread_attr_destroy(&attr);
	pthread_mutex_destroy(&mutex);
	pthread_detach(pthread_self());
	pthread_exit(NULL);
}