rizin/librz/util/sdb/src/main.c

// SPDX-FileCopyrightText: 2011-2020 pancake <pancake@nopcode.org>
// SPDX-License-Identifier: MIT

#include <signal.h>
#include <stdio.h>
#include <string.h>
#include <fcntl.h>
#include "sdb.h"
#include "sdb_private.h"

#define MODE_ZERO '0'
#define MODE_JSON 'j'
#define MODE_DFLT 0

static int save = 0;
static Sdb *s = NULL;
static ut32 options = SDB_OPTION_FS | SDB_OPTION_NOSTAMP;

static void terminate(RZ_UNUSED int sig) {
	if (!s) {
		return;
	}
	if (save && !sdb_sync(s)) {
		sdb_free(s);
		s = NULL;
		exit(1);
	}
	sdb_free(s);
	exit(sig < 2 ? sig : 0);
}

static void write_null(void) {
	write_(1, "", 1);
}

#define BS          128
#define USE_SLURPIN 1

static char *slurp(FILE *f, size_t *sz) {
	int blocksize = BS;
	static int bufsize = BS;
	static char *next = NULL;
	static size_t nextlen = 0;
	size_t len, rr, rr2;
	char *tmp, *buf = NULL;
	if (sz) {
		*sz = 0;
	}
#if USE_SLURPIN
	if (!sz) {
		/* this is faster but have limits */
		/* run test/add10k.sh script to benchmark */
		const int buf_size = 96096;

		buf = calloc(1, buf_size);
		if (!buf) {
			return NULL;
		}

		if (!fgets(buf, buf_size, f)) {
			free(buf);
			return NULL;
		}
		if (feof(f)) {
			free(buf);
			return NULL;
		}

		size_t buf_len = strlen(buf);
		if (buf_len > 0) {
			buf[buf_len - 1] = '\0';
		}

		char *newbuf = realloc(buf, buf_len + 1);
		// realloc behaves like free if buf_len is 0
		if (!newbuf) {
			return buf;
		}
		return newbuf;
	}
#endif
	buf = calloc(BS + 1, 1);
	if (!buf) {
		return NULL;
	}

	len = 0;
	for (;;) {
		if (next) {
			free(buf);
			buf = next;
			bufsize = nextlen + blocksize;
			// len = nextlen;
			rr = nextlen;
			rr2 = fread(buf + nextlen, 1, blocksize, f);
			if (rr2 > 0) {
				rr += rr2;
				bufsize += rr2;
			}
			next = NULL;
			nextlen = 0;
		} else {
			rr = fread(buf + len, 1, blocksize, f);
		}
		if (rr < 1) { // EOF
			buf[len] = 0;
			next = NULL;
			break;
		}
		len += rr;
		// buf[len] = 0;
#if !USE_SLURPIN
		if (!sz) {
			char *nl = strchr(buf, '\n');
			if (nl) {
				*nl++ = 0;
				int nlen = nl - buf;
				nextlen = len - nlen;
				if (nextlen > 0) {
					next = malloc(nextlen + blocksize + 1);
					if (!next) {
						eprintf("Cannot malloc %d\n", nextlen);
						break;
					}
					memcpy(next, nl, nextlen);
					if (!*next) {
						next = NULL;
					} else {
						//	continue;
					}
				} else {
					next = NULL;
					nextlen = 0; // strlen (next);;
				}
				break;
			}
		}
#endif
		bufsize += blocksize;
		tmp = realloc(buf, bufsize + 1);
		if (!tmp) {
			bufsize -= blocksize;
			break;
		}
		memset(tmp + bufsize - blocksize, 0, blocksize);
		buf = tmp;
	}
	if (sz) {
		*sz = len;
	}
	if (len < 1) {
		free(buf);
		return buf = NULL;
	}
	buf[len] = 0;
	return buf;
}

#if HAVE_HEADER_SYS_MMAN_H
static void synchronize(RZ_UNUSED int sig) {
	// TODO: must be in sdb_sync() or wat?
	sdb_sync(s);
	Sdb *n = sdb_new(s->path, s->name, s->lock);
	if (n) {
		sdb_config(n, options);
		sdb_free(s);
		s = n;
	}
}
#endif

