/* This file is part of vmod-dbrw
   Copyright (C) 2013-2018 Sergey Poznyakoff

   Vmod-dbrw 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.

   Vmod-dbrw 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 vmod-dbrw.  If not, see <http://www.gnu.org/licenses/>.
*/
#include "dbrw.h"
#include <stdarg.h>
#include "wordsplit.h"
#include "pthread.h"

enum {
	QDISP_NONE,
	QDISP_APPEND,
	QDISP_DISCARD
};

void
dbrw_debug(const char *fmt, ...)
{
	va_list ap;
	va_start(ap, fmt);
	vsyslog(LOG_DAEMON|LOG_DEBUG, fmt, ap);
	va_end(ap);
}

void
dbrw_error(const char *fmt, ...)
{
	va_list ap;
	va_start(ap, fmt);
	vsyslog(LOG_DAEMON|LOG_ERR, fmt, ap);
	va_end(ap);
}

static pthread_mutex_t config_pool_mtx = PTHREAD_MUTEX_INITIALIZER;
static VTAILQ_HEAD(, dbrw_config) config_pool =
	VTAILQ_HEAD_INITIALIZER(config_pool);

static pthread_mutex_t connect_pool_mtx = PTHREAD_MUTEX_INITIALIZER;
static VTAILQ_HEAD(, dbrw_connection) connect_pool =
	VTAILQ_HEAD_INITIALIZER(connect_pool);

static struct dbrw_connection *
dbrw_connection_get(struct dbrw_config *cfg)
{
	struct dbrw_connection *cp;
	
	AZ(pthread_mutex_lock(&connect_pool_mtx));
	VTAILQ_FOREACH(cp, &connect_pool, list) {
		CHECK_OBJ_NOTNULL(cp, DBRW_CONNECTION_MAGIC);
		if (cp->conf == cfg && !cp->busy)
			break;
	}
	if (!cp) {
		ALLOC_OBJ(cp, DBRW_CONNECTION_MAGIC);
		AN(cp);
		cp->state = state_init;
		cp->conf = cfg;
		cp->data = NULL;
		cp->matches = NULL;
		cp->matchsize = 0;

		if (sql_init(cp)) {
			free(cp);
			cp = NULL;
		} else
			VTAILQ_INSERT_HEAD(&connect_pool, cp, list);
	} else if (cp->state == state_connected
		   && cfg->idle_timeout >= 0
		   && cp->timestamp + cfg->idle_timeout < time(NULL)) {
		sql_disconnect(cp);
	}
	
	cp->busy = 1;

	if (cfg->idle_timeout == -2) {
		cfg->idle_timeout = sql_idle_timeout(cp);
		if (cfg->idle_timeout < -1)
			cfg->idle_timeout = -1;
	}
	
	pthread_mutex_unlock(&connect_pool_mtx);
	return cp;
}

static void
dbrw_connection_release(struct dbrw_connection *cp)
{
	AZ(pthread_mutex_lock(&connect_pool_mtx));
	cp->busy = 0;
	pthread_mutex_unlock(&connect_pool_mtx);
}

static void
disconnect(void)
{
	AZ(pthread_mutex_lock(&connect_pool_mtx));
	while (!VTAILQ_EMPTY(&connect_pool)) {
		struct dbrw_connection *cp = VTAILQ_FIRST(&connect_pool);
		VTAILQ_REMOVE(&connect_pool, cp, list);
		sql_disconnect(cp);
		sql_destroy(cp);
	}
	pthread_mutex_unlock(&connect_pool_mtx);
}

int
dbrw_event(VRT_CTX, struct vmod_priv *priv, enum vcl_event_e e)
{
	switch (e) {
	case VCL_EVENT_LOAD:
		atexit(disconnect);
		break;
	case VCL_EVENT_DISCARD:
		break;

	default:
		/* ignore */
		break;
	}
	return 0;
}

static void
argv_free(char **argv)
{
	if (!argv)
		return;
	while (*argv) {
		free(*argv);
		++*argv;
	}
}

static int
is_http_status(const char *arg)
{
	if (strlen(arg) != HTTP_STATUS_LEN)
		return 0;
	return  '1' <= arg[0] && arg[0] <= '5' &&
		'0' <= arg[1] && arg[1] <= '9' &&
		'0' <= arg[2] && arg[2] <= '9';
}

