// User-mode regression harness. The runner extracts the production functions
// unchanged; pool, file-I/O, and CI calls are test doubles. BCrypt is real CNG.
#define WIN32_NO_STATUS
#include <windows.h>
#undef WIN32_NO_STATUS
#include <ntstatus.h>
#include <bcrypt.h>
#include <wintrust.h>
#include <mssip.h>
#include <cstdio>
#include <cstdlib>
#include <vector>
#include <cstring>
#include <string>

#define NT_SUCCESS(s) (((NTSTATUS)(s)) >= 0)
#define ASSERT(x) Check(!!(x), #x)
#define DOMITO_POOL_TAG 0x446f6d69
#define NonPagedPoolNx 0
#define PagedPool 1
#define PAGED_CODE() ((void)0)
#define MINCRYPT_SHA1_LENGTH 20
#define MINCRYPT_SHA256_LENGTH 32
#define DOMITO_CI_LEGACY_CATALOG_POLICY_BITS 0x0002007FUL
#define DOMITO_CI_SOURCE_EMBEDDED 1UL
#define DOMITO_CI_SOURCE_CATALOG 2UL
typedef const UCHAR* PCUCHAR;

static unsigned checks;
static unsigned allocations;
static unsigned ciCalls;
static unsigned hashCalls;
static int failOperation;
static int failHashCall;
static ULONG_PTR bytesRead;
static bool allocatePolicies;
static NTSTATUS embeddedResult = STATUS_SUCCESS;
static NTSTATUS catalogResults[2] = {STATUS_SUCCESS, STATUS_SUCCESS};
static unsigned catalogCalls;
static ULONG expectedCatalogMask = DOMITO_CI_LEGACY_CATALOG_POLICY_BITS;
static BOOLEAN expectedTrustedMarker;
static LONGLONG observedFileSize = 1056;

static void
Check(
    bool ok,
    const char* message
)
{
    ++checks;

    if (!ok)
    {
        std::fprintf(stderr, "FAIL: %s\n", message);
        std::exit(1);
    }
}

static PVOID
ExAllocatePoolWithTag(
    int,
    SIZE_T size,
    ULONG
)
{
    if (failOperation == 4)
    {
        return nullptr;
    }

    void* p = HeapAlloc(GetProcessHeap(), 0, size);

    if (p)
    {
        ++allocations;
    }

    return p;
}

static void
ExFreePoolWithTag(
    PVOID p,
    ULONG
)
{
    Check(p != nullptr && allocations != 0, "allocation ownership");
    --allocations;
    Check(HeapFree(GetProcessHeap(), 0, p) != FALSE, "heap free");
}

static NTSTATUS
TestCreateHash(
    BCRYPT_ALG_HANDLE a,
    BCRYPT_HASH_HANDLE* h,
    PUCHAR o,
    ULONG os,
    PUCHAR s,
    ULONG ss,
    ULONG f
)
{
    if (failOperation == 1)
    {
        return STATUS_UNSUCCESSFUL;
    }

    return BCryptCreateHash(a, h, o, os, s, ss, f);
}

static NTSTATUS
TestHashData(
    BCRYPT_HASH_HANDLE h,
    PUCHAR b,
    ULONG n,
    ULONG f
)
{
    ++hashCalls;

    if (failOperation == 2 && hashCalls == (unsigned)failHashCall)
    {
        return STATUS_UNSUCCESSFUL;
    }

    return BCryptHashData(h, b, n, f);
}

static NTSTATUS
TestFinishHash(
    BCRYPT_HASH_HANDLE h,
    PUCHAR b,
    ULONG n,
    ULONG f
)
{
    if (failOperation == 3)
    {
        return STATUS_UNSUCCESSFUL;
    }

    return BCryptFinishHash(h, b, n, f);
}

// Minimal dependencies for the legacy wrapper; no actual CI code is invoked.
struct MINCRYPT_POLICY_INFO
{
    ULONG Size;
    PVOID ChainInfo;
};

struct KAPC_STATE
{
    ULONG unused;
};

struct IO_STATUS_BLOCK
{
    NTSTATUS Status;
    ULONG_PTR Information;
};

static PVOID PsInitialSystemProcess = nullptr;

static void
CiFreePolicyInfo(
    MINCRYPT_POLICY_INFO* p
)
{
    if (p->ChainInfo)
    {
        ExFreePoolWithTag(p->ChainInfo, DOMITO_POOL_TAG);
    }

    *p = {};
}

static void
PopulatePolicies(
    MINCRYPT_POLICY_INFO* policy,
    MINCRYPT_POLICY_INFO* timestamp
)
{
    Check(!policy->ChainInfo && !timestamp->ChainInfo, "previous CI attempt released before reuse");

    if (allocatePolicies)
    {
        policy->Size = timestamp->Size = sizeof(*policy);
        policy->ChainInfo = ExAllocatePoolWithTag(PagedPool, 1, DOMITO_POOL_TAG);
        timestamp->ChainInfo = ExAllocatePoolWithTag(PagedPool, 1, DOMITO_POOL_TAG);
    }
}

static void
KeStackAttachProcess(
    PVOID,
    KAPC_STATE*
)
{
}

static void
KeUnstackDetachProcess(
    KAPC_STATE*
)
{
}

typedef PVOID PFILE_OBJECT;
static PVOID fileObjectType;
static PVOID* IoFileObjectType = &fileObjectType;
#define KernelMode 0

