#include <ntos.h>
#include <zwapi.h>
#include <ntconfig.h>
#include "vdmp.h"

#ifdef ALLOC_PRAGMA
#pragma alloc_text(PAGE, VdmpInitialize)
#endif

#define KEY_VALUE_BUFFER_SIZE 1024

#if DEVL
ULONG VdmBopCount;
#endif


NTSTATUS
VdmpInitialize(
    PVDMICAUSERDATA pIcaUserData
    )

/*++

Routine Description:

    Initialize the address space of a VDM.

Arguments:

    None,

Return Value:

    NTSTATUS.

--*/

{
    NTSTATUS Status;
    OBJECT_ATTRIBUTES ObjectAttributes;
    UNICODE_STRING SectionName;
    UNICODE_STRING WorkString;
    ULONG ViewSize;
    LARGE_INTEGER ViewBase;
    PVOID BaseAddress;
    PVOID destination;
    HANDLE SectionHandle, RegistryHandle;
    PEPROCESS Process;
    ULONG ResultLength;
    ULONG Index;
    PCM_FULL_RESOURCE_DESCRIPTOR ResourceDescriptor;
    PCM_PARTIAL_RESOURCE_DESCRIPTOR PartialResourceDescriptor;
    PKEY_VALUE_FULL_INFORMATION KeyValueBuffer;
    PCM_ROM_BLOCK BiosBlock;
    ULONG LastMappedAddress;
    PVDM_PROCESS_OBJECTS pVdmObjects;
    USHORT PagedQuotaCharged = 0;
    USHORT NonPagedQuotaCharged = 0;

    PAGED_CODE();

    if ((KeI386MachineType & MACHINE_TYPE_PC_9800_COMPATIBLE) == 0) {

        //
        // This is PC/AT (and FMR in Japan) VDM.
        //

        RtlInitUnicodeString(
            &SectionName,
            L"\\Device\\PhysicalMemory"
            );

        InitializeObjectAttributes(
            &ObjectAttributes,
            &SectionName,
            OBJ_CASE_INSENSITIVE,
            (HANDLE) NULL,
            (PSECURITY_DESCRIPTOR) NULL
            );

        Status = ZwOpenSection(
            &SectionHandle,
            SECTION_ALL_ACCESS,
            &ObjectAttributes
            );

        if (!NT_SUCCESS(Status)) {

            return Status;

        }

        //
        // Copy the first page of memory into the VDM's address space
        //

        BaseAddress = 0;
        destination = 0;
        ViewSize = 0x1000;
        ViewBase.LowPart = 0;
        ViewBase.HighPart = 0;

        Status = ZwMapViewOfSection(
            SectionHandle,
            NtCurrentProcess(),
            &BaseAddress,
            0,
            ViewSize,
            &ViewBase,
            &ViewSize,
            ViewUnmap,
            0,
            PAGE_READWRITE
            );

        if (!NT_SUCCESS(Status)) {

            return Status;

        }

        RtlMoveMemory(
            destination,
            BaseAddress,
            ViewSize
            );

        Status = ZwUnmapViewOfSection(
            NtCurrentProcess(),
            BaseAddress
            );

        if (!NT_SUCCESS(Status)) {

            return Status;

        }

        //
        // Map Rom into address space
        //

        BaseAddress = (PVOID) 0x000C0000;
        ViewSize = 0x40000;
        ViewBase.LowPart = 0x000C0000;
        ViewBase.HighPart = 0;


        //
        // First unmap the reserved memory.  This must be done here to prevent
        // the virtual memory in question from being consumed by some other
        // alloc vm call.
        //

        Status = ZwFreeVirtualMemory(
            NtCurrentProcess(),
            &BaseAddress,
            &ViewSize,
            MEM_RELEASE
            );

        // N.B.  This should probably take into account the fact that there are
        // a handfull of error conditions that are ok.  (such as no memory to
        // release.)

        if (!NT_SUCCESS(Status)) {

            return Status;

        }

        //
        // Set up and open KeyPath
        //

        InitializeObjectAttributes(
            &ObjectAttributes,
            &CmRegistryMachineHardwareDescriptionSystemName,
            OBJ_CASE_INSENSITIVE,
            (HANDLE)NULL,
            NULL
            );

        Status = ZwOpenKey(
            &RegistryHandle,
            KEY_READ,
            &ObjectAttributes
            );

        if (!NT_SUCCESS(Status)) {
            return Status;
        }

        //
        // Allocate space for the data
        //

        KeyValueBuffer = ExAllocatePool(PagedPool, KEY_VALUE_BUFFER_SIZE);
        if (KeyValueBuffer == NULL) {
            ZwClose(RegistryHandle);
            return STATUS_NO_MEMORY;
        }

        //
        // Get the data for the rom information
        //

        RtlInitUnicodeString(
            &WorkString,
            L"Configuration Data"
            );

        Status = ZwQueryValueKey(
            RegistryHandle,
            &WorkString,
            KeyValueFullInformation,
            KeyValueBuffer,
            KEY_VALUE_BUFFER_SIZE,
            &ResultLength
            );

        if (!NT_SUCCESS(Status)) {
            ZwClose(RegistryHandle);
            ExFreePool(KeyValueBuffer);
            return Status;
        }

        ResourceDescriptor = (PCM_FULL_RESOURCE_DESCRIPTOR)
            ((PUCHAR) KeyValueBuffer + KeyValueBuffer->DataOffset);

        if ((KeyValueBuffer->DataLength < sizeof(CM_FULL_RESOURCE_DESCRIPTOR)) ||
            (ResourceDescriptor->PartialResourceList.Count < 2)
        ) {
            ZwClose(RegistryHandle);
            ExFreePool(KeyValueBuffer);
            // No rom blocks.
            return STATUS_SUCCESS;
        }

        PartialResourceDescriptor = (PCM_PARTIAL_RESOURCE_DESCRIPTOR)
            ((PUCHAR)ResourceDescriptor +
            sizeof(CM_FULL_RESOURCE_DESCRIPTOR) +
            ResourceDescriptor->PartialResourceList.PartialDescriptors[0]
                .u.DeviceSpecificData.DataSize);


        if (KeyValueBuffer->DataLength < ((PUCHAR)PartialResourceDescriptor -
            (PUCHAR)ResourceDescriptor + sizeof(CM_PARTIAL_RESOURCE_DESCRIPTOR)
            + sizeof(CM_ROM_BLOCK))
        ) {
            ZwClose(RegistryHandle);
            ExFreePool(KeyValueBuffer);
            return STATUS_ILL_FORMED_SERVICE_ENTRY;
        }


        BiosBlock = (PCM_ROM_BLOCK)((PUCHAR)PartialResourceDescriptor +
            sizeof(CM_PARTIAL_RESOURCE_DESCRIPTOR));

        Index = PartialResourceDescriptor->u.DeviceSpecificData.DataSize /
            sizeof(CM_ROM_BLOCK);

        //
        // N.B.  Rom blocks begin on 2K (not necessarily page) boundaries
        //       They end on 512 byte boundaries.  This means that we have
        //       to keep track of the last page mapped, and round the next
        //       Rom block up to the next page boundary if necessary.
        //

        LastMappedAddress = 0xC0000;

        while (Index) {
    #if 0
            DbgPrint(
                "Bios Block, PhysAddr = %lx, size = %lx\n",
                BiosBlock->Address,
                BiosBlock->Size
                );
    #endif
            if ((Index > 1) &&
                ((BiosBlock->Address + BiosBlock->Size) == BiosBlock[1].Address)
            ) {
                //
                // Coalesce adjacent blocks
                //
                BiosBlock[1].Address = BiosBlock[0].Address;
                BiosBlock[1].Size += BiosBlock[0].Size;
                Index--;
                BiosBlock++;
                continue;
            }

            BaseAddress = (PVOID)(BiosBlock->Address);
            ViewSize = BiosBlock->Size;

            if ((ULONG)BaseAddress < LastMappedAddress) {
                if (ViewSize > (LastMappedAddress - (ULONG)BaseAddress)) {
                    ViewSize = ViewSize - (LastMappedAddress - (ULONG)BaseAddress);
                    BaseAddress = (PVOID)LastMappedAddress;
                } else {
                    ViewSize = 0;
                }
            }

            ViewBase.LowPart = (ULONG)BaseAddress;

            if (ViewSize > 0) {

                Status = ZwMapViewOfSection(
                    SectionHandle,
                    NtCurrentProcess(),
                    &BaseAddress,
                    0,
                    ViewSize,
                    &ViewBase,
                    &ViewSize,
                    ViewUnmap,
                    MEM_DOS_LIM,
                    PAGE_READWRITE
                    );

                if (!NT_SUCCESS(Status)) {
                    break;
                }

                LastMappedAddress = (ULONG)BaseAddress + ViewSize;
            }

            Index--;
            BiosBlock++;
        }

        //
        // Free up the handles
        //

        ZwClose(SectionHandle);
        ZwClose(RegistryHandle);
        ExFreePool(KeyValueBuffer);

    } else {

        //
        // This is PC-9800 Series VDM.
        //

        Status = STATUS_SUCCESS;
    }

    //
    // Mark the process as a vdm
    //

    Process = PsGetCurrentProcess();
    Process->Pcb.VdmFlag = TRUE;


    //
    // Create VdmObjects structure
    //
    // N.B.  We don't use ExAllocatePoolWithQuota because it
    //       takes a reference to the process (which ExFreePool
    //       dereferences).  Since we expect to clean up on
    //       process deletion, we don't need or want the reference
    //       (which will prevent the process from being deleted)
    //

    try {

        Process->VdmObjects = ExAllocatePool(
            NonPagedPool,
            sizeof(VDM_PROCESS_OBJECTS)
            );

        if (Process->VdmObjects == NULL) {
            return STATUS_NO_MEMORY;
        }

        //
        // We use NonPagedQuotaCharged to keep track of the quota to return
        // if this function fails
        //
        PsChargePoolQuota(Process, NonPagedPool, sizeof(VDM_PROCESS_OBJECTS));
        NonPagedQuotaCharged = sizeof(VDM_PROCESS_OBJECTS);

        pVdmObjects = Process->VdmObjects;
        RtlZeroMemory( Process->VdmObjects, sizeof(VDM_PROCESS_OBJECTS));
        ExInitializeFastMutex(&pVdmObjects->DelayIntFastMutex);
        KeInitializeSpinLock(&pVdmObjects->DelayIntSpinLock);
        InitializeListHead(&pVdmObjects->DelayIntListHead);


        pVdmObjects->pIcaUserData = ExAllocatePool(
            PagedPool,
            sizeof(VDMICAUSERDATA)
            );

        if (pVdmObjects->pIcaUserData == NULL) {
            Status = STATUS_NO_MEMORY;
        } else {

            //
            // We use NonPagedQuotaCharged to keep track of the quota to return
            // if this function fails
            //
            PsChargePoolQuota(
                Process,
                PagedPool,
                sizeof(VDMICAUSERDATA)
                );
            PagedQuotaCharged = sizeof(VDMICAUSERDATA);

            RtlZeroMemory( pVdmObjects->pIcaUserData, sizeof(VDMICAUSERDATA));


            //
            // Copy Ica addresses from service data (in user space) into
            // pVdmObjects->pIcaUserData
            //

            ProbeForRead(pIcaUserData, sizeof(VDMICAUSERDATA), sizeof(UCHAR));
            *pVdmObjects->pIcaUserData  = *pIcaUserData;


            //
            // Probe static addresses in IcaUserData.
            //
            ProbeForWriteHandle(pIcaUserData->phWowIdleEvent);

            ProbeForWrite(
                pIcaUserData->pIcaLock,
                sizeof(RTL_CRITICAL_SECTION),
                sizeof(UCHAR)
                );

            ProbeForWrite(
                pIcaUserData->pIcaMaster,
                sizeof(VDMVIRTUALICA),
                sizeof(UCHAR)
                );

            ProbeForWrite(
                pIcaUserData->pIcaSlave,
                sizeof(VDMVIRTUALICA),
                sizeof(UCHAR)
                );

            ProbeForWriteUlong(pIcaUserData->pIretHooked);
            ProbeForWriteUlong(pIcaUserData->pDelayIrq);
            ProbeForWriteUlong(pIcaUserData->pUndelayIrq);
            ProbeForWriteUlong(pIcaUserData->pDelayIret);
        }

    } except(ExSystemExceptionFilter()) {
        Status = GetExceptionCode();
    }


    if (!NT_SUCCESS(Status)) {
        if (pVdmObjects->pIcaUserData)
            ExFreePool(pVdmObjects->pIcaUserData);
        ExFreePool(pVdmObjects);
        Process->VdmObjects = NULL;

        //
        // Return Quota charged
        //
        PsReturnPoolQuota(Process, NonPagedPool, NonPagedQuotaCharged);
        PsReturnPoolQuota(Process, PagedPool, PagedQuotaCharged);
        }

    //
    // following codepath only for PC/AT (and FMR in Japan) vdm
    //

    if (KeI386MachineType & MACHINE_TYPE_PC_9800_COMPATIBLE == 0) {

#ifdef WHEN_IO_DISPATCHING_IMPROVED
        // Sudeepb - Once we improve the IO dispatching we should use this
        // routine. Currently we are dispatching the printer ports directly
        // from emv86.asm and instemul.asm

        VdmInitializePrinter ();
#endif
    }

    return Status;

} // end InitializeVDM()