//-----------------------------------------------------------------------------
// Author - Marco Martinelli - https://github.com/martinellimarco
//-----------------------------------------------------------------------------

#pragma warning (disable : 4127)

#include "scream.h"
#include "ivshmemsavedata.h"

//=============================================================================
// Defines
//=============================================================================

#define MAGIC_NUM           0x11112014;

#pragma code_seg("PAGE")
//=============================================================================
// CIVSHMEMSaveData
//=============================================================================

//=============================================================================
BOOLEAN CIVSHMEMSaveData::RequestMMAP() {
	if (!m_ivshmem.devObj) {
		return FALSE;
	}

    IO_STATUS_BLOCK ioStatus = { 0 };
    KEVENT event;
    KeInitializeEvent(&event, NotificationEvent, FALSE);
    PIRP irp = IoBuildDeviceIoControlRequest(
        IOCTL_IVSHMEM_REQUEST_MMAP,
        m_ivshmem.devObj,
        NULL, 0,
        &m_ivshmem.mmap, sizeof(IVSHMEM_MMAP),
        FALSE,
        &event,
        &ioStatus
    );
    if (irp) {
        if (IoCallDriver(m_ivshmem.devObj, irp) == STATUS_PENDING) {
            KeWaitForSingleObject(&event, Executive, KernelMode, FALSE, NULL);
        }
        if (ioStatus.Status == 0) {
            return TRUE;
        }
        else {
            DPF_ENTER(("IVSHMEM: IRP for REQUEST_MMAP failed with status %u\n", ioStatus.Status));
        }
    }
    else {
        DPF_ENTER(("IVSHMEM: Failed to get the IRP for REQUEST_MMAP\n"));
    }
    return FALSE;
}

//=============================================================================
void CIVSHMEMSaveData::ReleaseMMAP() {
	if (!m_ivshmem.devObj) {
		return;
	}

    IO_STATUS_BLOCK ioStatus = { 0 };
    KEVENT event;
    KeInitializeEvent(&event, NotificationEvent, FALSE);
    PIRP irp = IoBuildDeviceIoControlRequest(
        IOCTL_IVSHMEM_RELEASE_MMAP,
        m_ivshmem.devObj,
        NULL, 0,
        NULL, 0,
        FALSE,
        &event,
        &ioStatus
    );
    if (irp) {
        if (IoCallDriver(m_ivshmem.devObj, irp) == STATUS_PENDING) {
            KeWaitForSingleObject(&event, Executive, KernelMode, FALSE, NULL);
        }
        if (ioStatus.Status != 0) {
            DPF_ENTER(("IVSHMEM: IRP for RELEASE_MMAP failed with status %u\n", ioStatus.Status));
        }
    }
    else {
        DPF_ENTER(("IVSHMEM: Failed to get the IRP for RELEASE_MMAP\n"));
    }
}

