[core] Added null-check for buffer extraction procedure (#1391)

Haivision · Jul 24, 2020 · 89902fb · 89902fb
1 parent ee870de
commit 89902fb
Showing 1 changed file with 38 additions and 13 deletions.
diff --git a/srtcore/buffer.cpp b/srtcore/buffer.cpp
@@ -898,28 +898,32 @@ int CRcvBuffer::readBuffer(char* data, int len)
             return -1;
         }
 
+        const CPacket& pkt = m_pUnit[p]->m_Packet;
+
         if (m_bTsbPdMode)
         {
             HLOGC(dlog.Debug,
                   log << CONID() << "readBuffer: chk if time2play:"
                       << " NOW=" << FormatTime(now)
-                      << " PKT TS=" << FormatTime(getPktTsbPdTime(m_pUnit[p]->m_Packet.getMsgTimeStamp())));
+                      << " PKT TS=" << FormatTime(getPktTsbPdTime(pkt.getMsgTimeStamp())));
 
-            if ((getPktTsbPdTime(m_pUnit[p]->m_Packet.getMsgTimeStamp()) > now))
+            if ((getPktTsbPdTime(pkt.getMsgTimeStamp()) > now))
                 break; /* too early for this unit, return whatever was copied */
         }
 
-        int unitsize = (int)m_pUnit[p]->m_Packet.getLength() - m_iNotch;
-        if (unitsize > rs)
-            unitsize = rs;
+        const int pktlen = pkt.getLength();
+        const int remain_pktlen = pktlen - m_iNotch;
+
+        const int unitsize = std::min(remain_pktlen, rs);
 
         HLOGC(dlog.Debug,
               log << CONID() << "readBuffer: copying buffer #" << p << " targetpos=" << int(data - begin)
                   << " sourcepos=" << m_iNotch << " size=" << unitsize << " left=" << (unitsize - rs));
-        memcpy((data), m_pUnit[p]->m_Packet.m_pcData + m_iNotch, unitsize);
+        memcpy((data), pkt.m_pcData + m_iNotch, unitsize);
+
         data += unitsize;
 
-        if ((rs > unitsize) || (rs == int(m_pUnit[p]->m_Packet.getLength()) - m_iNotch))
+        if (rs >= remain_pktlen)
         {
             freeUnitAt(p);
             p = shiftFwd(p);
@@ -945,20 +949,41 @@ int CRcvBuffer::readBufferToFile(fstream& ofs, int len)
     int lastack = m_iLastAckPos;
     int rs      = len;
 
+    int32_t trace_seq ATR_UNUSED = SRT_SEQNO_NONE;
+    int trace_shift ATR_UNUSED = -1;
+
     while ((p != lastack) && (rs > 0))
     {
-        int unitsize = (int)m_pUnit[p]->m_Packet.getLength() - m_iNotch;
-        if (unitsize > rs)
-            unitsize = rs;
+#if ENABLE_LOGGING
+        ++trace_shift;
+#endif
+        // Skip empty units. Note that this shouldn't happen
+        // in case of a file transfer.
+        if (!m_pUnit[p])
+        {
+            p = shiftFwd(p);
+            LOGC(mglog.Error, log << "readBufferToFile: IPE: NULL unit found in file transmission, last good %"
+                    << trace_seq << " + " << trace_shift);
+            continue;
+        }
 
-        ofs.write(m_pUnit[p]->m_Packet.m_pcData + m_iNotch, unitsize);
+        const CPacket& pkt = m_pUnit[p]->m_Packet;
+
+#if ENABLE_LOGGING
+        trace_seq = pkt.getSeqNo();
+#endif
+        const int pktlen = pkt.getLength();
+        const int remain_pktlen = pktlen - m_iNotch;
+
+        const int unitsize = std::min(remain_pktlen, rs);
+
+        ofs.write(pkt.m_pcData + m_iNotch, unitsize);
         if (ofs.fail())
             break;
 
-        if ((rs > unitsize) || (rs == int(m_pUnit[p]->m_Packet.getLength()) - m_iNotch))
+        if (rs >= remain_pktlen)
         {
             freeUnitAt(p);
-
             p = shiftFwd(p);
 
             m_iNotch = 0;