diff options
author | omgronny <omgronny@yandex-team.com> | 2024-08-08 13:44:16 +0300 |
---|---|---|
committer | omgronny <omgronny@yandex-team.com> | 2024-08-08 13:56:38 +0300 |
commit | 74f9dee36c760bd41ac8077b406660be3815aae1 (patch) | |
tree | 560c1bb062958ff059d27c467aaebf129ff7401d | |
parent | ed77962244e7018d6b7bd14af51dd890de3ea13f (diff) | |
download | ydb-74f9dee36c760bd41ac8077b406660be3815aae1.tar.gz |
YT-22275: Check map and list sizes when parsing operation spec
88b5f4f5b2e3ed7046d9ac90088e16a83ae7a765
-rw-r--r-- | yt/yt/core/ytree/convert-inl.h | 5 | ||||
-rw-r--r-- | yt/yt/core/ytree/convert.h | 3 | ||||
-rw-r--r-- | yt/yt/core/ytree/tree_builder.cpp | 17 | ||||
-rw-r--r-- | yt/yt/core/ytree/tree_builder.h | 4 |
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()); //////////////////////////////////////////////////////////////////////////////// |