diff --git a/node/network/protocol/src/lib.rs b/node/network/protocol/src/lib.rs index 5d0702c524a5..075a0d9decfe 100644 --- a/node/network/protocol/src/lib.rs +++ b/node/network/protocol/src/lib.rs @@ -357,73 +357,6 @@ pub mod v1 { Approvals(Vec), } - #[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)] - #[allow(missing_docs)] - pub enum CompressedPoVError { - #[error("Failed to compress a PoV")] - Compress, - #[error("Failed to decompress a PoV")] - Decompress, - #[error("Failed to decode the uncompressed PoV")] - Decode, - #[error("Architecture is not supported")] - NotSupported, - } - - /// SCALE and Zstd encoded [`PoV`]. - #[derive(Clone, Encode, Decode, PartialEq, Eq)] - pub struct CompressedPoV(Vec); - - impl CompressedPoV { - /// Compress the given [`PoV`] and returns a [`CompressedPoV`]. - #[cfg(not(target_os = "unknown"))] - pub fn compress(pov: &PoV) -> Result { - zstd::encode_all(pov.encode().as_slice(), 3).map_err(|_| CompressedPoVError::Compress).map(Self) - } - - /// Compress the given [`PoV`] and returns a [`CompressedPoV`]. - #[cfg(target_os = "unknown")] - pub fn compress(_: &PoV) -> Result { - Err(CompressedPoVError::NotSupported) - } - - /// Decompress `self` and returns the [`PoV`] on success. - #[cfg(not(target_os = "unknown"))] - pub fn decompress(&self) -> Result { - use std::io::Read; - const MAX_POV_BLOCK_SIZE: usize = 32 * 1024 * 1024; - - struct InputDecoder<'a, T: std::io::BufRead>(&'a mut zstd::Decoder, usize); - impl<'a, T: std::io::BufRead> parity_scale_codec::Input for InputDecoder<'a, T> { - fn read(&mut self, into: &mut [u8]) -> Result<(), parity_scale_codec::Error> { - self.1 = self.1.saturating_add(into.len()); - if self.1 > MAX_POV_BLOCK_SIZE { - return Err("pov block too big".into()) - } - self.0.read_exact(into).map_err(Into::into) - } - fn remaining_len(&mut self) -> Result, parity_scale_codec::Error> { - Ok(None) - } - } - - let mut decoder = zstd::Decoder::new(self.0.as_slice()).map_err(|_| CompressedPoVError::Decompress)?; - PoV::decode(&mut InputDecoder(&mut decoder, 0)).map_err(|_| CompressedPoVError::Decode) - } - - /// Decompress `self` and returns the [`PoV`] on success. - #[cfg(target_os = "unknown")] - pub fn decompress(&self) -> Result { - Err(CompressedPoVError::NotSupported) - } - } - - impl std::fmt::Debug for CompressedPoV { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "CompressedPoV({} bytes)", self.0.len()) - } - } - /// Network messages used by the collator protocol subsystem #[derive(Debug, Clone, Encode, Decode, PartialEq, Eq)] pub enum CollatorProtocolMessage { @@ -481,17 +414,3 @@ pub mod v1 { impl_try_from!(CollationProtocol, CollatorProtocol, CollatorProtocolMessage); } - -#[cfg(test)] -mod tests { - use polkadot_primitives::v1::PoV; - use super::v1::{CompressedPoV, CompressedPoVError}; - - #[test] - fn decompress_huge_pov_block_fails() { - let pov = PoV { block_data: vec![0; 63 * 1024 * 1024].into() }; - - let compressed = CompressedPoV::compress(&pov).unwrap(); - assert_eq!(CompressedPoVError::Decode, compressed.decompress().unwrap_err()); - } -} diff --git a/primitives/Cargo.toml b/primitives/Cargo.toml index 8f12f6dcaa32..44875c16672a 100644 --- a/primitives/Cargo.toml +++ b/primitives/Cargo.toml @@ -26,6 +26,9 @@ bitvec = { version = "0.20.1", default-features = false, features = ["alloc"] } frame-system = { git = "https://github.com/paritytech/substrate", branch = "master", default-features = false } hex-literal = "0.3.1" parity-util-mem = { version = "0.9.0", default-features = false, optional = true } +thiserror = "1.0.23" +[target.'cfg(not(target_os = "unknown"))'.dependencies] +zstd = "0.5.0" [dev-dependencies] sp-serializer = { git = "https://github.com/paritytech/substrate", branch = "master" } diff --git a/primitives/src/v1.rs b/primitives/src/v1.rs index c3d171338b3f..4987d8e2c956 100644 --- a/primitives/src/v1.rs +++ b/primitives/src/v1.rs @@ -454,6 +454,73 @@ impl PoV { } } +/// SCALE and Zstd encoded [`PoV`]. +#[derive(Clone, Encode, Decode, PartialEq, Eq)] +pub struct CompressedPoV(Vec); + +#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)] +#[allow(missing_docs)] +pub enum CompressedPoVError { + #[error("Failed to compress a PoV")] + Compress, + #[error("Failed to decompress a PoV")] + Decompress, + #[error("Failed to decode the uncompressed PoV")] + Decode, + #[error("Architecture is not supported")] + NotSupported, +} + +impl CompressedPoV { + /// Compress the given [`PoV`] and returns a [`CompressedPoV`]. + #[cfg(not(target_os = "unknown"))] + pub fn compress(pov: &PoV) -> Result { + zstd::encode_all(pov.encode().as_slice(), 3).map_err(|_| CompressedPoVError::Compress).map(Self) + } + + /// Compress the given [`PoV`] and returns a [`CompressedPoV`]. + #[cfg(target_os = "unknown")] + pub fn compress(_: &PoV) -> Result { + Err(CompressedPoVError::NotSupported) + } + + /// Decompress `self` and returns the [`PoV`] on success. + #[cfg(not(target_os = "unknown"))] + pub fn decompress(&self) -> Result { + use std::io::Read; + const MAX_POV_BLOCK_SIZE: usize = 32 * 1024 * 1024; + + struct InputDecoder<'a, T: std::io::BufRead>(&'a mut zstd::Decoder, usize); + impl<'a, T: std::io::BufRead> parity_scale_codec::Input for InputDecoder<'a, T> { + fn read(&mut self, into: &mut [u8]) -> Result<(), parity_scale_codec::Error> { + self.1 = self.1.saturating_add(into.len()); + if self.1 > MAX_POV_BLOCK_SIZE { + return Err("pov block too big".into()) + } + self.0.read_exact(into).map_err(Into::into) + } + fn remaining_len(&mut self) -> Result, parity_scale_codec::Error> { + Ok(None) + } + } + + let mut decoder = zstd::Decoder::new(self.0.as_slice()).map_err(|_| CompressedPoVError::Decompress)?; + PoV::decode(&mut InputDecoder(&mut decoder, 0)).map_err(|_| CompressedPoVError::Decode) + } + + /// Decompress `self` and returns the [`PoV`] on success. + #[cfg(target_os = "unknown")] + pub fn decompress(&self) -> Result { + Err(CompressedPoVError::NotSupported) + } +} + +impl std::fmt::Debug for CompressedPoV { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "CompressedPoV({} bytes)", self.0.len()) + } +} + /// A bitfield concerning availability of backed candidates. #[derive(PartialEq, Eq, Clone, Encode, Decode, RuntimeDebug)] pub struct AvailabilityBitfield(pub BitVec); @@ -982,6 +1049,7 @@ pub struct AbridgedHrmpChannel { #[cfg(test)] mod tests { use super::*; + use super::{CompressedPoV, CompressedPoVError, PoV}; #[test] fn group_rotation_info_calculations() { @@ -1008,4 +1076,13 @@ mod tests { &Hash::repeat_byte(3), ); } + + + #[test] + fn decompress_huge_pov_block_fails() { + let pov = PoV { block_data: vec![0; 63 * 1024 * 1024].into() }; + + let compressed = CompressedPoV::compress(&pov).unwrap(); + assert_eq!(CompressedPoVError::Decode, compressed.decompress().unwrap_err()); + } }