diff --git a/AK/ConstrainedStream.cpp b/AK/ConstrainedStream.cpp index 5de5c4d7fcd..8645b5779ca 100644 --- a/AK/ConstrainedStream.cpp +++ b/AK/ConstrainedStream.cpp @@ -29,7 +29,7 @@ ErrorOr ConstrainedStream::read_some(Bytes bytes) ErrorOr ConstrainedStream::discard(size_t discarded_bytes) { - if (discarded_bytes >= m_limit) + 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. diff --git a/Tests/AK/CMakeLists.txt b/Tests/AK/CMakeLists.txt index b857096d24f..40914ac0b63 100644 --- a/Tests/AK/CMakeLists.txt +++ b/Tests/AK/CMakeLists.txt @@ -21,6 +21,7 @@ set(AK_TEST_SOURCES TestCircularDeque.cpp TestCircularQueue.cpp TestComplex.cpp + TestConstrainedStream.cpp TestCoroutine.cpp TestDisjointChunks.cpp TestDistinctNumeric.cpp diff --git a/Tests/AK/TestConstrainedStream.cpp b/Tests/AK/TestConstrainedStream.cpp new file mode 100644 index 00000000000..3c2e7ee923e --- /dev/null +++ b/Tests/AK/TestConstrainedStream.cpp @@ -0,0 +1,36 @@ +/* + * Copyright (c) 2025, Lucas Chollet + * + * SPDX-License-Identifier: BSD-2-Clause + */ + +#include + +#include +#include +#include + +static Array base_data { 1, 2, 3, 4 }; + +TEST_CASE(basic_constraint) +{ + auto memory_stream = make(base_data.span()); + ConstrainedStream constrained_stream { move(memory_stream), 2 }; + Array read; + TRY_OR_FAIL(constrained_stream.read_until_filled(read)); + EXPECT_EQ(read.span(), base_data.span().trim(2)); + + EXPECT(constrained_stream.read_until_filled(read).is_error()); +} + +TEST_CASE(discard_until_constraint) +{ + auto memory_stream = make(base_data.span()); + ConstrainedStream constrained_stream { move(memory_stream), 3 }; + Array read; + TRY_OR_FAIL(constrained_stream.read_until_filled(read)); + EXPECT_EQ(read.span(), base_data.span().trim(2)); + + EXPECT(!constrained_stream.discard(1).is_error()); + EXPECT(constrained_stream.discard(1).is_error()); +}