dispatch.cpp

USB转comDriver的驱动

    //UCHAR * pBuf = (UCHAR*)ExAllocatePool(NonPagedPool, (writeLen+1) * sizeof(UCHAR));
    UCHAR * pBuf = dx->pTxArray;
    RtlZeroMemory(pBuf, (TX_BUFFER_SIZE+1) * sizeof(UCHAR));
    
    if (!pBuf)
		return Vcp4usbCompleteIrp(Irp, STATUS_INSUFFICIENT_RESOURCES, 0);

    KIRQL oldIrql;
    KeAcquireSpinLock(&dx->spinLock, &oldIrql);
	LARGE_INTEGER totalTimeout;
        
    IoAcquireCancelSpinLock(&Irp->CancelIrql);
    if (Irp->Cancel) // need cancel
    {
        IoReleaseCancelSpinLock(Irp->CancelIrql);
        Irp->IoStatus.Status = STATUS_CANCELLED;
        Irp->IoStatus.Information = 0;
        KeReleaseSpinLock(&dx->spinLock, oldIrql);
        DPrint(DBG_WRITE, ("write return STATUS_CANCELLED @ %d\n", KeGetCurrentIrql()));
        
        KeReleaseSpinLock(&dx->spinLock, oldIrql);
        return STATUS_CANCELLED;
    }
    IoSetCancelRoutine(Irp, CancelRoutineWrite);
    IoReleaseCancelSpinLock(Irp->CancelIrql);
    
    IoMarkIrpPending(Irp);
    Irp->IoStatus.Status = STATUS_PENDING;
    Irp->IoStatus.Information = 0;

    RtlCopyMemory(pBuf, Irp->AssociatedIrp.SystemBuffer, writeLen * sizeof(UCHAR));
    dx->irpWritePending = Irp;
    dx->curLenTx = writeLen;
    dx->pTxBuf = pBuf;
    
    // start write timer
    if (!CalculateWriteTimeout(dx, writeLen * 10, &totalTimeout))
    {
        if (totalTimeout.QuadPart != 0)
        {
            DPrint(DBG_WRITE, ("start write timer totalTimeout\n"));
            StartTimer(&dx->writeTotalTimer, totalTimeout);
        }
    }
    
    KeSetEvent(&dx->eventWriteStart, 1, FALSE);
    DPrintHex(DBG_DATALOG, ("Stage write"), dx->pTxBuf, dx->curLenTx);
    
    KeReleaseSpinLock(&dx->spinLock, oldIrql);

    DPrint(DBG_WRITE, ("Vcp4usbWrite() return STATUS_PENDING @ %d\n", KeGetCurrentIrql()));
    return STATUS_PENDING;
}

#define Free_Res() \
    do{ \
    if (hRead){ZwClose(hRead);hRead = NULL;} \
    if (hWrite){ZwClose(hWrite);hWrite = NULL;} \
    }while(0)
    //if (dx->pRxBuf){ExFreePool(dx->pRxBuf);dx->pRxBuf = NULL;}
    //if (dx->pRxBuf){dx->pRxBuf = NULL;}

#define TIMEOUT_1MS (-1i64 * 1000i64 * 10i64)
#define TOTALEVENT 3

