multistream max (#229)

Signed-off-by: turuslan <[email protected]>
libp2p · Jan 12, 2024 · aa51f77 · aa51f77
1 parent e697206
commit aa51f77
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 0 deletions.
diff --git 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
@@ -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;
@@ -43,7 +44,15 @@ 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) {
+            state_ = kOverflow;
+            break;
+          }
           continue;
         }
         if (s != VarintPrefixReader::kReady) {