#define WIN32_NO_STATUS
#include <windows.h>
#undef WIN32_NO_STATUS
#include <ntstatus.h>
#include <bcrypt.h>
#include <wincrypt.h>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <string>
#include <vector>
#include <initializer_list>
#define NT_SUCCESS(s) (((NTSTATUS)(s)) >= 0)
#define PASSIVE_LEVEL 0
#define PAGED_CODE() ((void)0)
#undef NTDDI_VERSION
#define NTDDI_VERSION 0x0a000008
#undef NTDDI_WIN10
#define NTDDI_WIN10 0x0a000000
#define DOMITO_POOL_TAG 0x446f6d69
#define MAXUSHORT 0xffff
#define PagedPool 1
#define KernelMode 0
#define ProcessImageFileName 27
#define OBJ_KERNEL_HANDLE 0x200
#define OBJ_CASE_INSENSITIVE 0x40
#define FILE_NON_DIRECTORY_FILE 0x40
#define FILE_SYNCHRONOUS_IO_NONALERT 0x20
typedef void* PEPROCESS;
typedef void* PFILE_OBJECT;
typedef const UCHAR* PCUCHAR;

typedef struct _UNICODE_STRING
{
    USHORT Length;
    USHORT MaximumLength;
    PWSTR Buffer;
} UNICODE_STRING, *PUNICODE_STRING;

typedef struct _ANSI_STRING
{
    USHORT Length;
    USHORT MaximumLength;
    PCHAR Buffer;
} ANSI_STRING, *PANSI_STRING;

struct OBJECT_ATTRIBUTES
{
    ULONG Attributes;
    PUNICODE_STRING ObjectName;
};

struct CLIENT_ID
{
    HANDLE UniqueProcess;
    HANDLE UniqueThread;
};

struct IO_STATUS_BLOCK
{
    NTSTATUS Status;
    ULONG_PTR Information;
};

#define InitializeObjectAttributes(a, n, f, r, s) ((a)->Attributes = (f), (a)->ObjectName = (n))
#define RTL_CONSTANT_STRING(s) {sizeof(s) - sizeof((s)[0]), sizeof(s), const_cast<char*>(s)}
typedef PVOID(NTAPI* PMINCRYPT_ALLOCATE_ROUTINE)(SIZE_T);
typedef PVOID NTAPI MINCRYPT_ALLOCATE_ROUTINE(SIZE_T);
#include "consumer-api.inc"

static unsigned checks;
static unsigned live;
static unsigned refs;
static unsigned derefs;
static unsigned closed;
static unsigned validated;
static unsigned opened;
static unsigned queried;
static int failure;
static int cryptoFailure;
static NTSTATUS moduleStatus = STATUS_SUCCESS;
static NTSTATUS exportStatus = STATUS_SUCCESS;
static NTSTATUS validateStatus = STATUS_SUCCESS;
static int fileToken;
static int processToken;
static void* fileTypeValue = &fileToken;
static void** IoFileObjectType = &fileTypeValue;

