[IP]
[reactos.git] / reactos / lib / drivers / ip / transport / tcp / tcp.c
index cfbca36..187c9a5 100644 (file)
 LONG TCP_IPIdentification = 0;
 static BOOLEAN TCPInitialized = FALSE;
 static NPAGED_LOOKASIDE_LIST TCPSegmentList;
-LIST_ENTRY SignalledConnections;
-LIST_ENTRY SleepingThreadsList;
-FAST_MUTEX SleepingThreadsLock;
-RECURSIVE_MUTEX TCPLock;
 PORT_SET TCPPorts;
+CLIENT_DATA ClientInfo;
 
-static VOID HandleSignalledConnection( PCONNECTION_ENDPOINT Connection,
-                                      ULONG NewState ) {
-    NTSTATUS Status = STATUS_SUCCESS;
-    PTCP_COMPLETION_ROUTINE Complete;
-    PTDI_BUCKET Bucket;
-    PLIST_ENTRY Entry;
-    PIRP Irp;
-    PMDL Mdl;
+VOID HandleSignalledConnection(PCONNECTION_ENDPOINT Connection)
+{
+        PTDI_BUCKET Bucket;
+        PLIST_ENTRY Entry;
+        NTSTATUS Status;
+        PIRP Irp;
+        PMDL Mdl;
+        ULONG SocketError = 0;
+        KIRQL OldIrql;
+        PTCP_COMPLETION_ROUTINE Complete;
 
-    TI_DbgPrint(MID_TRACE,("Handling signalled state on %x (%x)\n",
-                           Connection, Connection->SocketContext));
+        if (ClientInfo.Unlocked)
+            LockObjectAtDpcLevel(Connection);
 
-    /* Things that can happen when we try the initial connection */
-    if( NewState & SEL_CONNECT ) {
-       while( !IsListEmpty( &Connection->ConnectRequest ) ) {
-            Entry = RemoveHeadList( &Connection->ConnectRequest );
-            TI_DbgPrint(DEBUG_TCP, ("Connect Event\n"));
+        TI_DbgPrint(MID_TRACE,("Handling signalled state on %x (%x)\n",
+                               Connection, Connection->SocketContext));
 
-            Bucket = CONTAINING_RECORD( Entry, TDI_BUCKET, Entry );
-            Complete = Bucket->Request.RequestNotifyObject;
-            TI_DbgPrint(DEBUG_TCP,
-                        ("Completing Request %x\n", Bucket->Request));
+        if( Connection->SignalState & SEL_FIN ) {
+            TI_DbgPrint(DEBUG_TCP, ("EOF From socket\n"));
 
-            if( (NewState & (SEL_CONNECT | SEL_FIN)) ==
-                (SEL_CONNECT | SEL_FIN) )
-                Status = STATUS_CONNECTION_REFUSED;
-            else
-                Status = STATUS_SUCCESS;
+            /* If OskitTCP initiated the disconnect, try to read the socket error that occurred */
+            if (Connection->SocketContext)
+                SocketError = TCPTranslateError(OskitTCPGetSocketError(Connection->SocketContext));
 
-            Complete( Bucket->Request.RequestContext, Status, 0 );
+            /* Default to STATUS_CANCELLED if we initiated the disconnect or no socket error was reported */
+            if (!Connection->SocketContext || !SocketError)
+                SocketError = STATUS_CANCELLED;
 
-            /* Frees the bucket allocated in TCPConnect */
-            exFreePool( Bucket );
-        }
-    }
+            while (!IsListEmpty(&Connection->ReceiveRequest))
+            {
+               Entry = RemoveHeadList( &Connection->ReceiveRequest );
 
-    if( NewState & SEL_ACCEPT ) {
-       /* Handle readable on a listening socket --
-        * TODO: Implement filtering
-        */
-
-       TI_DbgPrint(DEBUG_TCP,("Accepting new connection on %x (Queue: %s)\n",
-                              Connection,
-                              IsListEmpty(&Connection->ListenRequest) ?
-                              "empty" : "nonempty"));
-
-       while( !IsListEmpty( &Connection->ListenRequest ) ) {
-           PIO_STACK_LOCATION IrpSp;
-
-           Entry = RemoveHeadList( &Connection->ListenRequest );
-           Bucket = CONTAINING_RECORD( Entry, TDI_BUCKET, Entry );
-           Complete = Bucket->Request.RequestNotifyObject;
-
-           Irp = Bucket->Request.RequestContext;
-           IrpSp = IoGetCurrentIrpStackLocation( Irp );
-
-           TI_DbgPrint(DEBUG_TCP,("Getting the socket\n"));
-           Status = TCPServiceListeningSocket
-               ( Connection->AddressFile->Listener,
-                 Bucket->AssociatedEndpoint,
-                 (PTDI_REQUEST_KERNEL)&IrpSp->Parameters );
-
-           TI_DbgPrint(DEBUG_TCP,("Socket: Status: %x\n"));
-
-           if( Status == STATUS_PENDING ) {
-               InsertHeadList( &Connection->ListenRequest, &Bucket->Entry );
-               break;
-           } else {
-               Complete( Bucket->Request.RequestContext, Status, 0 );
-               exFreePool( Bucket );
-           }
-       }
-    }
+               Bucket = CONTAINING_RECORD( Entry, TDI_BUCKET, Entry );
 
-    /* Things that happen after we're connected */
-    if( NewState & SEL_READ ) {
-       TI_DbgPrint(DEBUG_TCP,("Readable: irp list %s\n",
-                              IsListEmpty(&Connection->ReceiveRequest) ?
-                              "empty" : "nonempty"));
-
-       while( !IsListEmpty( &Connection->ReceiveRequest ) ) {
-           OSK_UINT RecvLen = 0, Received = 0;
-           OSK_PCHAR RecvBuffer = 0;
-
-           Entry = RemoveHeadList( &Connection->ReceiveRequest );
-           Bucket = CONTAINING_RECORD( Entry, TDI_BUCKET, Entry );
-           Complete = Bucket->Request.RequestNotifyObject;
-
-           Irp = Bucket->Request.RequestContext;
-           Mdl = Irp->MdlAddress;
-
-           TI_DbgPrint(DEBUG_TCP,
-                       ("Getting the user buffer from %x\n", Mdl));
-
-           NdisQueryBuffer( Mdl, &RecvBuffer, &RecvLen );
-
-           TI_DbgPrint(DEBUG_TCP,
-                       ("Reading %d bytes to %x\n", RecvLen, RecvBuffer));
-
-           TI_DbgPrint(DEBUG_TCP, ("Connection: %x\n", Connection));
-           TI_DbgPrint
-               (DEBUG_TCP,
-                ("Connection->SocketContext: %x\n",
-                 Connection->SocketContext));
-           TI_DbgPrint(DEBUG_TCP, ("RecvBuffer: %x\n", RecvBuffer));
-
-           Status = TCPTranslateError
-               ( OskitTCPRecv( Connection->SocketContext,
-                               RecvBuffer,
-                               RecvLen,
-                               &Received,
-                               0 ) );
-
-           TI_DbgPrint(DEBUG_TCP,("TCP Bytes: %d\n", Received));
-
-           if( Status == STATUS_SUCCESS ) {
-               TI_DbgPrint(DEBUG_TCP,("Received %d bytes with status %x\n",
-                                      Received, Status));
-
-               Complete( Bucket->Request.RequestContext,
-                         STATUS_SUCCESS, Received );
-               exFreePool( Bucket );
-           } else if( Status == STATUS_PENDING ) {
-               InsertHeadList
-                   ( &Connection->ReceiveRequest, &Bucket->Entry );
-               break;
-           } else {
-               TI_DbgPrint(DEBUG_TCP,
-                           ("Completing Receive request: %x %x\n",
-                            Bucket->Request, Status));
-               Complete( Bucket->Request.RequestContext, Status, 0 );
-               exFreePool( Bucket );
-           }
-       }
-    }
-    if( NewState & SEL_WRITE ) {
-       TI_DbgPrint(DEBUG_TCP,("Writeable: irp list %s\n",
-                              IsListEmpty(&Connection->ReceiveRequest) ?
-                              "empty" : "nonempty"));
-
-       while( !IsListEmpty( &Connection->SendRequest ) ) {
-           OSK_UINT SendLen = 0, Sent = 0;
-           OSK_PCHAR SendBuffer = 0;
-
-           Entry = RemoveHeadList( &Connection->SendRequest );
-           Bucket = CONTAINING_RECORD( Entry, TDI_BUCKET, Entry );
-           Complete = Bucket->Request.RequestNotifyObject;
-
-           Irp = Bucket->Request.RequestContext;
-           Mdl = Irp->MdlAddress;
-
-           TI_DbgPrint(DEBUG_TCP,
-                       ("Getting the user buffer from %x\n", Mdl));
-
-           NdisQueryBuffer( Mdl, &SendBuffer, &SendLen );
-
-           TI_DbgPrint(DEBUG_TCP,
-                       ("Writing %d bytes to %x\n", SendLen, SendBuffer));
-
-           TI_DbgPrint(DEBUG_TCP, ("Connection: %x\n", Connection));
-           TI_DbgPrint
-               (DEBUG_TCP,
-                ("Connection->SocketContext: %x\n",
-                 Connection->SocketContext));
-
-           Status = TCPTranslateError
-               ( OskitTCPSend( Connection->SocketContext,
-                               SendBuffer,
-                               SendLen,
-                               &Sent,
-                               0 ) );
-
-           TI_DbgPrint(DEBUG_TCP,("TCP Bytes: %d\n", Sent));
-
-           if( Status == STATUS_SUCCESS ) {
-               TI_DbgPrint(DEBUG_TCP,("Sent %d bytes with status %x\n",
-                                      Sent, Status));
-
-               Complete( Bucket->Request.RequestContext,
-                         STATUS_SUCCESS, Sent );
-               exFreePool( Bucket );
-           } else if( Status == STATUS_PENDING ) {
-               InsertHeadList
-                   ( &Connection->SendRequest, &Bucket->Entry );
-               break;
-           } else {
-               TI_DbgPrint(DEBUG_TCP,
-                           ("Completing Send request: %x %x\n",
-                            Bucket->Request, Status));
-               Complete( Bucket->Request.RequestContext, Status, 0 );
-               exFreePool( Bucket );
-           }
-       }
-    }
+               Bucket->Status = SocketError;
+               Bucket->Information = 0;
+
+               InsertTailList(&Connection->CompletionQueue, &Bucket->Entry);
+            }
+
+            while (!IsListEmpty(&Connection->SendRequest))
+            {
+               Entry = RemoveHeadList( &Connection->SendRequest );
+
+               Bucket = CONTAINING_RECORD( Entry, TDI_BUCKET, Entry );
+
+               Bucket->Status = SocketError;
+               Bucket->Information = 0;
+
+               InsertTailList(&Connection->CompletionQueue, &Bucket->Entry);
+            }
 
-    if( NewState & SEL_FIN ) {
-        PLIST_ENTRY ListsToErase[4];
-        NTSTATUS    IrpStatus[4];
-        UINT i;
-
-       TI_DbgPrint(DEBUG_TCP, ("EOF From socket\n"));
-
-        ListsToErase[0] = &Connection->ReceiveRequest;
-        IrpStatus   [0] = STATUS_SUCCESS;
-        ListsToErase[1] = &Connection->ListenRequest;
-        IrpStatus   [1] = STATUS_UNSUCCESSFUL;
-        ListsToErase[2] = &Connection->ConnectRequest;
-        IrpStatus   [2] = STATUS_UNSUCCESSFUL;
-        ListsToErase[3] = 0;
-       IrpStatus   [3] = 0;
-
-        for( i = 0; ListsToErase[i]; i++ ) {
-            while( !IsListEmpty( ListsToErase[i] ) ) {
-                Entry = RemoveHeadList( ListsToErase[i] );
-                Bucket = CONTAINING_RECORD( Entry, TDI_BUCKET, Entry );
-                Complete = Bucket->Request.RequestNotifyObject;
-                Complete( Bucket->Request.RequestContext, IrpStatus[i], 0 );
-               exFreePool( Bucket );
+            while (!IsListEmpty(&Connection->ListenRequest))
+            {
+               Entry = RemoveHeadList( &Connection->ListenRequest );
+
+               Bucket = CONTAINING_RECORD( Entry, TDI_BUCKET, Entry );
+
+               Bucket->Status = SocketError;
+               Bucket->Information = 0;
+               DereferenceObject(Bucket->AssociatedEndpoint);
+
+               InsertTailList(&Connection->CompletionQueue, &Bucket->Entry);
+            }
+
+            while (!IsListEmpty(&Connection->ConnectRequest))
+            {
+               Entry = RemoveHeadList( &Connection->ConnectRequest );
+
+               Bucket = CONTAINING_RECORD( Entry, TDI_BUCKET, Entry );
+
+               Bucket->Status = SocketError;
+               Bucket->Information = 0;
+
+               InsertTailList(&Connection->CompletionQueue, &Bucket->Entry);
             }
+
+            Connection->SignalState = SEL_FIN;
         }
-    }
 
-    Connection->Signalled = FALSE;
+        /* Things that can happen when we try the initial connection */
+        if( Connection->SignalState & SEL_CONNECT ) {
+            while (!IsListEmpty(&Connection->ConnectRequest)) {
+               Entry = RemoveHeadList( &Connection->ConnectRequest );
+
+               Bucket = CONTAINING_RECORD( Entry, TDI_BUCKET, Entry );
+
+               Bucket->Status = STATUS_SUCCESS;
+               Bucket->Information = 0;
+
+               InsertTailList(&Connection->CompletionQueue, &Bucket->Entry);
+           }
+       }
+
+       if( Connection->SignalState & SEL_ACCEPT ) {
+           /* Handle readable on a listening socket --
+            * TODO: Implement filtering
+            */
+           TI_DbgPrint(DEBUG_TCP,("Accepting new connection on %x (Queue: %s)\n",
+                                  Connection,
+                                  IsListEmpty(&Connection->ListenRequest) ?
+                                  "empty" : "nonempty"));
+
+           while (!IsListEmpty(&Connection->ListenRequest)) {
+               PIO_STACK_LOCATION IrpSp;
+
+               Entry = RemoveHeadList( &Connection->ListenRequest );
+
+               Bucket = CONTAINING_RECORD( Entry, TDI_BUCKET, Entry );
+
+               Irp = Bucket->Request.RequestContext;
+               IrpSp = IoGetCurrentIrpStackLocation( Irp );
+
+               TI_DbgPrint(DEBUG_TCP,("Getting the socket\n"));
+
+               Status = TCPServiceListeningSocket
+                   ( Connection->AddressFile->Listener,
+                     Bucket->AssociatedEndpoint,
+                     (PTDI_REQUEST_KERNEL)&IrpSp->Parameters );
+
+               TI_DbgPrint(DEBUG_TCP,("Socket: Status: %x\n"));
+
+               if( Status == STATUS_PENDING ) {
+                   InsertHeadList( &Connection->ListenRequest, &Bucket->Entry );
+                   break;
+               } else {
+                   Bucket->Status = Status;
+                   Bucket->Information = 0;
+                   DereferenceObject(Bucket->AssociatedEndpoint);
+
+                   InsertTailList(&Connection->CompletionQueue, &Bucket->Entry);
+               }
+          }
+      }
+
+      /* Things that happen after we're connected */
+      if( Connection->SignalState & SEL_READ ) {
+          TI_DbgPrint(DEBUG_TCP,("Readable: irp list %s\n",
+                                 IsListEmpty(&Connection->ReceiveRequest) ?
+                                 "empty" : "nonempty"));
+
+           while (!IsListEmpty(&Connection->ReceiveRequest)) {
+               OSK_UINT RecvLen = 0, Received = 0;
+               PVOID RecvBuffer = 0;
+
+               Entry = RemoveHeadList( &Connection->ReceiveRequest );
+
+               Bucket = CONTAINING_RECORD( Entry, TDI_BUCKET, Entry );
+
+               Irp = Bucket->Request.RequestContext;
+               Mdl = Irp->MdlAddress;
+
+               TI_DbgPrint(DEBUG_TCP,
+                           ("Getting the user buffer from %x\n", Mdl));
+
+               NdisQueryBuffer( Mdl, &RecvBuffer, &RecvLen );
+
+               TI_DbgPrint(DEBUG_TCP,
+                           ("Reading %d bytes to %x\n", RecvLen, RecvBuffer));
+
+               TI_DbgPrint(DEBUG_TCP, ("Connection: %x\n", Connection));
+               TI_DbgPrint
+                   (DEBUG_TCP,
+                    ("Connection->SocketContext: %x\n",
+                     Connection->SocketContext));
+               TI_DbgPrint(DEBUG_TCP, ("RecvBuffer: %x\n", RecvBuffer));
+
+               Status = TCPTranslateError
+                    ( OskitTCPRecv( Connection->SocketContext,
+                                    RecvBuffer,
+                                    RecvLen,
+                                    &Received,
+                                    0 ) );
+
+               TI_DbgPrint(DEBUG_TCP,("TCP Bytes: %d\n", Received));
+
+               if( Status == STATUS_PENDING ) {
+                   InsertHeadList( &Connection->ReceiveRequest, &Bucket->Entry );
+                   break;
+               } else {
+                   TI_DbgPrint(DEBUG_TCP,
+                               ("Completing Receive request: %x %x\n",
+                                Bucket->Request, Status));
+
+                   Bucket->Status = Status;
+                   Bucket->Information = (Status == STATUS_SUCCESS) ? Received : 0;
+
+                   InsertTailList(&Connection->CompletionQueue, &Bucket->Entry);
+               }
+           }
+       }
+       if( Connection->SignalState & SEL_WRITE ) {
+           TI_DbgPrint(DEBUG_TCP,("Writeable: irp list %s\n",
+                                  IsListEmpty(&Connection->SendRequest) ?
+                                  "empty" : "nonempty"));
+
+           while (!IsListEmpty(&Connection->SendRequest)) {
+               OSK_UINT SendLen = 0, Sent = 0;
+               PVOID SendBuffer = 0;
+
+               Entry = RemoveHeadList( &Connection->SendRequest );
+
+               Bucket = CONTAINING_RECORD( Entry, TDI_BUCKET, Entry );
+
+               Irp = Bucket->Request.RequestContext;
+               Mdl = Irp->MdlAddress;
+
+               TI_DbgPrint(DEBUG_TCP,
+                           ("Getting the user buffer from %x\n", Mdl));
+
+               NdisQueryBuffer( Mdl, &SendBuffer, &SendLen );
+
+               TI_DbgPrint(DEBUG_TCP,
+                           ("Writing %d bytes to %x\n", SendLen, SendBuffer));
+
+               TI_DbgPrint(DEBUG_TCP, ("Connection: %x\n", Connection));
+               TI_DbgPrint
+                (DEBUG_TCP,
+                 ("Connection->SocketContext: %x\n",
+                  Connection->SocketContext));
+
+               Status = TCPTranslateError
+                   ( OskitTCPSend( Connection->SocketContext,
+                                   SendBuffer,
+                                   SendLen,
+                                   &Sent,
+                                   0 ) );
+
+               TI_DbgPrint(DEBUG_TCP,("TCP Bytes: %d\n", Sent));
+
+               if( Status == STATUS_PENDING ) {
+                   InsertHeadList( &Connection->SendRequest, &Bucket->Entry );
+                   break;
+               } else {
+                   TI_DbgPrint(DEBUG_TCP,
+                               ("Completing Send request: %x %x\n",
+                               Bucket->Request, Status));
+
+                   Bucket->Status = Status;
+                   Bucket->Information = (Status == STATUS_SUCCESS) ? Sent : 0;
+
+                   InsertTailList(&Connection->CompletionQueue, &Bucket->Entry);
+               }
+           }
+       }
+
+       ReferenceObject(Connection);
+       if (ClientInfo.Unlocked)
+       {
+           UnlockObjectFromDpcLevel(Connection);
+           KeReleaseSpinLock(&ClientInfo.Lock, ClientInfo.OldIrql);
+       }
+       else
+       {
+           UnlockObject(Connection, Connection->OldIrql);
+       }
+
+       while ((Entry = ExInterlockedRemoveHeadList(&Connection->CompletionQueue,
+                                                   &Connection->Lock)))
+       {
+           Bucket = CONTAINING_RECORD(Entry, TDI_BUCKET, Entry);
+           Complete = Bucket->Request.RequestNotifyObject;
+
+           Complete(Bucket->Request.RequestContext, Bucket->Status, Bucket->Information);
+
+           ExFreePoolWithTag(Bucket, TDI_BUCKET_TAG);
+       }
+
+       if (!ClientInfo.Unlocked)
+       {
+           LockObject(Connection, &OldIrql);
+       }
+       else
+       {
+           KeAcquireSpinLock(&ClientInfo.Lock, &ClientInfo.OldIrql);
+       }
+       DereferenceObject(Connection);
+
+       /* If the socket is dead, remove the reference we added for oskit */
+       if (Connection->SignalState & SEL_FIN)
+           DereferenceObject(Connection);
 }
 
-VOID DrainSignals() {
-    PCONNECTION_ENDPOINT Connection;
-    PLIST_ENTRY ListEntry;
+VOID ConnectionFree(PVOID Object) {
+    PCONNECTION_ENDPOINT Connection = Object;
+    KIRQL OldIrql;
 
-    while( !IsListEmpty( &SignalledConnections ) ) {
-       ListEntry = RemoveHeadList( &SignalledConnections );
-       Connection = CONTAINING_RECORD( ListEntry, CONNECTION_ENDPOINT,
-                                       SignalList );
-       HandleSignalledConnection( Connection, Connection->SignalState );
-    }
+    TI_DbgPrint(DEBUG_TCP, ("Freeing TCP Endpoint\n"));
+
+    TcpipAcquireSpinLock(&ConnectionEndpointListLock, &OldIrql);
+    RemoveEntryList(&Connection->ListEntry);
+    TcpipReleaseSpinLock(&ConnectionEndpointListLock, OldIrql);
+
+    ExFreePoolWithTag( Connection, CONN_ENDPT_TAG );
 }
 
 PCONNECTION_ENDPOINT TCPAllocateConnectionEndpoint( PVOID ClientContext ) {
     PCONNECTION_ENDPOINT Connection =
-       exAllocatePool(NonPagedPool, sizeof(CONNECTION_ENDPOINT));
+        ExAllocatePoolWithTag(NonPagedPool, sizeof(CONNECTION_ENDPOINT),
+                              CONN_ENDPT_TAG);
     if (!Connection)
-       return Connection;
+        return Connection;
 
     TI_DbgPrint(DEBUG_CPOINT, ("Connection point file object allocated at (0x%X).\n", Connection));
 
     RtlZeroMemory(Connection, sizeof(CONNECTION_ENDPOINT));
 
     /* Initialize spin lock that protects the connection endpoint file object */
-    TcpipInitializeSpinLock(&Connection->Lock);
+    KeInitializeSpinLock(&Connection->Lock);
     InitializeListHead(&Connection->ConnectRequest);
     InitializeListHead(&Connection->ListenRequest);
     InitializeListHead(&Connection->ReceiveRequest);
     InitializeListHead(&Connection->SendRequest);
+    InitializeListHead(&Connection->CompletionQueue);
 
     /* Save client context pointer */
     Connection->ClientContext = ClientContext;
 
-    return Connection;
-}
+    /* Add an extra reference for oskit */
+    Connection->RefCount = 2;
+    Connection->Free = ConnectionFree;
 
-VOID TCPFreeConnectionEndpoint( PCONNECTION_ENDPOINT Connection ) {
-    TCPClose( Connection );
-    exFreePool( Connection );
+    /* Add connection endpoint to global list */
+    ExInterlockedInsertTailList(&ConnectionEndpointListHead,
+                                &Connection->ListEntry,
+                                &ConnectionEndpointListLock);
+
+    return Connection;
 }
 
 NTSTATUS TCPSocket( PCONNECTION_ENDPOINT Connection,
-                   UINT Family, UINT Type, UINT Proto ) {
+                    UINT Family, UINT Type, UINT Proto ) {
     NTSTATUS Status;
+    KIRQL OldIrql;
+
+    LockObject(Connection, &OldIrql);
 
     TI_DbgPrint(DEBUG_TCP,("Called: Connection %x, Family %d, Type %d, "
-                          "Proto %d\n",
-                          Connection, Family, Type, Proto));
+                           "Proto %d\n",
+                           Connection, Family, Type, Proto));
 
-    TcpipRecursiveMutexEnter( &TCPLock, TRUE );
     Status = TCPTranslateError( OskitTCPSocket( Connection,
-                                               &Connection->SocketContext,
-                                               Family,
-                                               Type,
-                                               Proto ) );
+                                                &Connection->SocketContext,
+                                                Family,
+                                                Type,
+                                                Proto ) );
 
     ASSERT_KM_POINTER(Connection->SocketContext);
 
     TI_DbgPrint(DEBUG_TCP,("Connection->SocketContext %x\n",
-                          Connection->SocketContext));
+                           Connection->SocketContext));
 
