Re: [Patch 1/2] MdeModulePkg: Replace ASSERT with error return code in PXE driver.

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

Reviewed-by: Wu Jiaxin <jiaxin.wu@intel.com>

Thanks,
Jiaxin

> -----Original Message-----
> From: Fu, Siyuan
> Sent: Friday, December 16, 2016 4:19 PM
> To: edk2-devel@lists.01.org
> Cc: Yao, Jiewen <jiewen.yao@intel.com>; Ye, Ting <ting.ye@intel.com>; Wu,
> Jiaxin <jiaxin.wu@intel.com>
> Subject: [Patch 1/2] MdeModulePkg: Replace ASSERT with error return code
> in PXE driver.
> 
> This patch remove the ASSERT when receive a DHCP packet large than the
> maximum cache buffer size.
> 
> Cc: Yao Jiewen <jiewen.yao@intel.com>
> Cc: Ye Ting <ting.ye@intel.com>
> Cc: Wu Jiaxin <jiaxin.wu@intel.com>
> Contributed-under: TianoCore Contribution Agreement 1.0
> Signed-off-by: Fu Siyuan <siyuan.fu@intel.com>
> ---
>  .../Universal/Network/UefiPxeBcDxe/PxeBcDhcp.c     | 99
> ++++++++++++++--------
>  .../Universal/Network/UefiPxeBcDxe/PxeBcDhcp.h     |  5 +-
>  2 files changed, 68 insertions(+), 36 deletions(-)
> 
> diff --git a/MdeModulePkg/Universal/Network/UefiPxeBcDxe/PxeBcDhcp.c
> b/MdeModulePkg/Universal/Network/UefiPxeBcDxe/PxeBcDhcp.c
> index f03176b..f0720e5 100644
> --- a/MdeModulePkg/Universal/Network/UefiPxeBcDxe/PxeBcDhcp.c
> +++ b/MdeModulePkg/Universal/Network/UefiPxeBcDxe/PxeBcDhcp.c
> @@ -69,20 +69,26 @@ PxeBcInitSeedPacket (
>  /**
>    Copy the DCHP4 packet from srouce to destination.
> 
> -  @param  Dst   Pointer to the EFI_DHCP4_PROTOCOL instance.
> -  @param  Src   Pointer to the EFI_DHCP4_PROTOCOL instance.
> +  @param[in]  Dst          Pointer to the cache buffer for DHCPv4 packet.
> +  @param[in]  Src          Pointer to the DHCPv4 packet to be cached.
> +
> +  @retval     EFI_SUCCESS                Packet is copied.
> +  @retval     EFI_BUFFER_TOO_SMALL       Cache buffer is not big enough to
> hold the packet.
> 
>  **/
> -VOID
> +EFI_STATUS
>  PxeBcCopyEfiDhcp4Packet (
>    IN EFI_DHCP4_PACKET  *Dst,
>    IN EFI_DHCP4_PACKET  *Src
>    )
>  {
> -  ASSERT (Dst->Size >= Src->Length);
> +  if (Dst->Size < Src->Length) {
> +    return EFI_BUFFER_TOO_SMALL;
> +  }
> 
>    CopyMem (&Dst->Dhcp4, &Src->Dhcp4, Src->Length);
>    Dst->Length = Src->Length;
> +  return EFI_SUCCESS;
>  }
> 
> 
> @@ -93,8 +99,11 @@ PxeBcCopyEfiDhcp4Packet (
>    @param  OfferIndex    Index of cached packets as complements of pxe
> mode data,
>                          the index is maximum offer number.
> 
> +  @retval     EFI_SUCCESS                Cache and parse the packet successfully.
> +  @retval     EFI_BUFFER_TOO_SMALL       Cache buffer is not big enough to
> hold the packet.
> +
>  **/
> -VOID
> +EFI_STATUS
>  PxeBcCopyProxyOffer (
>    IN PXEBC_PRIVATE_DATA  *Private,
>    IN UINT32              OfferIndex
> @@ -102,6 +111,7 @@ PxeBcCopyProxyOffer (  {
>    EFI_PXE_BASE_CODE_MODE  *Mode;
>    EFI_DHCP4_PACKET        *Offer;
> +  EFI_STATUS              Status;
> 
>    ASSERT (OfferIndex < Private->NumOffers);
>    ASSERT (OfferIndex < PXEBC_MAX_OFFER_NUM); @@ -109,11 +119,15
> @@ PxeBcCopyProxyOffer (
>    Mode  = Private->PxeBc.Mode;
>    Offer = &Private->Dhcp4Offers[OfferIndex].Packet.Offer;
> 
> -  PxeBcCopyEfiDhcp4Packet (&Private->ProxyOffer.Packet.Offer, Offer);
> +  Status = PxeBcCopyEfiDhcp4Packet (&Private->ProxyOffer.Packet.Offer,
> + Offer);  if (EFI_ERROR(Status)) {
> +    return Status;
> +  }
>    CopyMem (&Mode->ProxyOffer, &Offer->Dhcp4, Offer->Length);
>    Mode->ProxyOfferReceived = TRUE;
> 
>    PxeBcParseCachedDhcpPacket (&Private->ProxyOffer);
> +  return EFI_SUCCESS;
>  }
> 
> 
> @@ -411,8 +425,9 @@ PxeBcTryBinlProxy (
> 
>    @param  Private          Pointer to PxeBc private data.
> 
> -  @retval EFI_SUCCESS      Operational successful.
> -  @retval EFI_NO_RESPONSE  Offer dhcp service failed.
> +  @retval EFI_SUCCESS                Operational successful.
> +  @retval EFI_NO_RESPONSE            Offer dhcp service failed.
> +  @retval EFI_BUFFER_TOO_SMALL       Failed to copy the packet to Pxe base
> code mode.
> 
>  **/
>  EFI_STATUS
> @@ -515,7 +530,7 @@ PxeBcCheckSelectedOffer (
>          //
>          // Copy the proxy offer to Mode and set the flag
>          //
> -        PxeBcCopyProxyOffer (Private, ProxyOfferIndex);
> +        Status = PxeBcCopyProxyOffer (Private, ProxyOfferIndex);
>        }
>      } else {
>        //
> @@ -543,7 +558,7 @@ PxeBcCheckSelectedOffer (
>        // Other type of ACK is already cached. Bootp is special that we should
>        // use the bootp reply as the ACK and put it into the DHCP_ONLY buffer.
>        //
> -      PxeBcCopyEfiDhcp4Packet (&Private->Dhcp4Ack.Packet.Ack, Offer);
> +      Status = PxeBcCopyEfiDhcp4Packet (&Private->Dhcp4Ack.Packet.Ack,
> + Offer);
>      }
> 
>      PxeBcParseCachedDhcpPacket (&Private->Dhcp4Ack); @@ -566,8 +581,11
> @@ PxeBcCheckSelectedOffer (
>    @param  Private    Pointer to PxeBc private data.
>    @param  RcvdOffer  Pointer to the received Dhcp proxy offer packet.
> 
> +  @retval     EFI_SUCCESS      Cache and parse the packet successfully.
> +  @retval     Others           Operation failed.
> +
>  **/
> -VOID
> +EFI_STATUS
>  PxeBcCacheDhcpOffer (
>    IN PXEBC_PRIVATE_DATA  *Private,
>    IN EFI_DHCP4_PACKET    *RcvdOffer
> @@ -576,6 +594,7 @@ PxeBcCacheDhcpOffer (
>    PXEBC_CACHED_DHCP4_PACKET *CachedOffer;
>    EFI_DHCP4_PACKET          *Offer;
>    UINT8                     OfferType;
> +  EFI_STATUS                Status;
> 
>    CachedOffer = &Private->Dhcp4Offers[Private->NumOffers];
>    Offer       = &CachedOffer->Packet.Offer;
> @@ -583,18 +602,21 @@ PxeBcCacheDhcpOffer (
>    //
>    // Cache the orignal dhcp packet
>    //
> -  PxeBcCopyEfiDhcp4Packet (Offer, RcvdOffer);
> +  Status = PxeBcCopyEfiDhcp4Packet (Offer, RcvdOffer);  if
> + (EFI_ERROR(Status)) {
> +    return Status;
> +  }
> 
>    //
>    // Parse and validate the options (including dhcp option and vendor option)
>    //
>    if (!PxeBcParseCachedDhcpPacket (CachedOffer)) {
> -    return ;
> +    return EFI_ABORTED;
>    }
> 
>    OfferType = CachedOffer->OfferType;
>    if (OfferType >= DHCP4_PACKET_TYPE_MAX) {
> -    return ;
> +    return EFI_ABORTED;
>    }
> 
>    if (OfferType == DHCP4_PACKET_TYPE_BOOTP) { @@ -603,7 +625,7 @@
> PxeBcCacheDhcpOffer (
>        //
>        // Only cache the first bootp offer, discard others.
>        //
> -      return ;
> +      return EFI_ABORTED;
>      } else {
>        //
>        // Take as a dhcp only offer, but record index specifically.
> @@ -628,7 +650,7 @@ PxeBcCacheDhcpOffer (
>          //
>          // Only cache the first pxe10/wfm11a offers each, discard the others.
>          //
> -        return ;
> +        return EFI_ABORTED;
>        } else {
>          //
>          // Record index of the proxy dhcp offer with type other than binl.
> @@ -649,6 +671,8 @@ PxeBcCacheDhcpOffer (
>    // Count the accepted offers.
>    //
>    Private->NumOffers++;
> +
> +  return EFI_SUCCESS;
>  }
> 
>  /**
> @@ -960,6 +984,7 @@ PxeBcDhcpCallBack (
>      if (Private->NumOffers < PXEBC_MAX_OFFER_NUM) {
>        //
>        // Cache the dhcp offers in Private->Dhcp4Offers[]
> +      // If error happens, just ignore this packet and continue to wait more
> offer.
>        //
>        PxeBcCacheDhcpOffer (Private, Packet);
>      }
> @@ -982,20 +1007,15 @@ PxeBcDhcpCallBack (
>      break;
> 
>    case Dhcp4RcvdAck:
> -    if (Packet->Length > PXEBC_DHCP4_MAX_PACKET_SIZE) {
> -      //
> -      // Abort the DHCP if the ACK packet exceeds the maximum length.
> -      //
> -	  Status = EFI_ABORTED;
> -      break;
> -    }
> -
>      //
>      // Cache Ack
>      //
>      ASSERT (Private->SelectedOffer != 0);
> 
> -    PxeBcCopyEfiDhcp4Packet (&Private->Dhcp4Ack.Packet.Ack, Packet);
> +    Status = PxeBcCopyEfiDhcp4Packet (&Private->Dhcp4Ack.Packet.Ack,
> Packet);
> +    if (EFI_ERROR (Status)) {
> +      return EFI_ABORTED;
> +    }
>      break;
> 
>    default:
> @@ -1340,6 +1360,12 @@ PxeBcDiscvBootService (
>      Response  = Token.ResponseList;
> 
>      while (RepIndex < Token.ResponseCount) {
> +      if (Response->Length > PXEBC_DHCP4_MAX_PACKET_SIZE) {
> +        SrvIndex = 0;
> +        RepIndex++;
> +        Response = (EFI_DHCP4_PACKET *) ((UINT8 *) Response + Response-
> >Size);
> +        continue;
> +      }
> 
>        while (SrvIndex < IpCount) {
> 
> @@ -1360,14 +1386,16 @@ PxeBcDiscvBootService (
> 
>        SrvIndex = 0;
>        RepIndex++;
> -
>        Response = (EFI_DHCP4_PACKET *) ((UINT8 *) Response + Response-
> >Size);
>      }
> 
>      if (RepIndex < Token.ResponseCount) {
> 
>        if (Reply != NULL) {
> -        PxeBcCopyEfiDhcp4Packet (Reply, Response);
> +        Status = PxeBcCopyEfiDhcp4Packet (Reply, Response);
> +        if (EFI_ERROR(Status)) {
> +          goto ON_EXIT;
> +        }
>        }
> 
>        if (IsDiscv) {
> @@ -1380,18 +1408,21 @@ PxeBcDiscvBootService (
>      } else {
>        Status = EFI_NOT_FOUND;
>      }
> +  }
> 
> -    //
> -    // free the responselist
> -    //
> -    if (Token.ResponseList != NULL) {
> -      FreePool (Token.ResponseList);
> -    }
> +ON_EXIT:
> +  //
> +  // free the responselist
> +  //
> +  if (Token.ResponseList != NULL) {
> +    FreePool (Token.ResponseList);
>    }
>    //
>    // Free the dhcp packet
>    //
> -  FreePool (Token.Packet);
> +  if (Token.Packet != NULL) {
> +    FreePool (Token.Packet);
> +  }
> 
>    return Status;
>  }
> diff --git a/MdeModulePkg/Universal/Network/UefiPxeBcDxe/PxeBcDhcp.h
> b/MdeModulePkg/Universal/Network/UefiPxeBcDxe/PxeBcDhcp.h
> index 614ea75..d19d2a3 100644
> --- a/MdeModulePkg/Universal/Network/UefiPxeBcDxe/PxeBcDhcp.h
> +++ b/MdeModulePkg/Universal/Network/UefiPxeBcDxe/PxeBcDhcp.h
> @@ -291,8 +291,9 @@ PxeBcParseCachedDhcpPacket (
> 
>    @param  Private          Pointer to PxeBc private data.
> 
> -  @retval EFI_SUCCESS      Operational successful.
> -  @retval EFI_NO_RESPONSE  Offer dhcp service failed.
> +  @retval EFI_SUCCESS                Operational successful.
> +  @retval EFI_NO_RESPONSE            Offer dhcp service failed.
> +  @retval EFI_BUFFER_TOO_SMALL       Failed to copy the packet to Pxe base
> code mode.
> 
>  **/
>  EFI_STATUS
> --
> 2.7.4.windows.1