From 471e80e69e813d6ab30d966a4b5e2c01fb8df3a2 Mon Sep 17 00:00:00 2001 From: Gary Yu Date: Thu, 30 Aug 2018 15:21:13 +0800 Subject: [PATCH] fix for peer connect fail on 'WouldBlock' Error (#1440) * fix for fail of peer connect: 'WouldBlock' Error. * reuse the blocking variation of read_exact instead of re-implementing the loop --- p2p/src/conn.rs | 2 +- p2p/src/msg.rs | 10 +++++++--- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/p2p/src/conn.rs b/p2p/src/conn.rs index ec0163ff..6bfc1cd5 100644 --- a/p2p/src/conn.rs +++ b/p2p/src/conn.rs @@ -200,7 +200,7 @@ fn poll( let mut retry_send = Err(()); loop { // check the read end - if let Some(h) = try_break!(error_tx, read_header(conn)) { + if let Some(h) = try_break!(error_tx, read_header(conn, None)) { let msg = Message::from_header(h, conn); trace!( LOGGER, diff --git a/p2p/src/msg.rs b/p2p/src/msg.rs index 90375daf..dbcd9a3b 100644 --- a/p2p/src/msg.rs +++ b/p2p/src/msg.rs @@ -188,9 +188,13 @@ pub fn write_all(conn: &mut Write, mut buf: &[u8], timeout: u32) -> io::Result<( /// Read a header from the provided connection without blocking if the /// underlying stream is async. Typically headers will be polled for, so /// we do not want to block. -pub fn read_header(conn: &mut TcpStream) -> Result { +pub fn read_header(conn: &mut TcpStream, msg_type: Option) -> Result { let mut head = vec![0u8; HEADER_LEN as usize]; - read_exact(conn, &mut head, 10000, false)?; + if Some(Type::Hand) == msg_type { + read_exact(conn, &mut head, 10, true)?; + } else { + read_exact(conn, &mut head, 10000, false)?; + } let header = ser::deserialize::(&mut &head[..])?; let max_len = max_msg_size(header.msg_type); // TODO 4x the limits for now to leave ourselves space to change things @@ -220,7 +224,7 @@ pub fn read_message(conn: &mut TcpStream, msg_type: Type) -> Result where T: Readable, { - let header = read_header(conn)?; + let header = read_header(conn, Some(msg_type))?; if header.msg_type != msg_type { return Err(Error::BadMessage); }