From 59dd8557c70387ef9d44ce9087925f97f5fa1577 Mon Sep 17 00:00:00 2001
From: Anton Romanov <Anton.Romanoff@gmail.com>
Date: Tue, 13 Feb 2024 16:57:43 +0100
Subject: [PATCH] Early validate json path.

---
 .../yql/core/type_ann/type_ann_core.cpp       |  4 +---
 ydb/library/yql/core/ya.make                  |  1 +
 .../yql/core/yql_expr_type_annotation.cpp     | 20 +++++++++++++++++++
 .../yql/core/yql_expr_type_annotation.h       |  1 +
 ydb/library/yql/minikql/jsonpath/jsonpath.cpp |  2 +-
 5 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/ydb/library/yql/core/type_ann/type_ann_core.cpp b/ydb/library/yql/core/type_ann/type_ann_core.cpp
index 5ad0621f1d9c..c8382b80afb9 100644
--- a/ydb/library/yql/core/type_ann/type_ann_core.cpp
+++ b/ydb/library/yql/core/type_ann/type_ann_core.cpp
@@ -493,9 +493,7 @@ namespace NTypeAnnImpl {
             return false;
         }
 
-        // second argument must be "Utf8" type
-        const auto& jsonPathArg = function.JsonPath().Ref();
-        if (!EnsureSpecificDataType(jsonPathArg, EDataSlot::Utf8, ctx.Expr)) {
+        if (!EnsureValidJsonPath(function.JsonPath().Ref(), ctx.Expr)) {
             return false;
         }
 
diff --git a/ydb/library/yql/core/ya.make b/ydb/library/yql/core/ya.make
index 088e2dae50f2..6b478f8a892e 100644
--- a/ydb/library/yql/core/ya.make
+++ b/ydb/library/yql/core/ya.make
@@ -71,6 +71,7 @@ PEERDIR(
     ydb/library/yql/core/url_lister/interface
     ydb/library/yql/core/url_preprocessing/interface
     ydb/library/yql/minikql
+    ydb/library/yql/minikql/jsonpath
     ydb/library/yql/protos
     ydb/library/yql/public/udf
     ydb/library/yql/public/udf/tz
diff --git a/ydb/library/yql/core/yql_expr_type_annotation.cpp b/ydb/library/yql/core/yql_expr_type_annotation.cpp
index c18059f8ec05..bc17ed4339d5 100644
--- a/ydb/library/yql/core/yql_expr_type_annotation.cpp
+++ b/ydb/library/yql/core/yql_expr_type_annotation.cpp
@@ -9,6 +9,7 @@
 #include <ydb/library/yql/public/udf/udf_data_type.h>
 #include <ydb/library/yql/minikql/dom/json.h>
 #include <ydb/library/yql/minikql/dom/yson.h>
+#include <ydb/library/yql/minikql/jsonpath/jsonpath.h>
 #include <ydb/library/yql/core/sql_types/simple_types.h>
 #include "ydb/library/yql/parser/pg_catalog/catalog.h"
 #include <ydb/library/yql/parser/pg_wrapper/interface/utils.h>
@@ -6047,6 +6048,25 @@ bool EnsureScalarType(TPositionHandle position, const TTypeAnnotationNode& type,
     return true;
 }
 
+bool EnsureValidJsonPath(const TExprNode& node, TExprContext& ctx) {
+    if (!EnsureSpecificDataType(node, EDataSlot::Utf8, ctx))
+        return false;
+
+    if (node.IsCallable("Utf8")) {
+        if (TIssues issues; !NJsonPath::ParseJsonPath(node.Tail().Content(), issues, 7U)) {
+            TIssue issue(ctx.GetPosition(node.Pos()), TStringBuilder() << "Invalid json path: " << node.Tail().Content());
+            if (bool(issues)) {
+                for (const auto& i : issues) {
+                    issue.AddSubIssue(new TIssue(i));
+                }
+            }
+            ctx.AddError(issue);
+            return false;
+        }
+    }
+    return true;
+}
+
 const TTypeAnnotationNode* GetBlockItemType(const TTypeAnnotationNode& type, bool& isScalar) {
     YQL_ENSURE(type.IsBlockOrScalar());
     const auto kind = type.GetKind();
diff --git a/ydb/library/yql/core/yql_expr_type_annotation.h b/ydb/library/yql/core/yql_expr_type_annotation.h
index c7405b39bef3..d61896a6338d 100644
--- a/ydb/library/yql/core/yql_expr_type_annotation.h
+++ b/ydb/library/yql/core/yql_expr_type_annotation.h
@@ -145,6 +145,7 @@ bool EnsureDryType(TPositionHandle position, const TTypeAnnotationNode& type, TE
 bool EnsureDryType(const TExprNode& node, TExprContext& ctx);
 bool EnsureDictType(const TExprNode& node, TExprContext& ctx);
 bool EnsureDictType(TPositionHandle position, const TTypeAnnotationNode& type, TExprContext& ctx);
+bool EnsureValidJsonPath(const TExprNode& node, TExprContext& ctx);
 
 bool IsVoidType(const TExprNode& node, TExprContext& ctx);
 bool EnsureVoidType(const TExprNode& node, TExprContext& ctx);
diff --git a/ydb/library/yql/minikql/jsonpath/jsonpath.cpp b/ydb/library/yql/minikql/jsonpath/jsonpath.cpp
index 985ea72c361e..86dd3b83d183 100644
--- a/ydb/library/yql/minikql/jsonpath/jsonpath.cpp
+++ b/ydb/library/yql/minikql/jsonpath/jsonpath.cpp
@@ -59,7 +59,7 @@ const TAstNodePtr ParseJsonPathAst(const TStringBuf path, TIssues& issues, size_
     }
 
     google::protobuf::Arena arena;
-    const google::protobuf::Message* rawAst;
+    const google::protobuf::Message* rawAst = nullptr;
     {
     #if defined(_tsan_enabled_)
         TGuard<TMutex> guard(SanitizerJsonPathTranslationMutex);