tcp.c

ReactOS是一些高手根据Windows XP的内核编写出的类XP。内核实现机理和API函数调用几乎相同。甚至可以兼容XP的程序。喜欢研究系统内核的人可以看一看。

{
    KeInitializeEvent(&TimerLoopEvent, NotificationEvent, FALSE);
    PsCreateSystemThread(&TimerThreadHandle, THREAD_ALL_ACCESS, 0, 0, 0,
                         TimerThread, NULL);
}


NTSTATUS TCPStartup(VOID)
/*
 * FUNCTION: Initializes the TCP subsystem
 * RETURNS:
 *     Status of operation
 */
{
    NTSTATUS Status;

    TcpipRecursiveMutexInit( &TCPLock );
    ExInitializeFastMutex( &SleepingThreadsLock );
    InitializeListHead( &SleepingThreadsList );
    InitializeListHead( &SignalledConnections );
    Status = TCPMemStartup();
    if ( ! NT_SUCCESS(Status) ) {
	return Status;
    }

    PortsStartup( &TCPPorts, 1, 0xfffe );

    RegisterOskitTCPEventHandlers( &EventHandlers );
    InitOskitTCP();

    /* Register this protocol with IP layer */
    IPRegisterProtocol(IPPROTO_TCP, TCPReceive);

    ExInitializeNPagedLookasideList(
	&TCPSegmentList,                /* Lookaside list */
	NULL,                           /* Allocate routine */
	NULL,                           /* Free routine */
	0,                              /* Flags */
	sizeof(TCP_SEGMENT),            /* Size of each entry */
	TAG('T','C','P','S'),           /* Tag */
	0);                             /* Depth */

    StartTimer();

    TCPInitialized = TRUE;

    return STATUS_SUCCESS;
}


NTSTATUS TCPShutdown(VOID)
/*
 * FUNCTION: Shuts down the TCP subsystem
 * RETURNS:
 *     Status of operation
 */
{
    LARGE_INTEGER WaitForThread;

    if (!TCPInitialized)
	return STATUS_SUCCESS;

    WaitForThread.QuadPart = -2500000; /* 250 ms */
    KeSetEvent(&TimerLoopEvent, IO_NO_INCREMENT, FALSE);
    ZwWaitForSingleObject(TimerThreadHandle, FALSE, &WaitForThread);

    /* Deregister this protocol with IP layer */
    IPRegisterProtocol(IPPROTO_TCP, NULL);

    ExDeleteNPagedLookasideList(&TCPSegmentList);

    TCPInitialized = FALSE;

    DeinitOskitTCP();

    PortsShutdown( &TCPPorts );

    TCPMemShutdown();

    return STATUS_SUCCESS;
}

NTSTATUS TCPTranslateError( int OskitError ) {
    NTSTATUS Status = STATUS_UNSUCCESSFUL;

    switch( OskitError ) {
    case 0: Status = STATUS_SUCCESS; break;
    case OSK_EADDRNOTAVAIL:
    case OSK_EAFNOSUPPORT: Status = STATUS_INVALID_CONNECTION; break;
    case OSK_ECONNREFUSED:
    case OSK_ECONNRESET: Status = STATUS_REMOTE_NOT_LISTENING; break;
    case OSK_EINPROGRESS:
    case OSK_EAGAIN: Status = STATUS_PENDING; break;
    default: Status = STATUS_INVALID_CONNECTION; break;
    }

    TI_DbgPrint(DEBUG_TCP,("Error %d -> %x\n", OskitError, Status));
    return Status;
}

