Add H264 bitstream rewriting to limit frame reordering marker in header

webrtc/modules/rtp_rtcp/source/h264/bitstream_rewriter.cc

+/*
+ *  Copyright (c) 2016 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/h264/bitstream_rewriter.h"
+
+#include <algorithm>
+#include <memory>
+
+#include "webrtc/base/bitbuffer.h"
+#include "webrtc/base/bytebuffer.h"
+#include "webrtc/base/checks.h"
+#include "webrtc/base/logging.h"
+
+#include "webrtc/modules/rtp_rtcp/source/h264/h264_common.h"
+#include "webrtc/modules/rtp_rtcp/source/h264/sps_parser.h"
+
+namespace webrtc {
+
+// The maximum expected growth from adding a VUI to the SPS. It's actually
+// closer to 24 or so, but better safe than sorry.
+const size_t kMaxVuiSpsIncrease = 64;
+const size_t kNaluTypeSize = 1;
+
+#define RETURN_FALSE_ON_FAIL(x)                                  \
+  if (!(x)) {                                                    \
+    LOG_F(LS_ERROR) << " (line:" << __LINE__ << ") FAILED: " #x; \
+    return false;                                                \
+  }
+
+#define COPY_UINT8(src, dest, tmp)                   \
+  do {                                               \
+    RETURN_FALSE_ON_FAIL((src)->ReadUInt8(&tmp));    \
+    if (dest)                                        \
+      RETURN_FALSE_ON_FAIL((dest)->WriteUInt8(tmp)); \
+  } while (0)
+
+#define COPY_EXP_GOLOMB(src, dest, tmp)                          \
+  do {                                                           \
+    RETURN_FALSE_ON_FAIL((src)->ReadExponentialGolomb(&tmp));    \
+    if (dest)                                                    \
+      RETURN_FALSE_ON_FAIL((dest)->WriteExponentialGolomb(tmp)); \
+  } while (0)
+
+#define COPY_BITS(src, dest, tmp, bits)                   \
+  do {                                                    \
+    RETURN_FALSE_ON_FAIL((src)->ReadBits(&tmp, bits));    \
+    if (dest)                                             \
+      RETURN_FALSE_ON_FAIL((dest)->WriteBits(tmp, bits)); \
+  } while (0)
+
+bool H264BitstreamRewriter::ParseAndRewriteSps(const uint8_t* buffer,
+                                               size_t length) {
+  if (length <= kNaluTypeSize)
+    return false;
+  static const uint8_t kTypeMask = 0x1F;
+  uint8_t nal_type = buffer[0] & kTypeMask;
+  if (nal_type != H264Common::kSps) {
+    LOG(LS_ERROR) << "NALU is not an SPS.";
+    return false;
+  }
+  // Parse out the SPS RBSP. It should be small, so it's ok that we create a
+  // copy. We'll eventually write this back.
+  std::unique_ptr<rtc::ByteBufferWriter> sps_rbsp(
+      H264Common::ParseRbsp(buffer + kNaluTypeSize, length - kNaluTypeSize));
+  rtc::BitBuffer parsed_sps(reinterpret_cast<const uint8_t*>(sps_rbsp->Data()),
+                            sps_rbsp->Length());
+  SpsParser sps_parser;
+  if (!sps_parser.ParseFromRbspDecodedBitBuffer(&parsed_sps))
+    return false;
+
+  sps_state_ = sps_parser.GetState();

[noahric, 2016/05/18 01:35:06]

If you want sps_state_ to be exposed, have it as an out parameter as well, along
with output_buffer_.

[sprang_webrtc, 2016/05/20 16:10:59]

Done.

+  if (sps_state_->pic_order_cnt_type >= 2) {
+    // No need to rewrite VUI in this case.
+    return false;
+  }
+
+  // Create a new SPS buffer, large enough to hold our modifications.
+  // Also make sure these are zeroed out, since the last byte, if unaligned,
+  // needs to end in zeroes.
+  std::unique_ptr<uint8_t[]> translated_sps_rbsp(
+      new uint8_t[length + kMaxVuiSpsIncrease]());
+  // We're going to completely muck up alignment, so we need a BitBuffer to
+  // write with.
+  rtc::BitBufferWriter sps_writer(translated_sps_rbsp.get(),
+                                  length + kMaxVuiSpsIncrease);
+
+  // Check how far the SpsParser has read, and copy that data in bulk.
+  size_t byte_offset;
+  size_t bit_offset;
+  parsed_sps.GetCurrentOffset(&byte_offset, &bit_offset);
+  memcpy(translated_sps_rbsp.get(), sps_rbsp->Data(),
+         byte_offset + (bit_offset > 0 ? 1 : 0));  // OK to copy the last bits.
+
+  // SpsParser will have read the vui_params_present flag, which we want to

[noahric, 2016/05/18 01:35:06]

That's a suuuuuper dangerous coupling. I think you need to have a design that
doesn't rely on the implementation detail of SpsParser if this isn't also a part
of SpsParser.

In the case you move it, you could have an internal helper like ParseSpsUpToVui,
and then have the qp parser parse the Vui and do the rest, and this
functionality use the function below. That's safer, especially when defined
("UpToVui").

[sprang_webrtc, 2016/05/20 16:10:59]

Yeah, this was a bit iffy. Screwups should be caught be tests though.
Anyway, made a protected internal method SpsParser::ParseSpsUpToVui() that is
used
instead.

+  // modify, so back off one bit;
+  if (bit_offset == 0) {
+    --byte_offset;
+    bit_offset = 7;
+  } else {
+    --bit_offset;
+  }
+  sps_writer.Seek(byte_offset, bit_offset);
+
+  bool vui_updated;
+  if (!CopyAndRewriteVui(&parsed_sps, &sps_writer, &vui_updated)) {
+    LOG(LS_ERROR) << "Failed to parse/copy SPS VUI.";
+    return false;
+  }
+
+  if (!vui_updated)
+    return false;
+
+  if (!CopyRemainingBits(&parsed_sps, &sps_writer))
+    return false;
+
+  // Pad up to next byte with zeros.
+  sps_writer.GetCurrentOffset(&byte_offset, &bit_offset);
+  if (bit_offset > 0) {
+    sps_writer.WriteBits(0, 8 - bit_offset);
+    ++byte_offset;
+    bit_offset = 0;
+  }
+  RTC_DCHECK(byte_offset <= length + kMaxVuiSpsIncrease);
+
+  // Parsed and rewritten successfully, so go ahead and start writing to
+  // destination and mark that we're rewriting the picture order count type.
+  // Start by writing over the NALU header.
+  output_buffer_.reset(new rtc::ByteBufferWriter());
+  output_buffer_->WriteBytes(reinterpret_cast<const char*>(buffer),
+                             kNaluTypeSize);
+  H264Common::WriteRbsp(translated_sps_rbsp.get(), byte_offset,
+                        output_buffer_.get());
+  return true;
+}
+
+bool H264BitstreamRewriter::CopyAndRewriteVui(rtc::BitBuffer* source,

[noahric, 2016/05/18 01:35:06]

If you pass sps_state as an arg, the class doesn't need to hold state anymore.
Generally, that seems better, since it's an implementation detail that it has
state versus just returning whatever it needs from the function.

[sprang_webrtc, 2016/05/20 16:10:59]

Done.

+                                              rtc::BitBufferWriter* destination,
+                                              bool* out_vui_rewritten) {
+  uint32_t golomb_tmp;
+  uint32_t bits_tmp;
+  *out_vui_rewritten = false;
+
+  //
+  // vui_parameters_present_flag: u(1)
+  //
+  RETURN_FALSE_ON_FAIL(destination->WriteBits(1, 1));
+
+  // ********* IMPORTANT! **********
+  // Now we're at the VUI, so we want to (1) add it if it isn't present, and
+  // (2) rewrite frame reordering values so no reordering is allowed.
+  if (!sps_state_->vui_params_present) {
+    // Write a simple VUI with the parameters we want and 0 for all other flags.
+    // There are 8 flags to be off before the bitstream restriction flag.
+    RETURN_FALSE_ON_FAIL(destination->WriteBits(0, 8));
+    // bitstream_restriction_flag: u(1)
+    RETURN_FALSE_ON_FAIL(destination->WriteBits(1, 1));
+    RETURN_FALSE_ON_FAIL(
+        AddBitstreamRestriction(destination, sps_state_->max_num_ref_frames));
+  } else {
+    // Parse out the full VUI.
+    // aspect_ratio_info_present_flag: u(1)
+    COPY_BITS(source, destination, bits_tmp, 1);
+    if (bits_tmp == 1) {
+      // aspect_ratio_idc: u(8)
+      COPY_BITS(source, destination, bits_tmp, 8);
+      if (bits_tmp == 255u) {  // Extended_SAR
+        // sar_width/sar_height: u(16) each.
+        COPY_BITS(source, destination, bits_tmp, 32);
+      }
+    }
+    // overscan_info_present_flag: u(1)
+    COPY_BITS(source, destination, bits_tmp, 1);
+    if (bits_tmp == 1) {
+      // overscan_appropriate_flag: u(1)
+      COPY_BITS(source, destination, bits_tmp, 1);
+    }
+    // video_signal_type_present_flag: u(1)
+    COPY_BITS(source, destination, bits_tmp, 1);
+    if (bits_tmp == 1) {
+      // video_format + video_full_range_flag: u(3) + u(1)
+      COPY_BITS(source, destination, bits_tmp, 4);
+      // colour_description_present_flag: u(1)
+      COPY_BITS(source, destination, bits_tmp, 1);
+      if (bits_tmp == 1) {
+        // colour_primaries, transfer_characteristics, matrix_coefficients:
+        // u(8) each.
+        COPY_BITS(source, destination, bits_tmp, 24);
+      }
+    }
+    // chroma_loc_info_present_flag: u(1)
+    COPY_BITS(source, destination, bits_tmp, 1);
+    if (bits_tmp == 1) {
+      // chroma_sample_loc_type_(top|bottom)_field: ue(v) each.
+      COPY_EXP_GOLOMB(source, destination, golomb_tmp);
+      COPY_EXP_GOLOMB(source, destination, golomb_tmp);
+    }
+    // timing_info_present_flag: u(1)
+    COPY_BITS(source, destination, bits_tmp, 1);
+    if (bits_tmp == 1) {
+      // num_units_in_tick, time_scale: u(32) each
+      COPY_BITS(source, destination, bits_tmp, 32);
+      COPY_BITS(source, destination, bits_tmp, 32);
+      // fixed_frame_rate_flag: u(1)
+      COPY_BITS(source, destination, bits_tmp, 1);
+    }
+    // nal_hrd_parameters_present_flag: u(1)
+    uint32_t nal_hrd_parameters_present_flag;
+    COPY_BITS(source, destination, nal_hrd_parameters_present_flag, 1);
+    if (nal_hrd_parameters_present_flag == 1) {
+      RETURN_FALSE_ON_FAIL(CopyHrdParameters(source, destination));
+    }
+    // vcl_hrd_parameters_present_flag: u(1)
+    uint32_t vcl_hrd_parameters_present_flag;
+    COPY_BITS(source, destination, vcl_hrd_parameters_present_flag, 1);
+    if (vcl_hrd_parameters_present_flag == 1) {
+      RETURN_FALSE_ON_FAIL(CopyHrdParameters(source, destination));
+    }
+    if (nal_hrd_parameters_present_flag == 1 ||
+        vcl_hrd_parameters_present_flag == 1) {
+      // low_delay_hrd_flag: u(1)
+      COPY_BITS(source, destination, bits_tmp, 1);
+    }
+    // pic_struct_present_flag: u(1)
+    COPY_BITS(source, destination, bits_tmp, 1);
+
+    // bitstream_restriction_flag: u(1)
+    uint32_t bitstream_restriction_flag;
+    RETURN_FALSE_ON_FAIL(source->ReadBits(&bitstream_restriction_flag, 1));
+    RETURN_FALSE_ON_FAIL(destination->WriteBits(1, 1));
+    if (bitstream_restriction_flag == 0) {
+      // We're adding one from scratch.
+      RETURN_FALSE_ON_FAIL(
+          AddBitstreamRestriction(destination, sps_state_->max_num_ref_frames));
+    } else {
+      // We're replacing.
+      // motion_vectors_over_pic_boundaries_flag: u(1)
+      COPY_BITS(source, destination, bits_tmp, 1);
+      // max_bytes_per_pic_denom: ue(v)
+      COPY_EXP_GOLOMB(source, destination, golomb_tmp);
+      // max_bits_per_mb_denom: ue(v)
+      COPY_EXP_GOLOMB(source, destination, golomb_tmp);
+      // log2_max_mv_length_horizontal: ue(v)
+      COPY_EXP_GOLOMB(source, destination, golomb_tmp);
+      // log2_max_mv_length_vertical: ue(v)
+      COPY_EXP_GOLOMB(source, destination, golomb_tmp);
+      // ********* IMPORTANT! **********
+      // The next two are the ones we need to set to low numbers:
+      // max_num_reorder_frames: ue(v)
+      // max_dec_frame_buffering: ue(v)
+      // However, if they are already set to no greater than the numbers we
+      // want, then we don't need to be rewriting.
+      uint32_t max_num_reorder_frames, max_dec_frame_buffering;
+      RETURN_FALSE_ON_FAIL(
+          source->ReadExponentialGolomb(&max_num_reorder_frames));
+      RETURN_FALSE_ON_FAIL(
+          source->ReadExponentialGolomb(&max_dec_frame_buffering));
+      if (max_num_reorder_frames == 0 &&
+          max_dec_frame_buffering <= sps_state_->max_num_ref_frames) {
+        LOG(LS_INFO) << "VUI bitstream already contains an optimal VUI.";
+        return true;
+      }
+      RETURN_FALSE_ON_FAIL(destination->WriteExponentialGolomb(0));
+      RETURN_FALSE_ON_FAIL(
+          destination->WriteExponentialGolomb(sps_state_->max_num_ref_frames));
+    }
+  }
+  *out_vui_rewritten = true;
+  return true;
+}
+
+// Copies a VUI HRD parameters segment.
+bool H264BitstreamRewriter::CopyHrdParameters(
+    rtc::BitBuffer* source,
+    rtc::BitBufferWriter* destination) {
+  uint32_t golomb_tmp;
+  uint32_t bits_tmp;
+
+  // cbp_cnt_minus1: ue(v)
+  uint32_t cbp_cnt_minus1;
+  COPY_EXP_GOLOMB(source, destination, cbp_cnt_minus1);
+  // bit_rate_scale and cbp_size_scale: u(4) each
+  COPY_BITS(source, destination, bits_tmp, 8);
+  for (size_t i = 0; i <= cbp_cnt_minus1; ++i) {
+    // bit_rate_value_minus1 and cbp_size_value_minus1: ue(v) each
+    COPY_EXP_GOLOMB(source, destination, golomb_tmp);
+    COPY_EXP_GOLOMB(source, destination, golomb_tmp);
+    // cbr_flag: u(1)
+    COPY_BITS(source, destination, bits_tmp, 1);
+  }
+  // initial_cbp_removal_delay_length_minus1: u(5)
+  COPY_BITS(source, destination, bits_tmp, 5);
+  // cbp_removal_delay_length_minus1: u(5)
+  COPY_BITS(source, destination, bits_tmp, 5);
+  // dbp_output_delay_length_minus1: u(5)
+  COPY_BITS(source, destination, bits_tmp, 5);
+  // time_offset_length: u(5)
+  COPY_BITS(source, destination, bits_tmp, 5);
+  return true;
+}
+
+// These functions are similar to webrtc::H264SpsParser::Parse, and based on the
+// same version of the H.264 standard. You can find it here:
+// http://www.itu.int/rec/T-REC-H.264
+
+// Adds a bitstream restriction VUI segment.
+bool H264BitstreamRewriter::AddBitstreamRestriction(
+    rtc::BitBufferWriter* destination,
+    uint32_t max_num_ref_frames) {
+  // motion_vectors_over_pic_boundaries_flag: u(1)
+  // Default is 1 when not present.
+  RETURN_FALSE_ON_FAIL(destination->WriteBits(1, 1));
+  // max_bytes_per_pic_denom: ue(v)
+  // Default is 2 when not present.
+  RETURN_FALSE_ON_FAIL(destination->WriteExponentialGolomb(2));
+  // max_bits_per_mb_denom: ue(v)
+  // Default is 1 when not present.
+  RETURN_FALSE_ON_FAIL(destination->WriteExponentialGolomb(1));
+  // log2_max_mv_length_horizontal: ue(v)
+  // log2_max_mv_length_vertical: ue(v)
+  // Both default to 16 when not present.
+  RETURN_FALSE_ON_FAIL(destination->WriteExponentialGolomb(16));
+  RETURN_FALSE_ON_FAIL(destination->WriteExponentialGolomb(16));
+
+  // ********* IMPORTANT! **********
+  // max_num_reorder_frames: ue(v)
+  RETURN_FALSE_ON_FAIL(destination->WriteExponentialGolomb(0));
+  // max_dec_frame_buffering: ue(v)
+  RETURN_FALSE_ON_FAIL(destination->WriteExponentialGolomb(max_num_ref_frames));
+  return true;
+}
+
+bool H264BitstreamRewriter::CopyRemainingBits(
+    rtc::BitBuffer* source,
+    rtc::BitBufferWriter* destination) {
+  uint32_t bits_tmp;
+  // Try to get at least the destination aligned.
+  if (source->RemainingBitCount() > 0 && source->RemainingBitCount() % 8 != 0) {
+    size_t misaligned_bits = source->RemainingBitCount() % 8;
+    COPY_BITS(source, destination, bits_tmp, misaligned_bits);
+  }
+  while (source->RemainingBitCount() > 0) {
+    size_t count = std::min(static_cast<size_t>(32u),
+                            static_cast<size_t>(source->RemainingBitCount()));
+    COPY_BITS(source, destination, bits_tmp, count);
+  }
+  // TODO(noahric): The last byte could be all zeroes now, which we should just
+  // strip.
+  return true;
+}
+
+}  // namespace webrtc