//=============================================================================
CIVSHMEMSaveData::CIVSHMEMSaveData() : m_pBuffer(NULL), m_ulOffset(0), m_ulSendOffset(0), m_fWriteDisabled(FALSE) {
    PAGED_CODE();

    DPF_ENTER(("[CIVSHMEMSaveData::CIVSHMEMSaveData]"));
    
    if (g_UseIVSHMEM) {
        // allocate work item for this stream
        m_pWorkItem = (PIVSHMEM_SAVEWORKER_PARAM)ExAllocatePoolWithTag(NonPagedPool, sizeof(IVSHMEM_SAVEWORKER_PARAM), MSVAD_POOLTAG);
        if (m_pWorkItem) {
            m_pWorkItem->WorkItem = IoAllocateWorkItem(GetDeviceObject());
            KeInitializeEvent(&(m_pWorkItem->EventDone), NotificationEvent, TRUE);
        }

        // Enumerate IVSHMEM devices
        GUID interfaceClassGuid = { 0xdf576976,0x569d,0x4672,0x95,0xa0,0xf5,0x7e,0x4e,0xa0,0xb2,0x10 };
        PWSTR symbolicLinkList;
        NTSTATUS ntStatus;
        ntStatus = IoGetDeviceInterfaces(
            &interfaceClassGuid,
            NULL,
            0,
            &symbolicLinkList
        );

        if (NT_SUCCESS(ntStatus) && NULL != symbolicLinkList) {
            PFILE_OBJECT fileObj;
            PDEVICE_OBJECT devObj;
            UNICODE_STRING objName;

            KEVENT event;
            IO_STATUS_BLOCK ioStatus = { 0 };

            PIRP irp;
            UINT64 ivshmemSize = 0;
            
            for (PWSTR symbolicLink = symbolicLinkList;
                symbolicLink[0] != NULL && symbolicLink[1] != NULL;
                symbolicLink += wcslen(symbolicLink) + 1) {
                RtlInitUnicodeString(&objName, symbolicLink),

                ntStatus = IoGetDeviceObjectPointer(
                    &objName,
                    FILE_ALL_ACCESS,
                    &fileObj,
                    &devObj
                );
                if (NT_SUCCESS(ntStatus)) {
                    ObDereferenceObject(fileObj);
                    ntStatus = ObReferenceObjectByPointer(devObj, FILE_ALL_ACCESS, NULL, KernelMode);
                    if (NT_SUCCESS(ntStatus)) {
                        ivshmemSize = 0;
                        KeInitializeEvent(&event, NotificationEvent, FALSE);
                        irp = IoBuildDeviceIoControlRequest(
                            IOCTL_IVSHMEM_REQUEST_SIZE,
                            devObj,
                            NULL, 0,
                            &ivshmemSize, sizeof(ivshmemSize),
                            FALSE,
                            &event,
                            &ioStatus
                        );
                        if (irp) {
                            ntStatus = IoCallDriver(devObj, irp);
                            if (ntStatus == STATUS_PENDING) {
                                KeWaitForSingleObject(&event, Executive, KernelMode, FALSE, NULL);
                            }
                            if (ioStatus.Status == 0 && ivshmemSize == g_UseIVSHMEM*1024*1024) { //we suppose that we are the only g_UseIVSHMEM*1MiB IVSHMEM device around
                                DPF_ENTER(("IVSHMEM Device Detected!\n"));
                                
                                RtlZeroMemory(&m_ivshmem, sizeof(IVSHMEM_OBJECT));
                                m_ivshmem.devObj = devObj;
                                // we have found a suitable device, we skip the others, and we DON'T dereference devObj, it will be dereferenced in the destructor
                                break;
                            }
                        }
                        else {
                            DPF_ENTER(("IVSHMEM: Failed to get the IRP for REQUEST_SIZE\n"));
                        }
                    }

                    ObDereferenceObject(devObj);
                }
            }
        }

        ExFreePool(symbolicLinkList);
    }
} // CIVSHMEMSaveData

//=============================================================================
CIVSHMEMSaveData::~CIVSHMEMSaveData() {
    PAGED_CODE();

    DPF_ENTER(("[CIVSHMEMSaveData::~CIVSHMEMSaveData]"));
    if (g_UseIVSHMEM) {
        // frees the work item
        if (m_pWorkItem->WorkItem != NULL) {
            IoFreeWorkItem(m_pWorkItem->WorkItem);
            m_pWorkItem->WorkItem = NULL;
        }

        //release mmap
        if (m_ivshmem.initialized) {
            m_ivshmem.initialized = false;
            if (RequestMMAP()) {
                //set to 0 the header so the receiver know we terminated
                RtlZeroMemory(m_ivshmem.mmap.ptr, sizeof(IVSHMEM_SCREAM_HEADER));
            }
            ReleaseMMAP();
        }
        if (m_ivshmem.devObj) {
            ObDereferenceObject(m_ivshmem.devObj);
        }

        if (m_pBuffer) {
            ExFreePoolWithTag(m_pBuffer, MSVAD_POOLTAG);
        }
    }
} // CIVSHMEMSaveData

//=============================================================================
void CIVSHMEMSaveData::DestroyWorkItems(void) {
    PAGED_CODE();
    
    DPF_ENTER(("[CIVSHMEMSaveData::DestroyWorkItems]"));

    if (m_pWorkItem) {
        ExFreePoolWithTag(m_pWorkItem, MSVAD_POOLTAG);
        m_pWorkItem = NULL;
    }

} // DestroyWorkItems

