From ad4594c29829bf03293fa2da13d938b0aea7edb9 Mon Sep 17 00:00:00 2001
From: ANDREI STAROVEROV <staroverovad@sandbox-407923068>
Date: Tue, 11 May 2021 01:35:22 +0300
Subject: [PATCH] Add validation of unknown identifiers in function

---
 src/Functions/FunctionFactory.cpp             |  2 +
 src/Functions/FunctionFactory.h               |  3 --
 .../InterpreterCreateFunctionQuery.cpp        | 41 +++++++++++++++++++
 .../InterpreterCreateFunctionQuery.h          |  4 ++
 src/Parsers/ParserCreateFunctionQuery.cpp     |  1 +
 src/Parsers/ParserCreateFunctionQuery.h       |  1 +
 ...function_with_unknown_variable_in_body.sql |  3 +-
 7 files changed, 50 insertions(+), 5 deletions(-)

diff --git a/src/Functions/FunctionFactory.cpp b/src/Functions/FunctionFactory.cpp
index 7f330d45c37..69a34d4d030 100644
--- a/src/Functions/FunctionFactory.cpp
+++ b/src/Functions/FunctionFactory.cpp
@@ -16,6 +16,8 @@
 namespace DB
 {
 
+class UserDefinedFunction;
+
 namespace ErrorCodes
 {
     extern const int UNKNOWN_FUNCTION;
diff --git a/src/Functions/FunctionFactory.h b/src/Functions/FunctionFactory.h
index 176178f7593..850323fa2df 100644
--- a/src/Functions/FunctionFactory.h
+++ b/src/Functions/FunctionFactory.h
@@ -1,7 +1,6 @@
 #pragma once
 
 #include <Functions/IFunctionAdaptors.h>
-#include <Functions/UserDefinedFunction.h>
 #include <Interpreters/Context_fwd.h>
 #include <Common/IFactoryWithAliases.h>
 #include <Parsers/ASTCreateFunctionQuery.h>
@@ -15,8 +14,6 @@
 namespace DB
 {
 
-class UserDefinedFunction;
-
 /** Creates function by name.
   * Function could use for initialization (take ownership of shared_ptr, for example)
   *  some dictionaries from Context.
diff --git a/src/Interpreters/InterpreterCreateFunctionQuery.cpp b/src/Interpreters/InterpreterCreateFunctionQuery.cpp
index 4fa524534f3..9de6a4aaeff 100644
--- a/src/Interpreters/InterpreterCreateFunctionQuery.cpp
+++ b/src/Interpreters/InterpreterCreateFunctionQuery.cpp
@@ -1,3 +1,4 @@
+#include <Parsers/ASTIdentifier.h>
 #include <Interpreters/ExpressionActions.h>
 #include <Interpreters/ExpressionAnalyzer.h>
 #include <Interpreters/InterpreterCreateFunctionQuery.h>
@@ -7,12 +8,52 @@
 namespace DB
 {
 
+namespace ErrorCodes
+{
+    extern const int UNKNOWN_IDENTIFIER;
+}
+
 BlockIO InterpreterCreateFunctionQuery::execute()
 {
     FunctionNameNormalizer().visit(query_ptr.get());
     auto & create_function_query = query_ptr->as<ASTCreateFunctionQuery &>();
+    validateFunction(create_function_query.function_core);
     FunctionFactory::instance().registerUserDefinedFunction(create_function_query);
     return {};
 }
 
+void InterpreterCreateFunctionQuery::validateFunction(ASTPtr function)
+{
+    const auto * args_tuple = function->as<ASTFunction>()->arguments->children.at(0)->as<ASTFunction>();
+    std::unordered_set<String> arguments;
+    for (const auto & argument : args_tuple->arguments->children)
+        arguments.insert(argument->as<ASTIdentifier>()->name());
+
+    std::vector<String> identifiers_in_body;
+    ASTPtr function_body = function->as<ASTFunction>()->children.at(0)->children.at(1);
+    getIdentifiers(function_body, identifiers_in_body);
+
+    for (const auto & identifier : identifiers_in_body)
+    {
+        if (!arguments.contains(identifier))
+        {
+            std::stringstream s;
+            s << "Identifier '" << identifier << "' does not exist in arguments";
+            throw Exception(s.str(), ErrorCodes::UNKNOWN_IDENTIFIER);
+        }
+    }
+}
+
+void InterpreterCreateFunctionQuery::getIdentifiers(ASTPtr node, std::vector<String> & identifiers)
+{
+    for (const auto & child : node->children)
+    {
+        auto identifier_name_opt = tryGetIdentifierName(child);
+        if (identifier_name_opt)
+            identifiers.push_back(identifier_name_opt.value());
+
+        getIdentifiers(child, identifiers);
+    }
+}
+
 }
diff --git a/src/Interpreters/InterpreterCreateFunctionQuery.h b/src/Interpreters/InterpreterCreateFunctionQuery.h
index 81347bcc711..79b7839116a 100644
--- a/src/Interpreters/InterpreterCreateFunctionQuery.h
+++ b/src/Interpreters/InterpreterCreateFunctionQuery.h
@@ -15,6 +15,10 @@ public:
 
     BlockIO execute() override;
 
+private:
+    static void validateFunction(ASTPtr function);
+    static void getIdentifiers(ASTPtr node, std::vector<String> & identifiers);
+
 private:
     ASTPtr query_ptr;
 };
diff --git a/src/Parsers/ParserCreateFunctionQuery.cpp b/src/Parsers/ParserCreateFunctionQuery.cpp
index 1fcce6cbf45..fbfd02415e7 100644
--- a/src/Parsers/ParserCreateFunctionQuery.cpp
+++ b/src/Parsers/ParserCreateFunctionQuery.cpp
@@ -8,6 +8,7 @@
 
 namespace DB
 {
+
 bool ParserCreateFunctionQuery::parseImpl(IParser::Pos & pos, ASTPtr & node, Expected & expected)
 {
     ParserKeyword s_create("CREATE");
diff --git a/src/Parsers/ParserCreateFunctionQuery.h b/src/Parsers/ParserCreateFunctionQuery.h
index a48bbdeb563..aac643b995d 100644
--- a/src/Parsers/ParserCreateFunctionQuery.h
+++ b/src/Parsers/ParserCreateFunctionQuery.h
@@ -4,6 +4,7 @@
 
 namespace DB
 {
+
 /// CREATE FUNCTION test AS x -> x || '1'
 class ParserCreateFunctionQuery : public IParserBase
 {
diff --git a/tests/queries/0_stateless/01856_create_function_with_unknown_variable_in_body.sql b/tests/queries/0_stateless/01856_create_function_with_unknown_variable_in_body.sql
index a563d1c4ab2..1fed9fe9103 100644
--- a/tests/queries/0_stateless/01856_create_function_with_unknown_variable_in_body.sql
+++ b/tests/queries/0_stateless/01856_create_function_with_unknown_variable_in_body.sql
@@ -1,2 +1 @@
-create function MyFunc2 as (a, b) -> a || b || c;
-select MyFunc2('1', '2'); -- { serverError 47 }
+create function MyFunc2 as (a, b) -> a || b || c; -- { serverError 47 }