static NTSTATUS
ObReferenceObjectByHandle(
    HANDLE,
    ULONG,
    PVOID,
    int,
    PVOID* object,
    PVOID
)
{
    *object = &fileObjectType;
    return STATUS_SUCCESS;
}

static void
ObDereferenceObject(
    PVOID
)
{
}

static NTSTATUS
FsRtlGetFileSize(
    PFILE_OBJECT,
    LARGE_INTEGER* size
)
{
    size->QuadPart = observedFileSize;
    return STATUS_SUCCESS;
}

static NTSTATUS
DomitoReadFileAt(
    HANDLE,
    PVOID buffer,
    ULONG size,
    const LARGE_INTEGER* offset
)
{
    Check(offset->QuadPart == 1024, "explicit certificate read offset");
    std::memset(buffer, 0, (size_t)min(bytesRead, (ULONG_PTR)size));
    return bytesRead == size ? STATUS_SUCCESS : STATUS_END_OF_FILE;
}

static NTSTATUS
CiCheckSignedFile(
    PVOID,
    UINT32,
    ALG_ID,
    PVOID,
    UINT32,
    MINCRYPT_POLICY_INFO* policy,
    LARGE_INTEGER* time,
    MINCRYPT_POLICY_INFO* timestamp
)
{
    ++ciCalls;
    Check(time->QuadPart == 0, "fresh signing time for embedded validation");
    PopulatePolicies(policy, timestamp);
    time->QuadPart = 123;
    return embeddedResult;
}

static NTSTATUS
CiVerifyHashInCatalog(
    PVOID,
    UINT32,
    ALG_ID,
    BOOLEAN reload,
    UINT32 trusted,
    UINT32 mask,
    MINCRYPT_POLICY_INFO* policy,
    PVOID,
    LARGE_INTEGER* time,
    MINCRYPT_POLICY_INFO* timestamp
)
{
    ++ciCalls;
    Check(mask == expectedCatalogMask && trusted == expectedTrustedMarker, "explicit catalog policy forwarded");
    Check(reload == (catalogCalls != 0), "catalog reload sequence");
    Check(time->QuadPart == 0, "fresh signing time for catalog attempt");
    PopulatePolicies(policy, timestamp);
    time->QuadPart = 456;
    Check(catalogCalls < 2, "bounded catalog retries");
    return catalogResults[catalogCalls++];
}

#define BCryptCreateHash TestCreateHash
#define BCryptHashData TestHashData
#define BCryptFinishHash TestFinishHash
#include "digest-under-test.inc"
#undef BCryptCreateHash
#undef BCryptHashData
#undef BCryptFinishHash

static std::vector<UCHAR>
Image(
    bool pe64,
    ULONG nt = 128,
    bool certificate = false
)
{
    std::vector<UCHAR> b((size_t)nt + 896 + (certificate ? 32 : 0), 0);
    IMAGE_DOS_HEADER dos = {};
    dos.e_magic = IMAGE_DOS_SIGNATURE;
    dos.e_lfanew = (LONG)nt;
    std::memcpy(b.data(), &dos, sizeof(dos));
    ULONG signature = IMAGE_NT_SIGNATURE;
    std::memcpy(b.data() + nt, &signature, sizeof(signature));
    IMAGE_FILE_HEADER f = {};
    f.NumberOfSections = 1;
    f.Machine = pe64 ? IMAGE_FILE_MACHINE_AMD64 : IMAGE_FILE_MACHINE_I386;
    f.SizeOfOptionalHeader = pe64 ? sizeof(IMAGE_OPTIONAL_HEADER64) : sizeof(IMAGE_OPTIONAL_HEADER32);
    std::memcpy(b.data() + nt + 4, &f, sizeof(f));
    IMAGE_OPTIONAL_HEADER64 h64 = {};
    IMAGE_OPTIONAL_HEADER32 h32 = {};
    h64.Magic = IMAGE_NT_OPTIONAL_HDR64_MAGIC;
    h32.Magic = IMAGE_NT_OPTIONAL_HDR32_MAGIC;
    h64.NumberOfRvaAndSizes = h32.NumberOfRvaAndSizes = 16;
    h64.SectionAlignment = h32.SectionAlignment = 4096;
    h64.FileAlignment = h32.FileAlignment = 512;
    h64.SizeOfImage = h32.SizeOfImage = 8192;
    h64.Subsystem = h32.Subsystem = IMAGE_SUBSYSTEM_WINDOWS_CUI;
    h64.SizeOfHeaders = h32.SizeOfHeaders = nt + 384;

    if (certificate)
    {
        h64.DataDirectory[4] = h32.DataDirectory[4] = {nt + 896, 32};
        ULONG length = 32;
        std::memcpy(b.data() + nt + 896, &length, sizeof(length));
        USHORT revision = 0x0200;
        USHORT type = WIN_CERT_TYPE_PKCS_SIGNED_DATA;
        std::memcpy(b.data() + nt + 900, &revision, 2);
        std::memcpy(b.data() + nt + 902, &type, 2);
    }

    std::memcpy(b.data() + nt + 24, pe64 ? (void*)&h64 : (void*)&h32, f.SizeOfOptionalHeader);
    IMAGE_SECTION_HEADER section = {};
    section.VirtualAddress = 4096;
    section.Misc.VirtualSize = 512;
    section.Characteristics = IMAGE_SCN_CNT_CODE | IMAGE_SCN_MEM_READ;
    section.PointerToRawData = nt + 384;
    section.SizeOfRawData = 512;
    std::memcpy(b.data() + nt + 24 + f.SizeOfOptionalHeader, &section, sizeof(section));

    for (ULONG i = nt + 384; i < nt + 896; ++i)
    {
        b[i] = (UCHAR)(i * 31);
    }

    return b;
}

