/*
   Newest version in http://imapwiki.org/ImapTest

   Place this file to compiled Dovecot sources' root directory and run:

   gcc imaptest10.c -o imaptest -g -Wall -W -I. -Isrc/lib -Isrc/lib-imap -Isrc/lib-storage/index/mbox -DHAVE_CONFIG_H src/lib-storage/index/mbox/mbox-from.o src/lib-imap/imap-date.o src/lib/liblib.a

   This version builds against Dovecot v1.0
*/
#include "lib.h"
#include "lib-signals.h"
#include "array.h"
#include "base64.h"
#include "buffer.h"
#include "hash.h"
#include "str.h"
#include "network.h"
#include "istream.h"
#include "ostream.h"

#include "imap-date.h"
#include "mbox-from.h"

#include <stdio.h>
#include <stdlib.h>
#include <ctype.h>
#include <fcntl.h>
#include <unistd.h>

/* host / port where to connect to */
#define HOST "127.0.0.1"
#define PORT 143
/* Username. You can give either a single user, or a number of users in which
   case the username gets randomized at each connection. */
//#define USERNAME_TEMPLATE "u%04d@d%04d.domain.org"
//#define USERNAME_TEMPLATE "cras%ld"
#define USERNAME_TEMPLATE "tss"
#define USER_RAND 99
#define DOMAIN_RAND 99
/* Password (for all users) */
#define PASSWORD "pass"
/* Number of simultaneous client connections */
#define CLIENTS_COUNT 10
/* Number of clients to create at startup. After each successful login a new
   client is created. */
#define INIT_CLIENT_COUNT 10
/* Try to keep around this many messages in mailbox (in expunge + append) */
#define MESSAGE_COUNT_THRESHOLD 30
/* Append messages from this mbox file to mailboxes */
#define MBOX_PATH "/home/cras/mail/dovecot-crlf"

/* Add random keywords with max. length n */
//#define RAND_KEYWORDS 40

#define DELAY 1000

struct {
	const char *host, *username_template, *password, *mbox_path;
	const char *mailbox, *copy_dest;
	unsigned int port;

	unsigned int clients_count;
	unsigned int message_count_threshold;

	bool random_states, no_pipelining;

	struct ip_addr ip;
} conf;

enum imap_capability {
	CAP_LITERALPLUS		= 0x01,
	CAP_MULTIAPPEND		= 0x02
};
#define DEFAULT_CAPABILITIES (CAP_LITERALPLUS | CAP_MULTIAPPEND)

struct imap_capability_name {
	const char *name;
	enum imap_capability capability;
};

static const struct imap_capability_name cap_names[] = {
	{ "LITERAL+", CAP_LITERALPLUS },
	{ "MULTIAPPEND", CAP_MULTIAPPEND },

	{ NULL, 0 }
};

enum client_state {
	STATE_BANNER,
	STATE_AUTHENTICATE,
	STATE_LOGIN,
	STATE_LIST,
        STATE_STATUS,
	STATE_SELECT,
	STATE_FETCH,
	STATE_FETCH2,
	STATE_SEARCH,
	STATE_COPY,
	STATE_STORE,
	STATE_STORE_DEL,
	STATE_EXPUNGE,
	STATE_APPEND,
        STATE_NOOP,
        STATE_LOGOUT,
        STATE_DISCONNECT,
        STATE_DELAY,

        STATE_COUNT
};

enum login_state {
	LSTATE_NONAUTH,
	LSTATE_AUTH,
	LSTATE_SELECTED
};

enum state_flags {
	FLAG_MSGSET			= 0x01,
	FLAG_EXPUNGES			= 0x02,
	FLAG_STATECHANGE		= 0x04,
	FLAG_STATECHANGE_NONAUTH	= 0x08,
	FLAG_STATECHANGE_AUTH		= 0x10,
	FLAG_STATECHANGE_SELECTED	= 0x20
};

struct state {
	const char *name;
	const char *short_name;
	enum login_state login_state;
	int probability;
	int probability_again;
	enum state_flags flags;
};

static struct state states[STATE_COUNT] = {
	{ "BANNER",	  "Bann", LSTATE_NONAUTH,  0,   0,  0 },
	{ "AUTHENTICATE", "Auth", LSTATE_NONAUTH,  0,   0,  FLAG_STATECHANGE | FLAG_STATECHANGE_AUTH },
	{ "LOGIN",	  "Logi", LSTATE_NONAUTH,  100, 0,  FLAG_STATECHANGE | FLAG_STATECHANGE_AUTH },
	{ "LIST",	  "List", LSTATE_AUTH,     50,  0,  FLAG_EXPUNGES },
	{ "STATUS",	  "Stat", LSTATE_AUTH,     50,  0,  FLAG_EXPUNGES },
	{ "SELECT",	  "Sele", LSTATE_AUTH,     100, 0,  FLAG_STATECHANGE | FLAG_STATECHANGE_SELECTED },
	{ "FETCH",	  "Fetc", LSTATE_SELECTED, 100, 0,  FLAG_MSGSET },
	{ "FETCH2",	  "Fet2", LSTATE_SELECTED, 100, 30, FLAG_MSGSET },
	{ "SEARCH",	  "Sear", LSTATE_SELECTED, 0,   0,  0 },
	{ "COPY",	  "Copy", LSTATE_SELECTED, 33,  5,  FLAG_MSGSET | FLAG_EXPUNGES },
	{ "STORE",	  "Stor", LSTATE_SELECTED, 50,  0,  FLAG_MSGSET },
	{ "DELETE",	  "Dele", LSTATE_SELECTED, 100, 0,  FLAG_MSGSET },
	{ "EXPUNGE",	  "Expu", LSTATE_SELECTED, 100, 0,  FLAG_EXPUNGES },
	{ "APPEND",	  "Appe", LSTATE_AUTH,     100, 5,  FLAG_EXPUNGES },
	{ "NOOP",	  "Noop", LSTATE_AUTH,     0,   0,  FLAG_EXPUNGES },
	{ "LOGOUT",	  "Logo", LSTATE_NONAUTH,  100, 0,  FLAG_STATECHANGE | FLAG_STATECHANGE_NONAUTH },
	{ "DISCONNECT",	  "Disc", LSTATE_NONAUTH,  0,   0,  0 },
	{ "DELAY",	  "Dela", LSTATE_NONAUTH,  0,   0,  0 }
};

