From c3107b7a609cd53d47d610091168adafc3232c8f Mon Sep 17 00:00:00 2001
From: turuslan <turuslan.devbox@gmail.com>
Date: Fri, 12 Jan 2024 15:02:24 +0300
Subject: [PATCH 1/2] multistream max

Signed-off-by: turuslan <turuslan.devbox@gmail.com>
---
 include/libp2p/basic/varint_prefix_reader.hpp | 4 ++++
 src/protocol_muxer/multiselect/parser.cpp     | 9 +++++++++
 2 files changed, 13 insertions(+)

diff --git a/include/libp2p/basic/varint_prefix_reader.hpp b/include/libp2p/basic/varint_prefix_reader.hpp
index 2f296b242..28eebc5b0 100644
--- a/include/libp2p/basic/varint_prefix_reader.hpp
+++ b/include/libp2p/basic/varint_prefix_reader.hpp
@@ -42,6 +42,10 @@ namespace libp2p::basic {
       return value_;
     }
 
+    size_t size() const {
+      return got_bytes_;
+    }
+
     /// Resets reader's state
     void reset();
 
diff --git a/src/protocol_muxer/multiselect/parser.cpp b/src/protocol_muxer/multiselect/parser.cpp
index 00dc186d8..39f35c91c 100644
--- a/src/protocol_muxer/multiselect/parser.cpp
+++ b/src/protocol_muxer/multiselect/parser.cpp
@@ -11,6 +11,7 @@
 namespace libp2p::protocol_muxer::multiselect::detail {
 
   constexpr size_t kMaxRecursionDepth = 3;
+  constexpr size_t kMaxLenBytes = 2;
 
   size_t Parser::bytesNeeded() const {
     size_t n = 0;
@@ -44,12 +45,20 @@ namespace libp2p::protocol_muxer::multiselect::detail {
       if (expected_msg_size_ == 0) {
         auto s = varint_reader_.consume(data);
         if (s == VarintPrefixReader::kUnderflow) {
+          if (varint_reader_.size() >= kMaxLenBytes) {
+            state_ = kOverflow;
+            break;
+          }
           continue;
         }
         if (s != VarintPrefixReader::kReady) {
           state_ = kOverflow;
           break;
         }
+        if (varint_reader_.size() > kMaxLenBytes) {
+          state_ = kOverflow;
+          break;
+        }
         expected_msg_size_ = static_cast<IndexType>(varint_reader_.value());
         if (expected_msg_size_ == 0) {
           // zero varint received, not acceptable, but not fatal

From ae8bc23738c342d20197cb4762b1f782ca9e32f7 Mon Sep 17 00:00:00 2001
From: turuslan <turuslan.devbox@gmail.com>
Date: Fri, 12 Jan 2024 15:16:33 +0300
Subject: [PATCH 2/2] pr comment

Signed-off-by: turuslan <turuslan.devbox@gmail.com>
---
 src/protocol_muxer/multiselect/parser.cpp | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/src/protocol_muxer/multiselect/parser.cpp b/src/protocol_muxer/multiselect/parser.cpp
index 39f35c91c..8c9049a23 100644
--- a/src/protocol_muxer/multiselect/parser.cpp
+++ b/src/protocol_muxer/multiselect/parser.cpp
@@ -44,8 +44,12 @@ namespace libp2p::protocol_muxer::multiselect::detail {
 
       if (expected_msg_size_ == 0) {
         auto s = varint_reader_.consume(data);
+        if (varint_reader_.size() > kMaxLenBytes) {
+          state_ = kOverflow;
+          break;
+        }
         if (s == VarintPrefixReader::kUnderflow) {
-          if (varint_reader_.size() >= kMaxLenBytes) {
+          if (varint_reader_.size() == kMaxLenBytes) {
             state_ = kOverflow;
             break;
           }
@@ -55,10 +59,6 @@ namespace libp2p::protocol_muxer::multiselect::detail {
           state_ = kOverflow;
           break;
         }
-        if (varint_reader_.size() > kMaxLenBytes) {
-          state_ = kOverflow;
-          break;
-        }
         expected_msg_size_ = static_cast<IndexType>(varint_reader_.value());
         if (expected_msg_size_ == 0) {
           // zero varint received, not acceptable, but not fatal