VOID Vcp4usbThread(IN PVOID Context)
{
    PDEVICE_OBJECT              fdo;
    PVCP4USB_DEVICE_EXTENSION   dx;
    LARGE_INTEGER               timeout;
    PVOID                       aryEvent[TOTALEVENT];
    NTSTATUS                    status;
    HANDLE                      hRead = NULL, hWrite = NULL;
    BOOLEAN                     kill = FALSE, bNeedReadData = FALSE;
    
    fdo = (PDEVICE_OBJECT)Context;
    dx = (PVCP4USB_DEVICE_EXTENSION)fdo->DeviceExtension;
    timeout.QuadPart = TIMEOUT_1MS; // 1ms timeout
    
    DPrint(DBG_THREAD, ("Vcp4usbThread() start, dx=0x%x @ %d\n", dx, KeGetCurrentIrql()));
    
    KeSetPriorityThread(KeGetCurrentThread(), LOW_REALTIME_PRIORITY);
    aryEvent[0] = &dx->eventThreadExiting;
    aryEvent[1] = &dx->eventWriteStart;
    aryEvent[2] = &dx->eventReadEnd;
    
    while (!kill)
    {
        // wait event
        status = KeWaitForMultipleObjects(TOTALEVENT, aryEvent, WaitAny, 
            Executive, KernelMode, FALSE, &timeout, NULL);
        //DPrint(DBG_THREAD, ("wait object return:0x%x\n", status));
        
        // clear event
        if (status == STATUS_TIMEOUT) // check connect status
        {
            //DPrint(DBG_THREAD, ("time out\n"));
        }
        else if (status == STATUS_WAIT_0) // eventThreadExiting
        {
            kill = TRUE;
            KeResetEvent((PRKEVENT)aryEvent[0]);
            break; // or continue;
        }
        else if (status == STATUS_WAIT_1) // eventWriteStart
        {
            KeResetEvent((PRKEVENT)aryEvent[1]);
        }
        else if (status == STATUS_WAIT_2) // eventReadEnd
        {
            KeResetEvent((PRKEVENT)aryEvent[2]);
        }
        //else if (status == STATUS_WAIT_3) // eventReadStart
        //{
        //    KeResetEvent((PRKEVENT)aryEvent[3]);
        //}
        else
        {
            DPrint(DBG_THREAD, ("Vcp4usbThread() not handled status=%d\n", status));
			// if (DBG) DbgBreakPoint();
        }
        
        // check usb device connect when eventWriteStart & eventReadEnd, don't change var (status)
        if ((status == STATUS_TIMEOUT) || (status == STATUS_WAIT_1) || (status == STATUS_WAIT_2))
        {
			NTSTATUS curStatus;
            BOOLEAN bCurConnect = FALSE;
            UNICODE_STRING usbDeviceName;
            curStatus = GetUsbDeviceName(&usbDeviceName, &GUID_CLASS_XXXX_BULK,
                &XXXX_SYMLINK_CMPSTR, XXXX_SYMLINK_STRLEN, 0);
            if (NT_SUCCESS(curStatus))
            {
                HANDLE hDev;
                if (NT_SUCCESS(createFile(&hDev, &usbDeviceName, L"", 0, GENERIC_READ)))
                {
                    ZwClose(hDev);
                    bCurConnect = TRUE;
                }
            }
            
            if (!dx->bIsConnected && bCurConnect)
            {
                NTSTATUS rdStatus, wtStatus;
                DPrint(DBG_THREAD, ("Vcp4usbThread() usb device connected.\n"));
                if (!hRead)
                    rdStatus = createFile(&hRead, &usbDeviceName, L"\\PIPE0", 6, GENERIC_READ);
                if (!hWrite)
                    wtStatus = createFile(&hWrite, &usbDeviceName, L"\\PIPE1", 6, GENERIC_WRITE);
                if (NT_SUCCESS(rdStatus) && NT_SUCCESS(wtStatus))
                {
                    if (!dx->pRxBuf)
                        //dx->pRxBuf = (UCHAR*)ExAllocatePool(NonPagedPool, RX_BUFFER_SIZE + 1);
                        dx->pRxBuf = dx->pRxArray;
                    if (dx->pRxBuf)
                        RtlZeroMemory(dx->pRxBuf, RX_BUFFER_SIZE + 1);
                }
                if (!NT_SUCCESS(rdStatus) || !NT_SUCCESS(wtStatus) || (!dx->pRxBuf))
                {
                    DPrint(DBG_THREAD, ("Vcp4usbThread() usb pipe can't opened.\n"));
                    Free_Res();
                }
                else
                {
                    DPrint(DBG_THREAD, ("Vcp4usbThread() dx->pRxBuf = 0x%x\n", dx->pRxBuf));
                    dx->curLenRx = 0;
                    dx->bReadPending = FALSE;
                    dx->curLenTx = 0;
                    dx->bWritePending = FALSE;
                    dx->bIsConnected = TRUE;
                }
            }
            else if (dx->bIsConnected && !bCurConnect)
            {
                dx->bIsConnected = FALSE;
                DPrint(DBG_THREAD, ("Vcp4usbThread() usb device disconnected.\n"));
                Free_Res();
            }
            
            if (NT_SUCCESS(curStatus))
            {
                if (usbDeviceName.Buffer)
                    ExFreePool(usbDeviceName.Buffer);
            }
        }
        
        // write start
        if (status == STATUS_WAIT_1)
        {
            if (dx->bIsConnected && hWrite && !dx->bWritePending && dx->pTxBuf && dx->curLenTx)
            {
                KIRQL oldIrql;
                KeAcquireSpinLock(&dx->spinLock, &oldIrql);
            
                IO_STATUS_BLOCK ioBlock;
                LARGE_INTEGER byteOffset;
                ULONG key = 0;
            
                byteOffset.QuadPart = 0;
                ioBlock.Information = 0;
                ioBlock.Status = 0;
            
                dx->bWritePending = TRUE;
            
                DPrint(DBG_WRITE, ("Vcp4usbThread() write start\n"));
                ZwWriteFile(hWrite, NULL, (PIO_APC_ROUTINE)ApcRoutineWrite, dx,
                    &ioBlock, dx->pTxBuf, dx->curLenTx, &byteOffset, NULL);//&key); // zero byte
                DPrint(DBG_WRITE, ("Vcp4usbThread() write ready\n"));
            
                KeReleaseSpinLock(&dx->spinLock, oldIrql);
            }
        } // if (status == STATUS_WAIT_1) // write start
        
		// userapp maybe wait mask
		if ((dx->pWaitCtrlIrp != NULL) && (dx->eventMask & SERIAL_EV_RXCHAR))
			bNeedReadData = TRUE;
		
		// read start
        if ((status == STATUS_WAIT_2) || (bNeedReadData == TRUE))
        {
            if ((dx->bIsConnected && hRead && !dx->bReadPending && dx->pRxBuf && (dx->curLenRx <= RX_BUFFER_SIZE)))
            {
                KIRQL oldIrql;
                KeAcquireSpinLock(&dx->spinLock, &oldIrql);
                
                IO_STATUS_BLOCK ioBlock;
                LARGE_INTEGER byteOffset;
                ULONG key = 0;
            
                byteOffset.QuadPart = 0;
                ioBlock.Information = 0;
                ioBlock.Status = 0;
                
                dx->bReadPending = TRUE;
                PUCHAR pCurBuf = dx->pRxBuf;
                ULONG needReadLen = RX_BUFFER_SIZE - dx->curLenRx;
				needReadLen = (needReadLen > 832) ? 832 : needReadLen;
				
				pCurBuf += dx->curLenRx;
                
				DPrint(DBG_READ, ("Vcp4usbThread() read start:needReadLen=%d\n", needReadLen));
                ZwReadFile(hRead, NULL, (PIO_APC_ROUTINE)ApcRoutineRead, dx,
                    &ioBlock, pCurBuf, needReadLen, &byteOffset, NULL);//&key); // zero byte
                DPrint(DBG_READ, ("Vcp4usbThread() read ready\n"));
            
                KeReleaseSpinLock(&dx->spinLock, oldIrql);
            }
        } // if (status == STATUS_WAIT_2) // read start
    } // end while (TRUE)
    
    if (dx->irpReadPending)
    {
        IoCancelIrp(dx->irpReadPending);
        dx->irpReadPending = NULL;
    }
    if (dx->irpWritePending)
    {
        IoCancelIrp(dx->irpWritePending);
        dx->irpWritePending = NULL;
    }
    
    if (dx->pThreadObj != NULL)
    {
        ObDereferenceObject(&dx->pThreadObj);
        dx->pThreadObj = NULL;
    }
    
    Free_Res();
    dx->bReadPending = FALSE;
    dx->bWritePending = FALSE;
    dx->bIsConnected = FALSE;
    
    DPrint(DBG_THREAD, ("Vcp4usbThread exit\n"));
    KeSetEvent(&dx->eventThreadExit, 0, TRUE);
    
    PsTerminateSystemThread(STATUS_SUCCESS);
}