-    TcpipRecursiveMutexLeave( &TCPLock );
+    UnlockObject(Connection, OldIrql);
 
     return Status;
 }
@@ -320,48 +384,46 @@ VOID TCPReceive(PIP_INTERFACE Interface, PIP_PACKET IPPacket)
  *     This is the low level interface for receiving TCP data
  */
 {
+    KIRQL OldIrql;
+
     TI_DbgPrint(DEBUG_TCP,("Sending packet %d (%d) to oskit\n",
-                          IPPacket->TotalSize,
-                          IPPacket->HeaderSize));
+                           IPPacket->TotalSize,
+                           IPPacket->HeaderSize));
 
-    TcpipRecursiveMutexEnter( &TCPLock, TRUE );
+    KeAcquireSpinLock(&ClientInfo.Lock, &OldIrql);
+    ClientInfo.Unlocked = TRUE;
+    ClientInfo.OldIrql = OldIrql;
 
     OskitTCPReceiveDatagram( IPPacket->Header,
-                            IPPacket->TotalSize,
-                            IPPacket->HeaderSize );
+                             IPPacket->TotalSize,
+                             IPPacket->HeaderSize );
 
-    DrainSignals();
-
-    TcpipRecursiveMutexLeave( &TCPLock );
+    ClientInfo.Unlocked = FALSE;
+    KeReleaseSpinLock(&ClientInfo.Lock, OldIrql);
 }
 
 /* event.c */
 int TCPSocketState( void *ClientData,
-                   void *WhichSocket,
-                   void *WhichConnection,
-                   OSK_UINT NewState );
+                    void *WhichSocket,
+                    void *WhichConnection,
+                    OSK_UINT NewState );
 
 int TCPPacketSend( void *ClientData,
-                  OSK_PCHAR Data,
-                  OSK_UINT Len );
+                   OSK_PCHAR Data,
+                   OSK_UINT Len );
 
 POSK_IFADDR TCPFindInterface( void *ClientData,
-                             OSK_UINT AddrType,
-                             OSK_UINT FindType,
-                             OSK_SOCKADDR *ReqAddr );
+                              OSK_UINT AddrType,
+                              OSK_UINT FindType,
+                              OSK_SOCKADDR *ReqAddr );
 
 NTSTATUS TCPMemStartup( void );
 void *TCPMalloc( void *ClientData,
-                OSK_UINT bytes, OSK_PCHAR file, OSK_UINT line );
+                 OSK_UINT bytes, OSK_PCHAR file, OSK_UINT line );
 void TCPFree( void *ClientData,
-             void *data, OSK_PCHAR file, OSK_UINT line );
+              void *data, OSK_PCHAR file, OSK_UINT line );
 void TCPMemShutdown( void );
 