//=============================================================================
void CIVSHMEMSaveData::Disable(BOOL fDisable) {
    PAGED_CODE();

    m_fWriteDisabled = fDisable;
} // Disable

//=============================================================================
NTSTATUS CIVSHMEMSaveData::SetDeviceObject(IN PDEVICE_OBJECT DeviceObject) {
    PAGED_CODE();

    ASSERT(DeviceObject);

    NTSTATUS ntStatus = STATUS_SUCCESS;
    
    m_pDeviceObject = DeviceObject;
    return ntStatus;
}

//=============================================================================
PDEVICE_OBJECT CIVSHMEMSaveData::GetDeviceObject(void) {
    PAGED_CODE();

    return m_pDeviceObject;
}

#pragma code_seg("PAGE")
//=============================================================================
NTSTATUS CIVSHMEMSaveData::Initialize(DWORD nSamplesPerSec, WORD wBitsPerSample, WORD nChannels, DWORD dwChannelMask) {
    PAGED_CODE();

    NTSTATUS          ntStatus = STATUS_SUCCESS;

    DPF_ENTER(("[CIVSHMEMSaveData::Initialize]"));

    m_ivshmem.chunkSize = (UINT32)((wBitsPerSample>>3)*nChannels*nSamplesPerSec/50);
    if (RequestMMAP()) {
        PIVSHMEM_SCREAM_HEADER hdr = (PIVSHMEM_SCREAM_HEADER)m_ivshmem.mmap.ptr;

        RtlZeroMemory(hdr, sizeof(IVSHMEM_SCREAM_HEADER));
        m_ivshmem.writeIdx = 0;
        hdr->offset = m_ivshmem.offset = (UINT8)sizeof(IVSHMEM_SCREAM_HEADER);
        hdr->chunkSize = m_ivshmem.chunkSize;
        hdr->maxChunks = m_ivshmem.maxChunks = (UINT16)(m_ivshmem.mmap.size / m_ivshmem.chunkSize);
        m_ivshmem.bufferSize = m_ivshmem.chunkSize*m_ivshmem.maxChunks;

        // Only multiples of 44100 and 48000 are supported
        hdr->sampleRate = (UINT8)((nSamplesPerSec % 44100) ? (0 + (nSamplesPerSec / 48000)) : (128 + (nSamplesPerSec / 44100)));
        hdr->sampleSize = (UINT8)(wBitsPerSample);
        hdr->channels = (UINT8)(nChannels);
        hdr->channelMap = (UINT16)(dwChannelMask);

        hdr->magic = MAGIC_NUM;
        m_ivshmem.initialized = true;
    }
    else {
        m_ivshmem.initialized = false;
    }
    ReleaseMMAP();

    // Allocate memory for data buffer.
    if (NT_SUCCESS(ntStatus)) {
        m_pBuffer = (PBYTE)ExAllocatePoolWithTag(NonPagedPool, m_ivshmem.bufferSize, MSVAD_POOLTAG);
        if (!m_pBuffer) {
            DPF(D_TERSE, ("[Could not allocate memory for sending data]"));
            ntStatus = STATUS_INSUFFICIENT_RESOURCES;
        }
    }

    return ntStatus;
} // Initialize

//=============================================================================
IO_WORKITEM_ROUTINE IVSHMEMSendDataWorkerCallback;

VOID IVSHMEMSendDataWorkerCallback(PDEVICE_OBJECT pDeviceObject, IN  PVOID  Context) {
    UNREFERENCED_PARAMETER(pDeviceObject);

    PAGED_CODE();

    ASSERT(Context);
    
    PIVSHMEM_SAVEWORKER_PARAM pParam = (PIVSHMEM_SAVEWORKER_PARAM) Context;
    PCIVSHMEMSaveData        pSaveData;

    ASSERT(pParam->pSaveData);

    if (pParam->WorkItem) {
        pSaveData = pParam->pSaveData;
        pSaveData->IVSHMEMSendData();
    }

    KeSetEvent(&(pParam->EventDone), 0, FALSE);
} // SendDataWorkerCallback

