/*
 * s_client command
 *
 * Copyright (C) 2019  Patrick McDermott
 *
 * This file is part of wolfssl-util.
 *
 * wolfssl-util 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.
 *
 * wolfssl-util 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 wolfssl-util.  If not, see <http://www.gnu.org/licenses/>.
 */

#include <errno.h>
#include <netdb.h>
#include <poll.h>
#include <stdbool.h>
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <sys/socket.h>
#include <unistd.h>

#include <wolfssl/ssl.h>
#include <wolfssl/wolfcrypt/settings.h>

#include "commands.h"

#define CA_CERTS "/etc/ssl/certs"

#define ARRAY_SIZE(a)  (sizeof(a) / sizeof((a)[0]))
#define MAX(a, b)      (((a) > (b)) ? (a) : (b))
#define MIN(a, b)      (((a) < (b)) ? (a) : (b))

static _Bool
parse_host_port(char *hostport, char **host, char **port)
{
	*host = hostport;
	/* XXX: Port is required.  Otherwise, this will mangle IPv6 addresses
	 * without some more intelligent (and larger) code. */
	*port = strrchr(*host, ':');
	if (!*port) {
		fputs("Port is required\n", stderr);
		return false;
	}
	**port = '\0';
	++*port;

	return true;
}

static int
connect_socket(const char *host, const char *port)
{
	struct addrinfo  hints = {0};
	struct addrinfo *result;
	int              s;
	struct addrinfo *rp;
	int              sfd;

	hints.ai_family   = AF_UNSPEC;
	hints.ai_socktype = SOCK_STREAM;
	hints.ai_flags    = AI_PASSIVE;
	hints.ai_protocol = 0;

	s = getaddrinfo(host, port, &hints, &result);
	if (s != 0) {
		fprintf(stderr, "Failed to resolve host and port: %s\n",
				gai_strerror(s));
		return -1;
	}

	for (rp = result; rp != NULL; rp = rp->ai_next) {
		sfd = socket(rp->ai_family, rp->ai_socktype, rp->ai_protocol);
		if (sfd == -1) {
			continue;
		}
		if (connect(sfd, rp->ai_addr, rp->ai_addrlen) == -1) {
			close(sfd);
			continue;
		}
		break;
	}
	if (rp == NULL) {
		fputs("Failed to connect\n", stderr);
		sfd = -1;
	}

	freeaddrinfo(result);

	return sfd;
}

static _Bool
write_all(int fd, const void *buf, size_t count)
{
	ssize_t ret;

	while (count > 0) {
		while ((ret = write(fd, buf, count)) < 0 && errno == EINTR) {
			continue;
		}
		if (ret < 0) {
			return false;
		}
		buf = ((const char *) buf) + ret;
		count -= ret;
	}

	return true;
}

static _Bool
poll_fds(int sfd, WOLFSSL *ssl)
{
	struct pollfd fds[2] = {
		{ .fd = -1, .events = POLLIN|POLLERR, .revents = 0 },
		{ .fd = -1, .events = POLLIN|POLLERR, .revents = 0 },
	};
	char          buf[MAX(8192, WOLFSSL_MAX_ERROR_SZ)];
	ssize_t       len;
	int           ret;

	fds[0].fd = STDIN_FILENO;
	fds[1].fd = sfd;

	for (;;) {
		while (poll(fds, ARRAY_SIZE(fds), -1) < 0 && (errno == EINTR ||
					errno == EAGAIN)) {
			continue;
		}
		if (fds[0].revents > 0) {  /* stdin */
			len = read(STDIN_FILENO, buf, sizeof(buf));
			if (len < 0) {
				fputs("Input read error\n", stderr);
				return false;
			} else if (len == 0) {
				fds[0].fd = -1;  /* Stop polling. */
			} else if ((ret = wolfSSL_write(ssl, buf, len)) <= 0) {
				wolfSSL_ERR_error_string(wolfSSL_get_error(ssl,
							ret), buf);
				fprintf(stderr, "Socket write error: %s\n",
						buf);
				return false;
			}
		}
		if (fds[1].revents > 0) {  /* socket */
			ret = wolfSSL_read(ssl, buf, MIN(sizeof(buf), 1024));
			if (ret < 0) {
				wolfSSL_ERR_error_string(wolfSSL_get_error(ssl,
							ret), buf);
				fprintf(stderr, "Socket read error: %s\n",
						buf);
				return false;
			} else if (ret == 0) {
				fds[1].fd = -1;  /* Stop polling. */
				close(STDOUT_FILENO);  /* Signal socket EOF. */
			} else if (write_all(STDOUT_FILENO, buf, (size_t) ret)
					== false) {
				return false;
			}
		}
		if (fds[0].fd == fds[1].fd) {  /* Both -1 (no longer polled) */
			return true;
		}
	}

	/* Unreached */
}

