sp_maybe_compressed_blob/
lib.rs

1// This file is part of Substrate.
2
3// Copyright (C) Parity Technologies (UK) Ltd.
4// SPDX-License-Identifier: Apache-2.0
5
6// Licensed under the Apache License, Version 2.0 (the "License");
7// you may not use this file except in compliance with the License.
8// You may obtain a copy of the License at
9//
10// 	http://www.apache.org/licenses/LICENSE-2.0
11//
12// Unless required by applicable law or agreed to in writing, software
13// distributed under the License is distributed on an "AS IS" BASIS,
14// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15// See the License for the specific language governing permissions and
16// limitations under the License.
17
18//! Handling of blobs that may be compressed, based on an 8-byte magic identifier
19//! at the head.
20
21use std::{
22	borrow::Cow,
23	io::{Read, Write},
24};
25
26// An arbitrary prefix, that indicates a blob beginning with should be decompressed with
27// Zstd compression.
28//
29// This differs from the WASM magic bytes, so real WASM blobs will not have this prefix.
30const ZSTD_PREFIX: [u8; 8] = [82, 188, 83, 118, 70, 219, 142, 5];
31
32/// A recommendation for the bomb limit for code blobs.
33///
34/// This may be adjusted upwards in the future, but is set much higher than the
35/// expected maximum code size. When adjusting upwards, nodes should be updated
36/// before performing a runtime upgrade to a blob with larger compressed size.
37pub const CODE_BLOB_BOMB_LIMIT: usize = 50 * 1024 * 1024;
38
39/// A possible bomb was encountered.
40#[derive(Debug, Clone, PartialEq, thiserror::Error)]
41pub enum Error {
42	/// Decoded size was too large, and the code payload may be a bomb.
43	#[error("Possible compression bomb encountered")]
44	PossibleBomb,
45	/// The compressed value had an invalid format.
46	#[error("Blob had invalid format")]
47	Invalid,
48}
49
50fn read_from_decoder(
51	decoder: impl Read,
52	blob_len: usize,
53	bomb_limit: usize,
54) -> Result<Vec<u8>, Error> {
55	let mut decoder = decoder.take((bomb_limit + 1) as u64);
56
57	let mut buf = Vec::with_capacity(blob_len);
58	decoder.read_to_end(&mut buf).map_err(|_| Error::Invalid)?;
59
60	if buf.len() <= bomb_limit {
61		Ok(buf)
62	} else {
63		Err(Error::PossibleBomb)
64	}
65}
66
67fn decompress_zstd(blob: &[u8], bomb_limit: usize) -> Result<Vec<u8>, Error> {
68	let decoder = zstd::Decoder::new(blob).map_err(|_| Error::Invalid)?;
69
70	read_from_decoder(decoder, blob.len(), bomb_limit)
71}
72
73/// Decode a blob, if it indicates that it is compressed. Provide a `bomb_limit`, which
74/// is the limit of bytes which should be decompressed from the blob.
75pub fn decompress(blob: &[u8], bomb_limit: usize) -> Result<Cow<'_, [u8]>, Error> {
76	if blob.starts_with(&ZSTD_PREFIX) {
77		decompress_zstd(&blob[ZSTD_PREFIX.len()..], bomb_limit).map(Into::into)
78	} else {
79		Ok(blob.into())
80	}
81}
82
83/// Weakly compress a blob who's size is limited by `bomb_limit`.
84///
85/// If the blob's size is over the bomb limit, this will not compress the blob,
86/// as the decoder will not be able to be able to differentiate it from a compression bomb.
87pub fn compress_weakly(blob: &[u8], bomb_limit: usize) -> Option<Vec<u8>> {
88	compress_with_level(blob, bomb_limit, 3)
89}
90
91/// Strongly compress a blob who's size is limited by `bomb_limit`.
92///
93/// If the blob's size is over the bomb limit, this will not compress the blob, as the decoder will
94/// not be able to be able to differentiate it from a compression bomb.
95pub fn compress_strongly(blob: &[u8], bomb_limit: usize) -> Option<Vec<u8>> {
96	compress_with_level(blob, bomb_limit, 22)
97}
98
99/// Compress a blob who's size is limited by `bomb_limit`.
100///
101/// If the blob's size is over the bomb limit, this will not compress the blob, as the decoder will
102/// not be able to be able to differentiate it from a compression bomb.
103#[deprecated(
104	note = "Will be removed after June 2026. Use compress_strongly, compress_weakly or compress_with_level instead"
105)]
106pub fn compress(blob: &[u8], bomb_limit: usize) -> Option<Vec<u8>> {
107	compress_with_level(blob, bomb_limit, 3)
108}
109
110/// Compress a blob who's size is limited by `bomb_limit` with adjustable compression level.
111///
112/// The levels are passed through to `zstd` and can be in range [1, 22] (weakest to strongest).
113///
114/// If the blob's size is over the bomb limit, this will not compress the blob, as the decoder will
115/// not be able to be able to differentiate it from a compression bomb.
116fn compress_with_level(blob: &[u8], bomb_limit: usize, level: i32) -> Option<Vec<u8>> {
117	if blob.len() > bomb_limit {
118		return None
119	}
120
121	let mut buf = ZSTD_PREFIX.to_vec();
122
123	{
124		let mut v = zstd::Encoder::new(&mut buf, level).ok()?.auto_finish();
125		v.write_all(blob).ok()?;
126	}
127
128	Some(buf)
129}
130
131#[cfg(test)]
132mod tests {
133	use super::*;
134
135	const BOMB_LIMIT: usize = 10;
136
137	#[test]
138	fn refuse_to_encode_over_limit() {
139		let mut v = vec![0; BOMB_LIMIT + 1];
140		assert!(compress_weakly(&v, BOMB_LIMIT).is_none());
141		assert!(compress_strongly(&v, BOMB_LIMIT).is_none());
142
143		let _ = v.pop();
144		assert!(compress_weakly(&v, BOMB_LIMIT).is_some());
145		assert!(compress_strongly(&v, BOMB_LIMIT).is_some());
146	}
147
148	#[test]
149	fn compress_and_decompress() {
150		let v = vec![0; BOMB_LIMIT];
151
152		let compressed_weakly = compress_weakly(&v, BOMB_LIMIT).unwrap();
153		let compressed_strongly = compress_strongly(&v, BOMB_LIMIT).unwrap();
154
155		assert!(compressed_weakly.starts_with(&ZSTD_PREFIX));
156		assert!(compressed_strongly.starts_with(&ZSTD_PREFIX));
157
158		assert_eq!(&decompress(&compressed_weakly, BOMB_LIMIT).unwrap()[..], &v[..]);
159		assert_eq!(&decompress(&compressed_strongly, BOMB_LIMIT).unwrap()[..], &v[..]);
160	}
161
162	#[test]
163	fn decompresses_only_when_magic() {
164		let v = vec![0; BOMB_LIMIT + 1];
165
166		assert_eq!(&decompress(&v, BOMB_LIMIT).unwrap()[..], &v[..]);
167	}
168
169	#[test]
170	fn possible_bomb_fails() {
171		let encoded_bigger_than_bomb = vec![0; BOMB_LIMIT + 1];
172		let mut buf = ZSTD_PREFIX.to_vec();
173
174		{
175			let mut v = zstd::Encoder::new(&mut buf, 3).unwrap().auto_finish();
176			v.write_all(&encoded_bigger_than_bomb[..]).unwrap();
177		}
178
179		assert_eq!(decompress(&buf[..], BOMB_LIMIT).err(), Some(Error::PossibleBomb));
180	}
181}
sp_maybe_compressed_blob/lib.rs

sp_maybe_compressed_blob/
lib.rs
