C++ TdiBuildInternalDeviceControlIrp函数代码示例

NTSTATUS tdi_associate_address(PFILE_OBJECT connectionFileObject, HANDLE addressHandle)
{
    PDEVICE_OBJECT  devObj;
    KEVENT          event;
    PIRP            irp;
    IO_STATUS_BLOCK iosb;
    NTSTATUS        status;

    devObj = IoGetRelatedDeviceObject(connectionFileObject);

    KeInitializeEvent(&event, NotificationEvent, FALSE);

    irp = TdiBuildInternalDeviceControlIrp(TDI_ASSOCIATE_ADDRESS, devObj, connectionFileObject, &event, &iosb);

    if (irp == NULL)
    {
        return STATUS_INSUFFICIENT_RESOURCES;
    }

    TdiBuildAssociateAddress(irp, devObj, connectionFileObject, NULL, NULL, addressHandle);

    status = IoCallDriver(devObj, irp);

    if (status == STATUS_PENDING)
    {
        KeWaitForSingleObject(&event, Executive, KernelMode, FALSE, NULL);
        status = iosb.Status;
    }

    return status;
}

NTSTATUS tdi_set_event_handler(PFILE_OBJECT addressFileObject, LONG eventType, PVOID eventHandler, PVOID eventContext)
{
    PDEVICE_OBJECT  devObj;
    KEVENT          event;
    PIRP            irp;
    IO_STATUS_BLOCK iosb;
    NTSTATUS        status;

    devObj = IoGetRelatedDeviceObject(addressFileObject);

    KeInitializeEvent(&event, NotificationEvent, FALSE);

    irp = TdiBuildInternalDeviceControlIrp(TDI_SET_EVENT_HANDLER, devObj, addressFileObject, &event, &iosb);

    if (irp == NULL)
    {
        return STATUS_INSUFFICIENT_RESOURCES;
    }

    TdiBuildSetEventHandler(irp, devObj, addressFileObject, NULL, NULL, eventType, eventHandler, eventContext);

    status = IoCallDriver(devObj, irp);

    if (status == STATUS_PENDING)
    {
        KeWaitForSingleObject(&event, Executive, KernelMode, FALSE, NULL);
        status = iosb.Status;
    }

    return status;
}

NTSTATUS tdi_recv_stream(PFILE_OBJECT connectionFileObject, char *buf, int len, ULONG flags)
{
    PDEVICE_OBJECT  devObj;
    KEVENT          event;
    PIRP            irp;
    PMDL            mdl;
    IO_STATUS_BLOCK iosb;
    NTSTATUS        status;

    devObj = IoGetRelatedDeviceObject(connectionFileObject);

    KeInitializeEvent(&event, NotificationEvent, FALSE);

    irp = TdiBuildInternalDeviceControlIrp(TDI_RECEIVE, devObj, connectionFileObject, &event, &iosb);

    if (irp == NULL)
    {
        return STATUS_INSUFFICIENT_RESOURCES;
    }

    if (len)
    {
        mdl = IoAllocateMdl((void*) buf, len, FALSE, FALSE, NULL);

        if (mdl == NULL)
        {
            IoFreeIrp(irp);
            return STATUS_INSUFFICIENT_RESOURCES;
        }

        __try
        {
            MmProbeAndLockPages(mdl, KernelMode, IoWriteAccess);
            status = STATUS_SUCCESS;
        }
        __except (EXCEPTION_EXECUTE_HANDLER)
        {
            IoFreeMdl(mdl);
            IoFreeIrp(irp);
            status = STATUS_INVALID_USER_BUFFER;
        }

        if (!NT_SUCCESS(status))
        {
            return status;
        }
    }

    TdiBuildReceive(irp, devObj, connectionFileObject, NULL, NULL, len ? mdl : 0, flags, len);

    status = IoCallDriver(devObj, irp);

    if (status == STATUS_PENDING)
    {
        KeWaitForSingleObject(&event, Executive, KernelMode, FALSE, NULL);
        status = iosb.Status;
    }

    return NT_SUCCESS(status) ? iosb.Information : status;
}

