[AFD] Introduce and use pool tags. Thanks go to Arty for assisting me with this....
[reactos.git] / drivers / network / afd / afd / write.c
index 003ab2e..6056761 100644 (file)
@@ -7,8 +7,10 @@
  * UPDATE HISTORY:
  * 20040708 Created
  */
+
 #include "afd.h"
 
+static IO_COMPLETION_ROUTINE SendComplete;
 static NTSTATUS NTAPI SendComplete
 ( PDEVICE_OBJECT DeviceObject,
   PIRP Irp,
@@ -21,19 +23,23 @@ static NTSTATUS NTAPI SendComplete
     PAFD_SEND_INFO SendReq = NULL;
     PAFD_MAPBUF Map;
     UINT TotalBytesCopied = 0, TotalBytesProcessed = 0, SpaceAvail, i;
+    UINT SendLength, BytesCopied;
+    BOOLEAN HaltSendQueue;
+
+    UNREFERENCED_PARAMETER(DeviceObject);
 
     /*
      * The Irp parameter passed in is the IRP of the stream between AFD and
-     * TDI driver. It's not very usefull to us. We need the IRPs of the stream
+     * TDI driver. It's not very useful to us. We need the IRPs of the stream
      * between usermode and AFD. Those are chained from
      * FCB->PendingIrpList[FUNCTION_SEND] and you'll see them in the code
      * below as "NextIrp" ('cause they are the next usermode IRP to be
      * processed).
      */
 
-    AFD_DbgPrint(MID_TRACE,("Called, status %x, %d bytes used\n",
-                                                       Irp->IoStatus.Status,
-                                                       Irp->IoStatus.Information));
+    AFD_DbgPrint(MID_TRACE,("Called, status %x, %u bytes used\n",
+                            Irp->IoStatus.Status,
+                            Irp->IoStatus.Information));
 
     if( !SocketAcquireStateLock( FCB ) )
         return STATUS_FILE_CLOSED;
@@ -45,160 +51,200 @@ static NTSTATUS NTAPI SendComplete
     if( FCB->State == SOCKET_STATE_CLOSED ) {
         /* Cleanup our IRP queue because the FCB is being destroyed */
         while( !IsListEmpty( &FCB->PendingIrpList[FUNCTION_SEND] ) ) {
-              NextIrpEntry = RemoveHeadList(&FCB->PendingIrpList[FUNCTION_SEND]);
-              NextIrp = CONTAINING_RECORD(NextIrpEntry, IRP, Tail.Overlay.ListEntry);
-              NextIrpSp = IoGetCurrentIrpStackLocation( NextIrp );
-           SendReq = GetLockedData(NextIrp, NextIrpSp);
-              NextIrp->IoStatus.Status = STATUS_FILE_CLOSED;
-              NextIrp->IoStatus.Information = 0;
-              UnlockBuffers(SendReq->BufferArray, SendReq->BufferCount, FALSE);
-              if( NextIrp->MdlAddress ) UnlockRequest( NextIrp, IoGetCurrentIrpStackLocation( NextIrp ) );
-               (void)IoSetCancelRoutine(NextIrp, NULL);
-              IoCompleteRequest( NextIrp, IO_NETWORK_INCREMENT );
+            NextIrpEntry = RemoveHeadList(&FCB->PendingIrpList[FUNCTION_SEND]);
+            NextIrp = CONTAINING_RECORD(NextIrpEntry, IRP, Tail.Overlay.ListEntry);
+            NextIrpSp = IoGetCurrentIrpStackLocation( NextIrp );
+            SendReq = GetLockedData(NextIrp, NextIrpSp);
+            NextIrp->IoStatus.Status = STATUS_FILE_CLOSED;
+            NextIrp->IoStatus.Information = 0;
+            UnlockBuffers(SendReq->BufferArray, SendReq->BufferCount, FALSE);
+            if( NextIrp->MdlAddress ) UnlockRequest( NextIrp, IoGetCurrentIrpStackLocation( NextIrp ) );
+            (void)IoSetCancelRoutine(NextIrp, NULL);
+            IoCompleteRequest( NextIrp, IO_NETWORK_INCREMENT );
         }
-        
+
         RetryDisconnectCompletion(FCB);
-        
+
         SocketStateUnlock( FCB );
         return STATUS_FILE_CLOSED;
     }
 
     if( !NT_SUCCESS(Status) ) {
-               /* Complete all following send IRPs with error */
-
-               while( !IsListEmpty( &FCB->PendingIrpList[FUNCTION_SEND] ) ) {
-                       NextIrpEntry =
-                               RemoveHeadList(&FCB->PendingIrpList[FUNCTION_SEND]);
-                       NextIrp =
-                               CONTAINING_RECORD(NextIrpEntry, IRP, Tail.Overlay.ListEntry);
-                       NextIrpSp = IoGetCurrentIrpStackLocation( NextIrp );
-                       SendReq = GetLockedData(NextIrp, NextIrpSp);
-
-                       UnlockBuffers( SendReq->BufferArray,
-                                                  SendReq->BufferCount,
-                                                  FALSE );
-
-                       NextIrp->IoStatus.Status = Status;
-                       NextIrp->IoStatus.Information = 0;
-
-                       if ( NextIrp->MdlAddress ) UnlockRequest( NextIrp, IoGetCurrentIrpStackLocation( NextIrp ) );
-                        (void)IoSetCancelRoutine(NextIrp, NULL);
-                       IoCompleteRequest( NextIrp, IO_NETWORK_INCREMENT );
-               }
-        
+        /* Complete all following send IRPs with error */
+
+        while( !IsListEmpty( &FCB->PendingIrpList[FUNCTION_SEND] ) ) {
+            NextIrpEntry =
+                RemoveHeadList(&FCB->PendingIrpList[FUNCTION_SEND]);
+            NextIrp =
+                CONTAINING_RECORD(NextIrpEntry, IRP, Tail.Overlay.ListEntry);
+            NextIrpSp = IoGetCurrentIrpStackLocation( NextIrp );
+            SendReq = GetLockedData(NextIrp, NextIrpSp);
+
+            UnlockBuffers( SendReq->BufferArray,
+                           SendReq->BufferCount,
+                           FALSE );
+
+            NextIrp->IoStatus.Status = Status;
+            NextIrp->IoStatus.Information = 0;
+
+            if ( NextIrp->MdlAddress ) UnlockRequest( NextIrp, IoGetCurrentIrpStackLocation( NextIrp ) );
+            (void)IoSetCancelRoutine(NextIrp, NULL);
+            IoCompleteRequest( NextIrp, IO_NETWORK_INCREMENT );
+        }
+
         RetryDisconnectCompletion(FCB);
 
-               SocketStateUnlock( FCB );
+        SocketStateUnlock( FCB );
 
-               return STATUS_SUCCESS;
+        return STATUS_SUCCESS;
     }
 
     RtlMoveMemory( FCB->Send.Window,
-                                  FCB->Send.Window + FCB->Send.BytesUsed,
-                                  FCB->Send.BytesUsed - Irp->IoStatus.Information );
+                   FCB->Send.Window + Irp->IoStatus.Information,
+                   FCB->Send.BytesUsed - Irp->IoStatus.Information );
 
     TotalBytesProcessed = 0;
-    while (!IsListEmpty(&FCB->PendingIrpList[FUNCTION_SEND]) &&
-           TotalBytesProcessed != Irp->IoStatus.Information) {
-               NextIrpEntry =
-        RemoveHeadList(&FCB->PendingIrpList[FUNCTION_SEND]);
-               NextIrp =
-        CONTAINING_RECORD(NextIrpEntry, IRP, Tail.Overlay.ListEntry);
-               NextIrpSp = IoGetCurrentIrpStackLocation( NextIrp );
-               SendReq = GetLockedData(NextIrp, NextIrpSp);
-               Map = (PAFD_MAPBUF)(SendReq->BufferArray + SendReq->BufferCount);
-        
-        TotalBytesCopied = 0;
-        
-               for( i = 0; i < SendReq->BufferCount; i++ )
-                       TotalBytesCopied += SendReq->BufferArray[i].len;
-            
+    SendLength = Irp->IoStatus.Information;
+    HaltSendQueue = FALSE;
+    while (!IsListEmpty(&FCB->PendingIrpList[FUNCTION_SEND]) && SendLength > 0) {
+        NextIrpEntry = RemoveHeadList(&FCB->PendingIrpList[FUNCTION_SEND]);
+        NextIrp = CONTAINING_RECORD(NextIrpEntry, IRP, Tail.Overlay.ListEntry);
+        NextIrpSp = IoGetCurrentIrpStackLocation( NextIrp );
+        SendReq = GetLockedData(NextIrp, NextIrpSp);
+        Map = (PAFD_MAPBUF)(SendReq->BufferArray + SendReq->BufferCount);
+
+        TotalBytesCopied = (ULONG_PTR)NextIrp->Tail.Overlay.DriverContext[3];
+        ASSERT(TotalBytesCopied != 0);
+
+        /* If we didn't get enough, keep waiting */
+        if (TotalBytesCopied > SendLength)
+        {
+            /* Update the bytes left to copy */
+            TotalBytesCopied -= SendLength;
+            NextIrp->Tail.Overlay.DriverContext[3] = (PVOID)TotalBytesCopied;
+
+            /* Update the state variables */
+            FCB->Send.BytesUsed -= SendLength;
+            TotalBytesProcessed += SendLength;
+            SendLength = 0;
+
+            /* Pend the IRP */
+            InsertHeadList(&FCB->PendingIrpList[FUNCTION_SEND],
+                           &NextIrp->Tail.Overlay.ListEntry);
+            HaltSendQueue = TRUE;
+            break;
+        }
+
+        ASSERT(NextIrp->IoStatus.Information != 0);
+
         NextIrp->IoStatus.Status = Irp->IoStatus.Status;
-        NextIrp->IoStatus.Information = TotalBytesCopied;
-        
+
+        FCB->Send.BytesUsed -= TotalBytesCopied;
         TotalBytesProcessed += TotalBytesCopied;
-            
+        SendLength -= TotalBytesCopied;
+
         (void)IoSetCancelRoutine(NextIrp, NULL);
-            
+
         UnlockBuffers( SendReq->BufferArray,
                        SendReq->BufferCount,
                        FALSE );
-            
+
         if (NextIrp->MdlAddress) UnlockRequest(NextIrp, NextIrpSp);
-            
+
         IoCompleteRequest(NextIrp, IO_NETWORK_INCREMENT);
     }
 
-    ASSERT(TotalBytesProcessed == Irp->IoStatus.Information);
-    
-    FCB->Send.BytesUsed -= TotalBytesProcessed;
+    ASSERT(SendLength == 0);
 
-    while( !IsListEmpty( &FCB->PendingIrpList[FUNCTION_SEND] ) ) {
-               NextIrpEntry =
-                       RemoveHeadList(&FCB->PendingIrpList[FUNCTION_SEND]);
-               NextIrp =
-                       CONTAINING_RECORD(NextIrpEntry, IRP, Tail.Overlay.ListEntry);
-               NextIrpSp = IoGetCurrentIrpStackLocation( NextIrp );
-               SendReq = GetLockedData(NextIrp, NextIrpSp);
-               Map = (PAFD_MAPBUF)(SendReq->BufferArray + SendReq->BufferCount);
+   if ( !HaltSendQueue && !IsListEmpty( &FCB->PendingIrpList[FUNCTION_SEND] ) ) {
+        NextIrpEntry = FCB->PendingIrpList[FUNCTION_SEND].Flink;
+        NextIrp = CONTAINING_RECORD(NextIrpEntry, IRP, Tail.Overlay.ListEntry);
+        NextIrpSp = IoGetCurrentIrpStackLocation( NextIrp );
+        SendReq = GetLockedData(NextIrp, NextIrpSp);
+        Map = (PAFD_MAPBUF)(SendReq->BufferArray + SendReq->BufferCount);
 
-               AFD_DbgPrint(MID_TRACE,("SendReq @ %x\n", SendReq));
+        AFD_DbgPrint(MID_TRACE,("SendReq @ %p\n", SendReq));
 
-               SpaceAvail = FCB->Send.Size - FCB->Send.BytesUsed;
+        SpaceAvail = FCB->Send.Size - FCB->Send.BytesUsed;
         TotalBytesCopied = 0;
 
-               for( i = 0; i < SendReq->BufferCount; i++ ) {
-            if (SpaceAvail < SendReq->BufferArray[i].len)
-            {
-                InsertHeadList(&FCB->PendingIrpList[FUNCTION_SEND],
-                               &NextIrp->Tail.Overlay.ListEntry);
-                NextIrp = NULL;
-                break;
-            }
-                       Map[i].BufferAddress =
-                               MmMapLockedPages( Map[i].Mdl, KernelMode );
+        /* Count the total transfer size */
+        SendLength = 0;
+        for (i = 0; i < SendReq->BufferCount; i++)
+        {
+            SendLength += SendReq->BufferArray[i].len;
+        }
 
-                       RtlCopyMemory( FCB->Send.Window + FCB->Send.BytesUsed,
-                                                  Map[i].BufferAddress,
-                                                  SendReq->BufferArray[i].len );
+        /* Make sure we've got the space */
+        if (SendLength > SpaceAvail)
+        {
+           /* Blocking sockets have to wait here */
+           if (SendLength <= FCB->Send.Size && !((SendReq->AfdFlags & AFD_IMMEDIATE) || (FCB->NonBlocking)))
+           {
+               FCB->PollState &= ~AFD_EVENT_SEND;
 
-                       MmUnmapLockedPages( Map[i].BufferAddress, Map[i].Mdl );
+               NextIrp = NULL;
+           }
 
-                       TotalBytesCopied += SendReq->BufferArray[i].len;
-                       SpaceAvail -= SendReq->BufferArray[i].len;
-               }
+           /* Check if we can send anything */
+           if (SpaceAvail == 0)
+           {
+               FCB->PollState &= ~AFD_EVENT_SEND;
+
+               /* We should never be non-overlapped and get to this point */
+               ASSERT(SendReq->AfdFlags & AFD_OVERLAPPED);
+
+               NextIrp = NULL;
+           }
+        }
 
         if (NextIrp != NULL)
         {
-            FCB->Send.BytesUsed += TotalBytesCopied;
+            for( i = 0; i < SendReq->BufferCount; i++ ) {
+                BytesCopied = MIN(SendReq->BufferArray[i].len, SpaceAvail);
+
+                Map[i].BufferAddress =
+                   MmMapLockedPages( Map[i].Mdl, KernelMode );
+
+                RtlCopyMemory( FCB->Send.Window + FCB->Send.BytesUsed,
+                               Map[i].BufferAddress,
+                               BytesCopied );
+
+                MmUnmapLockedPages( Map[i].BufferAddress, Map[i].Mdl );
+
+                TotalBytesCopied += BytesCopied;
+                SpaceAvail -= BytesCopied;
+                FCB->Send.BytesUsed += BytesCopied;
+            }
+
+            NextIrp->IoStatus.Information = TotalBytesCopied;
+            NextIrp->Tail.Overlay.DriverContext[3] = (PVOID)NextIrp->IoStatus.Information;
         }
-        else
-            break;
     }
-    
-    if (FCB->Send.Size - FCB->Send.BytesUsed != 0 &&
-        !FCB->SendClosed)
+
+    if (FCB->Send.Size - FCB->Send.BytesUsed != 0 && !FCB->SendClosed &&
+        IsListEmpty(&FCB->PendingIrpList[FUNCTION_SEND]))
     {
-               FCB->PollState |= AFD_EVENT_SEND;
-               FCB->PollStatus[FD_WRITE_BIT] = STATUS_SUCCESS;
-               PollReeval( FCB->DeviceExt, FCB->FileObject );
+        FCB->PollState |= AFD_EVENT_SEND;
+        FCB->PollStatus[FD_WRITE_BIT] = STATUS_SUCCESS;
+        PollReeval( FCB->DeviceExt, FCB->FileObject );
     }
     else
     {
         FCB->PollState &= ~AFD_EVENT_SEND;
     }
 
+
     /* Some data is still waiting */
     if( FCB->Send.BytesUsed )
     {
-               Status = TdiSend( &FCB->SendIrp.InFlightRequest,
-                                                 FCB->Connection.Object,
-                                                 0,
-                                                 FCB->Send.Window,
-                                                 FCB->Send.BytesUsed,
-                                                 &FCB->SendIrp.Iosb,
-                                                 SendComplete,
-                                                 FCB );
+        Status = TdiSend( &FCB->SendIrp.InFlightRequest,
+                          FCB->Connection.Object,
+                          0,
+                          FCB->Send.Window,
+                          FCB->Send.BytesUsed,
+                          SendComplete,
+                          FCB );
     }
     else
     {
@@ -211,6 +257,7 @@ static NTSTATUS NTAPI SendComplete
     return STATUS_SUCCESS;
 }
 
+static IO_COMPLETION_ROUTINE PacketSocketSendComplete;
 static NTSTATUS NTAPI PacketSocketSendComplete
 ( PDEVICE_OBJECT DeviceObject,
   PIRP Irp,
@@ -220,9 +267,11 @@ static NTSTATUS NTAPI PacketSocketSendComplete
     PIRP NextIrp;
     PAFD_SEND_INFO SendReq;
 
-    AFD_DbgPrint(MID_TRACE,("Called, status %x, %d bytes used\n",
-                                                       Irp->IoStatus.Status,
-                                                       Irp->IoStatus.Information));
+    UNREFERENCED_PARAMETER(DeviceObject);
+
+    AFD_DbgPrint(MID_TRACE,("Called, status %x, %u bytes used\n",
+                            Irp->IoStatus.Status,
+                            Irp->IoStatus.Information));
 
     if( !SocketAcquireStateLock( FCB ) )
         return STATUS_FILE_CLOSED;
@@ -247,24 +296,24 @@ static NTSTATUS NTAPI PacketSocketSendComplete
         SocketStateUnlock( FCB );
         return STATUS_FILE_CLOSED;
     }
-    
+
     ASSERT(!IsListEmpty(&FCB->PendingIrpList[FUNCTION_SEND]));
-    
+
     /* TDI spec guarantees FIFO ordering on IRPs */
     NextIrpEntry = RemoveHeadList(&FCB->PendingIrpList[FUNCTION_SEND]);
     NextIrp = CONTAINING_RECORD(NextIrpEntry, IRP, Tail.Overlay.ListEntry);
-    
+
     SendReq = GetLockedData(NextIrp, IoGetCurrentIrpStackLocation(NextIrp));
-    
+
     NextIrp->IoStatus.Status = Irp->IoStatus.Status;
     NextIrp->IoStatus.Information = Irp->IoStatus.Information;
-    
+
     (void)IoSetCancelRoutine(NextIrp, NULL);
-    
+
     UnlockBuffers(SendReq->BufferArray, SendReq->BufferCount, FALSE);
-    
+
     UnlockRequest(NextIrp, IoGetCurrentIrpStackLocation(NextIrp));
-    
+
     IoCompleteRequest(NextIrp, IO_NETWORK_INCREMENT);
 
     FCB->PollState |= AFD_EVENT_SEND;
@@ -278,18 +327,23 @@ static NTSTATUS NTAPI PacketSocketSendComplete
 
 NTSTATUS NTAPI
 AfdConnectedSocketWriteData(PDEVICE_OBJECT DeviceObject, PIRP Irp,
-                                                       PIO_STACK_LOCATION IrpSp, BOOLEAN Short) {
+                            PIO_STACK_LOCATION IrpSp, BOOLEAN Short) {
     NTSTATUS Status = STATUS_SUCCESS;
     PFILE_OBJECT FileObject = IrpSp->FileObject;
     PAFD_FCB FCB = FileObject->FsContext;
     PAFD_SEND_INFO SendReq;
-    UINT TotalBytesCopied = 0, i, SpaceAvail = 0;
-    BOOLEAN NoSpace = FALSE;
+    UINT TotalBytesCopied = 0, i, SpaceAvail = 0, BytesCopied, SendLength;
+    KPROCESSOR_MODE LockMode;
+
+    UNREFERENCED_PARAMETER(DeviceObject);
+    UNREFERENCED_PARAMETER(Short);
 
-    AFD_DbgPrint(MID_TRACE,("Called on %x\n", FCB));
+    AFD_DbgPrint(MID_TRACE,("Called on %p\n", FCB));
 
     if( !SocketAcquireStateLock( FCB ) ) return LostSocket( Irp );
 
+    FCB->EventSelectDisabled &= ~AFD_EVENT_SEND;
+
     if( FCB->Flags & AFD_ENDPOINT_CONNECTIONLESS )
     {
         PAFD_SEND_INFO_UDP SendReq;
@@ -303,43 +357,51 @@ AfdConnectedSocketWriteData(PDEVICE_OBJECT DeviceObject, PIRP Irp,
                                            0 );
         }
 
-        if( !(SendReq = LockRequest( Irp, IrpSp )) )
+        if( !(SendReq = LockRequest( Irp, IrpSp, FALSE, &LockMode )) )
             return UnlockAndMaybeComplete( FCB, STATUS_NO_MEMORY, Irp, 0 );
 
         /* Must lock buffers before handing off user data */
         SendReq->BufferArray = LockBuffers( SendReq->BufferArray,
                                             SendReq->BufferCount,
                                             NULL, NULL,
-                                            FALSE, FALSE );
+                                            FALSE, FALSE, LockMode );
 
-               if( !SendReq->BufferArray ) {
-                       return UnlockAndMaybeComplete( FCB, STATUS_ACCESS_VIOLATION,
+        if( !SendReq->BufferArray ) {
+            return UnlockAndMaybeComplete( FCB, STATUS_ACCESS_VIOLATION,
                                            Irp, 0 );
-               }
+        }
 
         Status = TdiBuildConnectionInfo( &TargetAddress, FCB->RemoteAddress );
 
         if( NT_SUCCESS(Status) ) {
-            FCB->EventSelectDisabled &= ~AFD_EVENT_SEND;
             FCB->PollState &= ~AFD_EVENT_SEND;
-            
+
             Status = QueueUserModeIrp(FCB, Irp, FUNCTION_SEND);
             if (Status == STATUS_PENDING)
             {
-                TdiSendDatagram(&FCB->SendIrp.InFlightRequest,
-                                FCB->AddressFile.Object,
-                                SendReq->BufferArray[0].buf,
-                                SendReq->BufferArray[0].len,
-                                TargetAddress,
-                                &FCB->SendIrp.Iosb,
-                                PacketSocketSendComplete,
-                                FCB);
+                Status = TdiSendDatagram(&FCB->SendIrp.InFlightRequest,
+                                         FCB->AddressFile.Object,
+                                         SendReq->BufferArray[0].buf,
+                                         SendReq->BufferArray[0].len,
+                                         TargetAddress,
+                                         PacketSocketSendComplete,
+                                         FCB);
+                if (Status != STATUS_PENDING)
+                {
+                    NT_VERIFY(RemoveHeadList(&FCB->PendingIrpList[FUNCTION_SEND]) == &Irp->Tail.Overlay.ListEntry);
+                    Irp->IoStatus.Status = Status;
+                    Irp->IoStatus.Information = 0;
+                    (void)IoSetCancelRoutine(Irp, NULL);
+                    UnlockBuffers(SendReq->BufferArray, SendReq->BufferCount, FALSE);
+                    UnlockRequest(Irp, IoGetCurrentIrpStackLocation(Irp));
+                    IoCompleteRequest(Irp, IO_NETWORK_INCREMENT);
+                }
             }
-            
-            ExFreePool( TargetAddress );
-            
+
+            ExFreePoolWithTag(TargetAddress, TAG_AFD_TDI_CONNECTION_INFORMATION);
+
             SocketStateUnlock(FCB);
-            
+
             return STATUS_PENDING;
         }
         else
@@ -348,7 +410,7 @@ AfdConnectedSocketWriteData(PDEVICE_OBJECT DeviceObject, PIRP Irp,
             return UnlockAndMaybeComplete( FCB, Status, Irp, 0 );
         }
     }
-    
+
     if (FCB->PollState & AFD_EVENT_CLOSE)
     {
         AFD_DbgPrint(MIN_TRACE,("Connection reset by remote peer\n"));
@@ -356,15 +418,15 @@ AfdConnectedSocketWriteData(PDEVICE_OBJECT DeviceObject, PIRP Irp,
         /* This is an unexpected remote disconnect */
         return UnlockAndMaybeComplete(FCB, FCB->PollStatus[FD_CLOSE_BIT], Irp, 0);
     }
-    
+
     if (FCB->PollState & AFD_EVENT_ABORT)
     {
         AFD_DbgPrint(MIN_TRACE,("Connection aborted\n"));
-        
+
         /* This is an abortive socket closure on our side */
         return UnlockAndMaybeComplete(FCB, FCB->PollStatus[FD_CLOSE_BIT], Irp, 0);
     }
-    
+
     if (FCB->SendClosed)
     {
         AFD_DbgPrint(MIN_TRACE,("No more sends\n"));
@@ -373,74 +435,93 @@ AfdConnectedSocketWriteData(PDEVICE_OBJECT DeviceObject, PIRP Irp,
         return UnlockAndMaybeComplete(FCB, STATUS_FILE_CLOSED, Irp, 0);
     }
 
-    if( !(SendReq = LockRequest( Irp, IrpSp )) )
-               return UnlockAndMaybeComplete
-                       ( FCB, STATUS_NO_MEMORY, Irp, TotalBytesCopied );
+    if( !(SendReq = LockRequest( Irp, IrpSp, FALSE, &LockMode )) )
+        return UnlockAndMaybeComplete
+            ( FCB, STATUS_NO_MEMORY, Irp, 0 );
 
     SendReq->BufferArray = LockBuffers( SendReq->BufferArray,
-                                                                               SendReq->BufferCount,
-                                                                               NULL, NULL,
-                                                                               FALSE, FALSE );
+                                        SendReq->BufferCount,
+                                        NULL, NULL,
+                                        FALSE, FALSE, LockMode );
 
     if( !SendReq->BufferArray ) {
         return UnlockAndMaybeComplete( FCB, STATUS_ACCESS_VIOLATION,
                                        Irp, 0 );
     }
 
-    AFD_DbgPrint(MID_TRACE,("Socket state %d\n", FCB->State));
+    AFD_DbgPrint(MID_TRACE,("Socket state %u\n", FCB->State));
 
     if( FCB->State != SOCKET_STATE_CONNECTED ) {
-               if( (SendReq->AfdFlags & AFD_IMMEDIATE) || (FCB->NonBlocking) ) {
-                       AFD_DbgPrint(MID_TRACE,("Nonblocking\n"));
-                       UnlockBuffers( SendReq->BufferArray, SendReq->BufferCount, FALSE );
-                       return UnlockAndMaybeComplete
-                               ( FCB, STATUS_CANT_WAIT, Irp, 0 );
-               } else {
-                       AFD_DbgPrint(MID_TRACE,("Queuing request\n"));
-                       return LeaveIrpUntilLater( FCB, Irp, FUNCTION_SEND );
-               }
+        AFD_DbgPrint(MID_TRACE,("Socket not connected\n"));
+        UnlockBuffers( SendReq->BufferArray, SendReq->BufferCount, FALSE );
+        return UnlockAndMaybeComplete( FCB, STATUS_INVALID_CONNECTION, Irp, 0 );
     }
 
-    AFD_DbgPrint(MID_TRACE,("FCB->Send.BytesUsed = %d\n",
-                                                       FCB->Send.BytesUsed));
+    AFD_DbgPrint(MID_TRACE,("FCB->Send.BytesUsed = %u\n",
+                            FCB->Send.BytesUsed));
 
     SpaceAvail = FCB->Send.Size - FCB->Send.BytesUsed;
-        
-    AFD_DbgPrint(MID_TRACE,("We can accept %d bytes\n",
+
+    AFD_DbgPrint(MID_TRACE,("We can accept %u bytes\n",
                             SpaceAvail));
-    
-    for( i = 0; FCB->Send.BytesUsed < FCB->Send.Size &&
-        i < SendReq->BufferCount; i++ ) {
-        
-        if (SpaceAvail < SendReq->BufferArray[i].len)
+
+    /* Count the total transfer size */
+    SendLength = 0;
+    for (i = 0; i < SendReq->BufferCount; i++)
+    {
+        SendLength += SendReq->BufferArray[i].len;
+    }
+
+    /* Make sure we've got the space */
+    if (SendLength > SpaceAvail)
+    {
+        /* Blocking sockets have to wait here */
+        if (SendLength <= FCB->Send.Size && !((SendReq->AfdFlags & AFD_IMMEDIATE) || (FCB->NonBlocking)))
         {
-            if (TotalBytesCopied + SendReq->BufferArray[i].len > FCB->Send.Size)
+            FCB->PollState &= ~AFD_EVENT_SEND;
+            return LeaveIrpUntilLater(FCB, Irp, FUNCTION_SEND);
+        }
+
+        /* Check if we can send anything */
+        if (SpaceAvail == 0)
+        {
+            FCB->PollState &= ~AFD_EVENT_SEND;
+
+            /* Non-overlapped sockets will fail if we can send nothing */
+            if (!(SendReq->AfdFlags & AFD_OVERLAPPED))
             {
                 UnlockBuffers( SendReq->BufferArray, SendReq->BufferCount, FALSE );
-                
-                return UnlockAndMaybeComplete(FCB, STATUS_BUFFER_OVERFLOW, Irp, 0);
+                return UnlockAndMaybeComplete( FCB, STATUS_CANT_WAIT, Irp, 0 );
+            }
+            else
+            {
+                /* Overlapped sockets just pend */
+                return LeaveIrpUntilLater(FCB, Irp, FUNCTION_SEND);
             }
-            SpaceAvail += TotalBytesCopied;
-            NoSpace = TRUE;
-            break;
         }
-        
-        AFD_DbgPrint(MID_TRACE,("Copying Buffer %d, %x:%d to %x\n",
+    }
+
+    for ( i = 0; SpaceAvail > 0 && i < SendReq->BufferCount; i++ )
+    {
+        BytesCopied = MIN(SendReq->BufferArray[i].len, SpaceAvail);
+
+        AFD_DbgPrint(MID_TRACE,("Copying Buffer %u, %p:%u to %p\n",
                                 i,
                                 SendReq->BufferArray[i].buf,
-                                SendReq->BufferArray[i].len,
+                                BytesCopied,
                                 FCB->Send.Window + FCB->Send.BytesUsed));
-        
-        RtlCopyMemory( FCB->Send.Window + FCB->Send.BytesUsed,
+
+        RtlCopyMemory(FCB->Send.Window + FCB->Send.BytesUsed,
                       SendReq->BufferArray[i].buf,
-                      SendReq->BufferArray[i].len );
-        
-        TotalBytesCopied += SendReq->BufferArray[i].len;
-        SpaceAvail -= SendReq->BufferArray[i].len;
+                      BytesCopied);
+
+        TotalBytesCopied += BytesCopied;
+        SpaceAvail -= BytesCopied;
+        FCB->Send.BytesUsed += BytesCopied;
     }
-    
-    FCB->EventSelectDisabled &= ~AFD_EVENT_SEND;
-    
+
+    Irp->IoStatus.Information = TotalBytesCopied;
+
     if( TotalBytesCopied == 0 ) {
         AFD_DbgPrint(MID_TRACE,("Empty send\n"));
         UnlockBuffers( SendReq->BufferArray, SendReq->BufferCount, FALSE );
@@ -451,70 +532,58 @@ AfdConnectedSocketWriteData(PDEVICE_OBJECT DeviceObject, PIRP Irp,
     if (SpaceAvail)
     {
         FCB->PollState |= AFD_EVENT_SEND;
-               FCB->PollStatus[FD_WRITE_BIT] = STATUS_SUCCESS;
-               PollReeval( FCB->DeviceExt, FCB->FileObject );
+        FCB->PollStatus[FD_WRITE_BIT] = STATUS_SUCCESS;
+        PollReeval( FCB->DeviceExt, FCB->FileObject );
     }
     else
     {
         FCB->PollState &= ~AFD_EVENT_SEND;
     }
 
-    if (!NoSpace)
-    {
-        FCB->Send.BytesUsed += TotalBytesCopied;
-        AFD_DbgPrint(MID_TRACE,("Copied %d bytes\n", TotalBytesCopied));
-        
-        Status = QueueUserModeIrp(FCB, Irp, FUNCTION_SEND);
-        if (Status == STATUS_PENDING && !FCB->SendIrp.InFlightRequest)
-        {
-            TdiSend(&FCB->SendIrp.InFlightRequest,
-                    FCB->Connection.Object,
-                    0,
-                    FCB->Send.Window,
-                    FCB->Send.BytesUsed,
-                    &FCB->SendIrp.Iosb,
-                    SendComplete,
-                    FCB);
-        }
-        SocketStateUnlock(FCB);
-        
-        return STATUS_PENDING;
-    }
-    else
+    /* We use the IRP tail for some temporary storage here */
+    Irp->Tail.Overlay.DriverContext[3] = (PVOID)Irp->IoStatus.Information;
+
+    Status = QueueUserModeIrp(FCB, Irp, FUNCTION_SEND);
+    if (Status == STATUS_PENDING && !FCB->SendIrp.InFlightRequest)
     {
-        FCB->PollState &= ~AFD_EVENT_SEND;
-        if( (SendReq->AfdFlags & AFD_IMMEDIATE) || (FCB->NonBlocking) ) {
-            AFD_DbgPrint(MID_TRACE,("Nonblocking\n"));
-            UnlockBuffers( SendReq->BufferArray, SendReq->BufferCount, FALSE );
-            return UnlockAndMaybeComplete
-                       ( FCB, STATUS_CANT_WAIT, Irp, 0 );
-        } else {
-            AFD_DbgPrint(MID_TRACE,("Queuing request\n"));
-            return LeaveIrpUntilLater( FCB, Irp, FUNCTION_SEND );
-        }
+        TdiSend(&FCB->SendIrp.InFlightRequest,
+                FCB->Connection.Object,
+                0,
+                FCB->Send.Window,
+                FCB->Send.BytesUsed,
+                SendComplete,
+                FCB);
     }
+
+    SocketStateUnlock(FCB);
+
+    return STATUS_PENDING;
 }
 
 NTSTATUS NTAPI
 AfdPacketSocketWriteData(PDEVICE_OBJECT DeviceObject, PIRP Irp,
-                                                PIO_STACK_LOCATION IrpSp) {
+                         PIO_STACK_LOCATION IrpSp) {
     NTSTATUS Status = STATUS_SUCCESS;
     PTDI_CONNECTION_INFORMATION TargetAddress;
     PFILE_OBJECT FileObject = IrpSp->FileObject;
     PAFD_FCB FCB = FileObject->FsContext;
     PAFD_SEND_INFO_UDP SendReq;
+    KPROCESSOR_MODE LockMode;
 
-    AFD_DbgPrint(MID_TRACE,("Called on %x\n", FCB));
+    UNREFERENCED_PARAMETER(DeviceObject);
+
+    AFD_DbgPrint(MID_TRACE,("Called on %p\n", FCB));
 
     if( !SocketAcquireStateLock( FCB ) ) return LostSocket( Irp );
 
+    FCB->EventSelectDisabled &= ~AFD_EVENT_SEND;
+
     /* Check that the socket is bound */
     if( FCB->State != SOCKET_STATE_BOUND &&
         FCB->State != SOCKET_STATE_CREATED)
     {
         AFD_DbgPrint(MIN_TRACE,("Invalid socket state\n"));
-               return UnlockAndMaybeComplete
-                       ( FCB, STATUS_INVALID_PARAMETER, Irp, 0 );
+        return UnlockAndMaybeComplete(FCB, STATUS_INVALID_PARAMETER, Irp, 0);
     }
 
     if (FCB->SendClosed)
@@ -523,20 +592,23 @@ AfdPacketSocketWriteData(PDEVICE_OBJECT DeviceObject, PIRP Irp,
         return UnlockAndMaybeComplete(FCB, STATUS_FILE_CLOSED, Irp, 0);
     }
 
-    if( !(SendReq = LockRequest( Irp, IrpSp )) )
-               return UnlockAndMaybeComplete
-                       ( FCB, STATUS_NO_MEMORY, Irp, 0 );
-    
+    if( !(SendReq = LockRequest( Irp, IrpSp, FALSE, &LockMode )) )
+        return UnlockAndMaybeComplete(FCB, STATUS_NO_MEMORY, Irp, 0);
+
     if (FCB->State == SOCKET_STATE_CREATED)
     {
-        if( FCB->LocalAddress ) ExFreePool( FCB->LocalAddress );
+        if (FCB->LocalAddress)
+        {
+            ExFreePoolWithTag(FCB->LocalAddress, TAG_AFD_TRANSPORT_ADDRESS);
+        }
+
         FCB->LocalAddress =
         TaBuildNullTransportAddress( ((PTRANSPORT_ADDRESS)SendReq->TdiConnection.RemoteAddress)->
                                       Address[0].AddressType );
-        
+
         if( FCB->LocalAddress ) {
-            Status = WarmSocketForBind( FCB );
-            
+            Status = WarmSocketForBind( FCB, AFD_SHARE_WILDCARD );
+
             if( NT_SUCCESS(Status) )
                 FCB->State = SOCKET_STATE_BOUND;
             else
@@ -549,45 +621,53 @@ AfdPacketSocketWriteData(PDEVICE_OBJECT DeviceObject, PIRP Irp,
     SendReq->BufferArray = LockBuffers( SendReq->BufferArray,
                                         SendReq->BufferCount,
                                         NULL, NULL,
-                                        FALSE, FALSE );
+                                        FALSE, FALSE, LockMode );
 
     if( !SendReq->BufferArray )
-               return UnlockAndMaybeComplete( FCB, STATUS_ACCESS_VIOLATION,
+        return UnlockAndMaybeComplete( FCB, STATUS_ACCESS_VIOLATION,
                                        Irp, 0 );
 
     AFD_DbgPrint
-               (MID_TRACE,("RemoteAddress #%d Type %d\n",
-                                       ((PTRANSPORT_ADDRESS)SendReq->TdiConnection.RemoteAddress)->
-                                       TAAddressCount,
-                                       ((PTRANSPORT_ADDRESS)SendReq->TdiConnection.RemoteAddress)->
-                                       Address[0].AddressType));
+        (MID_TRACE,("RemoteAddress #%d Type %u\n",
+                    ((PTRANSPORT_ADDRESS)SendReq->TdiConnection.RemoteAddress)->
+                    TAAddressCount,
+                    ((PTRANSPORT_ADDRESS)SendReq->TdiConnection.RemoteAddress)->
+                    Address[0].AddressType));
 
     Status = TdiBuildConnectionInfo( &TargetAddress,
-                                                       ((PTRANSPORT_ADDRESS)SendReq->TdiConnection.RemoteAddress) );
+                            ((PTRANSPORT_ADDRESS)SendReq->TdiConnection.RemoteAddress) );
 
     /* Check the size of the Address given ... */
 
     if( NT_SUCCESS(Status) ) {
-        FCB->EventSelectDisabled &= ~AFD_EVENT_SEND;
-               FCB->PollState &= ~AFD_EVENT_SEND;
+        FCB->PollState &= ~AFD_EVENT_SEND;
 
         Status = QueueUserModeIrp(FCB, Irp, FUNCTION_SEND);
         if (Status == STATUS_PENDING)
         {
-            TdiSendDatagram(&FCB->SendIrp.InFlightRequest,
-                            FCB->AddressFile.Object,
-                            SendReq->BufferArray[0].buf,
-                            SendReq->BufferArray[0].len,
-                            TargetAddress,
-                            &FCB->SendIrp.Iosb,
-                            PacketSocketSendComplete,
-                            FCB);
+            Status = TdiSendDatagram(&FCB->SendIrp.InFlightRequest,
+                                     FCB->AddressFile.Object,
+                                     SendReq->BufferArray[0].buf,
+                                     SendReq->BufferArray[0].len,
+                                     TargetAddress,
+                                     PacketSocketSendComplete,
+                                     FCB);
+            if (Status != STATUS_PENDING)
+            {
+                NT_VERIFY(RemoveHeadList(&FCB->PendingIrpList[FUNCTION_SEND]) == &Irp->Tail.Overlay.ListEntry);
+                Irp->IoStatus.Status = Status;
+                Irp->IoStatus.Information = 0;
+                (void)IoSetCancelRoutine(Irp, NULL);
+                UnlockBuffers(SendReq->BufferArray, SendReq->BufferCount, FALSE);
+                UnlockRequest(Irp, IoGetCurrentIrpStackLocation(Irp));
+                IoCompleteRequest(Irp, IO_NETWORK_INCREMENT);
+            }
         }
 
-               ExFreePool( TargetAddress );
-        
+        ExFreePoolWithTag(TargetAddress, TAG_AFD_TDI_CONNECTION_INFORMATION);
+
         SocketStateUnlock(FCB);
-        
+
         return STATUS_PENDING;
     }
     else
@@ -596,4 +676,3 @@ AfdPacketSocketWriteData(PDEVICE_OBJECT DeviceObject, PIRP Irp,
         return UnlockAndMaybeComplete( FCB, Status, Irp, 0 );
     }
 }
-