ST_PointN returns a wrong result or causes a server crash (#7580)

* Add test * Add index bound checking logic * Address comment * Address comments Signed-off-by: Misiu Godfrey <misiu.godfrey@kraken.mapd.com>
SINHASantos · Aug 26, 2024 · 255f002 · 255f002
1 parent b2158a7
commit 255f002
Show file tree

Hide file tree

Showing 3 changed files with 83 additions and 23 deletions.
diff --git a/QueryEngine/GeoOperators/PointN.h b/QueryEngine/GeoOperators/PointN.h
@@ -51,6 +51,24 @@ class PointN : public Codegen {
     return SQLTypeInfo(kBOOLEAN);
   }
 
+  llvm::Value* codegenIndexOutOfBoundCheck(CgenState* cgen_state,
+                                           llvm::Value* index_lv,
+                                           llvm::Value* geosize_lv) {
+    llvm::Value* is_null_lv = cgen_state->llBool(false);
+    is_null_lv = cgen_state->ir_builder_.CreateOr(
+        is_null_lv,
+        cgen_state->ir_builder_.CreateICmp(llvm::ICmpInst::ICMP_SLT,
+                                           index_lv,
+                                           cgen_state->llInt(static_cast<int32_t>(0))));
+    return cgen_state->ir_builder_.CreateOr(
+        is_null_lv,
+        cgen_state->ir_builder_.CreateICmp(
+            llvm::ICmpInst::ICMP_SGE,
+            cgen_state->ir_builder_.CreateMul(index_lv,
+                                              cgen_state->llInt(static_cast<int32_t>(8))),
+            geosize_lv));
+  }
+
   // returns arguments lvs and null lv
   std::tuple<std::vector<llvm::Value*>, llvm::Value*> codegenLoads(
       const std::vector<llvm::Value*>& arg_lvs,
@@ -70,6 +88,7 @@ class PointN : public Codegen {
     auto index_lv = builder.CreateMul(
         builder.CreateSub(arg_lvs.back(), cgen_state->llInt(static_cast<int32_t>(1))),
         cgen_state->llInt(static_cast<int32_t>(2)));
+
     llvm::Value* is_null_lv{nullptr};
     if (arg_lvs.size() == 2) {
       // col byte stream from column on disk
@@ -95,10 +114,9 @@ class PointN : public Codegen {
 
       auto geo_size_lv = array_operand_lvs.back();
       // convert the index to a byte index
-      const auto outside_linestring_bounds_lv = builder.CreateNot(builder.CreateICmp(
-          llvm::ICmpInst::ICMP_SLT,
-          builder.CreateMul(index_lv, cgen_state->llInt(static_cast<int32_t>(8))),
-          geo_size_lv));
+      index_lv = builder.CreateMul(index_lv, cgen_state->llInt(static_cast<int32_t>(8)));
+      const auto outside_linestring_bounds_lv =
+          codegenIndexOutOfBoundCheck(cgen_state, index_lv, geo_size_lv);
       outside_linestring_bounds_lv->setName("outside_linestring_bounds");
       const auto input_is_null_lv = builder.CreateICmp(
           llvm::ICmpInst::ICMP_EQ,
@@ -114,8 +132,7 @@ class PointN : public Codegen {
       const auto geo_size_lv = arg_lvs[1];
       // TODO: bounds indices are 64 bits but should be 32 bits, as array length is
       // limited to 32 bits
-      is_null_lv = builder.CreateNot(
-          builder.CreateICmp(llvm::ICmpInst::ICMP_SLT, index_lv, geo_size_lv));
+      is_null_lv = codegenIndexOutOfBoundCheck(cgen_state, index_lv, geo_size_lv);
     }
     array_operand_lvs.push_back(index_lv);
     return std::make_tuple(array_operand_lvs, is_null_lv);

diff --git a/QueryEngine/RelAlgTranslatorGeo.cpp b/QueryEngine/RelAlgTranslatorGeo.cpp
@@ -548,24 +548,10 @@ std::vector<std::shared_ptr<Analyzer::Expr>> RelAlgTranslator::translateGeoFunct
                                 ": second argument is expected to be a literal");
       }
       const auto e = translateLiteral(rex_literal);
-      auto ce = std::dynamic_pointer_cast<Analyzer::Constant>(e);
-      if (!ce || !e->get_type_info().is_integer()) {
+      if (!e ||
+          !shared::is_any<kSMALLINT, kTINYINT, kINT>(e->get_type_info().get_type())) {
         throw QueryNotSupported(rex_function->getName() +
-                                ": expecting integer index as second argument");
-      }
-      int32_t index = 0;
-      if (e->get_type_info().get_type() == kSMALLINT) {
-        index = static_cast<int32_t>(ce->get_constval().smallintval);
-      } else if (e->get_type_info().get_type() == kTINYINT) {
-        index = static_cast<int32_t>(ce->get_constval().tinyintval);
-      } else if (e->get_type_info().get_type() == kINT) {
-        index = static_cast<int32_t>(ce->get_constval().intval);
-      } else {
-        throw QueryNotSupported(rex_function->getName() + " expecting integer index");
-      }
-      if (index == 0) {
-        // maybe we will just return NULL here?
-        throw QueryNotSupported(rex_function->getName() + ": invalid index");
+                                " expecting integer index as second argument");
       }
       arg0.push_back(e);
       auto oper_ti =

diff --git a/Tests/GeospatialTest.cpp b/Tests/GeospatialTest.cpp
@@ -2579,6 +2579,63 @@ TEST(GeoSpatial, Projections) {
   }
 }
 
+TEST(GeoSpatial, PointNOutOfBoundCheckNegativeIndex) {
+  for (auto dt : {ExecutorDeviceType::CPU, ExecutorDeviceType::GPU}) {
+    SKIP_NO_GPU();
+    ASSERT_EQ(
+        v<int64_t>(run_simple_agg(
+            "SELECT COUNT(*) FROM (SELECT ST_PointN(ST_GeomFromText('LINESTRING(-1 1, 1 "
+            "1)', 4326), -1) IS NULL as v) WHERE v IS TRUE;",
+            dt,
+            false)),
+        static_cast<int64_t>(1));
+  }
+}
+
+TEST(GeoSpatial, PointNIndexOverflow) {
+  for (auto dt : {ExecutorDeviceType::CPU, ExecutorDeviceType::GPU}) {
+    SKIP_NO_GPU();
+    ASSERT_EQ(
+        v<int64_t>(run_simple_agg(
+            "SELECT COUNT(*) FROM (select ST_PointN(ST_GeomFromText('LINESTRING(-1 1, 1 "
+            "1)', 4326), 1241231231) IS NULL as v) WHERE v IS TRUE;",
+            dt,
+            false)),
+        static_cast<int64_t>(1));
+  }
+}
+
+TEST(GeoSpatial, PointNInvalidIndex) {
+  for (auto dt : {ExecutorDeviceType::CPU, ExecutorDeviceType::GPU}) {
+    SKIP_NO_GPU();
+    ASSERT_EQ(
+        v<int64_t>(run_simple_agg(
+            "SELECT COUNT(*) FROM (select ST_PointN(ST_GeomFromText('LINESTRING(-1 1, 1 "
+            "1)', 4326), 0) IS NULL as v) WHERE v IS TRUE;",
+            dt,
+            false)),
+        static_cast<int64_t>(1));
+    ASSERT_EQ(
+        v<int64_t>(run_simple_agg(
+            "SELECT COUNT(*) FROM (select ST_PointN(ST_GeomFromText('LINESTRING(-1 1, 1 "
+            "1)', 4326), 3) IS NULL as v) WHERE v IS TRUE;",
+            dt,
+            false)),
+        static_cast<int64_t>(1));
+  }
+}
+
+TEST(GeoSpatial, PointNIndexLargerThanInt) {
+  for (auto dt : {ExecutorDeviceType::CPU, ExecutorDeviceType::GPU}) {
+    SKIP_NO_GPU();
+    // ST_PointN expecting integer index
+    ASSERT_ANY_THROW(run_simple_agg(
+        "SELECT ST_PointN(ST_GeomFromText('LINESTRING(-1 1, 1 1)', 4326), 2147483648);",
+        dt,
+        false));
+  }
+}
+
 class GeoSpatialTempTables : public ::testing::Test {
  protected:
   void SetUp() override { import_geospatial_test(/*with_temporary_tables=*/false); }