aws.c

// SPDX-License-Identifier: BSD-3-Clause

#define _GNU_SOURCE

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <assert.h>
#include <sys/types.h>
#include <unistd.h>
#include <fcntl.h>
#include <sys/stat.h>
#include <sys/epoll.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <sys/sendfile.h>
#include <sys/eventfd.h>
#include <libaio.h>
#include <errno.h>
#include <time.h>

#include "aws.h"
#include "utils/util.h"
#include "utils/debug.h"
#include "utils/sock_util.h"
#include "utils/w_epoll.h"

/* server socket file descriptor */
static int listenfd;

/* epoll file descriptor */
static int epollfd;

static io_context_t ctx;


static int aws_on_path_cb(http_parser *p, const char *buf, size_t len)
{
	struct connection *conn = (struct connection *)p->data;

	memcpy(conn->request_path, buf, len);
	conn->request_path[len] = '\0';
	conn->have_path = 1;

	return 0;
}


/* Prepare the connection buffer to send the reply header. */
static void connection_prepare_send_reply_header(struct connection *conn)
{
	/* Get current date and time */
	char date[100];
	time_t now = time(NULL);
	struct tm tm;

	gmtime_r(&now, &tm);
	strftime(date, sizeof(date), "%a, %d %b %Y %H:%M:%S GMT", &tm);

	conn->send_len = snprintf(conn->send_buffer, BUFSIZ,
							  "HTTP/1.1 200 OK\r\n"
							  "Date: %s\r\n"
							  "Content-Length: %lu\r\n"
							  "Connection: close\r\n"
							  "\r\n",
							  date, conn->file_size);
}


/* Prepare the connection buffer to send the 404 header. */
static void connection_prepare_send_404(struct connection *conn)
{
	/* Get current date and time */
	char date[100];
	time_t now = time(NULL);
	struct tm tm;

	gmtime_r(&now, &tm);
	strftime(date, sizeof(date), "%a, %d %b %Y %H:%M:%S GMT", &tm);

	conn->send_len = snprintf(conn->send_buffer, BUFSIZ,
							  "HTTP/1.1 404 Not Found\r\n"
							  "Date: %s\r\n"
							  "Content-Length: 0\r\n"
							  "Connection: close\r\n"
							  "\r\n",
							  date);
}


/* Get resource type depending on request path/filename.
 * Filename should point to the static or dynamic folder.
 */
static enum resource_type connection_get_resource_type(struct connection *conn)
{
	if (strstr(conn->request_path, "/static/"))
		return RESOURCE_TYPE_STATIC;

	if (strstr(conn->request_path, "/dynamic/"))
		return RESOURCE_TYPE_DYNAMIC;

	return RESOURCE_TYPE_NONE;
}


/* Initialize connection structure on given socket. */
struct connection *connection_create(int sockfd)
{
	struct connection *conn = malloc(sizeof(struct connection));

	DIE(conn == NULL, "conn malloc");

	conn->fd = -1;
	memset(conn->filename, 0, BUFSIZ);

	conn->eventfd = -1;
	conn->sockfd = sockfd;

	conn->ctx = ctx;
	conn->file_size = 0;

	memset(conn->recv_buffer, 0, BUFSIZ);
	conn->recv_len = 0;

	memset(conn->send_buffer, 0, BUFSIZ);
	conn->send_len = 0;
	conn->send_pos = 0;
	conn->file_pos = 0;
	conn->async_read_len = 0;

	conn->have_path = 0;
	memset(conn->request_path, 0, BUFSIZ);
	conn->res_type = RESOURCE_TYPE_NONE;
	conn->state = STATE_INITIAL;

	http_parser_init(&conn->request_parser, HTTP_REQUEST);
	conn->request_parser.data = conn;

	return conn;
}

/* Start asynchronous operation (read from file).
 */
