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

webrtc/modules/rtp_rtcp/source/rtp_sender_video.cc

 /*
  *  Copyright (c) 2012 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.
  */
 
 #include "webrtc/modules/rtp_rtcp/source/rtp_sender_video.h"
 
 #include <stdlib.h>
 #include <string.h>
 
 #include <memory>
 #include <vector>
 #include <utility>
 
 #include "webrtc/base/checks.h"
 #include "webrtc/base/logging.h"
+#include "webrtc/base/ptr_util.h"
 #include "webrtc/base/trace_event.h"
 #include "webrtc/modules/rtp_rtcp/include/rtp_rtcp_defines.h"
 #include "webrtc/modules/rtp_rtcp/source/byte_io.h"
 #include "webrtc/modules/rtp_rtcp/source/rtp_format_video_generic.h"
 #include "webrtc/modules/rtp_rtcp/source/rtp_format_vp8.h"
 #include "webrtc/modules/rtp_rtcp/source/rtp_format_vp9.h"
 #include "webrtc/modules/rtp_rtcp/source/rtp_header_extensions.h"
 #include "webrtc/modules/rtp_rtcp/source/rtp_packet_to_send.h"
 
 namespace webrtc {
@@ skipping 261 matching lines @@
                                const RTPFragmentationHeader* fragmentation,
                                const RTPVideoHeader* video_header) {
   if (payload_size == 0)
     return false;
 
   // Create header that will be reused in all packets.
   std::unique_ptr<RtpPacketToSend> rtp_header = rtp_sender_->AllocatePacket();
   rtp_header->SetPayloadType(payload_type);
   rtp_header->SetTimestamp(rtp_timestamp);
   rtp_header->set_capture_time_ms(capture_time_ms);
+  auto last_packet = rtc::MakeUnique<RtpPacketToSend>(*rtp_header);
 
   size_t fec_packet_overhead;
   bool red_enabled;
   int32_t retransmission_settings;
   {
     rtc::CritScope cs(&crit_);
     // According to
     // http://www.etsi.org/deliver/etsi_ts/126100_126199/126114/12.07.00_60/
     // ts_126114v120700p.pdf Section 7.4.5:
     // The MTSI client shall add the payload bytes as defined in this clause
     // onto the last RTP packet in each group of packets which make up a key
     // frame (I-frame or IDR frame in H.264 (AVC), or an IRAP picture in H.265
     // (HEVC)). The MTSI client may also add the payload bytes onto the last RTP
     // packet in each group of packets which make up another type of frame
     // (e.g. a P-Frame) only if the current value is different from the previous
     // value sent.
     if (video_header) {
       // Set rotation when key frame or when changed (to follow standard).
       // Or when different from 0 (to follow current receiver implementation).
       VideoRotation current_rotation = video_header->rotation;
       if (frame_type == kVideoFrameKey || current_rotation != last_rotation_ ||
           current_rotation != kVideoRotation_0)
-        rtp_header->SetExtension<VideoOrientation>(current_rotation);
+        last_packet->SetExtension<VideoOrientation>(current_rotation);
       last_rotation_ = current_rotation;
       // Report content type only for key frames.
       if (frame_type == kVideoFrameKey &&
           video_header->content_type != VideoContentType::UNSPECIFIED) {
-        rtp_header->SetExtension<VideoContentTypeExtension>(
+        last_packet->SetExtension<VideoContentTypeExtension>(
             video_header->content_type);
       }
     }
 
     // FEC settings.
     const FecProtectionParams& fec_params =
         frame_type == kVideoFrameKey ? key_fec_params_ : delta_fec_params_;
     if (flexfec_enabled())
       flexfec_sender_->SetFecParameters(fec_params);
     if (ulpfec_enabled())
       ulpfec_generator_.SetFecParameters(fec_params);
 
     fec_packet_overhead = CalculateFecPacketOverhead();
     red_enabled = this->red_enabled();
     retransmission_settings = retransmission_settings_;
   }
 
   size_t packet_capacity = rtp_sender_->MaxRtpPacketSize() -
                            fec_packet_overhead -
                            (rtp_sender_->RtxStatus() ? kRtxHeaderSize : 0);
   RTC_DCHECK_LE(packet_capacity, rtp_header->capacity());
   RTC_DCHECK_GT(packet_capacity, rtp_header->headers_size());
+  RTC_DCHECK_GT(packet_capacity, last_packet->headers_size());
   size_t max_data_payload_length = packet_capacity - rtp_header->headers_size();
+  RTC_DCHECK_GE(last_packet->headers_size(), rtp_header->headers_size());
+  size_t last_packet_reduction_len =
+      last_packet->headers_size() - rtp_header->headers_size();
 
   std::unique_ptr<RtpPacketizer> packetizer(RtpPacketizer::Create(
-      video_type, max_data_payload_length,
+      video_type, max_data_payload_length, last_packet_reduction_len,
       video_header ? &(video_header->codecHeader) : nullptr, frame_type));
   // Media packet storage.
   StorageType storage = packetizer->GetStorageType(retransmission_settings);
 
   // TODO(changbin): we currently don't support to configure the codec to
   // output multiple partitions for VP8. Should remove below check after the
   // issue is fixed.
   const RTPFragmentationHeader* frag =
       (video_type == kRtpVideoVp8) ? nullptr : fragmentation;
-  packetizer->SetPayloadData(payload_data, payload_size, frag);
+  size_t num_packets =
+      packetizer->SetPayloadData(payload_data, payload_size, frag);
+
+  if (num_packets == 0)
+    return false;
 
   bool first_frame = first_frame_sent_();
-  bool first = true;
-  bool last = false;
-  while (!last) {
-    std::unique_ptr<RtpPacketToSend> packet(new RtpPacketToSend(*rtp_header));
-
-    if (!packetizer->NextPacket(packet.get(), &last))
+  for (size_t i = 0; i < num_packets; ++i) {
+    bool last = (i + 1) == num_packets;
+    auto packet = last ? std::move(last_packet)
+                       : rtc::MakeUnique<RtpPacketToSend>(*rtp_header);
+    if (!packetizer->NextPacket(packet.get()))
       return false;
-    RTC_DCHECK_LE(packet->payload_size(), max_data_payload_length);
-
+    RTC_DCHECK_LE(packet->payload_size(),
+                  last ? max_data_payload_length - last_packet_reduction_len
+                       : max_data_payload_length);
     if (!rtp_sender_->AssignSequenceNumber(packet.get()))
       return false;
 
     const bool protect_packet =
         (packetizer->GetProtectionType() == kProtectedPacket);
     if (flexfec_enabled()) {
       // TODO(brandtr): Remove the FlexFEC code path when FlexfecSender
       // is wired up to PacedSender instead.
       SendVideoPacketWithFlexfec(std::move(packet), storage, protect_packet);
     } else if (red_enabled) {
       SendVideoPacketAsRedMaybeWithUlpfec(std::move(packet), storage,
                                           protect_packet);
     } else {
       SendVideoPacket(std::move(packet), storage);
     }
 
     if (first_frame) {
-      if (first) {
+      if (i == 0) {
         LOG(LS_INFO)
             << "Sent first RTP packet of the first video frame (pre-pacer)";
       }
       if (last) {
         LOG(LS_INFO)
             << "Sent last RTP packet of the first video frame (pre-pacer)";
       }
     }
-    first = false;
   }
 
   TRACE_EVENT_ASYNC_END1("webrtc", "Video", capture_time_ms, "timestamp",
                          rtp_timestamp);
   return true;
 }
 
 uint32_t RTPSenderVideo::VideoBitrateSent() const {
   rtc::CritScope cs(&stats_crit_);
   return video_bitrate_.Rate(clock_->TimeInMilliseconds()).value_or(0);
 }
 
 uint32_t RTPSenderVideo::FecOverheadRate() const {
   rtc::CritScope cs(&stats_crit_);
   return fec_bitrate_.Rate(clock_->TimeInMilliseconds()).value_or(0);
 }
 
 int RTPSenderVideo::SelectiveRetransmissions() const {
   rtc::CritScope cs(&crit_);
   return retransmission_settings_;
 }
 
 void RTPSenderVideo::SetSelectiveRetransmissions(uint8_t settings) {
   rtc::CritScope cs(&crit_);
   retransmission_settings_ = settings;
 }
 
 }  // namespace webrtc