pointer verification (not NULL) and buffer overrun fixes.

[mirror_edk2.git] / ShellPkg / Library / UefiShellDebug1CommandsLib / Bcfg.c

diff --git a/ShellPkg/Library/UefiShellDebug1CommandsLib/Bcfg.c b/ShellPkg/Library/UefiShellDebug1CommandsLib/Bcfg.c
index 73e8c576372af350d0e60b438c076724cd896ee5..d278ee60e6696d7d717670a1fb99cb53a089821f 100644 (file)
--- a/ShellPkg/Library/UefiShellDebug1CommandsLib/Bcfg.c
+++ b/ShellPkg/Library/UefiShellDebug1CommandsLib/Bcfg.c
@@ -265,27 +265,31 @@ BcfgAddDebug1(
     FilePathSize = GetDevicePathSize (FilePath);\r
 \r
     TempByteBuffer = AllocateZeroPool(sizeof(UINT32) + sizeof(UINT16) + DescSize + FilePathSize);\r
-    TempByteStart  = TempByteBuffer;\r
-    *((UINT32 *) TempByteBuffer) = LOAD_OPTION_ACTIVE;      // Attributes\r
-    TempByteBuffer += sizeof (UINT32);\r
+    if (TempByteBuffer != NULL) {\r
+      TempByteStart  = TempByteBuffer;\r
+      *((UINT32 *) TempByteBuffer) = LOAD_OPTION_ACTIVE;      // Attributes\r
+      TempByteBuffer += sizeof (UINT32);\r
 \r
-    *((UINT16 *) TempByteBuffer) = (UINT16)FilePathSize;    // FilePathListLength\r
-    TempByteBuffer += sizeof (UINT16);\r
+      *((UINT16 *) TempByteBuffer) = (UINT16)FilePathSize;    // FilePathListLength\r
+      TempByteBuffer += sizeof (UINT16);\r
 \r
-    CopyMem (TempByteBuffer, Desc, DescSize);\r
-    TempByteBuffer += DescSize;\r
-    CopyMem (TempByteBuffer, FilePath, FilePathSize);\r
+      CopyMem (TempByteBuffer, Desc, DescSize);\r
+      TempByteBuffer += DescSize;\r
+      CopyMem (TempByteBuffer, FilePath, FilePathSize);\r
 \r
-    UnicodeSPrint (OptionStr, sizeof(OptionStr), L"%s%04x", Target == BcfgTargetBootOrder?L"Boot":L"Driver", TargetLocation);\r
-    Status = gRT->SetVariable (\r
-          OptionStr,\r
-          &gEfiGlobalVariableGuid,\r
-          EFI_VARIABLE_NON_VOLATILE|EFI_VARIABLE_BOOTSERVICE_ACCESS|EFI_VARIABLE_RUNTIME_ACCESS,\r
-          sizeof(UINT32) + sizeof(UINT16) + DescSize + FilePathSize,\r
-          TempByteStart\r
-         );\r
-\r
-    FreePool(TempByteStart);\r
+      UnicodeSPrint (OptionStr, sizeof(OptionStr), L"%s%04x", Target == BcfgTargetBootOrder?L"Boot":L"Driver", TargetLocation);\r
+      Status = gRT->SetVariable (\r
+            OptionStr,\r
+            &gEfiGlobalVariableGuid,\r
+            EFI_VARIABLE_NON_VOLATILE|EFI_VARIABLE_BOOTSERVICE_ACCESS|EFI_VARIABLE_RUNTIME_ACCESS,\r
+            sizeof(UINT32) + sizeof(UINT16) + DescSize + FilePathSize,\r
+            TempByteStart\r
+           );\r
+\r
+      FreePool(TempByteStart);\r
+    } else {\r
+      Status = EFI_OUT_OF_RESOURCES;\r
+    }\r
 \r
     if (EFI_ERROR(Status)) {\r
       ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_BCFG_SET_VAR_FAIL), gShellDebug1HiiHandle, OptionStr, Status);\r
@@ -385,22 +389,25 @@ BcfgRemoveDebug1(
     return (SHELL_INVALID_PARAMETER);\r
   }\r
   NewOrder = AllocateZeroPool(OrderCount*sizeof(CurrentOrder[0]));\r
-  NewCount = OrderCount;\r
-  CopyMem(NewOrder, CurrentOrder, OrderCount*sizeof(CurrentOrder[0]));\r
-  for (LoopVar = 0 ; LoopVar < OrderCount ; LoopVar++){\r
-    if (NewOrder[LoopVar] == Location) {\r
-      CopyMem(NewOrder+LoopVar, NewOrder+LoopVar+1, (OrderCount - LoopVar - 1)*sizeof(CurrentOrder[0]));\r
-      NewCount--;\r
+  if (NewOrder != NULL) {\r
+    NewCount = OrderCount;\r
+    CopyMem(NewOrder, CurrentOrder, OrderCount*sizeof(CurrentOrder[0]));\r
+    for (LoopVar = 0 ; LoopVar < OrderCount ; LoopVar++){\r
+      if (NewOrder[LoopVar] == Location) {\r
+        CopyMem(NewOrder+LoopVar, NewOrder+LoopVar+1, (OrderCount - LoopVar - 1)*sizeof(CurrentOrder[0]));\r
+        NewCount--;\r
+      }\r
     }\r
+    Status = gRT->SetVariable(\r
+      Target == BcfgTargetBootOrder?(CHAR16*)L"BootOrder":(CHAR16*)L"DriverOrder",\r
+      (EFI_GUID*)&gEfiGlobalVariableGuid,\r
+      EFI_VARIABLE_NON_VOLATILE|EFI_VARIABLE_BOOTSERVICE_ACCESS|EFI_VARIABLE_RUNTIME_ACCESS,\r
+      NewCount*sizeof(NewOrder[0]),\r
+      NewOrder);\r
+    FreePool(NewOrder);\r
+  } else {\r
+    Status = EFI_OUT_OF_RESOURCES;\r
   }\r
-  Status = gRT->SetVariable(\r
-    Target == BcfgTargetBootOrder?(CHAR16*)L"BootOrder":(CHAR16*)L"DriverOrder",\r
-    (EFI_GUID*)&gEfiGlobalVariableGuid,\r
-    EFI_VARIABLE_NON_VOLATILE|EFI_VARIABLE_BOOTSERVICE_ACCESS|EFI_VARIABLE_RUNTIME_ACCESS,\r
-    NewCount*sizeof(NewOrder[0]),\r
-    NewOrder);\r
-  FreePool(NewOrder);\r
-\r
   if (EFI_ERROR(Status)) {\r
     ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_BCFG_WRITE_FAIL), gShellDebug1HiiHandle, Target == BcfgTargetBootOrder?(CHAR16*)L"BootOrder":(CHAR16*)L"DriverOrder", Status);\r
     return (SHELL_INVALID_PARAMETER);\r