enum command_reply {
	REPLY_BAD,
	REPLY_OK,
	REPLY_NO,
	REPLY_CONT
};

struct client;
struct command;

typedef void command_callback_t(struct client *client, struct command *cmd,
				const char *line, enum command_reply reply);

struct command {
	char *cmdline;
	enum client_state state;
	unsigned int tag;
	command_callback_t *callback;
};

struct client {
	int refcount;

        unsigned int idx;
        unsigned int cur, messages;
	struct istream *input;
	struct ostream *output;
	struct io *io;
	struct timeout *to;
	enum client_state state;
	enum login_state login_state;
	enum imap_capability capabilities;

	/* plan[0] contains always the next state we move to. */
	enum client_state plan[STATE_COUNT];
	unsigned int plan_size;

	struct mbox *mbox;
	uoff_t append_offset, append_size;

	array_t ARRAY_DEFINE(commands, struct command *);
	struct command *last_cmd;
	unsigned int tag_counter;

	uoff_t literal_length;
        int literal_nextline;
        time_t last_io;

	char *username;
	unsigned int delayed:1;
	unsigned int seen_banner:1;
	unsigned int append_unfinished:1;
};

struct mbox {
	int fd;
	char *path;
	struct istream *input;
	uoff_t next_offset;
};

static struct ioloop *ioloop;
static unsigned int counters[STATE_COUNT], total_counters[STATE_COUNT];

static bool stalled = FALSE;
static int clients_count = 0;
static array_t ARRAY_DEFINE(clients, struct client *);
static array_t ARRAY_DEFINE(stalled_clients, unsigned int);

static struct client *client_new(unsigned int idx, struct mbox *mbox);
static bool client_unref(struct client *client);

static void client_input(void *context);
static void state_callback(struct client *client, struct command *cmd,
			   const char *line, enum command_reply reply);
static int client_send_next_cmd(struct client *client);

static struct mbox *mbox_open(const char *path)
{
	struct mbox *mbox;

	mbox = i_new(struct mbox, 1);
	mbox->path = i_strdup(path);
	mbox->fd = open(path, O_RDONLY);
	if (mbox->fd == -1)
		i_fatal("open(%s) failed: %m", path);
	mbox->input =
		i_stream_create_file(mbox->fd, default_pool, (size_t)-1, FALSE);
	return mbox;
}

static void mbox_close(struct mbox *mbox)
{
	i_stream_unref(&mbox->input);
	(void)close(mbox->fd);
	i_free(mbox->path);
	i_free(mbox);
}

static void mbox_get_next_size(struct mbox *mbox, uoff_t *size_r,
			       time_t *time_r)
{
	const char *line;
	char *sender;
	uoff_t offset, last_offset;
	time_t next_time;

	i_stream_seek(mbox->input, mbox->next_offset);

	line = i_stream_read_next_line(mbox->input);
	if (line == NULL) {
		if (mbox->input->v_offset == 0)
			i_fatal("Empty mbox file: %s", mbox->path);

		mbox->next_offset = 0;
		mbox_get_next_size(mbox, size_r, time_r);
		return;
	}

	/* should be From-line */
	if (strncmp(line, "From ", 5) != 0 ||
	    mbox_from_parse((const void *)line+5, strlen(line+5),
			    time_r, &sender) < 0) {
		if (mbox->input->v_offset == 0)
			i_fatal("Not a valid mbox file: %s", mbox->path);
		i_panic("From-line not found at %"PRIuUOFF_T,
			mbox->input->v_offset);
	}
	i_free(sender);

        offset = last_offset = mbox->input->v_offset;
        while ((line = i_stream_read_next_line(mbox->input)) != NULL) {
		if (strncmp(line, "From ", 5) == 0 &&
		    mbox_from_parse((const void *)line+5, strlen(line+5),
				    &next_time, &sender) == 0) {
			i_free(sender);
			if (offset != last_offset)
                                break;

                        /* empty body */
                        offset = last_offset;
                }
                last_offset = mbox->input->v_offset;
        }
        if (offset == last_offset)
                i_fatal("mbox file ends with From-line: %s", mbox->path);

        *size_r = last_offset - offset;
	i_stream_seek(mbox->input, offset);

	mbox->next_offset = last_offset;
}

static int client_parse_literal(struct client *client)
{
	const unsigned char *data;
	size_t size;

	data = i_stream_get_data(client->input, &size);
	if (size < client->literal_length) {
		i_stream_skip(client->input, size);
		client->literal_length -= size;
		return FALSE;
	}

	i_stream_skip(client->input, client->literal_length);
	client->literal_length = 0;
	client->literal_nextline = TRUE;
	return TRUE;
}

static int client_skip_line(struct client *client, const char *line)
{
	size_t len = strlen(line);

	if (line[len-1] != '}')
		return TRUE;
	len--;

	/* skip literal */
	while (len > 0 && i_isdigit(line[len-1]))
		len--;
	if (len == 0 || line[len-1] != '{') {
		i_error("Invalid literal in: %s", line);
		client_unref(client);
		return FALSE;
	}
	client->literal_length = strtoul(t_strcut(line+len, '}'), NULL, 10);
	return client_parse_literal(client);
}

