Add ext_time_quota_acl helper

[thirdparty/squid.git] / helpers / external_acl / time_quota / ext_time_quota_acl.cc

/*
 * ext_time_quota_acl: Squid external acl helper for quota on usage.
 *
 * Copyright (C) 2011 Dr. Tilmann Bubeck <t.bubeck@reinform.de>
 *
 * 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 2 of the License, 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, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111, USA.
 */

#if HAVE_CONFIG_H
#include "config.h"
#endif
#include "helpers/defines.h"

#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <stdio.h>
#include <stdlib.h>
#include <stdarg.h>
#if HAVE_UNISTD_H
#include <unistd.h>
#endif
#include <string.h>
#include <time.h>
#if HAVE_GETOPT_H
#include <getopt.h>
#endif

/* At this point all Bit Types are already defined, so we must
   protect from multiple type definition on platform where
   __BIT_TYPES_DEFINED__ is not defined.
 */
#ifndef        __BIT_TYPES_DEFINED__
#define        __BIT_TYPES_DEFINED__
#endif

#if HAVE_DB_185_H
#include <db_185.h>
#elif HAVE_DB_H
#include <db.h>
#endif

#ifndef DEFAULT_QUOTA_DB
#error "Please define DEFAULT_QUOTA_DB preprocessor constant."
#endif

const char *db_path = DEFAULT_QUOTA_DB;
const char *program_name;

DB *db = NULL;

#define KEY_LAST_ACTIVITY            "last-activity"
#define KEY_PERIOD_START             "period-start"
#define KEY_PERIOD_LENGTH_CONFIGURED "period-length-configured"
#define KEY_TIME_BUDGET_LEFT         "time-budget-left"
#define KEY_TIME_BUDGET_CONFIGURED   "time-budget-configured"

/** Maximum size of buffers used to read or display lines. */
#define TQ_BUFFERSIZE                     1024

/** If there is more than this given number of seconds between two
 * sucessive requests, than the second request will be treated as a 
 * new request and the time between first and seconds request will
 * be treated as a activity pause.
 *
 * Otherwise the following request will be treated as belonging to the
 * same activity and the quota will be reduced.
 */
static int pauseLength = 300;

static FILE *logfile = stderr;
static int tq_debug_enabled = false;

static void open_log(const char *logfilename) {
    logfile = fopen(logfilename, "a");
    if ( logfile == NULL ) {
	perror(logfilename);
	logfile = stderr;
    }
}

static void vlog(const char *level, const char *format, va_list args)
{
    time_t now = time(NULL);

    fprintf(logfile, "%ld %s| %s: ", now, program_name, level);
    vfprintf (logfile, format, args);
    fflush(logfile);
}

static void log_debug(const char *format, ...)
{
    va_list args;

    if ( tq_debug_enabled ) {
	va_start (args, format);
	vlog("DEBUG", format, args);
	va_end (args);
    }
}

static void log_info(const char *format, ...)
{
    va_list args;

    va_start (args, format);
    vlog("INFO", format, args);
    va_end (args);
}

static void log_error(const char *format, ...)
{
    va_list args;

    va_start (args, format);
    vlog("ERROR", format, args);
    va_end (args);
}

static void log_fatal(const char *format, ...)
{
    va_list args;

    va_start (args, format);
    vlog("FATAL", format, args);
    va_end (args);
}

static void init_db(void)
{
    log_info("opening time quota database \"%s\".\n", db_path);
    db = dbopen(db_path, O_CREAT | O_RDWR, 0666, DB_BTREE, NULL);
    if (!db) {
        log_fatal("Failed to open time_quota db '%s'\n", db_path);
        exit(1);
    }
}

static void shutdown_db(void)
{
    db->close(db);
}

static void writeTime(const char *user_key, const char *sub_key, time_t t)
{
    char keybuffer[TQ_BUFFERSIZE];
    DBT key, data;

    if ( strlen(user_key) + strlen(sub_key) + 1 + 1 > sizeof(keybuffer) ) {
	log_error("key too long (%s,%s)\n", user_key, sub_key);
    } else {
	snprintf(keybuffer, sizeof(keybuffer), "%s-%s", user_key, sub_key);

	key.data = (void *)keybuffer;
	key.size = strlen(keybuffer);
	data.data = &t;
	data.size = sizeof(t);
	db->put(db, &key, &data, 0);
	log_debug("writeTime(\"%s\", %d)\n", keybuffer, t);
    }
}

