YT-22275: Check map and list sizes when parsing operation spec

88b5f4f5b2e3ed7046d9ac90088e16a83ae7a765

4 files changed, 22 insertions, 7 deletions

diff --git a/yt/yt/core/ytree/convert-inl.h b/yt/yt/core/ytree/convert-inl.h
index 340d8a8686..2b67d693a9 100644
--- a/yt/yt/core/ytree/convert-inl.h
+++ b/yt/yt/core/ytree/convert-inl.h
@@ -101,11 +101,12 @@ NYson::TYsonProducer ConvertToProducer(T&& value)
 template <class T>
 INodePtr ConvertToNode(
     const T& value,
-    INodeFactory* factory)
+    INodeFactory* factory,
+    int treeSizeLimit)
 {
     auto type = GetYsonType(value);
 
-    auto builder = CreateBuilderFromFactory(factory);
+    auto builder = CreateBuilderFromFactory(factory, treeSizeLimit);
     builder->BeginTree();
 
     switch (type) {
diff --git a/yt/yt/core/ytree/convert.h b/yt/yt/core/ytree/convert.h
index 6f9bdaa3ca..383d9bb945 100644
--- a/yt/yt/core/ytree/convert.h
+++ b/yt/yt/core/ytree/convert.h
@@ -36,7 +36,8 @@ NYson::TYsonProducer ConvertToProducer(T&& value);
 template <class T>
 INodePtr ConvertToNode(
     const T& value,
-    INodeFactory* factory = GetEphemeralNodeFactory());
+    INodeFactory* factory = GetEphemeralNodeFactory(),
+    int treeSizeLimit = std::numeric_limits<int>::max());
 
 template <class T>
 IAttributeDictionaryPtr ConvertToAttributes(const T& value);
diff --git a/yt/yt/core/ytree/tree_builder.cpp b/yt/yt/core/ytree/tree_builder.cpp
index af27368182..4f4603b176 100644
--- a/yt/yt/core/ytree/tree_builder.cpp
+++ b/yt/yt/core/ytree/tree_builder.cpp
@@ -22,8 +22,9 @@ class TTreeBuilder
     , public ITreeBuilder
 {
 public:
-    explicit TTreeBuilder(INodeFactory* factory)
+    TTreeBuilder(INodeFactory* factory, int treeSizeLimit)
         : Factory_(factory)
+        , TreeSizeLimit_(treeSizeLimit)
     {
         YT_ASSERT(Factory_);
     }
@@ -143,6 +144,9 @@ private:
     std::unique_ptr<TAttributeConsumer> AttributeConsumer_;
     IAttributeDictionaryPtr Attributes_;
 
+    const int TreeSizeLimit_;
+    int TreeSize_ = 0;
+
     void AddNode(INodePtr node, bool push)
     {
         if (Attributes_) {
@@ -150,6 +154,11 @@ private:
             Attributes_ = nullptr;
         }
 
+        if (++TreeSize_ > TreeSizeLimit_) {
+            THROW_ERROR_EXCEPTION("Tree size limit exceeded")
+                << TErrorAttribute("tree_size_limit", TreeSizeLimit_);
+        }
+
         if (NodeStack_.empty()) {
             ResultNode_ = node;
         } else {
@@ -170,9 +179,11 @@ private:
     }
 };
 
-std::unique_ptr<ITreeBuilder> CreateBuilderFromFactory(INodeFactory* factory)
+std::unique_ptr<ITreeBuilder> CreateBuilderFromFactory(
+    INodeFactory* factory,
+    int treeSizeLimit)
 {
-    return std::unique_ptr<ITreeBuilder>(new TTreeBuilder(std::move(factory)));
+    return std::unique_ptr<ITreeBuilder>(new TTreeBuilder(std::move(factory), treeSizeLimit));
 }
 
 ////////////////////////////////////////////////////////////////////////////////
diff --git a/yt/yt/core/ytree/tree_builder.h b/yt/yt/core/ytree/tree_builder.h
index 4fcbe25cc9..9154752a22 100644
--- a/yt/yt/core/ytree/tree_builder.h
+++ b/yt/yt/core/ytree/tree_builder.h
@@ -35,7 +35,9 @@ struct ITreeBuilder
 /*!
  *  \param factory A factory used for materializing the nodes.
  */
-std::unique_ptr<ITreeBuilder> CreateBuilderFromFactory(INodeFactory* factory);
+std::unique_ptr<ITreeBuilder> CreateBuilderFromFactory(
+    INodeFactory* factory,
+    int treeSizeLimit = std::numeric_limits<int>::max());
 
 ////////////////////////////////////////////////////////////////////////////////