From 19164badf9d8167f7cde13f52542f4070eca1fdc Mon Sep 17 00:00:00 2001 From: Roman Gershman Date: Tue, 17 Dec 2024 13:26:33 +0200 Subject: [PATCH] fix: potential OOM when first request sent in small bits (#4325) Before: if socket data arrived in small bits, then CheckForHttpProto would grow io_buf_ capacity exponentially with each iteration. For example, test_match_http test easily causes OOM. This PR ensures that there is always a buffer available - but it grows linearly with the input size. Currently, the total input in CheckForHttpProto is limited to 1024. Signed-off-by: Roman Gershman --- src/facade/dragonfly_connection.cc | 7 ++----- tests/dragonfly/connection_test.py | 8 ++++++++ 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/src/facade/dragonfly_connection.cc b/src/facade/dragonfly_connection.cc index e43c78f8c..39725001d 100644 --- a/src/facade/dragonfly_connection.cc +++ b/src/facade/dragonfly_connection.cc @@ -928,7 +928,7 @@ io::Result Connection::CheckForHttpProto() { return MatchHttp11Line(ib); } last_len = io_buf_.InputLen(); - UpdateIoBufCapacity(io_buf_, stats_, [&]() { io_buf_.EnsureCapacity(io_buf_.Capacity()); }); + UpdateIoBufCapacity(io_buf_, stats_, [&]() { io_buf_.EnsureCapacity(128); }); } while (last_len < 1024); return false; @@ -959,10 +959,7 @@ void Connection::ConnectionFlow() { // Main loop. if (parse_status != ERROR && !ec) { - if (io_buf_.AppendLen() < 64) { - UpdateIoBufCapacity(io_buf_, stats_, - [&]() { io_buf_.EnsureCapacity(io_buf_.Capacity() * 2); }); - } + UpdateIoBufCapacity(io_buf_, stats_, [&]() { io_buf_.EnsureCapacity(64); }); auto res = IoLoop(); if (holds_alternative(res)) { diff --git a/tests/dragonfly/connection_test.py b/tests/dragonfly/connection_test.py index a59c4a39c..cdd10c958 100755 --- a/tests/dragonfly/connection_test.py +++ b/tests/dragonfly/connection_test.py @@ -606,6 +606,14 @@ async def test_subscribe_in_pipeline(async_client: aioredis.Redis): assert res == ["one", ["subscribe", "ch1", 1], "two", ["subscribe", "ch2", 2], "three"] +async def test_match_http(df_server: DflyInstance): + client = df_server.client() + reader, writer = await asyncio.open_connection("localhost", df_server.port) + for i in range(2000): + writer.write(f"foo bar ".encode()) + await writer.drain() + + """ This test makes sure that Dragonfly can receive blocks of pipelined commands even while a script is still executing. This is a dangerous scenario because both the dispatch fiber