[Patch 1/2] MdeModulePke/Mtftp4Dxe: Correct the total received and saved block number.

[Jiaxin Wu <Jiaxin.wu@intel.com>]

The block returned from Mtftp4RemoveBlockNum is not the total received and
saved block number if it works in passive (Slave) mode.

The issue was exposed by the EMS test.

Cc: Ye Ting <ting.ye@intel.com>
Cc: Fu Siyuan <siyuan.fu@intel.com>
Contributed-under: TianoCore Contribution Agreement 1.1
Signed-off-by: Wu Jiaxin <jiaxin.wu@intel.com>
---
 .../Universal/Network/Mtftp4Dxe/Mtftp4Impl.h     |  6 +++++-
 .../Universal/Network/Mtftp4Dxe/Mtftp4Rrq.c      | 16 +++++++++++-----
 .../Universal/Network/Mtftp4Dxe/Mtftp4Support.c  | 10 +++++-----
 .../Universal/Network/Mtftp4Dxe/Mtftp4Support.h  |  6 +++---
 .../Universal/Network/Mtftp4Dxe/Mtftp4Wrq.c      |  6 +++---
 5 files changed, 27 insertions(+), 17 deletions(-)

diff --git a/MdeModulePkg/Universal/Network/Mtftp4Dxe/Mtftp4Impl.h b/MdeModulePkg/Universal/Network/Mtftp4Dxe/Mtftp4Impl.h
index de304f4e70..be2f8af6e4 100644
--- a/MdeModulePkg/Universal/Network/Mtftp4Dxe/Mtftp4Impl.h
+++ b/MdeModulePkg/Universal/Network/Mtftp4Dxe/Mtftp4Impl.h
@@ -124,13 +124,17 @@ struct _MTFTP4_PROTOCOL {
   LIST_ENTRY                    Blocks;
 
   UINT16                        WindowSize;
 
   //
-  // Record the total received block number and the already acked block number.
+  // Record the total received and saved block number.
   //
   UINT64                        TotalBlock;
+  
+  //
+  // Record the acked block number.
+  //
   UINT64                        AckedBlock;
 
   //
   // The server's communication end point: IP and two ports. one for
   // initial request, one for its selected port.
diff --git a/MdeModulePkg/Universal/Network/Mtftp4Dxe/Mtftp4Rrq.c b/MdeModulePkg/Universal/Network/Mtftp4Dxe/Mtftp4Rrq.c
index fedf1cde46..6960e322a5 100644
--- a/MdeModulePkg/Universal/Network/Mtftp4Dxe/Mtftp4Rrq.c
+++ b/MdeModulePkg/Universal/Network/Mtftp4Dxe/Mtftp4Rrq.c
@@ -152,10 +152,11 @@ Mtftp4RrqSaveBlock (
   EFI_MTFTP4_TOKEN          *Token;
   EFI_STATUS                Status;
   UINT16                    Block;
   UINT64                    Start;
   UINT32                    DataLen;
+  UINT64                    BlockCounter;
   BOOLEAN                   Completed;
 
   Completed = FALSE;
   Token     = Instance->Token;
   Block     = NTOHS (Packet->Data.Block);
@@ -172,14 +173,14 @@ Mtftp4RrqSaveBlock (
 
   //
   // Remove this block number from the file hole. If Mtftp4RemoveBlockNum
   // returns EFI_NOT_FOUND, the block has been saved, don't save it again.
   // Note that : For bigger files, allowing the block counter to roll over
-  // to accept transfers of unlimited size. So TotalBlock is memorised as
+  // to accept transfers of unlimited size. So BlockCounter is memorised as
   // continuous block counter.
   //
-  Status = Mtftp4RemoveBlockNum (&Instance->Blocks, Block, Completed, &Instance->TotalBlock);
+  Status = Mtftp4RemoveBlockNum (&Instance->Blocks, Block, Completed, &BlockCounter);
 
   if (Status == EFI_NOT_FOUND) {
     return EFI_SUCCESS;
   } else if (EFI_ERROR (Status)) {
     return Status;
@@ -198,11 +199,11 @@ Mtftp4RrqSaveBlock (
       return EFI_ABORTED;
     }
   }
 
   if (Token->Buffer != NULL) {
-     Start = MultU64x32 (Instance->TotalBlock - 1, Instance->BlkSize);
+     Start = MultU64x32 (BlockCounter - 1, Instance->BlkSize);
 
     if (Start + DataLen <= Token->BufferSize) {
       CopyMem ((UINT8 *) Token->Buffer + Start, Packet->Data.Data, DataLen);
 
       //
@@ -269,13 +270,13 @@ Mtftp4RrqHandleData (
   Expected    = Mtftp4GetNextBlockNum (&Instance->Blocks);
 
   ASSERT (Expected >= 0);
 
   //
-  // If we are active and received an unexpected packet, transmit
+  // If we are active (Master) and received an unexpected packet, transmit
   // the ACK for the block we received, then restart receiving the
-  // expected one. If we are passive, save the block.
+  // expected one. If we are passive (Slave), save the block.
   //
   if (Instance->Master && (Expected != BlockNum)) {
     //
     // If Expected is 0, (UINT16) (Expected - 1) is also the expected Ack number (65535).
     //
@@ -286,10 +287,15 @@ Mtftp4RrqHandleData (
 
   if (EFI_ERROR (Status)) {
     return Status;
   }
 
+  //
+  // Record the total received and saved block number.
+  //
+  Instance->TotalBlock ++;
+
   //
   // Reset the passive client's timer whenever it received a
   // valid data packet.
   //
   if (!Instance->Master) {
diff --git a/MdeModulePkg/Universal/Network/Mtftp4Dxe/Mtftp4Support.c b/MdeModulePkg/Universal/Network/Mtftp4Dxe/Mtftp4Support.c
index 71fd979b3a..5e282e9c4b 100644
--- a/MdeModulePkg/Universal/Network/Mtftp4Dxe/Mtftp4Support.c
+++ b/MdeModulePkg/Universal/Network/Mtftp4Dxe/Mtftp4Support.c
@@ -156,12 +156,12 @@ Mtftp4SetLastBlockNum (
 /**
   Remove the block number from the block range list.
 
   @param  Head                  The block range list to remove from
   @param  Num                   The block number to remove
-  @param  Completed             Whether Num is the last block number
-  @param  TotalBlock            The continuous block number in all
+  @param  Completed             Whether Num is the last block number.
+  @param  BlockCounter          The continuous block counter instead of the value after roll-over.
 
   @retval EFI_NOT_FOUND         The block number isn't in the block range list
   @retval EFI_SUCCESS           The block number has been removed from the list
   @retval EFI_OUT_OF_RESOURCES  Failed to allocate resource
 
@@ -169,11 +169,11 @@ Mtftp4SetLastBlockNum (
 EFI_STATUS
 Mtftp4RemoveBlockNum (
   IN LIST_ENTRY             *Head,
   IN UINT16                 Num,
   IN BOOLEAN                Completed,
-  OUT UINT64                *TotalBlock
+  OUT UINT64                *BlockCounter
   )
 {
   MTFTP4_BLOCK_RANGE        *Range;
   MTFTP4_BLOCK_RANGE        *NewRange;
   LIST_ENTRY                *Entry;
@@ -218,14 +218,14 @@ Mtftp4RemoveBlockNum (
       // transfers of unlimited size. There is no consensus, however, whether
       // the counter should wrap around to zero or to one. Many implementations
       // wrap to zero, because this is the simplest to implement. Here we choose
       // this solution.
       //
-      *TotalBlock  = Num;
+      *BlockCounter  = Num;
 
       if (Range->Round > 0) {
-        *TotalBlock += Range->Bound +  MultU64x32 ((UINTN) (Range->Round -1), (UINT32) (Range->Bound + 1)) + 1;
+        *BlockCounter += Range->Bound +  MultU64x32 ((UINTN) (Range->Round -1), (UINT32) (Range->Bound + 1)) + 1;
       }
 
       if (Range->Start > Range->Bound) {
         Range->Start = 0;
         Range->Round ++;
diff --git a/MdeModulePkg/Universal/Network/Mtftp4Dxe/Mtftp4Support.h b/MdeModulePkg/Universal/Network/Mtftp4Dxe/Mtftp4Support.h
index 6cc2756bc8..f7a6755fe8 100644
--- a/MdeModulePkg/Universal/Network/Mtftp4Dxe/Mtftp4Support.h
+++ b/MdeModulePkg/Universal/Network/Mtftp4Dxe/Mtftp4Support.h
@@ -90,12 +90,12 @@ Mtftp4SetLastBlockNum (
 /**
   Remove the block number from the block range list.
 
   @param  Head                  The block range list to remove from
   @param  Num                   The block number to remove
-  @param  Completed             Wether Num is the last block number
-  @param  TotalBlock            The continuous block number in all
+  @param  Completed             Whether Num is the last block number.
+  @param  BlockCounter          The continuous block counter instead of the value after roll-over.
 
   @retval EFI_NOT_FOUND         The block number isn't in the block range list
   @retval EFI_SUCCESS           The block number has been removed from the list
   @retval EFI_OUT_OF_RESOURCES  Failed to allocate resource
 
@@ -103,11 +103,11 @@ Mtftp4SetLastBlockNum (
 EFI_STATUS
 Mtftp4RemoveBlockNum (
   IN LIST_ENTRY             *Head,
   IN UINT16                 Num,
   IN BOOLEAN                Completed,
-  OUT UINT64                *TotalBlock
+  OUT UINT64                *BlockCounter
   );
 
 /**
   Set the timeout for the instance. User a longer time for passive instances.
 
diff --git a/MdeModulePkg/Universal/Network/Mtftp4Dxe/Mtftp4Wrq.c b/MdeModulePkg/Universal/Network/Mtftp4Dxe/Mtftp4Wrq.c
index ea309e2d6b..ee70accbcd 100644
--- a/MdeModulePkg/Universal/Network/Mtftp4Dxe/Mtftp4Wrq.c
+++ b/MdeModulePkg/Universal/Network/Mtftp4Dxe/Mtftp4Wrq.c
@@ -147,11 +147,11 @@ Mtftp4WrqHandleAck (
      OUT BOOLEAN               *Completed
   )
 {
   UINT16                    AckNum;
   INTN                      Expected;
-  UINT64                    TotalBlock;
+  UINT64                    BlockCounter;
 
   *Completed  = FALSE;
   AckNum      = NTOHS (Packet->Ack.Block[0]);
   Expected    = Mtftp4GetNextBlockNum (&Instance->Blocks);
 
@@ -166,13 +166,13 @@ Mtftp4WrqHandleAck (
   }
 
   //
   // Remove the acked block number, if this is the last block number,
   // tell the Mtftp4WrqInput to finish the transfer. This is the last
-  // block number if the block range are empty..
+  // block number if the block range are empty.
   //
-  Mtftp4RemoveBlockNum (&Instance->Blocks, AckNum, *Completed,&TotalBlock);
+  Mtftp4RemoveBlockNum (&Instance->Blocks, AckNum, *Completed, &BlockCounter);
 
   Expected = Mtftp4GetNextBlockNum (&Instance->Blocks);
 
   if (Expected < 0) {
 
-- 
2.17.1.windows.2