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

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);
 
   Status = gIfrLibHiiDatabase->ExportPackageLists (gIfrLibHiiDatabase, Handle, &BufferSize, HiiPackageList);
-  }
   if (EFI_ERROR (Status)) {
     return Status;
   }

MdeModulePkg/Library/ExtendedIfrSupportLib/Form.c

@@ -376,7 +376,14 @@ IfrLibUpdateForm (
   BufferSize = 0;
   HiiPackageList   = NULL;
   Status = HiiDatabase->ExportPackageLists (HiiDatabase, Handle, &BufferSize, HiiPackageList);
-  if (Status == EFI_BUFFER_TOO_SMALL) {
+  //
+  // 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);
 
@@ -385,7 +392,6 @@ IfrLibUpdateForm (
     FreePool (HiiPackageList);
     return Status;
   }
-  }
 
   //
   // Calculate and allocate space for retrieval of IFR data

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;
 }

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++) {
@@ -523,6 +523,13 @@ ConstructConfigAltResp (
     return EFI_INVALID_PARAMETER;
   }
   
+  DescHdr = NULL;
+  StringPtr = NULL;
+  AltCfg    = NULL;
+  ConfigResp = NULL;
+  BlockName = NULL;
+  NeedFreeConfigRequest = FALSE;
+
   //
   // Construct <ConfigHdr> : "GUID=...&NAME=...&PATH=..."
   //
@@ -535,8 +542,12 @@ ConstructConfigAltResp (
              Name,
              DriverHandle
              );
-  if (Status == EFI_BUFFER_TOO_SMALL) {
+  ASSERT (Status == EFI_BUFFER_TOO_SMALL);
   ConfigHdr = AllocateZeroPool (StrBufferLen);
+  if (ConfigHdr == NULL) {
+    Status = EFI_OUT_OF_RESOURCES;
+    goto Exit;
+  }
   Status = ConstructConfigHdr (
              ConfigHdr,
              &StrBufferLen,
@@ -544,30 +555,34 @@ ConstructConfigAltResp (
              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,13 +602,24 @@ ConstructConfigAltResp (
     return Status;
   }
 
+  AltRespLen = 0;
   //
   // Construct <AltResp>
   //
+  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 *));
-  AltRespLen = 0;
+    if (AltCfg == NULL) {
+      Status = EFI_OUT_OF_RESOURCES;
+      goto Exit;
+    }
+
     VA_START (Args, NumberAltCfg);
     for (Index = 0; Index < NumberAltCfg; Index++) {
       AltCfgId = (UINT16) VA_ARG (Args, UINT16);
@@ -612,11 +638,12 @@ ConstructConfigAltResp (
       BufToHexString (StringPtr + 8, &StrBufferLen, (UINT8 *) &AltCfgId, sizeof (UINT16));
       Status = ExtractBlockConfig (DefaultValueArray, &AltCfg[Index]);
       if (EFI_ERROR (Status)) {
-      return Status;
+        goto Exit;
       }
       AltRespLen += StrLen (AltCfg[Index]);
     }
     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);
+  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);
-    }
 
     ConfigResp = AllocateZeroPool (BufferLen + HeaderLen);
     if (ConfigResp == NULL) {