pgpmemlock.c

vc环境下的pgp源码

	pmli = pde->pLockedList;
	while (pmli) 
	{
		if ((pmli->peProcess == peProcess) &&
			(pmli->pfo == pfo))
		{
			if (pPrev)
			{
				pPrev->next = pmli->next;
			}
			else 
			{
				pde->pLockedList = pmli->next;
			}
			pNext = pmli->next;
			MmUnlockPages (pmli->pmdl);
			IoFreeMdl (pmli->pmdl);
			PGPmlFreeItem (pmli);
			pmli = pNext;
		}
		else 
		{
			pPrev = pmli;
			pmli = pmli->next;
		}
	}

	KeReleaseSpinLock (&pde->slListLock, irql);
}


//----------------------------------------------------|
//  remove all locked memory blocks from 
//	list of blocks and unlock the memory
//
//	Arguments:
//		pDevO				- pointer to a device object
//

VOID
PGPmlUnlockAllBlocks (
    IN PDEVICE_OBJECT	pDevO)
{
	PDEVEXTENSION		pde			= NULL;
	PPGPMEMLOCKITEM		pmli		= NULL;
	PPGPMEMLOCKITEM		pNext		= NULL;
	KIRQL				irql;

	pde = (PDEVEXTENSION)(pDevO->DeviceExtension);

	// first wipe all blocks 
	pmli = pde->pLockedList;
	while (pmli) 
	{
		PGPmlWipeBlock (pmli->ulPage, pmli->ulNumPages);
		pmli = pmli->next;
	}

	KeAcquireSpinLock (&pde->slListLock, &irql);

	// now unlock blocks and destroy list
	pmli = pde->pLockedList;
	while (pmli) 
	{
		pNext = pmli->next;
		MmUnlockPages (pmli->pmdl);
		IoFreeMdl (pmli->pmdl);
		PGPmlFreeItem (pmli);
		pmli = pNext;
	}

	pde->pLockedList = NULL;

	KeReleaseSpinLock (&pde->slListLock, irql);
}


//----------------------------------------------------|
//  lock the pages in the specified virtual address range
//
//	Arguments:
//		pDevO				- pointer to a device object
//		IoBuffer			- pointer to the I/O buffer
//		InputBufferLength	- input buffer length
//		OutputBufferLength	- output buffer length
//
//	Return Value:
//		STATUS_SUCCESS if sucessful, otherwise
//		STATUS_*
//

NTSTATUS
PGPmlLockTheMemory(
    IN PDEVICE_OBJECT pDevO,
	IN PIRP			  pIrp,
    IN OUT PVOID      IoBuffer,
    IN ULONG          InputBufferLength,
    IN ULONG          OutputBufferLength)
{
    PPGPMEMLOCKSTRUCT	ppmls			= (PPGPMEMLOCKSTRUCT) IoBuffer;
    NTSTATUS			ntStatus		= 0;
	PMDL				pmdl			= NULL;
	ULONG				ulPage			= 0;
	ULONG				ulNumPages		= 0;

	PGPmlPrint (("PGPmemlock: PGPmlLockTheMemory.\n"));

	if ((InputBufferLength < sizeof (PGPMEMLOCKSTRUCT)) ||
		(OutputBufferLength < sizeof (PGPMEMLOCKSTRUCT)) ||
		(IoBuffer == NULL))
    {
       PGPmlPrint (("PGPmemlock: Err: Invalid input parameters.\n"));
       ntStatus = STATUS_INVALID_PARAMETER;
       goto done;
    }

	// assume no error
	ppmls->ulError = 0;

	// allocate and initialize a Memory Descriptor List 
	pmdl = IoAllocateMdl(ppmls->pMem, ppmls->ulNumBytes, FALSE, FALSE, NULL); 
	if (pmdl == NULL) {
		PGPmlPrint (("PGPmemlock: Err: IoAllocateMdl failed.\n"));
		ntStatus = STATUS_INSUFFICIENT_RESOURCES;
		ppmls->ulError = kPGPMemLockError_AllocError;
		goto done;
	}

	// calculate page numbers
	ulPage = ((ULONG)(ppmls->pMem)) >> WIN32PAGESIZE;
	ulNumPages = (((ppmls->ulNumBytes)-1) >> WIN32PAGESIZE ) +1;
	PGPmlPrint (("PGPmemlock: locking page(s) %x - %x.\n", 
						ulPage, ulPage+ulNumPages-1));

	// add this block to the array of locked blocks
	ntStatus = PGPmlAddBlockToList (pDevO, pIrp, ulPage, ulNumPages, pmdl);
	if (ntStatus != STATUS_SUCCESS) 
	{
		// we couldn't allocate a new list item
		PGPmlPrint (("PGPmemlock: Err: PGPmlAddBlockToList failed.\n"));
		ppmls->ulError = kPGPMemLockError_ListError;

		// free the mdl we allocated,
		IoFreeMdl(pmdl);
		goto done;
	}

	// try to lock down the client's area
	__try {
		MmProbeAndLockPages(pmdl, UserMode, IoModifyAccess);
	} __except(EXCEPTION_EXECUTE_HANDLER) {
		// client must have passed a bad address, or it exceeded its quota
		PGPmlPrint (("PGPmemlock: Err: MmProbeAndLockPages failed.\n"));
		ntStatus = GetExceptionCode();
		ppmls->ulError = kPGPMemLockError_LockError;

		// remove block from array and free the MDL we allocated,
		PGPmlUnlockBlock (pDevO, pIrp, ulPage, ulNumPages, FALSE);
		IoFreeMdl(pmdl);
		goto done;
	}

done:
    return ntStatus;
}


