1 files changed, 67 insertions, 43 deletions

diff --git a/src/vmod_dbrw.c b/src/vmod_dbrw.c
index 63d4ea1..beeb9af 100644
--- a/src/vmod_dbrw.c
+++ b/src/vmod_dbrw.c
@@ -1,14 +1,14 @@
 /* This file is part of vmod-dbrw
-   Copyright (C) 2013-2018 Sergey Poznyakoff
+   Copyright (C) 2013-2019 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
@@ -16,24 +16,29 @@
 */
 #include "dbrw.h"
 #include <stdarg.h>
 #include "wordsplit.h"
 #include "pthread.h"
 
 enum {
 	QDISP_NONE,
 	QDISP_APPEND,
 	QDISP_DISCARD
 };
 
+enum {
+	MATCH_REGEX,
+	MATCH_EQ
+};
+
 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, ...)
 {
@@ -145,25 +150,25 @@ argv_free(char **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[])
+parse_flags(const char *arg, int *qdisp, int *flags, int *match, 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)) {
@@ -195,24 +200,28 @@ parse_flags(const char *arg, int *qdisp, int *flags, char status[])
 				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 if (strcmp(ws.ws_wordv[i], "regex") == 0) {
+			*match = MATCH_REGEX;
+		} else if (strcmp(ws.ws_wordv[i], "eq") == 0) {
+			*match = MATCH_EQ;
 		} else {
 			dbrw_error("unrecognized flag: %s", ws.ws_wordv[i]);
 			rc = 1;
 		}
 	}		   
 	
 	wordsplit_free(&ws);
 	return rc;
 }
 
 /* Configure the module.
    BACKEND  - "mysql" or "pgsql"
@@ -268,31 +277,33 @@ vmod_config(VRT_CTX, struct vmod_priv *priv,
 		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->match_type = MATCH_REGEX;
 		conf->status[0] = 0;
 
 		s = findparam(conf->param, "flags");
 		if (s)
 			parse_flags(s,
 				    &conf->qdisp,
 				    &conf->regflags,
+				    &conf->match_type,
 				    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)
@@ -404,99 +415,107 @@ findmatch(VRT_CTX, struct dbrw_connection *conn, char **param)
 	*/	     
 	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;
+		int match_type = conn->conf->match_type;
 		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))
+					&qdisp, &rflags, &match_type, 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);
+		if (match_type == MATCH_EQ) {
+			rc = ((rflags & REG_ICASE)
+			      ? strcasecmp : strcmp) (pat, val);
+		} else {
+			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;
 			}
-			conn->matches = p;
-			conn->matchsize = matchcount;
-		}
 
-		rc = regexec(&re, val, matchcount, conn->matches, 0);
+			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);
+			if (rc == 0)
+				res = expand_backref(ctx, val,
+						     sql_get_column(conn, i, 0),
+						     matchcount,
+						     conn->matches, qry);
 
-		regfree(&re);
+			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;
 		}
 	}
@@ -597,99 +616,104 @@ query_command_expand(char **ret, const char *cmd, size_t len, char **argv,
 {
 	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)
+static int
+do_rewrite(VRT_CTX, struct dbrw_connection *cp, VCL_STRING arg, char **result)
 {
 	struct wordsplit ws, wsenv;
 	int i, rc;
- 	char *res;
 	int wsflags;
+
+	*result = NULL;
 	
 	if (sql_connect(cp) || cp->state != state_connected)
-		return NULL;
+		return -1;
 
 	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;
+		return -1;
 	}
 
 	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;
+				return -1;
 			}
 			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;
+		return -1;
 	}
 
 	rc = sql_query(cp, ws.ws_wordv[0]);
 	wordsplit_free(&ws);
 
 	if (rc) {
-		debug(cp->conf, 1, ("vmod_rewrite: query failed"));
+		dbrw_error("vmod_rewrite: query failed");
 		wordsplit_free(&wsenv);
-		return NULL;
+		return -1;
 	}
 	
-	res = findmatch(ctx, cp, wsenv.ws_wordv);
+	*result = findmatch(ctx, cp, wsenv.ws_wordv);
 	wordsplit_free(&wsenv);
 
 	sql_free_result(cp);
 
-	return res;
+	return 0;
 }
 
 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;
-
+	int rc;
+	
 	AN(conf);
 	debug(conf, 2, ("vmod_rewrite(%s) begin", arg));
 	cp = dbrw_connection_get(conf);
 	AN(cp);
-	res = do_rewrite(ctx, cp, arg);
+	rc = do_rewrite(ctx, cp, arg, &res);
 	dbrw_connection_release(cp);
-	debug(conf, 1, ("vmod_rewrite: res=%s", res ? res : "(NULL)"));	
+	debug(conf, 1, ("vmod_rewrite: %d, res=%s", rc, res ? res : "(NULL)"));	
+	if (rc) {
+		dbrw_sethdr(ctx, HDR_REQ, "\022X-VMOD-DBRW-Error:", "1");
+	}
 	return res;
 }