AK/LibAudio: Add stream read operators to AK::BufferStream, and use it in AWavLoader

At the same time, we allow chaining of streaming operators, and add a
way to check for partial reads (also used in WAV parsing).

2 files changed, 43 insertions, 82 deletions

diff --git a/Libraries/LibAudio/AWavLoader.cpp b/Libraries/LibAudio/AWavLoader.cpp
index 16a0a73607..ff1048f17b 100644
--- a/Libraries/LibAudio/AWavLoader.cpp
+++ b/Libraries/LibAudio/AWavLoader.cpp
@@ -1,40 +1,10 @@
 #include <LibCore/CFile.h>
+#include <AK/BufferStream.h>
+#include <limits>
 
 #include "AWavLoader.h"
 #include "AWavFile.h"
 
-static u32 read_u32(const ByteBuffer& buf, u32& off)
-{
-    ASSERT(buf.size() - off >= 4);
-    u32 b0 = buf[off + 0];
-    u32 b1 = buf[off + 1];
-    u32 b2 = buf[off + 2];
-    u32 b3 = buf[off + 3];
-
-    u32 ret = 0;
-    ret |= (u8(b3) << 24);
-    ret |= (u8(b2) << 16);
-    ret |= (u8(b1) << 8);
-    ret |= (u8(b0));
-
-    off += 4;
-    return ret;
-}
-
-static u16 read_u16(const ByteBuffer& buf, u32& off)
-{
-    ASSERT(buf.size() - off >= 2);
-    u16 b0 = buf[off + 0];
-    u16 b1 = buf[off + 1];
-
-    u16 ret = 0;
-    ret |= (u8(b1) << 8);
-    ret |= (u8(b0));
-
-    off += 2;
-    return ret;
-}
-
 RefPtr<AWavFile> AWavLoader::load_wav(const StringView& path)
 {
     m_error_string = {};
@@ -45,25 +15,22 @@ RefPtr<AWavFile> AWavLoader::load_wav(const StringView& path)
         return nullptr;
     }
 
-    const auto& contents = wav.read_all();
+    auto contents = wav.read_all();
     return parse_wav(contents);
 }
 
 // TODO: A streaming parser might be better than forcing a ByteBuffer