NTSTATUS TCPConnect
( PCONNECTION_ENDPOINT Connection,
  PTDI_CONNECTION_INFORMATION ConnInfo,
  PTDI_CONNECTION_INFORMATION ReturnInfo,
  PTCP_COMPLETION_ROUTINE Complete,
  PVOID Context ) {
    NTSTATUS Status;
    SOCKADDR_IN AddressToConnect = { 0 }, AddressToBind = { 0 };
    IP_ADDRESS RemoteAddress;
    USHORT RemotePort;
    PTDI_BUCKET Bucket;

    TI_DbgPrint(DEBUG_TCP,("TCPConnect: Called\n"));

    Bucket = ExAllocatePool( NonPagedPool, sizeof(*Bucket) );
    if( !Bucket ) return STATUS_NO_MEMORY;

    TcpipRecursiveMutexEnter( &TCPLock, TRUE );

    /* Freed in TCPSocketState */
    Bucket->Request.RequestNotifyObject = (PVOID)Complete;
    Bucket->Request.RequestContext = Context;

    InsertHeadList( &Connection->ConnectRequest, &Bucket->Entry );

    Status = AddrBuildAddress
	((PTRANSPORT_ADDRESS)ConnInfo->RemoteAddress,
	 &RemoteAddress,
	 &RemotePort);

    TI_DbgPrint(DEBUG_TCP,
                ("Connecting to address %x:%x\n",
                 RemoteAddress.Address.IPv4Address,
                 RemotePort));

    if (!NT_SUCCESS(Status)) {
	TI_DbgPrint(DEBUG_TCP, ("Could not AddrBuildAddress in TCPConnect\n"));
	return Status;
    }

    AddressToConnect.sin_family = AF_INET;
    AddressToBind = AddressToConnect;

    OskitTCPBind( Connection->SocketContext,
		  Connection,
		  &AddressToBind,
		  sizeof(AddressToBind) );

    memcpy( &AddressToConnect.sin_addr,
	    &RemoteAddress.Address.IPv4Address,
	    sizeof(AddressToConnect.sin_addr) );
    AddressToConnect.sin_port = RemotePort;

    Status = TCPTranslateError
	( OskitTCPConnect( Connection->SocketContext,
			   Connection,
			   &AddressToConnect,
			   sizeof(AddressToConnect) ) );

    TcpipRecursiveMutexLeave( &TCPLock );

    if( Status == OSK_EINPROGRESS )
	return STATUS_PENDING;
    else
	return Status;
}

NTSTATUS TCPDisconnect
( PCONNECTION_ENDPOINT Connection,
  UINT Flags,
  PTDI_CONNECTION_INFORMATION ConnInfo,
  PTDI_CONNECTION_INFORMATION ReturnInfo,
  PTCP_COMPLETION_ROUTINE Complete,
  PVOID Context ) {
    NTSTATUS Status;

    TI_DbgPrint(DEBUG_TCP,("started\n"));

    TcpipRecursiveMutexEnter( &TCPLock, TRUE );

    switch( Flags & (TDI_DISCONNECT_ABORT | TDI_DISCONNECT_RELEASE) ) {
    case 0:
    case TDI_DISCONNECT_ABORT:
	Flags = 0;
	break;

    case TDI_DISCONNECT_ABORT | TDI_DISCONNECT_RELEASE:
	Flags = 2;
	break;

    case TDI_DISCONNECT_RELEASE:
	Flags = 1;
	break;
    }

    Status = TCPTranslateError
	( OskitTCPShutdown( Connection->SocketContext, Flags ) );

    TcpipRecursiveMutexLeave( &TCPLock );

    TI_DbgPrint(DEBUG_TCP,("finished %x\n", Status));

    return Status;
}

NTSTATUS TCPClose
( PCONNECTION_ENDPOINT Connection ) {
    NTSTATUS Status;

    TI_DbgPrint(DEBUG_TCP,("TCPClose started\n"));

    TcpipRecursiveMutexEnter( &TCPLock, TRUE );

    Status = TCPTranslateError( OskitTCPClose( Connection->SocketContext ) );

    /* Make our code remove all pending IRPs */
    Connection->State |= SEL_FIN;
    DrainSignals();

    TcpipRecursiveMutexLeave( &TCPLock );

    TI_DbgPrint(DEBUG_TCP,("TCPClose finished %x\n", Status));

    return Status;
}

