NetworkPkg/HttpBootDxe: Update device path node to include DNS information

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>
Reviewed-by: Ye Ting <ting.ye@intel.com>

3 files changed, 121 insertions, 24 deletions

diff --git a/NetworkPkg/HttpBootDxe/HttpBootClient.c b/NetworkPkg/HttpBootDxe/HttpBootClient.c
index 68f5a49..5c87171 100644
--- a/NetworkPkg/HttpBootDxe/HttpBootClient.c
+++ b/NetworkPkg/HttpBootDxe/HttpBootClient.c
@@ -16,7 +16,7 @@ WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
 #include "HttpBootDxe.h"

 

 /**

-  Update the IP and URL device path node to include the boot resource information.

+  Update the device path node to include the boot resource information.

 

   @param[in]    Private            The pointer to the driver's private data.

 

@@ -31,12 +31,14 @@ HttpBootUpdateDevicePath (
   )

 {

   EFI_DEV_PATH               *Node;

-  EFI_DEVICE_PATH_PROTOCOL   *TmpDevicePath;

+  EFI_DEVICE_PATH_PROTOCOL   *TmpIpDevicePath;

+  EFI_DEVICE_PATH_PROTOCOL   *TmpDnsDevicePath;

   EFI_DEVICE_PATH_PROTOCOL   *NewDevicePath;

   UINTN                      Length;

   EFI_STATUS                 Status;

 

-  TmpDevicePath = NULL;

+  TmpIpDevicePath  = NULL;

+  TmpDnsDevicePath = NULL;

   

   //

   // Update the IP node with DHCP assigned information.

@@ -72,29 +74,65 @@ HttpBootUpdateDevicePath (
     CopyMem (&Node->Ipv6.GatewayIpAddress, &Private->GatewayIp.v6, sizeof (EFI_IPv6_ADDRESS));

   }

   

-  TmpDevicePath = AppendDevicePathNode (Private->ParentDevicePath, (EFI_DEVICE_PATH_PROTOCOL*) Node);

+  TmpIpDevicePath = AppendDevicePathNode (Private->ParentDevicePath, (EFI_DEVICE_PATH_PROTOCOL*) Node);

   FreePool (Node);

-  if (TmpDevicePath == NULL) {

+  if (TmpIpDevicePath == NULL) {

     return EFI_OUT_OF_RESOURCES;

   }

 

   //

+  // Update the DNS node with DNS server IP list if existed.

+  //

+  if (Private->DnsServerIp != NULL) {

+    Length = sizeof (EFI_DEVICE_PATH_PROTOCOL) + sizeof (Node->Dns.IsIPv6) + Private->DnsServerCount * sizeof (EFI_IP_ADDRESS);

+    Node = AllocatePool (Length);

+    if (Node == NULL) {

+      FreePool (TmpIpDevicePath);

+      return EFI_OUT_OF_RESOURCES;

+    }

+    Node->DevPath.Type    = MESSAGING_DEVICE_PATH;

+    Node->DevPath.SubType = MSG_DNS_DP;

+    SetDevicePathNodeLength (Node, Length);

+    Node->Dns.IsIPv6 = Private->UsingIpv6 ? 0x01 : 0x00;

+    CopyMem ((UINT8*) Node + sizeof (EFI_DEVICE_PATH_PROTOCOL) + sizeof (Node->Dns.IsIPv6), Private->DnsServerIp, Private->DnsServerCount * sizeof (EFI_IP_ADDRESS));

+    

+    TmpDnsDevicePath = AppendDevicePathNode (TmpIpDevicePath, (EFI_DEVICE_PATH_PROTOCOL*) Node);

+    FreePool (Node);

+    FreePool (TmpIpDevicePath);

+    TmpIpDevicePath = NULL;

+    if (TmpDnsDevicePath == NULL) {

+      return EFI_OUT_OF_RESOURCES;

+    }

+  }

+

+  //

   // Update the URI node with the boot file URI.

   //

   Length = sizeof (EFI_DEVICE_PATH_PROTOCOL) + AsciiStrSize (Private->BootFileUri);

   Node = AllocatePool (Length);

   if (Node == NULL) {

-    FreePool (TmpDevicePath);

+    if (TmpIpDevicePath != NULL) {

+      FreePool (TmpIpDevicePath);

+    }

+    if (TmpDnsDevicePath != NULL) {

+      FreePool (TmpDnsDevicePath);

+    }

     return EFI_OUT_OF_RESOURCES;

   }

   Node->DevPath.Type    = MESSAGING_DEVICE_PATH;

   Node->DevPath.SubType = MSG_URI_DP;

   SetDevicePathNodeLength (Node, Length);

   CopyMem ((UINT8*) Node + sizeof (EFI_DEVICE_PATH_PROTOCOL), Private->BootFileUri, AsciiStrSize (Private->BootFileUri));

-  

-  NewDevicePath = AppendDevicePathNode (TmpDevicePath, (EFI_DEVICE_PATH_PROTOCOL*) Node);

+

+  if (TmpDnsDevicePath != NULL) {

+    NewDevicePath = AppendDevicePathNode (TmpDnsDevicePath, (EFI_DEVICE_PATH_PROTOCOL*) Node);

+    FreePool (TmpDnsDevicePath);

+  } else {

+    ASSERT (TmpIpDevicePath != NULL);

+    NewDevicePath = AppendDevicePathNode (TmpIpDevicePath, (EFI_DEVICE_PATH_PROTOCOL*) Node);

+    FreePool (TmpIpDevicePath);

+  }

   FreePool (Node);

-  FreePool (TmpDevicePath);

   if (NewDevicePath == NULL) {

     return EFI_OUT_OF_RESOURCES;

   }

@@ -153,6 +191,7 @@ HttpBootDhcp4ExtractUriInfo (
   HTTP_BOOT_DHCP4_PACKET_CACHE    *HttpOffer;

   UINT32                          SelectIndex;

   UINT32                          ProxyIndex;

+  UINT32                          DnsServerIndex;

   EFI_DHCP4_PACKET_OPTION         *Option;

   EFI_STATUS                      Status;

 

@@ -161,6 +200,8 @@ HttpBootDhcp4ExtractUriInfo (
   SelectIndex = Private->SelectIndex - 1;

   ASSERT (SelectIndex < HTTP_BOOT_OFFER_MAX_NUM);

 

+  DnsServerIndex = 0;

+  

   Status = EFI_SUCCESS;

 

   //

@@ -200,20 +241,37 @@ HttpBootDhcp4ExtractUriInfo (
     return Status;

   }

 

-  //

-  // Configure the default DNS server if server assigned.

-  //

   if ((SelectOffer->OfferType == HttpOfferTypeDhcpNameUriDns) || 

       (SelectOffer->OfferType == HttpOfferTypeDhcpDns) ||

       (SelectOffer->OfferType == HttpOfferTypeDhcpIpUriDns)) {

     Option = SelectOffer->OptList[HTTP_BOOT_DHCP4_TAG_INDEX_DNS_SERVER];

     ASSERT (Option != NULL);

+

+    //

+    // Record the Dns Server address list.

+    //

+    Private->DnsServerCount = (Option->Length) / sizeof (EFI_IPv4_ADDRESS);

+

+    Private->DnsServerIp = AllocateZeroPool (Private->DnsServerCount * sizeof (EFI_IP_ADDRESS));

+    if (Private->DnsServerIp == NULL) {

+      return EFI_OUT_OF_RESOURCES;

+    }

+

+    for (DnsServerIndex = 0; DnsServerIndex < Private->DnsServerCount; DnsServerIndex++) {

+      CopyMem (&(Private->DnsServerIp[DnsServerIndex].v4), &(((EFI_IPv4_ADDRESS *) Option->Data)[DnsServerIndex]), sizeof (EFI_IPv4_ADDRESS));

+    }

+    

+    //

+    // Configure the default DNS server if server assigned.

+    //    

     Status = HttpBootRegisterIp4Dns (

                Private,

                Option->Length,

                Option->Data

                );

     if (EFI_ERROR (Status)) {

+      FreePool (Private->DnsServerIp);

+      Private->DnsServerIp = NULL;

       return Status;

     }

   }

@@ -235,9 +293,13 @@ HttpBootDhcp4ExtractUriInfo (
   //

 

   //

-  // Update the device path to include the IP and boot URI information.

+  // Update the device path to include the boot resource information.

   //

   Status = HttpBootUpdateDevicePath (Private);

+  if (EFI_ERROR (Status) && Private->DnsServerIp != NULL) {

+    FreePool (Private->DnsServerIp);

+    Private->DnsServerIp = NULL;

+  }

 

   return Status;

 }

@@ -260,6 +322,7 @@ HttpBootDhcp6ExtractUriInfo (
   HTTP_BOOT_DHCP6_PACKET_CACHE    *HttpOffer;

   UINT32                          SelectIndex;

   UINT32                          ProxyIndex;

+  UINT32                          DnsServerIndex;

   EFI_DHCP6_PACKET_OPTION         *Option;

   EFI_IPv6_ADDRESS                IpAddr;

   CHAR8                           *HostName;

@@ -272,6 +335,8 @@ HttpBootDhcp6ExtractUriInfo (
   SelectIndex = Private->SelectIndex - 1;

   ASSERT (SelectIndex < HTTP_BOOT_OFFER_MAX_NUM);

 

+  DnsServerIndex = 0;

+

   Status   = EFI_SUCCESS;

   HostName = NULL;

   //

@@ -326,22 +391,37 @@ HttpBootDhcp6ExtractUriInfo (
   if (EFI_ERROR (Status)) {

     return Status;

   }

-  

-  //

-  // Configure the default DNS server if server assigned.

-  //

+

   if ((SelectOffer->OfferType == HttpOfferTypeDhcpNameUriDns) || 

       (SelectOffer->OfferType == HttpOfferTypeDhcpDns) ||

       (SelectOffer->OfferType == HttpOfferTypeDhcpIpUriDns)) {

     Option = SelectOffer->OptList[HTTP_BOOT_DHCP6_IDX_DNS_SERVER];

     ASSERT (Option != NULL);

+

+    //

+    // Record the Dns Server address list.

+    //

+    Private->DnsServerCount = HTONS (Option->OpLen) / sizeof (EFI_IPv6_ADDRESS);

+

+    Private->DnsServerIp = AllocateZeroPool (Private->DnsServerCount * sizeof (EFI_IP_ADDRESS));

+    if (Private->DnsServerIp == NULL) {

+      return EFI_OUT_OF_RESOURCES;

+    }

+

+    for (DnsServerIndex = 0; DnsServerIndex < Private->DnsServerCount; DnsServerIndex++) {

+      CopyMem (&(Private->DnsServerIp[DnsServerIndex].v6), &(((EFI_IPv6_ADDRESS *) Option->Data)[DnsServerIndex]), sizeof (EFI_IPv6_ADDRESS));

+    }

+

+    //

+    // Configure the default DNS server if server assigned.

+    //

     Status = HttpBootSetIp6Dns (

                Private,

                HTONS (Option->OpLen),

                Option->Data

                );

     if (EFI_ERROR (Status)) {

-      return Status;

+      goto Error;

     }

   }

   

@@ -365,7 +445,7 @@ HttpBootDhcp6ExtractUriInfo (
                &HostName

                );

     if (EFI_ERROR (Status)) {

-      return Status;

+      goto Error;

     }

 

     HostNameSize = AsciiStrSize (HostName);

@@ -376,6 +456,11 @@ HttpBootDhcp6ExtractUriInfo (
     }

     

     AsciiStrToUnicodeStrS (HostName, HostNameStr, HostNameSize);

+

+    if (HostName != NULL) {

+      FreePool (HostName);

+    }

+    

     Status = HttpBootDns (Private, HostNameStr, &IpAddr);

     FreePool (HostNameStr);

     if (EFI_ERROR (Status)) {

@@ -402,16 +487,21 @@ HttpBootDhcp6ExtractUriInfo (
   //

 

   //

-  // Update the device path to include the IP and boot URI information.

+  // Update the device path to include the boot resource information.

   //

   Status = HttpBootUpdateDevicePath (Private);

+  if (EFI_ERROR (Status)) {

+    goto Error;

+  }

+  

+  return Status;

 

 Error:

-  

-  if (HostName != NULL) {

-    FreePool (HostName);

+  if (Private->DnsServerIp != NULL) {

+    FreePool (Private->DnsServerIp);

+    Private->DnsServerIp = NULL;

   }

-    

+  

   return Status;

 }

 

diff --git a/NetworkPkg/HttpBootDxe/HttpBootDxe.h b/NetworkPkg/HttpBootDxe/HttpBootDxe.h
index 4632ee2..2c07cf7 100644
--- a/NetworkPkg/HttpBootDxe/HttpBootDxe.h
+++ b/NetworkPkg/HttpBootDxe/HttpBootDxe.h
@@ -202,6 +202,8 @@ struct _HTTP_BOOT_PRIVATE_DATA {
   EFI_IP_ADDRESS                            GatewayIp;

   EFI_IP_ADDRESS                            ServerIp;

   UINT16                                    Port;

+  UINT32                                    DnsServerCount;

+  EFI_IP_ADDRESS                            *DnsServerIp;

 

   //

   // The URI string attempt to download through HTTP, may point to

diff --git a/NetworkPkg/HttpBootDxe/HttpBootImpl.c b/NetworkPkg/HttpBootDxe/HttpBootImpl.c
index dfd0e23..06a8a6a 100644
--- a/NetworkPkg/HttpBootDxe/HttpBootImpl.c
+++ b/NetworkPkg/HttpBootDxe/HttpBootImpl.c
@@ -465,6 +465,11 @@ HttpBootStop (
     }

   }

 

+  if (Private->DnsServerIp != NULL) {

+    FreePool (Private->DnsServerIp);

+    Private->DnsServerIp = NULL;

+  }

+

   if (Private->FilePathUri!= NULL) {

     FreePool (Private->FilePathUri);

     HttpUrlFreeParser (Private->FilePathUriParser);