diff --git a/AK/CMakeLists.txt b/AK/CMakeLists.txt index 5d6539ffbe0..964f672a8f4 100644 --- a/AK/CMakeLists.txt +++ b/AK/CMakeLists.txt @@ -2,6 +2,7 @@ set(AK_SOURCES Assertions.cpp Base64.cpp CircularBuffer.cpp + ConstrainedStream.cpp DOSPackedTime.cpp DeprecatedFlyString.cpp DeprecatedString.cpp diff --git a/AK/ConstrainedStream.cpp b/AK/ConstrainedStream.cpp new file mode 100644 index 00000000000..5de5c4d7fcd --- /dev/null +++ b/AK/ConstrainedStream.cpp @@ -0,0 +1,62 @@ +/* + * Copyright (c) 2021, Ali Mohammad Pur + * Copyright (c) 2023, Tim Schumacher + * + * SPDX-License-Identifier: BSD-2-Clause + */ + +#include + +namespace AK { + +ConstrainedStream::ConstrainedStream(MaybeOwned stream, u64 limit) + : m_stream(move(stream)) + , m_limit(limit) +{ +} + +ErrorOr 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 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 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() +{ +} + +} diff --git a/AK/ConstrainedStream.h b/AK/ConstrainedStream.h new file mode 100644 index 00000000000..087c32594ff --- /dev/null +++ b/AK/ConstrainedStream.h @@ -0,0 +1,31 @@ +/* + * Copyright (c) 2021, Ali Mohammad Pur + * Copyright (c) 2023, Tim Schumacher + * + * SPDX-License-Identifier: BSD-2-Clause + */ + +#pragma once + +#include +#include + +namespace AK { + +class ConstrainedStream : public Stream { +public: + ConstrainedStream(MaybeOwned, u64 limit); + + virtual ErrorOr read_some(Bytes) override; + virtual ErrorOr discard(size_t discarded_bytes) override; + virtual ErrorOr write_some(ReadonlyBytes) override; + virtual bool is_eof() const override; + virtual bool is_open() const override; + virtual void close() override; + +private: + MaybeOwned m_stream; + u64 m_limit; +}; + +} diff --git a/AK/Forward.h b/AK/Forward.h index be76e64db95..e9b7dde8457 100644 --- a/AK/Forward.h +++ b/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; diff --git a/Userland/Libraries/LibWasm/Parser/Parser.cpp b/Userland/Libraries/LibWasm/Parser/Parser.cpp index c2582ddfa8b..d0a5dadc5d0 100644 --- a/Userland/Libraries/LibWasm/Parser/Parser.cpp +++ b/Userland/Libraries/LibWasm/Parser/Parser.cpp @@ -4,6 +4,7 @@ * SPDX-License-Identifier: BSD-2-Clause */ +#include #include #include #include @@ -1186,7 +1187,7 @@ ParseResult 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), size }; auto func = Func::parse(constrained_stream); if (func.is_error()) @@ -1301,7 +1302,7 @@ ParseResult 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), section_size }; switch (section_id) { case CustomSection::section_id: diff --git a/Userland/Libraries/LibWasm/Types.h b/Userland/Libraries/LibWasm/Types.h index 7d5b3303139..e653b71ea13 100644 --- a/Userland/Libraries/LibWasm/Types.h +++ b/Userland/Libraries/LibWasm/Types.h @@ -135,55 +135,6 @@ private: Vector m_buffer; }; -class ConstrainedStream : public Stream { -public: - explicit ConstrainedStream(Stream& stream, size_t size) - : m_stream(stream) - , m_bytes_left(size) - { - } - -private: - ErrorOr 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 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 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: