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/webrtc@6881 4adac7df-926f-26a2-2b94-8c16560cd09d
diff --git a/modules/rtp_rtcp/source/rtp_format_h264.cc b/modules/rtp_rtcp/source/rtp_format_h264.cc
index 01ca426..85bd946 100644
--- a/modules/rtp_rtcp/source/rtp_format_h264.cc
+++ b/modules/rtp_rtcp/source/rtp_format_h264.cc
@@ -29,6 +29,7 @@
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 @@
// 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) {
- if (fragment_length > 0) {
- assert(fragment_length > 0);
- uint8_t header = payload_data_[fragment_offset];
- packets_.push(Packet(fragment_offset,
- fragment_length,
- aggregated_fragments == 0,
- false,
- true,
- header));
- payload_size_left -= fragment_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;
- }
+ while (payload_size_left >= fragment_length + fragment_headers_length) {
+ assert(fragment_length > 0);
+ uint8_t header = payload_data_[fragment_offset];
+ packets_.push(Packet(fragment_offset,
+ fragment_length,
+ aggregated_fragments == 0,
+ false,
+ true,
+ header));
+ payload_size_left -= fragment_length;
+ payload_size_left -= fragment_headers_length;
+
// 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 @@
*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 @@
while (packet.aggregated) {
// Add NAL unit length field.
RtpUtility::AssignUWord16ToBuffer(&buffer[index], packet.size);
- index += 2;
- *bytes_to_send += 2;
+ index += kLengthFieldSize;
+ *bytes_to_send += kLengthFieldSize;
// Add NAL unit.
memcpy(&buffer[index], &payload_data_[packet.offset], packet.size);
index += packet.size;
diff --git a/modules/rtp_rtcp/source/rtp_format_h264_unittest.cc b/modules/rtp_rtcp/source/rtp_format_h264_unittest.cc
index bf9d381..55b9712 100644
--- a/modules/rtp_rtcp/source/rtp_format_h264_unittest.cc
+++ b/modules/rtp_rtcp/source/rtp_format_h264_unittest.cc
@@ -211,7 +211,8 @@
}
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 @@
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;