static int
parse_flags(const char *arg, int *qdisp, int *flags, char status[])
{
	struct wordsplit ws;
	int rc = 0;
	int i;
	
	if (!arg)
		return 0;
	ws.ws_delim = ",";
	if (wordsplit(arg, &ws, 
		      WRDSF_NOVAR | WRDSF_NOCMD | WRDSF_QUOTE |
		      WRDSF_CESCAPES |
		      WRDSF_DELIM)) {
		dbrw_error("cannot split string `%s': %s",
			   arg, wordsplit_strerror (&ws));
		return 1;
	}

	for (i = 0; i < ws.ws_wordc; i++) {
		if (strcmp(ws.ws_wordv[i], "nocase") == 0 ||
		    strcmp(ws.ws_wordv[i], "NC") == 0)
			*flags |= REG_ICASE;
		else if (strcmp(ws.ws_wordv[i], "case") == 0)
			*flags &= ~REG_ICASE;
		else if (strcmp(ws.ws_wordv[i], "qsappend") == 0 ||
			 strcmp(ws.ws_wordv[i], "QSA") == 0)
			*qdisp = QDISP_APPEND;
		else if (strcmp(ws.ws_wordv[i], "qsdiscard") == 0 ||
			 strcmp(ws.ws_wordv[i], "QSD") == 0)
			*qdisp = QDISP_DISCARD;
		else if (strncmp(ws.ws_wordv[i], "redirect=", 9) == 0) {
			if (!is_http_status(ws.ws_wordv[i] + 9)) {
				dbrw_error("invalid status code: %s",
					   ws.ws_wordv[i] + 9);
				rc = 1;
			} else {
				strncpy(status, ws.ws_wordv[i] + 9,
					HTTP_STATUS_LEN);
				status[HTTP_STATUS_LEN] = 0;
			}
		} else if (strncmp(ws.ws_wordv[i], "R=", 2) == 0) {
			if (!is_http_status(ws.ws_wordv[i] + 2)) {
				dbrw_error("invalid status code: %s",
					   ws.ws_wordv[i] + 2);
				rc = 1;
			} else {
				strncpy(status, ws.ws_wordv[i] + 2,
					HTTP_STATUS_LEN);
				status[HTTP_STATUS_LEN] = 0;
			}
		} else {
			dbrw_error("unrecognized flag: %s", ws.ws_wordv[i]);
			rc = 1;
		}
	}		   
	
	wordsplit_free(&ws);
	return rc;
}

/* Configure the module.
   BACKEND  - "mysql" or "pgsql"
   PARAM    - VAR=VALUE*
   QUERY    - Query to obtain the redirection target   
 */
VCL_VOID
vmod_config(VRT_CTX, struct vmod_priv *priv,
	    VCL_STRING bkname, VCL_STRING param, VCL_STRING query)
{
	struct dbrw_config *conf;
	struct dbrw_backend *backend;

	AZ(priv->priv);
	backend = dbrw_backend_select(bkname);
	if (!backend) {
		dbrw_error("unsupported backend: %s", bkname);
		abort();
	}

	AZ(pthread_mutex_lock(&config_pool_mtx));
	VTAILQ_FOREACH(conf, &config_pool, list) {
		CHECK_OBJ_NOTNULL(conf, DBRW_CONFIG_MAGIC);
		if (conf->backend == backend
		    && strcmp(conf->query, query) == 0
		    && strcmp(conf->param_str, param) == 0)
			break;
	}
	pthread_mutex_unlock(&config_pool_mtx);

	if (!conf) {
		char *s;
		struct wordsplit ws;
		ws.ws_delim = ";";
		if (wordsplit(param, &ws, 
			      WRDSF_NOVAR | WRDSF_NOCMD | WRDSF_QUOTE |
			      WRDSF_CESCAPES |
			      WRDSF_DELIM)) {
			dbrw_error("cannot split string `%s': %s",
				   param, wordsplit_strerror (&ws));
			abort();
		}
		ALLOC_OBJ(conf, DBRW_CONFIG_MAGIC);
		AN(conf);
		conf->param = ws.ws_wordv;
		ws.ws_wordv = NULL;
		ws.ws_wordc = 0;
		wordsplit_free(&ws);
		conf->param_str = strdup(param);
		AN(conf->param_str);
		conf->query = strdup(query);
		AN(conf->query);
		conf->backend = backend;

		if ((s = findparam(conf->param, "debug")) != NULL)
			conf->debug_level = atoi(s);

		if ((s = findparam(conf->param, "timeout")) != NULL)
			conf->idle_timeout = atoi(s);
		else
			conf->idle_timeout = -2;

		conf->qdisp = QDISP_NONE;
		conf->regflags = REG_EXTENDED;
		conf->status[0] = 0;

		s = findparam(conf->param, "flags");
		if (s)
			parse_flags(s,
				    &conf->qdisp,
				    &conf->regflags,
				    conf->status);
	
		AZ(pthread_mutex_lock(&config_pool_mtx));
		VTAILQ_INSERT_HEAD(&config_pool, conf, list);
		pthread_mutex_unlock(&config_pool_mtx);
	}
	priv->priv = conf;
}

