From e84bd0226834f9f5f5dbbbe8435081e26d832ebb Mon Sep 17 00:00:00 2001
From: Cameron Gutman <aicommander@gmail.com>
Date: Mon, 30 Mar 2009 22:20:12 +0000
Subject: [PATCH]  - Validate the output buffer size before writing to it  -
 Fix some potential memory leaks  - Lock the FCB in AfdCloseSocket

svn path=/trunk/; revision=40300
---
 reactos/drivers/network/afd/afd/info.c   |  6 +++--
 reactos/drivers/network/afd/afd/listen.c | 31 ++++++++++++++++++++++++
 reactos/drivers/network/afd/afd/main.c   | 19 ++++++++++-----
 3 files changed, 48 insertions(+), 8 deletions(-)

diff --git a/reactos/drivers/network/afd/afd/info.c b/reactos/drivers/network/afd/afd/info.c
index a47164365e1..561235634d3 100644
--- a/reactos/drivers/network/afd/afd/info.c
+++ b/reactos/drivers/network/afd/afd/info.c
@@ -154,8 +154,10 @@ AfdGetPeerName( PDEVICE_OBJECT DeviceObject, PIRP Irp,
 
                 if (NT_SUCCESS(Status))
                 {
-                    RtlCopyMemory(Irp->UserBuffer, ConnInfo->RemoteAddress, TaLengthOfTransportAddress
-                                                                                  (ConnInfo->RemoteAddress));
+                    if (IrpSp->Parameters.DeviceIoControl.OutputBufferLength >= TaLengthOfTransportAddress(ConnInfo->RemoteAddress))
+                        RtlCopyMemory(Irp->UserBuffer, ConnInfo->RemoteAddress, TaLengthOfTransportAddress(ConnInfo->RemoteAddress));
+                    else
+                        Status = STATUS_BUFFER_TOO_SMALL;
                 }
             }
          }
diff --git a/reactos/drivers/network/afd/afd/listen.c b/reactos/drivers/network/afd/afd/listen.c
index 0f840d1e179..ec15c0cc3e3 100644
--- a/reactos/drivers/network/afd/afd/listen.c
+++ b/reactos/drivers/network/afd/afd/listen.c
@@ -215,7 +215,21 @@ NTSTATUS AfdListenSocket(PDEVICE_OBJECT DeviceObject, PIRP Irp,
 	  FCB->LocalAddress->Address[0].AddressType );
 
     if( !FCB->ListenIrp.ConnectionReturnInfo || !FCB->ListenIrp.ConnectionCallInfo )
+    {
+        if (FCB->ListenIrp.ConnectionReturnInfo)
+        {
+            ExFreePool(FCB->ListenIrp.ConnectionReturnInfo);
+            FCB->ListenIrp.ConnectionReturnInfo = NULL;
+        }
+
+        if (FCB->ListenIrp.ConnectionCallInfo)
+        {
+            ExFreePool(FCB->ListenIrp.ConnectionCallInfo);
+            FCB->ListenIrp.ConnectionCallInfo = NULL;
+        }
+
 	return UnlockAndMaybeComplete( FCB, STATUS_NO_MEMORY, Irp, 0 );
+    }
 
     FCB->State = SOCKET_STATE_LISTENING;
 
@@ -230,6 +244,9 @@ NTSTATUS AfdListenSocket(PDEVICE_OBJECT DeviceObject, PIRP Irp,
     if( Status == STATUS_PENDING )
 	Status = STATUS_SUCCESS;
 
+    if (NT_SUCCESS(Status))
+        FCB->NeedsNewListen = FALSE;
+
     AFD_DbgPrint(MID_TRACE,("Returning %x\n", Status));
     return UnlockAndMaybeComplete( FCB, Status, Irp, 0 );
 }
@@ -298,7 +315,21 @@ NTSTATUS AfdAccept( PDEVICE_OBJECT DeviceObject, PIRP Irp,
                 FCB->LocalAddress->Address[0].AddressType );
 
             if( !FCB->ListenIrp.ConnectionReturnInfo || !FCB->ListenIrp.ConnectionCallInfo )
+            {
+                if (FCB->ListenIrp.ConnectionReturnInfo)
+                {
+                    ExFreePool(FCB->ListenIrp.ConnectionReturnInfo);
+                    FCB->ListenIrp.ConnectionReturnInfo = NULL;
+                }
+
+                if (FCB->ListenIrp.ConnectionCallInfo)
+                {
+                    ExFreePool(FCB->ListenIrp.ConnectionCallInfo);
+                    FCB->ListenIrp.ConnectionCallInfo = NULL;
+                }
+
 	        return UnlockAndMaybeComplete( FCB, STATUS_NO_MEMORY, Irp, 0 );
+            }
 
 	    Status = TdiListen( &FCB->ListenIrp.InFlightRequest,
 				FCB->Connection.Object,
diff --git a/reactos/drivers/network/afd/afd/main.c b/reactos/drivers/network/afd/afd/main.c
index 264080ab804..e8fe79dc076 100644
--- a/reactos/drivers/network/afd/afd/main.c
+++ b/reactos/drivers/network/afd/afd/main.c
@@ -142,11 +142,14 @@ AfdCreateSocket(PDEVICE_OBJECT DeviceObject, PIRP Irp,
 	/* Allocate our backup buffer */
 	FCB->Recv.Window = ExAllocatePool( NonPagedPool, FCB->Recv.Size );
 	if( !FCB->Recv.Window ) Status = STATUS_NO_MEMORY;
-        FCB->Send.Window = ExAllocatePool( NonPagedPool, FCB->Send.Size );
-	if( !FCB->Send.Window ) {
-	    if( FCB->Recv.Window ) ExFreePool( FCB->Recv.Window );
-	    Status = STATUS_NO_MEMORY;
-	}
+        if( NT_SUCCESS(Status) )
+        {
+            FCB->Send.Window = ExAllocatePool( NonPagedPool, FCB->Send.Size );
+	    if( !FCB->Send.Window ) {
+	         if( FCB->Recv.Window ) ExFreePool( FCB->Recv.Window );
+	         Status = STATUS_NO_MEMORY;
+            }
+	     }
 	/* A datagram socket is always sendable */
 	FCB->PollState |= AFD_EVENT_SEND;
         PollReeval( FCB->DeviceExt, FCB->FileObject );
@@ -235,6 +238,8 @@ AfdCloseSocket(PDEVICE_OBJECT DeviceObject, PIRP Irp,
     AFD_DbgPrint(MID_TRACE,
 		 ("AfdClose(DeviceObject %p Irp %p)\n", DeviceObject, Irp));
 
+    if( !SocketAcquireStateLock( FCB ) ) return LostSocket( Irp );
+
     AFD_DbgPrint(MID_TRACE,("FCB %x\n", FCB));
 
     FCB->PollState |= AFD_EVENT_CLOSE;
@@ -244,11 +249,13 @@ AfdCloseSocket(PDEVICE_OBJECT DeviceObject, PIRP Irp,
     if( FCB->EventSelect ) ObDereferenceObject( FCB->EventSelect );
 
     FileObject->FsContext = NULL;
+    SocketStateUnlock( FCB );
+
     DestroySocket( FCB );
 
     Irp->IoStatus.Status = STATUS_SUCCESS;
     Irp->IoStatus.Information = 0;
-    IoCompleteRequest(Irp, IO_NO_INCREMENT);
+    IoCompleteRequest(Irp, IO_NETWORK_INCREMENT);
 
     AFD_DbgPrint(MID_TRACE, ("Returning success.\n"));