// Read/Write request processors for usbiso driver
// Copyright (C) 1999 by Walter Oney
// All rights reserved

#include "stddcls.h"
#include "driver.h"

#if DBG
	#define MSGUSBSTRING(d,s,i) { \
		UNICODE_STRING sd; \
		if (i && NT_SUCCESS(GetStringDescriptor(d,i,&sd))) { \
			DbgPrint(s, sd.Buffer); \
			RtlFreeUnicodeString(&sd); \
		}}
#else
	#define MSGUSBSTRING(d,i,s)
#endif

typedef struct _RWCONTEXT {
	LIST_ENTRY list;			// list for queuing context structures
	PDEVICE_EXTENSION pdx;		// our device extension address
	PIRP mainirp;				// the main r/w IRP
	NTSTATUS status;			// ending status for main IRP
	ULONG numxfer;				// total number of bytes transferred
	ULONG numirps;				// total number of sub-irps
	LONG numpending;			// number of sub-irps not yet finished
	LONG refcnt;				// reference count (initially 2)
	struct {
		PIRP irp;				// subsidiary IRP
		PURB urb;				// URB packaged with sub-IRP
		PMDL mdl;				// partial MDL for the URB
		} sub[1];
	} RWCONTEXT, *PRWCONTEXT;

BOOLEAN DestroyContextStructure(PRWCONTEXT ctx);
ULONG GetCurrentFrame(PDEVICE_EXTENSION pdx);
VOID OnCancelReadWrite(PDEVICE_OBJECT fdo, PIRP Irp);
NTSTATUS OnReadWriteComplete(PDEVICE_OBJECT fdo, PIRP Irp, PRWCONTEXT context);
NTSTATUS OnStageComplete(PDEVICE_OBJECT fdo, PIRP Irp, PRWCONTEXT context);

NTSTATUS SelectAlternateInterface(PDEVICE_OBJECT fdo);
NTSTATUS SelectDefaultInterface(PDEVICE_OBJECT fdo);

///////////////////////////////////////////////////////////////////////////////

#pragma LOCKEDCODE

BOOLEAN DestroyContextStructure(PRWCONTEXT ctx)
	{							// DestroyContextStructure

	// Decrement the reference count. If it goes to zero, delete
	// all the subsidiary IRPs and the context structure itself.

	if (InterlockedDecrement(&ctx->refcnt) > 0)
		return FALSE;					// still potentially in use
	for (ULONG i = 0; i < ctx->numirps; ++i)
		if (ctx->sub[i].irp)
			IoFreeIrp(ctx->sub[i].irp);
	ExFreePool(ctx);
	return TRUE;
	}							// DestroyContextStructure

///////////////////////////////////////////////////////////////////////////////

#pragma PAGEDCODE

NTSTATUS DispatchCreate(PDEVICE_OBJECT fdo, PIRP Irp)
	{							// DispatchCreate
	PAGED_CODE();
	KdPrint((DRIVERNAME " - IRP_MJ_CREATE\n"));
	PDEVICE_EXTENSION pdx = (PDEVICE_EXTENSION) fdo->DeviceExtension;

	PIO_STACK_LOCATION stack = IoGetCurrentIrpStackLocation(Irp);

	NTSTATUS status = STATUS_SUCCESS;
	
	if (InterlockedIncrement(&pdx->handles) == 1)
		{					// no handles previously open
		status = SelectAlternateInterface(fdo);
		if (!NT_SUCCESS(status))
			{				// error selecting alternate interface
			InterlockedDecrement(&pdx->handles);
			IoReleaseRemoveLock(&pdx->RemoveLock, stack->FileObject);
			}				// error selecting alternate interface
		}					// no handles previously open

	return CompleteRequest(Irp, status, 0);
	}							// DispatchCreate

///////////////////////////////////////////////////////////////////////////////

#pragma PAGEDCODE

NTSTATUS DispatchClose(PDEVICE_OBJECT fdo, PIRP Irp)
	{							// DispatchClose
	PAGED_CODE();
	KdPrint((DRIVERNAME " - IRP_MJ_CLOSE\n"));
	PDEVICE_EXTENSION pdx = (PDEVICE_EXTENSION) fdo->DeviceExtension;
	PIO_STACK_LOCATION stack = IoGetCurrentIrpStackLocation(Irp);
	if (InterlockedDecrement(&pdx->handles) == 0)
		SelectDefaultInterface(fdo);
	return CompleteRequest(Irp, STATUS_SUCCESS, 0);
	}							// DispatchClose

///////////////////////////////////////////////////////////////////////////////

#pragma PAGEDCODE

