ShellPkg/TftpDynamicCommand: Add one option for tftp command to specify windowsize.

[mirror_edk2.git] / ShellPkg / DynamicCommand / TftpDynamicCommand / Tftp.c

diff --git a/ShellPkg/DynamicCommand/TftpDynamicCommand/Tftp.c b/ShellPkg/DynamicCommand/TftpDynamicCommand/Tftp.c
index 44be6d4e76f4bd754b007ba00d024797664a684f..c66be6b9d97e5eb001373be22f08d813f5e6b7d5 100644 (file)
--- a/ShellPkg/DynamicCommand/TftpDynamicCommand/Tftp.c
+++ b/ShellPkg/DynamicCommand/TftpDynamicCommand/Tftp.c
@@ -183,6 +183,7 @@ DownloadFile (
   IN   CONST CHAR8          *AsciiFilePath,\r
   IN   UINTN                FileSize,\r
   IN   UINT16               BlockSize,\r
+  IN   UINT16               WindowSize,\r
   OUT  VOID                 **Data\r
   );\r
 \r
@@ -227,6 +228,7 @@ STATIC CONST SHELL_PARAM_ITEM ParamList[] = {
   {L"-c", TypeValue},\r
   {L"-t", TypeValue},\r
   {L"-s", TypeValue},\r
+  {L"-w", TypeValue},\r
   {NULL , TypeMax}\r
   };\r
 \r
