AK: Allow ConstrainedStream to discard bytes until the limit

AK/ConstrainedStream.cpp

@@ -29,7 +29,7 @@ ErrorOr<Bytes> ConstrainedStream::read_some(Bytes bytes)
 
 ErrorOr<void> 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.

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

Tests/AK/TestConstrainedStream.cpp

@@ -0,0 +1,36 @@
+/*
+ * Copyright (c) 2025, Lucas Chollet <[email protected]>
+ *
+ * SPDX-License-Identifier: BSD-2-Clause
+ */
+
+#include <LibTest/TestCase.h>
+
+#include <AK/Array.h>
+#include <AK/ConstrainedStream.h>
+#include <AK/MemoryStream.h>
+
+static Array<u8, 4> base_data { 1, 2, 3, 4 };
+
+TEST_CASE(basic_constraint)
+{
+    auto memory_stream = make<FixedMemoryStream>(base_data.span());
+    ConstrainedStream constrained_stream { move(memory_stream), 2 };
+    Array<u8, 2> 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<FixedMemoryStream>(base_data.span());
+    ConstrainedStream constrained_stream { move(memory_stream), 3 };
+    Array<u8, 2> 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());
+}