1use crate::Error;
6
7pub trait Padding {
9 fn padded_buffer_length(data_length: usize, block_size: usize) -> usize;
12 fn pad(buffer: &mut [u8], data_length: usize, block_size: usize) -> Result<&[u8], Error>;
15 fn unpad(buffer: &[u8]) -> Result<&[u8], Error>;
18}
19
20pub enum ZeroPadding {}
24
25impl Padding for ZeroPadding {
26 fn padded_buffer_length(data_length: usize, block_size: usize) -> usize {
27 if data_length % block_size == 0 {
28 data_length
29 } else {
30 block_size + block_size * (data_length / block_size)
31 }
32 }
33
34 fn pad(buffer: &mut [u8], data_length: usize, block_size: usize) -> Result<&[u8], Error> {
35 let padded_length = Self::padded_buffer_length(data_length, block_size);
36 if buffer.len() < padded_length {
37 return Err(Error::PaddedBufferTooSmall);
38 }
39
40 if data_length < padded_length {
41 buffer[data_length..padded_length].fill(0u8);
42 }
43
44 Ok(&buffer[0..padded_length])
45 }
46
47 fn unpad(buffer: &[u8]) -> Result<&[u8], Error> {
48 let mut n = buffer.len() - 1;
49 while n > 0 {
50 if buffer[n] != 0 {
51 break;
52 }
53 n -= 1;
54 }
55 Ok(&buffer[0..=n])
56 }
57}
58
59pub enum Pkcs7 {}
61
62impl Padding for Pkcs7 {
63 fn padded_buffer_length(data_length: usize, block_size: usize) -> usize {
64 block_size + block_size * (data_length / block_size)
65 }
66
67 fn pad(buffer: &mut [u8], data_length: usize, block_size: usize) -> Result<&[u8], Error> {
68 let padded_length = Self::padded_buffer_length(data_length, block_size);
69 if buffer.len() < padded_length {
70 return Err(Error::PaddedBufferTooSmall);
71 }
72
73 let padding_length = padded_length - data_length;
74
75 buffer[data_length..padded_length].fill(padding_length as u8);
76 Ok(&buffer[0..padded_length])
77 }
78
79 fn unpad(buffer: &[u8]) -> Result<&[u8], Error> {
80 let padding_length = match buffer.last() {
81 None => return Err(Error::UnpadError("Buffer to unpad is empty".to_string())),
82 Some(l) => *l,
83 };
84 if buffer.len() < padding_length as usize {
85 return Err(Error::UnpadError(
86 "Buffer smaller than PKCS7 recorded padded length".to_string(),
87 ));
88 }
89 if padding_length == 0 {
90 return Err(Error::UnpadError("PKCS7 padding length can't be zero".to_string()));
91 }
92
93 Ok(&buffer[0..buffer.len() - padding_length as usize])
94 }
95}
96
97#[cfg(test)]
98mod test {
99 use crate::padding::{Padding, Pkcs7, ZeroPadding};
100 use crate::Error;
101
102 #[test]
103 fn test_zero_padding() {
104 const BLOCK_SIZE: usize = 16;
105 assert_eq!(ZeroPadding::padded_buffer_length(20, BLOCK_SIZE), 2 * BLOCK_SIZE);
106 assert_eq!(ZeroPadding::padded_buffer_length(2 * BLOCK_SIZE, BLOCK_SIZE), 2 * BLOCK_SIZE);
107
108 const DATA_LENGTH: usize = 22;
109 let mut buffer = [2u8; 2 * BLOCK_SIZE + 1];
110
111 buffer[0..2 * BLOCK_SIZE].fill(1u8);
112
113 let padded_buffer =
114 ZeroPadding::pad(buffer.as_mut_slice(), DATA_LENGTH, BLOCK_SIZE).unwrap();
115
116 assert_eq!(&padded_buffer[0..DATA_LENGTH], [1u8; DATA_LENGTH].as_slice());
117
118 assert_eq!(
119 &padded_buffer[DATA_LENGTH..],
120 vec![0u8; padded_buffer.len() - DATA_LENGTH].as_slice()
121 );
122
123 let unpadded_padded_buffer = ZeroPadding::unpad(padded_buffer).unwrap();
124 assert_eq!(unpadded_padded_buffer, [1u8; DATA_LENGTH].as_slice());
125
126 assert_eq!(buffer[2 * BLOCK_SIZE], 2u8);
127
128 assert!(matches!(
129 ZeroPadding::pad(buffer.as_mut_slice(), 2 * BLOCK_SIZE + 1, BLOCK_SIZE),
130 Err(Error::PaddedBufferTooSmall)
131 ));
132 }
133
134 #[test]
135 fn test_pkcs7_padding() {
136 const BLOCK_SIZE: usize = 16;
137 assert_eq!(Pkcs7::padded_buffer_length(20, BLOCK_SIZE), 2 * BLOCK_SIZE);
138 assert_eq!(Pkcs7::padded_buffer_length(2 * BLOCK_SIZE, BLOCK_SIZE), 3 * BLOCK_SIZE);
139
140 const DATA_LENGTH: usize = 22;
141 let mut buffer = [2u8; 2 * BLOCK_SIZE + 1];
142
143 buffer[0..2 * BLOCK_SIZE].fill(1u8);
144
145 let padded_buffer = Pkcs7::pad(buffer.as_mut_slice(), DATA_LENGTH, BLOCK_SIZE).unwrap();
146
147 assert_eq!(&padded_buffer[0..DATA_LENGTH], [1u8; DATA_LENGTH].as_slice());
148
149 assert_eq!(
150 &padded_buffer[DATA_LENGTH..],
151 vec![10u8; padded_buffer.len() - DATA_LENGTH].as_slice()
152 );
153
154 let unpadded_padded_buffer = Pkcs7::unpad(padded_buffer).unwrap();
155 assert_eq!(unpadded_padded_buffer, [1u8; DATA_LENGTH].as_slice());
156
157 assert_eq!(buffer[2 * BLOCK_SIZE], 2u8);
158
159 assert!(matches!(
160 Pkcs7::pad(buffer.as_mut_slice(), 2 * BLOCK_SIZE + 1, BLOCK_SIZE),
161 Err(Error::PaddedBufferTooSmall)
162 ));
163 }
164}