/* Functions */
void pgsql_plugin(int pipe_fd, struct configuration *cfgptr, void *ptr)
{
struct pkt_data *data;
struct ports_table pt;
struct pollfd pfd;
struct insert_data idata;
time_t refresh_deadline;
int timeout, refresh_timeout, amqp_timeout;
int ret, num;
struct ring *rg = &((struct channels_list_entry *)ptr)->rg;
struct ch_status *status = ((struct channels_list_entry *)ptr)->status;
struct plugins_list_entry *plugin_data = ((struct channels_list_entry *)ptr)->plugin;
int datasize = ((struct channels_list_entry *)ptr)->datasize;
u_int32_t bufsz = ((struct channels_list_entry *)ptr)->bufsize;
struct networks_file_data nfd;
char *dataptr;
unsigned char *rgptr;
int pollagain = TRUE;
u_int32_t seq = 1, rg_err_count = 0;
struct extra_primitives extras;
struct primitives_ptrs prim_ptrs;
#ifdef WITH_RABBITMQ
struct p_amqp_host *amqp_host = &((struct channels_list_entry *)ptr)->amqp_host;
#endif
memcpy(&config, cfgptr, sizeof(struct configuration));
memcpy(&extras, &((struct channels_list_entry *)ptr)->extras, sizeof(struct extra_primitives));
recollect_pipe_memory(ptr);
pm_setproctitle("%s [%s]", "PostgreSQL Plugin", config.name);
memset(&idata, 0, sizeof(idata));
if (config.pidfile) write_pid_file_plugin(config.pidfile, config.type, config.name);
if (config.logfile) {
fclose(config.logfile_fd);
config.logfile_fd = open_logfile(config.logfile, "a");
}
sql_set_signals();
sql_init_default_values(&extras);
PG_init_default_values(&idata);
PG_set_callbacks(&sqlfunc_cbr);
sql_set_insert_func();
/* some LOCAL initialization AFTER setting some default values */
reload_map = FALSE;
idata.now = time(NULL);
refresh_deadline = idata.now;
idata.cfg = &config;
sql_init_maps(&extras, &prim_ptrs, &nt, &nc, &pt);
sql_init_global_buffers();
sql_init_historical_acct(idata.now, &idata);
sql_init_triggers(idata.now, &idata);
sql_init_refresh_deadline(&refresh_deadline);
if (config.pipe_amqp) {
plugin_pipe_amqp_compile_check();
#ifdef WITH_RABBITMQ
pipe_fd = plugin_pipe_amqp_connect_to_consume(amqp_host, plugin_data);
amqp_timeout = plugin_pipe_amqp_set_poll_timeout(amqp_host, pipe_fd);
#endif
}
else setnonblocking(pipe_fd);
/* building up static SQL clauses */
idata.num_primitives = PG_compose_static_queries();
glob_num_primitives = idata.num_primitives;
/* handling logfile template stuff */
te = sql_init_logfile_template(&th);
INIT_BUF(logbuf);
/* setting up environment variables */
SQL_SetENV();
sql_link_backend_descriptors(&bed, &p, &b);
/* plugin main loop */
for(;;) {
poll_again:
status->wakeup = TRUE;
calc_refresh_timeout(refresh_deadline, idata.now, &refresh_timeout);
pfd.fd = pipe_fd;
pfd.events = POLLIN;
timeout = MIN(refresh_timeout, (amqp_timeout ? amqp_timeout : INT_MAX));
ret = poll(&pfd, (pfd.fd == ERR ? 0 : 1), timeout);
if (ret <= 0) {
if (getppid() == 1) {
Log(LOG_ERR, "ERROR ( %s/%s ): Core process *seems* gone. Exiting.\n", config.name, config.type);
exit_plugin(1);
}
if (ret < 0) goto poll_again;
}
idata.now = time(NULL);
now = idata.now;
if (config.sql_history) {
while (idata.now > (idata.basetime + idata.timeslot)) {
time_t saved_basetime = idata.basetime;
idata.basetime += idata.timeslot;
if (config.sql_history == COUNT_MONTHLY)
idata.timeslot = calc_monthly_timeslot(idata.basetime, config.sql_history_howmany, ADD);
glob_basetime = idata.basetime;
idata.new_basetime = saved_basetime;
glob_new_basetime = saved_basetime;
}
}
#ifdef WITH_RABBITMQ
if (config.pipe_amqp && pipe_fd == ERR) {
if (timeout == amqp_timeout) {
pipe_fd = plugin_pipe_amqp_connect_to_consume(amqp_host, plugin_data);
amqp_timeout = plugin_pipe_amqp_set_poll_timeout(amqp_host, pipe_fd);
}
else amqp_timeout = plugin_pipe_amqp_calc_poll_timeout_diff(amqp_host, idata.now);
}
#endif
switch (ret) {
case 0: /* poll(): timeout */
if (qq_ptr) sql_cache_flush(queries_queue, qq_ptr, &idata, FALSE);
sql_cache_handle_flush_event(&idata, &refresh_deadline, &pt);
break;
default: /* poll(): received data */
read_data:
if (!config.pipe_amqp) {
if (!pollagain) {
seq++;
seq %= MAX_SEQNUM;
if (seq == 0) rg_err_count = FALSE;
idata.now = time(NULL);
now = idata.now;
}
else {
if ((ret = read(pipe_fd, &rgptr, sizeof(rgptr))) == 0)
exit_plugin(1); /* we exit silently; something happened at the write end */
}
if ((rg->ptr + bufsz) > rg->end) rg->ptr = rg->base;
if (((struct ch_buf_hdr *)rg->ptr)->seq != seq) {
if (!pollagain) {
pollagain = TRUE;
goto poll_again;
}
else {
rg_err_count++;
if (config.debug || (rg_err_count > MAX_RG_COUNT_ERR)) {
Log(LOG_ERR, "ERROR ( %s/%s ): We are missing data.\n", config.name, config.type);
Log(LOG_ERR, "If you see this message once in a while, discard it. Otherwise some solutions follow:\n");
Log(LOG_ERR, "- increase shared memory size, 'plugin_pipe_size'; now: '%u'.\n", config.pipe_size);
Log(LOG_ERR, "- increase buffer size, 'plugin_buffer_size'; now: '%u'.\n", config.buffer_size);
Log(LOG_ERR, "- increase system maximum socket size.\n\n");
}
seq = ((struct ch_buf_hdr *)rg->ptr)->seq;
}
}
pollagain = FALSE;
memcpy(pipebuf, rg->ptr, bufsz);
rg->ptr += bufsz;
}
#ifdef WITH_RABBITMQ
else {
ret = p_amqp_consume_binary(amqp_host, pipebuf, config.buffer_size);
if (ret) pipe_fd = ERR;
seq = ((struct ch_buf_hdr *)pipebuf)->seq;
amqp_timeout = plugin_pipe_amqp_set_poll_timeout(amqp_host, pipe_fd);
}
#endif
/* lazy sql refresh handling */
if (idata.now > refresh_deadline) {
if (qq_ptr) sql_cache_flush(queries_queue, qq_ptr, &idata, FALSE);
sql_cache_handle_flush_event(&idata, &refresh_deadline, &pt);
}
else {
if (config.sql_trigger_exec) {
while (idata.now > idata.triggertime && idata.t_timeslot > 0) {
sql_trigger_exec(config.sql_trigger_exec);
idata.triggertime += idata.t_timeslot;
if (config.sql_trigger_time == COUNT_MONTHLY)
idata.t_timeslot = calc_monthly_timeslot(idata.triggertime, config.sql_trigger_time_howmany, ADD);
}
}
}
data = (struct pkt_data *) (pipebuf+sizeof(struct ch_buf_hdr));
Log(LOG_DEBUG, "DEBUG ( %s/%s ): buffer received seq=%u num_entries=%u\n", config.name, config.type, seq, ((struct ch_buf_hdr *)pipebuf)->num);
while (((struct ch_buf_hdr *)pipebuf)->num > 0) {
for (num = 0; primptrs_funcs[num]; num++)
(*primptrs_funcs[num])((u_char *)data, &extras, &prim_ptrs);
for (num = 0; net_funcs[num]; num++)
(*net_funcs[num])(&nt, &nc, &data->primitives, prim_ptrs.pbgp, &nfd);
if (config.ports_file) {
if (!pt.table[data->primitives.src_port]) data->primitives.src_port = 0;
if (!pt.table[data->primitives.dst_port]) data->primitives.dst_port = 0;
}
if (config.pkt_len_distrib_bins_str &&
config.what_to_count_2 & COUNT_PKT_LEN_DISTRIB)
evaluate_pkt_len_distrib(data);
prim_ptrs.data = data;
(*insert_func)(&prim_ptrs, &idata);
((struct ch_buf_hdr *)pipebuf)->num--;
if (((struct ch_buf_hdr *)pipebuf)->num) {
dataptr = (unsigned char *) data;
if (!prim_ptrs.vlen_next_off) dataptr += datasize;
else dataptr += prim_ptrs.vlen_next_off;
data = (struct pkt_data *) dataptr;
}
}
if (!config.pipe_amqp) goto read_data;
}
}
}
int main(int argc, char **argv)
{
struct insert_data idata;
PGresult *ret;
FILE *f;
unsigned char fbuf[SRVBUFLEN];
char logfile[SRVBUFLEN];
char default_pwd[] = "arealsmartpwd";
int have_pwd = 0, have_logfile = 0, n;
int result = 0, position = 0, howmany = 0;
int do_nothing = 0;
char *cl_sql_host = NULL, *cl_sql_user = NULL, *cl_sql_db = NULL, *cl_sql_table = NULL;
char *sql_host;
struct template_entry *teptr;
int tot_size = 0, cnt = 0;
u_char *te;
struct template_header th;
struct db_cache data;
/* getopt() stuff */
extern char *optarg;
extern int optind, opterr, optopt;
int errflag = 0, cp;
memset(&idata, 0, sizeof(idata));
memset(sql_data, 0, sizeof(sql_data));
memset(update_clause, 0, sizeof(update_clause));
memset(insert_clause, 0, sizeof(insert_clause));
memset(lock_clause, 0, sizeof(lock_clause));
memset(where, 0, sizeof(where));
memset(values, 0, sizeof(values));
memset(&data, 0, sizeof(data));
memset(timebuf, 0, sizeof(timebuf));
pp_size = sizeof(struct db_cache);
while (!errflag && ((cp = getopt(argc, argv, ARGS)) != -1)) {
switch (cp) {
case 'd':
debug = TRUE;
break;
case 'f':
strlcpy(logfile, optarg, sizeof(logfile));
have_logfile = TRUE;
break;
case 'o':
position = atoi(optarg);
if (!position) {
printf("ERROR: invalid offset. Exiting.\n");
exit(1);
}
break;
case 'n':
howmany = atoi(optarg);
if (!howmany) {
printf("ERROR: invalid number of elements. Exiting.\n");
exit(1);
}
break;
case 't':
do_nothing = TRUE;
break;
case 'i':
sql_dont_try_update = TRUE;
break;
case 'e':
sql_history_since_epoch = TRUE;
break;
case 'P':
strlcpy(sql_pwd, optarg, sizeof(sql_pwd));
have_pwd = TRUE;
break;
case 'U':
cl_sql_user = malloc(SRVBUFLEN);
memset(cl_sql_user, 0, SRVBUFLEN);
strlcpy(cl_sql_user, optarg, SRVBUFLEN);
break;
case 'D':
cl_sql_db = malloc(SRVBUFLEN);
memset(cl_sql_db, 0, SRVBUFLEN);
strlcpy(cl_sql_db, optarg, SRVBUFLEN);
break;
case 'H':
cl_sql_host = malloc(SRVBUFLEN);
memset(cl_sql_host, 0, SRVBUFLEN);
strlcpy(cl_sql_host, optarg, SRVBUFLEN);
break;
case 'T':
cl_sql_table = malloc(SRVBUFLEN);
memset(cl_sql_table, 0, SRVBUFLEN);
strlcpy(cl_sql_table, optarg, SRVBUFLEN);
break;
case 'h':
usage(argv[0]);
exit(0);
break;
default:
usage(argv[0]);
exit(1);
}
}
/* searching for user supplied values */
if (!howmany) howmany = -1;
if (!have_pwd) memcpy(sql_pwd, default_pwd, sizeof(default_pwd));
if (!have_logfile) {
usage(argv[0]);
printf("\nERROR: missing logfile (-f)\nExiting...\n");
exit(1);
}
f = fopen(logfile, "r");
if (!f) {
printf("ERROR: %s does not exists\nExiting...\n", logfile);
exit(1);
}
fread(&lh, sizeof(lh), 1, f);
lh.sql_table_version = ntohs(lh.sql_table_version);
lh.sql_optimize_clauses = ntohs(lh.sql_optimize_clauses);
lh.sql_history = ntohs(lh.sql_history);
lh.what_to_count = ntohl(lh.what_to_count);
lh.magic = ntohl(lh.magic);
if (lh.magic == MAGIC) {
if (debug) printf("OK: Valid logfile header read.\n");
printf("sql_db: %s\n", lh.sql_db);
printf("sql_table: %s\n", lh.sql_table);
printf("sql_user: %s\n", lh.sql_user);
printf("sql_host: %s\n", lh.sql_host);
if (cl_sql_db||cl_sql_table||cl_sql_user||cl_sql_host)
printf("OK: Overrided by commandline options:\n");
if (cl_sql_db) printf("sql_db: %s\n", cl_sql_db);
if (cl_sql_table) printf("sql_table: %s\n", cl_sql_table);
if (cl_sql_user) printf("sql_user: %s\n", cl_sql_user);
if (cl_sql_host) printf("sql_host: %s\n", cl_sql_host);
}
else {
printf("ERROR: Invalid magic number. Exiting.\n");
exit(1);
}
/* binding SQL stuff */
if (cl_sql_db) sql_db = cl_sql_db;
else sql_db = lh.sql_db;
if (cl_sql_table) sql_table = cl_sql_table;
else sql_table = lh.sql_table;
if (cl_sql_user) sql_user = cl_sql_user;
else sql_user = lh.sql_user;
if (cl_sql_host) sql_host = cl_sql_host;
else sql_host = lh.sql_host;
fread(&th, sizeof(th), 1, f);
th.magic = ntohl(th.magic);
th.num = ntohs(th.num);
th.sz = ntohs(th.sz);
if (th.magic == TH_MAGIC) {
if (debug) printf("OK: Valid template header read.\n");
if (th.num > N_PRIMITIVES) {
printf("ERROR: maximum number of primitives exceeded. Exiting.\n");
exit(1);
}
te = malloc(th.num*sizeof(struct template_entry));
memset(te, 0, th.num*sizeof(struct template_entry));
fread(te, th.num*sizeof(struct template_entry), 1, f);
}
else {
if (debug) printf("ERROR: no template header found.\n");
exit(1);
}
/* checking template */
if (th.sz >= sizeof(fbuf)) {
printf("ERROR: Objects are too big. Exiting.\n");
exit(1);
}
teptr = (struct template_entry *) te;
for (tot_size = 0, cnt = 0; cnt < th.num; cnt++, teptr++)
tot_size += teptr->size;
if (tot_size != th.sz) {
printf("ERROR: malformed template header. Size mismatch. Exiting.\n");
exit(1);
}
TPL_check_sizes(&th, &data, te);
if (!do_nothing) {
PG_compose_conn_string(&p, sql_host);
if (!PG_DB_Connect2(&p)) {
printf("ALERT: PG_DB_Connect2(): PGSQL daemon failed.\n");
exit(1);
}
}
else {
if (debug) print_header();
}
/* composing the proper (filled with primitives used during
the current execution) SQL strings */
idata.num_primitives = PG_compose_static_queries();
idata.now = time(NULL);
/* handling offset */
if (position) n = fseek(f, (th.sz*position), SEEK_CUR);
/* handling single or iterative request */
if (!do_nothing) ret = PQexec(p.desc, lock_clause);
while(!feof(f)) {
if (!howmany) break;
else if (howmany > 0) howmany--;
memset(fbuf, 0, th.sz);
n = fread(fbuf, th.sz, 1, f);
if (n) {
re++;
TPL_pop(fbuf, &data, &th, te);
if (!do_nothing) result = PG_cache_dbop(&p, &data, &idata);
else {
if (debug) print_data(&data, lh.what_to_count, (position+re));
}
if (!result) we++;
if (re != we) printf("WARN: unable to write element %u.\n", re);
}
}
if (!do_nothing) {
ret = PQexec(p.desc, "COMMIT");
if (PQresultStatus(ret) != PGRES_COMMAND_OK) {
we = 0; /* if we fail to commit, no elements will be written */
PQclear(ret);
}
printf("\nOK: written [%u/%u] elements.\n", we, re);
}
else printf("OK: read [%u] elements.\n", re);
PQfinish(p.desc);
fclose(f);
return 0;
}
