diff --git a/AK/MemoryStream.cpp b/AK/MemoryStream.cpp index c0d2be82c4..19448bfbbe 100644 --- a/AK/MemoryStream.cpp +++ b/AK/MemoryStream.cpp @@ -210,8 +210,16 @@ ErrorOr> AllocatingMemoryStream::offset_of(ReadonlyBytes needle search_spans[i] = m_chunks[i].span(); } + auto used_size_of_last_chunk = m_write_offset % CHUNK_SIZE; + + // The case where the stored write offset is actually the used space is the only case where a result of zero + // actually is zero. In other cases (i.e. our write offset is beyond the size of a chunk) the write offset + // already points to the beginning of the next chunk, in that case a result of zero indicates "use the last chunk in full". + if (m_write_offset >= CHUNK_SIZE && used_size_of_last_chunk == 0) + used_size_of_last_chunk = CHUNK_SIZE; + // Trimming is done first to ensure that we don't unintentionally shift around if the first and last chunks are the same. - search_spans[chunk_count - 1] = search_spans[chunk_count - 1].trim(m_write_offset % CHUNK_SIZE); + search_spans[chunk_count - 1] = search_spans[chunk_count - 1].trim(used_size_of_last_chunk); search_spans[0] = search_spans[0].slice(m_read_offset); return AK::memmem(search_spans.begin(), search_spans.end(), needle); diff --git a/Tests/AK/TestMemoryStream.cpp b/Tests/AK/TestMemoryStream.cpp index 073be8fd32..8d7e5ae42f 100644 --- a/Tests/AK/TestMemoryStream.cpp +++ b/Tests/AK/TestMemoryStream.cpp @@ -123,3 +123,31 @@ TEST_CASE(allocating_memory_stream_offset_of_after_chunk_reorder) EXPECT_EQ(offset.value(), 15ul); } } + +TEST_CASE(allocating_memory_stream_offset_of_with_write_offset_multiple_of_chunk_size) +{ + // This tests a specific edge case where we would erroneously trim the last searched block + // to size 0 if the current write offset is a multiple of the chunk size. + + AllocatingMemoryStream stream; + + // First, fill exactly one chunk (in groups of 16 bytes). + for (size_t i = 0; i < (AllocatingMemoryStream::CHUNK_SIZE / 16) - 1; ++i) + MUST(stream.write_until_depleted("AAAAAAAAAAAAAAAA"sv.bytes())); + MUST(stream.write_until_depleted("BCDEFGHIJKLMNOPQ"sv.bytes())); + + // Read a few bytes from the beginning to ensure that we are trying to slice into the zero-sized block. + MUST(stream.discard(32)); + + { + auto offset = MUST(stream.offset_of("B"sv.bytes())); + EXPECT(offset.has_value()); + EXPECT_EQ(offset.value(), AllocatingMemoryStream::CHUNK_SIZE - 32 - 16); + } + + { + auto offset = MUST(stream.offset_of("Q"sv.bytes())); + EXPECT(offset.has_value()); + EXPECT_EQ(offset.value(), AllocatingMemoryStream::CHUNK_SIZE - 32 - 1); + } +}