/*
 * Copyright (c) 2019 Tomohiro Kusumi <tkusumi@netbsd.org>
 * Copyright (c) 2019 The DragonFly Project
 * All rights reserved.
 *
 * This code is derived from software contributed to The DragonFly Project
 * by Matthew Dillon <dillon@dragonflybsd.org>
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in
 *    the documentation and/or other materials provided with the
 *    distribution.
 * 3. Neither the name of The DragonFly Project nor the names of its
 *    contributors may be used to endorse or promote products derived
 *    from this software without specific, prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * ``AS IS'' AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
 * FOR A PARTICULAR PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE
 * COPYRIGHT HOLDERS OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
 * INCIDENTAL, SPECIAL, EXEMPLARY OR CONSEQUENTIAL DAMAGES (INCLUDING,
 * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED
 * AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
 * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT
 * OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
 * SUCH DAMAGE.
 */

// # gcc -Wall -g -I../../sys -I../hammer2 -I../../crypto/libressl/include ../../sys/vfs/hammer2/xxhash/xxhash.c ../../sys/libkern/icrc32.c ../hammer2/subs.c ../hammer2/ondisk.c ./reconstruct.c -o reconstruct

#include <sys/types.h>
#include <sys/stat.h>
#include <unistd.h>
#include <fcntl.h>
#include <stdio.h>
#include <stdlib.h>
#include <stdbool.h>
#include <string.h>
#include <assert.h>

#include <openssl/sha.h>

#include <vfs/hammer2/hammer2_disk.h>
#include <vfs/hammer2/hammer2_xxhash.h>

#include "hammer2_subs.h"

static int modify_volume_header(hammer2_volume_data_t *,
    const hammer2_blockref_t *);
static int modify_blockref(const hammer2_volume_data_t *, int,
    hammer2_blockref_t *, hammer2_blockref_t *, int);
static int modify_check(int, hammer2_blockref_t *, const hammer2_blockref_t *,
    hammer2_media_data_t *, size_t, int);

static bool ForceOpt = false;

static int
reconstruct_volume_header(void)
{
	bool failed = false;
	int i;

	for (i = 0; i < HAMMER2_NUM_VOLHDRS; ++i) {
		hammer2_volume_data_t voldata;
		hammer2_blockref_t broot;
		hammer2_off_t off;
		ssize_t ret;

		memset(&broot, 0, sizeof(broot));
		broot.data_off = (i * HAMMER2_ZONE_BYTES64) | HAMMER2_PBUFRADIX;
		off = broot.data_off & ~HAMMER2_OFF_MASK_RADIX;
		if (lseek(hammer2_get_root_volume_fd(),
		    off - hammer2_get_root_volume_offset(), SEEK_SET) == -1) {
			perror("lseek");
			return -1;
		}

		ret = read(hammer2_get_root_volume_fd(), &voldata,
		    HAMMER2_VOLUME_BYTES);
		if (ret == HAMMER2_VOLUME_BYTES) {
			fprintf(stdout, "zone.%d %016jx\n",
			    i, (uintmax_t)broot.data_off);
			if (modify_volume_header(&voldata, &broot) == -1)
				failed = true;
		} else if (ret == -1) {
			perror("read");
			return -1;
		} else {
			fprintf(stderr, "Failed to read volume header\n");
			return -1;
		}
	}

	return failed ? -1 : 0;
}

static int
reconstruct_blockref(uint8_t type)
{
	bool failed = false;
	int i;

	for (i = 0; i < HAMMER2_NUM_VOLHDRS; ++i) {
		hammer2_volume_data_t voldata;
		hammer2_blockref_t broot;
		hammer2_off_t off;
		ssize_t ret;

		memset(&broot, 0, sizeof(broot));
		broot.type = type;
		broot.data_off = (i * HAMMER2_ZONE_BYTES64) | HAMMER2_PBUFRADIX;
		off = broot.data_off & ~HAMMER2_OFF_MASK_RADIX;
		if (lseek(hammer2_get_root_volume_fd(),
		    off - hammer2_get_root_volume_offset(), SEEK_SET) == -1) {
			perror("lseek");
			return -1;
		}

		ret = read(hammer2_get_root_volume_fd(), &voldata,
		    HAMMER2_VOLUME_BYTES);
		if (ret == HAMMER2_VOLUME_BYTES) {
			fprintf(stdout, "zone.%d %016jx\n",
			    i, (uintmax_t)broot.data_off);
			if (modify_blockref(&voldata, -1, &broot, NULL, -1) ==
			    -1)
				failed = true;
		} else if (ret == -1) {
			perror("read");
			return -1;
		} else {
			fprintf(stderr, "Failed to read volume header\n");
			return -1;
		}
	}

	return failed ? -1 : 0;
}