//----------------------------------------------------|
//  unlock the pages in the specified virtual address range
//
//	Arguments:
//		pDevO				- pointer to a device object
//		IoBuffer			- pointer to the I/O buffer
//		InputBufferLength	- input buffer length
//		OutputBufferLength	- output buffer length
//
//	Return Value:
//		STATUS_SUCCESS if sucessful, otherwise
//		STATUS_* 
//

NTSTATUS
PGPmlUnlockTheMemory(
    IN PDEVICE_OBJECT	pDevO,
	IN PIRP				pIrp,
    IN OUT PVOID		IoBuffer,
    IN ULONG			InputBufferLength,
    IN ULONG			OutputBufferLength)
{
	PPGPMEMLOCKSTRUCT	ppmls			= (PPGPMEMLOCKSTRUCT) IoBuffer;
	NTSTATUS			ntStatus		= STATUS_SUCCESS;
	ULONG				ulPage			= 0;
	ULONG				ulNumPages		= 0;

	PGPmlPrint (("PGPmemlock: PGPmlUnlockTheMemory.\n"));

	if ((InputBufferLength < sizeof (PGPMEMLOCKSTRUCT)) ||
		(OutputBufferLength < sizeof (PGPMEMLOCKSTRUCT)) ||
		(IoBuffer == NULL))
	{
		PGPmlPrint (("PGPmemlock: Err: Invalid input parameters.\n"));
		ntStatus = STATUS_INVALID_PARAMETER;
    }
	else 
	{
		// assume no error
		ppmls->ulError = 0;

		// calculate page numbers
		ulPage = ((ULONG)(ppmls->pMem)) >> WIN32PAGESIZE;
		ulNumPages = (((ppmls->ulNumBytes)-1) >> WIN32PAGESIZE ) +1;
		PGPmlPrint (("PGPmemlock: unlocking page(s) %x - %x.\n", 
						ulPage, ulPage+ulNumPages-1));

		// unlock the pages 
		ntStatus = PGPmlUnlockBlock (pDevO, pIrp, ulPage, ulNumPages, TRUE);
		if (ntStatus != STATUS_SUCCESS) 
		{
			// we couldn't unlock the memory
			PGPmlPrint (("PGPmemlock: Err: PGPmlUnlockBlock failed.\n"));
			ppmls->ulError = kPGPMemLockError_UnlockError;
		}
	}

	return ntStatus;
}


//----------------------------------------------------|
//  delete the associated device
//
//	Arguments:
//		pDrvO		- pointer to a driver object
//