-RefPtr<AWavFile> AWavLoader::parse_wav(const ByteBuffer& buffer)
+RefPtr<AWavFile> AWavLoader::parse_wav(ByteBuffer& buffer)
 {
-    u32 off = 0;
-
-    if (buffer.size() - off < 12) {
-        dbgprintf("WAV is too small (no header, %d bytes)\n", buffer.size());
-        return {};
-    }
-
-    dbgprintf("Trying to parse %d bytes of wav\n", buffer.size());
+    BufferStream stream(buffer);
 
 #define CHECK_OK(msg)                      \
     do {                                   \
         ASSERT(ok);                        \
+        if (stream.handle_read_failure()) { \
+            m_error_string = String::format("Premature stream EOF at %s", msg); \
+            return {}; \
+        } \
         if (!ok) {                         \
             m_error_string = String::format("Parsing failed: %s", msg); \
             return {};                     \
@@ -73,94 +40,88 @@ RefPtr<AWavFile> AWavLoader::parse_wav(const ByteBuffer& buffer)
     } while (0);
 
     bool ok = true;
-    u32 riff = read_u32(buffer, off);
+    u32 riff; stream >> riff;
     ok = ok && riff == 0x46464952; // "RIFF"
     CHECK_OK("RIFF header");
 
-    u32 sz = read_u32(buffer, off);
-    ASSERT(sz < 1024 * 1024 * 42);
+    u32 sz; stream >> sz;
     ok = ok && sz < 1024 * 1024 * 42; // arbitrary
     CHECK_OK("File size");
+    ASSERT(sz < 1024 * 1024 * 42);
 
-    u32 wave = read_u32(buffer, off);
+    u32 wave; stream >> wave;
     ok = ok && wave == 0x45564157; // "WAVE"
     CHECK_OK("WAVE header");
 
-    if (buffer.size() - off < 8) {
-        dbgprintf("WAV is too small (no fmt, %d bytes)\n", buffer.size());
-        return {};
-    }
-
-    u32 fmt_id = read_u32(buffer, off);
+    u32 fmt_id; stream >> fmt_id;
     ok = ok && fmt_id == 0x20746D66; // "FMT"
     CHECK_OK("FMT header");
 
-    u32 fmt_size = read_u32(buffer, off);
+    u32 fmt_size; stream >> fmt_size;
     ok = ok && fmt_size == 16;
-    ASSERT(fmt_size == 16);
     CHECK_OK("FMT size");
-
-    if (buffer.size() - off < 16) {
-        dbgprintf("WAV is too small (fmt chunk, %d bytes)\n", buffer.size());
-        return {};
-    }
+    ASSERT(fmt_size == 16);
 
     auto ret = adopt(*new AWavFile);
-    u16 audio_format = read_u16(buffer, off);
+    u16 audio_format; stream >> audio_format;
+    CHECK_OK("Audio format"); // incomplete read check
     ok = ok && audio_format == 1; // WAVE_FORMAT_PCM
     ASSERT(audio_format == 1);
-    CHECK_OK("Audio format");
+    CHECK_OK("Audio format"); // value check
     ret->m_format = AWavFile::Format::PCM;
 
-    u16 num_channels = read_u16(buffer, off);
+    u16 num_channels; stream >> num_channels;
     CHECK_OK("Channel count");
     ret->m_channel_count = num_channels;
 
-    u32 sample_rate = read_u32(buffer, off);
+    u32 sample_rate; stream >> sample_rate;
     CHECK_OK("Sample rate");
     ret->m_sample_rate = sample_rate;
 
-    u32 byte_rate = read_u32(buffer, off);
+    u32 byte_rate; stream >> byte_rate;
     CHECK_OK("Byte rate");
     ret->m_byte_rate = byte_rate;
 
-    u16 block_align = read_u16(buffer, off);
+    u16 block_align; stream >> block_align;
     CHECK_OK("Block align");
     ret->m_block_align = block_align;
 
-    u16 bits_per_sample = read_u16(buffer, off);
+    u16 bits_per_sample; stream >> bits_per_sample;
+    CHECK_OK("Bits per sample"); // incomplete read check
     ok = ok && (bits_per_sample == 8 || bits_per_sample == 16);
     ASSERT(bits_per_sample == 8 || bits_per_sample == 16);
-    CHECK_OK("Bits per sample");
+    CHECK_OK("Bits per sample"); // value check
     ret->m_bits_per_sample = bits_per_sample;
 
     // Read chunks until we find DATA
-    if (off >= u32(buffer.size()) - 8) {
-        ok = ok && false;
-        ASSERT_NOT_REACHED();
-        CHECK_OK("Premature EOF without DATA");
-    }
-
     bool found_data = false;
     u32 data_sz = 0;
-    while (off < u32(buffer.size()) - 8) {
-        u32 chunk_id = read_u32(buffer, off);
-        data_sz = read_u32(buffer, off);
+    while (true) {
+        u32 chunk_id; stream >> chunk_id;
+        CHECK_OK("Reading chunk ID searching for data");
+        stream >> data_sz;
+        CHECK_OK("Reading chunk size searching for data");
         if (chunk_id == 0x61746164) { // DATA
             found_data = true;
             break;
         }
-        off += data_sz;
     }
 
     ok = ok && found_data;
-    ASSERT(found_data);
     CHECK_OK("Found no data chunk");
+    ASSERT(found_data);
+
+    ok = ok && data_sz < std::numeric_limits<int>::max();
+    CHECK_OK("Data was too large");
+
+    // ### consider using BufferStream to do this for us
+    ok = ok && int(data_sz) <= (buffer.size() - stream.offset());
+    CHECK_OK("Bad DATA (truncated)");
 
-    ok = ok && data_sz <= (buffer.size() - off);
-    CHECK_OK("Bad DATA size");
+    ret->m_sample_data = buffer.slice(stream.offset(), data_sz);
 
-    ret->m_sample_data = buffer.slice(off, data_sz);
+    // At this point there should be no read failures!
+    ASSERT(!stream.handle_read_failure());
     return ret;
 }
 
diff --git a/Libraries/LibAudio/AWavLoader.h b/Libraries/LibAudio/AWavLoader.h
index be4facab45..72dee9f347 100644
--- a/Libraries/LibAudio/AWavLoader.h
+++ b/Libraries/LibAudio/AWavLoader.h
@@ -9,6 +9,6 @@ public:
     RefPtr<AWavFile> load_wav(const StringView& path);
     const char* error_string() { return m_error_string.characters(); }
 private:
-    RefPtr<AWavFile> parse_wav(const ByteBuffer& buffer);
+    RefPtr<AWavFile> parse_wav(ByteBuffer& buffer);
     String m_error_string;
 };