static void client_handle_untagged(struct client *client, const char *line)
{
	if (i_isdigit(line[2])) {
		const char *p = strchr(line+2, ' ');
		unsigned int num;

		if (p++ == NULL) {
			i_error("%s: Unexpected: %s", client->username, line);
			return;
		}

		num = strtoul(t_strcut(line+2, ' '), NULL, 10);
		if (strcasecmp(p, "EXISTS") == 0)
			client->messages = num;

                if (num > client->messages &&
		    client->last_cmd->state > STATE_SELECT) {
			i_fatal("%s: seq too high (%u > %u, state=%d): %s",
                                client->username, num, client->messages,
                                client->last_cmd->state, line);
		}

		if (strcasecmp(p, "EXPUNGE") == 0)
			client->messages--;
	} else if (strncasecmp(line+2, "BYE ", 4) == 0) {
		if (client->last_cmd == NULL ||
		    client->last_cmd->state != STATE_LOGOUT)
			i_info("%s: %s", client->username, line + 6);
		else
			counters[client->last_cmd->state]++;
		client->login_state = LSTATE_NONAUTH;
	} else if (strncasecmp(line+2, "BAD ", 4) == 0) {
		i_info("%s: %s", client->username, line + 6);
	}
}

static const char *get_random_flags(void)
{
	static const char *flags[] = {
		"\\Seen", "\\Flagged", "\\Draft", "\\Answered",
		"$Label1", "$Label2", "$Label3", "$Label4", "$Label5"
	};
#define FLAGS_COUNT (sizeof(flags) / sizeof(flags[0]))
	unsigned int i;
	string_t *str;

	str = t_str_new(128);
	for (i = 0; i < FLAGS_COUNT; i++) {
		if ((rand() % 4) == 0) {
			if (str_len(str) != 0)
				str_append_c(str, ' ');
			str_append(str, flags[i]);
		}
	}
#ifdef RAND_KEYWORDS
	if ((rand() % 10) == 0) {
		unsigned int j, len = (rand() % RAND_KEYWORDS) + 1;

		if (str_len(str) != 0)
			str_append_c(str, ' ');
		for (j = 0; j < len; j++)
			str_append_c(str, (rand() % 26) + 'A');
	}
#endif
	return str_c(str);
}

static void command_send(struct client *client, const char *cmdline,
			 command_callback_t *callback)
{
	struct command *cmd;
	unsigned int tag = client->tag_counter++;

	i_assert(!client->append_unfinished);

	cmd = i_new(struct command, 1);
	cmd->cmdline = i_strdup(cmdline);
	cmd->state = client->state;
	cmd->tag = tag;
	cmd->callback = callback;

	o_stream_send_str(client->output,
			  t_strdup_printf("%u %s\r\n", tag, cmdline));
	array_append(&client->commands, &cmd, 1);
	client->last_cmd = cmd;
}

static int client_append(struct client *client, bool continued)
{
	string_t *cmd;
	struct istream *input;
	time_t t;
	off_t ret;

	if (!continued) {
		mbox_get_next_size(client->mbox, &client->append_size, &t);
		client->append_offset = client->mbox->input->v_offset;

		cmd = t_str_new(128);
		if (!client->append_unfinished)
			str_printfa(cmd, "APPEND \"%s\"", conf.mailbox);
		if ((rand() % 2) == 0)
			str_printfa(cmd, " (%s)", get_random_flags());
		if ((rand() % 2) == 0)
			str_printfa(cmd, " \"%s\"", imap_to_datetime(t));
		str_printfa(cmd, " {%"PRIuUOFF_T, client->append_size);
		if ((client->capabilities & CAP_LITERALPLUS) != 0)
			str_append_c(cmd, '+');
		str_append_c(cmd, '}');

		if (client->append_unfinished) {
			/* continues the last APPEND call */
			str_append(cmd, "\r\n");
			o_stream_send_str(client->output, str_c(cmd));
		} else {
			client->state = STATE_APPEND;
			command_send(client, str_c(cmd), state_callback);
			client->append_unfinished = TRUE;
		}

		if ((client->capabilities & CAP_LITERALPLUS) == 0) {
			/* we'll have to wait for "+" */
			return 0;
		}
	} else {
		i_stream_seek(client->mbox->input, client->append_offset);
	}

	input = i_stream_create_limit(default_pool, client->mbox->input,
				      client->mbox->input->v_offset,
				      client->append_size);
	ret = o_stream_send_istream(client->output, input);
        i_stream_unref(&input);

	if (ret < 0) {
		i_error("APPEND failed: %m");
		return -1;
	}
	client->append_size -= ret;
	client->append_offset += ret;

	if (client->append_size != 0) {
		/* unfinished */
		o_stream_set_flush_pending(client->output, TRUE);
		return 0;
	}

	if ((client->capabilities & CAP_MULTIAPPEND) != 0 &&
	    client->plan_size > 0 && client->plan[0] == STATE_APPEND) {
		/* we want to append another message.
		   do it in the same transaction. */
		return client_send_next_cmd(client);
	}

	client->append_unfinished = FALSE;
	o_stream_send_str(client->output, "\r\n");
	return 0;
}

static void command_unlink(struct client *client, struct command *cmd)
{
	struct command *const *cmds;
	unsigned int i, count;

	cmds = array_get(&client->commands, &count);
	for (i = 0; i < count; i++) {
		if (cmds[i] == cmd) {
			array_delete(&client->commands, i, 1);
			break;
		}
	}
	i_assert(i < count);

	if (client->last_cmd == cmd)
		client->last_cmd = NULL;
}

static void command_free(struct command *cmd)
{
	i_free(cmd->cmdline);
	i_free(cmd);
}

static void client_delay_timeout(void *context)
{
	struct client *client = context;

	i_assert(client->io == NULL);

	client->delayed = FALSE;
	client->last_io = ioloop_time;

	timeout_remove(&client->to);
	client->io = io_add(i_stream_get_fd(client->input),
			    IO_READ, client_input, client);
}

static void client_delay(struct client *client, unsigned int msecs)
{
	i_assert(client->to == NULL);

	client->delayed = TRUE;
	io_remove(&client->io);
	client->to = timeout_add(msecs, client_delay_timeout, client);
}

static bool do_rand(enum client_state state)
{
	return (rand() % 100) < states[state].probability;
}

