diff --git a/MdeModulePkg/Universal/Network/Mtftp4Dxe/Mtftp4Rrq.c b/MdeModulePkg/Universal/Network/Mtftp4Dxe/Mtftp4Rrq.c index c574f315df..264598f49a 100644 --- a/MdeModulePkg/Universal/Network/Mtftp4Dxe/Mtftp4Rrq.c +++ b/MdeModulePkg/Universal/Network/Mtftp4Dxe/Mtftp4Rrq.c @@ -169,7 +169,7 @@ Mtftp4RrqSaveBlock ( // to accept transfers of unlimited size. So TotalBlock is memorised as // continuous block counter. // - Status = Mtftp4RemoveBlockNum (&Instance->Blocks, Block, &TotalBlock); + Status = Mtftp4RemoveBlockNum (&Instance->Blocks, Block, Completed, &TotalBlock); if (Status == EFI_NOT_FOUND) { return EFI_SUCCESS; diff --git a/MdeModulePkg/Universal/Network/Mtftp4Dxe/Mtftp4Support.c b/MdeModulePkg/Universal/Network/Mtftp4Dxe/Mtftp4Support.c index 4dcdf827ea..bb4ce0ae65 100644 --- a/MdeModulePkg/Universal/Network/Mtftp4Dxe/Mtftp4Support.c +++ b/MdeModulePkg/Universal/Network/Mtftp4Dxe/Mtftp4Support.c @@ -158,6 +158,7 @@ Mtftp4SetLastBlockNum ( @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 @retval EFI_NOT_FOUND The block number isn't in the block range list @@ -169,6 +170,7 @@ EFI_STATUS Mtftp4RemoveBlockNum ( IN LIST_ENTRY *Head, IN UINT16 Num, + IN BOOLEAN Completed, OUT UINT64 *TotalBlock ) { @@ -229,7 +231,7 @@ Mtftp4RemoveBlockNum ( Range->Round ++; } - if (Range->Start > Range->End) { + if ((Range->Start > Range->End) || Completed) { RemoveEntryList (&Range->Link); FreePool (Range); } diff --git a/MdeModulePkg/Universal/Network/Mtftp4Dxe/Mtftp4Support.h b/MdeModulePkg/Universal/Network/Mtftp4Dxe/Mtftp4Support.h index 7a712a1a6d..802e3867db 100644 --- a/MdeModulePkg/Universal/Network/Mtftp4Dxe/Mtftp4Support.h +++ b/MdeModulePkg/Universal/Network/Mtftp4Dxe/Mtftp4Support.h @@ -92,6 +92,7 @@ Mtftp4SetLastBlockNum ( @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 @retval EFI_NOT_FOUND The block number isn't in the block range list @@ -103,6 +104,7 @@ EFI_STATUS Mtftp4RemoveBlockNum ( IN LIST_ENTRY *Head, IN UINT16 Num, + IN BOOLEAN Completed, OUT UINT64 *TotalBlock ); diff --git a/MdeModulePkg/Universal/Network/Mtftp4Dxe/Mtftp4Wrq.c b/MdeModulePkg/Universal/Network/Mtftp4Dxe/Mtftp4Wrq.c index 4312e18c6f..412c472e42 100644 --- a/MdeModulePkg/Universal/Network/Mtftp4Dxe/Mtftp4Wrq.c +++ b/MdeModulePkg/Universal/Network/Mtftp4Dxe/Mtftp4Wrq.c @@ -166,7 +166,7 @@ Mtftp4WrqHandleAck ( // tell the Mtftp4WrqInput to finish the transfer. This is the last // block number if the block range are empty.. // - Mtftp4RemoveBlockNum (&Instance->Blocks, AckNum, &TotalBlock); + Mtftp4RemoveBlockNum (&Instance->Blocks, AckNum, *Completed,&TotalBlock); Expected = Mtftp4GetNextBlockNum (&Instance->Blocks);