/* Copyright (C) 2010 Timo Sirainen, LGPLv2.1 */

/*
   eval `cat /usr/local/lib/dovecot/dovecot-config`
   gcc -fPIC -shared -DHAVE_CONFIG_H \
     `echo $DOVECOT_CFLAGS $LIBDOVECOT_INCLUDE $LIBDOVECOT_STORAGE_INCLUDE` \
     pop3-throttle-plugin.c -o pop3_throttle_plugin.so
*/

#include "lib.h"
#include "array.h"
#include "hex-dec.h"
#include "read-full.h"
#include "write-full.h"
#include "file-dotlock.h"
#include "mail-storage-private.h"
#include "mail-namespace.h"
#include "mail-search-build.h"

#include <stdlib.h>
#include <fcntl.h>
#include <unistd.h>
#include <time.h>

void pop3_throttle_plugin_init(struct module *module);
void pop3_throttle_plugin_deinit(void);

#define POP3_THROTTLE_ENABLED_PATH "/etc/dovecot/pop3-throttle-enabled"
#define POP3_THROTTLE_STATE_FNAME "pop3-throttle.db"
/* reset state every n seconds */
#define POP3_THROTTLE_STATE_RESET_SECS (60*15)

#define POP3_THROTTLE_USER_CONTEXT(obj) \
	MODULE_CONTEXT(obj, pop3_throttle_user_module)
#define POP3_THROTTLE_CONTEXT(obj) \
	MODULE_CONTEXT(obj, pop3_throttle_storage_module)

struct pop3_throttle_state {
	uint32_t uid_validity;
	uint32_t highest_visible_uid;

	uint32_t start_time;
	uint32_t fetch_msgs;
	uint32_t fetch_kbytes;
};
#define POP3_THROTTLE_FILE_SIZE \
	(sizeof(struct pop3_throttle_state) / \
		sizeof(uint32_t) * (sizeof(uint32_t)*2 + 1))

struct pop3_throttle_mailbox {
	union mailbox_module_context module_ctx;

	uint32_t session_highest_visible_seq;
	struct pop3_throttle_state orig_state, state;
};

struct pop3_throttle_user {
	union mail_user_module_context module_ctx;

	/* allow max. n new messages per every state unit
	   (until state is reset) */
	unsigned int max_msgs_per_unit;
	/* allow max. n kbytes of new messages per every state unit */
	unsigned int max_kbytes_per_unit;
};

const char *pop3_throttle_plugin_version = PACKAGE_VERSION;

static const struct dotlock_settings dotlock_set = {
	.timeout = 2,
	.stale_timeout = 60
};

static MODULE_CONTEXT_DEFINE_INIT(pop3_throttle_user_module,
				  &mail_user_module_register);
static MODULE_CONTEXT_DEFINE_INIT(pop3_throttle_storage_module,
				  &mail_storage_module_register);

static const char *pop3_throttle_get_state_path(struct mailbox *box)
{
	const char *dir;

	dir = mailbox_list_get_path(box->list, NULL,
				    MAILBOX_LIST_PATH_TYPE_CONTROL);
	return t_strconcat(dir, "/"POP3_THROTTLE_STATE_FNAME, NULL);
}

static void pop3_throttle_read_state(const char *path,
				     struct pop3_throttle_state *state_r)
{
#define BUF_IDX_OFFSET(idx) \
	((idx)*(sizeof(uint32_t)*2+1))
#define BUF_GET_HEX(buf, idx) \
	hex2dec(buf + BUF_IDX_OFFSET(idx++), sizeof(uint32_t)*2)
	unsigned char buf[POP3_THROTTLE_FILE_SIZE];
	unsigned int idx;
	int fd, ret;

	memset(state_r, 0, sizeof(*state_r));

	fd = open(path, O_RDONLY);
	if (fd == -1) {
		if (errno != ENOENT)
			i_error("open(%s) failed: %m", path);
		return;
	}
	if ((ret = read_full(fd, buf, sizeof(buf))) <= 0) {
		if (ret < 0)
			i_error("read(%s) failed: %m", path);
		else
			i_error("read(%s) failed: Unexpected EOF", path);
	} else {
		idx = 0;
		state_r->uid_validity = BUF_GET_HEX(buf, idx);
		state_r->highest_visible_uid = BUF_GET_HEX(buf, idx);
		state_r->start_time = BUF_GET_HEX(buf, idx);
		state_r->fetch_msgs = BUF_GET_HEX(buf, idx);
		state_r->fetch_kbytes = BUF_GET_HEX(buf, idx);
		i_assert(BUF_IDX_OFFSET(idx) == sizeof(buf));
	}
	if (close(fd) < 0)
		i_error("close(%s) failed: %m", path);
}