-int TCPSleep( void *ClientData, void *token, int priority, char *msg,
-             int tmio );
-
-void TCPWakeup( void *ClientData, void *token );
-
 OSKITTCP_EVENT_HANDLERS EventHandlers = {
     NULL,             /* Client Data */
     TCPSocketState,   /* SocketState */
@@ -369,8 +431,8 @@ OSKITTCP_EVENT_HANDLERS EventHandlers = {
     TCPFindInterface, /* FindInterface */
     TCPMalloc,        /* Malloc */
     TCPFree,          /* Free */
-    TCPSleep,         /* Sleep */
-    TCPWakeup         /* Wakeup */
+    NULL,             /* Sleep */
+    NULL,             /* Wakeup */
 };
 
 static KEVENT TimerLoopEvent;
@@ -396,7 +458,7 @@ TimerThread(PVOID Context)
     while ( 1 ) {
         if (Next == NextFast) {
             NextFast += 2;
-        }
+       }
         if (Next == NextSlow) {
             NextSlow += 5;
         }
@@ -408,12 +470,7 @@ TimerThread(PVOID Context)
             PsTerminateSystemThread(Status);
         }
 
-        TcpipRecursiveMutexEnter( &TCPLock, TRUE );
         TimerOskitTCP( Next == NextFast, Next == NextSlow );
