//! Host tests for the image builder. Every case here encodes a rule the ESP32-P4 ROM bootloader
//! actually enforces, each of which was learned by flashing a deliberately broken image at the
//! board and reading the error off the serial port (see 04-report/evidence/).

const std = @import("std");
const image = @import("image.zig");

const testing = std.testing;

/// Build a 32-bit little-endian ELF with the given PT_LOAD segments, in memory.
const Load = struct { addr: u32, len: usize };

fn synthElf(gpa: std.mem.Allocator, entry: u32, loads: []const Load) ![]u8 {
    const ehsize = 52;
    const phentsize = 32;
    var out: std.ArrayList(u8) = .empty;
    errdefer out.deinit(gpa);

    const phoff = ehsize;
    var data_off = phoff + phentsize * loads.len;

    try out.appendSlice(gpa, &.{ 0x7F, 'E', 'L', 'F', 1, 1, 1, 0 }); // magic, 32-bit, LE, v1
    try out.appendNTimes(gpa, 0, 8); // padding
    try out.appendSlice(gpa, &std.mem.toBytes(@as(u16, 2))); // ET_EXEC
    try out.appendSlice(gpa, &std.mem.toBytes(@as(u16, 243))); // EM_RISCV
    try out.appendSlice(gpa, &std.mem.toBytes(@as(u32, 1))); // version
    try out.appendSlice(gpa, &std.mem.toBytes(entry));
    try out.appendSlice(gpa, &std.mem.toBytes(@as(u32, phoff)));
    try out.appendSlice(gpa, &std.mem.toBytes(@as(u32, 0))); // shoff
    try out.appendSlice(gpa, &std.mem.toBytes(@as(u32, 0))); // flags
    try out.appendSlice(gpa, &std.mem.toBytes(@as(u16, ehsize)));
    try out.appendSlice(gpa, &std.mem.toBytes(@as(u16, phentsize)));
    try out.appendSlice(gpa, &std.mem.toBytes(@as(u16, @intCast(loads.len))));
    try out.appendSlice(gpa, &std.mem.toBytes(@as(u16, 0))); // shentsize
    try out.appendSlice(gpa, &std.mem.toBytes(@as(u16, 0))); // shnum
    try out.appendSlice(gpa, &std.mem.toBytes(@as(u16, 0))); // shstrndx
    std.debug.assert(out.items.len == ehsize);

    for (loads) |l| {
        try out.appendSlice(gpa, &std.mem.toBytes(@as(u32, 1))); // PT_LOAD
        try out.appendSlice(gpa, &std.mem.toBytes(@as(u32, @intCast(data_off))));
        try out.appendSlice(gpa, &std.mem.toBytes(l.addr)); // vaddr
        try out.appendSlice(gpa, &std.mem.toBytes(l.addr)); // paddr
        try out.appendSlice(gpa, &std.mem.toBytes(@as(u32, @intCast(l.len)))); // filesz
        try out.appendSlice(gpa, &std.mem.toBytes(@as(u32, @intCast(l.len)))); // memsz
        try out.appendSlice(gpa, &std.mem.toBytes(@as(u32, 4))); // flags
        try out.appendSlice(gpa, &std.mem.toBytes(@as(u32, 0x1000))); // align
        data_off += l.len;
    }
    for (loads, 0..) |l, i| {
        try out.appendNTimes(gpa, @intCast('A' + i), l.len);
    }
    return out.toOwnedSlice(gpa);
}

fn segmentHeaders(bytes: []const u8) []const u8 {
    return bytes[24..];
}

test "two mapped segments in one MMU page produce a valid, tiny image" {
    const gpa = testing.allocator;
    // rodata at 0x40000020 (600 B) then text at 0x40000280: 0x20+600+8 == 0x280, congruent
    const elf = try synthElf(gpa, 0x40000280, &.{
        .{ .addr = 0x40000020, .len = 600 },
        .{ .addr = 0x40000280, .len = 760 },
    });
    defer gpa.free(elf);

    var layout = try image.fromElf(gpa, elf, .{});
    defer layout.deinit(gpa);

    try layout.validate(.{});
    try testing.expectEqual(@as(usize, 2), layout.segments.len); // no pad segment needed
    try testing.expectEqual(@as(u32, 0x40000280), layout.entry);
    // 24 B header + 2*(8 B segment header) + payload + checksum pad + 32 B digest
    try testing.expectEqual(@as(usize, 1440), layout.bytes.len);
    try testing.expectEqual(@as(u8, 0xE9), layout.bytes[0]);
    try testing.expectEqual(@as(u8, 2), layout.bytes[1]);
    try testing.expectEqual(@as(u8, 1), layout.bytes[0x17]); // hash_appended
}

