diff --git a/p2p/helper.go b/p2p/helper.go index 18fa6da63..e4d8b3a9c 100644 --- a/p2p/helper.go +++ b/p2p/helper.go @@ -57,36 +57,16 @@ func ReadMessageContent(s Stream) ([]byte, error) { log.Error("Failed reading the p2p message size field: only read", "Num of bytes", n) return contentBuf.Bytes(), err } - //log.Print(fourBytes) - // Number of bytes for the message content - bytesToRead := binary.BigEndian.Uint32(fourBytes) - //log.Printf("The content size is %d bytes.", bytesToRead) - //// Read the content in chunk of 16 * 1024 bytes - tmpBuf := make([]byte, BatchSizeInByte) -ILOOP: - for { - timeoutDuration := 10 * time.Second - s.SetReadDeadline(time.Now().Add(timeoutDuration)) - if bytesToRead < BatchSizeInByte { - // Read the last number of bytes less than 1024 - tmpBuf = make([]byte, bytesToRead) - } - n, err := r.Read(tmpBuf) - contentBuf.Write(tmpBuf[:n]) - switch err { - case io.EOF: - // TODO: should we return error here, or just ignore it? - log.Error("EOF reached while reading p2p message") - break ILOOP - case nil: - bytesToRead -= uint32(n) // TODO: think about avoid the casting in every loop - if bytesToRead <= 0 { - break ILOOP - } - default: - log.Error("Error reading p2p message") - return []byte{}, err - } + + contentLength := int(binary.BigEndian.Uint32(fourBytes)) + tmpBuf := make([]byte, contentLength) + timeoutDuration = 20 * time.Second + s.SetReadDeadline(time.Now().Add(timeoutDuration)) + m, err := io.ReadFull(r, tmpBuf) + if err != nil || m < contentLength { + log.Error("Read %v bytes, we need %v bytes", m, contentLength) + return []byte{}, err } + contentBuf.Write(tmpBuf) return contentBuf.Bytes(), nil }