-        if (Next == NextSlow) {
-            DrainSignals();
-        }
-        TcpipRecursiveMutexLeave( &TCPLock );
 
         Current = Next;
         if (10 <= Current) {
@@ -433,7 +490,6 @@ StartTimer(VOID)
                          TimerThread, NULL);
 }
 
-
 NTSTATUS TCPStartup(VOID)
 /*
  * FUNCTION: Initializes the TCP subsystem
@@ -443,21 +499,20 @@ NTSTATUS TCPStartup(VOID)
 {
     NTSTATUS Status;
 
-    TcpipRecursiveMutexInit( &TCPLock );
-    ExInitializeFastMutex( &SleepingThreadsLock );
-    InitializeListHead( &SleepingThreadsList );
-    InitializeListHead( &SignalledConnections );
     Status = TCPMemStartup();
     if ( ! NT_SUCCESS(Status) ) {
-       return Status;
+        return Status;
     }
 
     Status = PortsStartup( &TCPPorts, 1, 0xfffe );
     if( !NT_SUCCESS(Status) ) {
-       TCPMemShutdown();
-       return Status;
+        TCPMemShutdown();
+        return Status;
     }
 
+    KeInitializeSpinLock(&ClientInfo.Lock);
+    ClientInfo.Unlocked = FALSE;
+
     RegisterOskitTCPEventHandlers( &EventHandlers );
     InitOskitTCP();
 
@@ -465,13 +520,13 @@ NTSTATUS TCPStartup(VOID)
     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 */