int
s_client(int argc, char **argv)
{
	char           *host       = NULL;
	char           *port       = NULL;
	const char     *servername = NULL;
	int             ret        = EXIT_SUCCESS;
	WOLFSSL_METHOD *method;
	WOLFSSL_CTX    *ctx;
	WOLFSSL        *ssl;
	int             sfd;

	for (; argc > 0; --argc, ++argv) {
		if (strcmp(*argv, "-quiet") == 0) {
			/* No-op */
		} else if (strcmp(*argv, "-connect") == 0) {
			--argc, ++argv;
			if (parse_host_port(*argv, &host, &port) == false) {
				return EXIT_FAILURE;
			}
		} else if (strcmp(*argv, "-servername") == 0) {
			--argc, ++argv;
			servername = *argv;
		} else if (**argv == '-') {
			fprintf(stderr, "Unsupported option \"%s\"\n", *argv);
		}
	}

	wolfSSL_Init();

	method = wolfSSLv23_client_method();
	if (method == NULL) {
		fputs("Out of memory\n", stderr);
		ret = EXIT_FAILURE;
		goto cleanup;
	}

	ctx = wolfSSL_CTX_new(method);
	if (ctx == NULL) {
		fputs("Out of memory\n", stderr);
		ret = EXIT_FAILURE;
		goto cleanup;
	}

	if (wolfSSL_CTX_load_verify_locations_ex(ctx, NULL, CA_CERTS,
				WOLFSSL_LOAD_FLAG_IGNORE_ERR) !=
			WOLFSSL_SUCCESS) {
		fputs("Failed to load CA certificates\n", stderr);
		ret = EXIT_FAILURE;
		goto ctx_free;
	}

#ifdef HAVE_OCSP
	if (wolfSSL_CTX_EnableOCSP(ctx, WOLFSSL_OCSP_CHECKALL) !=
			WOLFSSL_SUCCESS) {
		fputs("Failed to enable OCSP\n", stderr);
		goto ctx_free;
	}
#endif
#ifdef HAVE_SNI
	if (servername != NULL) {
		if (wolfSSL_CTX_UseSNI(ctx, WOLFSSL_SNI_HOST_NAME, servername,
					strlen(servername)) != WOLFSSL_SUCCESS){
			fputs("Out of memory\n", stderr);
			ret = EXIT_FAILURE;
			goto ctx_free;
		}
	}
#else
	(void) servername;
#endif

	ssl = wolfSSL_new(ctx);
	if (ssl == NULL) {
		fputs("Out of memory\n", stderr);
		ret = EXIT_FAILURE;
		goto ctx_free;
	}

	sfd = connect_socket(host, port);
	if (sfd == -1) {
		ret = EXIT_FAILURE;
		goto ssl_free;
	}
	wolfSSL_set_fd(ssl, sfd);

	if (poll_fds(sfd, ssl) == false) {
		ret = EXIT_FAILURE;
	}

	close(sfd);

ssl_free:
	wolfSSL_free(ssl);
ctx_free:
	wolfSSL_CTX_free(ctx);
cleanup:
	wolfSSL_Cleanup();

	return ret;
}