NTSTATUS DispatchRead(PDEVICE_OBJECT fdo, PIRP Irp)
	{							// DispatchRead
	PAGED_CODE();
	PDEVICE_EXTENSION pdx = (PDEVICE_EXTENSION) fdo->DeviceExtension;

	ASSERT(pdx->hinpipe);

	NTSTATUS status = IoAcquireRemoveLock(&pdx->RemoveLock, Irp);
	if (!NT_SUCCESS(status))
		return CompleteRequest(Irp, status, 0);

	ULONG length = Irp->MdlAddress ?  MmGetMdlByteCount(Irp->MdlAddress) : 0;
	if (!length)
		{						// zero-length read
		IoReleaseRemoveLock(&pdx->RemoveLock, Irp);
		return CompleteRequest(Irp, STATUS_SUCCESS, 0);
		}						// zero-length read

	// Calculate the number of IRPs that we need to submit in order to fulfill the
	// request. One of the ingredients of the calculation is the "packet size" the
	// endpoint uses during each frame. We have to know this number a priori, because
	// it doesn't have a fixed relationship to the endpoint's transfer size. Another
	// ingredient of the calculation is the fact that one URB can accomodate only 255
	// packets. The last ingredient is the maximum transfer size we declared when we
	// first configured the pipe. This driver left the default value in place. Given
	// the value of 16 used here for the packet size, we'll end up with a "segsize" of
	// 4080 bytes in fact.

	ULONG packsize = 16;		// known because we write the firmware
	ULONG segsize = USBD_DEFAULT_MAXIMUM_TRANSFER_SIZE; // default we left in place
	if (segsize / packsize > 255)
		segsize = 255 * packsize;	// maximum based on # packets allowed
	ULONG numirps = (length + segsize - 1) / segsize;

	// Allocate a context structure for use in the main IRP's completion and cancel routines.
	// Build as many IRPs (and URBs) as are required to fulfill the request.

	ULONG ctxsize = sizeof(RWCONTEXT) + (numirps - 1) * sizeof(((PRWCONTEXT) 0)->sub);
	PRWCONTEXT ctx = (PRWCONTEXT) ExAllocatePool(NonPagedPool, ctxsize);
	if (!ctx)
		{
		KdPrint((DRIVERNAME	" - unable to allocate %d bytes for record keeping\n", ctxsize));
		IoReleaseRemoveLock(&pdx->RemoveLock, Irp);
		return CompleteRequest(Irp, STATUS_INSUFFICIENT_RESOURCES, 0);
		}
	RtlZeroMemory(ctx, ctxsize);

	ctx->numirps = ctx->numpending = numirps;
	ctx->pdx = pdx;
	ctx->mainirp = Irp;
	ctx->refcnt = 2;			// one for OnReadWriteComplete, one forOnCancelReadWrite
	Irp->Tail.Overlay.DriverContext[0] = (PVOID) ctx; // for cancel routine

	ULONG frame = GetCurrentFrame(pdx) + 2;		// start first URB 2 frames from now
	KdPrint((DRIVERNAME " - First transfer scheduled for frame %8.8lX\n", frame));

	ULONG i;
	CCHAR stacksize = pdx->LowerDeviceObject->StackSize + 1; // 1 extra for OnStageComplete
	PUCHAR va = (PUCHAR) MmGetMdlVirtualAddress(Irp->MdlAddress);

	for (i = 0; i < numirps; ++i)
		{						// for each IRP

		// Create a subsidiary IRP

		PIRP subirp = IoAllocateIrp(stacksize, FALSE);
		if (!subirp)
			break;				// can't create subsidiary IRP
		ctx->sub[i].irp = subirp;

		// Determine how many packets are required and allocate an URB big enough for all of them

		if (segsize > length)
			segsize = length;	// last transfer is short

		ULONG npackets = (segsize + packsize - 1) / packsize; // # packets needed in this stage
		ASSERT(npackets <= 255);	// else we didn't calculate segsize correctly
		ULONG size = GET_ISO_URB_SIZE(npackets); // size of URB needed for that many packets
		PURB urb = (PURB) ExAllocatePool(NonPagedPool, size);
		if (!urb)
			break;

		ctx->sub[i].urb = urb;

		// Create a partial MDL to map this segment. This is necessary because the offset fields in
		// the packet descriptors we'll build must begin at zero for each URB. (It would have been better
		// if they were just relative to whatever MDL in in the URB.)

		PMDL mdl = IoAllocateMdl((PVOID) va, segsize, FALSE, FALSE, NULL);
		if (!mdl)
			break;				// can't allocate memory for MDL
		IoBuildPartialMdl(Irp->MdlAddress, mdl, (PVOID) va, segsize);
		ctx->sub[i].mdl = mdl;
		va += segsize;			// for next iteration

		// Initialize the URB. Use a specific frame number rather than specifying the ASAP
		// flag. Much to my surprise, the bus driver will schedule a second ASAP transfer immediately
		// instead of waiting for current transfers to finish. This is contrary to what the DDK
		// doc says and is, I believe, a bug in XP.

		RtlZeroMemory(urb, size);
		urb->UrbIsochronousTransfer.Hdr.Length = (USHORT) size;
		urb->UrbIsochronousTransfer.Hdr.Function = URB_FUNCTION_ISOCH_TRANSFER;
		urb->UrbIsochronousTransfer.PipeHandle = pdx->hinpipe;
		urb->UrbIsochronousTransfer.TransferFlags = USBD_TRANSFER_DIRECTION_IN | USBD_SHORT_TRANSFER_OK;
		urb->UrbIsochronousTransfer.TransferBufferLength = segsize;
		urb->UrbIsochronousTransfer.TransferBufferMDL = mdl;
		urb->UrbIsochronousTransfer.NumberOfPackets = npackets;
		urb->UrbIsochronousTransfer.StartFrame = frame;

		// Initialize the buffer offset in each packet. The host controller will fill in the packet
		// length in one of two ways. For an input operation, it gets set to whatever the device
		// supplies. For an output operation, the host controller driver sets it based on the buffer
		// offsets.

		ULONG offset = 0;			// packet offset within subsidiary MDL

		for (ULONG j = 0; j < npackets; ++j)
			{						// for each packet
			urb->UrbIsochronousTransfer.IsoPacket[j].Offset = offset;
			ULONG packlen = length > packsize ? packsize : length;
			urb->UrbIsochronousTransfer.IsoPacket[j].Length = 0;

			length -= packlen;
			offset += packlen;
			}						// for each packet

		frame += npackets;			// frame for next URB

		// Initialize the subsidiary IRP. The first stack location is for use by OnStageComplete,
		// which needs several pieces of information it's inconvenient to pass in a single
		// context pointer. The second stack location will describe an internal IOCTL for USBD.

		IoSetNextIrpStackLocation(subirp);
		PIO_STACK_LOCATION stack = IoGetCurrentIrpStackLocation(subirp);
		stack->DeviceObject = fdo;
		stack->Parameters.Others.Argument1 = (PVOID) urb;
		stack->Parameters.Others.Argument2 = (PVOID) mdl;

		stack = IoGetNextIrpStackLocation(subirp);
		stack->MajorFunction = IRP_MJ_INTERNAL_DEVICE_CONTROL;
		stack->Parameters.Others.Argument1 = (PVOID) urb;
		stack->Parameters.DeviceIoControl.IoControlCode = IOCTL_INTERNAL_USB_SUBMIT_URB;
		
		IoSetCompletionRoutine(subirp, (PIO_COMPLETION_ROUTINE) OnStageComplete,
			(PVOID) ctx, TRUE, TRUE, TRUE);
		}						// for each IRP

	// If the IRP/URB construction loop finished early, there must have been a memory allocation
	// error.

	if (i < numirps)
		status = STATUS_INSUFFICIENT_RESOURCES;

	// See if the main IRP has already been cancelled.

	else
		{						// check for cancellation
		IoSetCancelRoutine(Irp, OnCancelReadWrite);
		if (Irp->Cancel)
			{
			status = STATUS_CANCELLED;

			// The main IRP has been cancelled. If the cancellation just happened and our
			// cancel routine was called, the following call to IoSetCancelRoutine will
			// return NULL. In this case, both our cleanup code (see below) and the cancel
			// routine will be calling DestroyContextStructure. One or the other of those
			// calls will finish the cleanup of the structure. Otherwise, it will no longer
			// be possible for our cancel routine to get called and we should get rid of
			// its now-unnecessary claim on the structure.

			if (IoSetCancelRoutine(Irp, NULL))
				--ctx->refcnt;
			}
		else
			status = STATUS_SUCCESS;
		}						// check for cancellation

	// Install a completion routine for the main IRP. It will normally be called when
	// the last stage IRP finishes. It might get called pretty soon, though, if we
	// fail the IRP before submitting the stage IRPs.

	IoSetCompletionRoutine(Irp, (PIO_COMPLETION_ROUTINE) OnReadWriteComplete,
		(PVOID) ctx, TRUE, TRUE, TRUE);
	IoMarkIrpPending(Irp);
	IoSetNextIrpStackLocation(Irp);	// so our completion routine will get called

	// If we need to fail the main IRP, cleanup and fail it now.

	if (!NT_SUCCESS(status))
		{						// abort now
		for (i = 0; i < numirps; ++i)
			{					// release memory we were able to allocate
			if (ctx->sub[i].urb)
				ExFreePool(ctx->sub[i].urb);
			if (ctx->sub[i].mdl)
				IoFreeMdl(ctx->sub[i].mdl);
			}					// release memory we were able to allocate
		CompleteRequest(Irp, status, 0); // will trigger completion routine
		return STATUS_PENDING;
		}						// abort now

	ASSERT(length == 0);		// IRP construction loop didn't work right if not true

	// Install a completion routine for the main IRP. Once this is in place, submit the
	// subsidiary IRPs. Note that we can't touch the main IRP pointer after we submit the
	// subsidiary IRPs because the sub-IRP completion routine (OnStageComplete) may complete it.

	for (i = 0; i < numirps; ++i)
		IoCallDriver(pdx->LowerDeviceObject, ctx->sub[i].irp);

	return STATUS_PENDING;
	}							// DispatchRead

