Refine the code logic, use dynamic allocate buffer instead of static array to fix potential buffer overflow.

Contributed-under: TianoCore Contribution Agreement 1.0
Signed-off-by: Eric Dong <eric.dong@intel.com>
Reviewed-by: Liming Gao <liming.gao@intel.com>

git-svn-id: https://svn.code.sf.net/p/edk2/code/trunk/edk2@15829 6f19259b-4bc3-4df7-8a09-765794883524

2 files changed, 100 insertions, 69 deletions

diff --git a/MdeModulePkg/Universal/PlatformDriOverrideDxe/PlatDriOverrideDxe.c b/MdeModulePkg/Universal/PlatformDriOverrideDxe/PlatDriOverrideDxe.c
index 0f0da166ad..d29c050a03 100644
--- a/MdeModulePkg/Universal/PlatformDriOverrideDxe/PlatDriOverrideDxe.c
+++ b/MdeModulePkg/Universal/PlatformDriOverrideDxe/PlatDriOverrideDxe.c
@@ -70,16 +70,16 @@ EFI_HANDLE                   *mDevicePathHandleBuffer;
 EFI_HANDLE                   *mDriverImageHandleBuffer;

 

 INTN                         mSelectedCtrIndex;

-EFI_STRING_ID                mControllerToken[MAX_CHOICE_NUM];

+EFI_STRING_ID                *mControllerToken;

 UINTN                        mDriverImageHandleCount;

-EFI_STRING_ID                mDriverImageToken[MAX_CHOICE_NUM];

-EFI_STRING_ID                mDriverImageFilePathToken[MAX_CHOICE_NUM];

-EFI_LOADED_IMAGE_PROTOCOL    *mDriverImageProtocol[MAX_CHOICE_NUM];

-EFI_DEVICE_PATH_PROTOCOL     *mControllerDevicePathProtocol[MAX_CHOICE_NUM];

+EFI_STRING_ID                *mDriverImageToken;

+EFI_DEVICE_PATH_PROTOCOL     **mControllerDevicePathProtocol;

 UINTN                        mSelectedDriverImageNum;

 UINTN                        mLastSavedDriverImageNum;

 UINT16                       mCurrentPage;

 EFI_CALLBACK_INFO           *mCallbackInfo;

+BOOLEAN                     *mDriSelection;

