Check the rand function parameter range. (vesoft-inc#2086)

Co-authored-by: dutor <440396+dutor@users.noreply.github.com>
grebokcn · May 19, 2020 · cb2a857 · cb2a857
1 parent 6153384
commit cb2a857
Show file tree

Hide file tree

Showing 6 changed files with 44 additions and 8 deletions.
diff --git a/src/common/base/Status.h b/src/common/base/Status.h
@@ -100,6 +100,7 @@ class Status final {
     STATUS_GENERATOR(Error);
     STATUS_GENERATOR(NoSuchFile);
     STATUS_GENERATOR(NotSupported);
+    STATUS_GENERATOR(InvalidParameter);
 
     // Graph engine errors
     STATUS_GENERATOR(SyntaxError);
@@ -141,6 +142,7 @@ class Status final {
         kError                  = 101,
         kNoSuchFile             = 102,
         kNotSupported           = 103,
+        kInvalidParameter       = 104,
         // 2xx, for graph engine errors
         kSyntaxError            = 201,
         kStatementEmpty         = 202,

diff --git a/src/common/filter/Expressions.cpp b/src/common/filter/Expressions.cpp
@@ -460,7 +460,7 @@ OptVariantType FunctionCallExpression::eval(Getters &getters) const {
 
     // TODO(simon.liu)
     auto r = function_(args);
-    return OptVariantType(r);
+    return r;
 }
 
 Status FunctionCallExpression::prepare() {

diff --git a/src/common/filter/Expressions.h b/src/common/filter/Expressions.h
@@ -10,6 +10,7 @@
 #include "base/StatusOr.h"
 #include "base/Status.h"
 #include "storage/client/StorageClient.h"
+#include "filter/FunctionManager.h"
 #include <boost/variant.hpp>
 #include <folly/futures/Future.h>
 
@@ -720,7 +721,7 @@ class FunctionCallExpression final : public Expression {
         }
     }
 
-    void setFunc(std::function<VariantType(const std::vector<VariantType>&)> func) {
+    void setFunc(FunctionManager::Function func) {
         function_ = func;
     }
 
@@ -732,7 +733,7 @@ class FunctionCallExpression final : public Expression {
 private:
     std::unique_ptr<std::string>                name_;
     std::vector<std::unique_ptr<Expression>>    args_;
-    std::function<VariantType(const std::vector<VariantType>&)> function_;
+    FunctionManager::Function                   function_;
 };
 
 // (uuid)expr

diff --git a/src/common/filter/FunctionManager.cpp b/src/common/filter/FunctionManager.cpp
@@ -6,6 +6,7 @@
 
 #include "base/Base.h"
 #include "filter/FunctionManager.h"
+#include "filter/Expressions.h"
 #include "time/WallClock.h"
 #include "filter/geo/GeoFilter.h"
 
@@ -193,18 +194,34 @@ FunctionManager::FunctionManager() {
         auto &attr = functions_["rand32"];
         attr.minArity_ = 0;
         attr.maxArity_ = 2;
-        attr.body_ = [] (const auto &args) {
+        attr.body_ = [] (const auto &args) -> OptVariantType {
             if (args.empty()) {
                 auto value = folly::Random::rand32();
                 return static_cast<int64_t>(static_cast<int32_t>(value));
             } else if (args.size() == 1UL) {
                 auto max = Expression::asInt(args[0]);
+                if (max < 0) {
+                    return Status::InvalidParameter("Invalid negative number");
+                }
+                if (max > std::numeric_limits<uint32_t>::max()) {
+                    return Status::InvalidParameter("Too large operand");
+                }
                 auto value = folly::Random::rand32(max);
                 return static_cast<int64_t>(static_cast<int32_t>(value));
             }
             DCHECK_EQ(2UL, args.size());
             auto min = Expression::asInt(args[0]);
             auto max = Expression::asInt(args[1]);
+            if (max < 0 || min < 0) {
+                return Status::InvalidParameter("Invalid negative number");
+            }
+            if (max > std::numeric_limits<uint32_t>::max() ||
+                min > std::numeric_limits<uint32_t>::max()) {
+                return Status::InvalidParameter("Too large operand");
+            }
+            if (min >= max) {
+                return Status::InvalidParameter("Invalid number range");
+            }
             return static_cast<int64_t>(folly::Random::rand32(min, max));
         };
     }
@@ -213,16 +230,25 @@ FunctionManager::FunctionManager() {
         auto &attr = functions_["rand64"];
         attr.minArity_ = 0;
         attr.maxArity_ = 2;
-        attr.body_ = [] (const auto &args) {
+        attr.body_ = [] (const auto &args) -> OptVariantType {
             if (args.empty()) {
                 return static_cast<int64_t>(folly::Random::rand64());
             } else if (args.size() == 1UL) {
                 auto max = Expression::asInt(args[0]);
+                if (max < 0) {
+                    return Status::InvalidParameter("Invalid negative number");
+                }
                 return static_cast<int64_t>(folly::Random::rand64(max));
             }
             DCHECK_EQ(2UL, args.size());
             auto min = Expression::asInt(args[0]);
             auto max = Expression::asInt(args[1]);
+            if (max < 0 || min < 0) {
+                return Status::InvalidParameter("Invalid negative number");
+            }
+            if (min >= max) {
+                return Status::InvalidParameter("Invalid number range");
+            }
             return static_cast<int64_t>(folly::Random::rand64(min, max));
         };
     }
@@ -527,7 +553,7 @@ FunctionManager::FunctionManager() {
             }
         };
     }
-}
+}  // NOLINT
 
 
 // static

diff --git a/src/common/filter/FunctionManager.h b/src/common/filter/FunctionManager.h
@@ -10,7 +10,6 @@
 #include "base/Base.h"
 #include "base/StatusOr.h"
 #include "base/Status.h"
-#include "filter/Expressions.h"
 #include <folly/futures/Future.h>
 
 /**
@@ -24,7 +23,7 @@ namespace nebula {
 
 class FunctionManager final {
 public:
-    using Function = std::function<VariantType(const std::vector<VariantType>&)>;
+    using Function = std::function<StatusOr<VariantType>(const std::vector<VariantType>&)>;
 
     /**
      * To obtain a function named `func', with the actual arity.

diff --git a/src/common/filter/test/ExpressionTest.cpp b/src/common/filter/test/ExpressionTest.cpp
@@ -630,8 +630,16 @@ TEST_F(ExpressionTest, FunctionCall) {
 
     TEST_EXPR_LT(rand32(1024), 1024);
     TEST_EXPR_GE(rand32(1024), 0);
+    TEST_EXPR_FAILED(rand32(-1));
+    TEST_EXPR_FAILED(rand32(-1, -2));
+    TEST_EXPR_FAILED(rand32(3, 2));
+    TEST_EXPR_FAILED(rand32(2, 2));
     TEST_EXPR_LT(rand64(1024, 4096), 4096);
     TEST_EXPR_GE(rand64(1024, 4096), 1024);
+    TEST_EXPR_FAILED(rand64(-1));
+    TEST_EXPR_FAILED(rand64(-1, -2));
+    TEST_EXPR_FAILED(rand64(3, 2));
+    TEST_EXPR_FAILED(rand64(2, 2));
 
     TEST_EXPR_GT(now(), 1554716753);
     TEST_EXPR_LE(now(), 4773548753);  // failed 102 years later