1use crate::errors::{Bee2Result, InvalidPaddingError, InvalidPaddingErrorKind};
2
3pub struct Block {
4 _private: (),
5}
6
7impl Block {
8 #[cfg(feature = "block-padding")]
9 fn block_pad(data: &[u8], block_size: u8) -> Vec<Box<[u8]>> {
10 let mut blocks: Vec<Box<[u8]>> = data
11 .chunks(block_size as usize)
12 .take(data.len() / block_size as usize)
13 .map(Box::from)
14 .collect();
15 let last_block_offset = data.len() / block_size as usize * block_size as usize;
16 let last_block = &data[last_block_offset..];
17 let unpadded_size = last_block.len() as u8;
18 let padding_size = block_size - unpadded_size;
19 let mut last_block_padded = last_block.to_vec();
20 last_block_padded.extend(vec![unpadded_size; padding_size as usize]);
21 blocks.push(last_block_padded.into_boxed_slice());
22 blocks
23 }
24
25 #[cfg(feature = "block-padding")]
26 fn block_unpad(data: Vec<Box<[u8]>>) -> Bee2Result<Box<[u8]>> {
27 let Some((last_block, block_size)) = data.last().map(|value| (value, value.len())) else {
28 return Err(InvalidPaddingError::new_box(
29 InvalidPaddingErrorKind::NoBlocks(data),
30 ));
31 };
32
33 let Some(&last_block_size) = last_block.last() else {
34 return Err(InvalidPaddingError::new_box(
35 InvalidPaddingErrorKind::NoData(data),
36 ));
37 };
38
39 if last_block_size as usize >= last_block.len() {
40 return Err(InvalidPaddingError::new_box(
41 InvalidPaddingErrorKind::NotEnoughData(
42 data.clone(),
43 last_block_size,
44 last_block.len() as u8,
45 ),
46 ));
47 }
48
49 let mut last_data_chunk: Vec<u8> = vec![];
50
51 for (i, byte) in last_block.iter().enumerate() {
52 if (i as u8) < last_block_size {
53 last_data_chunk.push(*byte);
54 } else if *byte != last_block_size {
55 return Err(InvalidPaddingError::new_box(
56 InvalidPaddingErrorKind::InvalidPaddingData(
57 data.clone(),
58 i,
59 *byte,
60 last_block_size,
61 ),
62 ));
63 }
64 }
65
66 let mut result_data: Vec<u8> = vec![];
67
68 for (i, block) in data.iter().take(data.len() - 1).enumerate() {
69 if block_size != block.len() {
70 return Err(InvalidPaddingError::new_box(
71 InvalidPaddingErrorKind::DifferentBlockSizes(
72 data.clone(),
73 i,
74 block_size as u8,
75 block.len() as u8,
76 ),
77 ));
78 }
79 result_data.extend(block);
80 }
81
82 result_data.extend(last_data_chunk);
83
84 Ok(result_data.into_boxed_slice())
85 }
86
87 #[cfg(feature = "block-padding")]
88 fn block_is_padding_correct(data: Vec<Box<[u8]>>) -> bool {
89 let Some((last_block, block_size)) = data.last().map(|value| (value, value.len())) else {
90 return false;
91 };
92
93 let Some(&last_block_size) = last_block.last() else {
94 return false;
95 };
96
97 if last_block_size as usize >= last_block.len() {
98 return false;
99 }
100
101 for byte in &last_block[last_block_size as usize..] {
102 if *byte != last_block_size {
103 return false;
104 }
105 }
106
107 for block in data.iter().take(data.len() - 1) {
108 if block_size != block.len() {
109 return false;
110 }
111 }
112
113 true
114 }
115
116 pub fn pad(data: &[u8], block_size: u8) -> Vec<Box<[u8]>> {
117 #[cfg(feature = "block-padding")]
118 {
119 Self::block_pad(data, block_size)
120 }
121 #[cfg(not(feature = "block-padding"))]
122 {
123 let _ = block_size;
124 vec![Box::from(data)]
125 }
126 }
127
128 pub fn unpad(data: Vec<Box<[u8]>>) -> Bee2Result<Box<[u8]>> {
129 #[cfg(feature = "block-padding")]
130 {
131 Self::block_unpad(data)
132 }
133 #[cfg(not(feature = "block-padding"))]
134 {
135 if let Some(result) = data.first() {
136 Ok(result).cloned()
137 } else {
138 Err(InvalidPaddingError::new_box(
139 InvalidPaddingErrorKind::NoBlocks(data),
140 ))
141 }
142 }
143 }
144
145 pub fn is_padding_correct(data: Vec<Box<[u8]>>) -> bool {
146 #[cfg(feature = "block-padding")]
147 {
148 Self::block_is_padding_correct(data)
149 }
150 #[cfg(not(feature = "block-padding"))]
151 {
152 !data.is_empty()
153 }
154 }
155}
156
157#[cfg(test)]
158mod tests {
159 use super::Block;
160
161 #[test]
162 fn test_multiple_data_sizes_16() {
163 for i in 0..200 {
164 let data = vec![255; i];
165 let padded = Block::pad(&data, 16);
166 assert!(Block::is_padding_correct(padded.clone()));
167 let unpadded = Block::unpad(padded).unwrap();
168 assert_eq!(*data, *unpadded);
169 }
170 }
171
172 #[test]
173 fn test_multiple_data_sizes_32() {
174 for i in 0..200 {
175 let data = vec![255; i];
176 let padded = Block::pad(&data, 32);
177 assert!(Block::is_padding_correct(padded.clone()));
178 let unpadded = Block::unpad(padded).unwrap();
179 assert_eq!(*data, *unpadded);
180 }
181 }
182
183 #[test]
184 fn test_various_bs_and_sizes() {
185 for block_size in 1..50 {
186 for data_size in 0..=2 * block_size + 5 {
187 let mut data = vec![];
188
189 for i in 0..data_size {
190 data.push(i);
191 }
192
193 let padded = Block::pad(&data, block_size);
194 assert!(Block::is_padding_correct(padded.clone()));
195 let unpadded = Block::unpad(padded).unwrap();
196 assert_eq!(*data, *unpadded);
197 }
198 }
199 }
200
201 #[test]
202 fn test_invalid_padding_16() {
203 for data_size in 0..200 {
204 let mut data = vec![];
205
206 for i in 0..data_size {
207 data.push(i);
208 }
209
210 let mut blocks: Vec<Box<[u8]>> = data
211 .chunks(16)
212 .take(data.len() / 16)
213 .map(Box::from)
214 .collect();
215 let last_block_offset = data.len() / 16 * 16;
216 let last_block = &data[last_block_offset..];
217 blocks.push(Box::from(last_block));
218
219 if data_size > 0 && data_size < 16 {
220 assert!(Block::is_padding_correct(blocks.clone()));
221 assert!(Block::unpad(blocks).is_ok());
222 } else {
223 assert!(!Block::is_padding_correct(blocks.clone()));
224 assert!(Block::unpad(blocks).is_err());
225 }
226 }
227 }
228}