/*++

Copyright (c) 2005  Changzhi Zhou All Rights Reserved

Module Name:

    TDIClient.c

Abstract:

    This module contains the routines relevant to TDI client calls.

Environment:

    Kernel mode

Revision History:

	Changzhi Zhou Feb 5, 2005

--*/
#include "TDIClient.h"
#include "main.h"

NTSTATUS InitializeConnection( PDEVICE_EXTENSION deviceExtension )
/*++

Routine Description:

	1. Open a transport address and get the highest deviceobject of the transport stack.
	2. Build RecvContext in the DeviceExtension
	3. SetTdiEventHandle, includes TDI_EVENT_RECEIVE_DATAGRAM and TDI_EVENT_ERROR

Arguments:

    deviceExtension - pointer to a device object extension.

Return Value:

    NT status code.

--*/
{
	NTSTATUS	status;
	static USHORT		Port = LOCAL_PORT;

	if ( ( deviceExtension->hTransAddr ) || ( deviceExtension->lpTransAddrFileObject ) )
		return STATUS_INVALID_DEVICE_STATE;
	status = TDIOpenTransportAddress( wcharNetDevName,
		&deviceExtension->hTransAddr,
		&deviceExtension->lpTransAddrFileObject,
		Port
		);
	DebugPrint(("OpenTransAddr: handle = %x, lpFileObj = %x, status = %x\n",
		deviceExtension->hTransAddr, deviceExtension->lpTransAddrFileObject, status ));
	if ( status != STATUS_SUCCESS )
	{
		deviceExtension->lpTransAddrFileObject = NULL;
		deviceExtension->hTransAddr				= NULL;
		return status;
	}
	Port += 0x100;
	deviceExtension->TDILowerDeviceObject = IoGetRelatedDeviceObject( deviceExtension->lpTransAddrFileObject );
	DebugPrint(("TDI Provider: 0x%x\n", deviceExtension->TDILowerDeviceObject ));
	if ( deviceExtension->TDILowerDeviceObject == NULL )
	{
		ObDereferenceObject ( deviceExtension->lpTransAddrFileObject );
		ZwClose ( deviceExtension->hTransAddr );
		deviceExtension->lpTransAddrFileObject = NULL;
		deviceExtension->hTransAddr				= NULL;
		return STATUS_UNSUCCESSFUL;
	}

	TDIBuildRecvContext( &deviceExtension->recvContext );

	TDIClnSetEventHandler( deviceExtension->lpTransAddrFileObject,
		deviceExtension->TDILowerDeviceObject, TDI_EVENT_RECEIVE_DATAGRAM,
		TDIEventRecvDatagram, deviceExtension );
	TDIClnSetEventHandler( deviceExtension->lpTransAddrFileObject,
		deviceExtension->TDILowerDeviceObject, TDI_EVENT_ERROR,
		TDIEventError, NULL );
	return status;
}
NTSTATUS Disconnection ( PDEVICE_EXTENSION deviceExtension )
/*++

Routine Description:

	1. SetTdiEventhandle TDI_EVENT_RECEIVE_DATAGRAM to NULL
	2. Free RecvContext in the DeviceExtension
	3. Close an opened transport address.
	3. SetTdiEventHandle, includes TDI_EVENT_RECEIVE_DATAGRAM and TDI_EVENT_ERROR

Arguments:

    deviceExtension - pointer to a device object extension.

Return Value:

    NT status code.

--*/
{
	NTSTATUS	status;
	if ( ( deviceExtension->hTransAddr == NULL ) || ( deviceExtension->lpTransAddrFileObject == NULL ) )
		return STATUS_INVALID_DEVICE_STATE;
	TDIClnSetEventHandler( deviceExtension->lpTransAddrFileObject,
		deviceExtension->TDILowerDeviceObject, TDI_EVENT_RECEIVE_DATAGRAM,
		NULL, deviceExtension );
	TDIFreeRecvContext( &deviceExtension->recvContext );

	ObDereferenceObject ( deviceExtension->lpTransAddrFileObject );
	deviceExtension->lpTransAddrFileObject = NULL;
	status = ZwClose ( deviceExtension->hTransAddr );
	if ( status == STATUS_SUCCESS )
		deviceExtension->hTransAddr = NULL;
	else
	{
		DebugPrint(("ZwClose hTransAddr failed with status %x !\n", status ));
	}
	return status;
}
NTSTATUS	TDISendDatagram(
    IN PDEVICE_OBJECT DeviceObject,
    IN PIRP Irp
	)
