Merge branch 'master' into grim

# Conflicts:
#	p2p/Cargo.toml

core/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "grin_core"
-version = "5.4.0"
+version = "5.4.1"
 authors = ["Grin Developers <mimblewimble@lists.launchpad.net>"]
 description = "Chain implementation for grin, a simple, private and scalable cryptocurrency implementation based on the Mimblewimble chain format."
 license = "Apache-2.0"
@@ -28,8 +28,8 @@ chrono = { version = "0.4.11", features = ["serde"] }
 zeroize = { version = "1.1", features =["zeroize_derive"] }
 bytes = "0.5"
 
-keychain = { package = "grin_keychain", path = "../keychain", version = "5.4.0" }
-util = { package = "grin_util", path = "../util", version = "5.4.0" }
+keychain = { package = "grin_keychain", path = "../keychain", version = "5.4.1" }
+util = { package = "grin_util", path = "../util", version = "5.4.1" }
 
 [dev-dependencies]
 serde_json = "1"

core/src/core/pmmr/segment.rs

@@ -21,6 +21,39 @@ use croaring::Bitmap;
 use std::cmp::min;
 use std::fmt::Debug;
 
+const MAX_SEGMENT_READ_ITEMS: u64 = 1_000_000;
+const SEGMENT_READ_PREALLOC_ITEMS: u64 = 1024;
+
+fn read_segment_item_count<R: Reader>(reader: &mut R) -> Result<u64, Error> {
+	let count = reader.read_u64()?;
+	if count > MAX_SEGMENT_READ_ITEMS {
+		return Err(Error::TooLargeReadErr);
+	}
+	Ok(count)
+}
+
+fn read_segment_positions<R: Reader>(reader: &mut R, count: u64) -> Result<Vec<u64>, Error> {
+	let mut positions = Vec::with_capacity(min(count, SEGMENT_READ_PREALLOC_ITEMS) as usize);
+	let mut last_pos = 0;
+	for _ in 0..count {
+		let pos = reader.read_u64()?;
+		if pos <= last_pos {
+			return Err(Error::SortError);
+		}
+		last_pos = pos;
+		positions.push(pos - 1);
+	}
+	Ok(positions)
+}
+
+fn read_segment_items<T: Readable, R: Reader>(reader: &mut R, count: u64) -> Result<Vec<T>, Error> {
+	let mut items = Vec::with_capacity(min(count, SEGMENT_READ_PREALLOC_ITEMS) as usize);
+	for _ in 0..count {
+		items.push(T::read(reader)?);
+	}
+	Ok(items)
+}
+
 #[derive(Clone, Debug, Eq, PartialEq)]
 /// Possible segment types, according to this desegmenter
 pub enum SegmentType {
@@ -568,39 +601,13 @@ impl<T: Readable> Readable for Segment<T> {
 	fn read<R: Reader>(reader: &mut R) -> Result<Self, Error> {
 		let identifier = Readable::read(reader)?;
 
-		let n_hashes = reader.read_u64()? as usize;
-		let mut hash_pos = Vec::with_capacity(n_hashes);
-		let mut last_pos = 0;
-		for _ in 0..n_hashes {
-			let pos = reader.read_u64()?;
-			if pos <= last_pos {
-				return Err(Error::SortError);
-			}
-			last_pos = pos;
-			hash_pos.push(pos - 1);
-		}
+		let n_hashes = read_segment_item_count(reader)?;
+		let hash_pos = read_segment_positions(reader, n_hashes)?;
+		let hashes = read_segment_items(reader, n_hashes)?;
 
-		let mut hashes = Vec::<Hash>::with_capacity(n_hashes);
-		for _ in 0..n_hashes {
-			hashes.push(Readable::read(reader)?);
-		}
-
-		let n_leaves = reader.read_u64()? as usize;
-		let mut leaf_pos = Vec::with_capacity(n_leaves);
-		last_pos = 0;
-		for _ in 0..n_leaves {
-			let pos = reader.read_u64()?;
-			if pos <= last_pos {
-				return Err(Error::SortError);
-			}
-			last_pos = pos;
-			leaf_pos.push(pos - 1);
-		}
-
-		let mut leaf_data = Vec::<T>::with_capacity(n_leaves);
-		for _ in 0..n_leaves {
-			leaf_data.push(Readable::read(reader)?);
-		}
+		let n_leaves = read_segment_item_count(reader)?;
+		let leaf_pos = read_segment_positions(reader, n_leaves)?;
+		let leaf_data = read_segment_items(reader, n_leaves)?;
 
 		let proof = Readable::read(reader)?;
 
@@ -823,12 +830,8 @@ impl SegmentProof {
 
 impl Readable for SegmentProof {
 	fn read<R: Reader>(reader: &mut R) -> Result<Self, Error> {
-		let n_hashes = reader.read_u64()? as usize;
-		let mut hashes = Vec::with_capacity(n_hashes);
-		for _ in 0..n_hashes {
-			let hash: Hash = Readable::read(reader)?;
-			hashes.push(hash);
-		}
+		let n_hashes = read_segment_item_count(reader)?;
+		let hashes = read_segment_items(reader, n_hashes)?;
 		Ok(Self { hashes })
 	}
 }

core/tests/segment.rs

@@ -16,10 +16,15 @@ mod common;
 
 use self::core::core::pmmr;
 use self::core::core::{Segment, SegmentIdentifier};
+use self::core::ser::{self, DeserializationMode, ProtocolVersion};
 use common::TestElem;
 use grin_core as core;
 use grin_core::core::pmmr::ReadablePMMR;
 
+fn push_u64(bytes: &mut Vec<u8>, n: u64) {
+	bytes.extend_from_slice(&n.to_be_bytes());
+}
+
 fn test_unprunable_size(height: u8, n_leaves: u32) {
 	let size = 1u64 << height;
 	let n_segments = (n_leaves as u64 + size - 1) / size;
@@ -59,3 +64,30 @@ fn unprunable_mmr() {
 		test_unprunable_size(3, i);
 	}
 }
+
+#[test]
+fn segment_read_rejects_large_hash_count() {
+	let mut bytes = vec![1];
+	push_u64(&mut bytes, 0);
+	push_u64(&mut bytes, 1_000_001);
+
+	let res: Result<Segment<TestElem>, _> = ser::deserialize(
+		&mut &bytes[..],
+		ProtocolVersion(1),
+		DeserializationMode::default(),
+	);
+	assert_eq!(res.err(), Some(ser::Error::TooLargeReadErr));
+}
+
+#[test]
+fn segment_proof_read_rejects_large_hash_count() {
+	let mut bytes = vec![];
+	push_u64(&mut bytes, 1_000_001);
+
+	let res: Result<self::core::core::SegmentProof, _> = ser::deserialize(
+		&mut &bytes[..],
+		ProtocolVersion(1),
+		DeserializationMode::default(),
+	);
+	assert_eq!(res.err(), Some(ser::Error::TooLargeReadErr));
+}