test "the previous segment grows when the next mapped segment is not congruent" {
    const gpa = testing.allocator;
    const elf = try synthElf(gpa, 0x40001000, &.{
        .{ .addr = 0x40000020, .len = 100 },
        .{ .addr = 0x40001000, .len = 64 }, // far away: needs padding to line up
    });
    defer gpa.free(elf);

    var layout = try image.fromElf(gpa, elf, .{});
    defer layout.deinit(gpa);

    try layout.validate(.{});
    // no extra segment: the first one carries filler instead of a pad segment paying a header
    try testing.expectEqual(@as(usize, 2), layout.segments.len);
    try testing.expect(layout.segments[0].filler > 0);
    try testing.expectEqual(@as(u32, 100), layout.payload - layout.segments[1].len);
}

test "segment lengths are padded to a multiple of four" {
    const gpa = testing.allocator;
    const elf = try synthElf(gpa, 0x40000040, &.{
        .{ .addr = 0x40000020, .len = 5 }, // 5 bytes: the loader would reject this as-is
        .{ .addr = 0x40000040, .len = 7 },
    });
    defer gpa.free(elf);

    var layout = try image.fromElf(gpa, elf, .{});
    defer layout.deinit(gpa);

    for (layout.segments) |s| try testing.expectEqual(@as(u32, 0), s.len % 4);
    try testing.expect(layout.segments[0].len >= 8); // 5 bytes rounded up, plus congruence filler
}

test "an image with one mapped segment is rejected before it can brick a board" {
    const gpa = testing.allocator;
    const elf = try synthElf(gpa, 0x40000020, &.{.{ .addr = 0x40000020, .len = 64 }});
    defer gpa.free(elf);

    var layout = try image.fromElf(gpa, elf, .{});
    defer layout.deinit(gpa);

    // Not a guess: shipping a one-segment image aborted the boot with
    // "Assert failed in unpack_load_app, bootloader_utility.c:842 (rom_index == 2)".
    try testing.expectError(error.NotTwoMappedSegments, layout.validate(.{}));
}

test "RAM-loaded segments do not count as mapped" {
    const gpa = testing.allocator;
    const elf = try synthElf(gpa, 0x40000020, &.{
        .{ .addr = 0x40000020, .len = 32 },
        .{ .addr = 0x40000060, .len = 32 },
        .{ .addr = 0x4FF00000, .len = 16 }, // L2MEM: loaded, not mapped
    });
    defer gpa.free(elf);

    var layout = try image.fromElf(gpa, elf, .{});
    defer layout.deinit(gpa);

    try layout.validate(.{});
    var mapped: usize = 0;
    var loaded: usize = 0;
    for (layout.segments) |s| switch (s.kind) {
        .mapped => mapped += 1,
        .loaded => loaded += 1,
        .pad => {},
    };
    try testing.expectEqual(@as(usize, 2), mapped);
    try testing.expectEqual(@as(usize, 1), loaded);
}

test "the checksum byte lands on a 16-byte boundary and the digest covers everything before it" {
    const gpa = testing.allocator;
    const elf = try synthElf(gpa, 0x40000280, &.{
        .{ .addr = 0x40000020, .len = 600 },
        .{ .addr = 0x40000280, .len = 760 },
    });
    defer gpa.free(elf);

    var layout = try image.fromElf(gpa, elf, .{});
    defer layout.deinit(gpa);

    const checksum_off = layout.bytes.len - 33;
    try testing.expectEqual(@as(usize, 15), checksum_off % 16);

    var expect: [32]u8 = undefined;
    std.crypto.hash.sha2.Sha256.hash(layout.bytes[0 .. layout.bytes.len - 32], &expect, .{});
    try testing.expectEqualSlices(u8, &expect, layout.bytes[layout.bytes.len - 32 ..]);

    // and the checksum itself is the XOR of every segment byte, seeded 0xEF (filler is zero,
    // so including it or not gives the same answer)
    var xor: u8 = 0xEF;
    var off: usize = 24;
    for (layout.segments) |s| {
        for (layout.bytes[off + 8 .. off + 8 + s.len]) |b| xor ^= b;
        off += 8 + s.len;
    }
    try testing.expectEqual(xor, layout.bytes[checksum_off]);
}

test "the header carries the revision window that keeps a pre-v3 die bootable" {
    const gpa = testing.allocator;
    const elf = try synthElf(gpa, 0x40000040, &.{
        .{ .addr = 0x40000020, .len = 16 },
        .{ .addr = 0x40000040, .len = 16 },
    });
    defer gpa.free(elf);

    var layout = try image.fromElf(gpa, elf, .{ .min_rev_full = 100, .max_rev_full = 199 });
    defer layout.deinit(gpa);

    try testing.expectEqual(@as(u16, 0x0012), std.mem.readInt(u16, layout.bytes[0x0C..0x0E], .little));
    try testing.expectEqual(@as(u16, 100), std.mem.readInt(u16, layout.bytes[0x0F..0x11], .little));
    try testing.expectEqual(@as(u16, 199), std.mem.readInt(u16, layout.bytes[0x11..0x13], .little));
}