static time_t readTime(const char *user_key, const char *sub_key)
{
    char keybuffer[TQ_BUFFERSIZE];
    DBT key, data;
    time_t t = 0;

    if ( strlen(user_key) + 1 + strlen(sub_key) + 1 > sizeof(keybuffer) ) {
	log_error("key too long (%s,%s)\n", user_key, sub_key);
    } else {
	snprintf(keybuffer, sizeof(keybuffer), "%s-%s", user_key, sub_key);

	key.data = (void *)keybuffer;
	key.size = strlen(keybuffer);
	if (db->get(db, &key, &data, 0) == 0) {
	    if (data.size != sizeof(t)) {
		log_error("CORRUPTED DATABASE (%s)\n", keybuffer);
	    } else {
		memcpy(&t, data.data, sizeof(t));
	    }
	}
	log_debug("readTime(\"%s\")=%d\n", keybuffer, t);
    }

    return t;
}

static void parseTime(const char *s, time_t *secs, time_t *start)
{
    double value;
    char unit;
    struct tm *ltime;
    int periodLength = 3600;

    *secs = 0;
    *start = time(NULL);
    ltime = localtime(start);

    sscanf(s, " %lf %c", &value, &unit);
    switch (unit) {
    case 's':
	periodLength = 1;
	break;
    case 'm':
	periodLength = 60;
	*start -= ltime->tm_sec;
	break;
    case 'h':
	periodLength = 3600;
	*start -= ltime->tm_min * 60 + ltime->tm_sec;
	break;
    case 'd':
	periodLength = 24 * 3600;
	*start -= ltime->tm_hour * 3600 + ltime->tm_min * 60 + ltime->tm_sec;
	break;
    case 'w':
	periodLength = 7 * 24 * 3600;
	*start -= ltime->tm_hour * 3600 + ltime->tm_min * 60 + ltime->tm_sec;
	*start -= ltime->tm_wday * 24 * 3600;
	*start += 24 * 3600;         // in europe, the week starts monday
	break;
    default:
	log_error("Wrong time unit \"%c\". Only \"m\", \"h\", \"d\", or \"w\" allowed.\n", unit);
	break;
    }

    *secs = (long)(periodLength * value);
}


/** This function parses the time quota file and stores it
 * in memory.
 */
static void readConfig(const char *filename) 
{
    char line[TQ_BUFFERSIZE];        /* the buffer for the lines read
				   from the dict file */
    char *cp;			/* a char pointer used to parse
				   each line */
    char *username;		/* for the username */
    char *budget;
    char *period;
    FILE *FH;
    time_t t;
    time_t budgetSecs, periodSecs;
    time_t start;

    log_info("reading config file \"%s\".\n", filename);

    FH = fopen(filename, "r");
    if ( FH ) {
	/* the pointer to the first entry in the linked list */
	while ((cp = fgets (line, sizeof(line), FH)) != NULL) {
	    if (line[0] == '#') {
		continue;
	    }
	    if ((cp = strchr (line, '\n')) != NULL) {
		/* chop \n characters */
		*cp = '\0';
	    }
	    log_debug("read config line \"%s\".\n", line);
	    if ((cp = strtok (line, "\t ")) != NULL) {
		username = cp;
		
		/* get the time budget */
		budget = strtok (NULL, "/");
		period = strtok (NULL, "/");

		parseTime(budget, &budgetSecs, &start);
		parseTime(period, &periodSecs, &start);

		log_debug("read time quota for user \"%s\": %lds / %lds starting %lds\n", 
		    username, budgetSecs, periodSecs, start);

		writeTime(username, KEY_PERIOD_START, start);
		writeTime(username, KEY_PERIOD_LENGTH_CONFIGURED, periodSecs);
		writeTime(username, KEY_TIME_BUDGET_CONFIGURED, budgetSecs);
		t = readTime(username, KEY_TIME_BUDGET_CONFIGURED);
		writeTime(username, KEY_TIME_BUDGET_LEFT, t);
	    }
	}
	fclose(FH);
    } else {
	perror(filename);
    }
}

