Support code blobs compressed with zstd (#8549)

* begin maybe-compressed-blob

* fix build

* implement blob compression / decompression

* add some tests

* decode -> decompress

* decompress code if compressed

* make API of compresseed blob crate take limit as parameter

* use new API in sc-executro

* wasm-builder: compress wasm

* fix typo

* simplify

* address review

* fix wasm_project.rs

* Update primitives/maybe-compressed-blob/Cargo.toml

Co-authored-by: Andronik Ordian <write@reusable.software>

Co-authored-by: Andronik Ordian <write@reusable.software>

substrate/Cargo.lock

@@ -6805,6 +6805,16 @@ version = "1.0.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "cb5d2a036dc6d2d8fd16fde3498b04306e29bd193bf306a57427019b823d5acd"
 
+[[package]]
+name = "ruzstd"
+version = "0.2.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3d425143485a37727c7a46e689bbe3b883a00f42b4a52c4ac0f44855c1009b00"
+dependencies = [
+ "byteorder",
+ "twox-hash",
+]
+
 [[package]]
 name = "rw-stream-sink"
 version = "0.2.1"
@@ -7337,6 +7347,7 @@ dependencies = [
  "sp-core",
  "sp-externalities",
  "sp-io",
+ "sp-maybe-compressed-blob",
  "sp-panic-handler",
  "sp-runtime",
  "sp-runtime-interface",
@@ -8876,6 +8887,14 @@ dependencies = [
  "sp-externalities",
 ]
 
+[[package]]
+name = "sp-maybe-compressed-blob"
+version = "3.0.0"
+dependencies = [
+ "ruzstd",
+ "zstd",
+]
+
 [[package]]
 name = "sp-npos-elections"
 version = "3.0.0"
@@ -9597,6 +9616,7 @@ dependencies = [
  "atty",
  "build-helper",
  "cargo_metadata",
+ "sp-maybe-compressed-blob",
  "tempfile",
  "toml",
  "walkdir",

substrate/Cargo.toml

@@ -154,6 +154,7 @@ members = [
 	"primitives/io",
 	"primitives/keyring",
 	"primitives/keystore",
+	"primitives/maybe-compressed-blob",
 	"primitives/npos-elections",
 	"primitives/npos-elections/compact",
 	"primitives/npos-elections/fuzzer",

substrate/client/executor/Cargo.toml

@@ -30,6 +30,7 @@ sp-api = { version = "3.0.0", path = "../../primitives/api" }
 sp-wasm-interface = { version = "3.0.0", path = "../../primitives/wasm-interface" }
 sp-runtime-interface = { version = "3.0.0", path = "../../primitives/runtime-interface" }
 sp-externalities = { version = "0.9.0", path = "../../primitives/externalities" }
+sp-maybe-compressed-blob = { version = "3.0.0", path = "../../primitives/maybe-compressed-blob" }
 sc-executor-common = { version = "0.9.0", path = "common" }
 sc-executor-wasmi = { version = "0.9.0", path = "wasmi" }
 sc-executor-wasmtime = { version = "0.9.0", path = "wasmtime", optional = true }

substrate/client/executor/src/wasm_runtime.rs

@@ -283,6 +283,11 @@ pub fn create_wasm_runtime_with_code(
 	allow_missing_func_imports: bool,
 	cache_path: Option<&Path>,
 ) -> Result<Arc<dyn WasmModule>, WasmError> {
+	use sp_maybe_compressed_blob::CODE_BLOB_BOMB_LIMIT;
+
+	let code = sp_maybe_compressed_blob::decompress(code, CODE_BLOB_BOMB_LIMIT)
+		.map_err(|e| WasmError::Other(format!("Decompression error: {:?}", e)))?;
+
 	match wasm_method {
 		WasmExecutionMethod::Interpreted => {
 			// Wasmi doesn't have any need in a cache directory.
@@ -292,7 +297,7 @@ pub fn create_wasm_runtime_with_code(
 			drop(cache_path);
 
 			sc_executor_wasmi::create_runtime(
-				code,
+				&code,
 				heap_pages,
 				host_functions,
 				allow_missing_func_imports,
@@ -301,7 +306,7 @@ pub fn create_wasm_runtime_with_code(
 		}
 		#[cfg(feature = "wasmtime")]
 		WasmExecutionMethod::Compiled => {
-			let blob = sc_executor_common::runtime_blob::RuntimeBlob::new(code)?;
+			let blob = sc_executor_common::runtime_blob::RuntimeBlob::new(&code)?;
 			sc_executor_wasmtime::create_runtime(
 				sc_executor_wasmtime::CodeSupplyMode::Verbatim { blob },
 				sc_executor_wasmtime::Config {

substrate/primitives/maybe-compressed-blob/Cargo.toml

@@ -0,0 +1,17 @@
+[package]
+name = "sp-maybe-compressed-blob"
+version = "3.0.0"
+authors = ["Parity Technologies <admin@parity.io>"]
+edition = "2018"
+license = "Apache-2.0"
+homepage = "https://substrate.dev"
+repository = "https://github.com/paritytech/substrate/"
+description = "Handling of blobs, usually Wasm code, which may be compresed"
+documentation = "https://docs.rs/sp-maybe-compressed-blob"
+readme = "README.md"
+
+[target.'cfg(not(target_os = "unknown"))'.dependencies]
+zstd = { version = "0.6.0", default-features = false }
+
+[target.'cfg(target_os = "unknown")'.dependencies]
+ruzstd = { version = "0.2.2" }

substrate/primitives/maybe-compressed-blob/README.md

@@ -0,0 +1,3 @@
+Handling of blobs, typicaly validation code, which may be compressed.
+
+License: Apache-2.0

substrate/primitives/maybe-compressed-blob/src/lib.rs

@@ -0,0 +1,166 @@
+// This file is part of Substrate.
+
+// Copyright (C) 2017-2021 Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: Apache-2.0
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// 	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Handling of blobs that may be compressed, based on an 8-byte magic identifier
+//! at the head.
+
+use std::borrow::Cow;
+use std::io::Read;
+
+// An arbitrary prefix, that indicates a blob beginning with should be decompressed with
+// Zstd compression.
+//
+// This differs from the WASM magic bytes, so real WASM blobs will not have this prefix.
+const ZSTD_PREFIX: [u8; 8] = [82, 188, 83, 118, 70, 219, 142, 5];
+
+/// A recommendation for the bomb limit for code blobs.
+///
+/// This may be adjusted upwards in the future, but is set much higher than the
+/// expected maximum code size. When adjusting upwards, nodes should be updated
+/// before performing a runtime upgrade to a blob with larger compressed size.
+pub const CODE_BLOB_BOMB_LIMIT: usize = 50 * 1024 * 1024;
+
+/// A possible bomb was encountered.
+#[derive(Debug, Clone, PartialEq)]
+pub enum Error {
+	/// Decoded size was too large, and the code payload may be a bomb.
+	PossibleBomb,
+	/// The compressed value had an invalid format.
+	Invalid,
+}
+
+impl std::fmt::Display for Error {
+	fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+		match *self {
+			Error::PossibleBomb => write!(f, "Possible compression bomb encountered"),
+			Error::Invalid => write!(f, "Blob had invalid format"),
+		}
+	}
+}
+
+impl std::error::Error for Error { }
+
+fn read_from_decoder(
+	decoder: impl Read,
+	blob_len: usize,
+	bomb_limit: usize,
+) -> Result<Vec<u8>, Error> {
+	let mut decoder = decoder.take((bomb_limit + 1) as u64);
+
+	let mut buf = Vec::with_capacity(blob_len);
+	decoder.read_to_end(&mut buf).map_err(|_| Error::Invalid)?;
+
+	if buf.len() <= bomb_limit {
+		Ok(buf)
+	} else {
+		Err(Error::PossibleBomb)
+	}
+}
+
+#[cfg(not(target_os = "unknown"))]
+fn decompress_zstd(blob: &[u8], bomb_limit: usize) -> Result<Vec<u8>, Error> {
+	let decoder = zstd::Decoder::new(blob).map_err(|_| Error::Invalid)?;
+
+	read_from_decoder(decoder, blob.len(), bomb_limit)
+}
+
+#[cfg(target_os = "unknown")]
+fn decompress_zstd(mut blob: &[u8], bomb_limit: usize) -> Result<Vec<u8>, Error> {
+	let blob_len = blob.len();
+	let decoder = ruzstd::streaming_decoder::StreamingDecoder::new(&mut blob)
+		.map_err(|_| Error::Invalid)?;
+
+	read_from_decoder(decoder, blob_len, bomb_limit)
+}
+
+/// Decode a blob, if it indicates that it is compressed. Provide a `bomb_limit`, which
+/// is the limit of bytes which should be decompressed from the blob.
+pub fn decompress(blob: &[u8], bomb_limit: usize) -> Result<Cow<[u8]>, Error> {
+	if blob.starts_with(&ZSTD_PREFIX) {
+		decompress_zstd(&blob[ZSTD_PREFIX.len()..], bomb_limit).map(Into::into)
+	} else {
+		Ok(blob.into())
+	}
+}
+
+/// Encode a blob as compressed. If the blob's size is over the bomb limit,
+/// this will not compress the blob, as the decoder will not be able to be
+/// able to differentiate it from a compression bomb.
+#[cfg(not(target_os = "unknown"))]
+pub fn compress(blob: &[u8], bomb_limit: usize) -> Option<Vec<u8>> {
+	use std::io::Write;
+
+	if blob.len() > bomb_limit {
+		return None;
+	}
+
+	let mut buf = ZSTD_PREFIX.to_vec();
+
+	{
+		let mut v = zstd::Encoder::new(&mut buf, 3).ok()?.auto_finish();
+		v.write_all(blob).ok()?;
+	}
+
+	Some(buf)
+}
+
+#[cfg(test)]
+mod tests {
+	use super::*;
+	use std::io::Write;
+
+	const BOMB_LIMIT: usize = 10;
+
+	#[test]
+	fn refuse_to_encode_over_limit() {
+		let mut v = vec![0; BOMB_LIMIT + 1];
+		assert!(compress(&v, BOMB_LIMIT).is_none());
+
+		let _ = v.pop();
+		assert!(compress(&v, BOMB_LIMIT).is_some());
+	}
+
+	#[test]
+	fn compress_and_decompress() {
+		let v = vec![0; BOMB_LIMIT];
+
+		let compressed = compress(&v, BOMB_LIMIT).unwrap();
+
+		assert!(compressed.starts_with(&ZSTD_PREFIX));
+		assert_eq!(&decompress(&compressed, BOMB_LIMIT).unwrap()[..], &v[..])
+	}
+
+	#[test]
+	fn decompresses_only_when_magic() {
+		let v = vec![0; BOMB_LIMIT + 1];
+
+		assert_eq!(&decompress(&v, BOMB_LIMIT).unwrap()[..], &v[..]);
+	}
+
+	#[test]
+	fn possible_bomb_fails() {
+		let encoded_bigger_than_bomb = vec![0; BOMB_LIMIT + 1];
+		let mut buf = ZSTD_PREFIX.to_vec();
+
+		{
+			let mut v = zstd::Encoder::new(&mut buf, 3).unwrap().auto_finish();
+			v.write_all(&encoded_bigger_than_bomb[..]).unwrap();
+		}
+
+		assert_eq!(decompress(&buf[..], BOMB_LIMIT).err(), Some(Error::PossibleBomb));
+	}
+}

substrate/utils/wasm-builder/Cargo.toml

@@ -21,3 +21,4 @@ walkdir = "2.3.1"
 wasm-gc-api = "0.1.11"
 atty = "0.2.13"
 ansi_term = "0.12.1"
+sp-maybe-compressed-blob = { version = "3.0.0", path = "../../primitives/maybe-compressed-blob" }

substrate/utils/wasm-builder/src/builder.rs

@@ -217,7 +217,7 @@ fn generate_rerun_if_changed_instructions() {
 /// `project_cargo_toml` - The path to the `Cargo.toml` of the project that should be built.
 /// `default_rustflags` - Default `RUSTFLAGS` that will always be set for the build.
 /// `features_to_enable` - Features that should be enabled for the project.
-/// `wasm_binary_name` - The optional wasm binary name that is extended with `.compact.wasm`.
+/// `wasm_binary_name` - The optional wasm binary name that is extended with `.compact.compressed.wasm`.
 ///                      If `None`, the project name will be used.
 fn build_project(
 	file_name: PathBuf,

substrate/utils/wasm-builder/src/wasm_project.rs

@@ -114,7 +114,7 @@ pub(crate) fn create_and_compile(
 	);
 
 	build_project(&project, default_rustflags, cargo_cmd);
-	let (wasm_binary, bloaty) = compact_wasm_file(
+	let (wasm_binary, wasm_binary_compressed, bloaty) = compact_wasm_file(
 		&project,
 		project_cargo_toml,
 		wasm_binary_name,
@@ -124,9 +124,13 @@ pub(crate) fn create_and_compile(
 		copy_wasm_to_target_directory(project_cargo_toml, wasm_binary)
 	);
 
+	wasm_binary_compressed.as_ref().map(|wasm_binary_compressed|
+		copy_wasm_to_target_directory(project_cargo_toml, wasm_binary_compressed)
+	);
+
 	generate_rerun_if_changed_instructions(project_cargo_toml, &project, &wasm_workspace);
 
-	(wasm_binary, bloaty)
+	(wasm_binary_compressed.or(wasm_binary), bloaty)
 }
 
 /// Find the `Cargo.lock` relative to the `OUT_DIR` environment variable.
@@ -441,12 +445,12 @@ fn build_project(project: &Path, default_rustflags: &str, cargo_cmd: CargoComman
 	}
 }
 
-/// Compact the WASM binary using `wasm-gc`. Returns the path to the bloaty WASM binary.
+/// Compact the WASM binary using `wasm-gc` and compress it using zstd.
 fn compact_wasm_file(
 	project: &Path,
 	cargo_manifest: &Path,
 	wasm_binary_name: Option<String>,
-) -> (Option<WasmBinary>, WasmBinaryBloaty) {
+) -> (Option<WasmBinary>, Option<WasmBinary>, WasmBinaryBloaty) {
 	let is_release_build = is_release_build();
 	let target = if is_release_build { "release" } else { "debug" };
 	let default_wasm_binary_name = get_wasm_binary_name(cargo_manifest);
@@ -468,6 +472,25 @@ fn compact_wasm_file(
 		None
 	};
 
+	let wasm_compact_compressed_file = wasm_compact_file.as_ref()
+		.and_then(|compact_binary| {
+			let file_name = wasm_binary_name.clone()
+				.unwrap_or_else(|| default_wasm_binary_name.clone());
+
+			let wasm_compact_compressed_file = project.join(
+				format!(
+					"{}.compact.compressed.wasm",
+					file_name,
+				)
+			);
+
+			if compress_wasm(&compact_binary.0, &wasm_compact_compressed_file) {
+				Some(WasmBinary(wasm_compact_compressed_file))
+			} else {
+				None
+			}
+		});
+
 	let bloaty_file_name = if let Some(name) = wasm_binary_name {
 		format!("{}.wasm", name)
 	} else {
@@ -477,7 +500,36 @@ fn compact_wasm_file(
 	let bloaty_file = project.join(bloaty_file_name);
 	fs::copy(wasm_file, &bloaty_file).expect("Copying the bloaty file to the project dir.");
 
-	(wasm_compact_file, WasmBinaryBloaty(bloaty_file))
+	(
+		wasm_compact_file,
+		wasm_compact_compressed_file,
+		WasmBinaryBloaty(bloaty_file),
+	)
+}
+
+fn compress_wasm(
+	wasm_binary_path: &Path,
+	compressed_binary_out_path: &Path,
+) -> bool {
+	use sp_maybe_compressed_blob::CODE_BLOB_BOMB_LIMIT;
+
+	let data = fs::read(wasm_binary_path).expect("Failed to read WASM binary");
+	if let Some(compressed) = sp_maybe_compressed_blob::compress(
+		&data,
+		CODE_BLOB_BOMB_LIMIT,
+	) {
+		fs::write(compressed_binary_out_path, &compressed[..])
+			.expect("Failed to write WASM binary");
+
+		true
+	} else {
+		println!(
+			"cargo:warning=Writing uncompressed wasm. Exceeded maximum size {}",
+			CODE_BLOB_BOMB_LIMIT,
+		);
+
+		false
+	}
 }
 
 /// Custom wrapper for a [`cargo_metadata::Package`] to store it in