From c3cd46d421cef32077ae6caf860b50d00ba2aa7f Mon Sep 17 00:00:00 2001
From: sfu5 <sfu5@6f19259b-4bc3-4df7-8a09-765794883524>
Date: Fri, 9 Sep 2011 08:31:37 +0000
Subject: [PATCH] 1. Fix bug in PXE driver UdpRead function to handle the IP
 fragmentation.

Signed-off-by: sfu5
Reviewed-by: xdu2
Reviewed-by: hhuan13

git-svn-id: https://edk2.svn.sourceforge.net/svnroot/edk2/trunk/edk2@12309 6f19259b-4bc3-4df7-8a09-765794883524
---
 .../Network/UefiPxeBcDxe/PxeBcImpl.c          | 59 ++++++++++++++-----
 1 file changed, 45 insertions(+), 14 deletions(-)

diff --git a/MdeModulePkg/Universal/Network/UefiPxeBcDxe/PxeBcImpl.c b/MdeModulePkg/Universal/Network/UefiPxeBcDxe/PxeBcImpl.c
index dfdfd35cb7..6f6d7f92ae 100644
--- a/MdeModulePkg/Universal/Network/UefiPxeBcDxe/PxeBcImpl.c
+++ b/MdeModulePkg/Universal/Network/UefiPxeBcDxe/PxeBcImpl.c
@@ -1631,7 +1631,13 @@ EfiPxeBcUdpRead (
   EFI_STATUS                Status;
   BOOLEAN                   IsDone;
   BOOLEAN                   Matched;
-  UINTN                     CopyLen;
+  UINTN                     CopiedLen;
+  UINTN                     HeaderLen;
+  UINTN                     HeaderCopiedLen;
+  UINTN                     BufferCopiedLen;
+  UINT32                    FragmentLength;
+  UINTN                     FragmentIndex;
+  UINT8                     *FragmentBuffer;
 
   if (This == NULL || DestIp == NULL || DestPort == NULL) {
     return EFI_INVALID_PARAMETER;
@@ -1788,26 +1794,51 @@ TRY_AGAIN:
     }
 
     if (Matched) {
+      ASSERT (RxData != NULL);
 
-      CopyLen = 0;
-
+      HeaderLen = 0;
       if (HeaderSize != NULL) {
-        CopyLen = MIN (*HeaderSize, RxData->DataLength);
-        CopyMem (HeaderPtr, RxData->FragmentTable[0].FragmentBuffer, CopyLen);
-        *HeaderSize = CopyLen;
+        HeaderLen = MIN (*HeaderSize, RxData->DataLength);
       }
 
-      if (RxData->DataLength - CopyLen > *BufferSize) {
-
+      if (RxData->DataLength - HeaderLen > *BufferSize) {
         Status = EFI_BUFFER_TOO_SMALL;
       } else {
+        *HeaderSize = HeaderLen;
+        *BufferSize = RxData->DataLength - HeaderLen;
 
-        *BufferSize = RxData->DataLength - CopyLen;
-        CopyMem (
-          BufferPtr,
-          (UINT8 *) RxData->FragmentTable[0].FragmentBuffer + CopyLen,
-          *BufferSize
-          );
+        HeaderCopiedLen = 0;
+        BufferCopiedLen = 0;
+        for (FragmentIndex = 0; FragmentIndex < RxData->FragmentCount; FragmentIndex++) {
+          FragmentLength = RxData->FragmentTable[FragmentIndex].FragmentLength;
+          FragmentBuffer = RxData->FragmentTable[FragmentIndex].FragmentBuffer;
+          if (HeaderCopiedLen + FragmentLength < HeaderLen) {
+            //
+            // Copy the header part of received data.
+            //
+            CopyMem ((UINT8 *) HeaderPtr + HeaderCopiedLen, FragmentBuffer, FragmentLength);
+            HeaderCopiedLen += FragmentLength;
+          } else if (HeaderCopiedLen < HeaderLen) {
+            //
+            // Copy the header part of received data.
+            //
+            CopiedLen = HeaderLen - HeaderCopiedLen;
+            CopyMem ((UINT8 *) HeaderPtr + HeaderCopiedLen, FragmentBuffer, CopiedLen);
+            HeaderCopiedLen += CopiedLen;
+
+            //
+            // Copy the other part of received data.
+            //
+            CopyMem ((UINT8 *) BufferPtr + BufferCopiedLen, FragmentBuffer + CopiedLen, FragmentLength - CopiedLen);
+            BufferCopiedLen += (FragmentLength - CopiedLen);
+          } else {
+            //
+            // Copy the other part of received data.
+            //
+            CopyMem ((UINT8 *) BufferPtr + BufferCopiedLen, FragmentBuffer, FragmentLength);
+            BufferCopiedLen += FragmentLength;
+          }
+        }
       }
     } else {