test "a RAM segment never has filler copied into memory" {
    const gpa = testing.allocator;
    // On a P4 memory map the RAM window sorts after the flash window, so a RAM segment can never
    // sit between two mapped ones - but if one ever did, growing it would copy filler into L2MEM
    // past the real data. Assert the property directly rather than the mechanism.
    const elf = try synthElf(gpa, 0x40001000, &.{
        .{ .addr = 0x40000020, .len = 64 },
        .{ .addr = 0x40001000, .len = 64 },
        .{ .addr = 0x4FF00000, .len = 40 },
    });
    defer gpa.free(elf);

    var layout = try image.fromElf(gpa, elf, .{});
    defer layout.deinit(gpa);

    try layout.validate(.{});
    for (layout.segments) |s| {
        if (s.kind == .loaded) try testing.expectEqual(@as(u32, 0), s.filler);
    }
}

test "sweep: every rodata length either builds a device-correct image or is refused" {
    // The test the second review round asked for. For a range of rodata lengths and text
    // placements, either the builder refuses, or the produced BYTES satisfy an independent
    // re-derivation of the device's rules - including the MMU invariant that the original solver
    // violated silently. This is the only test that looks at the output rather than at an error
    // name, and it is what would catch a regression in the solver, the linker-script hole, or the
    // checksum layout.
    const gpa = testing.allocator;
    var built: usize = 0;
    var refused: usize = 0;

    var len: usize = 1;
    while (len <= 300) : (len += 7) {
        var text: u32 = 0x40;
        while (text <= 0x400) : (text += 0x20) {
            const elf = try synthElf(gpa, 0x40000000 + text, &.{
                .{ .addr = 0x40000020, .len = len },
                .{ .addr = 0x40000000 + text, .len = 64 },
            });
            defer gpa.free(elf);

            var layout = image.fromElf(gpa, elf, .{}) catch {
                refused += 1;
                continue;
            };
            defer layout.deinit(gpa);
            try layout.validate(.{});
            try checkBytes(layout.bytes, 0x10000);
            built += 1;
        }
    }
    try testing.expect(built > 100);
    try testing.expect(refused > 0); // the impossible layouts really are refused
}

/// Re-derive the device's rules from a finished image, sharing no code with the builder.
fn checkBytes(bytes: []const u8, flash_offset: u32) !void {
    try testing.expectEqual(@as(u8, 0xE9), bytes[0]);
    const count = bytes[1];

    var mapped: usize = 0;
    var deltas: [16]i64 = undefined;
    var pages: [16]u64 = undefined;
    var off: usize = 24;
    var xor: u8 = 0xEF;

    for (0..count) |_| {
        const addr = std.mem.readInt(u32, bytes[off..][0..4], .little);
        const len = std.mem.readInt(u32, bytes[off + 4 ..][0..4], .little);
        try testing.expectEqual(@as(u32, 0), len % 4); // esp_image_format.c:857
        const data = bytes[off + 8 ..][0..len];
        for (data) |b| xor ^= b;

        if (addr >= 0x40000000 and addr < 0x44000000) {
            const flash = @as(i64, flash_offset) + @as(i64, @intCast(off + 8));
            try testing.expectEqual(@mod(@as(i64, addr), 0x10000), @mod(flash, 0x10000));
            deltas[mapped] = flash - @as(i64, addr);
            pages[mapped] = addr / 0x10000;
            mapped += 1;
        }
        off += 8 + len;
    }
    // Exactly two: the SOC_MMU_DI_VADDR_SHARED branch asserts rom_index == 2.
    try testing.expectEqual(@as(usize, 2), mapped); // bootloader_utility.c:842

    // Two mapped segments sharing a vaddr page must share the flash page: one MMU entry each.
    if (pages[0] == pages[1]) try testing.expectEqual(deltas[0], deltas[1]);

    const checksum_off = bytes.len - 33;
    try testing.expectEqual(@as(usize, 15), checksum_off % 16);
    try testing.expectEqual(xor, bytes[checksum_off]);
    for (bytes[off..checksum_off]) |b| try testing.expectEqual(@as(u8, 0), b);

    var digest: [32]u8 = undefined;
    std.crypto.hash.sha2.Sha256.hash(bytes[0 .. bytes.len - 32], &digest, .{});
    try testing.expectEqualSlices(u8, &digest, bytes[bytes.len - 32 ..]);
}