static int sdb_grep_dump(const char *dbname, int fmt, bool grep,
	const char *expgrep) {
	char *v, k[SDB_CDB_MAX_KEY] = { 0 };
	const char *comma = "";
	// local db beacuse is readonly and we dont need to finalize in case of ^C
	Sdb *db = sdb_new(NULL, dbname, 0);
	if (!db) {
		return 1;
	}
	sdb_config(db, options);
	sdb_dump_begin(db);
	if (fmt == MODE_JSON) {
		printf("{");
	}
	while (sdb_dump_dupnext(db, k, &v, NULL)) {
		if (grep && !strstr(k, expgrep) && !strstr(v, expgrep)) {
			free(v);
			continue;
		}
		switch (fmt) {
		case MODE_JSON:
			if (!strcmp(v, "true") || !strcmp(v, "false")) {
				printf("%s\"%s\":%s", comma, k, v);
			} else if (sdb_isnum(v)) {
				printf("%s\"%s\":%" PFMT64x "u", comma, k, sdb_atoi(v));
			} else if (*v == '{' || *v == '[') {
				printf("%s\"%s\":%s", comma, k, v);
			} else {
				printf("%s\"%s\":\"%s\"", comma, k, v);
			}
			comma = ",";
			break;
		case MODE_ZERO:
			printf("%s=%s", k, v);
			break;
		default:
			printf("%s=%s\n", k, v);
			break;
		}
		free(v);
	}
	switch (fmt) {
	case MODE_ZERO:
		fflush(stdout);
		write_null();
		break;
	case MODE_JSON:
		printf("}\n");
		break;
	}
	sdb_free(db);
	return 0;
}

static int sdb_grep(const char *db, int fmt, const char *grep) {
	return sdb_grep_dump(db, fmt, true, grep);
}

static int sdb_dump(const char *db, int fmt) {
	return sdb_grep_dump(db, fmt, false, NULL);
}

static int insertkeys(Sdb *s, const char **args, int nargs, int mode) {
	int must_save = 0;
	if (args && nargs > 0) {
		int i;
		for (i = 0; i < nargs; i++) {
			switch (mode) {
			case '-':
				must_save |= sdb_query(s, args[i]);
				break;
			case '=':
				if (strchr(args[i], '=')) {
					char *v, *kv = (char *)strdup(args[i]);
					v = strchr(kv, '=');
					if (v) {
						*v++ = 0;
						sdb_disk_insert(s, kv, v);
					}
					free(kv);
				}
				break;
			}
		}
	}
	return must_save;
}

static int createdb(const char *f, const char **args, int nargs) {
	s = sdb_new(NULL, f, 0);
	if (!s) {
		eprintf("Cannot create database\n");
		return 1;
	}
	sdb_config(s, options);
	int ret = 0;
	if (args) {
		int i;
		for (i = 0; i < nargs; i++) {
			if (!sdb_text_load(s, args[i])) {
				eprintf("Failed to load text sdb from %s\n", args[i]);
			}
		}
	} else {
		size_t len;
		char *in = slurp(stdin, &len);
		if (!in) {
			return 0;
		}
		if (!sdb_text_load_buf(s, in, len)) {
			eprintf("Failed to read text sdb from stdin\n");
		}
		free(in);
	}
	sdb_sync(s);
	return ret;
}

static int showusage(int o) {
	printf("usage: sdb [-0cdehjJv|-D A B] [-|db] "
	       "[.file]|[-=]|==||[-+][(idx)key[:json|=value] ..]\n");
	if (o == 2) {
		printf("  -0      terminate results with \\x00\n"
		       "  -c      count the number of keys database\n"
		       "  -d      decode base64 from stdin\n"
		       "  -D      diff two databases\n"
		       "  -e      encode stdin as base64\n"
		       "  -h      show this help\n"
		       "  -j      output in json\n"
		       "  -J      enable journaling\n"
		       "  -v      show version information\n");
		return 0;
	}
	return o;
}

static int showversion(void) {
	fflush(stdout);
	return 0;
}

static int jsonIndent(void) {
	size_t len;
	char *out;
	char *in = slurp(stdin, &len);
	if (!in) {
		return 0;
	}
	out = sdb_json_indent(in, "  ");
	if (!out) {
		free(in);
		return 1;
	}
	puts(out);
	free(out);
	free(in);
	return 0;
}

static int base64encode(void) {
	char *out;
	size_t len = 0;
	ut8 *in = (ut8 *)slurp(stdin, &len);
	if (!in) {
		return 0;
	}
	out = sdb_encode(in, (int)len);
	if (!out) {
		free(in);
		return 1;
	}
	puts(out);
	free(out);
	free(in);
	return 0;
}

static int base64decode(void) {
	ut8 *out;
	size_t len, ret = 1;
	char *in = slurp(stdin, &len);
	if (in) {
		int declen;
		out = sdb_decode(in, &declen);
		if (out && declen >= 0) {
			write_(1, out, declen);
			ret = 0;
		}
		free(out);
		free(in);
	}
	return ret;
}

