From c33443706009e5002912d7596c9529055439cad1 Mon Sep 17 00:00:00 2001
From: jimingquan <mingquan.ji@vesoft.com>
Date: Fri, 10 Sep 2021 10:29:35 +0800
Subject: [PATCH] Add max allowed query size (#2813)

* add max allowed query size
* address comment
---
 conf/nebula-graphd.conf.default               |  2 ++
 conf/nebula-graphd.conf.production            |  2 ++
 src/graph/service/GraphFlags.cpp              |  1 +
 .../validator/test/QueryValidatorTest.cpp     | 21 +++++++++++++++++++
 src/parser/GQLParser.h                        |  9 ++++++--
 5 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/conf/nebula-graphd.conf.default b/conf/nebula-graphd.conf.default
index adee924bd85..9951cf37896 100644
--- a/conf/nebula-graphd.conf.default
+++ b/conf/nebula-graphd.conf.default
@@ -33,6 +33,8 @@
 # Whether to treat partial success as an error.
 # This flag is only used for Read-only access, and Modify access always treats partial success as an error.
 --accept_partial_success=false
+# Maximum sentence length, unit byte
+--max_allowed_query_size=4194304
 
 ########## networking ##########
 # Comma separated Meta Server Addresses
diff --git a/conf/nebula-graphd.conf.production b/conf/nebula-graphd.conf.production
index 6c4eeb96405..b3eebaf8dbf 100644
--- a/conf/nebula-graphd.conf.production
+++ b/conf/nebula-graphd.conf.production
@@ -31,6 +31,8 @@
 # Whether to treat partial success as an error.
 # This flag is only used for Read-only access, and Modify access always treats partial success as an error.
 --accept_partial_success=false
+# Maximum sentence length, unit byte
+--max_allowed_query_size=4194304
 
 ########## networking ##########
 # Comma separated Meta Server Addresses
diff --git a/src/graph/service/GraphFlags.cpp b/src/graph/service/GraphFlags.cpp
index 08e3164ae75..20e60daba8c 100644
--- a/src/graph/service/GraphFlags.cpp
+++ b/src/graph/service/GraphFlags.cpp
@@ -45,6 +45,7 @@ DEFINE_string(auth_type,
 
 DEFINE_string(cloud_http_url, "", "cloud http url including ip, port, url path");
 DEFINE_uint32(max_allowed_statements, 512, "Max allowed sequential statements");
+DEFINE_uint32(max_allowed_query_size, 4194304, "Max allowed sequential query size");
 
 DEFINE_int64(max_allowed_connections,
              std::numeric_limits<int64_t>::max(),
diff --git a/src/graph/validator/test/QueryValidatorTest.cpp b/src/graph/validator/test/QueryValidatorTest.cpp
index 8593b17f55e..025e3475f5a 100644
--- a/src/graph/validator/test/QueryValidatorTest.cpp
+++ b/src/graph/validator/test/QueryValidatorTest.cpp
@@ -8,6 +8,7 @@
 #include "graph/validator/test/ValidatorTestBase.h"
 
 DECLARE_uint32(max_allowed_statements);
+DECLARE_uint32(max_allowed_query_size);
 
 namespace nebula {
 namespace graph {
@@ -1118,6 +1119,26 @@ TEST_F(QueryValidatorTest, TestMaxAllowedStatements) {
             "exceeded");
 }
 
+TEST_F(QueryValidatorTest, TestMaxAllowedQuerySize) {
+  FLAGS_max_allowed_query_size = 256;
+  std::string query = "INSERT VERTEX person(name, age) VALUES ";
+  std::string value = "\"person_1\":(\"person_1\", 1),";
+  int count = (FLAGS_max_allowed_query_size - query.size()) / value.size();
+  std::string values;
+  values.reserve(FLAGS_max_allowed_query_size);
+  for (int i = 0; i < count; ++i) {
+    values.append(value);
+  }
+  values.erase(values.size() - 1);
+  query += values;
+  EXPECT_TRUE(checkResult(query));
+  query.append(",\"person_2\":(\"person_2\", 2);");
+  auto result = checkResult(query);
+  EXPECT_FALSE(result);
+  EXPECT_EQ(std::string(result.message()), "SyntaxError: Query is too large (282 > 256).");
+  FLAGS_max_allowed_query_size = 4194304;
+}
+
 TEST_F(QueryValidatorTest, TestMatch) {
   {
     std::string query =
diff --git a/src/parser/GQLParser.h b/src/parser/GQLParser.h
index 10d1432d44c..d050fd71d27 100644
--- a/src/parser/GQLParser.h
+++ b/src/parser/GQLParser.h
@@ -11,6 +11,7 @@
 #include "parser/GraphParser.hpp"
 #include "parser/GraphScanner.h"
 
+DECLARE_uint32(max_allowed_query_size);
 namespace nebula {
 
 class GQLParser {
@@ -39,8 +40,12 @@ class GQLParser {
   }
 
   StatusOr<std::unique_ptr<Sentence>> parse(std::string query) {
-    // Since GraphScanner needs a writable buffer, we have to copy the query
-    // string
+    // Since GraphScanner needs a writable buffer, we have to copy the query string
+    size_t querySize = query.size();
+    size_t maxAllowedQuerySize = static_cast<size_t>(FLAGS_max_allowed_query_size);
+    if (querySize > maxAllowedQuerySize) {
+      return Status::SyntaxError("Query is too large (%ld > %ld).", querySize, maxAllowedQuerySize);
+    }
     buffer_ = std::move(query);
     pos_ = &buffer_[0];
     end_ = pos_ + buffer_.size();