/*++

Routine Description:

  This function is called in SampleWrite. It builds a internel irp and send down 
  to the underlying transport. Until the irp completes or timeout, it will return.
  The remote peer's IP address and port are indicated by RemoteAddress and 
  RemotePort in deviceExtension.

	1. TdiBuildInternalDeviceControlIrp for TDI_SEND_DATAGRAM
	2. Allocate and lockpage a MDL based on Irp->MdlAddress
	3. TdiBuildSendDatagram and pass down to the underlying transport device.
	4. Wait 3 seconds for the completion. If timeout, cancel the irp
	5. Set the Irp according with the completed irp.

Note:
	This routine is synchronous.
	It must be running at IRQL = PASSIVE_LEVEL

Arguments:

   DeviceObject - pointer to a device object.

   Irp - pointer to an I/O Request Packet for IRP_MJ_WRITE.

Return Value:

    NT status code.

--*/
{
	NTSTATUS	status;
	PIRP		pIrp;
	PDEVICE_EXTENSION	deviceExtension;
	KEVENT		Event;
	IO_STATUS_BLOCK	IoStatus;
	PMDL		pMdl;
	PVOID		VirtualAddress;
	ULONG		Length;
	EXCEPTION_POINTERS        * pExceptionInfo;
	ULONG                       lclExceptionCode;
	PVOID                       lclExceptionAddr;
	PVOID		SystemAddress;

	TA_IP_ADDRESS	RmtAddress = {1, {TDI_ADDRESS_LENGTH_IP, TDI_ADDRESS_TYPE_IP, { REMOTE_PORT, REMOTE_ADDRESS}}};
	TDI_CONNECTION_INFORMATION RmtInfo = {0, 0, 0, 0, sizeof(RmtAddress), &RmtAddress};

	deviceExtension = ( PDEVICE_EXTENSION )DeviceObject->DeviceExtension;
	RmtAddress.Address->Address->sin_port	= deviceExtension->RemotePort;
	RmtAddress.Address->Address->in_addr	= deviceExtension->RemoteAddress;
	status		= STATUS_SUCCESS;
	KeInitializeEvent( &Event, NotificationEvent, FALSE);
	pIrp = TdiBuildInternalDeviceControlIrp(
		TDI_SEND_DATAGRAM,
		deviceExtension->TDILowerDeviceObject,
		deviceExtension->lpTransAddrFileObject,
		&Event,
		&IoStatus );
	if ( NULL == pIrp )
	{
		DebugPrint(("TDISendDatagram: Allocate Irp failed\n"));
		return STATUS_INSUFFICIENT_RESOURCES;
	}
	VirtualAddress	= MmGetMdlVirtualAddress ( Irp->MdlAddress );
	Length			= MmGetMdlByteCount( Irp->MdlAddress );
	ASSERT ( VirtualAddress && Length );

	pMdl = IoAllocateMdl ( VirtualAddress, Length, FALSE, FALSE, pIrp );
	if ( NULL == pMdl )
	{
		DebugPrint(("TDISendDatagram: Allocate Mdl failed\n"));
		return STATUS_INSUFFICIENT_RESOURCES;
	}
	_try
	{
		MmProbeAndLockPages(pMdl,                     // (Try to) fix buffer.
			KernelMode,
			IoModifyAccess
			);
	}
	_except(
		pExceptionInfo = GetExceptionInformation(),                                                                                                                                                                                                    
		lclExceptionCode = pExceptionInfo->ExceptionRecord->ExceptionCode,                                                                                                                                                                             
		lclExceptionAddr = pExceptionInfo->ExceptionRecord->ExceptionAddress,                                                                                                                                                                          
		EXCEPTION_EXECUTE_HANDLER                                                                                                                                                                                                                      
		)                                                                                                                                                                                                                                               
	{                                                                                                                                                                                                                                                    
		DebugPrint((".TDISendDatagram:  MmProbeAndLockPages() failed.  Error = 0x%08x at 0x%08x\n",
			lclExceptionCode, lclExceptionAddr));
		status = STATUS_UNSUCCESSFUL;
		IoFreeMdl ( pMdl );
		CompleteRequest ( pIrp, status, 0 );
		Irp->IoStatus.Information = 0;
		return status;
	}
	TdiBuildSendDatagram(
		pIrp,
		deviceExtension->TDILowerDeviceObject,
		deviceExtension->lpTransAddrFileObject,
		NULL,
		NULL,
		pMdl,
		Length,
		&RmtInfo
		);
	status = IoCallDriver ( deviceExtension->TDILowerDeviceObject, pIrp );
	if( ( status != STATUS_SUCCESS ) || ( STATUS_PENDING != status ) )
	{
		DebugPrint(("TDISendDatagram: Problem in IoCallDriver with status %x\n", status ));
	}
	if ( status == STATUS_PENDING )
	{
        LARGE_INTEGER TimeOut;
		NTSTATUS		waitStatus;
		DebugPrint(("Send data pending... wait for 3 seconds\n"));
        TimeOut.QuadPart = -30000 * 1000;         // Calculate timeout value (5 seconds from now).
		waitStatus = KeWaitForSingleObject( &Event, Executive, KernelMode, FALSE, &TimeOut);
		if ( STATUS_TIMEOUT == waitStatus )
		{
			DebugPrint(("TDISendDatagram: Send data time out...\n"));
			IoCancelIrp ( pIrp );
			KeWaitForSingleObject( &Event, Executive, KernelMode, FALSE, 0);
		}
	}
	Irp->IoStatus.Information = IoStatus.Information;
	DebugPrint(("TDISendDatagram: Send %d bytes\n", Irp->IoStatus.Information ));
	status = ((STATUS_SUCCESS==status) || (STATUS_PENDING== status)) ? IoStatus.Status : status;
	return status;
}
NTSTATUS TDIEventError(
   IN PVOID TdiEventContext,  // The endpoint's file object.
   IN NTSTATUS Status         // Status code indicating error type.
   )
/*++
	Tdi error handle.
--*/
{
   DebugPrint(("Sample ErrorEventHandler: Status: 0x%8.8X\n", Status) );

   return( STATUS_SUCCESS );
}

NTSTATUS TDIEventRecvDatagram(
    IN PVOID  TdiEventContext,
    IN LONG  SourceAddressLength,
    IN PVOID  SourceAddress,
    IN LONG  OptionsLength,