Use std::stoi to detect invalid cases

zanmato1984 · Apr 22, 2024 · 5458b17 · 5458b17
1 parent 5f6bb7c
commit 5458b17
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 2 deletions.
diff --git a/cpp/src/gandiva/cache.cc b/cpp/src/gandiva/cache.cc
@@ -35,8 +35,15 @@ int GetCapacityFromEnvVar() {
   if (env_value.empty()) {
     return DEFAULT_CACHE_SIZE;
   }
-  int capacity = std::atoi(env_value.c_str());
-  if (capacity <= 0) {
+  int capacity = 0;
+  size_t length = 0;
+  bool exception = false;
+  try {
+    capacity = std::stoi(env_value.c_str(), &length);
+  } catch (const std::exception&) {
+    exception = true;
+  }
+  if (length != env_value.length() || exception || capacity <= 0) {
     ARROW_LOG(WARNING) << "Invalid cache size provided in GANDIVA_CACHE_SIZE. "
                        << "Using default cache size: " << DEFAULT_CACHE_SIZE;
     return DEFAULT_CACHE_SIZE;

diff --git a/cpp/src/gandiva/cache_test.cc b/cpp/src/gandiva/cache_test.cc
@@ -76,6 +76,12 @@ TEST(TestCache, TestGetCacheCapacityEnvVar) {
     ASSERT_EQ(internal::GetCapacityFromEnvVar(), default_capacity);
   }
 
+  // Number with invalid suffix.
+  {
+    ScopedEnvVar env(capacity_env_var, "42MB");
+    ASSERT_EQ(internal::GetCapacityFromEnvVar(), default_capacity);
+  }
+
   // Valid positive number.
   {
     ScopedEnvVar env(capacity_env_var, "42");
@@ -100,6 +106,13 @@ TEST(TestCache, TestGetCacheCapacityEnvVar) {
     ScopedEnvVar env(capacity_env_var, "-1");
     ASSERT_EQ(internal::GetCapacityFromEnvVar(), default_capacity);
   }
+
+  // Over int max.
+  {
+    auto str = std::to_string(static_cast<int64_t>(std::numeric_limits<int>::max()) + 1);
+    ScopedEnvVar env(capacity_env_var, str.c_str());
+    ASSERT_EQ(internal::GetCapacityFromEnvVar(), default_capacity);
+  }
 }
 
 }  // namespace gandiva