/*
 * Copyright (C) 2022 Camden Dixie O'Brien
 *
 * This program is free software: you can redistribute it and/or
 * modify it under the terms of the GNU Affero General Public License
 * as published by the Free Software Foundation, either version 3 of
 * the License, or (at your option) any later version.
 *
 * This program 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
 * Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public
 * License along with this program.  If not, see
 * <https://www.gnu.org/licenses/>.
 */

#include <dirent.h>
#include <errno.h>
#include <fcntl.h>
#include <netinet/in.h>
#include <signal.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <sys/stat.h>
#include <unistd.h>

#define HOST "::1"
#define PORT 7070

#define RESP_TERM ".\r\n"
#define RESP_TERM_LEN 3 /* strlen(".\r\n") */

#define PBUF_SIZE 1024
#define RBUF_SIZE 1024
#define SBUF_SIZE 1024
#define FBUF_SIZE 1024

/*
 * Enumeration of directory entry types.
 *
 * The numeric values are significant here (hence them being set
 * explicitly): they match those defined in RFC 1436.
 */
enum etype { ETYPE_FILE = 0, ETYPE_DIR = 1 };

static bool exit_requested = false;

void handle_exit_signal(int signum)
{
	(void)signum;
	exit_requested = true;
}

static int retrying_write(int fd, const char *buf, size_t len)
{
	int n;
	do {
		errno = 0;
		n = write(fd, buf, len);
	} while (errno == EINTR);
	return n;
}

static int append_selector_part(const char *part, size_t partlen, char *buf,
                                size_t bufsize)
{
	if (partlen == 0)
		return 0;

	if (strncmp(part, "..", partlen) == 0) {
		fprintf(stderr, "Selector string contains \"..\"\n");
		return -1;
	} else if (partlen + 1 > bufsize) {
		fprintf(stderr, "Path buffer too small\n");
		return -1;
	}

	*buf++ = '/';
	memcpy(buf, part, partlen);

	return partlen + 1;
}