test "parse round-trips what fromElf produced" {
    // Nothing tested image.parse, and the flash and size steps both depend on it.
    const gpa = testing.allocator;
    const elf = try synthElf(gpa, 0x400002c0, &.{
        .{ .addr = 0x40000020, .len = 600 },
        .{ .addr = 0x400002c0, .len = 380 },
    });
    defer gpa.free(elf);

    var built = try image.fromElf(gpa, elf, .{});
    defer built.deinit(gpa);
    var read_back = try image.parse(gpa, built.bytes, .{});
    defer read_back.deinit(gpa);

    try testing.expectEqual(built.entry, read_back.entry);
    try testing.expectEqual(built.segments.len, read_back.segments.len);
    for (built.segments, read_back.segments) |a, b| {
        try testing.expectEqual(a.addr, b.addr);
        try testing.expectEqual(a.len, b.len);
        try testing.expectEqual(a.kind, b.kind);
    }
    try testing.expectEqualSlices(u8, built.bytes, read_back.bytes);
}

test "a gap that would push a mapped segment into the next flash page is refused, not padded" {
    const gpa = testing.allocator;
    // The old solver shifted a whole MMU page forward to satisfy congruence modulo the page. That
    // kept both segments in one *vaddr* page while putting their data in two different *flash*
    // pages, so the bootloader's second MMU write replaced the first and every rodata read
    // resolved to filler zeros. Found by adversarial review, reproduced by -Ddescriptor=full.
    const elf = try synthElf(gpa, 0x40000140, &.{
        .{ .addr = 0x40000020, .len = 268 }, // ends at 0x12C; text at 0x140 needs data@0x140,
        .{ .addr = 0x40000140, .len = 236 }, // but the next data offset is 0x134: gap 12, fine
    });
    defer gpa.free(elf);
    var ok_layout = try image.fromElf(gpa, elf, .{});
    defer ok_layout.deinit(gpa);
    try ok_layout.validate(.{});
    try testing.expect(ok_layout.bytes.len < 1024); // no 64 KiB page jump

    // Now the pathological direction: the second mapped segment sits *before* where the first one
    // already reaches, so no amount of filler can line it up.
    const bad = try synthElf(gpa, 0x40000030, &.{
        .{ .addr = 0x40000020, .len = 512 },
        .{ .addr = 0x40000030, .len = 16 },
    });
    defer gpa.free(bad);
    try testing.expectError(error.MappedSegmentsTooClose, image.fromElf(gpa, bad, .{}));
}

test "validate rejects two mapped segments that would fight over one MMU entry" {
    // Hand-built because the solver now refuses to produce this: both segments are congruent and
    // both live in vaddr page 0x4000, but their data sits in two different flash pages, so the
    // bootloader's second MMU write would replace the first. This is the shape that boots with
    // every constant reading as zero.
    var segs = [_]image.Segment{
        .{ .addr = 0x40000020, .len = 0x10008, .filler = 0, .kind = .mapped },
        .{ .addr = 0x40000030, .len = 16, .filler = 0, .kind = .mapped },
    };
    const layout: image.Layout = .{
        .bytes = &.{},
        .segments = &segs,
        .entry = 0x40000030,
        .payload = 0,
        .filler = 0,
        .overhead = 0,
    };
    // segment 1's data lands at flash 0x10000 + (24 + 8 + 0x10008) + 8 = 0x20030: congruent
    // (0x30 == 0x40000030 % 64K) but one page further along than segment 0's 0x10020.
    try testing.expectError(error.MmuEntryConflict, layout.validate(.{}));
}

test "a partition that is not MMU-page aligned is refused" {
    const gpa = testing.allocator;
    const elf = try synthElf(gpa, 0x40000040, &.{
        .{ .addr = 0x40000020, .len = 16 },
        .{ .addr = 0x40000040, .len = 16 },
    });
    defer gpa.free(elf);
    // The device checks congruence against the absolute flash address, so an image built for
    // 0x11000 needs different padding from one built for 0x10000 - and the anchor cannot be a
    // whole number of pages, which means no layout satisfies the rule.
    try testing.expectError(error.PartitionNotPageAligned, image.fromElf(gpa, elf, .{ .flash_offset = 0x11000 }));
}

test "more than sixteen segments is refused, because the loader stops there" {
    const gpa = testing.allocator;
    var loads: [20]Load = undefined;
    for (&loads, 0..) |*l, i| l.* = .{ .addr = @intCast(0x4FF00000 + i * 0x100), .len = 16 };
    loads[0] = .{ .addr = 0x40000020, .len = 16 };
    loads[1] = .{ .addr = 0x40000040, .len = 16 };
    const elf = try synthElf(gpa, 0x40000040, &loads);
    defer gpa.free(elf);
    try testing.expectError(error.TooManySegments, image.fromElf(gpa, elf, .{}));
}