From 13ccde8637bb4d71bbf7ce65c00a844792028d0f Mon Sep 17 00:00:00 2001
From: Zaggy1024 <zaggy1024@gmail.com>
Date: Sat, 8 Oct 2022 20:06:15 -0500
Subject: [PATCH] LibVideo: Improve error reporting for VP9 range decoder

init_bool will now check whether there is enough data in the bitstream
for the range coding size to be fully read.

exit_bool will now read the entire padding element regardless of size,
which the spec does not specify a limit on.
---
 Userland/Libraries/LibVideo/VP9/BitStream.cpp | 25 +++++++++++++------
 Userland/Libraries/LibVideo/VP9/BitStream.h   |  3 ++-
 2 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/Userland/Libraries/LibVideo/VP9/BitStream.cpp b/Userland/Libraries/LibVideo/VP9/BitStream.cpp
index ffade4ef02..5ad26bbd0c 100644
--- a/Userland/Libraries/LibVideo/VP9/BitStream.cpp
+++ b/Userland/Libraries/LibVideo/VP9/BitStream.cpp
@@ -71,13 +71,13 @@ ErrorOr<u16> BitStream::read_f16()
 /* 9.2.1 */
 ErrorOr<void> BitStream::init_bool(size_t bytes)
 {
-    if (bytes == 0)
-        return Error::from_string_literal("Range coder size cannot be zero.");
+    if (bytes > bytes_remaining())
+        return Error::from_string_literal("Available data is too small for range decoder");
     m_bool_value = TRY(read_f8());
     m_bool_range = 255;
     m_bool_max_bits = (8 * bytes) - 8;
     if (TRY(read_bool(128)))
-        return Error::from_string_literal("Range coder's first bool was non-zero.");
+        return Error::from_string_literal("Range decoder marker was non-zero");
     return {};
 }
 
@@ -114,17 +114,26 @@ ErrorOr<bool> BitStream::read_bool(u8 probability)
 /* 9.2.3 */
 ErrorOr<void> BitStream::exit_bool()
 {
-    // FIXME: I'm not sure if this call to min is spec compliant, or if there is an issue elsewhere earlier in the parser.
-    auto padding_element = TRY(read_bits(min(m_bool_max_bits, (u64)bits_remaining())));
+    while (m_bool_max_bits > 0) {
+        auto padding_read_size = min(m_bool_max_bits, sizeof(m_reservoir) * 8);
+        auto padding_bits = TRY(read_bits(padding_read_size));
+        m_bool_max_bits -= padding_read_size;
+
+        if (padding_bits != 0)
+            return Error::from_string_literal("Range decoder has non-zero padding element");
+    }
 
     // FIXME: It is a requirement of bitstream conformance that enough padding bits are inserted to ensure that the final coded byte of a frame is not equal to a superframe marker.
     //  A byte b is equal to a superframe marker if and only if (b & 0xe0)is equal to 0xc0, i.e. if the most significant 3 bits are equal to 0b110.
-    if (padding_element != 0)
-        return Error::from_string_literal("Range coder padding was non-zero.");
     return {};
 }
 
-ErrorOr<u8> BitStream::read_literal(size_t n)
+size_t BitStream::range_coding_bits_remaining()
+{
+    return m_bool_max_bits;
+}
+
+ErrorOr<u8> BitStream::read_literal(u8 n)
 {
     u8 return_value = 0;
     for (size_t i = 0; i < n; i++) {
diff --git a/Userland/Libraries/LibVideo/VP9/BitStream.h b/Userland/Libraries/LibVideo/VP9/BitStream.h
index 0085d974b3..7c51c315e1 100644
--- a/Userland/Libraries/LibVideo/VP9/BitStream.h
+++ b/Userland/Libraries/LibVideo/VP9/BitStream.h
@@ -32,7 +32,8 @@ public:
     ErrorOr<void> init_bool(size_t bytes);
     ErrorOr<bool> read_bool(u8 probability);
     ErrorOr<void> exit_bool();
-    ErrorOr<u8> read_literal(size_t n);
+    ErrorOr<u8> read_literal(u8 bit_count);
+    size_t range_coding_bits_remaining();
 
     /* (4.9.2) */
     ErrorOr<i8> read_s(size_t n);