static char *
expand_backref(VRT_CTX, const char *str, const char *val,
	       size_t matchcount, regmatch_t *matches, char *qry)
{
	unsigned u;
        char *b, *p;
	
        u = WS_Reserve(WSPTR(ctx), 0);
        p = b = WSPTR(ctx)->f;

	while (*val) {
		if (u == 0) {
			WS_Release(WSPTR(ctx), 0);
			return NULL;
		}
		if (*val == '\\' || *val == '$') {
			if (val[1] >= 0 && val[1] <= '9') {
				int n = val[1] - '0';
				if (n < matchcount) {
					size_t len;

					len = matches[n].rm_eo -
						  matches[n].rm_so;

					if (len > u) {
						WS_Release(WSPTR(ctx), 0);
						return NULL;
					}

					memcpy(p, str + matches[n].rm_so, len);

					p += len;
					u -= len;
					val += 2;
					continue;
				}
			} else if (*val == '\\' && val[1]) {
				*p++ = val[1];
				u--;
				val += 2;
				continue;
			}
		}

		*p++ = *val++;
		u--;
	}

	if (qry) {
		size_t len = strlen(qry);
		if (len + 1 >= u)  {
			WS_Release(WSPTR(ctx), 0);
			return NULL;
		}
		if (strchr(b, '?'))
			*p++ = '&';
		else
			*p++ = '?';
		memcpy(p, qry, len);
		p += len;
		u -= len;
	}
	
	if (u == 0) {
		WS_Release(WSPTR(ctx), 0);
		return NULL;
	}
	*p++ = 0;
	
	WS_ReleaseP(WSPTR(ctx), p);

	return b;
}

#define ISEMPTY(s) ((s) == NULL || (s)[0] == 0)

static void
dbrw_sethdr(VRT_CTX, int where, const char *what, const char *value)
{
	struct gethdr_s s = { where, what };
	VRT_SetHdr(ctx, &s, value, vrt_magic_string_end);
}	