NTSTATUS TCPReceiveData
( PCONNECTION_ENDPOINT Connection,
  PNDIS_BUFFER Buffer,
  ULONG ReceiveLength,
  PULONG BytesReceived,
  ULONG ReceiveFlags,
  PTCP_COMPLETION_ROUTINE Complete,
  PVOID Context ) {
    OSK_PCHAR DataBuffer;
    UINT DataLen, Received = 0;
    NTSTATUS Status;
    PTDI_BUCKET Bucket;

    TI_DbgPrint(DEBUG_TCP,("Called for %d bytes (on socket %x)\n",
                           ReceiveLength, Connection->SocketContext));

    ASSERT_KM_POINTER(Connection->SocketContext);

    TcpipRecursiveMutexEnter( &TCPLock, TRUE );

    NdisQueryBuffer( Buffer, &DataBuffer, &DataLen );

    TI_DbgPrint(DEBUG_TCP,("TCP>|< Got an MDL %x (%x:%d)\n", Buffer, DataBuffer, DataLen));

    Status = TCPTranslateError
	( OskitTCPRecv
	  ( Connection->SocketContext,
	    DataBuffer,
	    DataLen,
	    &Received,
	    ReceiveFlags ) );

    TI_DbgPrint(DEBUG_TCP,("OskitTCPReceive: %x, %d\n", Status, Received));

    /* Keep this request around ... there was no data yet */
    if( Status == STATUS_PENDING ) {
	/* Freed in TCPSocketState */
	Bucket = ExAllocatePool( NonPagedPool, sizeof(*Bucket) );
	if( !Bucket ) {
	    TI_DbgPrint(DEBUG_TCP,("Failed to allocate bucket\n"));
	    TcpipRecursiveMutexLeave( &TCPLock );
	    return STATUS_NO_MEMORY;
	}

	Bucket->Request.RequestNotifyObject = Complete;
	Bucket->Request.RequestContext = Context;
	*BytesReceived = 0;

	InsertHeadList( &Connection->ReceiveRequest, &Bucket->Entry );
	Status = STATUS_PENDING;
	TI_DbgPrint(DEBUG_TCP,("Queued read irp\n"));
    } else {
	TI_DbgPrint(DEBUG_TCP,("Got status %x, bytes %d\n", Status, Received));
	*BytesReceived = Received;
    }

    TcpipRecursiveMutexLeave( &TCPLock );

    TI_DbgPrint(DEBUG_TCP,("Status %x\n", Status));

    return Status;
}

NTSTATUS TCPSendData
( PCONNECTION_ENDPOINT Connection,
  PCHAR BufferData,
  ULONG SendLength,
  PULONG BytesSent,
  ULONG Flags,
  PTCP_COMPLETION_ROUTINE Complete,
  PVOID Context ) {
    UINT Sent = 0;
    NTSTATUS Status;
    PTDI_BUCKET Bucket;

    TI_DbgPrint(DEBUG_TCP,("Called for %d bytes (on socket %x)\n",
                           SendLength, Connection->SocketContext));

    ASSERT_KM_POINTER(Connection->SocketContext);

    TcpipRecursiveMutexEnter( &TCPLock, TRUE );

    TI_DbgPrint(DEBUG_TCP,("Connection = %x\n", Connection));
    TI_DbgPrint(DEBUG_TCP,("Connection->SocketContext = %x\n",
			   Connection->SocketContext));

    Status = TCPTranslateError
	( OskitTCPSend( Connection->SocketContext,
			(OSK_PCHAR)BufferData, SendLength,
			&Sent, 0 ) );

    TI_DbgPrint(DEBUG_TCP,("OskitTCPSend: %x, %d\n", Status, Sent));

    /* Keep this request around ... there was no data yet */
    if( Status == STATUS_PENDING ) {
	/* Freed in TCPSocketState */
	Bucket = ExAllocatePool( NonPagedPool, sizeof(*Bucket) );
	if( !Bucket ) {
	    TI_DbgPrint(DEBUG_TCP,("Failed to allocate bucket\n"));
	    TcpipRecursiveMutexLeave( &TCPLock );
	    return STATUS_NO_MEMORY;
	}

	Bucket->Request.RequestNotifyObject = Complete;
	Bucket->Request.RequestContext = Context;
	*BytesSent = 0;

	InsertHeadList( &Connection->SendRequest, &Bucket->Entry );
	TI_DbgPrint(DEBUG_TCP,("Queued write irp\n"));
    } else {
	TI_DbgPrint(DEBUG_TCP,("Got status %x, bytes %d\n", Status, Sent));
	*BytesSent = Sent;
    }

    TcpipRecursiveMutexLeave( &TCPLock );

    TI_DbgPrint(DEBUG_TCP,("Status %x\n", Status));

    return Status;
}