static int
modify_volume_header(hammer2_volume_data_t *voldata,
    const hammer2_blockref_t *bref)
{
	hammer2_crc32_t crc0, crc1;
	const char *s = NULL;
	bool found = false;

	if ((voldata->magic != HAMMER2_VOLUME_ID_HBO) &&
	    (voldata->magic != HAMMER2_VOLUME_ID_ABO)) {
		fprintf(stderr, "Bad magic %jX\n", voldata->magic);
		return -1;
	}

	if (voldata->magic == HAMMER2_VOLUME_ID_ABO)
		fprintf(stderr, "Reverse endian\n");

	/* Need to test HAMMER2_VOL_ICRC_SECT1 first. */
	crc0 = voldata->icrc_sects[HAMMER2_VOL_ICRC_SECT1];
	crc1 = hammer2_icrc32((char*)voldata + HAMMER2_VOLUME_ICRC1_OFF,
	    HAMMER2_VOLUME_ICRC1_SIZE);
	if (crc0 != crc1) {
		if (ForceOpt)
			voldata->icrc_sects[HAMMER2_VOL_ICRC_SECT1] = crc1;
		found = true;
		s = "HAMMER2_VOL_ICRC_SECT1";
		printf("%s%016jx %s\n", ForceOpt ? "Modified " : "",
		    (uintmax_t)bref->data_off, s);
	}

	crc0 = voldata->icrc_sects[HAMMER2_VOL_ICRC_SECT0];
	crc1 = hammer2_icrc32((char*)voldata + HAMMER2_VOLUME_ICRC0_OFF,
	    HAMMER2_VOLUME_ICRC0_SIZE);
	if (crc0 != crc1) {
		if (ForceOpt)
			voldata->icrc_sects[HAMMER2_VOL_ICRC_SECT0] = crc1;
		found = true;
		s = "HAMMER2_VOL_ICRC_SECT0";
		printf("%s%016jx %s\n", ForceOpt ? "Modified " : "",
		    (uintmax_t)bref->data_off, s);
	}

	crc0 = voldata->icrc_volheader;
	crc1 = hammer2_icrc32((char*)voldata + HAMMER2_VOLUME_ICRCVH_OFF,
	    HAMMER2_VOLUME_ICRCVH_SIZE);
	if (crc0 != crc1) {
		if (ForceOpt)
			voldata->icrc_volheader = crc1;
		found = true;
		s = "volume header CRC";
		printf("%s%016jx %s\n", ForceOpt ? "Modified " : "",
		    (uintmax_t)bref->data_off, s);
	}

	if (found && ForceOpt) {
		ssize_t ret;
		int fd = hammer2_get_root_volume_fd();
		hammer2_off_t off = bref->data_off & ~HAMMER2_OFF_MASK_RADIX;
		if (lseek(fd, off - hammer2_get_root_volume_offset(), SEEK_SET)
		    == -1) {
			perror("lseek");
			return -1;
		}
		ret = write(fd, voldata, HAMMER2_PBUFSIZE);
		if (ret == -1) {
			perror("write");
			return -1;
		} else if (ret != (ssize_t)HAMMER2_PBUFSIZE) {
			fprintf(stderr, "Failed to write volume header\n");
			return -1;
		}
		if (fsync(fd) == -1) {
			perror("fsync");
			return -1;
		}
	}

	return 0;
}

static int
read_media(const hammer2_blockref_t *bref, hammer2_media_data_t *media,
    size_t *media_bytes)
{
	hammer2_off_t io_off, io_base;
	size_t bytes, io_bytes, boff;
	ssize_t ret;
	int fd;

	bytes = (bref->data_off & HAMMER2_OFF_MASK_RADIX);
	if (bytes)
		bytes = (size_t)1 << bytes;
	if (media_bytes)
		*media_bytes = bytes;

	if (!bytes)
		return 0;

	io_off = bref->data_off & ~HAMMER2_OFF_MASK_RADIX;
	io_base = io_off & ~(hammer2_off_t)(HAMMER2_LBUFSIZE - 1);
	boff = io_off - io_base;

	io_bytes = HAMMER2_LBUFSIZE;
	while (io_bytes + boff < bytes)
		io_bytes <<= 1;

	if (io_bytes > sizeof(*media)) {
		fprintf(stderr, "Bad I/O bytes\n");
		return -1;
	}
	fd = hammer2_get_volume_fd(io_off);
	if (lseek(fd, io_base - hammer2_get_volume_offset(io_base), SEEK_SET)
	    == -1) {
		perror("lseek");
		return -1;
	}
	ret = read(fd, media, io_bytes);
	if (ret == -1) {
		perror("read");
		return -1;
	} else if (ret != (ssize_t)io_bytes) {
		fprintf(stderr, "Failed to read media\n");
		return -1;
	}
	if (boff)
		memmove(media, (char *)media + boff, bytes);

	return 0;
}