static void
Set32(
    std::vector<UCHAR>& b,
    size_t at,
    ULONG value
)
{
    std::memcpy(b.data() + at, &value, sizeof(value));
}

static void
Set16(
    std::vector<UCHAR>& b,
    size_t at,
    USHORT value
)
{
    std::memcpy(b.data() + at, &value, sizeof(value));
}

static std::vector<UCHAR>
Digest(
    PUCHAR b,
    ULONG size,
    NTSTATUS expected = STATUS_SUCCESS
)
{
    UINT32 algorithm = 0xdead;
    ULONG length = 0xdead;
    ULONG certSize = 0xdead;
    PVOID digest = (void*)1;
    LPWIN_CERTIFICATE cert = (LPWIN_CERTIFICATE)1;
    hashCalls = 0;
    NTSTATUS s = DomitoCalculatePortableExecutableDigest(b, size, &algorithm, &length, &digest, &cert, &certSize);
    Check(s == expected, "digest status");
    std::vector<UCHAR> result;

    if (NT_SUCCESS(s))
    {
        Check(algorithm == CALG_SHA1 && length == 20 && digest != nullptr, "SHA1 output");
        Check((cert == nullptr) == (certSize == 0), "certificate outputs agree");

        if (cert)
        {
            Check((PUCHAR)cert >= b && (PUCHAR)cert + certSize <= b + size, "borrowed certificate span");
        }

        result.assign((PUCHAR)digest, (PUCHAR)digest + length);
        DomitoFreePortableExecutableDigest(digest);
    }
    else
    {
        Check(!algorithm && !length && !digest && !cert && !certSize, "failure outputs cleared");
    }

    Check(allocations == 0, "no leaked pool allocations");
    return result;
}

static std::vector<UCHAR>
Reference(
    const std::vector<UCHAR>& b,
    ULONG checksum,
    ULONG securitySlot,
    ULONG end
)
{
    // Explicit fixture ranges, independent of the production parser.
    std::vector<UCHAR> bytes(b.begin(), b.begin() + checksum);
    bytes.insert(bytes.end(), b.begin() + checksum + 4, b.begin() + securitySlot);

    if (securitySlot != end)
    {
        bytes.insert(bytes.end(), b.begin() + securitySlot + 8, b.begin() + end);
    }

    std::vector<UCHAR> digest(20);
    BCRYPT_ALG_HANDLE algorithm = nullptr;
    Check(
        NT_SUCCESS(BCryptOpenAlgorithmProvider(&algorithm, BCRYPT_SHA1_ALGORITHM, nullptr, 0)), "reference algorithm"
    );
    Check(
        NT_SUCCESS(BCryptHash(algorithm, nullptr, 0, bytes.data(), (ULONG)bytes.size(), digest.data(), 20)),
        "reference digest"
    );
    BCryptCloseAlgorithmProvider(algorithm, 0);
    return digest;
}

static void
GuardedInvalid(
    const std::vector<UCHAR>& b,
    NTSTATUS expected
)
{
    SYSTEM_INFO si;
    GetSystemInfo(&si);
    const SIZE_T span = ((b.size() + si.dwPageSize - 1) / si.dwPageSize) * si.dwPageSize;
    UCHAR* region = (UCHAR*)VirtualAlloc(nullptr, span + si.dwPageSize, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE);
    Check(region != nullptr, "guard allocation");
    DWORD old;
    Check(VirtualProtect(region + span, si.dwPageSize, PAGE_NOACCESS, &old) != FALSE, "guard protection");
    UCHAR* input = region + span - b.size();
    std::memcpy(input, b.data(), b.size());
    Digest(input, (ULONG)b.size(), expected);
    ULONG strictSize = 0;
    ULONG securitySize = 0;
    PVOID strictDigest = nullptr;
    LPWIN_CERTIFICATE certificate = nullptr;
    Check(
        !NT_SUCCESS(DomitoCalculatePortableExecutableDigestEx(
            input, (ULONG)b.size(), CALG_SHA256, &strictSize, &strictDigest, &certificate, &securitySize
        )),
        "strict guard-page malformed input rejected"
    );
    Check(!strictSize && !strictDigest && !certificate && !securitySize, "strict malformed outputs cleared");
    VirtualFree(region, 0, MEM_RELEASE);
}

static std::vector<UCHAR>
ExplicitDigest(
    const std::vector<UCHAR>& b,
    UINT32 calg,
    NTSTATUS expected = STATUS_SUCCESS
)
{
    ULONG length = 0xdead;
    ULONG certSize = 0xdead;
    PVOID digest = (PVOID)1;
    LPWIN_CERTIFICATE cert = (LPWIN_CERTIFICATE)1;
    NTSTATUS status =
        DomitoCalculatePortableExecutableDigestEx(b.data(), (ULONG)b.size(), calg, &length, &digest, &cert, &certSize);
    Check(status == expected, "explicit digest status");
    std::vector<UCHAR> result;

    if (NT_SUCCESS(status))
    {
        Check(
            length == (calg == CALG_SHA1     ? 20u
                       : calg == CALG_SHA256 ? 32u
                       : calg == CALG_SHA384 ? 48u
                                             : 64u),
            "explicit algorithm digest length"
        );
        result.assign((UCHAR*)digest, (UCHAR*)digest + length);
        DomitoFreePortableExecutableDigest(digest);
    }
    else
    {
        Check(!length && !digest && !cert && !certSize, "explicit failure outputs cleared");
    }

    Check(allocations == 0, "explicit digest allocation cleanup");
    return result;
}

