jam_std_common/util/
injective_decode.rs

1use crate::{CountedOutput, RecordedInput};
2use arrayvec::ArrayVec;
3use codec::{ConstEncodedLen, Decode, Encode, Error, Input, MaxEncodedLen, Output};
4
5/// `T` wrapper with a constant encoded length (`N`) and an injective decode implementation.
6///
7/// Note that the decode implementation will _not_ be injective if `T::decode` accepts
8/// non-canonical encodings.
9#[derive(Copy, Clone, Debug, Eq, PartialEq)]
10pub enum InjectiveDecode<T, const N: usize> {
11	Valid(T),
12	Invalid([u8; N]),
13}
14
15impl<T, const N: usize> InjectiveDecode<T, N> {
16	pub fn valid(&self) -> Option<&T> {
17		match self {
18			Self::Valid(value) => Some(value),
19			Self::Invalid(_) => None,
20		}
21	}
22}
23
24impl<T: MaxEncodedLen, const N: usize> Encode for InjectiveDecode<T, N> {
25	fn size_hint(&self) -> usize {
26		N
27	}
28
29	fn encode_to<O: Output + ?Sized>(&self, output: &mut O) {
30		match self {
31			Self::Valid(value) => {
32				let mut counted_output = CountedOutput::new(output);
33				debug_assert!(T::max_encoded_len() <= N);
34				value.encode_to(&mut counted_output);
35				// Pad with zeroes
36				while counted_output.num_bytes_written() < N {
37					counted_output.push_byte(0);
38				}
39				debug_assert_eq!(counted_output.num_bytes_written(), N);
40			},
41			Self::Invalid(encoded) => output.write(encoded),
42		}
43	}
44
45	fn using_encoded<R, F: FnOnce(&[u8]) -> R>(&self, f: F) -> R {
46		match self {
47			Self::Valid(_) => f(&self.encode()),
48			Self::Invalid(encoded) => f(encoded),
49		}
50	}
51
52	fn encoded_size(&self) -> usize {
53		N
54	}
55}
56
57impl<T: MaxEncodedLen, const N: usize> MaxEncodedLen for InjectiveDecode<T, N> {
58	fn max_encoded_len() -> usize {
59		N
60	}
61}
62
63impl<T: MaxEncodedLen, const N: usize> ConstEncodedLen for InjectiveDecode<T, N> {}
64
65impl<T: MaxEncodedLen + Decode, const N: usize> Decode for InjectiveDecode<T, N> {
66	fn decode<I: Input>(input: &mut I) -> Result<Self, Error> {
67		let mut recorded_input = RecordedInput::with_record(input, ArrayVec::<u8, N>::new());
68
69		debug_assert!(T::max_encoded_len() <= N);
70		let mut value = T::decode(&mut recorded_input);
71		if recorded_input.failed() {
72			value?;
73			// Expect error to have been forwarded by Decode impl...
74			debug_assert!(false);
75			return Err("Unforwarded input failure".into())
76		}
77
78		// Consume padding bytes; these must be zero
79		while !recorded_input.record().is_full() {
80			if (recorded_input.read_byte()? != 0) && value.is_ok() {
81				value = Err("Non-zero padding".into());
82			}
83		}
84
85		match value {
86			Ok(value) => Ok(Self::Valid(value)),
87			Err(_) =>
88				Ok(Self::Invalid(recorded_input.into_record().into_inner().expect("Filled above"))),
89		}
90	}
91
92	fn skip<I: Input>(input: &mut I) -> Result<(), Error> {
93		for _ in 0..N {
94			input.read_byte()?;
95		}
96		Ok(())
97	}
98
99	fn encoded_fixed_size() -> Option<usize> {
100		Some(N)
101	}
102}
103
104#[cfg(test)]
105mod tests {
106	use super::*;
107
108	#[test]
109	fn encode_valid_pads_with_zeroes() {
110		let x = InjectiveDecode::<u16, 4>::Valid(0x0201);
111		assert_eq!(&x.encode(), &[0x01, 0x02, 0x00, 0x00]);
112		assert_eq!(x.encoded_size(), 4);
113		assert_eq!(x.size_hint(), 4);
114	}
115
116	#[test]
117	fn encode_invalid_writes_raw_bytes() {
118		let x = InjectiveDecode::<u16, 4>::Invalid([1, 2, 3, 4]);
119		assert_eq!(&x.encode(), &[1, 2, 3, 4]);
120	}
121
122	#[test]
123	fn decode_valid_with_zero_padding() {
124		let encoded = [0x01, 0x02, 0x00, 0x00];
125		let decoded = InjectiveDecode::<u16, 4>::decode(&mut &encoded[..]).unwrap();
126		assert_eq!(decoded, InjectiveDecode::Valid(0x0201));
127	}
128
129	#[test]
130	fn decode_nonzero_padding_is_invalid() {
131		// The inner u16 decodes fine, but a non-zero padding byte makes the whole value invalid,
132		// capturing the raw bytes verbatim
133		let encoded = [0x01, 0x02, 0x00, 0x05];
134		let decoded = InjectiveDecode::<u16, 4>::decode(&mut &encoded[..]).unwrap();
135		assert_eq!(decoded, InjectiveDecode::Invalid(encoded));
136	}
137
138	#[test]
139	fn decode_invalid_value_is_captured() {
140		// 0x02 is not a valid bool encoding, so the value decode fails and the raw bytes are
141		// captured rather than the error being propagated
142		let encoded = [0x02, 0x00];
143		let decoded = InjectiveDecode::<bool, 2>::decode(&mut &encoded[..]).unwrap();
144		assert_eq!(decoded, InjectiveDecode::Invalid(encoded));
145	}
146
147	#[test]
148	fn decode_is_injective() {
149		// Re-encoding any decoded value must reproduce the original bytes exactly
150		for a in 0..=4u8 {
151			for b in 0..=4u8 {
152				for c in 0..=1u8 {
153					for d in 0..=1u8 {
154						let encoded = [a, b, c, d];
155						let decoded = InjectiveDecode::<u16, 4>::decode(&mut &encoded[..]).unwrap();
156						assert_eq!(&decoded.encode(), &encoded);
157					}
158				}
159			}
160		}
161	}
162
163	#[test]
164	fn decode_fails_on_truncated_input() {
165		let encoded = [0x01];
166		assert!(InjectiveDecode::<u16, 4>::decode(&mut &encoded[..]).is_err());
167	}
168
169	#[test]
170	fn skip() {
171		let mut input = [1, 2, 3, 4, 5, 6].as_slice();
172		InjectiveDecode::<u16, 4>::skip(&mut input).unwrap();
173		assert_eq!(input, &[5, 6]);
174		assert!(InjectiveDecode::<u16, 4>::skip(&mut input).is_err());
175	}
176}
jam_std_common/util/injective_decode.rs

jam_std_common/util/
injective_decode.rs
