From 1d51d9e8b63e673cc0502a50f5e8a1c8fa8fc00c Mon Sep 17 00:00:00 2001
From: Alexey Milovidov <milovidov@yandex-team.ru>
Date: Sat, 15 Aug 2020 11:49:54 +0300
Subject: [PATCH] Range check for h3KRing

---
 src/Functions/h3kRing.cpp                            | 12 ++++++++++++
 .../0_stateless/01442_h3kring_range_check.reference  |  1 +
 .../0_stateless/01442_h3kring_range_check.sql        |  4 ++++
 3 files changed, 17 insertions(+)
 create mode 100644 tests/queries/0_stateless/01442_h3kring_range_check.reference
 create mode 100644 tests/queries/0_stateless/01442_h3kring_range_check.sql

diff --git a/src/Functions/h3kRing.cpp b/src/Functions/h3kRing.cpp
index 0fb9c4b6a64..9702edf7079 100644
--- a/src/Functions/h3kRing.cpp
+++ b/src/Functions/h3kRing.cpp
@@ -14,10 +14,13 @@
 
 namespace DB
 {
+
 namespace ErrorCodes
 {
     extern const int ILLEGAL_TYPE_OF_ARGUMENT;
+    extern const int PARAMETER_OUT_OF_BOUND;
 }
+
 class FunctionH3KRing : public IFunction
 {
 public:
@@ -65,6 +68,15 @@ public:
             const H3Index origin_hindex = col_hindex->getUInt(row);
             const int k = col_k->getInt(row);
 
+            /// Overflow is possible. The function maxKringSize does not check for overflow.
+            /// The calculation is similar to square of k but several times more.
+            /// Let's use huge underestimation as the safe bound. We should not allow to generate too large arrays nevertheless.
+            constexpr auto max_k = 10000;
+            if (k > max_k)
+                throw Exception(ErrorCodes::PARAMETER_OUT_OF_BOUND, "Too large 'k' argument for {} function, maximum {}", getName(), max_k);
+            if (k < 0)
+                throw Exception(ErrorCodes::PARAMETER_OUT_OF_BOUND, "Argument 'k' for {} function must be non negative", getName());
+
             const auto vec_size = maxKringSize(k);
             hindex_vec.resize(vec_size);
             kRing(origin_hindex, k, hindex_vec.data());
diff --git a/tests/queries/0_stateless/01442_h3kring_range_check.reference b/tests/queries/0_stateless/01442_h3kring_range_check.reference
new file mode 100644
index 00000000000..9f54fe3133b
--- /dev/null
+++ b/tests/queries/0_stateless/01442_h3kring_range_check.reference
@@ -0,0 +1 @@
+122
diff --git a/tests/queries/0_stateless/01442_h3kring_range_check.sql b/tests/queries/0_stateless/01442_h3kring_range_check.sql
new file mode 100644
index 00000000000..e580beb17a3
--- /dev/null
+++ b/tests/queries/0_stateless/01442_h3kring_range_check.sql
@@ -0,0 +1,4 @@
+SELECT h3kRing(581276613233082367, 65535); -- { serverError 12 }
+SELECT h3kRing(581276613233082367, -1); -- { serverError 12 }
+SELECT length(h3kRing(111111111111, 1000));
+SELECT h3kRing(581276613233082367, nan); -- { serverError 43 }