static std::vector<UCHAR>
SipDigest(
    const std::vector<UCHAR>& b,
    UINT32 calg
)
{
    WCHAR tempPath[MAX_PATH];
    WCHAR tempName[MAX_PATH];
    Check(GetTempPathW(MAX_PATH, tempPath) != 0, "temp path");
    Check(GetTempFileNameW(tempPath, L"dci", 0, tempName) != 0, "temp name");
    HANDLE file = CreateFileW(
        tempName,
        GENERIC_READ | GENERIC_WRITE,
        FILE_SHARE_READ,
        nullptr,
        CREATE_ALWAYS,
        FILE_ATTRIBUTE_TEMPORARY | FILE_FLAG_DELETE_ON_CLOSE,
        nullptr
    );
    Check(file != INVALID_HANDLE_VALUE, "SIP fixture file");
    DWORD written = 0;
    Check(WriteFile(file, b.data(), (DWORD)b.size(), &written, nullptr) && written == b.size(), "write SIP fixture");
    GUID subject;
    Check(CryptSIPRetrieveSubjectGuid(tempName, file, &subject) != FALSE, "PE SIP subject");
    SIP_SUBJECTINFO info = {};
    info.cbSize = sizeof(info);
    info.pgSubjectType = &subject;
    info.hFile = file;
    info.pwsFileName = tempName;
    info.dwEncodingType = X509_ASN_ENCODING | PKCS_7_ASN_ENCODING;
    info.DigestAlgorithm.pszObjId = const_cast<char*>(
        calg == CALG_SHA1     ? szOID_OIWSEC_sha1
        : calg == CALG_SHA256 ? szOID_NIST_sha256
        : calg == CALG_SHA384 ? szOID_NIST_sha384
                              : szOID_NIST_sha512
    );
    DWORD size = 0;
    Check(CryptSIPCreateIndirectData(&info, &size, nullptr) != FALSE, "SIP size query");
    std::vector<UCHAR> buffer(size);
    auto indirect = (SIP_INDIRECT_DATA*)buffer.data();
    Check(CryptSIPCreateIndirectData(&info, &size, indirect) != FALSE, "Windows SIP digest");
    std::vector<UCHAR> result(indirect->Digest.pbData, indirect->Digest.pbData + indirect->Digest.cbData);
    CloseHandle(file);
    return result;
}

static std::vector<UCHAR>
Der(
    UCHAR tag,
    const std::vector<UCHAR>& value
)
{
    std::vector<UCHAR> out = {tag};

    if (value.size() < 128)
    {
        out.push_back((UCHAR)value.size());
    }
    else if (value.size() < 256)
    {
        out.push_back(0x81);
        out.push_back((UCHAR)value.size());
    }
    else
    {
        out.push_back(0x82);
        out.push_back((UCHAR)(value.size() >> 8));
        out.push_back((UCHAR)value.size());
    }

    out.insert(out.end(), value.begin(), value.end());
    return out;
}

static std::vector<UCHAR>
Join(
    std::initializer_list<std::vector<UCHAR>> values
)
{
    std::vector<UCHAR> out;

    for (auto& value : values)
    {
        out.insert(out.end(), value.begin(), value.end());
    }

    return out;
}

static std::vector<UCHAR>
SignedContent(
    UINT32 calg,
    bool nullParameters = true
)
{
    std::vector<UCHAR> oid = calg == CALG_SHA1 ? std::vector<UCHAR>{0x2B, 0x0E, 0x03, 0x02, 0x1A}
                                               : std::vector<UCHAR>{
                                                     0x60,
                                                     0x86,
                                                     0x48,
                                                     0x01,
                                                     0x65,
                                                     0x03,
                                                     0x04,
                                                     0x02,
                                                     (UCHAR)(calg == CALG_SHA256   ? 1
                                                             : calg == CALG_SHA384 ? 2
                                                                                   : 3)
                                                 };
    auto algorithm = Der(0x30, Join({Der(0x06, oid), nullParameters ? Der(0x05, {}) : std::vector<UCHAR>{}}));
    auto digest =
        Der(0x30,
            Join(
                {algorithm,
                 Der(0x04,
                     std::vector<UCHAR>(
                         calg == CALG_SHA1     ? 20
                         : calg == CALG_SHA256 ? 32
                         : calg == CALG_SHA384 ? 48
                                               : 64,
                         0xAA
                     ))}
            ));
    // A complete SHA-1 DigestInfo decoy in the ignored data attribute must not
    // override the selected SpcIndirectDataContent digest AlgorithmIdentifier.
    auto data =
        Der(0x30, Der(0x04, {0x30, 0x21, 0x30, 0x09, 0x06, 0x05, 0x2B, 0x0E, 0x03, 0x02, 0x1A, 0x05, 0, 0x04, 0x14}));
    auto indirect = Der(0x30, Join({data, digest}));
    auto content =
        Der(0x30, Join({Der(0x06, {0x2B, 0x06, 0x01, 0x04, 0x01, 0x82, 0x37, 0x02, 0x01, 0x04}), Der(0xA0, indirect)}));
    auto signedData = Der(0x30, Join({Der(0x02, {1}), Der(0x31, algorithm), content, Der(0x31, {})}));
    return Der(0x30, Join({Der(0x06, {0x2A, 0x86, 0x48, 0x86, 0xF7, 0x0D, 0x01, 0x07, 0x02}), Der(0xA0, signedData)}));
}

