[Patch 2/2] ShellPkg/tftp: Add one option for tftp command to specify windowsize

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

Define -w option for tftp shell command to specify the TFTP windowsize
option. That will benefit the big file download for tftp server.

Cc: Ye Ting <ting.ye@intel.com>
Cc: Fu Siyuan <siyuan.fu@intel.com>
Cc: Carsey Jaben <jaben.carsey@intel.com>
Contributed-under: TianoCore Contribution Agreement 1.0
Signed-off-by: Wu Jiaxin <jiaxin.wu@intel.com>
---
 ShellPkg/Library/UefiShellTftpCommandLib/Tftp.c    | 60 ++++++++++++++++++----
 .../UefiShellTftpCommandLib.uni                    |  6 ++-
 2 files changed, 54 insertions(+), 12 deletions(-)

diff --git a/ShellPkg/Library/UefiShellTftpCommandLib/Tftp.c b/ShellPkg/Library/UefiShellTftpCommandLib/Tftp.c
index 5c50797..26f73ea 100755
--- a/ShellPkg/Library/UefiShellTftpCommandLib/Tftp.c
+++ b/ShellPkg/Library/UefiShellTftpCommandLib/Tftp.c
@@ -1,10 +1,10 @@
 /** @file
   The implementation for the 'tftp' Shell command.
 
   Copyright (c) 2015, ARM Ltd. All rights reserved.<BR>
-  Copyright (c) 2015 - 2016, Intel Corporation. All rights reserved. <BR>
+  Copyright (c) 2015 - 2017, Intel Corporation. All rights reserved. <BR>
   (C) Copyright 2015 Hewlett Packard Enterprise Development LP<BR>
 
   This program and the accompanying materials
   are licensed and made available under the terms and conditions of the BSD License
   which accompanies this distribution.  The full text of the license may be found at
@@ -180,10 +180,11 @@ DownloadFile (
   IN   EFI_MTFTP4_PROTOCOL  *Mtftp4,
   IN   CONST CHAR16         *FilePath,
   IN   CONST CHAR8          *AsciiFilePath,
   IN   UINTN                FileSize,
   IN   UINT16               BlockSize,
+  IN   UINT16               WindowSize,
   OUT  VOID                 **Data
   );
 
 /**
   Update the progress of a file download
@@ -224,10 +225,11 @@ STATIC CONST SHELL_PARAM_ITEM ParamList[] = {
   {L"-l", TypeValue},
   {L"-r", TypeValue},
   {L"-c", TypeValue},
   {L"-t", TypeValue},
   {L"-s", TypeValue},
+  {L"-w", TypeValue},
   {NULL , TypeMax}
   };
 
 ///
 /// The default block size (512) of tftp is defined in the RFC1350.
@@ -236,11 +238,16 @@ STATIC CONST SHELL_PARAM_ITEM ParamList[] = {
 ///
 /// The valid range of block size option is defined in the RFC2348.
 ///
 #define MTFTP_MIN_BLKSIZE          8
 #define MTFTP_MAX_BLKSIZE          65464
+///
+/// The The default windowsize.
+///
+#define MTFTP_DEFAULT_WINDOWSIZE   1
 
+#define MTFTP_MAX_WINDOWSIZE       64
 
 /**
   Function for 'tftp' command.
 
   @param[in] ImageHandle  Handle to the Image (NULL if Internal).
@@ -285,18 +292,20 @@ ShellCommandRunTftp (
   EFI_MTFTP4_PROTOCOL     *Mtftp4;
   UINTN                   FileSize;
   VOID                    *Data;
   SHELL_FILE_HANDLE       FileHandle;
   UINT16                  BlockSize;
+  UINT16                  WindowSize;
 
   ShellStatus         = SHELL_INVALID_PARAMETER;
   ProblemParam        = NULL;
   NicFound            = FALSE;
   AsciiRemoteFilePath = NULL;
   Handles             = NULL;
   FileSize            = 0;
   BlockSize           = MTFTP_DEFAULT_BLKSIZE;
+  WindowSize          = 1;
 
   //
   // Initialize the Shell library (we must be in non-auto-init...)
   //
   Status = ShellInitialize ();
@@ -432,10 +441,24 @@ ShellCommandRunTftp (
       );
       goto Error;
     }
   }
 
+  ValueStr = ShellCommandLineGetValue (CheckPackage, L"-w");
+  if (ValueStr != NULL) {
+    if (!StringToUint16 (ValueStr, &WindowSize)) {
+      goto Error;
+    }
+    if (WindowSize < MTFTP_DEFAULT_WINDOWSIZE || WindowSize > MTFTP_MAX_WINDOWSIZE) {
+      ShellPrintHiiEx (
+        -1, -1, NULL, STRING_TOKEN (STR_GEN_PARAM_INV),
+        gShellTftpHiiHandle, L"tftp", ValueStr
+      );
+      goto Error;
+    }
+  }
+
   //
   // Locate all MTFTP4 Service Binding protocols
   //
   ShellStatus = SHELL_NOT_FOUND;
   Status = gBS->LocateHandleBuffer (
@@ -506,11 +529,11 @@ ShellCommandRunTftp (
         gShellTftpHiiHandle, RemoteFilePath, NicName, Status
       );
       goto NextHandle;
     }
 
-    Status = DownloadFile (Mtftp4, RemoteFilePath, AsciiRemoteFilePath, FileSize, BlockSize, &Data);
+    Status = DownloadFile (Mtftp4, RemoteFilePath, AsciiRemoteFilePath, FileSize, BlockSize, WindowSize, &Data);
     if (EFI_ERROR (Status)) {
       ShellPrintHiiEx (
         -1, -1, NULL, STRING_TOKEN (STR_TFTP_ERR_DOWNLOAD),
         gShellTftpHiiHandle, RemoteFilePath, NicName, Status
       );
@@ -890,20 +913,21 @@ DownloadFile (
   IN   EFI_MTFTP4_PROTOCOL  *Mtftp4,
   IN   CONST CHAR16         *FilePath,
   IN   CONST CHAR8          *AsciiFilePath,
   IN   UINTN                FileSize,
   IN   UINT16               BlockSize,
+  IN   UINT16               WindowSize,
   OUT  VOID                 **Data
   )
 {
   EFI_STATUS            Status;
   EFI_PHYSICAL_ADDRESS  PagesAddress;
   VOID                  *Buffer;
   DOWNLOAD_CONTEXT      *TftpContext;
   EFI_MTFTP4_TOKEN      Mtftp4Token;
-  EFI_MTFTP4_OPTION     ReqOpt;
-  UINT8                 OptBuf[10];
+  UINT8                 BlksizeBuf[10];
+  UINT8                 WindowsizeBuf[10];
 
   // Downloaded file can be large. BS.AllocatePages() is more faster
   // than AllocatePool() and avoid fragmentation.
   // The downloaded file could be an EFI application. Marking the
   // allocated page as EfiBootServicesCode would allow to execute a
@@ -932,17 +956,29 @@ DownloadFile (
   Mtftp4Token.Filename    = (UINT8*)AsciiFilePath;
   Mtftp4Token.BufferSize  = FileSize;
   Mtftp4Token.Buffer      = Buffer;
   Mtftp4Token.CheckPacket = CheckPacket;
   Mtftp4Token.Context     = (VOID*)TftpContext;
+  Mtftp4Token.OptionCount = 0;
+  Mtftp4Token.OptionList  = AllocatePool (sizeof (EFI_MTFTP4_OPTION) * 2);
+  if (Mtftp4Token.OptionList == NULL) {
+    Status = EFI_OUT_OF_RESOURCES;
+    goto Error;
+  }
+
   if (BlockSize != MTFTP_DEFAULT_BLKSIZE) {
-    ReqOpt.OptionStr = (UINT8 *) "blksize";
-    AsciiSPrint ((CHAR8 *)OptBuf, sizeof (OptBuf), "%d", BlockSize);
-    ReqOpt.ValueStr  = OptBuf;
+    Mtftp4Token.OptionList[Mtftp4Token.OptionCount].OptionStr = (UINT8 *) "blksize";
+    AsciiSPrint ((CHAR8 *) BlksizeBuf, sizeof (BlksizeBuf), "%d", BlockSize);
+    Mtftp4Token.OptionList[Mtftp4Token.OptionCount].ValueStr  = BlksizeBuf;
+    Mtftp4Token.OptionCount ++;
+  }
 
-    Mtftp4Token.OptionCount = 1;
-    Mtftp4Token.OptionList  = &ReqOpt;
+  if (WindowSize != MTFTP_DEFAULT_WINDOWSIZE) {
+    Mtftp4Token.OptionList[Mtftp4Token.OptionCount].OptionStr = (UINT8 *) "windowsize";
+    AsciiSPrint ((CHAR8 *) WindowsizeBuf, sizeof (WindowsizeBuf), "%d", WindowSize);
+    Mtftp4Token.OptionList[Mtftp4Token.OptionCount].ValueStr  = WindowsizeBuf;
+    Mtftp4Token.OptionCount ++;
   }
 
   ShellPrintHiiEx (
     -1, -1, NULL, STRING_TOKEN (STR_TFTP_DOWNLOADING),
     gShellTftpHiiHandle, FilePath
@@ -954,14 +990,18 @@ DownloadFile (
     gShellTftpHiiHandle
     );
 
 Error :
 
-  if (TftpContext == NULL) {
+  if (TftpContext != NULL) {
     FreePool (TftpContext);
   }
 
+  if (Mtftp4Token.OptionList != NULL) {
+    FreePool (Mtftp4Token.OptionList);
+  }
+
   if (EFI_ERROR (Status)) {
     gBS->FreePages (PagesAddress, EFI_SIZE_TO_PAGES (FileSize));
     return Status;
   }
 
diff --git a/ShellPkg/Library/UefiShellTftpCommandLib/UefiShellTftpCommandLib.uni b/ShellPkg/Library/UefiShellTftpCommandLib/UefiShellTftpCommandLib.uni
index 4f4447d..ca6fa44 100644
--- a/ShellPkg/Library/UefiShellTftpCommandLib/UefiShellTftpCommandLib.uni
+++ b/ShellPkg/Library/UefiShellTftpCommandLib/UefiShellTftpCommandLib.uni
@@ -1,9 +1,9 @@
 // /**
 //
 // (C) Copyright 2015-2016 Hewlett Packard Enterprise Development LP<BR>
-// Copyright (c) 2010 - 2016, Intel Corporation. All rights reserved. <BR>
+// Copyright (c) 2010 - 2017, Intel Corporation. All rights reserved. <BR>
 // This program and the accompanying materials
 // are licensed and made available under the terms and conditions of the BSD License
 // which accompanies this distribution. The full text of the license may be found at
 // http://opensource.org/licenses/bsd-license.php
 //
@@ -48,11 +48,11 @@
 ".SH NAME\r\n"
 "Download a file from TFTP server.\r\n"
 ".SH SYNOPSIS\r\n"
 " \r\n"
 "TFTP [-i interface] [-l <port>] [-r <port>] [-c <retry count>] [-t <timeout>]\r\n"
-"     [-s <block size>] host remotefilepath [localfilepath]\r\n"
+"     [-s <block size>] [-w <window size>] host remotefilepath [localfilepath]\r\n"
 ".SH OPTIONS\r\n"
 " \r\n"
 "  -i interface     - Specifies an adapter name, i.e., eth0.\r\n"
 "  -l port          - Specifies the local port number. Default value is 0\r\n"
 "                     and the port number is automatically assigned.\r\n"
@@ -61,10 +61,12 @@
 "                     wait for a response. The default value is 6.\r\n"
 "  -t <timeout>     - The number of seconds to wait for a response after\r\n"
 "                     sending a request packet. Default value is 4s.\r\n"
 "  -s <block size>  - Specifies the TFTP blksize option as defined in RFC 2348.\r\n"
 "                     Valid range is between 8 and 65464, default value is 512.\r\n"
+"  -w <window size> - Specifies the TFTP windowsize option as defined in RFC 7440.\r\n"
+"                     Valid range is between 1 and 64, default value is 1.\r\n"
 "  host             - Specify TFTP Server IPv4 address.\r\n"
 "  remotefilepath   - TFTP server file path to download the file.\r\n"
 "  localfilepath    - Local destination file path.\r\n"
 ".SH DESCRIPTION\r\n"
 " \r\n"
-- 
1.9.5.msysgit.1