diff --git a/webrtc/modules/rtp_rtcp/source/rtp_format_h264.cc b/webrtc/modules/rtp_rtcp/source/rtp_format_h264.cc index 01ca426e9d..85bd94688d 100644 --- a/webrtc/modules/rtp_rtcp/source/rtp_format_h264.cc +++ b/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) { - 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 @@ 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; - *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/webrtc/modules/rtp_rtcp/source/rtp_format_h264_unittest.cc b/webrtc/modules/rtp_rtcp/source/rtp_format_h264_unittest.cc index bf9d381a9d..55b9712bd1 100644 --- a/webrtc/modules/rtp_rtcp/source/rtp_format_h264_unittest.cc +++ b/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 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;