@@ -239,7 +241,17 @@ STATIC CONST SHELL_PARAM_ITEM ParamList[] = {
 ///\r
 #define MTFTP_MIN_BLKSIZE          8\r
 #define MTFTP_MAX_BLKSIZE          65464\r
-\r
+///\r
+/// The default windowsize (1) of tftp.\r
+///\r
+#define MTFTP_DEFAULT_WINDOWSIZE   1\r
+///\r
+/// The valid range of window size option.\r
+/// Note that: RFC 7440 does not mention max window size value, but for the\r
+/// stability reason, the value is limited to 64.\r
+///\r
+#define MTFTP_MIN_WINDOWSIZE       1\r
+#define MTFTP_MAX_WINDOWSIZE       64\r
 \r
 /**\r
   Function for 'tftp' command.\r
@@ -288,6 +300,7 @@ RunTftp (
   VOID                    *Data;\r
   SHELL_FILE_HANDLE       FileHandle;\r
   UINT16                  BlockSize;\r
+  UINT16                  WindowSize;\r
 \r
   ShellStatus         = SHELL_INVALID_PARAMETER;\r
   ProblemParam        = NULL;\r
@@ -297,6 +310,7 @@ RunTftp (
   FileSize            = 0;\r
   DataSize            = 0;\r
   BlockSize           = MTFTP_DEFAULT_BLKSIZE;\r
+  WindowSize          = MTFTP_DEFAULT_WINDOWSIZE;\r
 \r
   //\r
   // Initialize the Shell library (we must be in non-auto-init...)\r
@@ -436,6 +450,20 @@ RunTftp (
     }\r
   }\r
 \r
+  ValueStr = ShellCommandLineGetValue (CheckPackage, L"-w");\r
+  if (ValueStr != NULL) {\r
+    if (!StringToUint16 (ValueStr, &WindowSize)) {\r
+      goto Error;\r
+    }\r
+    if (WindowSize < MTFTP_MIN_WINDOWSIZE || WindowSize > MTFTP_MAX_WINDOWSIZE) {\r
+      ShellPrintHiiEx (\r
+        -1, -1, NULL, STRING_TOKEN (STR_GEN_PARAM_INV),\r
+        mTftpHiiHandle, L"tftp", ValueStr\r
+      );\r
+      goto Error;\r
+    }\r
+  }\r
+\r
   //\r
   // Locate all MTFTP4 Service Binding protocols\r
   //\r
@@ -510,7 +538,7 @@ RunTftp (
       goto NextHandle;\r
     }\r
 \r
-    Status = DownloadFile (Mtftp4, RemoteFilePath, AsciiRemoteFilePath, FileSize, BlockSize, &Data);\r
+    Status = DownloadFile (Mtftp4, RemoteFilePath, AsciiRemoteFilePath, FileSize, BlockSize, WindowSize, &Data);\r
     if (EFI_ERROR (Status)) {\r
       ShellPrintHiiEx (\r
         -1, -1, NULL, STRING_TOKEN (STR_TFTP_ERR_DOWNLOAD),\r
@@ -896,6 +924,7 @@ DownloadFile (
   IN   CONST CHAR8          *AsciiFilePath,\r
   IN   UINTN                FileSize,\r
   IN   UINT16               BlockSize,\r
+  IN   UINT16               WindowSize,\r
   OUT  VOID                 **Data\r
   )\r
 {\r
@@ -904,8 +933,8 @@ DownloadFile (
   VOID                  *Buffer;\r
   DOWNLOAD_CONTEXT      *TftpContext;\r
   EFI_MTFTP4_TOKEN      Mtftp4Token;\r
-  EFI_MTFTP4_OPTION     ReqOpt;\r
-  UINT8                 OptBuf[10];\r
+  UINT8                 BlksizeBuf[10];\r
+  UINT8                 WindowsizeBuf[10];\r
 \r
   // Downloaded file can be large. BS.AllocatePages() is more faster\r
   // than AllocatePool() and avoid fragmentation.\r
@@ -938,13 +967,25 @@ DownloadFile (
   Mtftp4Token.Buffer      = Buffer;\r
   Mtftp4Token.CheckPacket = CheckPacket;\r
   Mtftp4Token.Context     = (VOID*)TftpContext;\r
+  Mtftp4Token.OptionCount = 0;\r
+  Mtftp4Token.OptionList  = AllocatePool (sizeof (EFI_MTFTP4_OPTION) * 2);\r
+  if (Mtftp4Token.OptionList == NULL) {\r
+    Status = EFI_OUT_OF_RESOURCES;\r
+    goto Error;\r
+  }\r
+\r
   if (BlockSize != MTFTP_DEFAULT_BLKSIZE) {\r
-    ReqOpt.OptionStr = (UINT8 *) "blksize";\r
-    AsciiSPrint ((CHAR8 *)OptBuf, sizeof (OptBuf), "%d", BlockSize);\r
-    ReqOpt.ValueStr  = OptBuf;\r
+    Mtftp4Token.OptionList[Mtftp4Token.OptionCount].OptionStr = (UINT8 *) "blksize";\r
+    AsciiSPrint ((CHAR8 *) BlksizeBuf, sizeof (BlksizeBuf), "%d", BlockSize);\r
+    Mtftp4Token.OptionList[Mtftp4Token.OptionCount].ValueStr  = BlksizeBuf;\r
+    Mtftp4Token.OptionCount ++;\r
+  }\r
 \r
-    Mtftp4Token.OptionCount = 1;\r
-    Mtftp4Token.OptionList  = &ReqOpt;\r
+  if (WindowSize != MTFTP_DEFAULT_WINDOWSIZE) {\r
+    Mtftp4Token.OptionList[Mtftp4Token.OptionCount].OptionStr = (UINT8 *) "windowsize";\r
+    AsciiSPrint ((CHAR8 *) WindowsizeBuf, sizeof (WindowsizeBuf), "%d", WindowSize);\r
+    Mtftp4Token.OptionList[Mtftp4Token.OptionCount].ValueStr  = WindowsizeBuf;\r
+    Mtftp4Token.OptionCount ++;\r
   }\r
 \r
   ShellPrintHiiEx (\r
@@ -960,10 +1001,14 @@ DownloadFile (
 \r
 Error :\r
 \r
-  if (TftpContext == NULL) {\r
+  if (TftpContext != NULL) {\r
     FreePool (TftpContext);\r
   }\r
 \r
+  if (Mtftp4Token.OptionList != NULL) {\r
+    FreePool (Mtftp4Token.OptionList);\r
+  }\r
+\r
   if (EFI_ERROR (Status)) {\r
     gBS->FreePages (PagesAddress, EFI_SIZE_TO_PAGES (FileSize));\r
     return Status;\r