NTSTATUS tdi_disconnect(PFILE_OBJECT connectionFileObject, ULONG flags)
{
    PDEVICE_OBJECT  devObj;
    KEVENT          event;
    PIRP            irp;
    IO_STATUS_BLOCK iosb;
    NTSTATUS        status;

    devObj = IoGetRelatedDeviceObject(connectionFileObject);

    KeInitializeEvent(&event, NotificationEvent, FALSE);

    irp = TdiBuildInternalDeviceControlIrp(TDI_DISCONNECT, devObj, connectionFileObject, &event, &iosb);

    if (irp == NULL)
    {
        return STATUS_INSUFFICIENT_RESOURCES;
    }

    TdiBuildDisconnect(irp, devObj, connectionFileObject, NULL, NULL, NULL, flags, NULL, NULL);

    status = IoCallDriver(devObj, irp);

    if (status == STATUS_PENDING)
    {
        KeWaitForSingleObject(&event, Executive, KernelMode, FALSE, NULL);
        status = iosb.Status;
    }

    return status;
}

NTSTATUS
	TdiQueryAddress(
		PFILE_OBJECT		AddressObject, 
		PTDI_ADDRESS_INFO	AddressInfo,
		PULONG				pInfoLength
		)
{
	NTSTATUS	ntStatus;
	PIRP		Irp;
	PMDL		Mdl;
    PDEVICE_OBJECT	DeviceObject = IoGetRelatedDeviceObject(AddressObject);
    IO_STATUS_BLOCK IoStatus;

    if (!(Irp = TdiBuildInternalDeviceControlIrp(TDI_QUERY_INFORMATION, DeviceObject, FileObject, NULL, &IoStatus)))
		return STATUS_INSUFFICIENT_RESOURCES;

	if (!(Mdl = IoAllocateMdl(AddressInfo, *pInfoLength, FALSE, FALSE, Irp)))
	{
		IoFreeIrp(Irp);
		return STATUS_INSUFFICIENT_RESOURCES;
	}

	MmProbeAndLockPages(Mdl, KernelMode, IoModifyAccess);

    TdiBuildQueryInformation(Irp, DeviceObject, AddressObject, NULL, NULL, TDI_QUERY_ADDRESS_INFO, Mdl);

    ntStatus = IoCallDriver(DeviceObject, Irp);

    return(ntStatus == STATUS_SUCCESS ? IoStatus.Status : ntStatus);

}

NTSTATUS
	TdiDisconnect(
		PFILE_OBJECT ConnectionObject
		)
{
	NTSTATUS	ntStatus;
    KEVENT		Event; 
	PIRP		Irp;
    PDEVICE_OBJECT DeviceObject = IoGetRelatedDeviceObject(ConnectionObject);
    IO_STATUS_BLOCK IoStatus;

	KeInitializeEvent(&Event, NotificationEvent, FALSE);

    if (!(Irp = TdiBuildInternalDeviceControlIrp(TDI_DISCONNECT, DeviceObject, ConnectionObject, &Event, &IoStatus)))
		return STATUS_INSUFFICIENT_RESOURCES;

    TdiBuildDisconnect(Irp, DeviceObject, ConnectionObject, 0, 0, 0, TDI_DISCONNECT_RELEASE, 0, 0);

    ntStatus = IoCallDriver(DeviceObject, Irp);

    if (ntStatus == STATUS_PENDING)
        ntStatus = KeWaitForSingleObject(&Event, UserRequest, KernelMode, FALSE, 0);

    return(ntStatus == STATUS_SUCCESS ? IoStatus.Status : ntStatus);

}

NTSTATUS
	TdiSetEventHandler(
		PFILE_OBJECT	AddressObject,
		LONG			EventType, 
		PVOID			EventHandler,
		PVOID			EventContext
		)
{
	NTSTATUS	ntStatus;
    KEVENT		Event; 
	PIRP		Irp;
    PDEVICE_OBJECT	DeviceObject = IoGetRelatedDeviceObject(AddressObject);
	IO_STATUS_BLOCK IoStatus;
    
	KeInitializeEvent(&Event, NotificationEvent, FALSE);

    if (!(Irp = TdiBuildInternalDeviceControlIrp(TDI_SET_EVENT_HANDLER, DeviceObject, AddressObject, &Event, &IoStatus)))
		return STATUS_INSUFFICIENT_RESOURCES;

    TdiBuildSetEventHandler(Irp, DeviceObject, AddressObject, 0, 0, EventType, EventHandler, EventContext);

    ntStatus = IoCallDriver(DeviceObject, Irp);

    if (ntStatus == STATUS_PENDING)
        ntStatus = KeWaitForSingleObject(&Event, UserRequest, KernelMode, FALSE, 0);

    return(ntStatus == STATUS_SUCCESS ? IoStatus.Status : ntStatus);
}

