[Patch 2/3] NetworkPkg/UefiPxeBcDxe: Correct the handle for PXE Base Code Callback Protocol.

[Jiaxin Wu <jiaxin.wu@intel.com>]

According UEFI Spec:
The PXE Base Code Callback Protocol must be on the same handle as the PXE
Base Code Protocol.

But current implementation doesn't follow that. This patch is fix that issue.

Cc: Ye Ting <ting.ye@intel.com>
Cc: Fu Siyuan <siyuan.fu@intel.com>
Contributed-under: TianoCore Contribution Agreement 1.0
Signed-off-by: Wu Jiaxin <jiaxin.wu@intel.com>
---
 NetworkPkg/UefiPxeBcDxe/PxeBcBoot.c | 6 +++---
 NetworkPkg/UefiPxeBcDxe/PxeBcImpl.c | 2 +-
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/NetworkPkg/UefiPxeBcDxe/PxeBcBoot.c b/NetworkPkg/UefiPxeBcDxe/PxeBcBoot.c
index 0bb1d66..1f8895f 100644
--- a/NetworkPkg/UefiPxeBcDxe/PxeBcBoot.c
+++ b/NetworkPkg/UefiPxeBcDxe/PxeBcBoot.c
@@ -992,11 +992,11 @@ PxeBcInstallCallback (
   //
   // Check whether PxeBaseCodeCallbackProtocol already installed.
   //
   PxeBc  = &Private->PxeBc;
   Status = gBS->HandleProtocol (
-                  Private->Controller,
+                  Private->Mode.UsingIpv6 ? Private->Ip6Nic->Controller : Private->Ip4Nic->Controller,
                   &gEfiPxeBaseCodeCallbackProtocolGuid,
                   (VOID **) &Private->PxeBcCallback
                   );
   if (Status == EFI_UNSUPPORTED) {
 
@@ -1008,11 +1008,11 @@ PxeBcInstallCallback (
 
     //
     // Install a default callback if user didn't offer one.
     //
     Status = gBS->InstallProtocolInterface (
-                    &Private->Controller,
+                    Private->Mode.UsingIpv6 ? &Private->Ip6Nic->Controller : &Private->Ip4Nic->Controller,
                     &gEfiPxeBaseCodeCallbackProtocolGuid,
                     EFI_NATIVE_INTERFACE,
                     &Private->LoadFileCallback
                     );
 
@@ -1052,11 +1052,11 @@ PxeBcUninstallCallback (
     NewMakeCallback = FALSE;
 
     PxeBc->SetParameters (PxeBc, NULL, NULL, NULL, NULL, &NewMakeCallback);
 
     gBS->UninstallProtocolInterface (
-          Private->Controller,
+          Private->Mode.UsingIpv6 ? Private->Ip6Nic->Controller : Private->Ip4Nic->Controller,
           &gEfiPxeBaseCodeCallbackProtocolGuid,
           &Private->LoadFileCallback
           );
   }
 }
diff --git a/NetworkPkg/UefiPxeBcDxe/PxeBcImpl.c b/NetworkPkg/UefiPxeBcDxe/PxeBcImpl.c
index ab9e494..1fc26c5 100644
--- a/NetworkPkg/UefiPxeBcDxe/PxeBcImpl.c
+++ b/NetworkPkg/UefiPxeBcDxe/PxeBcImpl.c
@@ -1923,11 +1923,11 @@ EfiPxeBcSetParameters (
     if (*NewMakeCallback) {
       //
       // Update the previous PxeBcCallback protocol.
       //
       Status = gBS->HandleProtocol (
-                      Private->Controller,
+                      Mode->UsingIpv6 ? Private->Ip6Nic->Controller : Private->Ip4Nic->Controller,
                       &gEfiPxeBaseCodeCallbackProtocolGuid,
                       (VOID **) &Private->PxeBcCallback
                       );
 
       if (EFI_ERROR (Status) || (Private->PxeBcCallback->Callback == NULL)) {
-- 
1.9.5.msysgit.1