+UINTN                        mMaxDeviceCount;

 

 HII_VENDOR_DEVICE_PATH  mHiiVendorDevicePath = {

   {

@@ -425,6 +425,12 @@ UpdateDeviceSelectPage (
     return EFI_SUCCESS;

   }

 

+  mMaxDeviceCount = DevicePathHandleCount;

+  mControllerDevicePathProtocol = AllocateZeroPool (DevicePathHandleCount * sizeof (EFI_DEVICE_PATH_PROTOCOL *));

+  ASSERT (mControllerDevicePathProtocol != NULL);

+  mControllerToken = AllocateZeroPool (DevicePathHandleCount * sizeof (EFI_STRING_ID));

+  ASSERT (mControllerToken != NULL);

+

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

     if (FakeNvData->PciDeviceFilter == 0x01) {

       //

@@ -630,6 +636,9 @@ UpdateBindingDriverSelectPage (
   VOID                                      *EndOpCodeHandle;

   EFI_IFR_GUID_LABEL                        *StartLabel;

   EFI_IFR_GUID_LABEL                        *EndLabel;

+  EFI_LOADED_IMAGE_PROTOCOL                 **DriverImageProtocol;

+  EFI_STRING_ID                             *DriverImageFilePathToken;

+  UINT8                                     CheckFlags;

 

   //

   // If user select a controller item in the first page  the following code will be run.

@@ -698,6 +707,16 @@ UpdateBindingDriverSelectPage (
     return EFI_NOT_FOUND;

   }

 

+  mDriverImageToken = AllocateZeroPool (DriverImageHandleCount * sizeof (EFI_STRING_ID));

+  ASSERT (mDriverImageToken != NULL);

+  mDriSelection = AllocateZeroPool (DriverImageHandleCount * sizeof (BOOLEAN));

+  ASSERT (mDriSelection != NULL);

+

+  DriverImageProtocol = AllocateZeroPool (DriverImageHandleCount * sizeof (EFI_LOADED_IMAGE_PROTOCOL *));

+  ASSERT (DriverImageProtocol != NULL);

+  DriverImageFilePathToken = AllocateZeroPool (DriverImageHandleCount * sizeof (EFI_STRING_ID));

+  ASSERT (DriverImageFilePathToken != NULL);

+

   mDriverImageHandleCount = DriverImageHandleCount;

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

     //

@@ -718,16 +737,16 @@ UpdateBindingDriverSelectPage (
                     EFI_OPEN_PROTOCOL_GET_PROTOCOL

                     );

     if (EFI_ERROR (Status)) {

-      FakeNvData->DriSelection[Index] = 0x00;

+      mDriSelection[Index] = FALSE;

       continue;

     }

-    mDriverImageProtocol[Index] = LoadedImage;

+    DriverImageProtocol[Index] = LoadedImage;

     //

     // Find its related driver binding protocol

     //

     DriverBindingHandle = GetDriverBindingHandleFromImageHandle (mDriverImageHandleBuffer[Index]);

     if (DriverBindingHandle == NULL) {

-      FakeNvData->DriSelection[Index] = 0x00;

+      mDriSelection[Index] = FALSE;

       continue;

     }

 

@@ -741,7 +760,7 @@ UpdateBindingDriverSelectPage (
                         (VOID **) &LoadedImageDevicePath

                         );

     if (LoadedImageDevicePath == NULL) {

-      FakeNvData->DriSelection[Index] = 0x00;

+      mDriSelection[Index] = FALSE;

       continue;

     }

 

@@ -757,11 +776,11 @@ UpdateBindingDriverSelectPage (
                          (VOID **) &BusSpecificDriverOverride

                          );

         if (EFI_ERROR (Status) || BusSpecificDriverOverride == NULL) {

-          FakeNvData->DriSelection[Index] = 0x00;

+          mDriSelection[Index] = FALSE;

           continue;

         }

       } else {

-        FakeNvData->DriSelection[Index] = 0x00;

+        mDriSelection[Index] = FALSE;

         continue;

       }

     }

@@ -798,9 +817,9 @@ UpdateBindingDriverSelectPage (
     NewString = AllocateZeroPool (StrSize (DriverName));

     ASSERT (NewString != NULL); 

     if (EFI_ERROR (CheckMapping (mControllerDevicePathProtocol[mSelectedCtrIndex], LoadedImageDevicePath, &mMappingDataBase, NULL, NULL))) {

-      FakeNvData->DriSelection[Index] = 0x00;

+      mDriSelection[Index] = FALSE;

     } else {

-      FakeNvData->DriSelection[Index] = 0x01;

+      mDriSelection[Index] = TRUE;

       mLastSavedDriverImageNum++;

     }

     StrCat (NewString, DriverName);

@@ -820,21 +839,26 @@ UpdateBindingDriverSelectPage (
     NewString = AllocateZeroPool (StrSize (DriverName));

     ASSERT (NewString != NULL); 

     StrCat (NewString, DriverName);

-    NewStringHelpToken = HiiSetString (Private->RegisteredHandle, mDriverImageFilePathToken[Index], NewString, NULL);

+    NewStringHelpToken = HiiSetString (Private->RegisteredHandle, DriverImageFilePathToken[Index], NewString, NULL);

     ASSERT (NewStringHelpToken != 0);

-    mDriverImageFilePathToken[Index] = NewStringHelpToken;

+    DriverImageFilePathToken[Index] = NewStringHelpToken;

     FreePool (NewString);

     FreePool (DriverName);

 

+    CheckFlags        = 0;

+    if (mDriSelection[Index]) {

+      CheckFlags |= EFI_IFR_CHECKBOX_DEFAULT;

+    }

+

     HiiCreateCheckBoxOpCode (

       StartOpCodeHandle,

-      (UINT16) (DRIVER_SELECTION_QUESTION_ID + Index),

-      VARSTORE_ID_PLAT_OVER_MNGR,

-      (UINT16) (DRIVER_SELECTION_VAR_OFFSET + Index),

-      NewStringToken,

-      NewStringHelpToken,

+      (UINT16) (KEY_VALUE_DRIVER_OFFSET + Index),

       0,

       0,

+      NewStringToken,

+      NewStringHelpToken,

+      EFI_IFR_FLAG_CALLBACK,

+      CheckFlags,

       NULL

       );

   }

@@ -852,6 +876,15 @@ UpdateBindingDriverSelectPage (
 

   HiiFreeOpCodeHandle (StartOpCodeHandle);

   HiiFreeOpCodeHandle (EndOpCodeHandle);

+

+  if (DriverImageProtocol != NULL) {

+    FreePool (DriverImageProtocol);

+  }

+

+  if (DriverImageFilePathToken != NULL) {

+    FreePool (DriverImageFilePathToken);

+  }

+

   return EFI_SUCCESS;

 }

 

@@ -932,7 +965,7 @@ UpdatePrioritySelectPage (
   //

   SelectedDriverImageNum = 0;

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

-    if (FakeNvData->DriSelection[Index] != 0) {

+    if (mDriSelection[Index]) {

       SelectedDriverImageNum ++;

     }

   }

@@ -950,7 +983,7 @@ UpdatePrioritySelectPage (
   //

   SelectedDriverImageNum = 0;

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

-    if (FakeNvData->DriSelection[Index] != 0) {

+    if (mDriSelection[Index]) {

       //

       // Use the NO. in driver binding buffer as value, will use it later

       //

@@ -1068,7 +1101,7 @@ UpdatePrioritySelectPage (
 

 **/

 EFI_STATUS

-CommintChanges (

+CommitChanges (

   IN EFI_CALLBACK_INFO                *Private,

   IN UINT16                           KeyValue,

   IN PLAT_OVER_MNGR_DATA              *FakeNvData

@@ -1263,21 +1296,10 @@ PlatOverMngrRouteConfig (
   }

 

   Status = EFI_SUCCESS;

-  if (mCurrentPage == FORM_ID_DRIVER) {

-    KeyValue = KEY_VALUE_DRIVER_GOTO_ORDER;

-    UpdatePrioritySelectPage (Private, KeyValue, FakeNvData);

-    KeyValue = KEY_VALUE_ORDER_SAVE_AND_EXIT;

-    Status = CommintChanges (Private, KeyValue, FakeNvData);

-    //

-    // Since UpdatePrioritySelectPage will change mCurrentPage,

-    // should ensure the mCurrentPage still indicate the second page here

-    //

-    mCurrentPage = FORM_ID_DRIVER;

-  }

 

   if (mCurrentPage == FORM_ID_ORDER) {

     KeyValue = KEY_VALUE_ORDER_SAVE_AND_EXIT;

-    Status = CommintChanges (Private, KeyValue, FakeNvData);

+    Status = CommitChanges (Private, KeyValue, FakeNvData);

   }

 

   return Status;

@@ -1350,7 +1372,7 @@ PlatOverMngrCallback (
       }

     }

 

-    if (((KeyValue >= KEY_VALUE_DEVICE_OFFSET) && (KeyValue < KEY_VALUE_DEVICE_MAX)) || (KeyValue == KEY_VALUE_ORDER_GOTO_PREVIOUS)) {

+    if (((KeyValue >= KEY_VALUE_DEVICE_OFFSET) && (KeyValue < KEY_VALUE_DEVICE_OFFSET + mMaxDeviceCount)) || (KeyValue == KEY_VALUE_ORDER_GOTO_PREVIOUS)) {

       if (KeyValue == KEY_VALUE_ORDER_GOTO_PREVIOUS) {

         KeyValue = (EFI_QUESTION_ID) (mSelectedCtrIndex + KEY_VALUE_DEVICE_OFFSET);

       }

@@ -1384,30 +1406,34 @@ PlatOverMngrCallback (
       UpdateDeviceSelectPage (Private, KeyValue, FakeNvData);

     }

   } else if (Action == EFI_BROWSER_ACTION_CHANGED) {

-    switch (KeyValue) {

-    case KEY_VALUE_DEVICE_REFRESH:

-    case KEY_VALUE_DEVICE_FILTER:

-      UpdateDeviceSelectPage (Private, KeyValue, FakeNvData);

-      //

-      // Update page title string

-      //

-      NewStringToken = STRING_TOKEN (STR_TITLE);

-      if (HiiSetString (Private->RegisteredHandle, NewStringToken, L"First, Select the controller by device path", NULL) == 0) {

-        ASSERT (FALSE);

-      }

-    break;

-    

-    case KEY_VALUE_ORDER_SAVE_AND_EXIT:

-      Status = CommintChanges (Private, KeyValue, FakeNvData);

-      *ActionRequest = EFI_BROWSER_ACTION_REQUEST_SUBMIT;

-      if (EFI_ERROR (Status)) {

-        CreatePopUp (EFI_LIGHTGRAY | EFI_BACKGROUND_BLUE, &Key, L"Single Override Info too large, Saving Error!", NULL);

-        return EFI_DEVICE_ERROR;

-      }

-    break;

+    if ((KeyValue >= KEY_VALUE_DRIVER_OFFSET) && (KeyValue < KEY_VALUE_DRIVER_OFFSET + mDriverImageHandleCount)) {

+      mDriSelection[KeyValue - KEY_VALUE_DRIVER_OFFSET] = Value->b;

+    } else {

+      switch (KeyValue) {

+      case KEY_VALUE_DEVICE_REFRESH:

+      case KEY_VALUE_DEVICE_FILTER:

+        UpdateDeviceSelectPage (Private, KeyValue, FakeNvData);

+        //

+        // Update page title string

+        //

+        NewStringToken = STRING_TOKEN (STR_TITLE);

+        if (HiiSetString (Private->RegisteredHandle, NewStringToken, L"First, Select the controller by device path", NULL) == 0) {

+          ASSERT (FALSE);

+        }

+      break;

+      

+      case KEY_VALUE_ORDER_SAVE_AND_EXIT:

+        Status = CommitChanges (Private, KeyValue, FakeNvData);

+        *ActionRequest = EFI_BROWSER_ACTION_REQUEST_SUBMIT;

+        if (EFI_ERROR (Status)) {

+          CreatePopUp (EFI_LIGHTGRAY | EFI_BACKGROUND_BLUE, &Key, L"Single Override Info too large, Saving Error!", NULL);

+          return EFI_DEVICE_ERROR;

+        }

+      break;

 

-    default:

-    break;

+      default:

+      break;

+      }

     }

   }

 

@@ -1659,10 +1685,6 @@ PlatDriOverrideDxeInit (
   //

   mDriverImageHandleCount = 0;

   mCurrentPage = 0;

-  ZeroMem (mDriverImageToken, MAX_CHOICE_NUM * sizeof (EFI_STRING_ID));

-  ZeroMem (mDriverImageFilePathToken, MAX_CHOICE_NUM * sizeof (EFI_STRING_ID));

-  ZeroMem (mControllerToken, MAX_CHOICE_NUM * sizeof (EFI_STRING_ID));

-  ZeroMem (mDriverImageProtocol, MAX_CHOICE_NUM * sizeof (EFI_LOADED_IMAGE_PROTOCOL *));

   

   return EFI_SUCCESS;

 

@@ -1706,5 +1728,17 @@ PlatDriOverrideDxeUnload (
 

   FreePool (mCallbackInfo);

 

+  if (mControllerToken != NULL) {

+    FreePool (mControllerToken);

+  }

+

+  if (mControllerDevicePathProtocol != NULL) {

+    FreePool (mControllerDevicePathProtocol);

+  }

+

+  if (mDriverImageToken != NULL) {

+    FreePool (mDriverImageToken);

+  }

+

   return EFI_SUCCESS;

 }

diff --git a/MdeModulePkg/Universal/PlatformDriOverrideDxe/PlatOverMngr.h b/MdeModulePkg/Universal/PlatformDriOverrideDxe/PlatOverMngr.h
index 1ecdaafacc..ca2094551d 100644
--- a/MdeModulePkg/Universal/PlatformDriOverrideDxe/PlatOverMngr.h
+++ b/MdeModulePkg/Universal/PlatformDriOverrideDxe/PlatOverMngr.h
@@ -3,7 +3,7 @@
   The defintions are required both by Source code and Vfr file.

   The PLAT_OVER_MNGR_DATA structure, form guid and Ifr question ID are defined. 

 

-Copyright (c) 2007 - 2011, Intel Corporation. All rights reserved.<BR>

+Copyright (c) 2007 - 2014, Intel Corporation. All rights reserved.<BR>

 This program and the accompanying materials

 are licensed and made available under the terms and conditions of the BSD License

 which accompanies this distribution.  The full text of the license may be found at

@@ -22,7 +22,7 @@ WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
 //

 // The max number of the supported driver list.

 //

-#define MAX_CHOICE_NUM    0x00ff

+#define MAX_CHOICE_NUM    0x00FF

 #define UPDATE_DATA_SIZE  0x1000

 

 #define FORM_ID_DEVICE                 0x1100

@@ -30,7 +30,7 @@ WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
 #define FORM_ID_ORDER                  0x1500

 

 #define KEY_VALUE_DEVICE_OFFSET        0x0100

-#define KEY_VALUE_DEVICE_MAX           (KEY_VALUE_DEVICE_OFFSET + MAX_CHOICE_NUM)

+#define KEY_VALUE_DRIVER_OFFSET        0x0300

 

 #define KEY_VALUE_DEVICE_REFRESH       0x1234

 #define KEY_VALUE_DEVICE_FILTER        0x1235

@@ -47,7 +47,6 @@ WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
 #define LABEL_END                      0xffff

 

 typedef struct {

-  UINT8   DriSelection[MAX_CHOICE_NUM];

   UINT8   DriOrder[MAX_CHOICE_NUM];

   UINT8   PciDeviceFilter;

 } PLAT_OVER_MNGR_DATA;

@@ -56,7 +55,6 @@ typedef struct {
 // Field offset of structure PLAT_OVER_MNGR_DATA

 //

 #define VAR_OFFSET(Field)              ((UINTN) &(((PLAT_OVER_MNGR_DATA *) 0)->Field))

-#define DRIVER_SELECTION_VAR_OFFSET     (VAR_OFFSET (DriSelection))

 #define DRIVER_ORDER_VAR_OFFSET         (VAR_OFFSET (DriOrder))

 

 //

@@ -64,7 +62,6 @@ typedef struct {
 // In order to avoid to conflict them, the Driver Selection and Order QuestionID offset is defined from 0x0500.

 //

 #define QUESTION_ID_OFFSET              0x0500

-#define DRIVER_SELECTION_QUESTION_ID    (VAR_OFFSET (DriSelection) + QUESTION_ID_OFFSET)

 #define DRIVER_ORDER_QUESTION_ID        (VAR_OFFSET (DriOrder) + QUESTION_ID_OFFSET)

 

 #endif