static void
ExtendedTests()
{
    for (bool pe64 : {false, true})
    {
        for (bool cert : {false, true})
        {
            for (UINT32 calg : std::initializer_list<UINT32>{CALG_SHA1, CALG_SHA256, CALG_SHA384, CALG_SHA512})
            {
                auto b = Image(pe64, 128, cert);
                Check(ExplicitDigest(b, calg) == SipDigest(b, calg), "Windows SIP differential fixture");

                if (!cert)
                {
                    b.insert(b.end(), 48, 0x42); // Non-section overlay.
                    Check(ExplicitDigest(b, calg) == SipDigest(b, calg), "Windows SIP overlay differential");
                }
            }
        }
    }

    for (bool pe64 : {false, true})
    {
        for (ULONG count : {0UL, 4UL})
        {
            auto absent = Image(pe64);
            Set32(absent, pe64 ? 260 : 244, count);
            ExplicitDigest(absent, CALG_SHA256, STATUS_NOT_SUPPORTED);
        }
    }

    auto b = Image(true);
    // Two contiguous sections whose headers are in reverse raw-offset order.
    Set16(b, 134, 2);
    IMAGE_SECTION_HEADER first = {};
    IMAGE_SECTION_HEADER second = {};
    first.PointerToRawData = 768;
    first.SizeOfRawData = 256;
    second.PointerToRawData = 512;
    second.SizeOfRawData = 256;
    std::memcpy(b.data() + 392, &first, sizeof(first));
    std::memcpy(b.data() + 432, &second, sizeof(second));
    Check(ExplicitDigest(b, CALG_SHA256) == SipDigest(b, CALG_SHA256), "section header order differential");
    Set32(b, 392 + 20, 800);
    ExplicitDigest(b, CALG_SHA256, STATUS_INVALID_IMAGE_FORMAT); // extends beyond EOF
    Set32(b, 392 + 20, 640);
    ExplicitDigest(b, CALG_SHA256, STATUS_NOT_SUPPORTED); // overlap
    b = Image(true);
    Set32(b, 392 + 20, 528);
    Set32(b, 392 + 16, 496);
    ExplicitDigest(b, CALG_SHA256, STATUS_NOT_SUPPORTED); // gap
    b = Image(true, 128, true);
    b.insert(b.end(), 8, 0);
    ExplicitDigest(b, CALG_SHA256, STATUS_NOT_SUPPORTED); // nonterminal certificate
    b = Image(true);
    Set32(b, 392 + 20, 0xFFFFFFFF);
    ExplicitDigest(b, CALG_SHA256, STATUS_INVALID_IMAGE_FORMAT);
    ExplicitDigest(Image(true), 0x1234, STATUS_NOT_SUPPORTED);
    b = Image(true, 128, true);
    b.resize(1088);
    Set32(b, 300, 64);
    Set32(b, 1056, 32);
    Set16(b, 1060, 0x0200);
    Set16(b, 1062, 2);
    Check(ExplicitDigest(b, CALG_SHA256) == SipDigest(b, CALG_SHA256), "multiple certificate entries excluded");
    Set32(b, 1056, 0);
    ExplicitDigest(b, CALG_SHA256, STATUS_INVALID_IMAGE_FORMAT);
    b = Image(true, 128, true);
    Set32(b, 1024, 31); // Last entry's alignment padding is outside dwLength.
    Check(ExplicitDigest(b, CALG_SHA256) == SipDigest(b, CALG_SHA256), "certificate alignment padding");

    for (int failure : {1, 2, 3, 4})
    {
        failOperation = failure;
        failHashCall = 1;
        hashCalls = 0;
        ExplicitDigest(Image(true), CALG_SHA256, failure == 4 ? STATUS_INSUFFICIENT_RESOURCES : STATUS_UNSUCCESSFUL);
    }

    failOperation = 0;

    for (UINT32 calg : std::initializer_list<UINT32>{CALG_SHA1, CALG_SHA256, CALG_SHA384, CALG_SHA512})
    {
        auto content = SignedContent(calg);
        auto signedImage = Image(true);
        ULONG certLength = 8 + (ULONG)content.size();
        ULONG tableSize = (certLength + 7) & ~7u;
        signedImage.resize(1024 + tableSize);
        Set32(signedImage, 296, 1024);
        Set32(signedImage, 300, tableSize);
        Set32(signedImage, 1024, certLength);
        Set16(signedImage, 1028, 0x0200);
        Set16(signedImage, 1030, 2);
        std::memcpy(signedImage.data() + 1032, content.data(), content.size());
        IMAGE_DATA_DIRECTORY directories[5] = {};
        directories[4] = {1024, tableSize};
        UINT32 actual = 0;
        Check(
            DomitoGetPortableExecutableDigestKindEx(
                signedImage.data(), (ULONG)signedImage.size(), directories, &actual
            ) == STATUS_SUCCESS &&
                actual == calg,
            "bounded SPC digest algorithm, decoy ignored"
        );

        for (ULONG cut = 0; cut < content.size(); ++cut)
        {
            Set32(signedImage, 1024, cut + 8);
            actual = 0xDEAD;
            Check(
                !NT_SUCCESS(DomitoGetPortableExecutableDigestKindEx(
                    signedImage.data(), (ULONG)signedImage.size(), directories, &actual
                )) &&
                    actual == 0,
                "truncated DER rejected"
            );
        }

        Set32(signedImage, 1024, tableSize);
        Check(
            DomitoGetPortableExecutableDigestKindEx(
                signedImage.data(), (ULONG)signedImage.size(), directories, &actual
            ) == STATUS_SUCCESS &&
                actual == calg,
            "bounded zero padding inside dwLength"
        );

        if (tableSize > certLength)
        {
            signedImage.back() = 0x7F;
            Check(
                DomitoGetPortableExecutableDigestKindEx(
                    signedImage.data(), (ULONG)signedImage.size(), directories, &actual
                ) == STATUS_INVALID_IMAGE_FORMAT &&
                    actual == 0,
                "nonzero DER trailing bytes rejected"
            );
            signedImage.back() = 0;
        }

        Set32(signedImage, 1024, certLength);
        signedImage[1033] = 0x80; // Indefinite length forbidden in this DER subset.
        Check(
            DomitoGetPortableExecutableDigestKindEx(
                signedImage.data(), (ULONG)signedImage.size(), directories, &actual
            ) == STATUS_INVALID_IMAGE_FORMAT,
            "indefinite DER rejected"
        );
    }

    for (bool nullParameters : {false, true})
    {
        auto content = SignedContent(CALG_SHA256, nullParameters);
        auto pe = Image(true);
        const ULONG length = 8 + (ULONG)content.size();
        const ULONG tableSize = (length + 7) & ~7u;
        pe.resize(1024 + tableSize);
        Set32(pe, 1024, length);
        Set16(pe, 1028, 0x0200);
        Set16(pe, 1030, 2);
        std::memcpy(pe.data() + 1032, content.data(), content.size());
        IMAGE_DATA_DIRECTORY dirs[5] = {};
        dirs[4] = {1024, tableSize};
        UINT32 algorithm = 0;
        Check(
            DomitoGetPortableExecutableDigestKindEx(pe.data(), (ULONG)pe.size(), dirs, &algorithm) == STATUS_SUCCESS &&
                algorithm == CALG_SHA256,
            "AlgorithmIdentifier NULL optional"
        );
        // Change both copies of the SHA-256 OID to an unsupported family member.
        const UCHAR oid[] = {0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x01};

        for (size_t i = 1032; i + sizeof(oid) <= pe.size(); ++i)
        {
            if (!std::memcmp(pe.data() + i, oid, sizeof(oid)))
            {
                pe[i + sizeof(oid) - 1] = 0x7F;
            }
        }

        Check(
            DomitoGetPortableExecutableDigestKindEx(pe.data(), (ULONG)pe.size(), dirs, &algorithm) ==
                    STATUS_NOT_SUPPORTED &&
                !algorithm,
            "unsupported SPC algorithm does not default to SHA1"
        );
    }

    IMAGE_DATA_DIRECTORY dirs[5] = {};
    UINT32 actual = 123;
    Check(
        DomitoGetPortableExecutableDigestKindEx(b.data(), (ULONG)b.size(), dirs, &actual) == STATUS_NOT_FOUND &&
            actual == 0,
        "absent DER does not silently select SHA1"
    );
}

