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

            
21
use 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.
30
const 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.
37
pub const CODE_BLOB_BOMB_LIMIT: usize = 50 * 1024 * 1024;
38

            
39
/// A possible bomb was encountered.
40
#[derive(Debug, Clone, PartialEq, thiserror::Error)]
41
pub 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

            
50
fn 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

            
67
fn 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.
75
pub 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
/// Encode a blob as compressed. If the blob's size is over the bomb limit,
84
/// this will not compress the blob, as the decoder will not be able to be
85
/// able to differentiate it from a compression bomb.
86
pub fn compress(blob: &[u8], bomb_limit: usize) -> Option<Vec<u8>> {
87
	if blob.len() > bomb_limit {
88
		return None
89
	}
90

            
91
	let mut buf = ZSTD_PREFIX.to_vec();
92

            
93
	{
94
		let mut v = zstd::Encoder::new(&mut buf, 3).ok()?.auto_finish();
95
		v.write_all(blob).ok()?;
96
	}
97

            
98
	Some(buf)
99
}
100

            
101
#[cfg(test)]
102
mod tests {
103
	use super::*;
104

            
105
	const BOMB_LIMIT: usize = 10;
106

            
107
	#[test]
108
	fn refuse_to_encode_over_limit() {
109
		let mut v = vec![0; BOMB_LIMIT + 1];
110
		assert!(compress(&v, BOMB_LIMIT).is_none());
111

            
112
		let _ = v.pop();
113
		assert!(compress(&v, BOMB_LIMIT).is_some());
114
	}
115

            
116
	#[test]
117
	fn compress_and_decompress() {
118
		let v = vec![0; BOMB_LIMIT];
119

            
120
		let compressed = compress(&v, BOMB_LIMIT).unwrap();
121

            
122
		assert!(compressed.starts_with(&ZSTD_PREFIX));
123
		assert_eq!(&decompress(&compressed, BOMB_LIMIT).unwrap()[..], &v[..])
124
	}
125

            
126
	#[test]
127
	fn decompresses_only_when_magic() {
128
		let v = vec![0; BOMB_LIMIT + 1];
129

            
130
		assert_eq!(&decompress(&v, BOMB_LIMIT).unwrap()[..], &v[..]);
131
	}
132

            
133
	#[test]
134
	fn possible_bomb_fails() {
135
		let encoded_bigger_than_bomb = vec![0; BOMB_LIMIT + 1];
136
		let mut buf = ZSTD_PREFIX.to_vec();
137

            
138
		{
139
			let mut v = zstd::Encoder::new(&mut buf, 3).unwrap().auto_finish();
140
			v.write_all(&encoded_bigger_than_bomb[..]).unwrap();
141
		}
142

            
143
		assert_eq!(decompress(&buf[..], BOMB_LIMIT).err(), Some(Error::PossibleBomb));
144
	}
145
}