VOID DestroyNodeEventHandlerVisitor(PVOID context, UINT key, PVOID value)
{
	NTSTATUS status;
	PIRP irp = NULL;
	unsigned int i = 0;
	KEVENT irp_complete_event;
	IO_STATUS_BLOCK io_status_block;
	NodeEventHandler* node_event_handler = NULL;
	EventHandlerManager* event_handler_manager = NULL;

	node_event_handler = (NodeEventHandler*)value;
	event_handler_manager = (EventHandlerManager*)context;

	if(node_event_handler == NULL ||
		event_handler_manager == NULL )
		return;

	KeInitializeEvent(&irp_complete_event, NotificationEvent, FALSE);

	// For each hooked event handler send an IRP to set the event handler
	// back to what it originally was
	for(i = 0; i < TDI_EVENT_HANDLERS_MAX; i++)
	{
		PFILE_OBJECT node = node_event_handler->handlers[i].node;
		TdiEventHandler* original_handler = &node_event_handler->handlers[i].original_handler;

		if(original_handler->EventHandler == NULL)
			continue;

		irp = TdiBuildInternalDeviceControlIrp(TDI_SET_EVENT_HANDLER, event_handler_manager->connection_manager->device, node, 
			&irp_complete_event, &io_status_block);

		if(irp == NULL)
			continue;

		TdiBuildSetEventHandler(irp, event_handler_manager->connection_manager->device, node, NULL, NULL, 
			original_handler->EventType, original_handler->EventHandler, original_handler->EventContext);

		status = IoCallDriver(event_handler_manager->connection_manager->next_device, irp);

		if(status == STATUS_PENDING)                                
		{
			// Wait for the IRP to finish
			KeWaitForSingleObject((PVOID)&irp_complete_event, Executive, KernelMode, TRUE, NULL); 
			status = io_status_block.Status; 
		}

		if(status == STATUS_SUCCESS)
		{
			DbgPrint("EventHandlerManager: Successfully set original event handler n=%08x t=%08x\n", node, original_handler->EventType);
		}
		else
		{
			DbgPrint("EventHandlerManager: Failed to set original event handler - code=%08x n=%08x t=%08x\n", status, node, original_handler->EventType);
		}

		KeResetEvent(&irp_complete_event);
	}
}

NTSTATUS TdiSetEventHandler(
    PFILE_OBJECT FileObject,
    LONG EventType,
    PVOID Handler,
    PVOID Context)
/*
 * FUNCTION: Sets or resets an event handler
 * ARGUMENTS:
 *     FileObject = Pointer to file object
 *     EventType  = Event code
 *     Handler    = Event handler to be called when the event occurs
 *     Context    = Context input to handler when the event occurs
 * RETURNS:
 *     Status of operation
 * NOTES:
 *     Specify NULL for Handler to stop calling event handler
 */
{
    PDEVICE_OBJECT DeviceObject;
    IO_STATUS_BLOCK Iosb;
    KEVENT Event;
    PIRP Irp;

    AFD_DbgPrint(MAX_TRACE, ("Called\n"));

    if (!FileObject) {
        AFD_DbgPrint(MIN_TRACE, ("Bad file object.\n"));
        return STATUS_INVALID_PARAMETER;
    }

    DeviceObject = IoGetRelatedDeviceObject(FileObject);
    if (!DeviceObject) {
        AFD_DbgPrint(MIN_TRACE, ("Bad device object.\n"));
        return STATUS_INVALID_PARAMETER;
    }

    KeInitializeEvent(&Event, NotificationEvent, FALSE);

    Irp = TdiBuildInternalDeviceControlIrp(TDI_SET_EVENT_HANDLER,   /* Sub function */
                                           DeviceObject,            /* Device object */
                                           FileObject,              /* File object */
                                           &Event,                  /* Event */
                                           &Iosb);                  /* Status */
    if (!Irp)
        return STATUS_INSUFFICIENT_RESOURCES;



    TdiBuildSetEventHandler(Irp,
                            DeviceObject,
                            FileObject,
                            NULL,
                            NULL,
                            EventType,
                            Handler,
                            Context);

    return TdiCall(Irp, DeviceObject, &Event, &Iosb);
}