VOID
PGPmlUnload(
    IN PDRIVER_OBJECT pDrvO)
{ 
    WCHAR           deviceLinkBuffer[]  = L"\\DosDevices\\PGPMEMLOCK";
    UNICODE_STRING	deviceLinkUnicodeString;
	PDEVICE_OBJECT	pDevO				= NULL;

	// get first (and only) device object
	pDevO = pDrvO->DeviceObject;

	// Unlock all blocks referenced in list and free list
	if (KeGetCurrentIrql () < DISPATCH_LEVEL) 
	{
		PGPmlUnlockAllBlocks (pDevO);
	}

    // Delete the symbolic link
    RtlInitUnicodeString (&deviceLinkUnicodeString,
                          deviceLinkBuffer);
    IoDeleteSymbolicLink (&deviceLinkUnicodeString);

    // Delete the device object
    PGPmlPrint (("PGPmemlock: unloading.\n"));
    IoDeleteDevice (pDrvO->DeviceObject);

	// Close down our item allocation system
	PGPmlHeapDestroy (); 
}


//----------------------------------------------------|
//  process the IRPs sent to this device
//
//	Arguments:
//		pDevO		- pointer to a device object
//		pIrp		- pointer to an I/O Request Packet
//
//	Return Value:
//		STATUS_SUCCESS if successful,
//		STATUS_* otherwise
//

NTSTATUS
PGPmlDispatch(
    IN PDEVICE_OBJECT pDevO,
    IN PIRP pIrp)
{
    PIO_STACK_LOCATION pIrpStack			= NULL;
    PVOID              ioBuffer				= NULL;
    ULONG              inBufferLength		= 0;
    ULONG              outBufferLength		= 0;
    ULONG              ioControlCode		= 0;
    NTSTATUS           ntStatus				= STATUS_SUCCESS;

    // Init to default settings
    pIrp->IoStatus.Status      = STATUS_SUCCESS;
    pIrp->IoStatus.Information = 0;

    // Get a pointer to the current location in the Irp. This is where
    //     the function codes and parameters are located.
    pIrpStack = IoGetCurrentIrpStackLocation (pIrp);

    // Get the pointer to the input/output buffer and it's length
    ioBuffer        = 
				pIrp->AssociatedIrp.SystemBuffer;
    inBufferLength  = 
				pIrpStack->Parameters.DeviceIoControl.InputBufferLength;
    outBufferLength = 
				pIrpStack->Parameters.DeviceIoControl.OutputBufferLength;

    switch (pIrpStack->MajorFunction)
    {
	// called when user app calls CreateFile
    case IRP_MJ_CREATE:
        PGPmlPrint (("PGPmemlock: IRP_MJ_CREATE\n"));
        break;

	// called when user app closes handle to driver --
	//	make sure all locked memory belonging to this handle gets unlocked
    case IRP_MJ_CLOSE:
		PGPmlPrint (("PGPmemlock: IRP_MJ_CLOSE\n"));
		if (KeGetCurrentIrql () < DISPATCH_LEVEL)
		{
			PGPmlUnlockHandleBlocks (pDevO, pIrp);
		}
		break;

	// called when user apps calls DeviceIoControl
    case IRP_MJ_DEVICE_CONTROL:
        ioControlCode = pIrpStack->Parameters.DeviceIoControl.IoControlCode;

        switch (ioControlCode)
        {
		// user app wants to lock some memory ...
        case IOCTL_PGPMEMLOCK_LOCK_MEMORY:

			// try locking it
            pIrp->IoStatus.Status = PGPmlLockTheMemory (pDevO,
														pIrp,
                                                        ioBuffer,
                                                        inBufferLength,
                                                        outBufferLength);

            if (NT_SUCCESS(pIrp->IoStatus.Status))
            {
                pIrp->IoStatus.Information = sizeof(PGPMEMLOCKSTRUCT);
                PGPmlPrint (("PGPmemlock: memory successfully locked.\n"));
            }
            else
            {
                PGPmlPrint (("PGPmemlock: Err: memory lock failed.\n"));
            }
            break;

		// user app wants to unlock some memory
        case IOCTL_PGPMEMLOCK_UNLOCK_MEMORY:

			// try unlocking it ...
            pIrp->IoStatus.Status = PGPmlUnlockTheMemory (pDevO,
														  pIrp,
                                         				  ioBuffer,
                                                          inBufferLength,
                                                          outBufferLength);

            if (NT_SUCCESS(pIrp->IoStatus.Status))
            {
               pIrp->IoStatus.Information = sizeof(PGPMEMLOCKSTRUCT);
               PGPmlPrint (("PGPmemlock: memory successfully unlocked.\n"));
            }
            else
            {
                PGPmlPrint (("PGPmemlock: Err: memory unlock failed.\n"));
            }
            break;

        default:
            pIrp->IoStatus.Status = STATUS_INVALID_PARAMETER;
            PGPmlPrint (
					("PGPmemlock: Err: unknown IRP_MJ_DEVICE_CONTROL.\n"));
            break;
        }

        break;
    }

    ntStatus = pIrp->IoStatus.Status;
    IoCompleteRequest(pIrp, IO_NO_INCREMENT);

    return ntStatus;
}