static char *
findmatch(VRT_CTX, struct dbrw_connection *conn, char **param)
{
	unsigned i;
	unsigned nt, nf;
	char *res;
	struct wordsplit ws;
	int wsflags = WRDSF_NOCMD | WRDSF_QUOTE |
		      WRDSF_NOSPLIT |
		      WRDSF_ENV | WRDSF_UNDEF;
	
	if ((nt = sql_num_tuples(conn)) == 0 ||
	    (nf = sql_num_fields(conn)) == 0)
		return NULL;

	debug(conn->conf, 2, ("query returned %u tuples, %u columns", nt, nf));
	
	if (nf < 3)
		return WS_Copy(WSPTR(ctx), sql_get_column(conn, 0, 0), -1);

	/* Three or four fields:
	     result
	     pattern
	     value
	     [flags]
	*/	     
	ws.ws_env = (const char **)param;
	res = NULL;
	for (i = 0; i < nt; i++) {
		const char *pat = sql_get_column(conn, i, 1);
		const char *val = sql_get_column(conn, i, 2);
		regex_t re;
		int rc;
		size_t matchcount;     /* Number of used entries in matches */
		int rflags = conn->conf->regflags;
		char status[HTTP_STATUS_LEN+1];
		int qdisp = conn->conf->qdisp;
		char *qry = NULL;

		if (ISEMPTY(pat)) {
			res = WS_Copy(WSPTR(ctx), sql_get_column(conn, i, 0), -1);
			break;
		}
		
		debug(conn->conf, 1, ("considering \"%s\" ~ \"%s\"", pat, val));
		if (wordsplit(val, &ws, wsflags)) {
			dbrw_error("cannot expand string `%s': %s",
				      val, wordsplit_strerror(&ws));
			continue;
		}
		wsflags |= WRDSF_REUSE;

		val = ws.ws_wordv[0];
		debug(conn->conf, 1, ("expanded to \"%s\" ~ \"%s\"",
				      pat, val));

		strcpy(status, conn->conf->status);
		if (nf == 4) {
			if (parse_flags(sql_get_column(conn, i, 3),
					&qdisp, &rflags, status))
				continue;
		}
		
		switch (qdisp) {
		case QDISP_NONE:
			break;
		case QDISP_APPEND:
			qry = strchr(val, '?');
			if (qry)
				*qry++ = 0;
			break;
		case QDISP_DISCARD:
			qry = strchr(val, '?');
			if (qry) {
				debug(conn->conf, 1,
				      ("discarding query part \"%s\"", qry));
				*qry = 0;
				qry = NULL;
			}
		}
			
		rc = regcomp(&re, pat, rflags);
		if (rc) {
			char errbuf[512];

			regerror(rc, &re, errbuf, sizeof(errbuf));
			dbrw_error("cannot compile regexp \"%s\": %s",
				      pat, errbuf);
			continue;
		}

		matchcount = re.re_nsub + 1;
		if (conn->matchsize < matchcount) {
			void *p = realloc(conn->matches,
					sizeof(conn->matches[0]) * matchcount);
			if (!p) {
				dbrw_error("not enough memory");
				regfree(&re);
				continue;
			}
			conn->matches = p;
			conn->matchsize = matchcount;
		}

		rc = regexec(&re, val, matchcount, conn->matches, 0);
		
		if (rc == 0)
			res = expand_backref(ctx, val,
					     sql_get_column(conn, i, 0),
					     matchcount, conn->matches, qry);

		regfree(&re);
		if (rc == 0) {
			debug(conn->conf, 1, ("match"));
			if (status[0]) {
				debug(conn->conf, 1,
				      ("setting status %s", status));
				dbrw_sethdr(ctx, HDR_REQ, 
					    "\023X-VMOD-DBRW-Status:",
					    status);
			}
			break;
		}
	}

	if (wsflags & WRDSF_REUSE)
		wordsplit_free(&ws);

	return res;
}

static int
expand_error(char **ret, char const *func, char const *msg)
{
	static char delim[] = ": ";
	*ret = malloc(strlen(func) + strlen(msg) + 1);
	if (*ret) {
		strcat(strcat(strcpy(*ret, func), delim), msg);
		return WRDSE_USERERR;
	} else
		return WRDSE_NOSPACE;
}

static int
expand_urlprefixes(struct dbrw_connection *cp, char **argv, char **ret)
{
	char *arg;
	size_t n, len, i, j;
	char *q, *res;
	
	if (argv[1] == NULL || argv[2] != NULL) 
		return expand_error(ret, argv[0], "bad arguments");

	/* Create a copy of the argument */
	if (cp->conf->backend->sql_escape) {
		arg = sql_escape(cp, argv[1]);
	} else {
		arg = strdup(argv[1]);
	}
	if (!arg)
		return WRDSE_NOSPACE;

	/* Cut off eventual query */
	i = j = strcspn(arg, "?");
	arg[i] = 0;

	/* Compute the resulting length */
	len = i;
	n = 1;
	for (; i > 0; i--) {
		if (arg[i] == '/') {
			len += i;
			n++;
		}
	}

	/* Count quotes around each member */
	len += n * 2 + n - 1;  

	/* Allocate the result */
	res = malloc(len + 1);
	if (!res) {
		free(arg);
		return WRDSE_NOSPACE;
	}

	/* Format the result */
	q = res;
	i = j;
	while (i) {
		if (q > res)
			*q++ = ',';
		*q++ = '\'';
		memcpy(q, arg, i);
		q += i;
		*q++ = '\'';
		i--;
		while (i > 0 && arg[i] != '/')
			i--;
	} 
	*q = 0;
	*ret = res;
	free(arg);
	
	return WRDSE_OK;
}