void connection_start_async_io(struct connection *conn)
{
	if (conn->eventfd < 0) {
		conn->eventfd = eventfd(0, EFD_NONBLOCK);
		DIE(conn->eventfd < 0, "eventfd");
		w_epoll_add_ptr_in(epollfd, conn->eventfd, conn);
	}

	io_prep_pread(&conn->iocb, conn->fd, conn->send_buffer, BUFSIZ, conn->file_pos);
	io_set_eventfd(&conn->iocb, conn->eventfd);
	conn->piocb[0] = &conn->iocb;

	int rc = io_submit(conn->ctx, 1, conn->piocb);

	DIE(rc < 0, "io_submit");
}


/* Remove connection handler. */
void connection_remove(struct connection *conn)
{
	int rc;

	rc = w_epoll_remove_ptr(epollfd, conn->sockfd, conn);
	DIE(rc < 0, "w_epoll_remove_ptr");

	/* Close fds */
	if (conn->sockfd >= 0) {
		rc = close(conn->sockfd);
		DIE(rc < 0, "close conn sockfd");
	}

	if (conn->fd >= 0) {
		rc = close(conn->fd);
		DIE(rc < 0, "close conn fd");
	}

	if (conn->eventfd >= 0) {
		rc = close(conn->eventfd);
		DIE(rc < 0, "close conn eventfd");
	}

	free(conn);
}


/* Handle a new connection request on the server socket. */
void handle_new_connection(void)
{
	int rc;

	/* Accept new connection. */
	int connectfd = accept(listenfd, NULL, NULL);

	DIE(connectfd < 0, "accept");

	/* Set socket to be non-blocking. */
	fcntl(connectfd, F_SETFL, fcntl(connectfd, F_GETFL, 0) | O_NONBLOCK);

	/* Instantiate new connection handler. */
	struct connection *conn = connection_create(connectfd);

	/* Add socket to epoll. */
	rc = w_epoll_add_ptr_in(epollfd, connectfd, conn);
	DIE(rc < 0, "w_epoll_add_ptr_in");
}

/* Receive message on socket.
 * Store message in recv_buffer in struct connection.
 */
void receive_data(struct connection *conn)
{
	/* Append new data to existing buffer */
	ssize_t bytes_read = recv(conn->sockfd,
							  conn->recv_buffer + conn->recv_len,
							  BUFSIZ - conn->recv_len, 0);

	if (bytes_read < 0) {
		if (errno == EAGAIN || errno == EWOULDBLOCK) {
			/* No data ready yet */
			return;
		}

		conn->state = STATE_CONNECTION_CLOSED;
		return;
	}

	if (bytes_read == 0) {
		conn->state = STATE_CONNECTION_CLOSED;
		return;
	}

	conn->recv_len += bytes_read;
}


/* Open file and update connection fields. */
int connection_open_file(struct connection *conn)
{
	conn->fd = open(conn->filename, O_RDONLY);
	if (conn->fd < 0)
		return -1;

	struct stat stats;
	int rc = fstat(conn->fd, &stats);

	if (rc < 0) {
		close(conn->fd);
		conn->fd = -1;
		return -1;
	}

	conn->file_size = stats.st_size;
	conn->file_pos = 0;

	return 0;
}


/* Complete asynchronous operation; operation returns successfully.
 * Prepare socket for sending.
 */
void connection_complete_async_io(struct connection *conn)
{
	struct io_event events[1];
	int n = io_getevents(conn->ctx, 1, 1, events, NULL);

	if (n <= 0) {
		conn->state = STATE_CONNECTION_CLOSED;
		return;
	}

	conn->async_read_len = events[0].res;

	if (conn->async_read_len <= 0) {
		conn->state = STATE_CONNECTION_CLOSED;
		return;
	}

	conn->file_pos += conn->async_read_len;
	conn->send_pos = 0;

	uint64_t tmp;

	read(conn->eventfd, &tmp, sizeof(tmp));

	conn->state = STATE_SENDING_DATA;

	w_epoll_update_ptr_out(epollfd, conn->sockfd, conn);
}


