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

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);
+  // Handle is a invalid handle. Check if Handle is corrupted.
-    ASSERT (HiiPackageList != NULL);
+  //
+  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);
+  Status = HiiDatabase->ExportPackageLists (HiiDatabase, Handle, &BufferSize, HiiPackageList);
-    if (EFI_ERROR (Status)) {
+  if (EFI_ERROR (Status)) {
-      FreePool (HiiPackageList);
+    FreePool (HiiPackageList);
-      return Status;
+    return Status;
-    }
   }
 
   //

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++) {
@@ -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) {
+  ASSERT (Status == EFI_BUFFER_TOO_SMALL);
-    ConfigHdr = AllocateZeroPool (StrBufferLen);
+  ConfigHdr = AllocateZeroPool (StrBufferLen);
-    Status = ConstructConfigHdr (
+  if (ConfigHdr == NULL) {
-               ConfigHdr,
+    Status = EFI_OUT_OF_RESOURCES;
-               &StrBufferLen,
+    goto Exit;
-               Guid,
-               Name,
-               DriverHandle
-               );
   }
+  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));
+  if (NumberAltCfg > 0) {
-  StringPtr = DescHdr;
+    DescHdr = AllocateZeroPool (NumberAltCfg * 16 * sizeof (CHAR16));
-  AltCfg = AllocateZeroPool (NumberAltCfg * sizeof (CHAR16 *));
+    if (DescHdr == NULL) {
-  AltRespLen = 0;
+      Status = EFI_OUT_OF_RESOURCES;
-  VA_START (Args, NumberAltCfg);
+      goto Exit;
-  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)) {
-      return Status;
     }
-    AltRespLen += StrLen (AltCfg[Index]);
+    
+    StringPtr = DescHdr;
+    AltCfg = AllocateZeroPool (NumberAltCfg * sizeof (CHAR16 *));
+    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);
+      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]);
+    }
+    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);
+  if (ConfigResp != NULL) {
-  FreePool (DescHdr);
+    FreePool (ConfigResp);
-  FreePool (AltCfg);
+  }
+
+  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) {