/* multiplex.h */

#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>

#include <sys/time.h>

#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <fcntl.h>
#include <unistd.h>
#include <errno.h>
#include <err.h>
#include <error.h>

#include <syslog.h>

#include <event.h>
#include <config/parse.h>
#include <stdarg.h>

/* This is a fairly large buffer, but it appears to be the only way
 * to keep libevent's buffered events sustained. And to avoid buffer
 * underruns. Libevent REALLY doesn't like buffer underruns. This size
 * allows for at least 160Mbit/s */
#define BUFFERSIZE 102400 
#define LOGBUFFERSIZE 100
#define CFGFILENAME SYSCONFIGDIR"/multiplex.conf"

struct client { 
	int client_fd; 
	int server_fd; 
	int first_packet; 
	struct sockaddr_in client_addr; 
	struct bufferevent *client_ev; 
	struct bufferevent *server_ev; 
	struct event timeout; 
	struct timeval tv; 
}; 

int setnonblock(int fd);
void log_message(int level, const char *message, ...);
int connect_client(struct client* myclient, int port);
void parse_config();
void on_server_read(struct bufferevent *bev, void *arg);
void on_write(struct bufferevent *bev, void *arg);
void on_error(struct bufferevent *bev, short what, void *arg);
void on_client_read(struct bufferevent *bev, void *arg);
void on_timeout(int fd, short event, void *arg);
void on_accept(int fd, short ev, void *arg

/***************************************************************************************/

/* multiplex.c
 *
 * Multiplex - a port multiplexer for SSH and HTTPS 
 * Copyright (C) 2007 Hein-Pieter van Braam 
 *
 * This program is distributed under the terms of the GNU GPL version 2.
 * See the file COPYING.
 * 
 * compile with
 * gcc multiplex.c -o multiplex -levent -ldebug -lconfig
 * 
*/

#include "multiplex.h"

/* some configuratoin globals */
struct conf {
	int https_port;
	int ssh_port;
	int listen_port;
	int daemonize;
	int client_delay;
} conf;

int setnonblock(int fd){
	int flags;

	flags = fcntl(fd, F_GETFL);
	if (flags < 0)
		return flags;
	flags |= O_NONBLOCK;
	if (fcntl(fd, F_SETFL, flags) < 0)
		return -1;

	return 0;
}

void log_message(int level, const char *message, ...) {
        char buf[LOGBUFFERSIZE];
        va_list va;
        va_start(va, message);
        vsnprintf(buf, LOGBUFFERSIZE, message, va);
        va_end(va);
        printf("WARNING: %i: %s", level, buf);
}

int connect_client(struct client* myclient, int port){
	int server_fd;
	struct sockaddr_in server_addr;
	int reuseaddr_on;

	server_fd = socket(AF_INET, SOCK_STREAM, 0);
	if (server_fd < 0)
		log_message(LOG_WARNING, "Listen failed.\n");
	memset(&server_addr, 0, sizeof(server_addr));
	server_addr.sin_family = AF_INET;
	server_addr.sin_addr.s_addr = INADDR_ANY;
	server_addr.sin_port = htons(port);
	reuseaddr_on = 1;

	if (connect (server_fd, (struct sockaddr *) &server_addr, sizeof (server_addr)) < 0){
		log_message(LOG_WARNING, "Can't connect to address: %i\n", port);
		bufferevent_free(myclient->client_ev);
		close(myclient->client_fd);
		log_message(LOG_NOTICE, "Disconnected client %s\n", inet_ntoa(myclient->client_addr.sin_addr));
		free(myclient);
		return -1;
	}

	myclient->server_fd = server_fd;
	myclient->server_ev = bufferevent_new(myclient->server_fd, on_server_read, on_write, on_error, myclient);
	bufferevent_enable(myclient->server_ev, EV_READ);
	myclient->first_packet = 0;
	return server_fd;
}

void on_write(struct bufferevent *bev, void *arg){
}

void on_error(struct bufferevent *bev, short what, void *arg){
	struct client* myclient = arg;

	if (what & EVBUFFER_EOF) {
		log_message(LOG_NOTICE, "Disconnected client %s\n", inet_ntoa(myclient->client_addr.sin_addr));
	}
	else {
		log_message(LOG_NOTICE, "Socket error on client %s\n", inet_ntoa(myclient->client_addr.sin_addr));
	}

	/* If we are still at the first packet, make sure we clean up the timeout event */
	if (myclient->first_packet == 1) 
		event_del(&myclient->timeout);

	bufferevent_free(myclient->client_ev);
	close(myclient->client_fd);

	if (myclient->server_fd > 0){
		bufferevent_free(myclient->server_ev);
		close(myclient->server_fd);
	}


	free(myclient);
}

void on_client_read(struct bufferevent *bev, void *arg){
	struct client* myclient = arg;
	unsigned char buf[BUFFERSIZE];
	int n, retval;
	
	/* If this is the fist packet, that means the client send us something. This means that it is an HTTP client */
	if (myclient->first_packet == 1) {
		event_del(&myclient->timeout);
		retval = connect_client(myclient, conf.https_port);
		if (retval == -1){
			return;
		}
	}

	n = bufferevent_read(myclient->client_ev, &buf, BUFFERSIZE);
	bufferevent_write(myclient->server_ev, buf, n);
}

void on_server_read(struct bufferevent *bev, void *arg){
	struct client* myclient = arg;
	unsigned char buf[BUFFERSIZE];
	int n;

	n = bufferevent_read(myclient->server_ev, &buf, BUFFERSIZE);
	//printf("Got %i bytes from server %i\n", n, myclient->server_fd);
	bufferevent_write(myclient->client_ev, buf, n);
	//printf("wrote %i bytes to client %i\n", n, myclient->server_fd);
}


/* Callback for the conf.client_delay timeout 
 * The theory is as follows: if the client hasn't responded with anything in 2 seconds
 * it is probably waiting for the server. Now, SSH server sends an hello string, HTTP(s) does not. 
 * So, if no response from the client in 2 seconds, we just assume it is waiting for an SSH hello*/
void on_timeout(int fd, short event, void *arg){
	struct client* myclient = arg;
	int retval;

	/* Race here... client might have disconnected during the timeout */
	if (myclient->first_packet == 1){
		retval = connect_client(myclient, conf.ssh_port);
		if (retval == -1){
			event_del(&myclient->timeout);
			return;
		}
	}
	event_del(&myclient->timeout);
}

/* Callback */
void on_accept(int fd, short ev, void *arg){
	int client_fd;
	struct sockaddr_in client_addr;
	socklen_t client_len = sizeof(client_addr);
	struct client *client;

	client_fd = accept(fd, (struct sockaddr *)&client_addr, &client_len);
	if (client_fd < 0) {
		log_message(LOG_WARNING, "accept failed");
		return;
	}

	if (setnonblock(client_fd) < 0)
		log_message(LOG_WARNING, "failed to set client socket non-blocking");

	client = calloc(1, sizeof(*client));
	if (client == NULL)
		err(1, "malloc failed");

	/* Set a timer */
	evtimer_set(&client->timeout, on_timeout, client);
	timerclear(&client->tv);
	client->tv.tv_sec = conf.client_delay; 
	event_add(&client->timeout, &client->tv);

	client->first_packet = 1;
	client->client_fd = client_fd;
	client->client_ev = bufferevent_new(client->client_fd, on_client_read, on_write, on_error, client);
	client->client_addr = client_addr;
	client->server_fd = 0;

	bufferevent_enable(client->client_ev, EV_READ);
	log_message(LOG_NOTICE, "Accepted connection from %s\n", inet_ntoa(client->client_addr.sin_addr));
}

void parse_config(){
	section_t *section;
	stmt_t *cur, *prev;

	conf.https_port = 0;
	conf.ssh_port = 0;
	conf.listen_port = 0;
	conf.client_delay = 2;
	conf.daemonize = 1;

	if ((section = parse (CFGFILENAME)) == NULL){
		error(1, 0, "Config error: Failed to open config file: %s", CFGFILENAME);
	}
	
	prev = section->stmt;
	do {
		cur = section->stmt->next; 
		if (! strncmp(cur->name, "listen_port", 11)){
			conf.listen_port = cur->args[0].integer;
		}
		if (! strncmp(cur->name, "https_port", 10)){
			conf.https_port = cur->args[0].integer;
		}
		if (! strncmp(cur->name, "ssh_port", 8)){
			conf.ssh_port = cur->args[0].integer;
		}
		if (! strncmp(cur->name, "client_delay", 12)){
			conf.client_delay = cur->args[0].integer;
		}
		if (! strncmp(cur->name, "daemonize", 9)){
			conf.daemonize = cur->args[0].integer;
		}
		section->stmt = cur;
	} while ( prev != section->stmt );

	parse_destroy (&section);

	if (conf.listen_port == 0)
		error(1, 0, "Config error: listen_port not defined");
	if (conf.https_port == 0)
		error(1, 0, "Config error: https_port not defined");
	if (conf.ssh_port == 0)
		error(1, 0, "Config error: ssh_port not defined");
	if (conf.client_delay == 0)
		error(1, 0, "Config error: client_delay not defined");
}

int main(int argc, char **argv){
	struct event ev_accept;
	struct sockaddr_in listen_addr;
	int reuseaddr_on, listen_fd, pid;

	parse_config();

	/* Initialize our listening socket, largely boilerplate */
	memset(&listen_addr, 0, sizeof(listen_addr));
	listen_addr.sin_family = AF_INET;
	listen_addr.sin_addr.s_addr = INADDR_ANY;
	listen_addr.sin_port = htons(conf.listen_port);
	reuseaddr_on = 1;
	listen_fd = socket(AF_INET, SOCK_STREAM, 0);

	if (listen_fd < 0)
		err(1, "listen failed");
	if (setsockopt(listen_fd, SOL_SOCKET, SO_REUSEADDR, &reuseaddr_on, sizeof(reuseaddr_on)) < 0)
		err(1, "setsockopts failed");
	if (bind(listen_fd, (struct sockaddr *)&listen_addr, sizeof(listen_addr)) < 0)
		err(1, "bind failed");
	if (listen(listen_fd, 5) < 0)
		err(1, "listen failed");
	if (setnonblock(listen_fd) < 0)
		err(1, "failed to set server socket to non-blocking");

	if (conf.daemonize == 1){
		pid = fork();
		if (pid < 0) {
			error(1, 0, "Failed to fork()");
		} else if (pid > 0) {
			exit(EXIT_SUCCESS);
		}
	}

	/* Initialize libevent and the new client event (on_accept) */
	event_init();
	event_set(&ev_accept, listen_fd, EV_READ|EV_PERSIST, on_accept, NULL);
	event_add(&ev_accept, NULL);

	/* libevent main loop */
	event_dispatch();

	return 0;
}

);