VOID TCPTimeout(VOID) {
    /* Now handled by TimerThread */
}

UINT TCPAllocatePort( UINT HintPort ) {
    if( HintPort ) {
	if( AllocatePort( &TCPPorts, HintPort ) ) return HintPort;
	else {
            TI_DbgPrint
                (MID_TRACE,("We got a hint port but couldn't allocate it\n"));
            return (UINT)-1;
        }
    } else return AllocatePortFromRange( &TCPPorts, 1024, 5000 );
}

VOID TCPFreePort( UINT Port ) {
    DeallocatePort( &TCPPorts, Port );
}

NTSTATUS TCPGetPeerAddress
( PCONNECTION_ENDPOINT Connection,
  PTRANSPORT_ADDRESS Address ) {
    OSK_UINT LocalAddress, RemoteAddress;
    OSK_UI16 LocalPort, RemotePort;
    PTA_IP_ADDRESS AddressIP = (PTA_IP_ADDRESS)Address;

    TcpipRecursiveMutexEnter( &TCPLock, TRUE );

    OskitTCPGetAddress
        ( Connection->SocketContext,
          &LocalAddress, &LocalPort,
          &RemoteAddress, &RemotePort );

    AddressIP->TAAddressCount = 1;
    AddressIP->Address[0].AddressLength = TDI_ADDRESS_LENGTH_IP;
    AddressIP->Address[0].AddressType = TDI_ADDRESS_TYPE_IP;
    AddressIP->Address[0].Address[0].sin_port = RemotePort;
    AddressIP->Address[0].Address[0].in_addr = RemoteAddress;

    TcpipRecursiveMutexLeave( &TCPLock );

    return STATUS_SUCCESS;
}

VOID TCPRemoveIRP( PCONNECTION_ENDPOINT Endpoint, PIRP Irp ) {
    PLIST_ENTRY Entry;
    PLIST_ENTRY ListHead[4];
    KIRQL OldIrql;
    PTDI_BUCKET Bucket;
    UINT i = 0;

    ListHead[0] = &Endpoint->SendRequest;
    ListHead[1] = &Endpoint->ReceiveRequest;
    ListHead[2] = &Endpoint->ConnectRequest;
    ListHead[3] = &Endpoint->ListenRequest;

    TcpipAcquireSpinLock( &Endpoint->Lock, &OldIrql );

    for( i = 0; i < sizeof( ListHead ) / sizeof( ListHead[0] ); i++ ) {
	for( Entry = ListHead[i]->Flink;
	     Entry != ListHead[i];
	     Entry = Entry->Flink ) {
	    Bucket = CONTAINING_RECORD( Entry, TDI_BUCKET, Entry );
	    
	    if( Bucket->Request.RequestContext == Irp ) {
		RemoveEntryList( &Bucket->Entry );
		break;
	    }
	}
    }

    TcpipReleaseSpinLock( &Endpoint->Lock, OldIrql );
}

/* EOF */