static bool do_rand_again(enum client_state state)
{
	return (rand() % 100) < states[state].probability_again;
}

static void auth_plain_callback(struct client *client, struct command *cmd,
				const char *line, enum command_reply reply)
{
	buffer_t *str, *buf;

	if (reply == REPLY_OK) {
		state_callback(client, cmd, line, reply);
		return;
	}
	if (reply != REPLY_CONT) {
		i_error("AUTHENTICATE failed: %s", line);
		client_unref(client);
		return;
	}

	counters[cmd->state]++;

	buf = t_str_new(512);
	str_append_c(buf, '\0');
	str_append(buf, client->username);
	str_append_c(buf, '\0');
	str_append(buf, conf.password);

	str = t_str_new(512);
	base64_encode(buf->data, buf->used, str);
	str_append(str, "\r\n");

	o_stream_send_str(client->output, str_c(str));
}

static int client_handle_cmd_reply(struct client *client, struct command *cmd,
				   const char *line, enum command_reply reply)
{
	const char *str;
	unsigned int i;

	switch (reply) {
	case REPLY_OK:
		if (cmd->state != STATE_DISCONNECT)
			counters[cmd->state]++;
		break;
	case REPLY_NO:
		if (cmd->state == STATE_COPY)
			break;

		i_error("%s failed: %s", states[cmd->state].name, line);
		break;

	case REPLY_BAD:
		i_error("%s replied BAD: %s", states[cmd->state].name, line);
		return -1;
	case REPLY_CONT:
		if (cmd->state == STATE_APPEND)
			break;

		i_error("%s: Unexpected continuation: %s",
			states[cmd->state].name, line);
		return -1;
	}

	switch (cmd->state) {
	case STATE_LOGIN:
		client->login_state = LSTATE_AUTH;
		if (reply != REPLY_OK) {
			/* authentication failed */
			return -1;
		}

		for (i = 0; i < 3 && !stalled; i++) {
			if (array_count(&clients) >= conf.clients_count)
				break;

			client_new(array_count(&clients), client->mbox);
		}
		break;
	case STATE_SELECT:
		client->login_state = LSTATE_SELECTED;
		break;
	case STATE_COPY:
		if (reply == REPLY_NO) {
			if (strncasecmp(line, "[TRYCREATE]", 11) == 0) {
				str = t_strdup_printf("CREATE %s",
						      conf.copy_dest);
				client->state = STATE_COPY;
				command_send(client, str, state_callback);
				break;
			}
			i_error("COPY failed: %s", line);
		}
		break;
	case STATE_APPEND:
		if (reply == REPLY_CONT) {
			/* finish appending */
			if (client_append(client, TRUE) < 0)
				return -1;
			break;
		}
		break;
	case STATE_LOGOUT:
		if (client->login_state != LSTATE_NONAUTH) {
			/* untagged bye set state to DISCONNECT, so we
			   shouldn't get here. */
			i_error("Server didn't send BYE but: %s", line);
		}
		return -1;
	case STATE_DISCONNECT:
		return -1;
	default:
		break;
	}

	return 0;
}

static enum client_state client_eat_first_plan(struct client *client)
{
	enum client_state state;

	i_assert(client->plan_size > 0);
	state = client->plan[0];

	client->plan_size--;
	memmove(client->plan, client->plan + 1,
		sizeof(client->plan[0]) * client->plan_size);
	return state;
}

static int client_send_next_cmd(struct client *client)
{
	enum client_state state;
	string_t *cmd;
	const char *str;
	unsigned int i, r, seq1, seq2, deleted;

	state = client_eat_first_plan(client);

	if (client->messages == 0 &&
	    states[state].login_state == LSTATE_SELECTED) {
		/* no messages, no point in doing this command */
		return 0;
	}

	client->state = state;
	switch (state) {
	case STATE_AUTHENTICATE:
		command_send(client, "AUTHENTICATE plain", auth_plain_callback);
		break;
	case STATE_LOGIN:
		str = t_strdup_printf("LOGIN \"%s\" \"%s\"",
				      client->username, conf.password);
		command_send(client, str, state_callback);
		break;
	case STATE_LIST:
		str = t_strdup_printf("LIST \"\" *");
		command_send(client, str, state_callback);
		break;
	case STATE_SELECT:
		str = t_strdup_printf("SELECT \"%s\"", conf.mailbox);
		command_send(client, str, state_callback);
		break;
	case STATE_FETCH:
		command_send(client, "FETCH 1:* (UID FLAGS ENVELOPE BODY)",
			     state_callback);
		break;
	case STATE_FETCH2:
		str = t_strdup_printf("FETCH %lu (BODY[])",
				      (random() % client->messages) + 1);
		command_send(client, str, state_callback);
		break;
	case STATE_SEARCH:
		command_send(client, "SEARCH BODY dovecot", state_callback);
		break;
	case STATE_COPY:
		i_assert(conf.copy_dest != NULL);

		seq1 = (rand() % client->messages) + 1;
		seq2 = (rand() % (client->messages - seq1 + 1));
		seq2 = seq1 + I_MIN(seq2, 5);
		str = t_strdup_printf("COPY %u:%u %s",
				      seq1, seq2, conf.copy_dest);
		command_send(client, str, state_callback);
		break;
	case STATE_STORE:
		cmd = t_str_new(512);
		for (i = 1; i <= client->messages; i++) {
			if ((random() % 5) == 0)
				str_printfa(cmd, "%u,", i);
		}
		if (str_len(cmd) == 0)
			break;

		str_insert(cmd, 0, "STORE ");
		str_truncate(cmd, str_len(cmd) - 1);
		str_append_c(cmd, ' ');
		switch (rand() % 3) {
		case 0:
			str_append_c(cmd, '+');
			break;
		case 1:
			str_append_c(cmd, '-');
			break;
		default:
			break;
		}
		str_printfa(cmd, "FLAGS.SILENT (%s)", get_random_flags());

		command_send(client, str_c(cmd), state_callback);
		break;
	case STATE_STORE_DEL:
		cmd = t_str_new(512);
		str_append(cmd, "STORE ");

		if (client->messages > conf.message_count_threshold) {
			r = 100 - conf.message_count_threshold * 100 /
				client->messages;
		} else {
			r = 0;
		}

		deleted = 0;
		for (i = 1; i <= client->messages; i++) {
			if ((random() % 100) <= (int)r) {
				str_printfa(cmd, "%u,", i);
				deleted++;
			}
		}
		if (deleted == 0)
			break;

		str_truncate(cmd, str_len(cmd) - 1);
		str_append(cmd, " +FLAGS.SILENT \\Deleted");

		command_send(client, str_c(cmd), state_callback);
		break;
	case STATE_EXPUNGE:
		command_send(client, "EXPUNGE", state_callback);
		break;
	case STATE_APPEND:
		if (client->messages >= conf.message_count_threshold)
			break;

		if (client_append(client, FALSE) < 0)
			return -1;
		break;
	case STATE_STATUS:
		str = t_strdup_printf("STATUS \"%s\" (MESSAGES UNSEEN RECENT)",
				      conf.mailbox);
		command_send(client, str, state_callback);
		break;
	case STATE_NOOP:
		command_send(client, "NOOP", state_callback);
		break;
	case STATE_LOGOUT:
		command_send(client, "LOGOUT", state_callback);
		break;

	case STATE_BANNER:
	case STATE_DISCONNECT:
	case STATE_DELAY:
	case STATE_COUNT:
		i_unreached();
	}
	return 0;
}

