K8:

1) Add in "check-for-null" before dereferencing a pointer.
2) unify the handling of EFI_OUT_OF_RESOURCE case.

git-svn-id: https://edk2.svn.sourceforge.net/svnroot/edk2/trunk/edk2@7310 6f19259b-4bc3-4df7-8a09-765794883524

4 files changed, 115 insertions, 70 deletions

diff --git a/MdeModulePkg/Library/ExtendedIfrSupportLib/Common.c b/MdeModulePkg/Library/ExtendedIfrSupportLib/Common.c
index 56f2b79..3853a2e 100644
--- a/MdeModulePkg/Library/ExtendedIfrSupportLib/Common.c
+++ b/MdeModulePkg/Library/ExtendedIfrSupportLib/Common.c
@@ -100,14 +100,19 @@ IfrLibExtractClassFromHiiHandle (
   BufferSize = 0;

   HiiPackageList = NULL;

   Status = gIfrLibHiiDatabase->ExportPackageLists (gIfrLibHiiDatabase, Handle, &BufferSize, HiiPackageList);

+  //

+  // Handle is a invalid handle. Check if Handle is corrupted.

+  //

   ASSERT (Status != EFI_NOT_FOUND);

+  //

+  // The return status should always be EFI_BUFFER_TOO_SMALL as input buffer's size is 0.

+  //

+  ASSERT (Status == EFI_BUFFER_TOO_SMALL);

   

-  if (Status == EFI_BUFFER_TOO_SMALL) {

-    HiiPackageList = AllocatePool (BufferSize);

-    ASSERT (HiiPackageList != NULL);

+  HiiPackageList = AllocatePool (BufferSize);

+  ASSERT (HiiPackageList != NULL);

 

-    Status = gIfrLibHiiDatabase->ExportPackageLists (gIfrLibHiiDatabase, Handle, &BufferSize, HiiPackageList);

-  }

+  Status = gIfrLibHiiDatabase->ExportPackageLists (gIfrLibHiiDatabase, Handle, &BufferSize, HiiPackageList);

   if (EFI_ERROR (Status)) {

     return Status;

   }

diff --git a/MdeModulePkg/Library/ExtendedIfrSupportLib/Form.c b/MdeModulePkg/Library/ExtendedIfrSupportLib/Form.c
index 7f7af9c..793efdf 100644
--- a/MdeModulePkg/Library/ExtendedIfrSupportLib/Form.c
+++ b/MdeModulePkg/Library/ExtendedIfrSupportLib/Form.c
@@ -376,15 +376,21 @@ IfrLibUpdateForm (
   BufferSize = 0;

   HiiPackageList   = NULL;

   Status = HiiDatabase->ExportPackageLists (HiiDatabase, Handle, &BufferSize, HiiPackageList);

-  if (Status == EFI_BUFFER_TOO_SMALL) {

-    HiiPackageList = AllocatePool (BufferSize);

-    ASSERT (HiiPackageList != NULL);

+  //

+  // Handle is a invalid handle. Check if Handle is corrupted.

+  //

+  ASSERT (Status != EFI_NOT_FOUND);

+  //

+  // The return status should always be EFI_BUFFER_TOO_SMALL as input buffer's size is 0.

+  //

+  ASSERT (Status == EFI_BUFFER_TOO_SMALL);

+  HiiPackageList = AllocatePool (BufferSize);

+  ASSERT (HiiPackageList != NULL);

 

-    Status = HiiDatabase->ExportPackageLists (HiiDatabase, Handle, &BufferSize, HiiPackageList);

-    if (EFI_ERROR (Status)) {

-      FreePool (HiiPackageList);

-      return Status;

-    }

+  Status = HiiDatabase->ExportPackageLists (HiiDatabase, Handle, &BufferSize, HiiPackageList);

+  if (EFI_ERROR (Status)) {

+    FreePool (HiiPackageList);

+    return Status;

   }

 

   //

diff --git a/MdeModulePkg/Library/UefiHiiLib/HiiLib.c b/MdeModulePkg/Library/UefiHiiLib/HiiLib.c
index f5544d9..83e33d2 100644
--- a/MdeModulePkg/Library/UefiHiiLib/HiiLib.c
+++ b/MdeModulePkg/Library/UefiHiiLib/HiiLib.c
@@ -281,13 +281,13 @@ HiiLibGetHiiHandles (
   OUT    EFI_HII_HANDLE            **HiiHandleBuffer

   )

 {

-  UINTN       BufferLength;

   EFI_STATUS  Status;

 

   ASSERT (HandleBufferLength != NULL);

   ASSERT (HiiHandleBuffer != NULL);

 

-  BufferLength = 0;

+  *HandleBufferLength = 0;

+  *HiiHandleBuffer = NULL;

 

   //

   // Try to find the actual buffer size for HiiHandle Buffer.

@@ -296,27 +296,26 @@ HiiLibGetHiiHandles (
                                  mHiiDatabaseProt,

                                  EFI_HII_PACKAGE_TYPE_ALL,

                                  NULL,

-                                 &BufferLength,

+                                 HandleBufferLength,

                                  *HiiHandleBuffer

                                  );

-

+  

   if (Status == EFI_BUFFER_TOO_SMALL) {

-      *HiiHandleBuffer = AllocateZeroPool (BufferLength);

+      *HiiHandleBuffer = AllocateZeroPool (*HandleBufferLength);

       ASSERT (*HiiHandleBuffer != NULL);

       Status = mHiiDatabaseProt->ListPackageLists (

                                      mHiiDatabaseProt,

                                      EFI_HII_PACKAGE_TYPE_ALL,

                                      NULL,

-                                     &BufferLength,

+                                     HandleBufferLength,

                                      *HiiHandleBuffer

                                      );

-      //

-      // we should not fail here.

-      //

-      ASSERT_EFI_ERROR (Status);

-  }

 

-  *HandleBufferLength = BufferLength;

+      if (EFI_ERROR (Status)) {

+        FreePool (*HiiHandleBuffer);

+        *HiiHandleBuffer = NULL;

+      }

+  }

 

   return Status;

 }

diff --git a/MdeModulePkg/Library/UefiIfrSupportLib/UefiIfrForm.c b/MdeModulePkg/Library/UefiIfrSupportLib/UefiIfrForm.c
index fb023c3..bfca72e 100644
--- a/MdeModulePkg/Library/UefiIfrSupportLib/UefiIfrForm.c
+++ b/MdeModulePkg/Library/UefiIfrSupportLib/UefiIfrForm.c
@@ -342,10 +342,10 @@ ExtractBlockName (
   //                 |   8  | 4 |  7   | 4 |

   //

   StringPtr = AllocateZeroPool ((BlockNameNumber * (8 + 4 + 7 + 4) + 1) * sizeof (CHAR16));

-  *BlockName = StringPtr;

   if (StringPtr == NULL) {

     return EFI_OUT_OF_RESOURCES;

   }

+  *BlockName = StringPtr;

 

   Buffer += sizeof (UINT32);

   for (Index = 0; Index < BlockNameNumber; Index++) {

@@ -522,6 +522,13 @@ ConstructConfigAltResp (
   if (ConfigAltResp == NULL) {

     return EFI_INVALID_PARAMETER;

   }

+  

+  DescHdr = NULL;

+  StringPtr = NULL;

+  AltCfg    = NULL;

+  ConfigResp = NULL;

+  BlockName = NULL;

+  NeedFreeConfigRequest = FALSE;

 

   //

   // Construct <ConfigHdr> : "GUID=...&NAME=...&PATH=..."

@@ -535,39 +542,47 @@ ConstructConfigAltResp (
              Name,

              DriverHandle

              );

-  if (Status == EFI_BUFFER_TOO_SMALL) {

-    ConfigHdr = AllocateZeroPool (StrBufferLen);

-    Status = ConstructConfigHdr (

-               ConfigHdr,

-               &StrBufferLen,

-               Guid,

-               Name,

-               DriverHandle

-               );

+  ASSERT (Status == EFI_BUFFER_TOO_SMALL);

+  ConfigHdr = AllocateZeroPool (StrBufferLen);

+  if (ConfigHdr == NULL) {

+    Status = EFI_OUT_OF_RESOURCES;

+    goto Exit;

   }

+  Status = ConstructConfigHdr (

+             ConfigHdr,

+             &StrBufferLen,

+             Guid,

+             Name,

+             DriverHandle

+             );

 

   if (EFI_ERROR (Status)) {

-    return Status;

+    goto Exit;

   }

 

   //

   // Construct <ConfigResp>

   //

-  NeedFreeConfigRequest = FALSE;

   if (ConfigRequest == NULL) {

     //

     // If ConfigRequest is set to NULL, export all configurable elements in BlockNameArray

     //

     Status = ExtractBlockName (BlockNameArray, &BlockName);

     if (EFI_ERROR (Status)) {

-      return Status;

+      goto Exit;

     }

     

     Len = StrSize (ConfigHdr);

     ConfigRequest = AllocateZeroPool (Len + StrSize (BlockName) - sizeof (CHAR16));

+    if (ConfigRequest == NULL) {

+      Status = EFI_OUT_OF_RESOURCES;

+      goto Exit;

+    }

+    

     StrCpy (ConfigRequest, ConfigHdr);

     StrCat (ConfigRequest, BlockName);

     NeedFreeConfigRequest = TRUE;

+

   }

 

   Status = gBS->LocateProtocol (&gEfiHiiConfigRoutingProtocolGuid, NULL, (VOID **) &HiiConfigRouting);

@@ -587,36 +602,48 @@ ConstructConfigAltResp (
     return Status;

   }

 

+  AltRespLen = 0;

   //

   // Construct <AltResp>

   //

-  DescHdr = AllocateZeroPool (NumberAltCfg * 16 * sizeof (CHAR16));

-  StringPtr = DescHdr;

-  AltCfg = AllocateZeroPool (NumberAltCfg * sizeof (CHAR16 *));

-  AltRespLen = 0;

-  VA_START (Args, NumberAltCfg);

-  for (Index = 0; Index < NumberAltCfg; Index++) {

-    AltCfgId = (UINT16) VA_ARG (Args, UINT16);

-    DefaultValueArray = (UINT8 *) VA_ARG (Args, VOID *);

-

-    //

-    // '&' <ConfigHdr>

-    //

-    AltRespLen += (StrLen (ConfigHdr) + 1);

-

-    StringPtr = DescHdr + Index * 16;

-    StrCpy (StringPtr, L"&ALTCFG=");

-    AltRespLen += (8 + sizeof (UINT16) * 2);

+  if (NumberAltCfg > 0) {

+    DescHdr = AllocateZeroPool (NumberAltCfg * 16 * sizeof (CHAR16));

+    if (DescHdr == NULL) {

+      Status = EFI_OUT_OF_RESOURCES;

+      goto Exit;

+    }

+    

+    StringPtr = DescHdr;

+    AltCfg = AllocateZeroPool (NumberAltCfg * sizeof (CHAR16 *));

+    if (AltCfg == NULL) {

+      Status = EFI_OUT_OF_RESOURCES;

+      goto Exit;

+    }

 

-    StrBufferLen = 5;

-    BufToHexString (StringPtr + 8, &StrBufferLen, (UINT8 *) &AltCfgId, sizeof (UINT16));

-    Status = ExtractBlockConfig (DefaultValueArray, &AltCfg[Index]);

-    if (EFI_ERROR (Status)) {

-      return Status;

+    VA_START (Args, NumberAltCfg);

+    for (Index = 0; Index < NumberAltCfg; Index++) {

+      AltCfgId = (UINT16) VA_ARG (Args, UINT16);

+      DefaultValueArray = (UINT8 *) VA_ARG (Args, VOID *);

+    

+      //

+      // '&' <ConfigHdr>

+      //

+      AltRespLen += (StrLen (ConfigHdr) + 1);

+    

+      StringPtr = DescHdr + Index * 16;

+      StrCpy (StringPtr, L"&ALTCFG=");

+      AltRespLen += (8 + sizeof (UINT16) * 2);

+    

+      StrBufferLen = 5;

+      BufToHexString (StringPtr + 8, &StrBufferLen, (UINT8 *) &AltCfgId, sizeof (UINT16));

+      Status = ExtractBlockConfig (DefaultValueArray, &AltCfg[Index]);

+      if (EFI_ERROR (Status)) {

+        goto Exit;

+      }

+      AltRespLen += StrLen (AltCfg[Index]);

     }

-    AltRespLen += StrLen (AltCfg[Index]);

+    VA_END (Args);

   }

-  VA_END (Args);

 

   //

   // Generate the final <ConfigAltResp>

@@ -625,7 +652,7 @@ ConstructConfigAltResp (
   TempStr = AllocateZeroPool (StrBufferLen);

   *ConfigAltResp = TempStr;

   if (TempStr == NULL) {

-    return EFI_OUT_OF_RESOURCES;

+    goto Exit;

   }

 

   //

@@ -641,13 +668,23 @@ ConstructConfigAltResp (
     FreePool (AltCfg[Index]);

   }

 

+Exit:

   if (NeedFreeConfigRequest) {

     FreePool (ConfigRequest);

   }

   FreePool (ConfigHdr);

-  FreePool (ConfigResp);

-  FreePool (DescHdr);

-  FreePool (AltCfg);

+  if (ConfigResp != NULL) {

+    FreePool (ConfigResp);

+  }

+

+  if (BlockName != NULL) {

+    FreePool (BlockName);

+  }

+

+  if (NumberAltCfg > 0) {

+    FreePool (DescHdr);

+    FreePool (AltCfg);

+  }

 

   return EFI_SUCCESS;

 }

@@ -1213,7 +1250,7 @@ GetBrowserData (
   BufferLen = 0x4000;

   ConfigResp = AllocateZeroPool (BufferLen + HeaderLen);

   if (ConfigResp == NULL) {

-    BufferLen = 0;

+    return EFI_OUT_OF_RESOURCES;

   }

 

   StringPtr = ConfigResp + HeaderLen;

@@ -1229,9 +1266,7 @@ GetBrowserData (
                            VariableName

                            );

   if (Status == EFI_BUFFER_TOO_SMALL) {

-    if (ConfigResp != NULL) {

-      FreePool (ConfigResp);

-    }

+    FreePool (ConfigResp);

 

     ConfigResp = AllocateZeroPool (BufferLen + HeaderLen);

     if (ConfigResp == NULL) {