MdeModulePkg/Network: Fix potential ASSERT if NetIp4IsUnicast is called

[mirror_edk2.git] / MdeModulePkg / Universal / Network / Mtftp4Dxe / Mtftp4Support.c

diff --git a/MdeModulePkg/Universal/Network/Mtftp4Dxe/Mtftp4Support.c b/MdeModulePkg/Universal/Network/Mtftp4Dxe/Mtftp4Support.c
index 26860e52bb4b9fc19013d7cc4bf6d750efbb62c7..c625fda020e81034aca5fdf4baf07616bc3eb03e 100644 (file)
--- a/MdeModulePkg/Universal/Network/Mtftp4Dxe/Mtftp4Support.c
+++ b/MdeModulePkg/Universal/Network/Mtftp4Dxe/Mtftp4Support.c
@@ -1,7 +1,7 @@
 /** @file\r
   Support routines for Mtftp.\r
   \r
-Copyright (c) 2006 - 2010, Intel Corporation. All rights reserved.<BR>\r
+Copyright (c) 2006 - 2015, Intel Corporation. All rights reserved.<BR>\r
 This program and the accompanying materials\r
 are licensed and made available under the terms and conditions of the BSD License\r
 which accompanies this distribution.  The full text of the license may be found at\r
@@ -278,13 +278,16 @@ Mtftp4SendRequest (
   EFI_MTFTP4_PACKET         *Packet;\r
   EFI_MTFTP4_OPTION         *Options;\r
   EFI_MTFTP4_TOKEN          *Token;\r
+  RETURN_STATUS             Status;\r
   NET_BUF                   *Nbuf;\r
   UINT8                     *Mode;\r
   UINT8                     *Cur;\r
-  UINT32                    Len;\r
   UINTN                     Index;\r
-  UINT32                    Len1;\r
-  UINT32                    Len2;\r
+  UINT32                    BufferLength;\r
+  UINTN                     FileNameLength;\r
+  UINTN                     ModeLength;\r
+  UINTN                     OptionStrLength;\r
+  UINTN                     ValueStrLength;\r
 \r
   Token   = Instance->Token;\r
   Options = Token->OptionList;\r
@@ -297,39 +300,52 @@ Mtftp4SendRequest (
   //\r
   // Compute the packet length\r
   //\r
-  Len1 = (UINT32) AsciiStrLen ((CHAR8 *) Token->Filename);\r
-  Len2 = (UINT32) AsciiStrLen ((CHAR8 *) Mode);\r
-  Len  = (Len1 + Len2 + 4);\r
+  FileNameLength = AsciiStrLen ((CHAR8 *) Token->Filename);\r
+  ModeLength     = AsciiStrLen ((CHAR8 *) Mode);\r
+  BufferLength   = (UINT32) FileNameLength + (UINT32) ModeLength + 4;\r
 \r
   for (Index = 0; Index < Token->OptionCount; Index++) {\r
-    Len1 = (UINT32) AsciiStrLen ((CHAR8 *) Options[Index].OptionStr);\r
-    Len2 = (UINT32) AsciiStrLen ((CHAR8 *) Options[Index].ValueStr);\r
-    Len += Len1 + Len2 + 2;\r
+    OptionStrLength = AsciiStrLen ((CHAR8 *) Options[Index].OptionStr);\r
+    ValueStrLength  = AsciiStrLen ((CHAR8 *) Options[Index].ValueStr);\r
+    BufferLength   += (UINT32) OptionStrLength + (UINT32) ValueStrLength + 2;\r
   }\r
-\r
   //\r
   // Allocate a packet then copy the data over\r
   //\r
-  if ((Nbuf = NetbufAlloc (Len)) == NULL) {\r
+  if ((Nbuf = NetbufAlloc (BufferLength)) == NULL) {\r
     return EFI_OUT_OF_RESOURCES;\r
   }\r
 \r
-  Packet         = (EFI_MTFTP4_PACKET *) NetbufAllocSpace (Nbuf, Len, FALSE);\r
+  Packet         = (EFI_MTFTP4_PACKET *) NetbufAllocSpace (Nbuf, BufferLength, FALSE);\r
   ASSERT (Packet != NULL);\r
 \r
   Packet->OpCode = HTONS (Instance->Operation);\r
+  BufferLength  -= sizeof (Packet->OpCode);\r
+  \r
   Cur            = Packet->Rrq.Filename;\r
-  Cur            = (UINT8 *) AsciiStrCpy ((CHAR8 *) Cur, (CHAR8 *) Token->Filename);\r
-  Cur           += AsciiStrLen ((CHAR8 *) Token->Filename) + 1;\r
-  Cur            = (UINT8 *) AsciiStrCpy ((CHAR8 *) Cur, (CHAR8 *) Mode);\r
-  Cur           += AsciiStrLen ((CHAR8 *) Mode) + 1;\r
+  Status         = AsciiStrCpyS ((CHAR8 *) Cur, BufferLength, (CHAR8 *) Token->Filename);\r
+  ASSERT_EFI_ERROR (Status);\r
+  BufferLength  -= (UINT32) (FileNameLength + 1);\r
+  Cur           += FileNameLength + 1;\r
+  Status         = AsciiStrCpyS ((CHAR8 *) Cur, BufferLength, (CHAR8 *) Mode);\r
+  ASSERT_EFI_ERROR (Status);\r
+  BufferLength  -= (UINT32) (ModeLength + 1);\r
+  Cur           += ModeLength + 1;\r
 \r
   for (Index = 0; Index < Token->OptionCount; ++Index) {\r
-    Cur  = (UINT8 *) AsciiStrCpy ((CHAR8 *) Cur, (CHAR8 *) Options[Index].OptionStr);\r
-    Cur += AsciiStrLen ((CHAR8 *) Options[Index].OptionStr) + 1;\r
-\r
-    Cur  = (UINT8 *) AsciiStrCpy ((CHAR8 *) Cur, (CHAR8 *) Options[Index].ValueStr);\r
-    Cur += AsciiStrLen ((CHAR8 *) (CHAR8 *) Options[Index].ValueStr) + 1;\r
+    OptionStrLength = AsciiStrLen ((CHAR8 *) Options[Index].OptionStr);\r
+    ValueStrLength  = AsciiStrLen ((CHAR8 *) Options[Index].ValueStr);\r
+    \r
+    Status          = AsciiStrCpyS ((CHAR8 *) Cur, BufferLength, (CHAR8 *) Options[Index].OptionStr);\r
+    ASSERT_EFI_ERROR (Status);\r
+    BufferLength   -= (UINT32) (OptionStrLength + 1);\r
+    Cur            += OptionStrLength + 1;\r
+    \r
+    Status          = AsciiStrCpyS ((CHAR8 *) Cur, BufferLength, (CHAR8 *) Options[Index].ValueStr);\r
+    ASSERT_EFI_ERROR (Status);\r
+    BufferLength   -= (UINT32) (ValueStrLength + 1);\r
+    Cur            += ValueStrLength + 1;\r
+    \r
   }\r
 \r
   return Mtftp4SendPacket (Instance, Nbuf);\r
@@ -371,7 +387,7 @@ Mtftp4SendError (
   TftpError->OpCode = HTONS (EFI_MTFTP4_OPCODE_ERROR);\r
   TftpError->Error.ErrorCode = HTONS (ErrCode);\r
 \r
-  AsciiStrCpy ((CHAR8 *) TftpError->Error.ErrorMessage, (CHAR8 *) ErrInfo);\r
+  AsciiStrCpyS ((CHAR8 *) TftpError->Error.ErrorMessage, Len, (CHAR8 *) ErrInfo);\r
 \r
   return Mtftp4SendPacket (Instance, Packet);\r
 }\r
@@ -443,7 +459,7 @@ Mtftp4SendPacket (
   UDP_END_POINT             UdpPoint;\r
   EFI_STATUS                Status;\r
   UINT16                    OpCode;\r
-  UINT16                    Value;\r
+  UINT8                     *Buffer;\r
 \r
   //\r
   // Save the packet for retransmission\r
@@ -464,8 +480,9 @@ Mtftp4SendPacket (
   // Send the requests to the listening port, other packets\r
   // to the connected port\r
   //\r
-  Value = *((UINT16 *) NetbufGetByte (Packet, 0, NULL));\r
-  OpCode = NTOHS (Value);\r
+  Buffer = NetbufGetByte (Packet, 0, NULL);\r
+  ASSERT (Buffer != NULL);\r
+  OpCode = NTOHS (*(UINT16 *)Buffer);\r
 \r
   if ((OpCode == EFI_MTFTP4_OPCODE_RRQ) || \r
       (OpCode == EFI_MTFTP4_OPCODE_DIR) ||\r
@@ -511,7 +528,7 @@ Mtftp4Retransmit (
   UDP_END_POINT             UdpPoint;\r
   EFI_STATUS                Status;\r
   UINT16                    OpCode;\r
-  UINT16                    Value;\r
+  UINT8                     *Buffer;\r
 \r
   ASSERT (Instance->LastPacket != NULL);\r
 \r
@@ -521,8 +538,9 @@ Mtftp4Retransmit (
   //\r
   // Set the requests to the listening port, other packets to the connected port\r
   //\r
-  Value = *(UINT16 *) NetbufGetByte (Instance->LastPacket, 0, NULL);\r
-  OpCode = NTOHS (Value);\r
+  Buffer = NetbufGetByte (Instance->LastPacket, 0, NULL);\r
+  ASSERT (Buffer != NULL);\r
+  OpCode = NTOHS (*(UINT16 *) Buffer);\r
 \r
   if ((OpCode == EFI_MTFTP4_OPCODE_RRQ) || (OpCode == EFI_MTFTP4_OPCODE_DIR) ||\r
       (OpCode == EFI_MTFTP4_OPCODE_WRQ)) {\r