static enum client_state client_get_next_state(enum client_state state)
{
	i_assert(state < STATE_LOGOUT);

	for (;;) {
		if (!conf.random_states)
			state++;
		else {
			/* if we're not in selected state, we'll randomly do
			   LIST, SELECT, APPEND or LOGOUT */
			state = STATE_LIST +
				(rand() % (STATE_LOGOUT - STATE_LIST + 1));
		}

		if (do_rand(state))
			break;

		if (state == STATE_LOGOUT) {
			/* logout skipped, wrap */
			state = STATE_AUTHENTICATE + 1;
		}
	}
	return state;
}

static struct command *command_lookup(struct client *client, unsigned int tag)
{
	struct command *const *cmds;
	unsigned int i, count;

	cmds = array_get(&client->commands, &count);
	for (i = 0; i < count; i++) {
		if (cmds[i]->tag == tag)
			return cmds[i];
	}
	return NULL;
}

static enum login_state flags2login_state(enum state_flags flags)
{
	if ((flags & FLAG_STATECHANGE_NONAUTH) != 0)
		return LSTATE_NONAUTH;
	else if ((flags & FLAG_STATECHANGE_AUTH) != 0)
		return LSTATE_AUTH;
	else if ((flags & FLAG_STATECHANGE_SELECTED) != 0)
		return LSTATE_SELECTED;

	i_unreached();
}

static enum state_flags
client_get_pending_cmd_flags(struct client *client,
			     enum login_state *new_lstate_r)
{
	enum state_flags state_flags = 0;
	struct command *const *cmds;
	unsigned int i, count;

	*new_lstate_r = client->login_state;
	cmds = array_get(&client->commands, &count);
	for (i = 0; i < count; i++) {
		enum state_flags flags = states[cmds[i]->state].flags;

		if ((flags & FLAG_STATECHANGE) != 0)
			*new_lstate_r = flags2login_state(flags);
		state_flags |= flags;
	}
	return state_flags;
}

static enum client_state client_update_plan(struct client *client)
{
	enum client_state state;
	enum login_state lstate;

	state = client->plan_size > 0 ?
		client->plan[client->plan_size - 1] : client->plan[0];
	if (client->plan_size > 0 &&
	    (states[state].flags & FLAG_STATECHANGE) != 0) {
		/* wait until the state change is done before making new
		   commands. */
		return client->plan[0];
	}
	if ((client_get_pending_cmd_flags(client, &lstate) &
	     FLAG_STATECHANGE) != 0)
		return state;

	if (state == STATE_LOGOUT)
		return state;

	while (client->plan_size <
	       sizeof(client->plan)/sizeof(client->plan[0])) {
		switch (client->login_state) {
		case LSTATE_NONAUTH:
			/* we begin with LOGIN/AUTHENTICATE commands */
			i_assert(client->plan_size == 0);
			state = do_rand(STATE_AUTHENTICATE) ?
				STATE_AUTHENTICATE : STATE_LOGIN;
			break;
		case LSTATE_AUTH:
		case LSTATE_SELECTED:
			if (!do_rand_again(state))
				state = client_get_next_state(state);
			break;
		}
		i_assert(state <= STATE_LOGOUT);

		if (states[state].login_state > client->login_state) {
			/* can't do this now */
			continue;
		}

		client->plan[client->plan_size++] = state;
		if ((states[state].flags & FLAG_STATECHANGE) != 0)
			break;
	}

	i_assert(client->plan_size > 0);
	state = client->plan[0];
	i_assert(states[state].login_state <= client->login_state);
	return state;
}

static bool client_pending_cmds_allow_statechange(struct client *client,
						  enum client_state state)
{
	enum login_state old_lstate, new_lstate = 0;
	struct command *const *cmds;
	unsigned int i, count;

	new_lstate = flags2login_state(states[state].flags);

	cmds = array_get(&client->commands, &count);
	for (i = 0; i < count; i++) {
		if ((states[cmds[i]->state].flags & FLAG_STATECHANGE) != 0)
			return FALSE;

		old_lstate = states[cmds[i]->state].login_state;
		if (new_lstate < old_lstate)
			return FALSE;
		if (new_lstate == old_lstate && new_lstate == LSTATE_SELECTED)
			return FALSE;
	}
	return TRUE;
}

