Check that remaining OBU payload size is non-negative

After removing ULEBs for trimming and extension header, the remaining payload must be, at a minimum, 0. PiperOrigin-RevId: 714121233
AOMediaCodec · Jan 13, 2025 · 8106b84 · 8106b84
1 parent b05301b
commit 8106b84
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 6 deletions.
diff --git a/iamf/obu/obu_header.cc b/iamf/obu/obu_header.cc
@@ -208,12 +208,12 @@ absl::Status GetObuSizeAndValidate(const LebGenerator& leb_generator,
 // Returns the size of the payload associated with the OBU, i.e. the number of
 // bytes that contain payload data. See
 // https://aomediacodec.github.io/iamf/#obu_size for more details.
-int64_t GetObuPayloadSize(const DecodedUleb128& obu_size,
-                          const uint8_t& num_samples_to_trim_at_end_size,
-                          const uint8_t& num_samples_to_trim_at_start_size,
-                          const uint8_t& extension_header_size_size,
-                          const uint8_t& extension_header_bytes_size) {
-  return obu_size -
+int64_t GetObuPayloadSize(DecodedUleb128 obu_size,
+                          uint8_t num_samples_to_trim_at_end_size,
+                          uint8_t num_samples_to_trim_at_start_size,
+                          uint8_t extension_header_size_size,
+                          uint8_t extension_header_bytes_size) {
+  return static_cast<int64_t>(obu_size) -
          (num_samples_to_trim_at_end_size + num_samples_to_trim_at_start_size +
           extension_header_size_size + extension_header_bytes_size);
 }
@@ -276,6 +276,10 @@ absl::Status ObuHeader::ReadAndValidate(
       obu_size, num_samples_to_trim_at_end_size,
       num_samples_to_trim_at_start_size, extension_header_size_size,
       extension_header_bytes.size());
+  if (output_payload_serialized_size < 0) {
+    return absl::InvalidArgumentError(
+        "obu_size not valid for OBU flags. Negative remaining payload size.");
+  }
 
   RETURN_IF_NOT_OK(Validate(*this));
 

diff --git a/iamf/obu/tests/obu_header_test.cc b/iamf/obu/tests/obu_header_test.cc
@@ -30,6 +30,7 @@ namespace iamf_tools {
 namespace {
 
 using ::absl_testing::IsOk;
+using ::testing::Not;
 
 // Max value of a decoded ULEB128.
 constexpr uint32_t kMaxUlebDecoded = UINT32_MAX;
@@ -925,5 +926,20 @@ TEST_F(ObuHeaderTest, ReadAndValidateTrimmingStatusFlagNonZeroBothTrims) {
   EXPECT_TRUE(obu_header_.extension_header_bytes.empty());
 }
 
+TEST_F(ObuHeaderTest, NegativePayloadSizeNotAcceptable) {
+  std::vector<uint8_t> source_data = {kAudioFrameId0WithTrim,
+                                      // `obu_size`
+                                      2,
+                                      // `num_samples_to_trim_at_end`.
+                                      0x80, 0x01,
+                                      // `num_samples_to_trim_at_start`.
+                                      0x02};
+  auto read_bit_buffer = MemoryBasedReadBitBuffer::CreateFromSpan(
+      1024, absl::MakeConstSpan(source_data));
+  EXPECT_THAT(
+      obu_header_.ReadAndValidate(*read_bit_buffer, payload_serialized_size_),
+      Not(IsOk()));
+}
+
 }  // namespace
 }  // namespace iamf_tools