/*
 * A very simple controlable traffic generator for TCP testing.
 *
 * Copyright 2007, Haiku, Inc. All Rights Reserved.
 * Distributed under the terms of the MIT License.
 *
 * Authors:
 *      Hugo Santos, hugosantos@gmail.com
 */

#include <OS.h>

#include <assert.h>
#include <ctype.h>
#include <stdio.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

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


struct context {
	int sock;
	uint8 generator;
	int index;
	int8_t buffer[256];
};

static int process_command(context *ctx);

static int
number(context *ctx)
{
	int result = 0;

	while (isdigit(ctx->buffer[ctx->index])) {
		result *= 10;
		result += ctx->buffer[ctx->index] - '0';
		ctx->index++;
	}

	return result;
}


static int
value(context *ctx)
{
	if (ctx->buffer[ctx->index] == '[') {
		ctx->index++;
		int upper, lower = number(ctx);
		if (ctx->buffer[ctx->index] == ',') {
			ctx->index++;
			upper = number(ctx);
		} else {
			upper = lower + 50;
			lower -= 50;
		}

		return lower + rand() % (upper - lower + 1);
	}

	return number(ctx);
}


static int
repeat(context *ctx)
{
	int max, saved, count = number(ctx);

	max = saved = ctx->index;
	for (int i = 0; i < count; i++) {
		ctx->index = saved;
		if (process_command(ctx) < 0)
			return -1;
		if (ctx->index > max)
			max = ctx->index;
	}

	ctx->index = max;
	return 0;
}


static void
send_packet(context *ctx, size_t bytes)
{
	uint8_t buffer[1024];
	uint8_t *ptr = buffer;

	if (bytes > sizeof(buffer))
		ptr = new uint8_t[bytes];

	for (size_t i = 0; i < bytes; i++) {
		ptr[i] = ctx->generator + '0';
		ctx->generator = (ctx->generator + 1) % 10;
	}

	send(ctx->sock, ptr, bytes, 0);

	if (ptr != buffer)
		delete [] ptr;
}


static int
process_command(context *ctx)
{
	while (ctx->buffer[ctx->index] != '.') {
		ctx->index++;

		switch (ctx->buffer[ctx->index - 1]) {
			case 'r':
				if (repeat(ctx) < 0)
					return -1;
				break;

			case 'b':
				send_packet(ctx, 1);
				break;

			case 'p':
				send_packet(ctx, value(ctx));
				break;

			case 's':
				usleep(value(ctx) * 1000);
				break;

			case 'W':
			{
				int value = number(ctx);
				setsockopt(ctx->sock, SOL_SOCKET, SO_SNDBUF, &value,
					sizeof(value));
				break;
			}

			case 'k':
				return -1;
		}
	}

	return 0;
}


static int
read_command(context *ctx)
{
	int index = 0;

	do {
		int size = recv(ctx->sock, ctx->buffer + index, 1, 0);
		if (size == 0)
			return -1;
		else if (size < 0)
			continue;

		index++;
	} while (ctx->buffer[index - 1] != '.');

	ctx->index = 0;
	return process_command(ctx);
}


static int32
handle_client(void *data)
{
	context ctx = { *(int *)data, 0 };

	while (read_command(&ctx) == 0);

	fprintf(stderr, "Client %d leaving.\n", ctx.sock);

	close(ctx.sock);

	return 0;
}


int
main(int argc, char *argv[])
{
	int port = 12345;

	for (int i = 1; i < argc; i++) {
		if (!strcmp(argv[i], "-p")) {
			i++;
			assert(i < argc);
			port = atoi(argv[i]);
		} else if (!strcmp(argv[i], "-h")) {
			fprintf(stderr, "tcptester [-p port]\n");
			return 1;
		}
	}

	int sock = socket(AF_INET, SOCK_STREAM, 0);

	if (sock < 0) {
		perror("socket()");
		return -1;
	}

	sockaddr_in sin;
	memset(&sin, 0, sizeof(sin));
	sin.sin_family = AF_INET;
	sin.sin_port = htons(port);

	if (bind(sock, (sockaddr *)&sin, sizeof(sockaddr_in)) < 0) {
		perror("bind()");
		return -1;
	}

	if (listen(sock, 5) < 0) {
		perror("listen()");
		return -1;
	}

	while (1) {
		sockaddr_in peer;
		socklen_t peerLen = sizeof(peer);

		int newSock = accept(sock, (sockaddr *)&peer, &peerLen);
		if (newSock < 0) {
			perror("accept()");
			return -1;
		}

		char buf[64];
		inet_ntop(AF_INET, &peer.sin_addr, buf, sizeof(buf));

		thread_id newThread = spawn_thread(handle_client, "client",
			B_NORMAL_PRIORITY, &newSock);

		fprintf(stderr, "New client %d from %s with thread id %ld.\n",
			newSock, buf, (int32)newThread);

		resume_thread(newThread);
	}

	return 0;
}