static void state_callback(struct client *client, struct command *cmd,
			   const char *line, enum command_reply reply)
{
	enum state_flags pending_flags;
	enum login_state new_lstate;
	enum client_state state;

	if (client_handle_cmd_reply(client, cmd, line, reply) < 0) {
		client_unref(client);
		return;
	}

	for (;;) {
		state = client_update_plan(client);
		i_assert(state <= STATE_LOGOUT);

		if (client->append_unfinished)
			break;
		if (conf.no_pipelining && array_count(&client->commands) > 0)
			break;

		if ((states[state].flags & FLAG_STATECHANGE) != 0) {
			/* this command would change the state. check if there
			   are any pending commands that don't like the
			   change */
			if (!client_pending_cmds_allow_statechange(client,
								   state))
				break;
		}
		pending_flags = client_get_pending_cmd_flags(client,
							     &new_lstate);
		if ((states[state].flags & FLAG_STATECHANGE) == 0 &&
		    (pending_flags & FLAG_STATECHANGE) != 0) {
			/* we're changing state. allow this command if its
			   required login_state is lower than the current
			   state or the state we're changing to. */
			if (new_lstate <= states[state].login_state ||
			    client->login_state < states[state].login_state)
				break;
		}
		if ((states[state].flags & FLAG_MSGSET) != 0 &&
		    (pending_flags & (FLAG_EXPUNGES | FLAG_STATECHANGE)) != 0) {
			/* msgset may become invalid if we send it now */
			break;
		}

		if (client_send_next_cmd(client) < 0) {
			client_unref(client);
			return;
		}
	}

	if (!client->delayed && do_rand(STATE_DELAY)) {
		counters[STATE_DELAY]++;
		client_delay(client, DELAY);
	}
}

static void client_capability_parse(struct client *client, const char *line)
{
	const char *const *tmp;
	unsigned int i;

	client->capabilities = 0;

	for (tmp = t_strsplit(line, " "); *tmp != NULL; tmp++) {
		for (i = 0; cap_names[i].name != NULL; i++) {
			if (strcasecmp(*tmp, cap_names[i].name) == 0) {
				client->capabilities |= cap_names[i].capability;
				break;
			}
		}
	}
}

static void client_input(void *context)
{
	struct client *client = context;
	const char *line, *p;
	struct command *cmd;
	enum command_reply reply;

        client->last_io = ioloop_time;

	switch (i_stream_read(client->input)) {
	case 0:
		return;
	case -1:
		/* disconnected */
		client_unref(client);
		return;
	case -2:
		/* buffer full */
		i_error("line too long");
		client_unref(client);
		return;
	}

	if (client->literal_length > 0) {
		if (!client_parse_literal(client))
			return;
	}

	if (!client->seen_banner) {
		/* we haven't received the banner yet */
		line = i_stream_next_line(client->input);
		if (line == NULL)
			return;
		client->seen_banner = TRUE;

		p = strstr(line, "[CAPABILITY ");
		if (p != NULL)
			client_capability_parse(client, t_strcut(p + 12, ']'));

		(void)client_update_plan(client);
		client_send_next_cmd(client);
	}

	while ((line = i_stream_next_line(client->input)) != NULL) {
		if (client->literal_nextline) {
			client->literal_nextline = FALSE;
			if (!client_skip_line(client, line))
				return;
			continue;
		}

		if (strncmp(line, "+ ", 2) == 0) {
			if (client->last_cmd == NULL) {
				i_error("Unexpected command contination: %s",
					line);
				client_unref(client);
				return;
			}
			client->last_cmd->
				callback(client, client->last_cmd,
					 line + 2, REPLY_CONT);
			continue;
		}

		if (strncmp(line, "* ", 2) == 0) {
			client_handle_untagged(client, line);
			if (!client_skip_line(client, line))
				return;
			continue;
		}

		cmd = command_lookup(client, atoi(t_strcut(line, ' ')));
		if (cmd == NULL) {
			i_error("Unexpected tagged reply: %s", line);
			client_unref(client);
			return;
		}
		p = strchr(line, ' ');
		if (p != NULL) {
			p++;
			if (strncmp(p, "OK ", 3) == 0) {
				reply = REPLY_OK;
				p += 3;
			} else if (strncmp(p, "NO ", 3) == 0) {
				reply = REPLY_NO;
				p += 3;
			} else if (strncmp(p, "BAD ", 4) == 0) {
				reply = REPLY_BAD;
				p += 4;
				i_error("BAD reply '%s' for command: %s",
					p, cmd->cmdline);
			} else {
				p = NULL;
			}
		}
		if (p == NULL) {
			i_error("Invalid input: %s", line);
			client_unref(client);
			return;
		}

		command_unlink(client, cmd);

		client->refcount++;
		cmd->callback(client, cmd, p, reply);
		command_free(cmd);

		if (!client_unref(client))
			return;
	}
	if (do_rand(STATE_DISCONNECT)) {
		/* random disconnection */
		counters[STATE_DISCONNECT]++;
		client_unref(client);
		return;
	}
}

static void client_wait_connect(void *context)
{
	struct client *client = context;
	int err, fd;

	fd = i_stream_get_fd(client->input);
	err = net_geterror(fd);
	if (err != 0) {
		i_error("connect() failed: %s", strerror(err));
		client_unref(client);
		return;
	}

	io_remove(&client->io);
	client->io = io_add(fd, IO_READ, client_input, client);
}

static int client_output(void *context)
{
        struct client *client = context;
	int ret;

	ret = o_stream_flush(client->output);
	client->last_io = ioloop_time;

	if (client->append_size > 0) {
		if (client_append(client, TRUE) < 0)
			client_unref(client);
	}

        return ret;
}

static struct client *client_new(unsigned int idx, struct mbox *mbox)
{
	struct client *client;
	int fd;

	if (stalled) {
		array_append(&stalled_clients, &idx, 1);
		return NULL;
	}

	fd = net_connect_ip(&conf.ip, conf.port, NULL);
	if (fd < 0) {
		i_error("connect() failed: %m");
		return NULL;
	}