static void processActivity(const char *user_key)
{
     time_t now = time(NULL);
     time_t lastActivity;
     time_t activityLength;
     time_t periodStart;
     time_t periodLength;
     time_t userPeriodLength;
     time_t timeBudgetCurrent;
     time_t timeBudgetConfigured;
     char message[TQ_BUFFERSIZE];

     log_debug("processActivity(\"%s\")\n", user_key);

     // [1] Reset period if over
     periodStart = readTime(user_key, KEY_PERIOD_START);
     if ( periodStart == 0 ) {
	 // This is the first period ever.
	 periodStart = now;
	 writeTime(user_key, KEY_PERIOD_START, periodStart);
     }

     periodLength = now - periodStart;
     userPeriodLength = readTime(user_key, KEY_PERIOD_LENGTH_CONFIGURED);
     if ( userPeriodLength == 0 ) {
	 // This user is not configured. Allow anything.
	 log_debug("No period length found for user \"%s\". Quota for this user disabled.\n", user_key);
	 writeTime(user_key, KEY_TIME_BUDGET_LEFT, pauseLength);
     } else {
	 if ( periodLength >= userPeriodLength ) {
	     // a new period has started.
	     log_debug("New time period started for user \"%s\".\n", user_key);
	     while ( periodStart < now ) {
		 periodStart += periodLength;
	     }
	     writeTime(user_key, KEY_PERIOD_START, periodStart);
	     timeBudgetConfigured = readTime(user_key, KEY_TIME_BUDGET_CONFIGURED);
	     if ( timeBudgetConfigured == 0 ) {
		 log_debug("No time budget configured for user \"%s\". Quota for this user disabled.\n", user_key);
		 writeTime(user_key, KEY_TIME_BUDGET_LEFT, pauseLength);
	     } else {
		 writeTime(user_key, KEY_TIME_BUDGET_LEFT, timeBudgetConfigured);
	     }
	 }
     }

     // [2] Decrease time budget iff activity
     lastActivity = readTime(user_key, KEY_LAST_ACTIVITY);
     if ( lastActivity == 0 ) {
	 // This is the first request ever
	 writeTime(user_key, KEY_LAST_ACTIVITY, now);
     } else {
	 activityLength = now - lastActivity;
	 if ( activityLength >= pauseLength ) {
	     // This is an activity pause.
	     log_debug("Activity pause detected for user \"%s\".\n", user_key);
	     writeTime(user_key, KEY_LAST_ACTIVITY, now);
	 } else {
	     // This is real usage.
	     writeTime(user_key, KEY_LAST_ACTIVITY, now);

	     log_debug("Time budget reduced by %ld for user \"%s\".\n", 
		 activityLength, user_key);
	     timeBudgetCurrent = readTime(user_key, KEY_TIME_BUDGET_LEFT);
	     timeBudgetCurrent -= activityLength;
	     writeTime(user_key, KEY_TIME_BUDGET_LEFT, timeBudgetCurrent);
	 }
     }

     timeBudgetCurrent = readTime(user_key, KEY_TIME_BUDGET_LEFT);
     snprintf(message, TQ_BUFFERSIZE, "message=\"Remaining quota for '%s' is %d seconds.\"", user_key, (int)timeBudgetCurrent);
     if ( timeBudgetCurrent > 0 ) {
	 log_debug("OK %s.\n", message);
	 SEND_OK(message);
     } else {
	 log_debug("ERR %s\n", message);
	 SEND_ERR("Time budget exceeded.");
     }

     db->sync(db, 0);
}

static void usage(void)
{
    log_error("Wrong usage. Please reconfigure in squid.conf.\n");

    fprintf(stderr, "Usage: %s [-d] [-l logfile] [-b dbpath] [-p pauselen] [-h] configfile\n", program_name);
    fprintf(stderr, "	-d            enable debugging output to logfile\n");
    fprintf(stderr, "	-l logfile    log messages to logfile\n");
    fprintf(stderr, "	-b dbpath     Path where persistent session database will be kept\n");
    fprintf(stderr, "	              If option is not used, then " DEFAULT_QUOTA_DB " will be used.\n");
    fprintf(stderr, "	-p pauselen   length in seconds to describe a pause between 2 requests.\n");
    fprintf(stderr, "	-h            show show command line help.\n");
    fprintf(stderr, "configfile is a file containing time quota definitions.\n");
}

int main(int argc, char **argv)
{
    char request[HELPER_INPUT_BUFFER];
    int opt;

    program_name = argv[0];

    while ((opt = getopt(argc, argv, "dp:l:b:h")) != -1) {
        switch (opt) {
        case 'd':
	    tq_debug_enabled = true;
            break;
        case 'l':
	    open_log(optarg);
            break;
        case 'b':
            db_path = optarg;
            break;
        case 'p':
            pauseLength = atoi(optarg);
            break;
        case 'h':
            usage();
            exit(0);
            break;
        }
    }

    log_info("Starting %s\n", __FILE__);
    setbuf(stdout, NULL);

    init_db();

    if ( optind + 1 != argc ) {
	usage();
	exit(1);
    } else {
	readConfig(argv[optind]);
    }

    log_info("Waiting for requests...\n");
    while (fgets(request, HELPER_INPUT_BUFFER, stdin)) {
	// we expect the following line syntax: "%LOGIN 
        const char *user_key = NULL;
        user_key = strtok(request, " \n");

	processActivity(user_key);
    }
    log_info("Ending %s\n", __FILE__);
    shutdown_db();
    return 0;
}