+        &TCPSegmentList,                /* Lookaside list */
+        NULL,                           /* Allocate routine */
+        NULL,                           /* Free routine */
+        0,                              /* Flags */
+        sizeof(TCP_SEGMENT),            /* Size of each entry */
+        'SPCT',                         /* Tag */
+        0);                             /* Depth */
 
     StartTimer();
 
@@ -491,7 +546,7 @@ NTSTATUS TCPShutdown(VOID)
     LARGE_INTEGER WaitForThread;
 
     if (!TCPInitialized)
-       return STATUS_SUCCESS;
+        return STATUS_SUCCESS;
 
     WaitForThread.QuadPart = -2500000; /* 250 ms */
     KeSetEvent(&TimerLoopEvent, IO_NO_INCREMENT, FALSE);
@@ -514,17 +569,29 @@ NTSTATUS TCPShutdown(VOID)
 }
 
 NTSTATUS TCPTranslateError( int OskitError ) {
-    NTSTATUS Status = STATUS_UNSUCCESSFUL;
+    NTSTATUS Status;
 
     switch( OskitError ) {
     case 0: Status = STATUS_SUCCESS; break;
-    case OSK_EADDRNOTAVAIL:
+    case OSK_EADDRNOTAVAIL: Status = STATUS_INVALID_ADDRESS; break;
     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;
+    case OSK_ECONNREFUSED: Status = STATUS_REMOTE_NOT_LISTENING; break;
+    case OSK_ECONNRESET:
+    case OSK_ECONNABORTED: Status = STATUS_REMOTE_DISCONNECT; break;
+    case OSK_EWOULDBLOCK:
+    case OSK_EINPROGRESS: Status = STATUS_PENDING; break;
+    case OSK_EINVAL: Status = STATUS_INVALID_PARAMETER; break;
+    case OSK_ENOMEM:
+    case OSK_ENOBUFS: Status = STATUS_INSUFFICIENT_RESOURCES; break;
+    case OSK_ESHUTDOWN: Status = STATUS_FILE_CLOSED; break;
+    case OSK_EMSGSIZE: Status = STATUS_BUFFER_TOO_SMALL; break;
+    case OSK_ETIMEDOUT: Status = STATUS_TIMEOUT; break;
+    case OSK_ENETUNREACH: Status = STATUS_NETWORK_UNREACHABLE; break;
+    case OSK_EFAULT: Status = STATUS_ACCESS_VIOLATION; break;
+    default:
+       DbgPrint("OskitTCP returned unhandled error code: %d\n", OskitError);
+       Status = STATUS_INVALID_CONNECTION;
+       break;
     }
 
     TI_DbgPrint(DEBUG_TCP,("Error %d -> %x\n", OskitError, Status));
@@ -542,59 +609,86 @@ NTSTATUS TCPConnect
     IP_ADDRESS RemoteAddress;
     USHORT RemotePort;
     PTDI_BUCKET Bucket;
+    PNEIGHBOR_CACHE_ENTRY NCE;
+    KIRQL OldIrql;
 
     TI_DbgPrint(DEBUG_TCP,("TCPConnect: Called\n"));
 
-    Bucket = exAllocatePool( NonPagedPool, sizeof(*Bucket) );
-    if( !Bucket ) return STATUS_NO_MEMORY;
+    Status = AddrBuildAddress
+        ((PTRANSPORT_ADDRESS)ConnInfo->RemoteAddress,
+         &RemoteAddress,
+         &RemotePort);
 
-    TcpipRecursiveMutexEnter( &TCPLock, TRUE );
+    if (!NT_SUCCESS(Status)) {
+        TI_DbgPrint(DEBUG_TCP, ("Could not AddrBuildAddress in TCPConnect\n"));
+        return Status;
+    }
 
     /* 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"));
-       TcpipRecursiveMutexLeave( &TCPLock );
-       return Status;
-    }
-
     AddressToConnect.sin_family = AF_INET;
     AddressToBind = AddressToConnect;
 
+    LockObject(Connection, &OldIrql);
+
+    if (!Connection->AddressFile)
+    {
+        UnlockObject(Connection, OldIrql);
+        return STATUS_INVALID_PARAMETER;
+    }
+
+    if (AddrIsUnspecified(&Connection->AddressFile->Address))
+    {
+        if (!(NCE = RouteGetRouteToDestination(&RemoteAddress)))
+        {
+            UnlockObject(Connection, OldIrql);
+            return STATUS_NETWORK_UNREACHABLE;
+        }
+
+        AddressToBind.sin_addr.s_addr = NCE->Interface->Unicast.Address.IPv4Address;
+    }
+    else
+    {
+        AddressToBind.sin_addr.s_addr = Connection->AddressFile->Address.Address.IPv4Address;
+    }
+
     Status = TCPTranslateError
         ( OskitTCPBind( Connection->SocketContext,
-                 Connection,
-                 &AddressToBind,
-                 sizeof(AddressToBind) ) );
+                        &AddressToBind,
+                        sizeof(AddressToBind) ) );
 
     if (NT_SUCCESS(Status)) {
         memcpy( &AddressToConnect.sin_addr,
-               &RemoteAddress.Address.IPv4Address,
-               sizeof(AddressToConnect.sin_addr) );
+                &RemoteAddress.Address.IPv4Address,
+                sizeof(AddressToConnect.sin_addr) );
         AddressToConnect.sin_port = RemotePort;
 
         Status = TCPTranslateError
-           ( OskitTCPConnect( Connection->SocketContext,
-                              Connection,
-                              &AddressToConnect,
-                              sizeof(AddressToConnect) ) );
+            ( OskitTCPConnect( Connection->SocketContext,
+                               &AddressToConnect,
+                               sizeof(AddressToConnect) ) );
+
+        if (Status == STATUS_PENDING)
+        {
+            Bucket = ExAllocatePoolWithTag( NonPagedPool, sizeof(*Bucket), TDI_BUCKET_TAG );
+            if( !Bucket )
+            {
+               UnlockObject(Connection, OldIrql);
+               return STATUS_NO_MEMORY;
+            }
+            
+            Bucket->Request.RequestNotifyObject = (PVOID)Complete;
+            Bucket->Request.RequestContext = Context;
+                       
+            InsertTailList( &Connection->ConnectRequest, &Bucket->Entry );
+        }
     }
 
-    TcpipRecursiveMutexLeave( &TCPLock );
+    UnlockObject(Connection, OldIrql);
 
     return Status;
 }
@@ -606,31 +700,20 @@ NTSTATUS TCPDisconnect
   PTDI_CONNECTION_INFORMATION ReturnInfo,
   PTCP_COMPLETION_ROUTINE Complete,
   PVOID Context ) {
-    NTSTATUS Status;
+    NTSTATUS Status = STATUS_INVALID_PARAMETER;
+    KIRQL OldIrql;
 
     TI_DbgPrint(DEBUG_TCP,("started\n"));
 
-    TcpipRecursiveMutexEnter( &TCPLock, TRUE );
+    LockObject(Connection, &OldIrql);
 
-    switch( Flags & (TDI_DISCONNECT_ABORT | TDI_DISCONNECT_RELEASE) ) {
-    case 0:
-    case TDI_DISCONNECT_ABORT:
-       Flags = 0;
-       break;
+    if (Flags & TDI_DISCONNECT_RELEASE)
+        Status = TCPTranslateError(OskitTCPDisconnect(Connection->SocketContext));
 
-    case TDI_DISCONNECT_ABORT | TDI_DISCONNECT_RELEASE:
-       Flags = 2;
-       break;
+    if ((Flags & TDI_DISCONNECT_ABORT) || !Flags)
+        Status = TCPTranslateError(OskitTCPShutdown(Connection->SocketContext, FWRITE | FREAD));
 
-    case TDI_DISCONNECT_RELEASE:
-       Flags = 1;
-       break;
-    }
-
-    Status = TCPTranslateError
-       ( OskitTCPShutdown( Connection->SocketContext, Flags ) );
-
-    TcpipRecursiveMutexLeave( &TCPLock );
+    UnlockObject(Connection, OldIrql);
 
     TI_DbgPrint(DEBUG_TCP,("finished %x\n", Status));
 
@@ -638,22 +721,46 @@ NTSTATUS TCPDisconnect
 }
 
 NTSTATUS TCPClose
-( PCONNECTION_ENDPOINT Connection ) {
+( PCONNECTION_ENDPOINT Connection )
+{
+    KIRQL OldIrql;
     NTSTATUS Status;
+    PVOID Socket;
+
+    /* We don't rely on SocketContext == NULL for socket
+     * closure anymore but we still need it to determine
+     * if we caused the closure
+     */
+    LockObject(Connection, &OldIrql);
+    Socket = Connection->SocketContext;
+    Connection->SocketContext = NULL;
+
+    /* Don't try to close again if the other side closed us already */
+    if (Connection->SignalState != SEL_FIN)
+    {
+       /* We need to close here otherwise oskit will never indicate
+        * SEL_FIN and we will never fully close the connection */
+       Status = TCPTranslateError( OskitTCPClose( Socket ) );
+
+       if (!NT_SUCCESS(Status))
+       {
+           Connection->SocketContext = Socket;
+           UnlockObject(Connection, OldIrql);
+           return Status;
+       }
+    }
+    else
+    {
+       /* We are already closed by the other end so return success */
+       Status = STATUS_SUCCESS;
+    }
 
-    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();
+    if (Connection->AddressFile)
+        DereferenceObject(Connection->AddressFile);
 
-    TcpipRecursiveMutexLeave( &TCPLock );
+    UnlockObject(Connection, OldIrql);
 
-    TI_DbgPrint(DEBUG_TCP,("TCPClose finished %x\n", Status));
+    DereferenceObject(Connection);
 
     return Status;
 }