	client = i_new(struct client, 1);
	client->refcount = 1;
	client->tag_counter = 1;
	client->capabilities = DEFAULT_CAPABILITIES;
	client->idx = idx;
        client->mbox = mbox;
	client->input = i_stream_create_file(fd, default_pool, 65536, TRUE);
	client->output = o_stream_create_file(fd, default_pool, (size_t)-1, FALSE);
	client->io = io_add(fd, IO_READ, client_wait_connect, client);
	client->username = i_strdup_printf(conf.username_template,
					   (random() % USER_RAND) + 1,
					   (random() % DOMAIN_RAND) + 1);
        client->last_io = ioloop_time;
	ARRAY_CREATE(&client->commands, default_pool, struct command *, 16);
	o_stream_set_flush_callback(client->output, client_output, client);
	clients_count++;

        array_idx_set(&clients, idx, &client);
        return client;
}

static bool client_unref(struct client *client)
{
	struct mbox *mbox = client->mbox;
	unsigned int idx = client->idx;
	struct command *const *cmds;
	struct client *null;
	unsigned int i, count;

	i_assert(client->refcount > 0);
	if (--client->refcount > 0)
		return TRUE;

	if (--clients_count == 0)
		stalled = FALSE;
	null = NULL;
	array_idx_set(&clients, idx, &null);

	cmds = array_get(&client->commands, &count);
	for (i = 0; i < count; i++)
		command_free(cmds[i]);

	/*if (--clients_count == 0)
		io_loop_stop(ioloop);*/
	if (client->io != NULL)
		io_remove(&client->io);
	if (client->to != NULL)
		timeout_remove(&client->to);
	o_stream_unref(&client->output);
	i_stream_unref(&client->input);
	i_free(client->username);
	i_free(client);

	if (io_loop_is_running(ioloop)) {
		client_new(idx, mbox);
		if (!stalled) {
			const unsigned int *indexes;
			unsigned int i, count;

			indexes = array_get(&stalled_clients, &count);
			for (i = 0; i < count && i < 3; i++)
				client_new(indexes[i], mbox);
			array_delete(&stalled_clients, 0, i);
		}
	}
	return FALSE;
}

#define STATE_IS_VISIBLE(state) \
	(states[i].probability != 0)

static void print_header(void)
{
	unsigned int i;
	bool have_agains = FALSE;

	for (i = 1; i < STATE_COUNT; i++) {
		if (!STATE_IS_VISIBLE(i))
			continue;
		printf("%s ", states[i].short_name);
	}
	printf("\n");
	for (i = 1; i < STATE_COUNT; i++) {
		if (!STATE_IS_VISIBLE(i))
			continue;
		if (states[i].probability_again)
			have_agains = TRUE;
		printf("%3d%% ", states[i].probability);
	}
	printf("\n");

	if (have_agains) {
		for (i = 1; i < STATE_COUNT; i++) {
			if (!STATE_IS_VISIBLE(i))
				continue;
			if (states[i].probability_again == 0)
				printf("     ");
			else
				printf("%3d%% ", states[i].probability_again);
		}
		printf("\n");
	}
}

static void print_timeout(void *context __attr_unused__)
{
	struct client *const *c;
        static int rowcount = 0;
        unsigned int i, count, banner_waits, stall_count;

        if ((rowcount++ % 10) == 0)
                print_header();

        for (i = 1; i < STATE_COUNT; i++) {
		if (!STATE_IS_VISIBLE(i))
			continue;
		printf("%4d ", counters[i]);
		total_counters[i] += counters[i];
		counters[i] = 0;
        }

	stalled = FALSE;
	banner_waits = 0;
	stall_count = 0;

	c = array_get(&clients, &count);
	for (i = 0; i < count; i++) {
		if (c[i] != NULL && c[i]->state == STATE_BANNER) {
			banner_waits++;

			if (c[i]->last_io < ioloop_time - 15) {
				stall_count++;
				stalled = TRUE;
			}
                }
        }

	printf("%3d/%3d", (clients_count - banner_waits), clients_count);
	if (stall_count > 0)
		printf(" (%u stalled)", stall_count);

	if (array_count(&clients) < conf.clients_count) {
		printf(" [%d%%]", array_count(&clients) * 100 /
		       conf.clients_count);
	}

	printf("\n");
	for (i = 0; i < count; i++) {
		if (c[i] != NULL && c[i]->state != STATE_BANNER &&
		    c[i]->to == NULL && c[i]->last_io < ioloop_time - 15) {
			stalled = TRUE;
                        fprintf(stderr, " - %d. stalled for %u secs in %s\n", i,
				(unsigned)(ioloop_time - c[i]->last_io),
				states[c[i]->state].name);
                }
        }
}

static void print_total(void)
{
        unsigned int i;

	printf("\nTotals:\n");
	print_header();

        for (i = 1; i < STATE_COUNT; i++) {
		if (!STATE_IS_VISIBLE(i))
			continue;

		total_counters[i] += counters[i];
		printf("%4d ", total_counters[i]);
	}
	printf("\n");
}

static void fix_probabilities(void)
{
	unsigned int i;

	if (conf.copy_dest == NULL)
		states[STATE_COPY].probability = 0;

	if (states[STATE_LOGIN].probability == 100)
		states[STATE_AUTHENTICATE].probability = 0;
	else if (states[STATE_AUTHENTICATE].probability == 100)
		states[STATE_LOGIN].probability = 0;

	for (i = STATE_LIST; i <= STATE_LOGOUT; i++) {
		if (states[i].probability > 0)
			break;
	}
	if (i > STATE_LOGOUT)
		i_fatal("Invalid probabilities");
}

static void sig_die(int signo __attr_unused__, void *context __attr_unused__)
{
	io_loop_stop(ioloop);
}

static void timeout_stop(void *context __attr_unused__)
{
	io_loop_stop(ioloop);
}

static struct state *state_find(const char *name)
{
	unsigned int i;

