Refine code to make it more safely.

Signed-off-by: Dong Eric <eric.dong@intel.com>
Reviewed-by: Ye Ting <ting.ye@intel.com>

git-svn-id: https://edk2.svn.sourceforge.net/svnroot/edk2/trunk/edk2@13931 6f19259b-4bc3-4df7-8a09-765794883524

2 files changed, 48 insertions, 5 deletions

diff --git a/MdeModulePkg/Universal/Network/IScsiDxe/IScsiCHAP.c b/MdeModulePkg/Universal/Network/IScsiDxe/IScsiCHAP.c
index be5ad88f15..7775ed993c 100644
--- a/MdeModulePkg/Universal/Network/IScsiDxe/IScsiCHAP.c
+++ b/MdeModulePkg/Universal/Network/IScsiDxe/IScsiCHAP.c
@@ -332,6 +332,9 @@ IScsiCHAPToSendReq (
   Session     = Conn->Session;

   AuthData    = &Session->AuthData;

   LoginReq    = (ISCSI_LOGIN_REQUEST *) NetbufGetByte (Pdu, 0, 0);

+  if (LoginReq == NULL) {

+    return EFI_PROTOCOL_ERROR;

+  }

   Status      = EFI_SUCCESS;

 

   RspLen      = 2 * ISCSI_CHAP_RSP_LEN + 3;

diff --git a/MdeModulePkg/Universal/Network/IScsiDxe/IScsiProto.c b/MdeModulePkg/Universal/Network/IScsiDxe/IScsiProto.c
index 0853f98767..b1f250b55e 100644
--- a/MdeModulePkg/Universal/Network/IScsiDxe/IScsiProto.c
+++ b/MdeModulePkg/Universal/Network/IScsiDxe/IScsiProto.c
@@ -441,6 +441,9 @@ IScsiAddKeyValuePair (
   CHAR8               *Data;

 

   LoginReq    = (ISCSI_LOGIN_REQUEST *) NetbufGetByte (Pdu, 0, NULL);

+  if (LoginReq == NULL) {

+    return EFI_PROTOCOL_ERROR;

+  }

   DataSegLen  = NTOH24 (LoginReq->DataSegmentLength);

 

   KeyLen      = (UINT32) AsciiStrLen (Key);

@@ -607,6 +610,9 @@ IScsiProcessLoginRsp (
   Session   = Conn->Session;

 

   LoginRsp  = (ISCSI_LOGIN_RESPONSE *) NetbufGetByte (Pdu, 0, NULL);

+  if (LoginRsp == NULL) {

+    return EFI_PROTOCOL_ERROR;

+  }

   if (!ISCSI_CHECK_OPCODE (LoginRsp, ISCSI_OPCODE_LOGIN_RSP)) {

     //

     // It's not a Login Response

@@ -2070,6 +2076,7 @@ IScsiGenerateDataOutPduSequence (
   NET_BUF             *DataOutPdu;

   ISCSI_CONNECTION    *Conn;

   ISCSI_XFER_CONTEXT  *XferContext;

+  UINT8               *DataOutPacket;

 

   PduList = AllocatePool (sizeof (LIST_ENTRY));

   if (PduList == NULL) {

@@ -2113,7 +2120,14 @@ IScsiGenerateDataOutPduSequence (
   //

   // Set the F bit for the last data out PDU in this sequence.

   //

-  ISCSI_SET_FLAG (NetbufGetByte (DataOutPdu, 0, NULL), ISCSI_BHS_FLAG_FINAL);

+  DataOutPacket = NetbufGetByte (DataOutPdu, 0, NULL);

+  if (DataOutPacket == NULL) {

+    IScsiFreeNbufList (PduList);

+    PduList = NULL;

+    goto ON_EXIT;

+  }

+

+  ISCSI_SET_FLAG (DataOutPacket, ISCSI_BHS_FLAG_FINAL);

 

 ON_EXIT:

 

@@ -2194,6 +2208,9 @@ IScsiOnDataInRcvd (
   EFI_STATUS          Status;

 

   DataInHdr                   = (ISCSI_SCSI_DATA_IN *) NetbufGetByte (Pdu, 0, NULL);

+  if (DataInHdr == NULL) {

+    return EFI_PROTOCOL_ERROR;

+  }

 

   DataInHdr->InitiatorTaskTag = NTOHL (DataInHdr->InitiatorTaskTag);

   DataInHdr->ExpCmdSN         = NTOHL (DataInHdr->ExpCmdSN);

@@ -2282,6 +2299,9 @@ IScsiOnR2TRcvd (
   UINT8                   *Data;

 

   R2THdr = (ISCSI_READY_TO_TRANSFER *) NetbufGetByte (Pdu, 0, NULL);

+  if (R2THdr == NULL) {

+    return EFI_PROTOCOL_ERROR;

+  }

 

   R2THdr->InitiatorTaskTag = NTOHL (R2THdr->InitiatorTaskTag);

   R2THdr->TargetTransferTag = NTOHL (R2THdr->TargetTransferTag);

@@ -2345,6 +2365,9 @@ IScsiOnScsiRspRcvd (
   UINT32            DataSegLen;

 

   ScsiRspHdr                    = (SCSI_RESPONSE *) NetbufGetByte (Pdu, 0, NULL);

+  if (ScsiRspHdr == NULL) {

+    return EFI_PROTOCOL_ERROR;

+  }

 

   ScsiRspHdr->InitiatorTaskTag  = NTOHL (ScsiRspHdr->InitiatorTaskTag);

   if (ScsiRspHdr->InitiatorTaskTag != Tcb->InitiatorTaskTag) {

@@ -2407,6 +2430,9 @@ IScsiOnScsiRspRcvd (
   DataSegLen = ISCSI_GET_DATASEG_LEN (ScsiRspHdr);

   if (DataSegLen != 0) {

     SenseData               = (ISCSI_SENSE_DATA *) NetbufGetByte (Pdu, sizeof (SCSI_RESPONSE), NULL);

+    if (SenseData == NULL) {

+      return EFI_PROTOCOL_ERROR;

+    }

 

     SenseData->Length       = NTOHS (SenseData->Length);

 

@@ -2441,6 +2467,9 @@ IScsiOnNopInRcvd (
   EFI_STATUS    Status;

 

   NopInHdr            = (ISCSI_NOP_IN *) NetbufGetByte (Pdu, 0, NULL);

+  if (NopInHdr == NULL) {

+    return EFI_PROTOCOL_ERROR;

+  }

 

   NopInHdr->StatSN    = NTOHL (NopInHdr->StatSN);

   NopInHdr->ExpCmdSN  = NTOHL (NopInHdr->ExpCmdSN);

@@ -2496,7 +2525,7 @@ IScsiExecuteScsiCommand (
   UINT8                   *Data;

   ISCSI_IN_BUFFER_CONTEXT InBufferContext;

   UINT64                  Timeout;

-  UINT8                   *Buffer;

+  UINT8                   *PduHdr;

 

   Private       = ISCSI_DRIVER_DATA_FROM_EXT_SCSI_PASS_THRU (PassThru);

   Session       = &Private->Session;

@@ -2534,8 +2563,13 @@ IScsiExecuteScsiCommand (
   }

 

   XferContext         = &Tcb->XferContext;

-  Buffer              = NetbufGetByte (Pdu, 0, NULL);

-  XferContext->Offset = ISCSI_GET_DATASEG_LEN (Buffer);

+  PduHdr              = NetbufGetByte (Pdu, 0, NULL);

+  if (PduHdr == NULL) {

+    Status = EFI_PROTOCOL_ERROR;

+    NetbufFree (Pdu);

+    goto ON_EXIT;

+  }

+  XferContext->Offset = ISCSI_GET_DATASEG_LEN (PduHdr);

 

   //

   // Transmit the SCSI Command PDU.

@@ -2591,7 +2625,13 @@ IScsiExecuteScsiCommand (
       goto ON_EXIT;

     }

 

-    switch (ISCSI_GET_OPCODE (NetbufGetByte (Pdu, 0, NULL))) {

+    PduHdr = NetbufGetByte (Pdu, 0, NULL);

+    if (PduHdr == NULL) {

+      Status = EFI_PROTOCOL_ERROR;

+      NetbufFree (Pdu);

+      goto ON_EXIT;

+    }

+    switch (ISCSI_GET_OPCODE (PduHdr)) {

     case ISCSI_OPCODE_SCSI_DATA_IN:

       Status = IScsiOnDataInRcvd (Pdu, Tcb, Packet);

       break;