@@ -666,55 +773,53 @@ NTSTATUS TCPReceiveData
   ULONG ReceiveFlags,
   PTCP_COMPLETION_ROUTINE Complete,
   PVOID Context ) {
-    OSK_PCHAR DataBuffer;
+    PVOID DataBuffer;
     UINT DataLen, Received = 0;
     NTSTATUS Status;
     PTDI_BUCKET Bucket;
+    KIRQL OldIrql;
 
     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));
 
+    LockObject(Connection, &OldIrql);
+
     Status = TCPTranslateError
-       ( OskitTCPRecv
-         ( Connection->SocketContext,
-           DataBuffer,
-           DataLen,
-           &Received,
-           ReceiveFlags ) );
+        ( 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"));
+        /* Freed in TCPSocketState */
+        Bucket = ExAllocatePoolWithTag( NonPagedPool, sizeof(*Bucket), TDI_BUCKET_TAG );
+        if( !Bucket ) {
+            TI_DbgPrint(DEBUG_TCP,("Failed to allocate bucket\n"));
+            UnlockObject(Connection, OldIrql);
+            return STATUS_NO_MEMORY;
+        }
+
+        Bucket->Request.RequestNotifyObject = Complete;
+        Bucket->Request.RequestContext = Context;
+        *BytesReceived = 0;
+
+        InsertTailList( &Connection->ReceiveRequest, &Bucket->Entry );
+        TI_DbgPrint(DEBUG_TCP,("Queued read irp\n"));
     } else {
-       TI_DbgPrint(DEBUG_TCP,("Got status %x, bytes %d\n", Status, Received));
-       *BytesReceived = Received;
+        TI_DbgPrint(DEBUG_TCP,("Got status %x, bytes %d\n", Status, Received));
+        *BytesReceived = Received;
     }
 