static void
SystemImageTests()
{
    WCHAR system[MAX_PATH];
    Check(GetSystemDirectoryW(system, MAX_PATH) != 0, "system directory");

    for (const WCHAR* name : {L"ci.dll", L"ntdll.dll", L"kernel32.dll", L"notepad.exe"})
    {
        std::wstring path = std::wstring(system) + L"\\" + name;
        HANDLE file = CreateFileW(
            path.c_str(), GENERIC_READ, FILE_SHARE_READ | FILE_SHARE_DELETE, nullptr, OPEN_EXISTING, 0, nullptr
        );
        Check(file != INVALID_HANDLE_VALUE, "system PE corpus file");
        LARGE_INTEGER size;
        Check(
            GetFileSizeEx(file, &size) && size.QuadPart > 0 && size.QuadPart < 32 * 1024 * 1024,
            "bounded system PE size"
        );
        std::vector<UCHAR> b((size_t)size.QuadPart);
        DWORD read;
        Check(ReadFile(file, b.data(), (DWORD)b.size(), &read, nullptr) && read == b.size(), "read system PE fixture");
        CloseHandle(file);

        for (UINT32 calg : std::initializer_list<UINT32>{CALG_SHA1, CALG_SHA256, CALG_SHA384, CALG_SHA512})
        {
            Check(ExplicitDigest(b, calg) == SipDigest(b, calg), "real system PE Windows SIP differential");
        }

        IMAGE_DOS_HEADER dos;
        std::memcpy(&dos, b.data(), sizeof(dos));
        USHORT magic;
        const ULONG optional = (ULONG)dos.e_lfanew + 24;
        std::memcpy(&magic, b.data() + optional, 2);
        const ULONG directory = magic == IMAGE_NT_OPTIONAL_HDR64_MAGIC
                                    ? FIELD_OFFSET(IMAGE_OPTIONAL_HEADER64, DataDirectory)
                                    : FIELD_OFFSET(IMAGE_OPTIONAL_HEADER32, DataDirectory);
        IMAGE_DATA_DIRECTORY dirs[5];
        std::memcpy(dirs, b.data() + optional + directory, sizeof(dirs));

        if (dirs[4].Size)
        {
            UINT32 actual;
            Check(
                NT_SUCCESS(DomitoGetPortableExecutableDigestKindEx(b.data(), (ULONG)b.size(), dirs, &actual)),
                "real embedded SPC digest algorithm"
            );
            // Independent Windows ASN.1 decoder: extract signed content with
            // CryptQueryObject/CryptMsg, then decode SPC_INDIRECT_DATA_OBJID.
            ULONG certLength;
            std::memcpy(&certLength, b.data() + dirs[4].VirtualAddress, 4);
            CRYPT_DATA_BLOB blob = {certLength - 8, b.data() + dirs[4].VirtualAddress + 8};
            HCRYPTMSG message = nullptr;
            Check(
                CryptQueryObject(
                    CERT_QUERY_OBJECT_BLOB,
                    &blob,
                    CERT_QUERY_CONTENT_FLAG_PKCS7_SIGNED,
                    CERT_QUERY_FORMAT_FLAG_BINARY,
                    0,
                    nullptr,
                    nullptr,
                    nullptr,
                    nullptr,
                    &message,
                    nullptr
                ) != FALSE,
                "Windows decode PKCS7"
            );
            DWORD contentSize = 0;
            Check(
                CryptMsgGetParam(message, CMSG_CONTENT_PARAM, 0, nullptr, &contentSize) != FALSE,
                "Windows SPC content size"
            );
            std::vector<UCHAR> content(contentSize);
            Check(
                CryptMsgGetParam(message, CMSG_CONTENT_PARAM, 0, content.data(), &contentSize) != FALSE,
                "Windows SPC content"
            );
            SIP_INDIRECT_DATA* indirect = nullptr;
            DWORD decodedSize = 0;
            Check(
                CryptDecodeObjectEx(
                    X509_ASN_ENCODING | PKCS_7_ASN_ENCODING,
                    SPC_INDIRECT_DATA_OBJID,
                    content.data(),
                    contentSize,
                    CRYPT_DECODE_ALLOC_FLAG,
                    nullptr,
                    &indirect,
                    &decodedSize
                ) != FALSE,
                "Windows decode indirect data"
            );
            const char* oid = indirect->DigestAlgorithm.pszObjId;
            const UINT32 expected = !std::strcmp(oid, szOID_OIWSEC_sha1)   ? CALG_SHA1
                                    : !std::strcmp(oid, szOID_NIST_sha256) ? CALG_SHA256
                                    : !std::strcmp(oid, szOID_NIST_sha384) ? CALG_SHA384
                                                                           : CALG_SHA512;
            Check(actual == expected, "SPC digest algorithm matches Windows ASN1 decoder");
            auto digest = ExplicitDigest(b, actual);
            Check(
                digest.size() == indirect->Digest.cbData &&
                    !std::memcmp(digest.data(), indirect->Digest.pbData, digest.size()),
                "computed digest matches real embedded SPC digest"
            );
            LocalFree(indirect);
            CryptMsgClose(message);
        }
    }
}