static void dbdiff_cb(const SdbDiff *diff, void *user) {
	char sbuf[512];
	int r = sdb_diff_format(sbuf, sizeof(sbuf), diff);
	if (r < 0) {
		return;
	}
	char *buf = sbuf;
	char *hbuf = NULL;
	if ((size_t)r >= sizeof(sbuf)) {
		hbuf = malloc(r + 1);
		if (!hbuf) {
			return;
		}
		r = sdb_diff_format(hbuf, r + 1, diff);
		if (r < 0) {
			goto beach;
		}
	}
	printf("\x1b[%sm%s\x1b[0m\n", diff->add ? "32" : "31", buf);
beach:
	free(hbuf);
}

static bool dbdiff(const char *a, const char *b) {
	Sdb *A = sdb_new(NULL, a, 0);
	Sdb *B = sdb_new(NULL, b, 0);
	bool equal = sdb_diff(A, B, dbdiff_cb, NULL);
	sdb_free(A);
	sdb_free(B);
	return equal;
}

int showcount(const char *db) {
	ut32 d;
	s = sdb_new(NULL, db, 0);
	if (sdb_stats(s, &d, NULL)) {
		printf("%d\n", d);
	}
	// TODO: show version, timestamp information
	sdb_free(s);
	return 0;
}

int main(int argc, const char **argv) {
	char *line;
	const char *arg, *grep = NULL;
	int i, fmt = MODE_DFLT;
	int db0 = 1, argi = 1;
	bool interactive = false;

	/* terminate flags */
	if (argc < 2) {
		return showusage(1);
	}
	arg = argv[1];

	if (arg[0] == '-') { // && arg[1] && arg[2]==0) {
		switch (arg[1]) {
		case 0:
			/* no-op */
			break;
		case '0':
			fmt = MODE_ZERO;
			db0++;
			argi++;
			if (db0 >= argc) {
				return showusage(1);
			}
			break;
		case 'g':
			db0 += 2;
			if (db0 >= argc) {
				return showusage(1);
			}
			grep = argv[2];
			argi += 2;
			break;
		case 'J':
			options |= SDB_OPTION_JOURNAL;
			db0++;
			argi++;
			if (db0 >= argc) {
				return showusage(1);
			}
			break;
		case 'c': return (argc < 3) ? showusage(1) : showcount(argv[2]);
		case 'v': return showversion();
		case 'h': return showusage(2);
		case 'e': return base64encode();
		case 'd': return base64decode();
		case 'D':
			if (argc == 4) {
				return dbdiff(argv[2], argv[3]) ? 0 : 1;
			}
			return showusage(0);
		case 'j':
			if (argc > 2) {
				return sdb_dump(argv[db0 + 1], MODE_JSON);
			}
			return jsonIndent();
		default:
			eprintf("Invalid flag %s\n", arg);
			break;
		}
	}

	/* sdb - */
	if (argi == 1 && !strcmp(argv[argi], "-")) {
		/* no database */
		argv[argi] = "";
		if (argc == db0 + 1) {
			interactive = true;
			/* if no argument passed */
			argv[argi] = "-";
			argc++;
			argi++;
		}
	}
	/* sdb dbname */
	if (argc - 1 == db0) {
		if (grep) {
			return sdb_grep(argv[db0], fmt, grep);
		}
		return sdb_dump(argv[db0], fmt);
	}
#if HAVE_HEADER_SYS_MMAN_H
	signal(SIGINT, terminate);
	signal(SIGHUP, synchronize);
#endif
	int ret = 0;
	if (interactive || !strcmp(argv[db0 + 1], "-")) {
		if ((s = sdb_new(NULL, argv[db0], 0))) {
			sdb_config(s, options);
			int kvs = db0 + 2;
			if (kvs < argc) {
				save |= insertkeys(s, argv + argi + 2, argc - kvs, '-');
			}
			for (; (line = slurp(stdin, NULL));) {
				save |= sdb_query(s, line);
				if (fmt) {
					fflush(stdout);
					write_null();
				}
				free(line);
			}
		}
	} else if (!strcmp(argv[db0 + 1], "=")) {
		ret = createdb(argv[db0], NULL, 0);
	} else if (!strcmp(argv[db0 + 1], "==")) {
		ret = createdb(argv[db0], argv + db0 + 2, argc - (db0 + 2));
	} else {
		s = sdb_new(NULL, argv[db0], 0);
		if (!s) {
			return 1;
		}
		sdb_config(s, options);
		for (i = db0 + 1; i < argc; i++) {
			save |= sdb_query(s, argv[i]);
			if (fmt) {
				fflush(stdout);
				write_null();
			}
		}
	}
	terminate(ret);
	return ret;
}