-    TcpipRecursiveMutexLeave( &TCPLock );
+    UnlockObject(Connection, OldIrql);
 
     TI_DbgPrint(DEBUG_TCP,("Status %x\n", Status));
 
@@ -732,61 +837,56 @@ NTSTATUS TCPSendData
     UINT Sent = 0;
     NTSTATUS Status;
     PTDI_BUCKET Bucket;
+    KIRQL OldIrql;
+
+    LockObject(Connection, &OldIrql);
 
     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));
+                           Connection->SocketContext));
 
     Status = TCPTranslateError
-       ( OskitTCPSend( Connection->SocketContext,
-                       (OSK_PCHAR)BufferData, SendLength,
-                       &Sent, 0 ) );
+        ( 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"));
+        /* Freed in TCPSocketState */
+        Bucket = ExAllocatePoolWithTag( NonPagedPool, sizeof(*Bucket), TDI_BUCKET_TAG );
+        if( !Bucket ) {
+            UnlockObject(Connection, OldIrql);
+            TI_DbgPrint(DEBUG_TCP,("Failed to allocate bucket\n"));
+            return STATUS_NO_MEMORY;
+        }
+        
+        Bucket->Request.RequestNotifyObject = Complete;
+        Bucket->Request.RequestContext = Context;
+        *BytesSent = 0;
+        
+        InsertTailList( &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;
+        TI_DbgPrint(DEBUG_TCP,("Got status %x, bytes %d\n", Status, Sent));
+        *BytesSent = Sent;
     }
 
-    TcpipRecursiveMutexLeave( &TCPLock );
+    UnlockObject(Connection, OldIrql);
 
     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 {
+        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;
@@ -798,60 +898,71 @@ VOID TCPFreePort( UINT Port ) {
     DeallocatePort( &TCPPorts, Port );
 }
 
-NTSTATUS TCPGetPeerAddress
+NTSTATUS TCPGetSockAddress
 ( PCONNECTION_ENDPOINT Connection,
-  PTRANSPORT_ADDRESS Address ) {
+  PTRANSPORT_ADDRESS Address,
+  BOOLEAN GetRemote ) {
     OSK_UINT LocalAddress, RemoteAddress;
     OSK_UI16 LocalPort, RemotePort;
     PTA_IP_ADDRESS AddressIP = (PTA_IP_ADDRESS)Address;
+    NTSTATUS Status;
+    KIRQL OldIrql;
+
+    LockObject(Connection, &OldIrql);
 
-    TcpipRecursiveMutexEnter( &TCPLock, TRUE );
+    Status = TCPTranslateError(OskitTCPGetAddress(Connection->SocketContext,
+                                                  &LocalAddress, &LocalPort,
+                                                  &RemoteAddress, &RemotePort));
 
-    OskitTCPGetAddress
-        ( Connection->SocketContext,
-          &LocalAddress, &LocalPort,
-          &RemoteAddress, &RemotePort );
+    UnlockObject(Connection, OldIrql);
+
+    if (!NT_SUCCESS(Status))
+        return Status;
 
     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 );
+    AddressIP->Address[0].Address[0].sin_port = GetRemote ? RemotePort : LocalPort;
+    AddressIP->Address[0].Address[0].in_addr = GetRemote ? RemoteAddress : LocalAddress;
 
-    return STATUS_SUCCESS;
+    return Status;
 }
 
-VOID TCPRemoveIRP( PCONNECTION_ENDPOINT Endpoint, PIRP Irp ) {
+BOOLEAN TCPRemoveIRP( PCONNECTION_ENDPOINT Endpoint, PIRP Irp ) {
     PLIST_ENTRY Entry;
     PLIST_ENTRY ListHead[4];
     KIRQL OldIrql;
     PTDI_BUCKET Bucket;
     UINT i = 0;
+    BOOLEAN Found = FALSE;
 
     ListHead[0] = &Endpoint->SendRequest;
     ListHead[1] = &Endpoint->ReceiveRequest;
     ListHead[2] = &Endpoint->ConnectRequest;
     ListHead[3] = &Endpoint->ListenRequest;
 
-    TcpipAcquireSpinLock( &Endpoint->Lock, &OldIrql );
+    LockObject(Endpoint, &OldIrql);
 
-    for( i = 0; i < 4; 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 );
-               exFreePool( Bucket );
-               break;
-           }
-       }
+    for( i = 0; i < 4; 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 );
+                ExFreePoolWithTag( Bucket, TDI_BUCKET_TAG );
+                Found = TRUE;
+                break;
+            }
+        }
     }
 
-    TcpipReleaseSpinLock( &Endpoint->Lock, OldIrql );
+    UnlockObject(Endpoint, OldIrql);
+
+    return Found;
 }
 
 /* EOF */