/* Parse the HTTP header and extract the file path. */
int parse_header(struct connection *conn)
{
	/* Use mostly null settings except for on_path callback. */
	http_parser_settings settings_on_path = {
		.on_message_begin = 0,
		.on_header_field = 0,
		.on_header_value = 0,
		.on_path = aws_on_path_cb,
		.on_url = 0,
		.on_fragment = 0,
		.on_query_string = 0,
		.on_body = 0,
		.on_headers_complete = 0,
		.on_message_complete = 0
	};

	return http_parser_execute(&conn->request_parser, &settings_on_path,
							   conn->recv_buffer, conn->recv_len);
}


/* Send static data using sendfile(2). */
enum connection_state connection_send_static(struct connection *conn)
{
	size_t remaining = conn->file_size - conn->file_pos;

	off_t offset = conn->file_pos;
	ssize_t sent = sendfile(conn->sockfd, conn->fd, &offset, remaining);

	conn->file_pos = offset;

	if (sent < 0) {
		if (errno == EAGAIN || errno == EWOULDBLOCK)
			return STATE_SENDING_DATA;

		return STATE_CONNECTION_CLOSED;
	}

	if (conn->file_pos >= conn->file_size)
		return STATE_CONNECTION_CLOSED;

	return STATE_SENDING_DATA;
}


/* Read data asynchronously. Returns 0 on success and -1 on error. */
int connection_send_dynamic(struct connection *conn)
{
	size_t remaining = conn->async_read_len - conn->send_pos;

	ssize_t sent = send(conn->sockfd, conn->send_buffer + conn->send_pos, remaining, 0);

	if (sent < 0) {
		if (errno == EAGAIN || errno == EWOULDBLOCK)
			return 0;

		return -1;
	}

	conn->send_pos += sent;

	if (conn->send_pos >= conn->async_read_len) {
		if (conn->file_pos >= conn->file_size) {
			/* File sent completely */
			conn->state = STATE_CONNECTION_CLOSED;
			return 0;
		}

		/* More data to read, start next async read */
		conn->state = STATE_ASYNC_ONGOING;
		connection_start_async_io(conn);
	}

	return 0;
}


/* Handle input information: may be a new message or notification of
 * completion of an asynchronous I/O operation.
 */
void handle_input(struct connection *conn)
{
	switch (conn->state) {
	case STATE_INITIAL:
		conn->state = STATE_RECEIVING_DATA;
		break;

	case STATE_RECEIVING_DATA:
		receive_data(conn);

		if (conn->recv_len > 0) {
			/* Check if we have complete HTTP header (ends with \r\n\r\n) */
			if (memmem(conn->recv_buffer, conn->recv_len, "\r\n\r\n", 4)) {
				parse_header(conn);
				if (conn->have_path) {
					conn->state = STATE_REQUEST_RECEIVED;
					handle_input(conn);
				}
			}
		}

		break;

	case STATE_REQUEST_RECEIVED:
		/* Get filename */
		snprintf(conn->filename, BUFSIZ, "%s%s",
				 AWS_DOCUMENT_ROOT, conn->request_path + 1);

		/* Get resource type */
		conn->res_type = connection_get_resource_type(conn);

		int rc = connection_open_file(conn);

		if (rc >= 0) {
			/* Success 200 */
			connection_prepare_send_reply_header(conn);
			conn->send_pos = 0;
			conn->state = STATE_SENDING_HEADER;
		} else {
			/* Error 404 */
			connection_prepare_send_404(conn);
			conn->send_pos = 0;
			conn->state = STATE_SENDING_404;
		}

		w_epoll_update_ptr_out(epollfd, conn->sockfd, conn);

		break;
	case STATE_ASYNC_ONGOING:
		connection_complete_async_io(conn);

		break;

	default:
		printf("shouldn't get here %d\n", conn->state);
	}
}

/* Handle output information: may be a new valid requests or notification of
 * completion of an asynchronous I/O operation or invalid requests.
 */
