Check for array size overflow in topK #14452

src/AggregateFunctions/AggregateFunctionTopK.cpp

@@ -85,12 +85,12 @@ AggregateFunctionPtr createAggregateFunctionTopK(const std::string & name, const
             load_factor = applyVisitor(FieldVisitorConvertToNumber<UInt64>(), params[1]);
 
             if (load_factor < 1)
-                throw Exception("Too small parameter for aggregate function " + name + ". Minimum: 1",
+                throw Exception("Too small parameter 'load_factor' for aggregate function " + name + ". Minimum: 1",
                     ErrorCodes::ARGUMENT_OUT_OF_BOUND);
         }
 
-        if (k > TOP_K_MAX_SIZE)
-            throw Exception("Too large parameter for aggregate function " + name + ". Maximum: " + toString(TOP_K_MAX_SIZE),
+        if (k > TOP_K_MAX_SIZE || load_factor > TOP_K_MAX_SIZE || k * load_factor > TOP_K_MAX_SIZE)
+            throw Exception("Too large parameter(s) for aggregate function " + name + ". Maximum: " + toString(TOP_K_MAX_SIZE),
                 ErrorCodes::ARGUMENT_OUT_OF_BOUND);
 
         if (k == 0)

src/Common/SpaceSaving.h

@@ -147,16 +147,17 @@ public:
     {
         // Increase weight of a key that already exists
         auto hash = counter_map.hash(key);
-        auto counter = findCounter(key, hash);
-        if (counter)
+
+        if (auto counter = findCounter(key, hash); counter)
         {
             counter->count += increment;
             counter->error += error;
             percolate(counter);
             return;
         }
+
         // Key doesn't exist, but can fit in the top K
-        else if (unlikely(size() < capacity()))
+        if (unlikely(size() < capacity()))
         {
             auto c = new Counter(arena.emplace(key), increment, error, hash);
             push(c);

tests/queries/0_stateless/01471_top_k_range_check.sql

+SELECT length(topKWeighted(2, -9223372036854775808)(number, 1025)) FROM system.numbers; -- { serverError 69 }