///////////////////////////////////////////////////////////////////////////////
// GetCurrentFrame obtains the current USB frame number

#pragma PAGEDCODE

ULONG GetCurrentFrame(PDEVICE_EXTENSION pdx)
	{							// GetCurrentFrame
	URB urb;
	urb.UrbGetCurrentFrameNumber.Hdr.Length = sizeof(struct _URB_GET_CURRENT_FRAME_NUMBER);
	urb.UrbGetCurrentFrameNumber.Hdr.Function = URB_FUNCTION_GET_CURRENT_FRAME_NUMBER;

	NTSTATUS status = SendAwaitUrb(pdx->DeviceObject, &urb);
	if (!NT_SUCCESS(status))
		{
		KdPrint((DRIVERNAME " - URB_FUNCTION_GET_CURRENT_FRAME_NUMBER failed - %X\n", status));
		return 0;
		}

	return urb.UrbGetCurrentFrameNumber.FrameNumber;
	}							// GetCurrentFrame

///////////////////////////////////////////////////////////////////////////////

#pragma PAGEDCODE

NTSTATUS GetStringDescriptor(PDEVICE_OBJECT fdo, UCHAR istring, PUNICODE_STRING s)
	{							// GetStringDescriptor
	NTSTATUS status;
	PDEVICE_EXTENSION pdx = (PDEVICE_EXTENSION) fdo->DeviceExtension;
	URB urb;

	UCHAR data[256];			// maximum-length buffer

	// If this is the first time here, read string descriptor zero and arbitrarily select
	// the first language identifer as the one to use in subsequent get-descriptor calls.

	if (!pdx->langid)
		{						// determine default language id
		UsbBuildGetDescriptorRequest(&urb, sizeof(_URB_CONTROL_DESCRIPTOR_REQUEST), USB_STRING_DESCRIPTOR_TYPE,
			0, 0, data, NULL, sizeof(data), NULL);
		status = SendAwaitUrb(fdo, &urb);
		if (!NT_SUCCESS(status))
			return status;
		pdx->langid = *(LANGID*)(data + 2);
		}						// determine default language id

	// Fetch the designated string descriptor.

	UsbBuildGetDescriptorRequest(&urb, sizeof(_URB_CONTROL_DESCRIPTOR_REQUEST), USB_STRING_DESCRIPTOR_TYPE,
		istring, pdx->langid, data, NULL, sizeof(data), NULL);
	status = SendAwaitUrb(fdo, &urb);
	if (!NT_SUCCESS(status))
		return status;

	ULONG nchars = (data[0] - sizeof(WCHAR)) / sizeof(WCHAR);
	if (nchars > 127)
		nchars = 127;
	PWSTR p = (PWSTR) ExAllocatePool(PagedPool, (nchars + 1) * sizeof(WCHAR));
	if (!p)
		return STATUS_INSUFFICIENT_RESOURCES;

	memcpy(p, data + 2, nchars * sizeof(WCHAR));
	p[nchars] = 0;

	s->Length = (USHORT) (sizeof(WCHAR) * nchars);
	s->MaximumLength = (USHORT) ((sizeof(WCHAR) * nchars) + sizeof(WCHAR));
	s->Buffer = p;

	return STATUS_SUCCESS;
	}							// GetStringDescriptor

///////////////////////////////////////////////////////////////////////////////

#pragma LOCKEDCODE

VOID OnCancelReadWrite(PDEVICE_OBJECT fdo, PIRP Irp)
	{							// OnCancelReadWrite
	
	IoReleaseCancelSpinLock(Irp->CancelIrql);

	// Attempt to cancel all the subsidiary IRPs for this IRP. The main IRP cannot
	// have progressed in the cancellation process past our own completion routine,