/*
 * The QUAFFLER library
 * Copyright (c) 2006 Clifford Wolf <clifford@clifford.at>
 *
 * Permission is hereby granted, free of charge, to any person obtaining a
 * copy of this software and associated documentation files (the "Software"),
 * to deal in the Software without restriction, including without limitation
 * the rights to use, copy, modify, merge, publish, distribute, sublicense,
 * and/or sell copies of the Software, and to permit persons to whom the
 * Software is furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included
 * in all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
 * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR
 * OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
 * ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
 * OTHER DEALINGS IN THE SOFTWARE.
 */

#include <stdlib.h>
#include <string.h>
#include <stdio.h>

#include "quaffler.h"

struct qfl_ctx *qfl_ctx_create(int port, int portrange_begin, int portrange_end)
{
	struct qfl_ctx *ctx = calloc(1, sizeof(struct qfl_ctx));

	ctx->listenport = port;
	ctx->portrange_begin = portrange_begin;
	ctx->portrange_end = portrange_end;

	if (qfl_os_listen_open(&ctx->listenport, &ctx->listenhdl, portrange_begin, portrange_end)) {
		free(ctx);
		return 0;
	}

	ctx->rate_up = QFL_DEFAULT_RATE_TOTAL / 2;
	ctx->rate_down = QFL_DEFAULT_RATE_TOTAL / 2;
	ctx->rate_total = QFL_DEFAULT_RATE_TOTAL;
	ctx->last_quota_update = qfl_os_time_ms();

	ctx->start_timestamp = qfl_os_time_ms();

	return ctx;
}

void qfl_ctx_destroy(struct qfl_ctx *ctx)
{
	qfl_os_listen_close(&ctx->listenhdl);

	while (ctx->conn_list)
		qfl_conn_destroy(ctx->conn_list);

	free(ctx);
}

struct qfl_conn *qfl_conn_create(struct qfl_ctx *ctx, const char *peer)
{
	int i;
	struct qfl_conn *conn = calloc(1, sizeof(struct qfl_conn));

	if (peer) {
		if (qfl_os_connection_open(peer, &conn->connhdl, ctx->portrange_begin, ctx->portrange_end)) {
			free(conn);
			return 0;
		}

		conn->initiator = 1;
		conn->crypto_state = 1024;
		conn->peer = strdup(peer);
	}

	for (i=0; i<QFL_FIFO_COUNT; i++)
		conn->fifo[i] = qfl_pkgfifo_create();

	conn->next = ctx->conn_list;
	if (conn->next) conn->next->last = conn;
	ctx->conn_list = conn;

	conn->ctx = ctx;
	conn->upload_priority = 100;

	return conn;
}

void qfl_conn_destroy(struct qfl_conn *conn)
{
	int i;

	if (conn->process_close)
		conn->process_close(conn);

	qfl_os_connection_close(&conn->connhdl);

	free(conn->peer);
	free(conn->mode);

	for (i=0; i<QFL_FIFO_COUNT; i++)
		if (conn->fifo[i])
			qfl_pkgfifo_destroy(conn->fifo[i]);

	free(conn->senddata);

	if (conn->next) conn->next->last = conn->last;
	if (conn->last) conn->last->next = conn->next;
	else conn->ctx->conn_list = conn->next;

	free(conn);
}

void qfl_ctx_io_wait(struct qfl_ctx *ctx, int timeout_ms)
{
	qfl_os_waitfornetwork(ctx, timeout_ms);
}

static void set_quota_on_conn_list(struct qfl_conn *conn, long long quota, int isupload)
{
	int points = 0;
	struct qfl_conn *c;

	for (c=conn; c; c=c->next)
		if (c->needs_quota)
			points += isupload ? c->upload_priority + 1 : 1;

	float quota_per_point = (float)quota / (float)points;

	for (c=conn; c; c=c->next)
		if (c->needs_quota)
			c->quota = (isupload ? c->upload_priority + 1 : 1) * quota_per_point + 1;
		else
			c->quota = 0;
}

