#include "usbloader.h"

NTSTATUS PnpHandleStartDevice(IN PDEVICE_OBJECT fdo,IN PIRP Irp);
NTSTATUS ForwardAndWait(IN PDEVICE_OBJECT fdo,IN PIRP Irp);
NTSTATUS ForwardIrpComplete(IN PDEVICE_OBJECT fdo,
						    IN PIRP Irp,
                            IN PKEVENT pev);
NTSTATUS UsbStartDevice(IN PDEVICE_OBJECT fdo);
NTSTATUS PnpHandleDefault(IN PDEVICE_OBJECT fdo,IN PIRP Irp);
NTSTATUS PnpHandleRemoveDevice(IN PDEVICE_OBJECT fdo,IN PIRP Irp);
NTSTATUS UsbRemoveDevice(IN PDEVICE_OBJECT fdo);
NTSTATUS Ezusb_8051Reset(PDEVICE_OBJECT fdo,UCHAR resetBit);
NTSTATUS Ezusb_DownloadIntelHex(PDEVICE_OBJECT fdo,PINTEL_HEX_RECORD hexRecord);


extern INTEL_HEX_RECORD firmware[];
extern INTEL_HEX_RECORD loader[];

NTSTATUS TestAddDevice(IN PDRIVER_OBJECT DriverObject,
                       IN PDEVICE_OBJECT PhysicalDeviceObject)
{
   NTSTATUS ntStatus = STATUS_SUCCESS;
   PDEVICE_OBJECT fdo = NULL;
   PDEVICE_EXTENSION pdx;

   ntStatus = IoCreateDevice (DriverObject,
                              sizeof (DEVICE_EXTENSION),
                              NULL,  
                              FILE_DEVICE_UNKNOWN,
                              0,
                              FALSE,
                              &fdo);

   if (NT_SUCCESS(ntStatus))
   {
      pdx = fdo->DeviceExtension;

      fdo->Flags &= ~DO_DEVICE_INITIALIZING;

      fdo->Flags |= DO_DIRECT_IO;

      pdx->PhysicalDeviceObject=PhysicalDeviceObject;

      pdx->LowerDeviceObject =
         IoAttachDeviceToDeviceStack(fdo, PhysicalDeviceObject);

	   pdx->Usages = 1;		
	   KeInitializeEvent(&pdx->evRemove,
                        NotificationEvent,
                        FALSE);  
   }

   return ntStatus;		
}

NTSTATUS TestPnpIrp(IN PDEVICE_OBJECT fdo, IN PIRP Irp)
{
	NTSTATUS ntStatus=STATUS_SUCCESS;
	PIO_STACK_LOCATION IrpStack;
    PDEVICE_EXTENSION pdx = fdo->DeviceExtension;
    ULONG MinorFunction;

    if (!LockDevice(fdo))
		return CompleteRequest(Irp, STATUS_DELETE_PENDING, 0);

    IrpStack = IoGetCurrentIrpStackLocation (Irp);
    MinorFunction = IrpStack->MinorFunction;

    switch (MinorFunction)
		{
		case IRP_MN_START_DEVICE:	
			ntStatus = PnpHandleStartDevice(fdo,Irp);
			break; 

         case IRP_MN_REMOVE_DEVICE:
		  ntStatus = PnpHandleRemoveDevice(fdo,Irp);
          break;

        case IRP_MN_QUERY_CAPABILITIES:
		  {
			  PDEVICE_CAPABILITIES pdc = IrpStack->Parameters.DeviceCapabilities.Capabilities;
		      if (pdc->Version < 1) {
				  ntStatus = PnpHandleDefault(fdo, Irp);
                  break;
			  }
              ntStatus = ForwardAndWait(fdo, Irp);
			  if (NT_SUCCESS(ntStatus)) { 
				  pdc = IrpStack->Parameters.DeviceCapabilities.Capabilities;
           	      pdc->SurpriseRemovalOK = TRUE;
			  }
   		      ntStatus = CompleteRequest(Irp, ntStatus, Irp->IoStatus.Information);
		  }
          break; 
       default:
		 ntStatus = PnpHandleDefault(fdo, Irp);
   } 
   if (MinorFunction != IRP_MN_REMOVE_DEVICE)
	   UnlockDevice(fdo);

   return ntStatus;
}

BOOLEAN LockDevice(IN PDEVICE_OBJECT fdo)
{
   PDEVICE_EXTENSION pdx = (PDEVICE_EXTENSION) fdo->DeviceExtension;

   LONG usage = InterlockedIncrement(&pdx->Usages);

   if (pdx->Removing) {
	   if (InterlockedDecrement(&pdx->Usages) == 0)
		   KeSetEvent(&pdx->evRemove, 0, FALSE);
	   return FALSE;
   }
   return TRUE;
}

