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

webrtc/modules/rtp_rtcp/source/rtp_format_h264.cc

 /*
  *  Copyright (c) 2014 The WebRTC project authors. All Rights Reserved.
  *
  *  Use of this source code is governed by a BSD-style license
  *  that can be found in the LICENSE file in the root of the source
  *  tree. An additional intellectual property rights grant can be found
  *  in the file PATENTS.  All contributing project authors may
  *  be found in the AUTHORS file in the root of the source tree.
  */
 
@@ skipping 61 matching lines @@
 
     offsets->push_back(offset + kStapAHeaderSize);
     offset += kLengthFieldSize + nalu_size;
   }
   return true;
 }
 
 }  // 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 ||
             packetization_mode == H264PacketizationMode::SingleNalUnit);
 }
 
 RtpPacketizerH264::~RtpPacketizerH264() {
 }
 
 RtpPacketizerH264::Fragment::Fragment(const uint8_t* buffer, size_t length)
     : buffer(buffer), length(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) {
   RTC_DCHECK(packets_.empty());
   RTC_DCHECK(input_fragments_.empty());
   RTC_DCHECK(fragmentation);
   for (int i = 0; i < fragmentation->fragmentationVectorSize; ++i) {
     const uint8_t* buffer =
         &payload_data[fragmentation->fragmentationOffset[i]];
     size_t length = fragmentation->fragmentationLength[i];
@@ skipping 48 matching lines @@
                                     SpsValidEvent::kSentSpsParseFailure,
                                     SpsValidEvent::kSpsRewrittenMax);
           break;
       }
     }
 
     if (!updated_sps)
       input_fragments_.push_back(Fragment(buffer, length));
   }
   GeneratePackets();
+  return total_packets_;
 }
 
 void RtpPacketizerH264::GeneratePackets() {
   for (size_t i = 0; i < input_fragments_.size();) {
     switch (packetization_mode_) {
       case H264PacketizationMode::SingleNalUnit:
         PacketizeSingleNalu(i);
         ++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 {
           i = PacketizeStapA(i);
         }
         break;
     }
   }
 }
 
 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;
 
-  const size_t avg_size = (fragment_length + num_fragments - 1) / num_fragments;
+  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;
+
   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);
 }
 
 size_t 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;
   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
     // NALU of this packet.
     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;
 }
 
 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 "
       << fragment->length << ", packetization mode " << packetization_mode_;
   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;
   }
 
   PacketUnit packet = packets_.front();
   if (packet.first_fragment && packet.last_fragment) {
     // Single NAL unit packet.
     size_t bytes_to_send = packet.source_fragment.length;
     uint8_t* buffer = rtp_packet->AllocatePayload(bytes_to_send);
     memcpy(buffer, packet.source_fragment.buffer, bytes_to_send);
     packets_.pop();
     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);
   // STAP-A NALU header.
   buffer[0] = (packet->header & (kFBit | kNriMask)) | H264::NaluType::kStapA;
   size_t index = kNalHeaderSize;
   bool is_last_fragment = packet->last_fragment;
   while (packet->aggregated) {
     const Fragment& fragment = packet->source_fragment;
     // Add NAL unit length field.
@@ skipping 319 matching lines @@
   h264->packetization_type = kH264FuA;
   h264->nalu_type = original_nal_type;
   if (first_fragment) {
     h264->nalus[h264->nalus_length] = nalu;
     h264->nalus_length = 1;
   }
   return true;
 }
 
 }  // namespace webrtc