	for (i = 0; i < STATE_COUNT; i++) {
		if (strcasecmp(states[i].name, name) == 0 ||
		    strcasecmp(states[i].short_name, name) == 0)
			return &states[i];
	}
	return NULL;
}

static void print_help(void)
{
	printf(
"imaptest [user=USER] [host=HOST] [port=PORT] [pass=PASSWORD] [mbox=MBOX] "
"         [clients=CC] [msgs=NMSG] [box=MAILBOX] [copybox=DESTBOX]\n"
"         [-] [<state>[=<n%%>[,<m%%>]]] [random] [no_pipelining] "
"\n"
" USER = template for username. \"u%%04d\" will generate users \"u0001\" to\n"
"        \"u0099\". \"u%%04d@d%%04d\" will generate also \"d0001\" to \"d0099\".\n"
" MBOX = path to mbox from which we read mails to append.\n"
" MAILBOX = Mailbox name where to do all the work (default = INBOX).\n"
" DESTBOX = Mailbox name where to copy messages.\n"
" CC   = number of concurrent clients. [%u]\n"
" NMSG = target number of messages in the mailbox. [%u]\n"
"\n"
" -    = Sets all probabilities to 0%% except for LOGIN, LOGOUT and SELECT\n"
" <state> = Sets state's probability to n%% and repeated probability to m%%\n",
	CLIENTS_COUNT, MESSAGE_COUNT_THRESHOLD);
}

int main(int argc __attr_unused__, char *argv[])
{
	struct mbox *mbox;
	struct timeout *to, *to_stop;
	struct client *const *c;
	struct ip_addr *ips;
	struct state *state;
	const char *key, *value;
	unsigned int i, count;
	int ret, ipcount;

	lib_init();
	ioloop = io_loop_create(system_pool);

	lib_signals_init();
        lib_signals_set_handler(SIGINT, TRUE, sig_die, NULL);

	conf.password = PASSWORD;
	conf.username_template = USERNAME_TEMPLATE;
	conf.host = HOST;
	conf.port = PORT;
	conf.mbox_path = MBOX_PATH;
	conf.mailbox = "INBOX";
	conf.clients_count = CLIENTS_COUNT;
	conf.message_count_threshold = MESSAGE_COUNT_THRESHOLD;
	to_stop = NULL;

	for (argv++; *argv != NULL; argv++) {
		value = strchr(*argv, '=');
		key = value == NULL ? *argv :
			t_strdup_until(*argv, value);
		if (value != NULL) value++;

		if (strcmp(*argv, "-h") == 0 ||
		    strcmp(*argv, "--help") == 0) {
			print_help();
			return 0;
		}
		if (strcmp(key, "secs") == 0) {
			to_stop = timeout_add(atoi(value) * 1000,
					      timeout_stop, NULL);
			continue;
		}
		if (strcmp(key, "seed") == 0) {
			srand(atoi(value));
			continue;
		}

		if (strcmp(*argv, "-") == 0) {
			for (i = STATE_LOGIN+1; i < STATE_LOGOUT; i++) {
				if (i != STATE_SELECT)
					states[i].probability = 0;
			}
			continue;
		}

		state = state_find(key);
		if (state != NULL) {
			/* [<probability>[,<probability_again>]] */
			const char *p;

			if (value == NULL) {
				state->probability = 100;
				continue;
			}
			p = strchr(value, ',');
			if (p != NULL)
				value = t_strdup_until(value, p++);

			state->probability = atoi(value);
			if (p != NULL)
				state->probability_again = atoi(p);
			continue;
		}

		if (strcmp(*argv, "random") == 0) {
			conf.random_states = TRUE;
			continue;
		}
		if (strcmp(*argv, "no_pipelining") == 0) {
			conf.no_pipelining = TRUE;
			continue;
		}

		/* pass=password */
		if (strcmp(key, "pass") == 0) {
			conf.password = value;
			continue;
		}

		/* mbox=path */
		if (strcmp(key, "mbox") == 0) {
			conf.mbox_path = value;
			continue;
		}

		/* clients=# */
		if (strcmp(key, "clients") == 0) {
			conf.clients_count = atoi(value);
			continue;
		}

		/* msgs=# */
		if (strcmp(key, "msgs") == 0) {
			conf.message_count_threshold = atoi(value);
			continue;
		}

		/* box=mailbox */
		if (strcmp(key, "box") == 0) {
			conf.mailbox = value;
			continue;
		}

		/* copybox=mailbox */
		if (strcmp(key, "copybox") == 0) {
			conf.copy_dest = value;
			continue;
		}
		if (strcmp(key, "user") == 0) {
			conf.username_template = value;
			continue;
		}
		if (strcmp(key, "host") == 0) {
			conf.host = value;
			continue;
		}
		if (strcmp(key, "port") == 0) {
			conf.port = atoi(value);
			continue;
		}

		printf("Unknown arg: %s\n", *argv);
		return 1;
	}

	if ((ret = net_gethostbyname(conf.host, &ips, &ipcount)) != 0) {
		i_error("net_gethostbyname(%s) failed: %s",
			conf.host, net_gethosterror(ret));
		return 1;
	}
	conf.ip = ips[0];

	fix_probabilities();

	mbox = mbox_open(conf.mbox_path);

	ARRAY_CREATE(&clients, default_pool, struct client *, CLIENTS_COUNT);
	ARRAY_CREATE(&stalled_clients, default_pool, unsigned int, CLIENTS_COUNT);
	to = timeout_add(1000, print_timeout, NULL);
	for (i = 0; i < INIT_CLIENT_COUNT && i < conf.clients_count; i++)
		client_new(i, mbox);
        io_loop_run(ioloop);

	c = array_get(&clients, &count);
	for (i = 0; i < count; i++) {
		if (c[i] != NULL)
			client_unref(c[i]);
        }

	print_total();
	mbox_close(mbox);
	timeout_remove(&to);
	if (to_stop != NULL)
		timeout_remove(&to_stop);

	lib_signals_deinit();
	io_loop_destroy(&ioloop);
	lib_deinit();
	return 0;
}