void qfl_ctx_io_process(struct qfl_ctx *ctx)
{
	unsigned long long now = qfl_os_time_ms();
	struct qfl_conn *conn;
	int i, didsomething;

	// any new connections?
	while ((conn = qfl_os_accept(ctx)) != 0)
		if (ctx->process_newconn) ctx->process_newconn(conn);

	// set new quotas
	ctx->quota_up += ctx->rate_up * (now - ctx->last_quota_update);
	ctx->quota_down += ctx->rate_down * (now - ctx->last_quota_update);
	ctx->quota_total += ctx->rate_total * (now - ctx->last_quota_update);
	if (ctx->quota_up > ctx->rate_up * 1000) ctx->quota_up = ctx->rate_up * 1000;
	if (ctx->quota_down > ctx->rate_down * 1000) ctx->quota_down = ctx->rate_down * 1000;
	if (ctx->quota_total > ctx->rate_total * 1000) ctx->quota_total = ctx->rate_total * 1000;
	ctx->last_quota_update = now;

	// receive from connections
	for (conn = ctx->conn_list; conn; conn = conn->next)
		conn->needs_quota = 1;
	do {
		didsomething = 0;

		long long quota = ctx->quota_down;
		if (quota > ctx->quota_total) quota = ctx->quota_total;
		if (quota <= 0) break;

		set_quota_on_conn_list(ctx->conn_list, quota, 0);

		for (conn = ctx->conn_list; conn; conn = conn->next)
		{
			unsigned char buffer[4096];

		try_receiving_more:;
			int recv_len = conn->quota > 4096 ? 4096 : conn->quota;
			int rc = qfl_os_recv(&conn->connhdl, buffer, recv_len);
			if (rc > 0 ) {
				unsigned char *buffer_p = buffer;
				int buffer_len = rc;

				while (conn->crypto_state >= 0 && buffer_len)
				{
					// no crypto?
					if (conn->crypto_state == 0 && *buffer_p == 0) {
						buffer_len--;
						buffer_p++;
						conn->crypto_state = -2;
						continue;
					}

					// responder, receiving RSA key
					if (conn->crypto_state >= 0 && conn->crypto_state <= 255) {
						conn->key_peer_pub[conn->crypto_state++] = *(buffer_p++);
						buffer_len--;
						continue;
					}

					// responder, receiving S-Box, OFB encrypted
					if (conn->crypto_state >= 512 && conn->crypto_state <= 767) {
						conn->output_sbox[conn->crypto_state - 512] = *(buffer_p++);
						if (conn->crypto_state++ == 767) {
							qfl_ofb_xor(conn->input_sbox, conn->output_sbox, 256);
							qfl_ofb_recover(conn->output_sbox);
							conn->crypto_state = -1;
						}
						buffer_len--;
						continue;
					}

					// initiator, receiving S-Box, RSA encrypted
					if (conn->crypto_state >= 1280 && conn->crypto_state <= 1535) {
						conn->output_sbox[conn->crypto_state - 1280] = *(buffer_p++);
						if (conn->crypto_state++ == 1535) {
							// fprintf(stderr, "[%d -> %.*s] Running RSA decryption..", conn->ctx->listenport, strcspn(conn->peer, " \t\n"), conn->peer);
							qfl_rsa_decrypt(2048, conn->output_sbox, conn->ctx->key_conn_pub, conn->ctx->key_conn_sec);
							qfl_ofb_recover(conn->output_sbox);
							// fprintf(stderr, " DONE.\n");
						}
						buffer_len--;
						continue;
					}
				}

				if (buffer_len && conn->crypto_state < 0) {
					if (conn->crypto_state == -1)
						qfl_ofb_xor(conn->input_sbox, buffer_p, buffer_len);
					qfl_pkgfifo_pushdata(conn->fifo[QFL_FIFO_IN], buffer_p, buffer_len);
				}

				ctx->traffic_down += rc;
				ctx->quota_total -= rc;
				ctx->quota_down -= rc;
				conn->quota -= rc;
				didsomething = 1;
				if (conn->quota)
					goto try_receiving_more;
			}
			if (rc == 0)
				conn->needs_quota = 0;
			if (rc < 0)
				conn->closeit = 1;

			while (conn->fifo[QFL_FIFO_IN]->first && conn->process_packages && !conn->closeit)
				conn->process_packages(conn);
		}
	} while (didsomething);

	// send to connections
	for (conn = ctx->conn_list; conn; conn = conn->next)
		conn->needs_quota = 1;
	do {
		didsomething = 0;

		long long quota = ctx->quota_up;
		if (quota > ctx->quota_total) quota = ctx->quota_total;
		if (quota <= 0) break;

		set_quota_on_conn_list(ctx->conn_list, quota, 1);

		for (conn = ctx->conn_list; conn; conn = conn->next)
		{
		try_sending_more:;
			if (!conn->senddata) {
				conn->senddata_offset = 0;

				// responder, sending S-Box, RSA encrypted
				if (conn->crypto_state == 256)
				{
					conn->senddata = malloc(256);
					conn->senddata_len = 256;
					qfl_ofb_init(conn->input_sbox);
					memcpy(conn->senddata, conn->input_sbox, 255);
					conn->senddata[255] = 0;
					qfl_rsa_encrypt(2048, conn->senddata, conn->key_peer_pub);
					conn->crypto_state += 256;
					goto got_new_senddata;
				}

				// initiator, sending RSA key
				if (conn->crypto_state == 1024)
				{
					if (conn->ctx->skip_crypto) {
						conn->senddata = malloc(1);
						conn->senddata_len = 1;
						conn->senddata[0] = 0;
						conn->crypto_state = -2;
						goto got_new_senddata;
					}

					conn->senddata = malloc(256);
					conn->senddata_len = 256;
					memcpy(conn->senddata, conn->ctx->key_conn_pub, 256);
					conn->crypto_state += 256;
					goto got_new_senddata;
				}

				// initiator, sending S-Box, OFB encrypted
				if (conn->crypto_state == 1536)
				{
					conn->senddata = malloc(256);
					conn->senddata_len = 256;
					qfl_ofb_init(conn->input_sbox);
					memcpy(conn->senddata, conn->input_sbox, 256);
					qfl_ofb_xor(conn->output_sbox, conn->senddata, 256);
					conn->crypto_state = -1;
					goto got_new_senddata;
				}

				if (conn->crypto_state >= 0)
					continue;

				for (i=1; i<QFL_FIFO_COUNT; i++) {
					if (conn->fifo[i]->first) {
						qfl_pkgfifo_shiftdata(conn->fifo[i],
								&conn->senddata,
								&conn->senddata_len);
						if (conn->senddata) {
							if (conn->crypto_state == -1)
								qfl_ofb_xor(conn->output_sbox, conn->senddata, conn->senddata_len);
							goto got_new_senddata;
						}
					}
				}
				continue;
			}

		got_new_senddata:;
			int send_len = conn->senddata_len - conn->senddata_offset;
			if (send_len > conn->quota) send_len = conn->quota;

			int rc = qfl_os_send(&conn->connhdl,
					conn->senddata + conn->senddata_offset,
					send_len);

			if (rc > 0) {
				ctx->traffic_up += rc;
				ctx->quota_total -= rc;
				ctx->quota_up -= rc;
				conn->quota -= rc;
				conn->senddata_offset += rc;
				didsomething = 1;
				if (conn->senddata_offset == conn->senddata_len) {
					free(conn->senddata);
					conn->senddata_len = 0;
					conn->senddata = 0;
					if (conn->quota)
						goto try_sending_more;
				}
			}
			if (rc == 0)
				conn->needs_quota = 0;
			if (rc < 0)
				conn->closeit = 1;
		}
	} while (didsomething);

	// close connections, call action callback, remember last conn
	conn = ctx->conn_list;
	struct qfl_conn *last_conn = 0;
	while (conn) {
		struct qfl_conn *next_conn = conn->next;
		if (conn->closeit)
			qfl_conn_destroy(conn);
		else {
			if (conn->process_action)
				conn->process_action(conn);
			last_conn = conn;
		}
		conn = next_conn;
	}

	// move first conn to the end of the list
	conn = ctx->conn_list;
	if (conn && conn != last_conn) {
		ctx->conn_list = conn->next;
		ctx->conn_list->last = 0;
		last_conn->next = conn;
		conn->last = last_conn;
		conn->next = 0;
	}
}