BOOLEAN KTdiStreamSocket::Unbind()
{
  //KLocker locker(&m_KSynchroObject);

  BOOLEAN         bRes = TRUE;
  PIRP            pIrp = NULL, pIrpError = NULL;
  PDEVICE_OBJECT  pDeviceObject;
  NTSTATUS        NtStatus;
  IO_STATUS_BLOCK IoStatusBlock;

  __try
  {
    if (m_bOpen == TRUE && m_bBind == TRUE && m_bConnected == FALSE && m_bListen == FALSE)
    {
      bRes = FALSE;

      pDeviceObject = IoGetRelatedDeviceObject(m_pTdiConnectionObject);
      
      pIrp = TdiBuildInternalDeviceControlIrp(
                    TDI_DISASSOCIATE_ADDRESS, 
                    pDeviceObject, 
                    m_pTdiConnectionObject,
                    NULL, 
                    NULL);
      pIrpError = pIrp;
      if (pIrp != NULL)
      {
        TdiBuildDisassociateAddress(
               pIrp, 
               pDeviceObject, 
               m_pTdiConnectionObject,
               NULL,
               NULL);
        
        pIrpError = NULL;
        NtStatus = TdiCall(pIrp, pDeviceObject, &IoStatusBlock);
        if (NT_SUCCESS(NtStatus))
        {
          m_bBind = FALSE;
          bRes = TRUE;
        }
        else
        {
          DbgPrint ("TdiUnbind: ERROR (%08x)!!!\n", NtStatus);
        }
      }
    }
  }
  __finally
  {
   if (pIrpError != NULL)
     IoFreeIrp(pIrpError);
  }

  return bRes;
}

NTSTATUS TdiListen(
    PIRP *Irp,
    PFILE_OBJECT ConnectionObject,
    PTDI_CONNECTION_INFORMATION *RequestConnectionInfo,
    PTDI_CONNECTION_INFORMATION *ReturnConnectionInfo,
    PIO_COMPLETION_ROUTINE  CompletionRoutine,
    PVOID CompletionContext)
/*
 * FUNCTION: Listen on a connection endpoint for a connection request from a remote peer
 * ARGUMENTS:
 *     CompletionRoutine = Routine to be called when IRP is completed
 *     CompletionContext = Context for CompletionRoutine
 * RETURNS:
 *     Status of operation
 *     May return STATUS_PENDING
 */
{
    PDEVICE_OBJECT DeviceObject;

    AFD_DbgPrint(MAX_TRACE, ("Called\n"));

    ASSERT(*Irp == NULL);

    if (!ConnectionObject) {
        AFD_DbgPrint(MIN_TRACE, ("Bad connection object.\n"));
        return STATUS_INVALID_PARAMETER;
    }

    DeviceObject = IoGetRelatedDeviceObject(ConnectionObject);
    if (!DeviceObject) {
        AFD_DbgPrint(MIN_TRACE, ("Bad device object.\n"));
        return STATUS_INVALID_PARAMETER;
    }

    *Irp = TdiBuildInternalDeviceControlIrp(TDI_LISTEN,              /* Sub function */
                                            DeviceObject,            /* Device object */
                                            ConnectionObject,        /* File object */
                                            NULL,                    /* Event */
                                            NULL);                   /* Status */
    if (*Irp == NULL)
        return STATUS_INSUFFICIENT_RESOURCES;

    TdiBuildListen(*Irp,                   /* IRP */
                   DeviceObject,           /* Device object */
                   ConnectionObject,       /* File object */
                   CompletionRoutine,      /* Completion routine */
                   CompletionContext,      /* Completion routine context */
                   0,                      /* Flags */
                   *RequestConnectionInfo, /* Request connection information */
                   *ReturnConnectionInfo);  /* Return connection information */

    TdiCall(*Irp, DeviceObject, NULL /* Don't wait for completion */, NULL);

    return STATUS_PENDING;
}