static int pop3_throttle_write_state(const char *path,
				     const struct pop3_throttle_state *state)
{
#define BUF_ADD_HEX(buf, dec) \
	dec2hex(buffer_append_space_unsafe(buf, sizeof(uint32_t)*2), \
		dec, sizeof(uint32_t)*2), buffer_append_c(buf, '\n')
	struct dotlock *dotlock;
	buffer_t *buf;
	int fd;

	if ((fd = file_dotlock_open(&dotlock_set, path, 0, &dotlock)) < 0)
		return -1;

	buf = buffer_create_dynamic(pool_datastack_create(),
				    POP3_THROTTLE_FILE_SIZE);
	BUF_ADD_HEX(buf, state->uid_validity);
	BUF_ADD_HEX(buf, state->highest_visible_uid);
	BUF_ADD_HEX(buf, state->start_time);
	BUF_ADD_HEX(buf, state->fetch_msgs);
	BUF_ADD_HEX(buf, state->fetch_kbytes);
	i_assert(buf->used == POP3_THROTTLE_FILE_SIZE);

	if (write_full(fd, buf->data, buf->used) < 0) {
		i_error("write(%s) failed: %m",
			file_dotlock_get_lock_path(dotlock));
		(void)file_dotlock_delete(&dotlock);
		return -1;
	} else {
		return file_dotlock_replace(&dotlock, 0);
	}
}

static void pop3_throttle_init_state(struct mailbox *box,
				     struct pop3_throttle_state *state)
{
	struct mailbox_status status;
	time_t now, expire_time;

	now = time(NULL);
	expire_time = state->start_time + POP3_THROTTLE_STATE_RESET_SECS;

	mailbox_get_status(box, STATUS_UIDVALIDITY | STATUS_UIDNEXT, &status);
	if (state->uid_validity != status.uidvalidity) {
		if (box->storage->user->mail_debug) {
			i_debug("pop3-throttle: uidvalidity changed %u -> %u",
				state->uid_validity, status.uidvalidity);
		}
		memset(state, 0, sizeof(*state));
		state->uid_validity = status.uidvalidity;
		state->start_time = now;
	} else if (expire_time < now) {
		if (box->storage->user->mail_debug) {
			i_debug("pop3-throttle: throttle expired %u secs ago, "
				"reseting state",
				(unsigned int)(now - expire_time));
		}
		state->start_time = now;
		state->fetch_msgs = 0;
		state->fetch_kbytes = 0;
	} else {
		if (box->storage->user->mail_debug) {
			i_debug("pop3-throttle: throttle expires in %u secs",
				(unsigned int)(expire_time - now));
		}
	}

	if (state->highest_visible_uid == 0) {
		/* we don't know how high uids pop3 clients have seen,
		   so play safe and assume they've seen everything */
		state->highest_visible_uid = status.uidnext - 1;
		if (box->storage->user->mail_debug) {
			i_debug("pop3-throttle: previous state unknown, "
				"assuming client has seen all messages");
		}
	}
}

static void pop3_throttle_get_visible_msgs(struct mailbox *box)
{
	struct pop3_throttle_mailbox *tbox = POP3_THROTTLE_CONTEXT(box);
	struct pop3_throttle_user *tuser =
		POP3_THROTTLE_USER_CONTEXT(box->storage->user);
	struct mail_search_args *search_args;
        struct mailbox_transaction_context *t;
	struct mail_search_context *ctx;
	struct mail *mail;
	struct stat st;
	uoff_t size;
	unsigned int max_msgs_per_unit, max_kbytes_per_unit;
	uint32_t seq1, seq2;

	(void)mailbox_sync(box, MAILBOX_SYNC_FLAG_FULL_READ);

	if (stat(POP3_THROTTLE_ENABLED_PATH, &st) == 0) {
		max_msgs_per_unit = tuser->max_msgs_per_unit;
		max_kbytes_per_unit = tuser->max_kbytes_per_unit;
		if (box->storage->user->mail_debug)
			i_debug("pop3-throttle: throttling active");
	} else {
		if (box->storage->user->mail_debug) {
			i_debug("pop3-throttle: throttling inactive ("
				POP3_THROTTLE_ENABLED_PATH" doesn't exist)");
		}
		max_msgs_per_unit = 0;
		max_kbytes_per_unit = 0;
	}

	search_args = mail_search_build_init();
	search_args->args = p_new(search_args->pool, struct mail_search_arg, 1);
	search_args->args->type = SEARCH_UIDSET;
	p_array_init(&search_args->args->value.seqset, search_args->pool, 1);
	seq_range_array_add_range(&search_args->args->value.seqset,
				  tbox->state.highest_visible_uid + 1,
				  (uint32_t)-1);

	t = mailbox_transaction_begin(box, 0);
	ctx = mailbox_search_init(t, search_args, NULL);
	mail_search_args_unref(&search_args);

	mail = mail_alloc(t, MAIL_FETCH_VIRTUAL_SIZE, NULL);
	while (mailbox_search_next(ctx, mail)) {
		if (mail_get_virtual_size(mail, &size) < 0)
			continue;

		if (tbox->state.fetch_msgs >= max_msgs_per_unit &&
		    max_msgs_per_unit > 0) {
			if (box->storage->user->mail_debug) {
				i_debug("pop3-throttle: max_msgs reached, "
					"uid=%u isn't visible", mail->uid);
			}
			break;
		}
		if (tbox->state.fetch_kbytes >= max_kbytes_per_unit &&
		    max_kbytes_per_unit > 0) {
			if (box->storage->user->mail_debug) {
				i_debug("pop3-throttle: max_kbytes reached, "
					"uid=%u isn't visible", mail->uid);
			}
			break;
		}

		tbox->state.fetch_msgs++;
		tbox->state.fetch_kbytes += (size + 1023) / 1024;
		tbox->state.highest_visible_uid = mail->uid;
	}
	mail_free(&mail);
	(void)mailbox_search_deinit(&ctx);
	(void)mailbox_transaction_commit(&t);

	if (tbox->state.highest_visible_uid == 0)
		seq2 = 0;
	else {
		mailbox_get_seq_range(box, 1, tbox->state.highest_visible_uid,
				      &seq1, &seq2);
	}
	tbox->session_highest_visible_seq = seq2;
}