void UnlockDevice(PDEVICE_OBJECT fdo)
{
	PDEVICE_EXTENSION pdx;
	LONG usage;

	pdx = (PDEVICE_EXTENSION) fdo->DeviceExtension;
    usage = InterlockedDecrement(&pdx->Usages);

    if (usage == 0) {		
		KeSetEvent(&pdx->evRemove, 0, FALSE);
   }						
}

NTSTATUS CompleteRequest(IN PIRP Irp,
						 IN NTSTATUS status,
						 IN ULONG info)
{
	Irp->IoStatus.Status = status;
	Irp->IoStatus.Information = info;
	IoCompleteRequest(Irp, IO_NO_INCREMENT);
	
	return status;
}

NTSTATUS PnpHandleStartDevice(IN PDEVICE_OBJECT fdo,IN PIRP Irp)
{
	NTSTATUS ntStatus;
	
    ntStatus = ForwardAndWait(fdo, Irp);
	if (!NT_SUCCESS(ntStatus))
		return CompleteRequest(Irp, ntStatus, Irp->IoStatus.Information);
    ntStatus = UsbStartDevice(fdo);
    
	return CompleteRequest(Irp, ntStatus, 0);
}

NTSTATUS ForwardAndWait(IN PDEVICE_OBJECT fdo,IN PIRP Irp)
{
	NTSTATUS ntStatus;
	PDEVICE_EXTENSION pdx = (PDEVICE_EXTENSION) fdo->DeviceExtension;
	KEVENT event;

	KeInitializeEvent(&event, NotificationEvent, FALSE);

	IoCopyCurrentIrpStackLocationToNext(Irp);
	IoSetCompletionRoutine(Irp, (PIO_COMPLETION_ROUTINE)ForwardIrpComplete,
		                  (PVOID) &event, TRUE, TRUE, TRUE);

	ntStatus = IoCallDriver(pdx->LowerDeviceObject, Irp);

	if (ntStatus == STATUS_PENDING)	{
		KeWaitForSingleObject(&event, Executive, KernelMode, FALSE, NULL);
        ntStatus = Irp->IoStatus.Status;
    }
	return ntStatus;
}

NTSTATUS ForwardIrpComplete(IN PDEVICE_OBJECT fdo,
						    IN PIRP Irp,
                            IN PKEVENT pev)
{
   KeSetEvent(pev, 0, FALSE);
   return STATUS_MORE_PROCESSING_REQUIRED;
}

NTSTATUS UsbStartDevice(IN PDEVICE_OBJECT fdo)
{
   Ezusb_8051Reset(fdo,1);
   Ezusb_DownloadIntelHex(fdo,loader);
   Ezusb_8051Reset(fdo,0);

   Ezusb_DownloadIntelHex(fdo,firmware);
   Ezusb_8051Reset(fdo,1);
   Ezusb_8051Reset(fdo,0);

   return STATUS_SUCCESS;
}

NTSTATUS Ezusb_8051Reset(PDEVICE_OBJECT fdo,UCHAR resetBit)
{
   NTSTATUS ntStatus;
   PURB urb = NULL;

   urb = ExAllocatePool(NonPagedPool, 
                       sizeof(struct _URB_CONTROL_VENDOR_OR_CLASS_REQUEST));

   if (urb)
   {
      RtlZeroMemory(urb,sizeof(struct  _URB_CONTROL_VENDOR_OR_CLASS_REQUEST));

      urb->UrbHeader.Length = sizeof(struct _URB_CONTROL_VENDOR_OR_CLASS_REQUEST);
      urb->UrbHeader.Function = URB_FUNCTION_VENDOR_DEVICE;

      urb->UrbControlVendorClassRequest.TransferBufferLength = 1;
      urb->UrbControlVendorClassRequest.TransferBuffer = &resetBit;
      urb->UrbControlVendorClassRequest.TransferBufferMDL = NULL;
      urb->UrbControlVendorClassRequest.Request = ANCHOR_LOAD_INTERNAL;
      urb->UrbControlVendorClassRequest.Value = CPUCS_REG;
      urb->UrbControlVendorClassRequest.Index = 0;

      ntStatus = UsbCallUSBDI(fdo, urb);
   }
   else
   {
      ntStatus = STATUS_NO_MEMORY;
   }

   if (urb)
      ExFreePool(urb);

   return ntStatus;
}

