fix(dot/network): resize bytes slice buffer if needed (#2291)

ChainSafe · Feb 17, 2022 · 8db8b2a · 8db8b2a
1 parent 67a9bbb
commit 8db8b2a
Show file tree

Hide file tree

Showing 6 changed files with 11 additions and 9 deletions.
diff --git a/dot/network/helpers_test.go b/dot/network/helpers_test.go
@@ -86,7 +86,7 @@ func (s *testStreamHandler) readStream(stream libp2pnetwork.Stream,
 	}()
 
 	for {
-		tot, err := readStream(stream, msgBytes)
+		tot, err := readStream(stream, &msgBytes)
 		if errors.Is(err, io.EOF) {
 			return
 		} else if err != nil {

diff --git a/dot/network/inbound.go b/dot/network/inbound.go
@@ -17,10 +17,9 @@ func (s *Service) readStream(stream libp2pnetwork.Stream, decoder messageDecoder
 	peer := stream.Conn().RemotePeer()
 	buffer := s.bufPool.Get().(*[]byte)
 	defer s.bufPool.Put(buffer)
-	msgBytes := *buffer
 
 	for {
-		n, err := readStream(stream, msgBytes[:])
+		n, err := readStream(stream, buffer)
 		if err != nil {
 			logger.Tracef(
 				"failed to read from stream id %s of peer %s using protocol %s: %s",
@@ -32,6 +31,7 @@ func (s *Service) readStream(stream libp2pnetwork.Stream, decoder messageDecoder
 
 		// decode message based on message type
 		// stream should always be inbound if it passes through service.readStream
+		msgBytes := *buffer
 		msg, err := decoder(msgBytes[:n], peer, isInbound(stream))
 		if err != nil {
 			logger.Tracef("failed to decode message from stream id %s using protocol %s: %s",

diff --git a/dot/network/notifications.go b/dot/network/notifications.go
@@ -431,14 +431,14 @@ func (s *Service) readHandshake(stream libp2pnetwork.Stream, decoder HandshakeDe
 
 		buffer := s.bufPool.Get().(*[]byte)
 		defer s.bufPool.Put(buffer)
-		msgBytes := *buffer
 
-		tot, err := readStream(stream, msgBytes[:])
+		tot, err := readStream(stream, buffer)
 		if err != nil {
 			hsC <- &handshakeReader{hs: nil, err: err}
 			return
 		}
 
+		msgBytes := *buffer
 		hs, err := decoder(msgBytes[:tot])
 		if err != nil {
 			s.host.cm.peerSetHandler.ReportPeer(peerset.ReputationChange{

diff --git a/dot/network/service.go b/dot/network/service.go
@@ -34,7 +34,7 @@ const (
 	blockAnnounceID = "/block-announces/1"
 	transactionsID  = "/transactions/1"
 
-	maxMessageSize = 1024 * 63 // 63kb for now
+	maxMessageSize = 1024 * 64 // 64kb for now
 )
 
 var (

diff --git a/dot/network/sync.go b/dot/network/sync.go
@@ -58,7 +58,7 @@ func (s *Service) receiveBlockResponse(stream libp2pnetwork.Stream) (*BlockRespo
 
 	buf := s.blockResponseBuf
 
-	n, err := readStream(stream, buf)
+	n, err := readStream(stream, &buf)
 	if err != nil {
 		return nil, fmt.Errorf("read stream error: %w", err)
 	}

diff --git a/dot/network/utils.go b/dot/network/utils.go
@@ -176,7 +176,7 @@ func readLEB128ToUint64(r io.Reader, buf []byte) (uint64, int, error) {
 }
 
 // readStream reads from the stream into the given buffer, returning the number of bytes read
-func readStream(stream libp2pnetwork.Stream, buf []byte) (int, error) {
+func readStream(stream libp2pnetwork.Stream, bufPointer *[]byte) (int, error) {
 	if stream == nil {
 		return 0, errors.New("stream is nil")
 	}
@@ -185,6 +185,7 @@ func readStream(stream libp2pnetwork.Stream, buf []byte) (int, error) {
 		tot int
 	)
 
+	buf := *bufPointer
 	length, bytesRead, err := readLEB128ToUint64(stream, buf[:1])
 	if err != nil {
 		return bytesRead, fmt.Errorf("failed to read length: %w", err)
@@ -195,8 +196,9 @@ func readStream(stream libp2pnetwork.Stream, buf []byte) (int, error) {
 	}
 
 	if length > uint64(len(buf)) {
+		extraBytes := int(length) - len(buf)
+		*bufPointer = append(buf, make([]byte, extraBytes)...) // TODO #2288 use bytes.Buffer instead
 		logger.Warnf("received message with size %d greater than allocated message buffer size %d", length, len(buf))
-		return 0, fmt.Errorf("message size greater than allocated message buffer: got %d", length)
 	}
 
 	if length > maxBlockResponseSize {