NTSTATUS Vcp4usbCreate(IN PDEVICE_OBJECT fdo, IN PIRP Irp)
{
    PVCP4USB_DEVICE_EXTENSION dx=(PVCP4USB_DEVICE_EXTENSION)fdo->DeviceExtension;
    PIO_STACK_LOCATION IrpStack = IoGetCurrentIrpStackLocation(Irp);
    NTSTATUS status = STATUS_SUCCESS;
    
    DPrint(DBG_CREATE, ("Vcp4usbCreate() @ %d\n", KeGetCurrentIrql()));
    
    Irp->IoStatus.Information = 0;
    
    // Create thread
    if (dx->pThreadObj != NULL)
    {
        status = STATUS_UNSUCCESSFUL;
    }
    else
    {
        HANDLE hThread;
        OBJECT_ATTRIBUTES oa;
        InitializeObjectAttributes(&oa, NULL, OBJ_KERNEL_HANDLE, NULL, NULL);
        status = PsCreateSystemThread(&hThread, 
            THREAD_ALL_ACCESS, &oa, NULL, NULL, Vcp4usbThread, fdo);
        if (!NT_SUCCESS(status))
        {
            DPrint(DBG_CREATE, ("CreateThread error:0x%x\n", status));
            return Vcp4usbCompleteIrp(Irp, status, Irp->IoStatus.Information);
        }
        status = ObReferenceObjectByHandle(hThread, 
            THREAD_ALL_ACCESS, NULL, KernelMode, &dx->pThreadObj, NULL);
        if (!NT_SUCCESS(status) || !dx->pThreadObj)
        {
            status = STATUS_UNSUCCESSFUL;
        }
        if (hThread)
            ZwClose(hThread);
    }
    
    DPrint(DBG_CREATE, ("Vcp4usbCreate() exit, return 0x%x\n", status));
    
    // Complete successfully
    return Vcp4usbCompleteIrp(Irp, status, Irp->IoStatus.Information);
}