av2/common/obu_util.c

/*
 * Copyright (c) 2021, Alliance for Open Media. All rights reserved
 *
 * This source code is subject to the terms of the BSD 3-Clause Clear License
 * and the Alliance for Open Media Patent License 1.0. If the BSD 3-Clause Clear
 * License was not distributed with this source code in the LICENSE file, you
 * can obtain it at aomedia.org/license/software-license/bsd-3-c-c/.  If the
 * Alliance for Open Media Patent License 1.0 was not distributed with this
 * source code in the PATENTS file, you can obtain it at
 * aomedia.org/license/patent-license/.
 */
#include <assert.h>

#include "av2/common/obu_util.h"

#include "avm_dsp/bitreader_buffer.h"
#include "av2/common/enums.h"

// If read_obu_size() returns AVM_CODEC_OK, it is guaranteed that
// *length_field_size <= bytes_available.
static avm_codec_err_t read_obu_size(const uint8_t *data,
                                     size_t bytes_available,
                                     size_t *const obu_size,
                                     size_t *const length_field_size) {
  uint64_t u_obu_size = 0;
  if (avm_uleb_decode(data, bytes_available, &u_obu_size, length_field_size) !=
      0) {
    return AVM_CODEC_CORRUPT_FRAME;
  }

  if (u_obu_size > UINT32_MAX) return AVM_CODEC_CORRUPT_FRAME;
  *obu_size = (size_t)u_obu_size;
  assert(*length_field_size <= bytes_available);
  return AVM_CODEC_OK;
}

// Parses OBU header and stores values in 'header'.
static avm_codec_err_t read_obu_header(struct avm_read_bit_buffer *rb,
                                       ObuHeader *header) {
  if (!rb || !header) return AVM_CODEC_INVALID_PARAM;

  const ptrdiff_t bit_buffer_byte_length = rb->bit_buffer_end - rb->bit_buffer;

  if (bit_buffer_byte_length < 1) return AVM_CODEC_CORRUPT_FRAME;
  header->size = 1;

  header->obu_header_extension_flag = avm_rb_read_bit(rb);
  header->type = (OBU_TYPE)avm_rb_read_literal(rb, 5);  // obu_type

  header->obu_tlayer_id = avm_rb_read_literal(rb, TLAYER_BITS);

  if (header->obu_header_extension_flag) {
    if (bit_buffer_byte_length == 1) return AVM_CODEC_CORRUPT_FRAME;
    header->size += 1;

    header->obu_mlayer_id = avm_rb_read_literal(rb, MLAYER_BITS);
    header->obu_xlayer_id = avm_rb_read_literal(rb, XLAYER_BITS);
  } else {
    header->obu_mlayer_id = 0;
    if (header->type == OBU_MULTI_STREAM_DECODER_OPERATION ||
        header->type == OBU_TEMPORAL_DELIMITER)
      header->obu_xlayer_id = GLOBAL_XLAYER_ID;
    else
      header->obu_xlayer_id = 0;
  }

  if (header->type == OBU_TEMPORAL_DELIMITER) {
    if (header->obu_xlayer_id != GLOBAL_XLAYER_ID)
      return AVM_CODEC_CORRUPT_FRAME;
  }

  return AVM_CODEC_OK;
}

avm_codec_err_t avm_read_obu_header_and_size(const uint8_t *data,
                                             size_t bytes_available,
                                             ObuHeader *obu_header,
                                             size_t *const payload_size,
                                             size_t *const bytes_read) {
  size_t length_field_size_obu = 0;
  size_t obu_size = 0;
  avm_codec_err_t status;

  // Size field comes before the OBU header, and includes the OBU header
  status =
      read_obu_size(data, bytes_available, &obu_size, &length_field_size_obu);

  if (status != AVM_CODEC_OK) return status;

  struct avm_read_bit_buffer rb = { data + length_field_size_obu,
                                    data + bytes_available, 0, NULL, NULL };

  status = read_obu_header(&rb, obu_header);
  if (status != AVM_CODEC_OK) return status;

  // Derive the payload size from the data we've already read
  if (obu_size < obu_header->size) return AVM_CODEC_CORRUPT_FRAME;
  *payload_size = obu_size - obu_header->size;
  *bytes_read = length_field_size_obu + obu_header->size;

  assert(*bytes_read <= bytes_available);
  return AVM_CODEC_OK;
}