static int
write_media(const hammer2_blockref_t *bref, const hammer2_media_data_t *media,
    size_t media_bytes)
{
	hammer2_off_t io_off, io_base;
	char buf[HAMMER2_PBUFSIZE];
	size_t bytes, io_bytes, boff;
	ssize_t ret;
	int fd;

	bytes = (bref->data_off & HAMMER2_OFF_MASK_RADIX);
	if (bytes)
		bytes = (size_t)1 << bytes;
	assert(bytes != 0);
	assert(bytes == media_bytes);

	io_off = bref->data_off & ~HAMMER2_OFF_MASK_RADIX;
	io_base = io_off & ~(hammer2_off_t)(HAMMER2_LBUFSIZE - 1);
	boff = io_off - io_base;

	io_bytes = HAMMER2_LBUFSIZE;
	while (io_bytes + boff < bytes)
		io_bytes <<= 1;

	if (io_bytes > sizeof(buf)) {
		fprintf(stderr, "Bad I/O bytes\n");
		return -1;
	}
	fd = hammer2_get_volume_fd(io_off);
	if (lseek(fd, io_base - hammer2_get_volume_offset(io_base), SEEK_SET)
	    == -1) {
		perror("lseek");
		return -1;
	}
	if (read(fd, buf, io_bytes) != (ssize_t)io_bytes) {
		perror("read");
		return -1;
	}

	memcpy(buf + boff, media, media_bytes);
	if (lseek(fd, io_base - hammer2_get_volume_offset(io_base), SEEK_SET)
	    == -1) {
		perror("lseek");
		return -1;
	}
	ret = write(fd, buf, io_bytes);
	if (ret == -1) {
		perror("write");
		return -1;
	} else if (ret != (ssize_t)io_bytes) {
		fprintf(stderr, "Failed to write media\n");
		return -1;
	}
	if (fsync(fd) == -1) {
		perror("fsync");
		return -1;
	}

	return 0;
}

static int
modify_blockref(const hammer2_volume_data_t *voldata, int bi,
    hammer2_blockref_t *bref, hammer2_blockref_t *prev_bref, int depth)
{
	hammer2_media_data_t media;
	hammer2_blockref_t *bscan;
	int i, bcount;
	size_t bytes;

	if (read_media(bref, &media, &bytes) == -1)
		return -1;

	if (!bytes)
		return 0;

	switch (bref->type) {
	case HAMMER2_BREF_TYPE_INODE:
		if (!(media.ipdata.meta.op_flags & HAMMER2_OPFLAG_DIRECTDATA)) {
			bscan = &media.ipdata.u.blockset.blockref[0];
			bcount = HAMMER2_SET_COUNT;
		} else {
			bscan = NULL;
			bcount = 0;
		}
		break;
	case HAMMER2_BREF_TYPE_INDIRECT:
		bscan = &media.npdata[0];
		bcount = bytes / sizeof(hammer2_blockref_t);
		break;
	case HAMMER2_BREF_TYPE_FREEMAP_NODE:
		bscan = &media.npdata[0];
		bcount = bytes / sizeof(hammer2_blockref_t);
		break;
	case HAMMER2_BREF_TYPE_VOLUME:
		bscan = &media.voldata.sroot_blockset.blockref[0];
		bcount = HAMMER2_SET_COUNT;
		break;
	case HAMMER2_BREF_TYPE_FREEMAP:
		bscan = &media.voldata.freemap_blockset.blockref[0];
		bcount = HAMMER2_SET_COUNT;
		break;
	default:
		bscan = NULL;
		bcount = 0;
		break;
	}

	for (i = 0; i < bcount; ++i)
		if (bscan[i].type != HAMMER2_BREF_TYPE_EMPTY)
			if (modify_blockref(voldata, i, &bscan[i], bref,
			    depth + 1) == -1)
				return -1;

	if (ForceOpt)
		if (read_media(bref, &media, &bytes) == -1)
			return -1;
	if (modify_check(bi, prev_bref, bref, &media, bytes, depth) == -1)
		return -1;

	return 0;
}

