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>

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,19 +74,49 @@ 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;
@@ -92,9 +124,15 @@ HttpBootUpdateDevicePath (
   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;
   //
@@ -327,21 +392,36 @@ HttpBootDhcp6ExtractUriInfo (
     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,14 +487,19 @@ 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 (Private->DnsServerIp != NULL) {
-  if (HostName != NULL) {
+    FreePool (Private->DnsServerIp);
-    FreePool (HostName);
+    Private->DnsServerIp = NULL;
   }
   
   return Status;

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

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