Fix STAP-A bug where we might overflow the packet buffer due to not accounting for the length of the length field.

BUG=3679
R=mflodman@webrtc.org

Review URL: https://webrtc-codereview.appspot.com/17079004

git-svn-id: http://webrtc.googlecode.com/svn/trunk@6881 4adac7df-926f-26a2-2b94-8c16560cd09d

webrtc/modules/rtp_rtcp/source/rtp_format_h264.cc

@@ -29,6 +29,7 @@ enum Nalu {
 
 static const size_t kNalHeaderSize = 1;
 static const size_t kFuAHeaderSize = 2;
+static const size_t kLengthFieldSize = 2;
 
 // Bit masks for FU (A and B) indicators.
 enum NalDefs { kFBit = 0x80, kNriMask = 0x60, kTypeMask = 0x1F };
@@ -157,30 +158,34 @@ int RtpPacketizerH264::PacketizeStapA(size_t fragment_index,
   // Aggregate fragments into one packet (STAP-A).
   size_t payload_size_left = max_payload_len_;
   int aggregated_fragments = 0;
+  size_t fragment_headers_length = 0;
   assert(payload_size_left >= fragment_length);
-  while (payload_size_left >= fragment_length) {
+  while (payload_size_left >= fragment_length + fragment_headers_length) {
-    if (fragment_length > 0) {
+    assert(fragment_length > 0);
-      assert(fragment_length > 0);
+    uint8_t header = payload_data_[fragment_offset];
-      uint8_t header = payload_data_[fragment_offset];
+    packets_.push(Packet(fragment_offset,
-      packets_.push(Packet(fragment_offset,
+                         fragment_length,
-                           fragment_length,
+                         aggregated_fragments == 0,
-                           aggregated_fragments == 0,
+                         false,
-                           false,
+                         true,
-                           true,
+                         header));
-                           header));
+    payload_size_left -= fragment_length;
-      payload_size_left -= fragment_length;
+    payload_size_left -= fragment_headers_length;
-      // If we are going to try to aggregate more fragments into this packet
+
-      // we need to add the STAP-A NALU header.
-      if (aggregated_fragments == 0 && payload_size_left > 0)
-        payload_size_left -= kNalHeaderSize;
-      ++aggregated_fragments;
-    }
     // Next fragment.
     ++fragment_index;
     if (fragment_index == fragmentation_.fragmentationVectorSize)
       break;
     fragment_offset = fragmentation_.fragmentationOffset[fragment_index];
     fragment_length = fragmentation_.fragmentationLength[fragment_index];
+
+    fragment_headers_length = kLengthFieldSize;
+    // If we are going to try to aggregate more fragments into this packet
+    // we need to add the STAP-A NALU header and a length field for the first
+    // NALU of this packet.
+    if (aggregated_fragments == 0)
+      fragment_headers_length += kNalHeaderSize + kLengthFieldSize;
+    ++aggregated_fragments;
   }
   packets_.back().last_fragment = true;
   return fragment_index;
@@ -203,13 +208,15 @@ bool RtpPacketizerH264::NextPacket(uint8_t* buffer,
     *bytes_to_send = packet.size;
     memcpy(buffer, &payload_data_[packet.offset], packet.size);
     packets_.pop();
+    assert(*bytes_to_send <= max_payload_len_);
   } else if (packet.aggregated) {
     NextAggregatePacket(buffer, bytes_to_send);
+    assert(*bytes_to_send <= max_payload_len_);
   } else {
     NextFragmentPacket(buffer, bytes_to_send);
+    assert(*bytes_to_send <= max_payload_len_);
   }
   *last_packet = packets_.empty();
-  assert(*bytes_to_send <= max_payload_len_);
   return true;
 }
 
@@ -224,8 +231,8 @@ void RtpPacketizerH264::NextAggregatePacket(uint8_t* buffer,
   while (packet.aggregated) {
     // Add NAL unit length field.
     RtpUtility::AssignUWord16ToBuffer(&buffer[index], packet.size);
-    index += 2;
+    index += kLengthFieldSize;
-    *bytes_to_send += 2;
+    *bytes_to_send += kLengthFieldSize;
     // Add NAL unit.
     memcpy(&buffer[index], &payload_data_[packet.offset], packet.size);
     index += packet.size;

webrtc/modules/rtp_rtcp/source/rtp_format_h264_unittest.cc

@@ -211,7 +211,8 @@ TEST(RtpPacketizerH264Test, TestSingleNaluTwoPackets) {
 }
 
 TEST(RtpPacketizerH264Test, TestStapA) {
-  const size_t kFrameSize = kMaxPayloadSize - 100;
+  const size_t kFrameSize = kMaxPayloadSize - 3 * kLengthFieldLength -
+      kNalHeaderSize;
   uint8_t frame[kFrameSize] = {0x07, 0xFF,  // F=0, NRI=0, Type=7.
                                0x08, 0xFF,  // F=0, NRI=0, Type=8.
                                0x05};       // F=0, NRI=0, Type=5.
@@ -245,6 +246,50 @@ TEST(RtpPacketizerH264Test, TestStapA) {
   EXPECT_FALSE(packetizer->NextPacket(packet, &length, &last));
 }
 
+TEST(RtpPacketizerH264Test, TestTooSmallForStapAHeaders) {
+  const size_t kFrameSize = kMaxPayloadSize - 1;
+  uint8_t frame[kFrameSize] = {0x07, 0xFF,  // F=0, NRI=0, Type=7.
+                               0x08, 0xFF,  // F=0, NRI=0, Type=8.
+                               0x05};       // F=0, NRI=0, Type=5.
+  const size_t kPayloadOffset = 5;
+  for (size_t i = 0; i < kFrameSize - kPayloadOffset; ++i)
+    frame[i + kPayloadOffset] = i;
+  RTPFragmentationHeader fragmentation;
+  fragmentation.VerifyAndAllocateFragmentationHeader(3);
+  fragmentation.fragmentationOffset[0] = 0;
+  fragmentation.fragmentationLength[0] = 2;
+  fragmentation.fragmentationOffset[1] = 2;
+  fragmentation.fragmentationLength[1] = 2;
+  fragmentation.fragmentationOffset[2] = 4;
+  fragmentation.fragmentationLength[2] =
+      kNalHeaderSize + kFrameSize - kPayloadOffset;
+  scoped_ptr<RtpPacketizer> packetizer(
+      RtpPacketizer::Create(kRtpVideoH264, kMaxPayloadSize));
+  packetizer->SetPayloadData(frame, kFrameSize, &fragmentation);
+
+  uint8_t packet[kMaxPayloadSize] = {0};
+  size_t length = 0;
+  bool last = false;
+  ASSERT_TRUE(packetizer->NextPacket(packet, &length, &last));
+  size_t expected_packet_size = kNalHeaderSize;
+  for (size_t i = 0; i < 2; ++i) {
+    expected_packet_size += kLengthFieldLength +
+        fragmentation.fragmentationLength[i];
+  }
+  ASSERT_EQ(expected_packet_size, length);
+  EXPECT_FALSE(last);
+  for (size_t i = 0; i < 2; ++i)
+    VerifyStapAPayload(fragmentation, 0, i, frame, kFrameSize, packet, length);
+
+  ASSERT_TRUE(packetizer->NextPacket(packet, &length, &last));
+  expected_packet_size = fragmentation.fragmentationLength[2];
+  ASSERT_EQ(expected_packet_size, length);
+  EXPECT_TRUE(last);
+  VerifySingleNaluPayload(fragmentation, 2, frame, kFrameSize, packet, length);
+
+  EXPECT_FALSE(packetizer->NextPacket(packet, &length, &last));
+}
+
 TEST(RtpPacketizerH264Test, TestMixedStapA_FUA) {
   const size_t kFuaNaluSize = 2 * (kMaxPayloadSize - 100);
   const size_t kStapANaluSize = 100;