Add support for 3-byte headers in VideoToolbox NALU parser.

webrtc/modules/video_coding/codecs/h264/h264_video_toolbox_nalu.cc

 /*
  *  Copyright (c) 2015 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/video_coding/codecs/h264/h264_video_toolbox_nalu.h"
 
 #if defined(WEBRTC_VIDEO_TOOLBOX_SUPPORTED)
 
 #include <CoreFoundation/CoreFoundation.h>
 #include <memory>
 #include <vector>
 
 #include "webrtc/base/checks.h"
 #include "webrtc/base/logging.h"
-#include "webrtc/common_video/h264/h264_common.h"
 
 namespace webrtc {
 
 using H264::NaluType;
 using H264::kAud;
 using H264::kSps;
 using H264::ParseNaluType;
+using H264::NaluIndex;

[tkchin_webrtc, 2016/09/27 11:29:35]

nit: consider alphabetizing

[kthelgason, 2016/09/28 08:48:27]

Done.

 
 const char kAnnexBHeaderBytes[4] = {0, 0, 0, 1};
 const size_t kAvccHeaderByteSize = sizeof(uint32_t);
 
 bool H264CMSampleBufferToAnnexBBuffer(
     CMSampleBufferRef avcc_sample_buffer,
     bool is_keyframe,
     rtc::Buffer* annexb_buffer,
     webrtc::RTPFragmentationHeader** out_header) {
   RTC_DCHECK(avcc_sample_buffer);
@@ skipping 10 matching lines @@
 
   // Get parameter set information.
   int nalu_header_size = 0;
   size_t param_set_count = 0;
   OSStatus status = CMVideoFormatDescriptionGetH264ParameterSetAtIndex(
       description, 0, nullptr, nullptr, &param_set_count, &nalu_header_size);
   if (status != noErr) {
     LOG(LS_ERROR) << "Failed to get parameter set.";
     return false;
   }
-  // TODO(tkchin): handle other potential sizes.
-  RTC_DCHECK_EQ(nalu_header_size, 4);

[tkchin_webrtc, 2016/09/27 11:29:35]

why is this removed? Perhaps should be checked that it is equal to
kAvccHeaderByteSize instead?

[kthelgason, 2016/09/28 08:48:27]

That doesn't make sense to me, won't kAvccHeaderByteSize always be 4?
nalu_header_size is supposed to be allowed to be either 3 or 4, right?

[tkchin_webrtc, 2016/09/28 11:33:24]

iiuc nalu_header_size here is not the 0001/001. It is the length byte for avcc
format. We expect it to only 4 bytes (the size of a 32-bit int) but who knows if
what's returned by the CMSampleBuffer changes.

   RTC_DCHECK_EQ(param_set_count, 2u);
 
   // Truncate any previous data in the buffer without changing its capacity.
   annexb_buffer->SetSize(0);
 
   size_t nalu_offset = 0;
   std::vector<size_t> frag_offsets;
   std::vector<size_t> frag_lengths;
 
   // Place all parameter sets at the front of buffer.
@@ skipping 61 matching lines @@
     uint32_t* uint32_data_ptr = reinterpret_cast<uint32_t*>(data_ptr);
     uint32_t packet_size = CFSwapInt32BigToHost(*uint32_data_ptr);
     // Update buffer.
     annexb_buffer->AppendData(kAnnexBHeaderBytes, sizeof(kAnnexBHeaderBytes));
     annexb_buffer->AppendData(data_ptr + nalu_header_size, packet_size);
     // Update fragmentation.
     frag_offsets.push_back(nalu_offset + sizeof(kAnnexBHeaderBytes));
     frag_lengths.push_back(packet_size);
     nalu_offset += sizeof(kAnnexBHeaderBytes) + packet_size;
 
-    size_t bytes_written = packet_size + nalu_header_size;
+    size_t bytes_written = packet_size + sizeof(kAnnexBHeaderBytes);

[tkchin_webrtc, 2016/09/27 11:29:35]

good catch

     bytes_remaining -= bytes_written;
     data_ptr += bytes_written;
   }
   RTC_DCHECK_EQ(bytes_remaining, (size_t)0);
 
   std::unique_ptr<webrtc::RTPFragmentationHeader> header;
   header.reset(new webrtc::RTPFragmentationHeader());
   header->VerifyAndAllocateFragmentationHeader(frag_offsets.size());
   RTC_DCHECK_EQ(frag_lengths.size(), frag_offsets.size());
   for (size_t i = 0; i < frag_offsets.size(); ++i) {
@@ skipping 151 matching lines @@
       &description);
   if (status != noErr) {
     LOG(LS_ERROR) << "Failed to create video format description.";
     return nullptr;
   }
   return description;
 }
 
 AnnexBBufferReader::AnnexBBufferReader(const uint8_t* annexb_buffer,
                                        size_t length)
-    : start_(annexb_buffer), offset_(0), next_offset_(0), length_(length) {
+    : start_(annexb_buffer), length_(length) {
   RTC_DCHECK(annexb_buffer);
-  offset_ = FindNextNaluHeader(start_, length_, 0);
-  next_offset_ =
-      FindNextNaluHeader(start_, length_, offset_ + sizeof(kAnnexBHeaderBytes));
+  offsets_ = H264::FindNaluIndices(annexb_buffer, length);

[tkchin_webrtc, 2016/09/27 11:29:34]

How large do we expect these buffers to be? Is it safe to put potentially heavy
work into the ctor like this?

[kthelgason, 2016/09/28 08:48:27]

I have no idea, didn't consider that. In what context would it represent an
issue? The work has to be done at some point.

[tkchin_webrtc, 2016/09/28 11:33:24]

It probably doesn't matter in this particular case. It would only be an issue if
the buffer was extremely large. Just following the general paradigm of not
putting heavy work into ctors.

+  offset_ = offsets_.begin();
 }
 
 bool AnnexBBufferReader::ReadNalu(const uint8_t** out_nalu,
                                   size_t* out_length) {
   RTC_DCHECK(out_nalu);
   RTC_DCHECK(out_length);
   *out_nalu = nullptr;
   *out_length = 0;
 
-  size_t data_offset = offset_ + sizeof(kAnnexBHeaderBytes);
-  if (data_offset > length_) {
+  if (offset_ == offsets_.end()) {
     return false;
   }
-  *out_nalu = start_ + data_offset;
-  *out_length = next_offset_ - data_offset;
-  offset_ = next_offset_;
-  next_offset_ =
-      FindNextNaluHeader(start_, length_, offset_ + sizeof(kAnnexBHeaderBytes));
+  *out_nalu = start_ + offset_->payload_start_offset;
+  *out_length = offset_->payload_size;
+  ++offset_;
   return true;
 }
 
 size_t AnnexBBufferReader::BytesRemaining() const {
-  return length_ - offset_;
-}
-
-size_t AnnexBBufferReader::FindNextNaluHeader(const uint8_t* start,
-                                              size_t length,
-                                              size_t offset) const {
-  RTC_DCHECK(start);
-  if (offset + sizeof(kAnnexBHeaderBytes) > length) {
-    return length;
+  if (offset_ == offsets_.end()) {
+    return 0;
   }
-  // NALUs are separated by an 00 00 00 01 header. Scan the byte stream
-  // starting from the offset for the next such sequence.
-  const uint8_t* current = start + offset;
-  // The loop reads sizeof(kAnnexBHeaderBytes) at a time, so stop when there
-  // aren't enough bytes remaining.
-  const uint8_t* const end = start + length - sizeof(kAnnexBHeaderBytes);
-  while (current < end) {
-    if (current[3] > 1) {
-      current += 4;
-    } else if (current[3] == 1 && current[2] == 0 && current[1] == 0 &&
-               current[0] == 0) {
-      return current - start;
-    } else {
-      ++current;
-    }
-  }
-  return length;
+  return length_ - offset_->start_offset;
 }
 
 AvccBufferWriter::AvccBufferWriter(uint8_t* const avcc_buffer, size_t length)
     : start_(avcc_buffer), offset_(0), length_(length) {
   RTC_DCHECK(avcc_buffer);
 }
 
 bool AvccBufferWriter::WriteNalu(const uint8_t* data, size_t data_size) {
   // Check if we can write this length of data.
   if (data_size + kAvccHeaderByteSize > BytesRemaining()) {
     return false;
   }
   // Write length header, which needs to be big endian.
   uint32_t big_endian_length = CFSwapInt32HostToBig(data_size);
   memcpy(start_ + offset_, &big_endian_length, sizeof(big_endian_length));
   offset_ += sizeof(big_endian_length);
   // Write data.
   memcpy(start_ + offset_, data, data_size);
   offset_ += data_size;
   return true;
 }
 
 size_t AvccBufferWriter::BytesRemaining() const {
   return length_ - offset_;
 }
 
 }  // namespace webrtc
 
 #endif  // defined(WEBRTC_VIDEO_TOOLBOX_SUPPORTED)