/*
 * Forked from OpenSSL sslecho demo
 *
 * OpenSSL: gcc server.c -o server -lssl -lcrypto
 * AWSLC: gcc server.c -DUSE_AWSLC -I/opt/awslc-1.29.0/include/ -o server -L/opt/awslc-1.29.0/lib/ -lssl -lcrypto -Wl,-rpath,/opt/awslc-1.29.0/lib/
 *
 * Test:
 *   ./server file.crt file.key
 *    openssl s_client -connect localhost:4433
 *
 */


#include <openssl/ssl.h>
#include <openssl/err.h>
#include <openssl/bio.h>


#include <stdbool.h>
#include <stdio.h>
#include <string.h>
#include <signal.h>
#include <fcntl.h>
#include <unistd.h>
#include <sys/socket.h>
#include <arpa/inet.h>
#include <netinet/in.h>

static const int server_port = 4433;

char *file_crt = NULL;
char *file_key = NULL;

static volatile bool    server_running = true;

#if defined(USE_AWSLC)
int switchctx_cbk(const struct ssl_early_callback_ctx *ctx)
{
	SSL *ssl = ctx->ssl;
#else
int switchctx_cbk(SSL *ssl, int *al, void *arg)
{
#endif
	const SSL_CIPHER *cipher;
	STACK_OF(SSL_CIPHER) *ha_ciphers; /* haproxy side ciphers */
	uint32_t cipher_id;
	size_t len;
	const uint8_t *cipher_suites;

	ha_ciphers = SSL_get_ciphers(ssl);
	for (size_t i = 0; i < sk_SSL_CIPHER_num(ha_ciphers); i++) {
		const SSL_CIPHER *c = sk_SSL_CIPHER_value(ha_ciphers, i);
		fprintf(stderr, "%s\n", SSL_CIPHER_get_name(c));
	}
	return 1;
}


static int create_socket(bool isServer)
{
	int s;
	int optval = 1;
	struct sockaddr_in addr;

	s = socket(AF_INET, SOCK_STREAM, 0);
	if (s < 0) {
		perror("Unable to create socket");
		exit(EXIT_FAILURE);
	}

	if (isServer) {
		addr.sin_family = AF_INET;
		addr.sin_port = htons(server_port);
		addr.sin_addr.s_addr = INADDR_ANY;

		/* Reuse the address; good for quick restarts */
		if (setsockopt(s, SOL_SOCKET, SO_REUSEADDR, &optval, sizeof(optval))
				< 0) {
			perror("setsockopt(SO_REUSEADDR) failed");
			exit(EXIT_FAILURE);
		}

		if (bind(s, (struct sockaddr*) &addr, sizeof(addr)) < 0) {
			perror("Unable to bind");
			exit(EXIT_FAILURE);
		}

		if (listen(s, 1) < 0) {
			perror("Unable to listen");
			exit(EXIT_FAILURE);
		}
	}

	return s;
}

static SSL_CTX* create_context()
{
	const SSL_METHOD *method;
	SSL_CTX *ctx;

	method = TLS_server_method();

	ctx = SSL_CTX_new(method);
	if (ctx == NULL) {
		perror("Unable to create SSL context");
		ERR_print_errors_fp(stderr);
		exit(EXIT_FAILURE);
	}

#if defined(USE_AWSLC)
	SSL_CTX_set_select_certificate_cb(ctx, switchctx_cbk);
#else
	SSL_CTX_set_client_hello_cb(ctx, switchctx_cbk, NULL);
#endif
	return ctx;
}

static void configure_server_context(SSL_CTX *ctx)
{
	/* Set the key and cert */
	if (SSL_CTX_use_certificate_chain_file(ctx, file_crt) <= 0) {
		ERR_print_errors_fp(stderr);
		exit(EXIT_FAILURE);
	}

	if (SSL_CTX_use_PrivateKey_file(ctx, file_key, SSL_FILETYPE_PEM) <= 0) {
		ERR_print_errors_fp(stderr);
		exit(EXIT_FAILURE);
	}
}

static void usage(char *argv[])
{
	printf("Usage: %s <crt> <key>\n", argv[0]);
	exit(EXIT_FAILURE);
}


#define BUFFERSIZE 1024
int main(int argc, char **argv)
{
	bool isServer = true;
	int result;

	SSL_CTX *ssl_ctx = NULL;
	SSL *ssl = NULL;

	int server_skt = -1;
	int client_skt = -1;

	/* used by fgets */
	char buffer[BUFFERSIZE];
	char *txbuf;

	char rxbuf[128];
	size_t rxcap = sizeof(rxbuf);
	int rxlen;

	char *rem_server_ip = NULL;

	struct sockaddr_in addr;
	unsigned int addr_len = sizeof(addr);

	/* ignore SIGPIPE so that server can continue running when client pipe closes abruptly */
	signal(SIGPIPE, SIG_IGN);

	/* Splash */
	printf("\nsslecho : Simple Echo Client/Server : %s : %s\n\n", __DATE__,  __TIME__);

	/* Need to know if client or server */
	if (argc < 3) {
		usage(argv);
		/* NOTREACHED */
	}

	file_crt = argv[1];
	file_key = argv[2];

	/* Create context used by both client and server */
	ssl_ctx = create_context();

	printf("We are the server on port: %d\n\n", server_port);

	/* Configure server context with appropriate key files */
	configure_server_context(ssl_ctx);

	/* Create server socket; will bind with server port and listen */
	server_skt = create_socket(true);

	/*
	 * Loop to accept clients.
	 * Need to implement timeouts on TCP & SSL connect/read functions
	 * before we can catch a CTRL-C and kill the server.
	 */
	while (server_running) {
		/* Wait for TCP connection from client */
		client_skt = accept(server_skt, (struct sockaddr*) &addr,
				&addr_len);
		if (client_skt < 0) {
			perror("Unable to accept");
			goto error;
		}

		printf("Client TCP connection accepted\n");

		/* Create server SSL structure using newly accepted client socket */
		ssl = SSL_new(ssl_ctx);
		if (!SSL_set_fd(ssl, client_skt)) {
			ERR_print_errors_fp(stderr);
			goto error;
		}

		/* Wait for SSL connection from the client */
		if (SSL_accept(ssl) <= 0) {
			ERR_print_errors_fp(stderr);
			goto error;
		} else {

			printf("Client SSL connection accepted\n\n");

			/* Echo loop */
			while (true) {
				/* Get message from client; will fail if client closes connection */
				if ((rxlen = SSL_read(ssl, rxbuf, rxcap)) <= 0) {
					if (rxlen == 0) {
						printf("Client closed connection\n");
					} else {
						printf("SSL_read returned %d\n", rxlen);
					}
					ERR_print_errors_fp(stderr);
					break;
				}
				/* Insure null terminated input */
				rxbuf[rxlen] = 0;
				/* Look for kill switch */
				if (strcmp(rxbuf, "kill\n") == 0) {
					/* Terminate...with extreme prejudice */
					printf("Server received 'kill' command\n");
					server_running = false;
					break;
				}
				/* Show received message */
				printf("Received: %s", rxbuf);
				/* Echo it back */
				if (SSL_write(ssl, rxbuf, rxlen) <= 0) {
					ERR_print_errors_fp(stderr);
					goto error;
				}
			}
		}
error:
		if (server_running) {
			/* Cleanup for next client */
			SSL_shutdown(ssl);
			SSL_free(ssl);
			close(client_skt);
		}
	}
	printf("Server exiting...\n");
exit:
	/* Close up */
	if (ssl != NULL) {
		SSL_shutdown(ssl);
		SSL_free(ssl);
	}
	SSL_CTX_free(ssl_ctx);

	if (server_skt != -1)
		close(server_skt);

	return EXIT_SUCCESS;
}