diff --git a/p2p/src/conn.rs b/p2p/src/conn.rs index e2e400a834..344ddef59c 100644 --- a/p2p/src/conn.rs +++ b/p2p/src/conn.rs @@ -28,7 +28,10 @@ use std::{cmp, thread, time}; use crate::core::ser; use crate::core::ser::FixedLength; -use crate::msg::{read_body, read_header, read_item, write_to_buf, MsgHeader, Type}; +use crate::msg::{ + read_body, read_discard, read_header, read_item, write_to_buf, MsgHeader, MsgHeaderWrapper, + Type, +}; use crate::types::Error; use crate::util::read_write::{read_exact, write_all}; use crate::util::{RateCounter, RwLock}; @@ -251,27 +254,34 @@ fn poll( let mut retry_send = Err(()); loop { // check the read end - if let Some(h) = try_break!(read_header(&mut reader, None)) { - let msg = Message::from_header(h, &mut reader); - - trace!( - "Received message header, type {:?}, len {}.", - msg.header.msg_type, - msg.header.msg_len - ); - - // Increase received bytes counter - let received = received_bytes.clone(); - { - let mut received_bytes = received_bytes.write(); - received_bytes.inc(MsgHeader::LEN as u64 + msg.header.msg_len); + match try_break!(read_header(&mut reader, None)) { + Some(MsgHeaderWrapper::Known(header)) => { + let msg = Message::from_header(header, &mut reader); + + trace!( + "Received message header, type {:?}, len {}.", + msg.header.msg_type, + msg.header.msg_len + ); + + // Increase received bytes counter + received_bytes + .write() + .inc(MsgHeader::LEN as u64 + msg.header.msg_len); + + if let Some(Some(resp)) = + try_break!(handler.consume(msg, &mut writer, received_bytes.clone())) + { + try_break!(resp.write(sent_bytes.clone())); + } } + Some(MsgHeaderWrapper::Unknown(msg_len)) => { + // Increase received bytes counter + received_bytes.write().inc(MsgHeader::LEN as u64 + msg_len); - if let Some(Some(resp)) = - try_break!(handler.consume(msg, &mut writer, received)) - { - try_break!(resp.write(sent_bytes.clone())); + try_break!(read_discard(msg_len, &mut reader)); } + None => {} } // check the write end, use or_else so try_recv is lazily eval'd diff --git a/p2p/src/msg.rs b/p2p/src/msg.rs index 30a7892af1..8a02f6d7f3 100644 --- a/p2p/src/msg.rs +++ b/p2p/src/msg.rs @@ -81,6 +81,11 @@ fn max_block_size() -> u64 { (global::max_block_weight() / consensus::BLOCK_OUTPUT_WEIGHT * 708) as u64 } +// Max msg size when msg type is unknown. +fn default_max_msg_size() -> u64 { + max_block_size() +} + // Max msg size for each msg type. fn max_msg_size(msg_type: Type) -> u64 { match msg_type { @@ -119,24 +124,20 @@ fn magic() -> [u8; 2] { /// Read a header from the provided stream 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(stream: &mut dyn Read, msg_type: Option) -> Result { +/// +/// Note: We return a MsgHeaderWrapper here as we may encounter an unknown msg type. +/// +pub fn read_header( + stream: &mut dyn Read, + msg_type: Option, +) -> Result { let mut head = vec![0u8; MsgHeader::LEN]; if Some(Type::Hand) == msg_type { read_exact(stream, &mut head, time::Duration::from_millis(10), true)?; } else { read_exact(stream, &mut head, time::Duration::from_secs(10), 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 - if header.msg_len > max_len * 4 { - error!( - "Too large read {}, had {}, wanted {}.", - header.msg_type as u8, max_len, header.msg_len - ); - return Err(Error::Serialization(ser::Error::TooLargeReadErr)); - } + let header = ser::deserialize::(&mut &head[..])?; Ok(header) } @@ -158,13 +159,28 @@ pub fn read_body(h: &MsgHeader, stream: &mut dyn Read) -> Result Result<(), Error> { + let mut buffer = vec![0u8; msg_len as usize]; + read_exact(stream, &mut buffer, time::Duration::from_secs(20), true)?; + Ok(()) +} + /// Reads a full message from the underlying stream. pub fn read_message(stream: &mut dyn Read, msg_type: Type) -> Result { - let header = read_header(stream, Some(msg_type))?; - if header.msg_type != msg_type { - return Err(Error::BadMessage); + match read_header(stream, Some(msg_type))? { + MsgHeaderWrapper::Known(header) => { + if header.msg_type == msg_type { + read_body(&header, stream) + } else { + Err(Error::BadMessage) + } + } + MsgHeaderWrapper::Unknown(msg_len) => { + read_discard(msg_len, stream)?; + Err(Error::BadMessage) + } } - read_body(&header, stream) } pub fn write_to_buf(msg: T, msg_type: Type) -> Result, Error> { @@ -191,7 +207,19 @@ pub fn write_message( Ok(()) } +/// A wrapper around a message header. If the header is for an unknown msg type +/// then we will be unable to parse the msg itself (just a bunch of random bytes). +/// But we need to know how many bytes to discard to discard the full message. +#[derive(Clone)] +pub enum MsgHeaderWrapper { + /// A "known" msg type with deserialized msg header. + Known(MsgHeader), + /// An unknown msg type with corresponding msg size in bytes. + Unknown(u64), +} + /// Header of any protocol message, used to identify incoming messages. +#[derive(Clone)] pub struct MsgHeader { magic: [u8; 2], /// Type of the message. @@ -212,7 +240,8 @@ impl MsgHeader { } impl FixedLength for MsgHeader { - const LEN: usize = 1 + 1 + 1 + 8; + // 2 magic bytes + 1 type byte + 8 bytes (msg_len) + const LEN: usize = 2 + 1 + 8; } impl Writeable for MsgHeader { @@ -228,19 +257,49 @@ impl Writeable for MsgHeader { } } -impl Readable for MsgHeader { - fn read(reader: &mut dyn Reader) -> Result { +impl Readable for MsgHeaderWrapper { + fn read(reader: &mut dyn Reader) -> Result { let m = magic(); reader.expect_u8(m[0])?; reader.expect_u8(m[1])?; - let (t, len) = ser_multiread!(reader, read_u8, read_u64); + + // Read the msg header. + // We do not yet know if the msg type is one we support locally. + let (t, msg_len) = ser_multiread!(reader, read_u8, read_u64); + + // Attempt to convert the msg type byte into one of our known msg type enum variants. + // Check the msg_len while we are at it. match Type::from_u8(t) { - Some(ty) => Ok(MsgHeader { - magic: m, - msg_type: ty, - msg_len: len, - }), - None => Err(ser::Error::CorruptedData), + Some(msg_type) => { + // TODO 4x the limits for now to leave ourselves space to change things. + let max_len = max_msg_size(msg_type) * 4; + if msg_len > max_len { + error!( + "Too large read {:?}, max_len: {}, msg_len: {}.", + msg_type, max_len, msg_len + ); + return Err(ser::Error::TooLargeReadErr); + } + + Ok(MsgHeaderWrapper::Known(MsgHeader { + magic: m, + msg_type, + msg_len, + })) + } + None => { + // Unknown msg type, but we still want to limit how big the msg is. + let max_len = default_max_msg_size() * 4; + if msg_len > max_len { + error!( + "Too large read (unknown msg type) {:?}, max_len: {}, msg_len: {}.", + t, max_len, msg_len + ); + return Err(ser::Error::TooLargeReadErr); + } + + Ok(MsgHeaderWrapper::Unknown(msg_len)) + } } } } diff --git a/p2p/src/protocol.rs b/p2p/src/protocol.rs index 200d1f3703..9b3100ce1f 100644 --- a/p2p/src/protocol.rs +++ b/p2p/src/protocol.rs @@ -353,9 +353,8 @@ impl MessageHandler for Protocol { Ok(None) } - - _ => { - debug!("unknown message type {:?}", msg.header.msg_type); + Type::Error | Type::Hand | Type::Shake => { + debug!("Received an unexpected msg: {:?}", msg.header.msg_type); Ok(None) } }