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

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;

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);
+  PduHdr              = NetbufGetByte (Pdu, 0, NULL);
-  XferContext->Offset = ISCSI_GET_DATASEG_LEN (Buffer);
+  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;