int main(int argc, char *argv[])
{
	(void)argc;
	(void)argv;

	int res = EXIT_SUCCESS;
	static char pbuf[PBUF_SIZE], rbuf[RBUF_SIZE], sbuf[SBUF_SIZE],
	    fbuf[FBUF_SIZE];

	/*
	 * Get srvroot path from arguments and copy into pbuf.
	 *
	 * The srvroot being at the start of pbuf should be maintained
	 * through the whole application. Any trailing '/' is removed if
	 * present.
	 */
	if (argc < 2) {
		fprintf(stderr, "Usage: %s srvroot\n", argv[0]);
		return EXIT_FAILURE;
	}
	size_t srvroot_len = strlen(argv[1]);
	if (argv[1][srvroot_len - 1] == '/')
		--srvroot_len;
	if (srvroot_len > PBUF_SIZE) {
		fprintf(stderr, "srvroot path is too long\n");
		return EXIT_FAILURE;
	}
	memcpy(pbuf, argv[1], srvroot_len);

	/*
	 * Register signal handler for SIGTERM and SIGINT.
	 *
	 * The behaviour of signal() is not entirely consistent across
	 * different implementations, so this should ideally be rewritten
	 * to use signaction() instead at some point.
	 */
	if (signal(SIGTERM, handle_exit_signal) == SIG_ERR) {
		fprintf(stderr, "Failed to register SIGTERM signal handler\n");
		return EXIT_FAILURE;
	}
	if (signal(SIGINT, handle_exit_signal) == SIG_ERR) {
		fprintf(stderr, "Failed to register SIGINT signal handler\n");
		return EXIT_FAILURE;
	}

	/*
	 * Initialise socket.
	 *
	 * Currently only supporting IPv6, and hard-coding address and
	 * port. Should eventully get address and port from arguments, and
	 * support IPv4 as well.
	 */
	int sfd = socket(AF_INET6, SOCK_STREAM, 0);
	if (sfd == -1) {
		fprintf(stderr, "Failed to open socket\n");
		return EXIT_FAILURE;
	}
	const struct sockaddr_in6 haddr = {
		.sin6_family = AF_INET6,
		.sin6_port = htons(PORT),
		.sin6_addr = IN6ADDR_LOOPBACK_INIT,
	};
	if (bind(sfd, (const struct sockaddr *)&haddr, sizeof(haddr)) == -1) {
		fprintf(stderr, "Error binding socket to address\n");
		res = EXIT_FAILURE;
		goto close_server_socket;
	}
	if (listen(sfd, 8) == -1) {
		fprintf(stderr, "Error attempting to listen on socket\n");
		res = EXIT_FAILURE;
		goto close_server_socket;
	}

	struct sockaddr_in6 paddr;
	socklen_t paddr_size = sizeof(paddr);
	int cfd;
	ssize_t n, slen;
	while (!exit_requested) {
		/*
		 * Accept incoming connection.
		 *
		 * If this is interrupted we go to the next loop iteration
		 * rather than retrying directly in case the received signal
		 * was requesting that the server exits.
		 *
		 * System calls after this point should be retried before
		 * checking if we need to exit, as we don't want to leave the
		 * client hanging.
		 */
		cfd = accept(sfd, (struct sockaddr *)&paddr, &paddr_size);
		if (cfd == -1) {
			if (errno != EINTR)
				fprintf(stderr, "Error accepting connection\n");
			continue;
		}

		/*
		 * Read selector from client socket.
		 *
		 * For now, assuming that the selector is less than RDBUF
		 * bytes long.
		 */
		do {
			errno = 0;
			n = read(cfd, sbuf, SBUF_SIZE);
		} while (errno == EINTR);
		if (n == -1) {
			fprintf(stderr, "Error reading selector from client\n");
			goto close_client_socket;
		}

		/*
		 * Locate the end of the selector.
		 *
		 * The end of the selector string is indicated with CRLF. Most
		 * of the time this will be at the end of the received data,
		 * so might as well start from there.
		 */
		slen = n;
		for (unsigned i = n - 2; i < n; --i) {
			if (sbuf[i] == 0x0d && sbuf[i + 1] == 0x0a)
				slen = i;
		}
		if (slen == n) {
			fprintf(stderr, "Received invalid selector (no CRLF)\n");
			goto close_client_socket;
		}

		/*
		 * Construct the full path to the requested resource.
		 *
		 * This is needed in order to pass it to stat() and fopen() or
		 * opendir() later.  It must be null-terminated.  We also need
		 * to make sure that none of the path's parts are "..", as
		 * this could be used to escape the srvroot directory.
		 */
		char *sp = sbuf;
		unsigned splen = 0, ppos = srvroot_len;
		for (unsigned i = 0; i < slen; ++i) {
			if (sbuf[i] != '/') {
				++splen;
				continue;
			}
			n = append_selector_part(sp, splen, &pbuf[ppos], PBUF_SIZE - ppos);
			if (n == -1)
				goto close_client_socket;
			ppos += n;
			sp = &sbuf[i + 1];
			splen = 0;
		}
		n = append_selector_part(sp, splen, &pbuf[ppos], PBUF_SIZE - ppos);
		if (n == -1)
			goto close_client_socket;
		ppos += n;

		/*
		 * Determine whether the resource is a file or a directory.
		 *
		 * The path passed to stat() has to be null-terminated, which
		 * pbuf is not, so a trailing null byte must first be added
		 * there.
		 */
		if (ppos + 1 > PBUF_SIZE) {
			fprintf(stderr, "Path buffer is too small\n");
			goto close_client_socket;
		}
		pbuf[ppos] = '\0';
		struct stat rstat;
		if (stat(pbuf, &rstat) == -1) {
			fprintf(stderr, "Failed to stat() path \"%s\"\n", pbuf);
			goto close_client_socket;
		}
		if (S_ISREG(rstat.st_mode)) {
			/*
			 * Open the file.
			 *
			 * The path is already null-terminated from the earlier call
			 * to stat().
			 */
			FILE *rf;
			do {
				errno = 0;
				rf = fopen(pbuf, "r");
			} while (errno == EINTR);
			if (rf == NULL) {
				fprintf(stderr, "Failed to open %s\n", pbuf);
				goto close_client_socket;
			}

			/*
			 * Construct response in rbuf.
			 *
			 * The response should be the file contents, but with CRLF
			 * line endings instead of just LF. Replacing the line
			 * endings is done with a buffered reading approach to
			 * minimize syscalls (as opposed to using fgetc()).
			 */
			unsigned rpos = 0;
			while (true) {
				/* Fill fbuf from file. */
				n = fread(fbuf, sizeof(*fbuf), FBUF_SIZE, rf);
				if (n != FBUF_SIZE && ferror(rf)) {
					fprintf(stderr, "Failed to read from file %s\n", pbuf);
					goto close_client_socket;
				}

				/* Copy from fbuf to rbuf, replacing LF with CRLF. */
				for (unsigned fpos = 0; fpos < n; ++fpos) {
					if (fbuf[fpos] == '\n') {
						if (RBUF_SIZE - rpos < 2) {
							fprintf(stderr, "Response buffer is too small");
							goto close_client_socket;
						}
						rbuf[rpos++] = '\r';
						rbuf[rpos++] = '\n';
					} else {
						if (RBUF_SIZE - rpos < 1) {
							fprintf(stderr, "Response buffer is too small");
							goto close_client_socket;
						}
						rbuf[rpos++] = fbuf[fpos];
					}
				}

				if (feof(rf))
					break;
			}

			/*
			 * Terminate and send the response.
			 */
			if (RBUF_SIZE - rpos < RESP_TERM_LEN) {
				fprintf(stderr, "Response buffer is too small");
				goto close_client_socket;
			}
			memcpy(&rbuf[rpos], RESP_TERM, RESP_TERM_LEN);
			rpos += RESP_TERM_LEN;
			if (retrying_write(cfd, rbuf, rpos) == -1) {
				fprintf(stderr, "Failed to write response to client\n");
				goto close_client_socket;
			}
		} else if (S_ISDIR(rstat.st_mode)) {
			/*
			 * Open the directory.
			 *
			 * The path is already null-terminated from the earlier call
			 * to stat().
			 */
			DIR *rdir;
			do {
				errno = 0;
				rdir = opendir(pbuf);
			} while (errno == EINTR);
			if (rdir == NULL) {
				fprintf(stderr, "Failed to open %s\n", pbuf);
				goto close_client_socket;
			}

			/*
			 * Write a line for each entry in the directory to the client.
			 */
			struct dirent *ent;
			unsigned namelen;
			while ((ent = readdir(rdir)) != NULL) {
				/*
				 * Skip . and .. entries.
				 */
				if (strcmp(ent->d_name, ".") == 0
				    || strcmp(ent->d_name, "..") == 0)
					continue;

				/*
				 * Construct full path to entry in pbuf.
				 *
				 * This is needed in order to pass it to stat() later. The
				 * path must be null-terminated.
				 */
				namelen = strlen(ent->d_name);
				if (ppos + namelen + 2 > PBUF_SIZE) {
					fprintf(stderr, "Path buffer is too small\n");
					goto close_client_socket;
				}
				pbuf[ppos] = '/';
				memcpy(&pbuf[ppos + 1], ent->d_name, namelen);
				pbuf[ppos + 1 + namelen] = '\0';

				/*
				 * Identify entry type from file inode information.
				 */
				if (stat(pbuf, &rstat) == -1) {
					fprintf(stderr,
					        "Failed to stat() path \"%s\", skipping entry\n",
					        pbuf);
					continue;
				}
				enum etype type;
				if (S_ISREG(rstat.st_mode))
					type = ETYPE_FILE;
				else if (S_ISDIR(rstat.st_mode))
					type = ETYPE_DIR;
				else
					continue;

				/*
				 * Format and send response line for current entry.
				 */
				n = snprintf(rbuf, RBUF_SIZE, "%1u%s\t%s\t%s\t%u\r\n", type,
				             ent->d_name, &pbuf[srvroot_len], HOST, PORT);
				if (n >= RBUF_SIZE) {
					fprintf(stderr,
					        "Response buffer was too small, skipping entry\n");
					continue;
				}
				if (retrying_write(cfd, rbuf, n) != n) {
					fprintf(stderr, "Error sending respose line to client\n");
					continue;
				}
			}
			if (retrying_write(cfd, ".\r\n", 3) != 3) {
				fprintf(stderr,
				        "Error sending response terminator to client\n");
				goto close_client_socket;
			}

			/*
			 * Close the resource.
			 */
			closedir(rdir);
		} else {
			fprintf(stderr,
			        "Requested resource \"%s\" was not a directory or a "
			        "regular file\n",
			        pbuf);
			goto close_client_socket;
		}

	close_client_socket:
		close(cfd);
	}

close_server_socket:
	close(sfd);
	return res;
}