/* 
   This file is part of Mailfromd.
   Copyright (C) 2005-2020 Sergey Poznyakoff
   Copyright (C) 2009 Netservers Ltd

   This program is free software; you can redistribute it and/or modify
   it under the terms of the GNU General Public License as published by
   the Free Software Foundation; either version 3, or (at your option)
   any later version.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU General Public License for more details.

   You should have received a copy of the GNU General Public License
   along with this program.  If not, see <http://www.gnu.org/licenses/>. */

#ifdef HAVE_CONFIG_H
# include <config.h>
#endif

#include <mailutils/mailutils.h>
#include "libmf.h"
#include "mfdb.h"
#include "filenames.h"
#include "inttypes.h"

static void tbf_rate_print_item(struct mu_dbm_datum const *key,
				struct mu_dbm_datum const *val);
static int tbf_rate_expire_item(struct mu_dbm_datum const *content);

static struct db_format tbf_rate_format_struct = {
	"tbf",
	"tbf.db",
	1,
	0,
	DEFAULT_EXPIRE_INTERVAL,
	tbf_rate_print_item,
	tbf_rate_expire_item
};

struct db_format *tbf_rate_format = &tbf_rate_format_struct;


/*
  This implements a classical token bucket filter algorithm, with
  the minor optimization that the number of tokens in each bucket
  is interpolated and thus only recalculated when needed.
  
  Tokens are added to the bucket indentified by the key 'email'
  at constant rate of 1 token per 'interval' microseconds, to a
  maximum of 'burst_size' tokens.

  Each call to check_tbf_rate() will:
     return ret=1 (true) if 'cost' tokens are available, and
                         remove 'cost' tokens from the bucket

  or return ret=0 (false) if there are less than 'cost' tokens
                          available.

  If no bucket is found for the specified key, a new bucket is
  created and initialized to contain 'burst_size' tokens.
  
  Ben McKeegan <ben@netservers.co.uk>
*/

struct tbf_rate_result {
	uint64_t timestamp;  /* microseconds since epoch */
	time_t expirytime;
	size_t tokens;       /* tokens available */
};

#ifndef USEC_PER_SEC
# define USEC_PER_SEC  1000000L
#endif

mf_status
check_tbf_rate(char *email, int *ret,
	       size_t cost, size_t interval, size_t burst_size)
{
	mu_dbm_file_t db;
	struct mu_dbm_datum key;
	struct mu_dbm_datum contents;
	int local_contents = 0;
	struct tbf_rate_result *rp, tbf_rate;
	struct timeval tv;
	uint64_t now;
	int res;

	if (interval == 0 || burst_size == 0)
		return mf_failure;

	if (!cost) {
		/* cost free, so don't waste time on database access */
		*ret = 1;
		return mf_success;
	}
	if (cost > burst_size) {
		/* impossibly expensive, so don't waste time on
		   database access */
		*ret = 0;
		return mf_success;
	}

	mu_debug(tbf_rate_format->debug_handle, MU_DEBUG_TRACE5,
		 ("getting TBF rate info for %s", email));
	db = mf_dbm_open(tbf_rate_format->dbname, MU_STREAM_RDWR);
	if (!db) {
		mu_error(_("mf_dbm_open(%s) failed: %s"),
			 tbf_rate_format->dbname,
			 mu_dbm_strerror(db));
		return mf_failure;
	}
  
	memset(&key, 0, sizeof key);
	memset(&contents, 0, sizeof contents);
	key.mu_dptr = email;
	key.mu_dsize = strlen(email) + 1;
	
	gettimeofday(&tv,NULL);
	now = (uint64_t)tv.tv_sec * USEC_PER_SEC + (uint64_t)tv.tv_usec;

	if ((res = mu_dbm_fetch(db, &key, &contents)) == 0) {
		uint64_t elapsed;
		uint64_t tokens;
		
		rp = (struct tbf_rate_result *) contents.mu_dptr;
		/* calculate elapsed time and number of new tokens since
		   last add */;
		elapsed = now - rp->timestamp;
		tokens = elapsed / interval; /* partial tokens ignored */
		/* timestamp set to time of most recent token */
		rp->timestamp += tokens * interval; 
		
		/* add existing tokens to 64bit counter to prevent overflow
		   in range check */
		tokens += rp->tokens;
		if (tokens >= burst_size)
			rp->tokens = burst_size;
		else
			rp->tokens = (size_t)tokens;
		
		mu_debug(tbf_rate_format->debug_handle, MU_DEBUG_TRACE5,
			 ("found, elapsed time: %"PRIu64
			  " us, new tokens: %"PRIu64", total: %lu ",
			  elapsed, tokens, (unsigned long) rp->tokens));
 	} else {
		if (res != MU_ERR_NOENT)
			mu_error(_("cannot fetch `%s' from `%s': %s"),
				 email, tbf_rate_format->dbname,
				 mu_dbm_strerror(db));	
		/* Initialize the structure */
		tbf_rate.timestamp = now;
		tbf_rate.tokens = burst_size;
		rp = &tbf_rate;
		local_contents = 1;
	}
  
	if (cost <= rp->tokens) {
		*ret = 1;
		rp->tokens -= cost;
		mu_debug(tbf_rate_format->debug_handle, MU_DEBUG_TRACE5,
			 ("tbf_rate matched %s", email));
	} else {
		*ret = 0;
		mu_debug(tbf_rate_format->debug_handle, MU_DEBUG_TRACE5,
			 ("tbf_rate overlimit on %s", email));
	}

	/* record may expire as soon as enough tokens have been
	   delivered to overflow the bucket, since a new
	   record may be created on demand with a full bucket */
	rp->expirytime = (time_t)
		(((uint64_t)interval *
		  (uint64_t)(1 + burst_size - rp->tokens) +
		  rp->timestamp) / USEC_PER_SEC) + 1;

	/* Update the db */
	contents.mu_dptr = (void*)rp;
	contents.mu_dsize = sizeof(*rp);
	res = mu_dbm_store(db, &key, &contents, 1);
	if (res)
		mu_error (_("cannot store datum `%s' into `%s': %s"),
			  email, tbf_rate_format->dbname,
			  res == MU_ERR_FAILURE ?
			   mu_dbm_strerror(db) : mu_strerror(res));

	if (!local_contents)
		mu_dbm_datum_free(&contents);
  
	mu_dbm_close(db);
	return mf_success;
}

static void
tbf_rate_print_item(struct mu_dbm_datum const *key,
		    struct mu_dbm_datum const *val)
{
	const struct tbf_rate_result *res = (const struct tbf_rate_result *)
		                              val->mu_dptr;
	int size = key->mu_dsize - 1; /* Size includes the trailing nul */

	printf("%*.*s tok:%lu@", size, size, key->mu_dptr,
	       (unsigned long) res->tokens);
	format_time_str(stdout, (time_t) (res->timestamp / USEC_PER_SEC));
	printf(" exp:");
	format_time_str(stdout, res->expirytime);
	printf("\n");
}

static int
tbf_rate_expire_item(struct mu_dbm_datum const *content)
{
	struct tbf_rate_result const *res =
		(struct tbf_rate_result const *) content->mu_dptr;
	return tbf_rate_format->expire_interval &&
		time(NULL) - res->expirytime >
		tbf_rate_format->expire_interval;
}