static void
pop3_throttle_get_status(struct mailbox *box,
			 enum mailbox_status_items items,
			 struct mailbox_status *status_r)
{
	struct pop3_throttle_mailbox *tbox = POP3_THROTTLE_CONTEXT(box);

	tbox->module_ctx.super.get_status(box, items, status_r);

	if (status_r->messages > tbox->session_highest_visible_seq)
		status_r->messages = tbox->session_highest_visible_seq;
}

static void pop3_throttle_mailbox_opened(struct mailbox *box)
{
	struct pop3_throttle_mailbox *tbox = POP3_THROTTLE_CONTEXT(box);
	const char *path;

	if (tbox == NULL)
		return;

	path = pop3_throttle_get_state_path(box);
	pop3_throttle_read_state(path, &tbox->orig_state);

	tbox->state = tbox->orig_state;
	pop3_throttle_init_state(box, &tbox->state);
	pop3_throttle_get_visible_msgs(box);

	if (memcmp(&tbox->state, &tbox->orig_state, sizeof(tbox->state)) != 0)
		pop3_throttle_write_state(path, &tbox->state);
}

static void pop3_throttle_mailbox_allocated(struct mailbox *box)
{
	struct pop3_throttle_user *tuser =
		POP3_THROTTLE_USER_CONTEXT(box->storage->user);
	struct pop3_throttle_mailbox *tbox;
	struct mail_namespace *ns;

	ns = mailbox_list_get_namespace(box->list);
	if (strcmp(box->name, "INBOX") != 0 ||
	    (ns->flags & NAMESPACE_FLAG_INBOX_USER) == 0)
		return;

	if (tuser->max_msgs_per_unit == 0 && tuser->max_kbytes_per_unit == 0) {
		if (box->storage->user->mail_debug)
			i_debug("pop3-throttle: No limits, disabling");
		return;
	}

	tbox = p_new(box->pool, struct pop3_throttle_mailbox, 1);
	tbox->module_ctx.super = box->v;
	tbox->session_highest_visible_seq = (uint32_t)-1;

	box->v.get_status = pop3_throttle_get_status;

	MODULE_CONTEXT_SET(box, pop3_throttle_storage_module, tbox);
}

static void pop3_throttle_mail_user_created(struct mail_user *user)
{
	struct pop3_throttle_user *tuser;
	const char *value;

	tuser = p_new(user->pool, struct pop3_throttle_user, 1);

	value = mail_user_plugin_getenv(user, "pop3_throttle_max_msgs");
	tuser->max_msgs_per_unit =
		value == NULL ? 0 : strtoul(value, NULL, 10);

	value = mail_user_plugin_getenv(user, "pop3_throttle_max_kbytes");
	tuser->max_kbytes_per_unit =
		value == NULL ? 0 : strtoul(value, NULL, 10);

	if (user->mail_debug) {
		i_debug("pop3-throttle: max_msgs=%u max_kbytes=%u",
			tuser->max_msgs_per_unit, tuser->max_kbytes_per_unit);
	}

	MODULE_CONTEXT_SET(user, pop3_throttle_user_module, tuser);
}

static struct mail_storage_hooks pop3_throttle_mail_storage_hooks = {
	.mailbox_opened = pop3_throttle_mailbox_opened,
	.mailbox_allocated = pop3_throttle_mailbox_allocated,
	.mail_user_created = pop3_throttle_mail_user_created
};

void pop3_throttle_plugin_init(struct module *module)
{
	mail_storage_hooks_add(module, &pop3_throttle_mail_storage_hooks);
}

void pop3_throttle_plugin_deinit(void)
{
	mail_storage_hooks_remove(&pop3_throttle_mail_storage_hooks);
}