static struct expcom {
	char *com;
	int (*exp) (struct dbrw_connection *, char **, char **);
} expcomtab[] = {
	{ "urlprefixes", expand_urlprefixes },
	{ NULL }
};

static int
query_command_expand(char **ret, const char *cmd, size_t len, char **argv,
		     void *clos)
{
	struct expcom *ec;
	static char diagmsg[] = "unknown command: ";

	for (ec = expcomtab; ec->com; ec++) {
		if (strcmp(ec->com, argv[0]) == 0)
			return ec->exp(clos, argv, ret);
	}

	return expand_error(ret, argv[0], "unknown command");
}

static char *
do_rewrite(VRT_CTX, struct dbrw_connection *cp, VCL_STRING arg)
{
	struct wordsplit ws, wsenv;
	int i, rc;
 	char *res;
	int wsflags;
	
	if (sql_connect(cp) || cp->state != state_connected)
		return NULL;

	debug(cp->conf, 2, ("vmod_rewrite: splitting arg"));
	wsenv.ws_delim = ";";
	if (wordsplit(arg, &wsenv, WRDSF_NOVAR|WRDSF_NOCMD|WRDSF_DELIM)) {
		dbrw_error("cannot split string `%s': %s",
			      arg, wordsplit_strerror(&wsenv));
		return NULL;
	}

	if (cp->conf->backend->sql_escape) {
		debug(cp->conf, 2, ("escaping variables"));
		for (i = 0; i < wsenv.ws_wordc; i++) {
			char *p = sql_escape(cp, wsenv.ws_wordv[i]);
			if (!p) {
				dbrw_error("cannot expand argument");
				wordsplit_free(&wsenv);
				return NULL;
			}
			free(wsenv.ws_wordv[i]);
			wsenv.ws_wordv[i] = p;
			debug(cp->conf, 3, ("%d: %s",i,p));
		}
	}

	debug(cp->conf, 2, ("expanding query {\"%s\"}", cp->conf->query));
	ws.ws_env = (const char **)wsenv.ws_wordv;
	ws.ws_command = query_command_expand;
	ws.ws_closure = cp;
	wsflags = WRDSF_NOSPLIT | WRDSF_CLOSURE | WRDSF_ENV | WRDSF_UNDEF;

	if (cp->conf->debug_level == 100) {
		ws.ws_debug = dbrw_debug;
		wsflags |= WRDSF_DEBUG | WRDSF_SHOWDBG;
	}

	rc = wordsplit(cp->conf->query, &ws, wsflags);

	if (rc) {
		dbrw_error("cannot expand query `%s': %s",
			   cp->conf->query, wordsplit_strerror(&ws));
		wordsplit_free(&wsenv);
		return NULL;
	}

	rc = sql_query(cp, ws.ws_wordv[0]);
	wordsplit_free(&ws);

	if (rc) {
		debug(cp->conf, 1, ("vmod_rewrite: query failed"));
		wordsplit_free(&wsenv);
		return NULL;
	}
	
	res = findmatch(ctx, cp, wsenv.ws_wordv);
	wordsplit_free(&wsenv);

	sql_free_result(cp);

	return res;
}

VCL_STRING
vmod_rewrite(VRT_CTX, struct vmod_priv *priv, VCL_STRING arg)
{
	struct dbrw_config *conf = priv->priv;
	struct dbrw_connection *cp;
 	char *res = NULL;

	AN(conf);
	debug(conf, 2, ("vmod_rewrite(%s) begin", arg));
	cp = dbrw_connection_get(conf);
	AN(cp);
	res = do_rewrite(ctx, cp, arg);
	dbrw_connection_release(cp);
	debug(conf, 1, ("vmod_rewrite: res=%s", res ? res : "(NULL)"));	
	return res;
}