NTSTATUS TdiConnect(
    PIRP *Irp,
    PFILE_OBJECT ConnectionObject,
    PTDI_CONNECTION_INFORMATION ConnectionCallInfo,
    PTDI_CONNECTION_INFORMATION ConnectionReturnInfo,
    PIO_COMPLETION_ROUTINE CompletionRoutine,
    PVOID CompletionContext)
/*
 * FUNCTION: Connect a connection endpoint to a remote peer
 * ARGUMENTS:
 *     ConnectionObject = Pointer to connection endpoint file object
 *     RemoteAddress    = Pointer to remote address
 * RETURNS:
 *     Status of operation
 */
{
    PDEVICE_OBJECT DeviceObject;

    AFD_DbgPrint(MAX_TRACE, ("Called\n"));

    ASSERT(*Irp == NULL);

    if (!ConnectionObject) {
        AFD_DbgPrint(MIN_TRACE, ("Bad connection object.\n"));
        return STATUS_INVALID_PARAMETER;
    }

    DeviceObject = IoGetRelatedDeviceObject(ConnectionObject);
    if (!DeviceObject) {
        AFD_DbgPrint(MIN_TRACE, ("Bad device object.\n"));
        return STATUS_INVALID_PARAMETER;
    }

    *Irp = TdiBuildInternalDeviceControlIrp(TDI_CONNECT,             /* Sub function */
                                            DeviceObject,            /* Device object */
                                            ConnectionObject,        /* File object */
                                            NULL,                    /* Event */
                                            NULL);                   /* Status */
    if (!*Irp) {
        return STATUS_INSUFFICIENT_RESOURCES;
    }

    TdiBuildConnect(*Irp,                   /* IRP */
                    DeviceObject,           /* Device object */
                    ConnectionObject,       /* File object */
                    CompletionRoutine,      /* Completion routine */
                    CompletionContext,      /* Completion routine context */
                    NULL,                   /* Time */
                    ConnectionCallInfo,     /* Request connection information */
                    ConnectionReturnInfo);  /* Return connection information */

    TdiCall(*Irp, DeviceObject, NULL, NULL);

    return STATUS_PENDING;
}

NTSTATUS TdiQueryInformation(
    PFILE_OBJECT FileObject,
    LONG QueryType,
    PMDL MdlBuffer)
/*
 * FUNCTION: Query for information
 * ARGUMENTS:
 *     FileObject   = Pointer to file object
 *     QueryType    = Query type
 *     MdlBuffer    = Pointer to MDL buffer specific for query type
 * RETURNS:
 *     Status of operation
 */
{
    PDEVICE_OBJECT DeviceObject;
    IO_STATUS_BLOCK Iosb;
    KEVENT Event;
    PIRP Irp;

    if (!FileObject) {
        AFD_DbgPrint(MIN_TRACE, ("Bad file object.\n"));
        return STATUS_INVALID_PARAMETER;
    }

    DeviceObject = IoGetRelatedDeviceObject(FileObject);
    if (!DeviceObject) {
        AFD_DbgPrint(MIN_TRACE, ("Bad device object.\n"));
        return STATUS_INVALID_PARAMETER;
    }

    KeInitializeEvent(&Event, NotificationEvent, FALSE);

    Irp = TdiBuildInternalDeviceControlIrp(TDI_QUERY_INFORMATION,       /* Sub function */
                                           DeviceObject,                /* Device object */
                                           ConnectionObject,            /* File object */
                                           &Event,                      /* Event */
                                           &Iosb);                      /* Status */
    if (!Irp) {
        return STATUS_INSUFFICIENT_RESOURCES;
    }

    TdiBuildQueryInformation(Irp,
                             DeviceObject,
                             FileObject,
                             NULL,
                             NULL,
                             QueryType,
                             MdlBuffer);

    return TdiCall(Irp, DeviceObject, &Event, &Iosb);
}