NTSTATUS Ezusb_DownloadIntelHex(PDEVICE_OBJECT fdo,PINTEL_HEX_RECORD hexRecord)
{
   NTSTATUS ntStatus;
   PURB urb = NULL;
   PINTEL_HEX_RECORD ptr = hexRecord;

   urb = ExAllocatePool(NonPagedPool,
                       sizeof(struct _URB_CONTROL_VENDOR_OR_CLASS_REQUEST));

   if (urb)
   {
      while (ptr->Type == 0)
      {
         if (!INTERNAL_RAM(ptr->Address))
         {
            RtlZeroMemory(urb,sizeof(struct  _URB_CONTROL_VENDOR_OR_CLASS_REQUEST));

            urb->UrbHeader.Length = sizeof(struct _URB_CONTROL_VENDOR_OR_CLASS_REQUEST);
            urb->UrbHeader.Function = URB_FUNCTION_VENDOR_DEVICE;
            urb->UrbControlVendorClassRequest.TransferBufferLength = ptr->Length;
            urb->UrbControlVendorClassRequest.TransferBuffer = ptr->Data;
            urb->UrbControlVendorClassRequest.Request = ANCHOR_LOAD_EXTERNAL;
            urb->UrbControlVendorClassRequest.Value = ptr->Address;
            urb->UrbControlVendorClassRequest.Index = 0;

            ntStatus = UsbCallUSBDI(fdo, urb);

            if (!NT_SUCCESS(ntStatus))
               break;
         }
         ptr++;
      }

      Ezusb_8051Reset(fdo,1);
      ptr = hexRecord;
      while (ptr->Type == 0)
      {
         if (INTERNAL_RAM(ptr->Address))
         {
            RtlZeroMemory(urb,sizeof(struct  _URB_CONTROL_VENDOR_OR_CLASS_REQUEST));

            urb->UrbHeader.Length = sizeof(struct _URB_CONTROL_VENDOR_OR_CLASS_REQUEST);
            urb->UrbHeader.Function = URB_FUNCTION_VENDOR_DEVICE;
            urb->UrbControlVendorClassRequest.TransferBufferLength = ptr->Length;
            urb->UrbControlVendorClassRequest.TransferBuffer = ptr->Data;
            urb->UrbControlVendorClassRequest.Request = ANCHOR_LOAD_INTERNAL;
            urb->UrbControlVendorClassRequest.Value = ptr->Address;
            urb->UrbControlVendorClassRequest.Index = 0;

            ntStatus = UsbCallUSBDI(fdo, urb);

            if (!NT_SUCCESS(ntStatus))
               break;
         }
         ptr++;
      }

   }
   else
   {
      ntStatus = STATUS_NO_MEMORY;
   }

   if (urb)
      ExFreePool(urb);

   return ntStatus;
}

NTSTATUS UsbCallUSBDI(IN PDEVICE_OBJECT fdo,IN PURB Urb)
{
    
    NTSTATUS ntStatus, status = STATUS_SUCCESS;
    PDEVICE_EXTENSION pdx;
    PIRP irp;
    KEVENT event;
    IO_STATUS_BLOCK ioStatus;
    PIO_STACK_LOCATION nextStack;

    pdx = fdo->DeviceExtension;

    KeInitializeEvent(&event, NotificationEvent, FALSE);
    irp = IoBuildDeviceIoControlRequest(
                                        IOCTL_INTERNAL_USB_SUBMIT_URB,
                                        pdx->LowerDeviceObject,
                                        NULL,
                                        0,
                                        NULL,
                                        0,
                                        TRUE, 
                                        &event,
                                        &ioStatus);

    nextStack = IoGetNextIrpStackLocation(irp);

    nextStack->Parameters.Others.Argument1 = Urb;

    ntStatus = IoCallDriver(pdx->LowerDeviceObject,irp);
    if (ntStatus == STATUS_PENDING) 
		KeWaitForSingleObject(&event,Suspended,KernelMode,FALSE,NULL);
	else
		ioStatus.Status = ntStatus;
	
	ntStatus=ioStatus.Status;
        
    return ntStatus;
}

NTSTATUS PnpHandleDefault(IN PDEVICE_OBJECT fdo,IN PIRP Irp)
{
   PDEVICE_EXTENSION pdx;
   pdx = (PDEVICE_EXTENSION) fdo->DeviceExtension;
   IoSkipCurrentIrpStackLocation(Irp);
   return IoCallDriver(pdx->LowerDeviceObject, Irp);
}

NTSTATUS PnpHandleRemoveDevice(IN PDEVICE_OBJECT fdo,IN PIRP Irp)
{
	NTSTATUS ntStatus;
    PDEVICE_EXTENSION pdx;
	
	pdx = (PDEVICE_EXTENSION) fdo->DeviceExtension;
	pdx->Removing = TRUE;
	UnlockDevice(fdo);	
	UnlockDevice(fdo);	
	KeWaitForSingleObject(&pdx->evRemove, Executive, KernelMode, FALSE, NULL);

	UsbRemoveDevice(fdo);
    ntStatus = PnpHandleDefault(fdo, Irp);
    return ntStatus;			
}

NTSTATUS UsbRemoveDevice(IN  PDEVICE_OBJECT fdo)
{
	NTSTATUS ntStatus = STATUS_SUCCESS;
	PDEVICE_EXTENSION pdx;

    pdx = fdo->DeviceExtension;
    IoDetachDevice(pdx->LowerDeviceObject);
    IoDeleteDevice (fdo);
    return ntStatus;
}