static void
Check(
    bool value,
    const char* msg
)
{
    ++checks;

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

static void*
Allocate(
    SIZE_T size
)
{
    if (failure == 1)
    {
        return nullptr;
    }

    void* p = std::malloc(size);
    Check(p != nullptr, "allocation");
    ++live;
    std::memset(p, 0xa5, size);
    return p;
}

static void
Release(
    void* p
)
{
    Check(p && live, "single owner release");
    --live;
    std::free(p);
}

static struct
{
    PVOID (*Allocate)(SIZE_T);
    void (*Free)(PVOID);
} G_Memory = {Allocate, Release};

static PVOID
ExAllocatePoolWithTag(
    int,
    SIZE_T size,
    ULONG
)
{
    return Allocate(size);
}

static void
ExFreePoolWithTag(
    void* p,
    ULONG
)
{
    Release(p);
}

static void
ExFreePool(
    void* p
)
{
    Release(p);
}

static NTSTATUS
FakePublisher(
    MINCERT_BLOB*,
    PMINCRYPT_ALLOCATE_ROUTINE allocator,
    PUNICODE_STRING name
)
{
    name->Buffer = static_cast<PWCHAR>(allocator(12));

    if (!name->Buffer)
    {
        return STATUS_INSUFFICIENT_RESOURCES;
    }

    std::memcpy(name->Buffer, L"hello", 12);
    name->Length = 10;
    name->MaximumLength = 12;

    if (failure == 8)
    {
        return STATUS_UNSUCCESSFUL;
    }

    if (failure == 9)
    {
        name->Length = 13;
    }

    return STATUS_SUCCESS;
}

static NTSTATUS
DomitoFindModuleBaseAddress(
    PANSI_STRING name,
    PVOID* module
)
{
    Check(name && name->Length, "module name");

    if (NT_SUCCESS(moduleStatus))
    {
        *module = &fileToken;
    }

    return moduleStatus;
}

static NTSTATUS
DomitoFindExportedFunctionAddress(
    PVOID module,
    PANSI_STRING name,
    PVOID* routine
)
{
    Check(module == &fileToken && name->Length, "export resolution");

    if (NT_SUCCESS(exportStatus))
    {
        *routine = reinterpret_cast<PVOID>(FakePublisher);
    }

    return exportStatus;
}

void
DomitoFreeCiValidationResult(
    DOMITO_CI_VALIDATION_RESULT* result
)
{
    if (result->PolicyInfo.ChainInfo)
    {
        Release(result->PolicyInfo.ChainInfo);
    }

    *result = {};
}

NTSTATUS
DomitoValidateFileObjectEx(
    PFILE_OBJECT object,
    UINT32 flags,
    UINT8 level,
    DOMITO_CI_SIGNATURE_SELECTION selection,
    DOMITO_CI_VALIDATION_RESULT* result
)
{
    ++validated;
    Check(
        object == &fileToken && flags == 0x20 && level == 7 && selection == DomitoCiSignatureSelectionNative,
        "CI arguments unchanged"
    );
    Check(!result->PolicyInfo.ChainInfo, "result reset before native call");
    result->PolicyInfo.ChainInfo = static_cast<MINCRYPT_CHAIN_INFO*>(Allocate(1));
    return validateStatus;
}

static NTSTATUS
ObReferenceObjectByHandle(
    HANDLE handle,
    ACCESS_MASK access,
    void* type,
    int mode,
    void** object,
    void*
)
{
    Check(
        handle == &fileToken && access == FILE_READ_DATA && type == *IoFileObjectType && mode == KernelMode,
        "trusted handle reference"
    );

    if (failure == 2)
    {
        return STATUS_ACCESS_DENIED;
    }

    ++refs;
    *object = &fileToken;
    return STATUS_SUCCESS;
}

static void
ObDereferenceObject(
    void* object
)
{
    Check(object == &fileToken, "file dereference");
    ++derefs;
}

static NTSTATUS
ZwOpenFile(
    HANDLE* handle,
    ACCESS_MASK access,
    OBJECT_ATTRIBUTES* a,
    IO_STATUS_BLOCK*,
    ULONG share,
    ULONG options
)
{
    ++opened;
    Check(
        access == (FILE_READ_DATA | SYNCHRONIZE) && a->Attributes == (OBJ_KERNEL_HANDLE | OBJ_CASE_INSENSITIVE) &&
            share == FILE_SHARE_READ && options == (FILE_NON_DIRECTORY_FILE | FILE_SYNCHRONOUS_IO_NONALERT),
        "read-only stable path open"
    );

    if (failure == 3)
    {
        return STATUS_OBJECT_NAME_NOT_FOUND;
    }

    *handle = &fileToken;
    return STATUS_SUCCESS;
}

static void
ZwClose(
    HANDLE handle
)
{
    Check(handle == &fileToken || handle == &processToken, "close owned handle");
    ++closed;
}

static void
FillName(
    PUNICODE_STRING name
)
{
    name->Length = 8;
    name->MaximumLength = 10;
    name->Buffer = reinterpret_cast<PWCHAR>(name + 1);
    std::memcpy(name->Buffer, L"test", 10);
}

static NTSTATUS
SeLocateProcessImageName(
    PEPROCESS process,
    PUNICODE_STRING* image
)
{
    Check(process == &processToken, "referenced process passed unchanged");

    if (failure == 4)
    {
        return STATUS_NOT_FOUND;
    }

    // OS allocations are independent of configurable Domito allocator failures.
    *image = static_cast<PUNICODE_STRING>(std::malloc(sizeof(UNICODE_STRING) + 10));
    ++live;
    FillName(*image);

    if (failure == 5)
    {
        (*image)->Length = 9;
    }

    return STATUS_SUCCESS;
}

static NTSTATUS
ZwOpenProcess(
    HANDLE* process,
    ACCESS_MASK access,
    OBJECT_ATTRIBUTES* a,
    CLIENT_ID* cid
)
{
    Check(
        access == 0x1000 && a->Attributes == OBJ_KERNEL_HANDLE && cid->UniqueProcess == ULongToHandle(123),
        "PID least privilege kernel open"
    );

    if (failure == 6)
    {
        return STATUS_ACCESS_DENIED;
    }

    *process = &processToken;
    return STATUS_SUCCESS;
}

static NTSTATUS
QueryProcess(
    HANDLE process,
    int info,
    PVOID buffer,
    ULONG size,
    PULONG returned
)
{
    ++queried;
    Check(process == &processToken && info == ProcessImageFileName, "PID query bound to one handle");
    *returned = sizeof(UNICODE_STRING) + 10;

    if (!buffer)
    {
        return failure == 10 ? STATUS_SUCCESS : STATUS_INFO_LENGTH_MISMATCH;
    }

    Check(size == *returned, "sized name allocation");
    FillName(static_cast<PUNICODE_STRING>(buffer));

    if (failure == 7)
    {
        return STATUS_UNSUCCESSFUL;
    }

    if (failure == 11)
    {
        static_cast<PUNICODE_STRING>(buffer)->Buffer = reinterpret_cast<PWCHAR>(1);
    }

    return STATUS_SUCCESS;
}

static struct
{
    NTSTATUS (*ZwQueryInformationProcess)(
        HANDLE,
        int,
        PVOID,
        ULONG,
        PULONG
    );
} G_Common = {QueryProcess};

static NTSTATUS
OpenAlgorithm(
    BCRYPT_ALG_HANDLE* a,
    LPCWSTR n,
    LPCWSTR p,
    ULONG f
)
{
    return cryptoFailure == 1 ? STATUS_UNSUCCESSFUL : BCryptOpenAlgorithmProvider(a, n, p, f);
}

static NTSTATUS
CreateHash(
    BCRYPT_ALG_HANDLE a,
    BCRYPT_HASH_HANDLE* h,
    PUCHAR o,
    ULONG os,
    PUCHAR s,
    ULONG ss,
    ULONG f
)
{
    return cryptoFailure == 2 ? STATUS_UNSUCCESSFUL : BCryptCreateHash(a, h, o, os, s, ss, f);
}

static NTSTATUS
HashData(
    BCRYPT_HASH_HANDLE h,
    PUCHAR p,
    ULONG n,
    ULONG f
)
{
    return cryptoFailure == 3 ? STATUS_UNSUCCESSFUL : BCryptHashData(h, p, n, f);
}

static NTSTATUS
FinishHash(
    BCRYPT_HASH_HANDLE h,
    PUCHAR p,
    ULONG n,
    ULONG f
)
{
    if (cryptoFailure == 4)
    {
        std::memset(p, 0xff, n);
        return STATUS_UNSUCCESSFUL;
    }

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

#define BCryptOpenAlgorithmProvider OpenAlgorithm
#define BCryptCreateHash CreateHash
#define BCryptHashData HashData
#define BCryptFinishHash FinishHash
#ifdef DOMITO_TEST_OTHER_ARCH
#undef _M_AMD64
#undef _AMD64_
#endif
#include "consumer-under-test.inc"
#undef BCryptOpenAlgorithmProvider
#undef BCryptCreateHash
#undef BCryptHashData
#undef BCryptFinishHash

using Bytes = std::vector<UCHAR>;

static Bytes
Der(
    UCHAR tag,
    const Bytes& contents
)
{
    Bytes b = {tag};

    if (contents.size() < 128)
    {
        b.push_back(static_cast<UCHAR>(contents.size()));
    }
    else
    {
        Check(contents.size() < 256, "fixture DER bound");
        b.push_back(0x81);
        b.push_back(static_cast<UCHAR>(contents.size()));
    }

    b.insert(b.end(), contents.begin(), contents.end());
    return b;
}

static Bytes
Join(
    std::initializer_list<Bytes> values
)
{
    Bytes b;

    for (const auto& v : values)
    {
        b.insert(b.end(), v.begin(), v.end());
    }

    return b;
}

static Bytes
Fixture(
    bool version
)
{
    Bytes sig = Der(0x30, Join({Der(6, {0x2a, 0x86, 0x48, 0x86, 0xf7, 0x0d, 1, 1, 11}), Der(5, {})}));
    Bytes rsa = Der(0x30, Join({Der(6, {0x2a, 0x86, 0x48, 0x86, 0xf7, 0x0d, 1, 1, 1}), Der(5, {})}));
    Bytes name = Der(0x30, Der(0x31, Der(0x30, Join({Der(6, {0x55, 4, 3}), Der(0x0c, {'t', 'e', 's', 't'})}))));
    const std::string date = "260101000000Z";
    Bytes time = Der(0x17, Bytes(date.begin(), date.end()));
    Bytes key = Der(0x30, Join({Der(2, {0, 0xab, 0xcd}), Der(2, {1, 0, 1})}));
    key.insert(key.begin(), 0);
    Bytes spki = Der(0x30, Join({rsa, Der(3, key)}));
    Bytes tbs = Join(
        {version ? Der(0xa0, Der(2, {2})) : Bytes{}, Der(2, {1}), sig, name, Der(0x30, Join({time, time})), name, spki}
    );
    return Der(0x30, Join({Der(0x30, tbs), sig, Der(3, {0, 0xab, 0xcd})}));
}

static void
TestCertificates()
{
    for (bool v : {false, true})
    {
        auto cert = Fixture(v);
        DOMITO_CERTIFICATE_SPKI spki = {};
        Check(
            NT_SUCCESS(DomitoGetCertificateSpki(cert.data(), static_cast<ULONG>(cert.size()), &spki)),
            "v1/v3 extraction"
        );
        Check(
            spki.Content.Data > spki.Encoded.Data && spki.Content.Size < spki.Encoded.Size,
            "content and complete TLV differ"
        );
        auto context = CertCreateCertificateContext(X509_ASN_ENCODING, cert.data(), static_cast<DWORD>(cert.size()));
        Check(context != nullptr, "Crypt32 independently decodes fixture certificate");
        BYTE* encoded = nullptr;
        DWORD size = 0;
        Check(
            CryptEncodeObjectEx(
                X509_ASN_ENCODING,
                X509_PUBLIC_KEY_INFO,
                &context->pCertInfo->SubjectPublicKeyInfo,
                CRYPT_ENCODE_ALLOC_FLAG,
                nullptr,
                &encoded,
                &size
            ) != FALSE,
            "Crypt32 independent SPKI encoding"
        );
        Check(
            size == spki.Encoded.Size && !std::memcmp(encoded, spki.Encoded.Data, size),
            "extracted complete SPKI equals Crypt32"
        );
        LocalFree(encoded);
        CertFreeCertificateContext(context);
        UCHAR full[64];
        UCHAR legacy[64];
        UCHAR whole[64];
        UCHAR oracle[64];
        ULONG a = 64;
        ULONG b = 64;
        ULONG c = 64;
        DWORD d = 64;
        Check(
            NT_SUCCESS(DomitoHashCertificate(
                cert.data(), static_cast<ULONG>(cert.size()), DomitoCertificateHashSpkiDer, CALG_SHA_256, full, &a
            )),
            "full SPKI hash"
        );
        Check(
            NT_SUCCESS(DomitoHashCertificate(
                cert.data(),
                static_cast<ULONG>(cert.size()),
                DomitoCertificateHashSpkiContentLegacy,
                CALG_SHA_256,
                legacy,
                &b
            )),
            "legacy content SPKI hash"
        );
        Check(a == 32 && b == 32 && std::memcmp(full, legacy, 32) != 0, "explicit SPKI conventions differ");
        Check(
            CryptHashCertificate2(
                BCRYPT_SHA256_ALGORITHM, 0, nullptr, spki.Content.Data, spki.Content.Size, oracle, &d
            ) != FALSE &&
                d == b && !std::memcmp(legacy, oracle, d),
            "legacy hash matches exact content-only oracle"
        );
        Check(
            NT_SUCCESS(DomitoHashCertificate(
                cert.data(), static_cast<ULONG>(cert.size()), DomitoCertificateHashDer, CALG_SHA_256, whole, &c
            )),
            "DER thumbprint"
        );
        d = 64;
        Check(
            CryptHashCertificate2(
                BCRYPT_SHA256_ALGORITHM, 0, nullptr, cert.data(), static_cast<DWORD>(cert.size()), oracle, &d
            ) != FALSE &&
                d == c && !std::memcmp(whole, oracle, d),
            "thumbprint matches Crypt32"
        );

        for (ULONG i = 0; i < cert.size(); ++i)
        {
            spki.Encoded = {cert.data(), 99};
            Check(!NT_SUCCESS(DomitoGetCertificateSpki(cert.data(), i, &spki)), "every truncated certificate rejected");
            Check(!spki.Encoded.Data && !spki.Content.Data, "failed extraction clears borrowed views");
        }

        cert.push_back(0);
        Check(
            !NT_SUCCESS(DomitoGetCertificateSpki(cert.data(), static_cast<ULONG>(cert.size()), &spki)),
            "trailing byte rejected"
        );
    }

    for (auto malformed :
         {Bytes{0x30, 0x80, 0, 0},
          Bytes{0x30, 0x81, 1, 0},
          Bytes{0x30, 0x85, 1, 2, 3, 4, 5},
          Bytes{0x30, 0x84, 0xff, 0xff, 0xff, 0xff}})
    {
        DOMITO_CERTIFICATE_SPKI spki = {};
        Check(
            !NT_SUCCESS(DomitoGetCertificateSpki(malformed.data(), static_cast<ULONG>(malformed.size()), &spki)),
            "bad DER lengths rejected"
        );
    }

    const UCHAR abc[] = {'a', 'b', 'c'};
    UCHAR hash[64];

    for (ALG_ID alg : {CALG_SHA1, CALG_SHA_256, CALG_SHA_384, CALG_SHA_512})
    {
        ULONG n = 0;
        Check(DomitoHashBuffer(abc, 3, alg, nullptr, &n) == STATUS_BUFFER_TOO_SMALL, "hash sizing");
        Check(n == 20 || n == 32 || n == 48 || n == 64, "known output size");
        ULONG cap = n - 1;
        Check(DomitoHashBuffer(abc, 3, alg, hash, &cap) == STATUS_BUFFER_TOO_SMALL && cap == n, "short hash capacity");
        cap = 64;
        Check(NT_SUCCESS(DomitoHashBuffer(nullptr, 0, alg, hash, &cap)), "empty hash");
    }

    for (int fail = 1; fail <= 4; ++fail)
    {
        cryptoFailure = fail;
        ULONG n = 64;
        std::memset(hash, 0xff, 64);
        Check(!NT_SUCCESS(DomitoHashBuffer(abc, 3, CALG_SHA_256, hash, &n)) && !n, "CNG failure propagated");

        for (int i = 0; i < 32; ++i)
        {
            Check(!hash[i], "CNG failure digest erased");
        }
    }

    cryptoFailure = 0;

    for (unsigned byte = 0; byte < 256; ++byte)
    {
        UCHAR in = static_cast<UCHAR>(byte);
        UCHAR out = 0;
        WCHAR hex[3];
        Check(NT_SUCCESS(DomitoEncodeHex(&in, 1, hex, 3)), "hex encoding");
        UNICODE_STRING str = {4, 6, hex};
        Check(NT_SUCCESS(DomitoDecodeHex(&str, &out, 1)) && out == in, "all byte values hex roundtrip");
    }

    for (WCHAR ch : {L'g', L'/', L':', L' ', L'\0', L'\u1234'})
    {
        WCHAR text[] = {L'0', ch};
        UNICODE_STRING str = {4, 4, text};
        UCHAR out = 9;
        Check(
            DomitoDecodeHex(&str, &out, 1) == STATUS_INVALID_PARAMETER && !out,
            "invalid hex digit rejected and output zeroed"
        );
    }

    WCHAR shortHex[2];
    Check(DomitoEncodeHex(abc, 3, shortHex, 2) == STATUS_BUFFER_TOO_SMALL && !shortHex[0], "short hex buffer cleared");

    // Put every truncated prefix immediately before an inaccessible page.
    SYSTEM_INFO system = {};
    GetSystemInfo(&system);
    const SIZE_T page = system.dwPageSize;
    auto guarded = static_cast<PUCHAR>(VirtualAlloc(nullptr, 2 * page, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE));
    Check(guarded != nullptr, "guard-page allocation");
    DWORD old = 0;
    Check(VirtualProtect(guarded + page, page, PAGE_NOACCESS, &old) != FALSE, "guard page installed");
    const auto fixture = Fixture(true);

    for (ULONG n = 0; n < fixture.size(); ++n)
    {
        UCHAR* input = guarded + page - n;
        std::memcpy(input, fixture.data(), n);
        DOMITO_CERTIFICATE_SPKI spki = {};
        Check(!NT_SUCCESS(DomitoGetCertificateSpki(input, n, &spki)), "truncated DER at guard page rejected");
    }

    VirtualFree(guarded, 0, MEM_RELEASE);

    // Read-only optional corpus: exercises real long-length RSA/ECC certificates.
    HCERTSTORE store = CertOpenSystemStoreW(0, L"ROOT");

    if (store)
    {
        PCCERT_CONTEXT context = nullptr;
        unsigned count = 0;
        while (count < 16 && (context = CertEnumCertificatesInStore(store, context)) != nullptr)
        {
            ++count;
            DOMITO_CERTIFICATE_SPKI spki = {};
            Check(
                NT_SUCCESS(DomitoGetCertificateSpki(context->pbCertEncoded, context->cbCertEncoded, &spki)),
                "real certificate SPKI extraction"
            );
            BYTE* encoded = nullptr;
            DWORD size = 0;
            Check(
                CryptEncodeObjectEx(
                    X509_ASN_ENCODING,
                    X509_PUBLIC_KEY_INFO,
                    &context->pCertInfo->SubjectPublicKeyInfo,
                    CRYPT_ENCODE_ALLOC_FLAG,
                    nullptr,
                    &encoded,
                    &size
                ) != FALSE,
                "real Crypt32 SPKI encoding"
            );
            Check(
                size == spki.Encoded.Size && !std::memcmp(encoded, spki.Encoded.Data, size),
                "real certificate SPKI equals Crypt32"
            );
            LocalFree(encoded);
        }

        if (context)
        {
            CertFreeCertificateContext(context);
        }

        CertCloseStore(store, 0);
    }
}

static void
TestOwnership()
{
    DOMITO_CI_VALIDATION_RESULT result = {};
    UNICODE_STRING path = {8, 10, const_cast<PWCHAR>(L"\\foo")};

    for (int fail : {0, 2, 3})
    {
        failure = fail;
        const auto r = refs, d = derefs, c = closed;
        result.PolicyInfo.ChainInfo = static_cast<MINCRYPT_CHAIN_INFO*>(Allocate(1));
        NTSTATUS s = DomitoValidateFilePathEx(&path, 0x20, 7, DomitoCiSignatureSelectionNative, &result);
        Check(NT_SUCCESS(s) == (fail == 0), "path/reference failure returned");
        Check(
            refs - r == derefs - d && closed - c == (fail == 3 ? 0u : 1u),
            "all references and owned path handles released"
        );
        DomitoFreeCiValidationResult(&result);
        Check(!live, "path no allocations leaked");
    }

    failure = 0;
    validateStatus = STATUS_INVALID_IMAGE_HASH;
    const auto c = closed;
    Check(
        DomitoValidateFileHandleEx(&fileToken, 0x20, 7, DomitoCiSignatureSelectionNative, &result) == validateStatus,
        "CI failure retained"
    );
    Check(closed == c && result.PolicyInfo.ChainInfo, "borrowed handle open and diagnostics retained");
    DomitoFreeCiValidationResult(&result);
    validateStatus = STATUS_SUCCESS;

    for (int fail : {0, 1, 4, 5})
    {
        failure = fail;
        PUNICODE_STRING name = reinterpret_cast<PUNICODE_STRING>(1);
        NTSTATUS s = DomitoGetProcessImageNameEx(&processToken, &name);
        Check(NT_SUCCESS(s) == (fail == 0), "PEPROCESS lookup failure");

        if (NT_SUCCESS(s))
        {
            Check(
                name->Buffer == reinterpret_cast<PWCHAR>(name + 1) && name->Length == 8,
                "one contiguous process-name allocation"
            );
        }
        else
        {
            Check(!name, "failed process name null");
        }

        DomitoFreeProcessImageName(name);
        Check(!live, "process-name allocations released");
    }

    for (int fail : {0, 1, 6, 7, 10, 11})
    {
        failure = fail;
        PUNICODE_STRING name = reinterpret_cast<PUNICODE_STRING>(1);
        NTSTATUS s = DomitoGetProcessImageName(123, &name);
        Check(NT_SUCCESS(s) == (fail == 0), "PID query failure");

        if (!NT_SUCCESS(s))
        {
            Check(!name, "PID failure never leaves freed pointer");
        }

        DomitoFreeProcessImageName(name);
        Check(!live, "PID query cleanup");
    }

    failure = 0;
    MINCERT_BLOB blob = {1, reinterpret_cast<PUCHAR>(&fileToken)};

    for (int fail : {0, 1, 8, 9})
    {
        failure = fail;
        UNICODE_STRING name = {};
        NTSTATUS s = DomitoGetCertificatePublisherName(&blob, &name);
        Check(NT_SUCCESS(s) == (fail == 0), "publisher error propagation");
        DomitoFreeCertificatePublisherName(&name);
        DomitoFreeCertificatePublisherName(&name);
        Check(!live && !name.Buffer, "publisher matching free including partial failure");
    }

    failure = 0;
    ANSI_STRING module = RTL_CONSTANT_STRING("module");
    ANSI_STRING name = RTL_CONSTANT_STRING("export");
    PVOID address = reinterpret_cast<PVOID>(1);
    moduleStatus = STATUS_NOT_FOUND;
    Check(
        DomitoFindModuleExportedFunctionAddress(&module, &name, &address) == STATUS_NOT_FOUND && !address,
        "missing module nulls export"
    );
    moduleStatus = STATUS_SUCCESS;
    exportStatus = STATUS_NOT_SUPPORTED;
    address = reinterpret_cast<PVOID>(1);
    Check(
        DomitoFindModuleExportedFunctionAddress(&module, &name, &address) == STATUS_NOT_SUPPORTED && !address,
        "unsupported forwarder status preserved"
    );
    exportStatus = STATUS_SUCCESS;
}

static void
TestChains()
{
    alignas(8) UCHAR storage[512] = {};
    auto chain = reinterpret_cast<MINCRYPT_CHAIN_INFO*>(storage);
    chain->Size = sizeof(storage);
    chain->ChainElements = reinterpret_cast<MINCRYPT_CHAIN_ELEMENT*>(storage + 56);
    chain->NumberOfChainElement = 2;

    for (ULONG i = 0; i < 2; ++i)
    {
        chain->ChainElements[i].Certificate = {1, storage + 400 + i};
        chain->ChainElements[i].HashSize = 32;
    }

    MINCRYPT_POLICY_INFO policy = {};
    policy.Size = sizeof(policy);
    policy.ChainInfo = chain;
    DOMITO_CI_CHAIN_VIEW view = {};
#ifdef DOMITO_TEST_OTHER_ARCH
    Check(
        DomitoGetCiChainView(&policy, DomitoCiChainLayoutObservedX64, &view) == STATUS_NOT_SUPPORTED && !view.Count,
        "non-x64 rejects unverified layout"
    );
#else
    Check(
        NT_SUCCESS(DomitoGetCiChainView(&policy, DomitoCiChainLayoutObservedX64, &view)) && view.Count == 2,
        "minimal historical 56-byte chain prefix"
    );
    // The last certificate is an arbitrary byte, not a proven self-signed root.
    chain->ChainElements[1].Certificate.Data = storage + 512;
    Check(
        !NT_SUCCESS(DomitoGetCiChainView(&policy, DomitoCiChainLayoutObservedX64, &view)) && !view.Count,
        "child span past allocation rejected"
    );
    chain->ChainElements[1].Certificate.Data = storage + 401;
    chain->ChainElements[0].Subject = {reinterpret_cast<PCHAR>(storage + 500), 13};
    Check(
        !NT_SUCCESS(DomitoGetCiChainView(&policy, DomitoCiChainLayoutObservedX64, &view)),
        "malformed name span rejected"
    );
    chain->ChainElements[0].Subject = {};
    chain->NumberOfChainElement = 0xffffffff;
    Check(
        !NT_SUCCESS(DomitoGetCiChainView(&policy, DomitoCiChainLayoutObservedX64, &view)),
        "element multiplication overflow rejected"
    );
    chain->NumberOfChainElement = 2;
    chain->ChainElements[0].HashSize = 65;
    Check(
        !NT_SUCCESS(DomitoGetCiChainView(&policy, DomitoCiChainLayoutObservedX64, &view)),
        "oversize CI element hash rejected"
    );
    chain->ChainElements[0].HashSize = 32;

    for (ULONG n = 0; n < 56; ++n)
    {
        chain->Size = n;
        Check(
            !NT_SUCCESS(DomitoGetCiChainView(&policy, DomitoCiChainLayoutObservedX64, &view)),
            "short historical header rejected"
        );
    }
#endif
}

int
main()
{
    TestCertificates();
    TestOwnership();
    TestChains();
    Check(!live, "all allocations released");
    std::printf("PASS: %u consumer-helper checks, no live allocations\n", checks);
}