Fix packetization logic to leave space for extensions in the last packet

webrtc/modules/rtp_rtcp/source/rtp_format_h264.cc

Index: webrtc/modules/rtp_rtcp/source/rtp_format_h264.cc
diff --git a/webrtc/modules/rtp_rtcp/source/rtp_format_h264.cc b/webrtc/modules/rtp_rtcp/source/rtp_format_h264.cc
index e190ea2f6aeea5f8e261cec86c6782477e867d32..ce6abaf9ec7b20a1edee9fd995935fe9fa5c6c54 100644
--- a/webrtc/modules/rtp_rtcp/source/rtp_format_h264.cc
+++ b/webrtc/modules/rtp_rtcp/source/rtp_format_h264.cc
@@ -79,8 +79,11 @@ bool ParseStapAStartOffsets(const uint8_t* nalu_ptr,
 }  // namespace
 
 RtpPacketizerH264::RtpPacketizerH264(size_t max_payload_len,
+                                     size_t last_packet_extensions_len,
                                      H264PacketizationMode packetization_mode)
     : max_payload_len_(max_payload_len),
+      last_packet_extensions_len_(last_packet_extensions_len),
+      total_packets_(0),
       packetization_mode_(packetization_mode) {
   // Guard against uninitialized memory in packetization_mode.
   RTC_CHECK(packetization_mode == H264PacketizationMode::NonInterleaved ||
@@ -95,7 +98,7 @@ RtpPacketizerH264::Fragment::Fragment(const uint8_t* buffer, size_t length)
 RtpPacketizerH264::Fragment::Fragment(const Fragment& fragment)
     : buffer(fragment.buffer), length(fragment.length) {}
 
-void RtpPacketizerH264::SetPayloadData(
+size_t RtpPacketizerH264::SetPayloadData(
     const uint8_t* payload_data,
     size_t payload_size,
     const RTPFragmentationHeader* fragmentation) {
@@ -164,6 +167,7 @@ void RtpPacketizerH264::SetPayloadData(
       input_fragments_.push_back(Fragment(buffer, length));
   }
   GeneratePackets();
+  return total_packets_;
 }
 
 void RtpPacketizerH264::GeneratePackets() {
@@ -174,7 +178,10 @@ void RtpPacketizerH264::GeneratePackets() {
         ++i;
         break;
       case H264PacketizationMode::NonInterleaved:
-        if (input_fragments_[i].length > max_payload_len_) {
+        if (input_fragments_[i].length > max_payload_len_ ||
+            (i + 1 == input_fragments_.size() &&
+             (input_fragments_[i].length + last_packet_extensions_len_ >
+              max_payload_len_))) {
           PacketizeFuA(i);
           ++i;
         } else {
@@ -189,24 +196,41 @@ void RtpPacketizerH264::PacketizeFuA(size_t fragment_index) {
   // Fragment payload into packets (FU-A).
   // Strip out the original header and leave room for the FU-A header.
   const Fragment& fragment = input_fragments_[fragment_index];
-
+  bool is_last_fragment = fragment_index + 1 == input_fragments_.size();
   size_t fragment_length = fragment.length - kNalHeaderSize;
   size_t offset = kNalHeaderSize;
   size_t bytes_available = max_payload_len_ - kFuAHeaderSize;
-  const size_t num_fragments =
-      (fragment_length + (bytes_available - 1)) / bytes_available;
+  size_t extra_len = is_last_fragment ? last_packet_extensions_len_ : 0;
+
+  size_t num_packets =
+      (fragment_length + extra_len + (bytes_available - 1)) / bytes_available;
+
+  const size_t avg_size =
+      (fragment_length + extra_len + num_packets - 1) / num_packets;
 
-  const size_t avg_size = (fragment_length + num_fragments - 1) / num_fragments;
   while (fragment_length > 0) {
     size_t packet_length = avg_size;
-    if (fragment_length < avg_size)
+    if (fragment_length <= packet_length) {  // Last portion of the payload
       packet_length = fragment_length;
+      // One additional packet may be used for extensions in the last packet.
+      // Together with last payload packet there may be at most 2 of them.
+      RTC_CHECK_LE(num_packets, 2);
+      // Whole payload fits in the first num_packets-1 packets but extra packet
+      // is used for extensions.
+      if (num_packets == 2) {
+        // Leave at least one byte of data for the last packet.
+        packet_length = packet_length - 1;
+      }
+    }
+    RTC_CHECK_GT(packet_length, 0);
     packets_.push(PacketUnit(Fragment(fragment.buffer + offset, packet_length),
                              offset - kNalHeaderSize == 0,
                              fragment_length == packet_length, false,
                              fragment.buffer[0]));
     offset += packet_length;
     fragment_length -= packet_length;
+    total_packets_++;
+    num_packets--;
   }
   RTC_CHECK_EQ(0, fragment_length);
 }
@@ -218,19 +242,17 @@ size_t RtpPacketizerH264::PacketizeStapA(size_t fragment_index) {
   size_t fragment_headers_length = 0;
   const Fragment* fragment = &input_fragments_[fragment_index];
   RTC_CHECK_GE(payload_size_left, fragment->length);
-  while (payload_size_left >= fragment->length + fragment_headers_length) {
+  total_packets_++;
+  while (payload_size_left >= fragment->length + fragment_headers_length &&
+         (fragment_index + 1 < input_fragments_.size() ||
+          payload_size_left >= fragment->length + fragment_headers_length +
+                                   last_packet_extensions_len_)) {
     RTC_CHECK_GT(fragment->length, 0);
     packets_.push(PacketUnit(*fragment, aggregated_fragments == 0, false, true,
                              fragment->buffer[0]));
     payload_size_left -= fragment->length;
     payload_size_left -= fragment_headers_length;
 
-    // Next fragment.
-    ++fragment_index;
-    if (fragment_index == input_fragments_.size())
-      break;
-    fragment = &input_fragments_[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
@@ -238,7 +260,14 @@ size_t RtpPacketizerH264::PacketizeStapA(size_t fragment_index) {
     if (aggregated_fragments == 0)
       fragment_headers_length += kNalHeaderSize + kLengthFieldSize;
     ++aggregated_fragments;
+
+    // Next fragment.
+    ++fragment_index;
+    if (fragment_index == input_fragments_.size())
+      break;
+    fragment = &input_fragments_[fragment_index];
   }
+  RTC_CHECK_GT(aggregated_fragments, 0);
   packets_.back().last_fragment = true;
   return fragment_index;
 }
@@ -246,6 +275,8 @@ size_t RtpPacketizerH264::PacketizeStapA(size_t fragment_index) {
 void RtpPacketizerH264::PacketizeSingleNalu(size_t fragment_index) {
   // Add a single NALU to the queue, no aggregation.
   size_t payload_size_left = max_payload_len_;
+  if (fragment_index + 1 == input_fragments_.size())
+    payload_size_left -= last_packet_extensions_len_;
   const Fragment* fragment = &input_fragments_[fragment_index];
   RTC_CHECK_GE(payload_size_left, fragment->length)
       << "Payload size left " << payload_size_left << ", fragment length "
@@ -253,14 +284,12 @@ void RtpPacketizerH264::PacketizeSingleNalu(size_t fragment_index) {
   RTC_CHECK_GT(fragment->length, 0u);
   packets_.push(PacketUnit(*fragment, true /* first */, true /* last */,
                            false /* aggregated */, fragment->buffer[0]));
+  total_packets_++;
 }
 
-bool RtpPacketizerH264::NextPacket(RtpPacketToSend* rtp_packet,
-                                   bool* last_packet) {
+bool RtpPacketizerH264::NextPacket(RtpPacketToSend* rtp_packet) {
   RTC_DCHECK(rtp_packet);
-  RTC_DCHECK(last_packet);
   if (packets_.empty()) {
-    *last_packet = true;
     return false;
   }
 
@@ -274,19 +303,25 @@ bool RtpPacketizerH264::NextPacket(RtpPacketToSend* rtp_packet,
     input_fragments_.pop_front();
   } else if (packet.aggregated) {
     RTC_CHECK_EQ(H264PacketizationMode::NonInterleaved, packetization_mode_);
-    NextAggregatePacket(rtp_packet);
+    NextAggregatePacket(rtp_packet, total_packets_ == 1);
   } else {
     RTC_CHECK_EQ(H264PacketizationMode::NonInterleaved, packetization_mode_);
     NextFragmentPacket(rtp_packet);
   }
   RTC_DCHECK_LE(rtp_packet->payload_size(), max_payload_len_);
-  *last_packet = packets_.empty();
-  rtp_packet->SetMarker(*last_packet);
+  if (packets_.empty()) {
+    RTC_DCHECK_LE(rtp_packet->payload_size(),
+                  max_payload_len_ - last_packet_extensions_len_);
+  }
+  rtp_packet->SetMarker(packets_.empty());
+  total_packets_--;
   return true;
 }
 
-void RtpPacketizerH264::NextAggregatePacket(RtpPacketToSend* rtp_packet) {
-  uint8_t* buffer = rtp_packet->AllocatePayload(max_payload_len_);
+void RtpPacketizerH264::NextAggregatePacket(RtpPacketToSend* rtp_packet,
+                                            bool last) {
+  uint8_t* buffer = rtp_packet->AllocatePayload(
+      last ? max_payload_len_ - last_packet_extensions_len_ : max_payload_len_);
   RTC_DCHECK(buffer);
   PacketUnit* packet = &packets_.front();
   RTC_CHECK(packet->first_fragment);