int
main()
{
    SystemImageTests();
    ExtendedTests();

    for (bool pe64 : {false, true})
    {
        for (bool certificate : {false, true})
        {
            auto b = Image(pe64, 128, certificate);
            auto result = Digest(b.data(), (ULONG)b.size());
            Check(
                result == Reference(b, 216, pe64 ? 296 : 280, 1024), "fixture digest equals independent CNG reference"
            );
            Set32(b, 216, 0x12345678);
            Check(result == Digest(b.data(), (ULONG)b.size()), "checksum excluded");

            if (certificate)
            {
                b.back() ^= 0x55;
                Check(result == Digest(b.data(), (ULONG)b.size()), "certificate contents excluded");
            }
        }
    }

    auto large = Image(true, 4 * 1024 * 1024);
    Check(
        Digest(large.data(), (ULONG)large.size()) == Reference(large, 4194392, 4194472, (ULONG)large.size()),
        "large DOS stub"
    );
    auto b = Image(true);

    for (ULONG size = 0; size < 392; ++size)
    {
        std::vector<UCHAR> prefix(b.begin(), b.begin() + size);

        if (size == 0)
        {
            Digest(b.data(), 0, STATUS_INVALID_IMAGE_NOT_MZ);
        }
        else
        {
            GuardedInvalid(prefix, size < 64 ? STATUS_INVALID_IMAGE_NOT_MZ : STATUS_INVALID_IMAGE_FORMAT);
        }
    }

    auto bad = b;
    Set32(bad, 60, 0xffffffff);
    GuardedInvalid(bad, STATUS_INVALID_IMAGE_FORMAT);
    bad = b;
    Set32(bad, 60, 0x7fffffff);
    GuardedInvalid(bad, STATUS_INVALID_IMAGE_FORMAT);
    bad = b;
    Set16(bad, 148, 1);
    GuardedInvalid(bad, STATUS_INVALID_IMAGE_FORMAT);
    bad = b;
    Set16(bad, 152, 0x999);
    GuardedInvalid(bad, STATUS_INVALID_IMAGE_FORMAT);
    bad = b;
    Set32(bad, 260, 17);
    GuardedInvalid(bad, STATUS_INVALID_IMAGE_FORMAT);

    for (ULONG offset : {200UL, 2048UL, 0xffffffffUL})
    {
        bad = b;
        Set32(bad, 296, offset);
        Set32(bad, 300, 32);
        GuardedInvalid(bad, STATUS_INVALID_IMAGE_FORMAT);
    }

    bad = b;
    Set32(bad, 300, 32);
    GuardedInvalid(bad, STATUS_INVALID_IMAGE_FORMAT);
    bad = b;
    Set32(bad, 296, 800);
    GuardedInvalid(bad, STATUS_INVALID_IMAGE_FORMAT);
    bad = b;
    Set32(bad, 296, 800);
    Set32(bad, 300, 0xffffffff);
    GuardedInvalid(bad, STATUS_INVALID_IMAGE_FORMAT);
    bad = Image(true, 128, true);
    Set32(bad, 1024, 33);
    GuardedInvalid(bad, STATUS_INVALID_IMAGE_FORMAT);
    bad = Image(true, 128, true);
    Set32(bad, 1024, 7);
    GuardedInvalid(bad, STATUS_INVALID_IMAGE_FORMAT);
    bad = Image(true);
    Set32(bad, 260, 0); // No directory entry: don't exclude its physical bytes.
    Check(Digest(bad.data(), (ULONG)bad.size()) == Reference(bad, 216, 1024, 1024), "absent security directory");

    for (int failure : {1, 2, 3, 4})
    {
        failOperation = failure;

        for (int call = 1; call <= (failure == 2 ? 3 : 1); ++call)
        {
            failHashCall = call;
            Digest(b.data(), (ULONG)b.size(), failure == 4 ? STATUS_INSUFFICIENT_RESOURCES : STATUS_UNSUCCESSFUL);
        }
    }

    failOperation = 0;
    Digest(nullptr, 1024, STATUS_INVALID_PARAMETER);
    IMAGE_DATA_DIRECTORY directory = {1024, 32};
    MINCRYPT_POLICY_INFO policy = {};
    MINCRYPT_POLICY_INFO timestamp = {};
    LARGE_INTEGER signing = {};
    UCHAR hash[20] = {};
    bytesRead = 16;
    ciCalls = 0;
    Check(
        DomitoValidateFileLegacyMode(nullptr, hash, 20, CALG_SHA1, &directory, &policy, &signing, &timestamp) ==
            STATUS_END_OF_FILE,
        "short certificate read rejected"
    );
    Check(ciCalls == 0 && allocations == 0, "short read does not reach CI or leak");
    bytesRead = 32;
    Check(
        NT_SUCCESS(
            DomitoValidateFileLegacyMode(nullptr, hash, 20, CALG_SHA1, &directory, &policy, &signing, &timestamp)
        ),
        "complete certificate read accepted"
    );
    Check(ciCalls == 1 && allocations == 0, "complete read reaches CI once");
    // Failure paths must release previous attempts before embedded/catalog/reload reuse.
    allocatePolicies = true;
    embeddedResult = STATUS_INVALID_IMAGE_HASH;
    catalogResults[0] = STATUS_INVALID_IMAGE_HASH;
    catalogResults[1] = STATUS_SUCCESS;
    catalogCalls = 0;
    expectedCatalogMask = 0x8000;
    expectedTrustedMarker = TRUE;
    ULONG source = 0xDEAD;
    Check(
        NT_SUCCESS(DomitoValidateFileLegacyModeEx(
            (HANDLE)1,
            hash,
            20,
            CALG_SHA1,
            &directory,
            &policy,
            &signing,
            &timestamp,
            expectedCatalogMask,
            TRUE,
            &source
        )),
        "explicit catalog fallback"
    );
    Check(
        source == DOMITO_CI_SOURCE_CATALOG && catalogCalls == 2 && allocations == 2,
        "final policy ownership and catalog source"
    );
    CiFreePolicyInfo(&policy);
    CiFreePolicyInfo(&timestamp);
    Check(allocations == 0, "final results freed once");
    source = 123;
    Check(
        DomitoValidateFileLegacyModeEx(
            (HANDLE)1, hash, 20, CALG_SHA256, &directory, &policy, &signing, &timestamp, 0, FALSE, &source
        ) == STATUS_INVALID_PARAMETER &&
            !source,
        "algorithm and length mismatch rejected"
    );
    Check(
        DomitoValidateFileLegacyModeEx(
            (HANDLE)1, hash, 20, CALG_SHA384, &directory, &policy, &signing, &timestamp, 0, FALSE, &source
        ) == STATUS_NOT_SUPPORTED,
        "legacy catalog algorithm capability explicit"
    );
    observedFileSize = 1040;
    Check(
        DomitoValidateFileLegacyModeEx(
            (HANDLE)1, hash, 20, CALG_SHA1, &directory, &policy, &signing, &timestamp, 0, FALSE, &source
        ) == STATUS_INVALID_IMAGE_FORMAT,
        "certificate file range checked before allocation"
    );
    Check(allocations == 0, "invalid file range never allocated");
    observedFileSize = 1056;
    expectedCatalogMask = 0x8000;
    expectedTrustedMarker = TRUE;
    catalogCalls = 0;
    catalogResults[1] = STATUS_INVALID_IMAGE_HASH;
    Check(
        DomitoValidateFileLegacyModeEx(
            (HANDLE)1,
            hash,
            20,
            CALG_SHA1,
            &directory,
            &policy,
            &signing,
            &timestamp,
            expectedCatalogMask,
            TRUE,
            &source
        ) == STATUS_INVALID_IMAGE_HASH,
        "final catalog failure retained"
    );
    Check(!source && allocations == 2, "final failed-call policies remain caller-owned");
    CiFreePolicyInfo(&policy);
    CiFreePolicyInfo(&timestamp);
    Check(allocations == 0, "failed final policies freed once");
    std::printf("PASS: %u digest/legacy checks (Windows SIP, Crypt32, guard pages, failures, ownership).\n", checks);
}