Backport #85440 to 25.5: Fix possible UB in case of MEMORY_LIMIT_EXCEEDED during String deserialization

robot-clickhouse · robot-clickhouse · commit 488c14188964 · 2025-08-15T09:12:10.000Z
diff --git a/src/DataTypes/Serializations/SerializationString.cpp b/src/DataTypes/Serializations/SerializationString.cpp
@@ -149,6 +149,7 @@ void SerializationString::serializeBinaryBulk(const IColumn & column, WriteBuffe
 
 template <int UNROLL_TIMES>
 static NO_INLINE void deserializeBinarySSE2(ColumnString::Chars & data, ColumnString::Offsets & offsets, ReadBuffer & istr, size_t limit)
+try
 {
     size_t offset = data.size();
     /// Avoiding calling resize in a loop improves the performance.
@@ -171,8 +172,6 @@ static NO_INLINE void deserializeBinarySSE2(ColumnString::Chars & data, ColumnSt
                 max_string_size);
 
         offset += size + 1;
-        offsets.push_back(offset);
-
         if (unlikely(offset > data.size()))
             data.resize_exact(roundUpToPowerOfTwoOrZero(std::max(offset, data.size() * 2)));
 
@@ -205,10 +204,18 @@ static NO_INLINE void deserializeBinarySSE2(ColumnString::Chars & data, ColumnSt
         }
 
         data[offset - 1] = 0;
+
+        offsets.push_back(offset);
     }
 
     data.resize_exact(offset);
 }
+catch (...)
+{
+    /// We are doing resize_exact() of bigger values than we have, let's make sure that it will be correct (even in case of exceptions)
+    data.resize_exact(offsets.back());
+    throw;
+}
 
 
 void SerializationString::deserializeBinaryBulk(IColumn & column, ReadBuffer & istr, size_t rows_offset, size_t limit, double avg_value_size_hint) const
diff --git a/src/DataTypes/Serializations/tests/gtest_string_serialization.cpp b/src/DataTypes/Serializations/tests/gtest_string_serialization.cpp
@@ -0,0 +1,82 @@
+#include <Columns/ColumnString.h>
+#include <DataTypes/DataTypeFactory.h>
+#include <DataTypes/Serializations/SerializationString.h>
+#include <IO/ReadBufferFromString.h>
+#include <Common/MemoryTracker.h>
+#include <Common/ThreadStatus.h>
+
+#include <gtest/gtest.h>
+
+namespace DB
+{
+    namespace ErrorCodes
+    {
+        extern const int MEMORY_LIMIT_EXCEEDED;
+    }
+}
+
+using namespace DB;
+
+TEST(StringSerialization, IncorrectStateAfterMemoryLimitExceeded)
+{
+    MainThreadStatus::getInstance();
+
+    constexpr size_t rows = 1'000'000;
+
+    WriteBufferFromOwnString out;
+
+    auto src_column = ColumnString::create();
+    src_column->insertMany("foobar", rows);
+
+    {
+        auto serialization = std::make_shared<SerializationString>();
+        ISerialization::SerializeBinaryBulkSettings settings;
+        ISerialization::SerializeBinaryBulkStatePtr state;
+        settings.position_independent_encoding = false;
+        settings.getter = [&out](const auto &) { return &out; };
+        serialization->serializeBinaryBulkWithMultipleStreams(*src_column, 0, src_column->size(), settings, state);
+    }
+
+    size_t memory_limit_exceeded_errors = 0;
+    auto run_with_memory_failures = [&](auto cb)
+    {
+        total_memory_tracker.setFaultProbability(0.2);
+        try
+        {
+            cb();
+        }
+        catch (Exception & e)
+        {
+            if (e.code() != ErrorCodes::MEMORY_LIMIT_EXCEEDED)
+                throw;
+
+            ++memory_limit_exceeded_errors;
+            total_memory_tracker.setFaultProbability(0);
+        }
+        total_memory_tracker.setFaultProbability(0);
+    };
+
+    auto type_string = std::make_shared<DataTypeString>();
+    size_t non_empty_result = 0;
+    while (memory_limit_exceeded_errors < 10 || non_empty_result < 10)
+    {
+        ColumnPtr result_column = type_string->createColumn();
+        ReadBufferFromOwnString in(out.str());
+
+        auto serialization = type_string->getDefaultSerialization();
+        ISerialization::DeserializeBinaryBulkSettings settings;
+        ISerialization::DeserializeBinaryBulkStatePtr state;
+        settings.position_independent_encoding = false;
+        settings.getter = [&in](const auto &) { return &in; };
+
+        run_with_memory_failures([&]() { serialization->deserializeBinaryBulkWithMultipleStreams(result_column, 0, src_column->size(), settings, state, nullptr); });
+
+        auto & result = assert_cast<ColumnString &>(*result_column->assumeMutable());
+        if (!result.empty())
+        {
+            ++non_empty_result;
+            ASSERT_EQ(result.getDataAt(0), "foobar");
+            ASSERT_EQ(result.getDataAt(result.size() - 1), "foobar");
+        }
+    }
+}
diff --git a/src/Interpreters/tests/gtest_filecache.cpp b/src/Interpreters/tests/gtest_filecache.cpp
@@ -9,6 +9,7 @@
 #include <thread>
 
 #include <Core/ServerUUID.h>
+#include <Common/ThreadStatus.h>
 #include <Common/iota.h>
 #include <Common/randomSeed.h>
 #include <DataTypes/DataTypesNumber.h>
@@ -304,7 +305,11 @@ void increasePriority(const HolderPtr & holder, size_t pos)
 class FileCacheTest : public ::testing::Test
 {
 public:
-    FileCacheTest() {
+    FileCacheTest()
+    {
+        /// Reset current_thread to avoid conflicts of ThreadStatus with MainThreadStatus
+        current_thread = nullptr;
+
         /// Context has to be created before calling cache.initialize();
         /// Otherwise the tests which run before FileCacheTest.get are failed
         /// It is logical to call destroyContext() at destructor.
@@ -313,6 +318,12 @@ class FileCacheTest : public ::testing::Test
         getContext();
     }
 
+    ~FileCacheTest() override
+    {
+        /// Reset current_thread back
+        current_thread = MainThreadStatus::get();
+    }
+
     static void setupLogs(const std::string & level)
     {
         Poco::AutoPtr<Poco::ConsoleChannel> channel(new Poco::ConsoleChannel(std::cerr));