void handle_output(struct connection *conn)
{
	size_t remaining;
	ssize_t sent;

	switch (conn->state) {
	case STATE_SENDING_HEADER:
		remaining = conn->send_len - conn->send_pos;

		sent = send(conn->sockfd, conn->send_buffer + conn->send_pos, remaining, 0);

		if (sent < 0) {
			if (errno == EAGAIN || errno == EWOULDBLOCK)
				return;

			conn->state = STATE_CONNECTION_CLOSED;
			return;
		}

		conn->send_pos += sent;

		/* Header is sent */
		if (conn->send_pos == conn->send_len) {
			if (conn->res_type == RESOURCE_TYPE_STATIC) {
				conn->state = STATE_SENDING_DATA;
				w_epoll_update_ptr_out(epollfd, conn->sockfd, conn);
			} else if (conn->res_type == RESOURCE_TYPE_DYNAMIC) {
				conn->state = STATE_ASYNC_ONGOING;
				connection_start_async_io(conn);
			} else {
				// 404
				conn->state = STATE_CONNECTION_CLOSED;
			}
		}

		break;

	case STATE_SENDING_404:
		remaining = conn->send_len - conn->send_pos;

		sent = send(conn->sockfd, conn->send_buffer + conn->send_pos, remaining, 0);

		if (sent < 0) {
			if (errno == EAGAIN || errno == EWOULDBLOCK)
				return;

			conn->state = STATE_CONNECTION_CLOSED;
			return;
		}

		conn->send_pos += sent;
		if (conn->send_pos == conn->send_len)
			conn->state = STATE_CONNECTION_CLOSED;

		break;

	case STATE_SENDING_DATA:
		if (conn->res_type == RESOURCE_TYPE_STATIC) {
			conn->state = connection_send_static(conn);

			if (conn->state == STATE_SENDING_DATA)
				w_epoll_update_ptr_out(epollfd, conn->sockfd, conn);
		} else if (conn->res_type == RESOURCE_TYPE_DYNAMIC) {
			int rc = connection_send_dynamic(conn);

			if (rc < 0)
				conn->state = STATE_CONNECTION_CLOSED;

			if (conn->state == STATE_SENDING_DATA)
				w_epoll_update_ptr_out(epollfd, conn->sockfd, conn);
		}

		break;

	case STATE_ASYNC_ONGOING:
		/* Waiting for async I/O to complete, ignore output events */
		break;

	default:
		ERR("Unexpected state\n");
		exit(1);
	}
}


/* Handle new client. There can be input and output connections.
 * Take care of what happened at the end of a connection.
 */
void handle_client(uint32_t event, struct connection *conn)
{
	if (event & EPOLLIN)
		handle_input(conn);

	if (event & EPOLLOUT)
		handle_output(conn);

	if (conn->state == STATE_CONNECTION_CLOSED)
		connection_remove(conn);
}


int main(void)
{
	int rc;

	/* Initialize asynchronous operations. */
	rc = io_setup(10, &ctx);
	DIE(rc < 0, "io_setup");

	/* Initialize multiplexing. */
	epollfd = w_epoll_create();
	DIE(epollfd < 0, "w_epoll_create");

	/* Create server socket. */
	listenfd = tcp_create_listener(AWS_LISTEN_PORT, 10);
	DIE(listenfd < 0, "tcp_create_listener");

	/* Add server socket to epoll object*/
	rc = w_epoll_add_fd_in(epollfd, listenfd);
	DIE(rc < 0, "w_epoll_add_fd_in");

	dlog(LOG_INFO, "Server waiting for connections on port %d\n", AWS_LISTEN_PORT);

	/* server main loop */
	while (1) {
		struct epoll_event rev;

		/* Wait for events. */
		int n = w_epoll_wait_infinite(epollfd, &rev);

		DIE(n < 0, "w_epoll_wait_infinite");

		/* Switch event types; consider
		 *   - new connection requests (on server socket)
		 *   - socket communication (on connection sockets)
		 */
		struct connection *conn = (struct connection *)rev.data.ptr;

		if (rev.data.fd == listenfd) {
			/* New connection */
			handle_new_connection();
		} else if (rev.data.fd == conn->eventfd) {
			/* Async I/O */
			handle_input(conn);
		} else {
			/* Socket */
			handle_client(rev.events, conn);
		}
	}

	return 0;
}