//=============================================================================
void CIVSHMEMSaveData::IVSHMEMSendData() {
    ULONG storeOffset;
    PBYTE mmap = NULL;
    if (RequestMMAP()) {
        mmap = (PBYTE)m_ivshmem.mmap.ptr + m_ivshmem.offset;
    }

    PIVSHMEM_SCREAM_HEADER hdr = (PIVSHMEM_SCREAM_HEADER)m_ivshmem.mmap.ptr;
    
    while (1) {
        // Read latest storeOffset. There might be new data.
        storeOffset = m_ulOffset;

        // Abort if there's nothing to send. Note: When storeOffset < sendOffset, we can always send a chunk.
        if ((storeOffset >= m_ulSendOffset) && ((storeOffset - m_ulSendOffset) < m_ivshmem.chunkSize))
            break;

        // Send a chunk
        if (mmap) {
            if (++m_ivshmem.writeIdx >= m_ivshmem.maxChunks) {
                m_ivshmem.writeIdx = 0;
            }
            RtlCopyMemory(&mmap[m_ivshmem.writeIdx*m_ivshmem.chunkSize], &m_pBuffer[m_ulSendOffset], m_ivshmem.chunkSize);
            hdr->writeIdx = m_ivshmem.writeIdx;
        }

        m_ulSendOffset += m_ivshmem.chunkSize; if (m_ulSendOffset >= m_ivshmem.bufferSize) m_ulSendOffset = 0;
    }
    ReleaseMMAP();
}

#pragma code_seg("PAGE")
//=============================================================================
void CIVSHMEMSaveData::WaitAllWorkItems(void) {
    PAGED_CODE();

    DPF_ENTER(("[CIVSHMEMSaveData::WaitAllWorkItems]"));

    DPF(D_VERBOSE, ("[Waiting for WorkItem]"));
    KeWaitForSingleObject(&(m_pWorkItem->EventDone), Executive, KernelMode, FALSE, NULL);
    
} // WaitAllWorkItems

#pragma code_seg()
//=============================================================================
void CIVSHMEMSaveData::WriteData(IN PBYTE pBuffer, IN ULONG ulByteCount) {
    ASSERT(pBuffer);

    if (!m_ivshmem.initialized) {
        return;
    }

    LARGE_INTEGER timeOut = { 0 };
    NTSTATUS ntStatus;
    ULONG offset;
    ULONG toWrite;
    ULONG w;

    if (m_fWriteDisabled) {
        return;
    }

    //DPF_ENTER(("[CIVSHMEMSaveData::WriteData ulByteCount=%lu]", ulByteCount));

    // Undersized (paranoia)
    if (0 == ulByteCount) {
        return;
    }

    // Oversized (paranoia)
    if (ulByteCount > (m_ivshmem.bufferSize / 2)) {
        return;
    }

    // Append to ring buffer. Don't write intermediate states to m_ulOffset,
    // but update it once at the end.
    offset = m_ulOffset;
    toWrite = ulByteCount;
    while (toWrite > 0) {
        w = offset % m_ivshmem.chunkSize;
        if (w > 0) {
            // Fill up last chunk
            w = (m_ivshmem.chunkSize - w);
            w = (toWrite < w) ? toWrite : w;
            RtlCopyMemory(&(m_pBuffer[offset]), &(pBuffer[ulByteCount - toWrite]), w);
        }
        else {
            // Start a new chunk
            w = ((m_ivshmem.bufferSize - offset) < toWrite) ? (m_ivshmem.bufferSize - offset) : toWrite;
            w = (w > m_ivshmem.chunkSize) ? m_ivshmem.chunkSize : w;
            RtlCopyMemory(&(m_pBuffer[offset]), &(pBuffer[ulByteCount - toWrite]), w);
        }
        toWrite -= w;
        offset += w;  if (offset >= m_ivshmem.bufferSize) offset = 0;
    }
    m_ulOffset = offset;

    // If I/O worker was done, relaunch it
    ntStatus = KeWaitForSingleObject(&(m_pWorkItem->EventDone), Executive, KernelMode, FALSE, &timeOut);
    if (STATUS_SUCCESS == ntStatus) {
        m_pWorkItem->pSaveData = this;
        KeResetEvent(&(m_pWorkItem->EventDone));
        IoQueueWorkItem(m_pWorkItem->WorkItem, IVSHMEMSendDataWorkerCallback, CriticalWorkQueue, (PVOID)m_pWorkItem);
    }
} // WriteData