NTSTATUS TdiAssociateAddressFile(
    HANDLE AddressHandle,
    PFILE_OBJECT ConnectionObject)
/*
 * FUNCTION: Associates a connection endpoint to an address file object
 * ARGUMENTS:
 *     AddressHandle    = Handle to address file object
 *     ConnectionObject = Connection endpoint file object
 * RETURNS:
 *     Status of operation
 */
{
    PDEVICE_OBJECT DeviceObject;
    IO_STATUS_BLOCK Iosb;
    KEVENT Event;
    PIRP Irp;

    AFD_DbgPrint(MAX_TRACE, ("Called. AddressHandle (%p)  ConnectionObject (%p)\n",
                             AddressHandle, ConnectionObject));

    if (!ConnectionObject) {
        AFD_DbgPrint(MIN_TRACE, ("Bad connection object.\n"));
        return STATUS_INVALID_PARAMETER;
    }

    DeviceObject = IoGetRelatedDeviceObject(ConnectionObject);
    if (!DeviceObject) {
        AFD_DbgPrint(MIN_TRACE, ("Bad device object.\n"));
        return STATUS_INVALID_PARAMETER;
    }

    KeInitializeEvent(&Event, NotificationEvent, FALSE);

    Irp = TdiBuildInternalDeviceControlIrp(TDI_ASSOCIATE_ADDRESS,   /* Sub function */
                                           DeviceObject,            /* Device object */
                                           ConnectionObject,        /* File object */
                                           &Event,                  /* Event */
                                           &Iosb);                  /* Status */
    if (!Irp)
        return STATUS_INSUFFICIENT_RESOURCES;

    TdiBuildAssociateAddress(Irp,
                             DeviceObject,
                             ConnectionObject,
                             NULL,
                             NULL,
                             AddressHandle);

    return TdiCall(Irp, DeviceObject, &Event, &Iosb);
}

NTSTATUS
	TdiRecvfrom(
		PFILE_OBJECT	AddressObject, 
		PVOID			Data, 
		ULONG			Length, 
		ULONG			Addr,
		USHORT			Port,
		PULONG			pRead,
		PKEVENT			ReadEvent
		)
{
	KEVENT	Event; 
	PIRP	Irp;
	PMDL	Mdl;
	NTSTATUS	ntStatus;
	IO_STATUS_BLOCK IoStatus;
    PDEVICE_OBJECT DeviceObject = IoGetRelatedDeviceObject(AddressObject);
	TA_IP_ADDRESS	SrcAddr = {1, {TDI_ADDRESS_LENGTH_IP, TDI_ADDRESS_TYPE_IP, {Port, Addr}}};
    TDI_CONNECTION_INFORMATION Src = {0, 0, 0, 0, sizeof SrcAddr, &SrcAddr};
	PKEVENT pEvent = ReadEvent;

	if (!ReadEvent)
	{
		KeInitializeEvent(&Event, NotificationEvent, FALSE);
		pEvent = &Event;
	}

    if (!(Irp = TdiBuildInternalDeviceControlIrp(TDI_RECEIVE_DATAGRAM, DeviceObject, AddressObject, pEvent, &IoStatus)))
		return STATUS_INSUFFICIENT_RESOURCES;

    if (!(Mdl = IoAllocateMdl(Data, Length, FALSE, FALSE, Irp)))
		return STATUS_INSUFFICIENT_RESOURCES;

    MmProbeAndLockPages(Mdl, KernelMode, IoModifyAccess);
	TdiBuildReceiveDatagram(Irp, DeviceObject, AddressObject, 0, 0, Mdl, Length, &Src, &Src, TDI_RECEIVE_NORMAL);

   	ntStatus = IoCallDriver(DeviceObject, Irp);

    if ((ntStatus == STATUS_PENDING) && (!ReadEvent))
        ntStatus = KeWaitForSingleObject(pEvent, UserRequest, KernelMode, FALSE, 0);

	if (ntStatus == STATUS_SUCCESS)
	{
		*pRead = IoStatus.Status;
	}

    return(ntStatus);

}