//----------------------------------------------------|
//  called upon initialization of driver
//
//	Routine Description:
//
//	Installable driver initialization entry point.
//	This entry point is called directly by the I/O system.
//
//	Arguments:
//		pDrvO		 -	pointer to the driver object
//		RegistryPath -	pointer to a unicode string representing the path
//						to driver-specific key in the registry
//
//	Return Value:
//		STATUS_SUCCESS if successful,
//		STATUS_* otherwise
//
	
NTSTATUS
DriverEntry(
    IN PDRIVER_OBJECT  pDrvO,
    IN PUNICODE_STRING RegistryPath)
{

    PDEVICE_OBJECT	pDevO				= NULL;
    WCHAR			deviceNameBuffer[]	= L"\\Device\\PGPmemlock";
    UNICODE_STRING	deviceNameUnicodeString;
    WCHAR			deviceLinkBuffer[]	= L"\\DosDevices\\PGPMEMLOCK";
    UNICODE_STRING	deviceLinkUnicodeString;
    NTSTATUS		ntStatus			= STATUS_SUCCESS;
	PDEVEXTENSION	pde					= NULL;

    PGPmlPrint (("PGPmemlock: entering DriverEntry.\n"));

	// Initialize internal list item allocation routines
	ntStatus = PGPmlHeapInit ();
	if (!NT_SUCCESS (ntStatus)) 
	{
		PGPmlPrint (("PGPmemlock: Err: PGPmlHeapInit failed.\n"));
		return ntStatus;
	}

    // Create an NON exclusive device object (multiple threads
    // can make requests to this device at a time)
    RtlInitUnicodeString (&deviceNameUnicodeString,
                          deviceNameBuffer);

	ntStatus = IoCreateDevice (pDrvO,
                               sizeof(DEVEXTENSION),
                               &deviceNameUnicodeString,
                               FILE_DEVICE_PGPMEMLOCK,
                               0,
                               FALSE,
                               &pDevO);

    if (NT_SUCCESS(ntStatus))
    {
        // Create dispatch points for device control, create, close.
        pDrvO->MajorFunction[IRP_MJ_CREATE]         =
        pDrvO->MajorFunction[IRP_MJ_CLOSE]          =
        pDrvO->MajorFunction[IRP_MJ_DEVICE_CONTROL] = PGPmlDispatch;
        pDrvO->DriverUnload                         = PGPmlUnload;

        // Create a symbolic link, e.g. a name that a Win32 app can specify
        // to open the device
        RtlInitUnicodeString(&deviceLinkUnicodeString,
                             deviceLinkBuffer);
        ntStatus = IoCreateSymbolicLink (&deviceLinkUnicodeString,
                                         &deviceNameUnicodeString);

        if (!NT_SUCCESS(ntStatus))
        {
            // Symbolic link creation failed- note this & then delete the
            // device object (it's useless if a Win32 app can't get at it).
            PGPmlPrint (("PGPmemlock: Err: IoCreateSymbolicLink failed.\n"));
            IoDeleteDevice (pDevO);
        }
		else 
		{
			// initialize the device extension data structure
			pde = (PDEVEXTENSION)(pDevO->DeviceExtension);
			memset (pde, 0, sizeof(DEVEXTENSION));
			KeInitializeSpinLock (&pde->slListLock);
		}
    }
    else
    {
        PGPmlPrint (("PGPmemlock: Err: IoCreateDevice failed.\n"));
    }

    return ntStatus;
}