static int
modify_check(int bi, hammer2_blockref_t *prev_bref,
    const hammer2_blockref_t *bref, hammer2_media_data_t *media,
    size_t media_bytes, int depth)
{
	hammer2_media_data_t bscan_media;
	hammer2_blockref_t *bscan;
	bool found = false;
	size_t bytes;
	uint32_t cv;
	uint64_t cv64;

	//SHA256_CTX hash_ctx;
	union {
		uint8_t digest[SHA256_DIGEST_LENGTH];
		uint64_t digest64[SHA256_DIGEST_LENGTH/8];
	} u;


	if (!prev_bref)
		return 0;
	if (read_media(prev_bref, &bscan_media, &bytes) == -1)
		return -1;
	assert(bytes);

	switch (prev_bref->type) {
	case HAMMER2_BREF_TYPE_INODE:
		if (!(bscan_media.ipdata.meta.op_flags &
		    HAMMER2_OPFLAG_DIRECTDATA))
			bscan = &bscan_media.ipdata.u.blockset.blockref[bi];
		else
			bscan = NULL;
		break;
	case HAMMER2_BREF_TYPE_INDIRECT:
		bscan = &bscan_media.npdata[bi];
		break;
	case HAMMER2_BREF_TYPE_VOLUME:
		bscan = &bscan_media.voldata.sroot_blockset.blockref[bi];
		break;
	case HAMMER2_BREF_TYPE_FREEMAP:
		bscan = &bscan_media.voldata.freemap_blockset.blockref[bi];
		break;
	case HAMMER2_BREF_TYPE_FREEMAP_NODE:
		bscan = &bscan_media.npdata[bi];
		break;
	default:
		assert(0);
		break;
	}

	if (memcmp(bref, bscan, sizeof(*bref))) {
		fprintf(stderr, "Blockref contents mismatch\n");
		return -1;
	}

	switch (HAMMER2_DEC_CHECK(bscan->methods)) {
	case HAMMER2_CHECK_ISCSI32:
		cv = hammer2_icrc32(media, media_bytes);
		if (bscan->check.iscsi32.value != cv) {
			if (ForceOpt)
				bscan->check.iscsi32.value = cv;
			found = true;
		}
		break;
	case HAMMER2_CHECK_XXHASH64:
		cv64 = XXH64(media, media_bytes, XXH_HAMMER2_SEED);
		if (bscan->check.xxhash64.value != cv64) {
			if (ForceOpt)
				bscan->check.xxhash64.value = cv64;
			found = true;
		}
		break;
	case HAMMER2_CHECK_SHA192:
#if 0
		SHA256_Init(&hash_ctx);
		SHA256_Update(&hash_ctx, &media, bytes);
		SHA256_Final(u.digest, &hash_ctx);
#endif
		u.digest64[2] ^= u.digest64[3];
		if (memcmp(u.digest, bscan->check.sha192.data,
		    sizeof(bscan->check.sha192.data))) {
			if (ForceOpt)
				memcpy(&bscan->check.sha192.data, u.digest,
				    sizeof(bscan->check.sha192.data));
			found = true;
		}
		fprintf(stderr, "HAMMER2_CHECK_SHA192 unsupported\n");
		assert(0);
		break;
	case HAMMER2_CHECK_FREEMAP:
		cv = hammer2_icrc32(media, media_bytes);
		if (bscan->check.freemap.icrc32 != cv) {
			if (ForceOpt)
				bscan->check.freemap.icrc32 = cv;
			found = true;
		}
		break;
	}

	if (found) {
		if (ForceOpt) {
			if (write_media(prev_bref, &bscan_media, bytes) == -1)
				return -1;
		}
		/* If !ForceOpt, only first bad blockref is printed. */
		printf("%s%2d %-8s blockref[%-3d] %016jx %02x %s\n",
		    ForceOpt ? "Modified " : "",
		    depth, hammer2_breftype_to_str(prev_bref->type), bi,
		    (uintmax_t)bscan->data_off, bscan->methods,
		    hammer2_breftype_to_str(bscan->type));
	}

	return 0;
}

int
main(int argc, char **argv)
{
	int ch;
	const char *binpath = argv[0];
	const char *devpath;

	while ((ch = getopt(argc, argv, "f")) != -1) {
		switch(ch) {
		case 'f':
			ForceOpt = true;
			break;
		default:
			break;
		}
	}
	argc -= optind;
	argv += optind;

	if (argc < 1) {
		fprintf(stderr, "%s [-f] special\n", binpath);
		exit(1);
	}
	devpath = argv[0];
	hammer2_init_volumes(devpath, 0);

	printf("freemap\n");
	if (reconstruct_blockref(HAMMER2_BREF_TYPE_FREEMAP) == -1)
		exit(1);
	printf("volume\n");
	if (reconstruct_blockref(HAMMER2_BREF_TYPE_VOLUME) == -1)
		exit(1);

	printf("volume header\n");
	if (reconstruct_volume_header() == -1)
		exit(1);

	hammer2_cleanup_volumes();

	return 0;
}