AK: Move ConstrainedStream from LibWasm and limit discarding

AK/CMakeLists.txt

@@ -2,6 +2,7 @@ set(AK_SOURCES
     Assertions.cpp
     Base64.cpp
     CircularBuffer.cpp
+    ConstrainedStream.cpp
     DOSPackedTime.cpp
     DeprecatedFlyString.cpp
     DeprecatedString.cpp

AK/ConstrainedStream.cpp

@@ -0,0 +1,62 @@
+/*
+ * Copyright (c) 2021, Ali Mohammad Pur <mpfard@serenityos.org>
+ * Copyright (c) 2023, Tim Schumacher <timschumi@gmx.de>
+ *
+ * SPDX-License-Identifier: BSD-2-Clause
+ */
+
+#include <AK/ConstrainedStream.h>
+
+namespace AK {
+
+ConstrainedStream::ConstrainedStream(MaybeOwned<Stream> stream, u64 limit)
+    : m_stream(move(stream))
+    , m_limit(limit)
+{
+}
+
+ErrorOr<Bytes> ConstrainedStream::read_some(Bytes bytes)
+{
+    if (bytes.size() >= m_limit)
+        bytes = bytes.trim(m_limit);
+
+    auto result = TRY(m_stream->read_some(bytes));
+
+    m_limit -= result.size();
+
+    return result;
+}
+
+ErrorOr<void> ConstrainedStream::discard(size_t discarded_bytes)
+{
+    if (discarded_bytes >= m_limit)
+        return Error::from_string_literal("Trying to discard more bytes than allowed");
+
+    // Note: This will remove more bytes from the limit than intended if a failing discard yields partial results.
+    //       However, stopping early is most likely better than running over the end of the stream.
+    m_limit -= discarded_bytes;
+    TRY(m_stream->discard(discarded_bytes));
+
+    return {};
+}
+
+ErrorOr<size_t> ConstrainedStream::write_some(ReadonlyBytes)
+{
+    return Error::from_errno(EBADF);
+}
+
+bool ConstrainedStream::is_eof() const
+{
+    return m_limit == 0 || m_stream->is_eof();
+}
+
+bool ConstrainedStream::is_open() const
+{
+    return m_stream->is_open();
+}
+
+void ConstrainedStream::close()
+{
+}
+
+}

AK/ConstrainedStream.h

@@ -0,0 +1,31 @@
+/*
+ * Copyright (c) 2021, Ali Mohammad Pur <mpfard@serenityos.org>
+ * Copyright (c) 2023, Tim Schumacher <timschumi@gmx.de>
+ *
+ * SPDX-License-Identifier: BSD-2-Clause
+ */
+
+#pragma once
+
+#include <AK/MaybeOwned.h>
+#include <AK/Stream.h>
+
+namespace AK {
+
+class ConstrainedStream : public Stream {
+public:
+    ConstrainedStream(MaybeOwned<Stream>, u64 limit);
+
+    virtual ErrorOr<Bytes> read_some(Bytes) override;
+    virtual ErrorOr<void> discard(size_t discarded_bytes) override;
+    virtual ErrorOr<size_t> write_some(ReadonlyBytes) override;
+    virtual bool is_eof() const override;
+    virtual bool is_open() const override;
+    virtual void close() override;
+
+private:
+    MaybeOwned<Stream> m_stream;
+    u64 m_limit;
+};
+
+}

AK/Forward.h

@@ -22,6 +22,7 @@ class BigEndianOutputBitStream;
 class Bitmap;
 using ByteBuffer = Detail::ByteBuffer<32>;
 class CircularBuffer;
+class ConstrainedStream;
 class DeprecatedFlyString;
 class DeprecatedString;
 class DeprecatedStringCodePointIterator;
@@ -155,6 +156,7 @@ using AK::ByteBuffer;
 using AK::Bytes;
 using AK::CircularBuffer;
 using AK::CircularQueue;
+using AK::ConstrainedStream;
 using AK::DeprecatedFlyString;
 using AK::DeprecatedString;
 using AK::DeprecatedStringCodePointIterator;

Userland/Libraries/LibWasm/Parser/Parser.cpp

@@ -4,6 +4,7 @@
  * SPDX-License-Identifier: BSD-2-Clause
  */
 
+#include <AK/ConstrainedStream.h>
 #include <AK/Debug.h>
 #include <AK/Endian.h>
 #include <AK/LEB128.h>
@@ -1186,7 +1187,7 @@ ParseResult<CodeSection::Code> CodeSection::Code::parse(Stream& stream)
         return with_eof_check(stream, ParseError::InvalidSize);
     size_t size = size_or_error.release_value();
 
-    auto constrained_stream = ConstrainedStream { stream, size };
+    auto constrained_stream = ConstrainedStream { MaybeOwned<Stream>(stream), size };
 
     auto func = Func::parse(constrained_stream);
     if (func.is_error())
@@ -1301,7 +1302,7 @@ ParseResult<Module> Module::parse(Stream& stream)
             return with_eof_check(stream, ParseError::ExpectedSize);
         size_t section_size = section_size_or_error.release_value();
 
-        auto section_stream = ConstrainedStream { stream, section_size };
+        auto section_stream = ConstrainedStream { MaybeOwned<Stream>(stream), section_size };
 
         switch (section_id) {
         case CustomSection::section_id:

Userland/Libraries/LibWasm/Types.h

@@ -135,55 +135,6 @@ private:
     Vector<u8, 8> m_buffer;
 };
 
-class ConstrainedStream : public Stream {
-public:
-    explicit ConstrainedStream(Stream& stream, size_t size)
-        : m_stream(stream)
-        , m_bytes_left(size)
-    {
-    }
-
-private:
-    ErrorOr<Bytes> read_some(Bytes bytes) override
-    {
-        auto to_read = min(m_bytes_left, bytes.size());
-        auto read_bytes = TRY(m_stream.read_some(bytes.slice(0, to_read)));
-        m_bytes_left -= read_bytes.size();
-        return read_bytes;
-    }
-
-    bool is_eof() const override
-    {
-        return m_bytes_left == 0 || m_stream.is_eof();
-    }
-
-    ErrorOr<void> discard(size_t count) override
-    {
-        if (count > m_bytes_left)
-            return Error::from_string_literal("Trying to discard more bytes than allowed");
-
-        return m_stream.discard(count);
-    }
-
-    virtual ErrorOr<size_t> write_some(ReadonlyBytes) override
-    {
-        return Error::from_errno(EBADF);
-    }
-
-    virtual bool is_open() const override
-    {
-        return m_stream.is_open();
-    }
-
-    virtual void close() override
-    {
-        m_stream.close();
-    }
-
-    Stream& m_stream;
-    size_t m_bytes_left { 0 };
-};
-
 // https://webassembly.github.io/spec/core/bikeshed/#value-types%E2%91%A2
 class ValueType {
 public: