@@ -17,21 +17,29 @@ pub(super) fn stretch_key(key: &Pin<Box<GenericArray<u8, U32>>>) -> Result<Aes25
1717}
1818
1919/// Pads bytes to a minimum length using PKCS7-like padding.
20- /// The last N bytes of the padded bytes all have the value N. Minimum of 1 padding byte.
21- /// For example, padded to size 4, the value 0,0 becomes 0,0,2,2.
22- pub ( crate ) fn pad_bytes ( bytes : & mut Vec < u8 > , min_length : usize ) {
20+ /// The last N bytes of the padded bytes all have the value N. Minimum of 1 padding byte. maximum of
21+ /// 255 bytes. For example, padded to size 4, the value 0,0 becomes 0,0,2,2.
22+ /// If the more than 255 bytes of padding is required, this will return an error.
23+ pub ( crate ) fn pad_bytes ( bytes : & mut Vec < u8 > , min_length : usize ) -> Result < ( ) , CryptoError > {
2324 // at least 1 byte of padding is required
2425 let pad_bytes = min_length. saturating_sub ( bytes. len ( ) ) . max ( 1 ) ;
26+ // Since each byte represents the padding size, the maximum padding is 255.
27+ // If more padding is required, this will return an error.
28+ if pad_bytes > 255 {
29+ return Err ( CryptoError :: InvalidPadding ) ;
30+ }
2531 let padded_length = max ( min_length, bytes. len ( ) + 1 ) ;
2632 bytes. resize ( padded_length, pad_bytes as u8 ) ;
33+ Ok ( ( ) )
2734}
2835
2936/// Unpads bytes that is padded using the PKCS7-like padding defined by [pad_bytes].
3037/// The last N bytes of the padded bytes all have the value N. Minimum of 1 padding byte.
3138/// For example, padded to size 4, the value 0,0 becomes 0,0,2,2.
3239pub ( crate ) fn unpad_bytes ( padded_bytes : & [ u8 ] ) -> Result < & [ u8 ] , CryptoError > {
3340 let pad_len = * padded_bytes. last ( ) . ok_or ( CryptoError :: InvalidPadding ) ? as usize ;
34- if pad_len >= padded_bytes. len ( ) {
41+ // The padding is at minimum 1 as noted in `pad_bytes`
42+ if pad_len == 0 || pad_len > padded_bytes. len ( ) {
3543 return Err ( CryptoError :: InvalidPadding ) ;
3644 }
3745 Ok ( padded_bytes[ ..( padded_bytes. len ( ) - pad_len) ] . as_ref ( ) )
@@ -68,16 +76,121 @@ mod tests {
6876 ) ;
6977 }
7078
79+ #[ test]
80+ fn test_pad_bytes_256_error ( ) {
81+ let mut bytes = vec ! [ 1u8 ; 0 ] ;
82+ let result = pad_bytes ( & mut bytes, 256 ) ;
83+ assert ! ( matches!( result, Err ( CryptoError :: InvalidPadding ) ) ) ;
84+ }
85+
7186 #[ test]
7287 fn test_pad_bytes_roundtrip ( ) {
7388 let original_bytes = vec ! [ 1u8 ; 10 ] ;
7489 let mut cloned_bytes = original_bytes. clone ( ) ;
7590 let mut encoded_bytes = vec ! [ 1u8 ; 12 ] ;
7691 encoded_bytes[ 10 ] = 2 ;
7792 encoded_bytes[ 11 ] = 2 ;
78- pad_bytes ( & mut cloned_bytes, 12 ) ;
93+ pad_bytes ( & mut cloned_bytes, 12 ) . expect ( "Padding failed" ) ;
7994 assert_eq ! ( encoded_bytes, cloned_bytes) ;
8095 let unpadded_bytes = unpad_bytes ( & cloned_bytes) . unwrap ( ) ;
8196 assert_eq ! ( original_bytes, unpadded_bytes) ;
8297 }
98+
99+ #[ test]
100+ fn test_pad_bytes_roundtrip_empty ( ) {
101+ let original_bytes = Vec :: new ( ) ;
102+ let mut cloned_bytes = original_bytes. clone ( ) ;
103+ pad_bytes ( & mut cloned_bytes, 32 ) . expect ( "Padding failed" ) ;
104+ let unpadded = unpad_bytes ( & cloned_bytes) . unwrap ( ) ;
105+ assert_eq ! ( Vec :: <u8 >:: new( ) , unpadded) ;
106+ }
107+
108+ #[ test]
109+ fn test_unpad_bytes_invalid_empty ( ) {
110+ let data: Vec < u8 > = vec ! [ ] ;
111+ let result = unpad_bytes ( & data) ;
112+ assert ! ( matches!( result, Err ( CryptoError :: InvalidPadding ) ) ) ;
113+ }
114+
115+ #[ test]
116+ fn test_unpad_bytes_invalid_too_large ( ) {
117+ // Last byte is 5, but only 4 bytes in total
118+ let data = vec ! [ 1 , 2 , 3 , 5 ] ;
119+ let result = unpad_bytes ( & data) ;
120+ assert ! ( matches!( result, Err ( CryptoError :: InvalidPadding ) ) ) ;
121+ }
122+
123+ #[ test]
124+ fn test_unpad_bytes_invalid_0_padding ( ) {
125+ // Padding value of 0 is invalid
126+ let data = vec ! [ 1 , 2 , 3 , 0 ] ;
127+ let result = unpad_bytes ( & data) ;
128+ assert ! ( matches!( result, Err ( CryptoError :: InvalidPadding ) ) ) ;
129+ }
130+
131+ #[ test]
132+ fn test_pad_and_unpad_bytes_range_0_to_1024 ( ) {
133+ let cases: Vec < _ > = ( 0 ..=1024 )
134+ . flat_map ( |data_size| ( 2 ..=1024 ) . map ( move |padding_size| ( data_size, padding_size) ) )
135+ . collect ( ) ;
136+
137+ let data_larger_than_padding_cases: Vec < _ > = cases
138+ . clone ( )
139+ . into_iter ( )
140+ . filter ( |( data_size, padding_size) | data_size > padding_size)
141+ . collect ( ) ;
142+ for ( data_size, padding_size) in data_larger_than_padding_cases {
143+ let mut data: Vec < u8 > = vec ! [ 0x12 ; data_size] ;
144+ let original = data. clone ( ) ;
145+ pad_bytes ( & mut data, padding_size) . expect ( "Padding failed" ) ;
146+ let unpadded = unpad_bytes ( & data) . expect ( "Unpadding failed" ) ;
147+ assert_eq ! (
148+ unpadded, original,
149+ "Failed at size {} and padding {}" ,
150+ data_size, padding_size
151+ ) ;
152+ }
153+
154+ let padding_larger_than_data_cases: Vec < _ > = cases
155+ . clone ( )
156+ . into_iter ( )
157+ . filter ( |( data_size, padding_size) | {
158+ data_size <= padding_size && ( padding_size - data_size) <= 255
159+ } )
160+ . collect ( ) ;
161+ for ( data_size, padding_size) in padding_larger_than_data_cases {
162+ println ! (
163+ "Testing data_size: {}, padding_size: {}" ,
164+ data_size, padding_size
165+ ) ;
166+ let data_original: Vec < u8 > = vec ! [ 0x12 ; data_size] ;
167+ let mut data = data_original. clone ( ) ;
168+
169+ pad_bytes ( & mut data, padding_size) . expect ( "Padding failed" ) ;
170+ let unpadded = unpad_bytes ( & data) . expect ( "Unpadding failed" ) ;
171+ assert_eq ! (
172+ unpadded, data_original,
173+ "Failed at size {} and padding {}" ,
174+ data_size, padding_size
175+ ) ;
176+ }
177+
178+ let padding_massively_larger_than_data_cases: Vec < _ > = cases
179+ . into_iter ( )
180+ . filter ( |( data_size, padding_size) | {
181+ data_size <= padding_size && ( padding_size - data_size) > 255
182+ } )
183+ . collect ( ) ;
184+ for ( data_size, padding_size) in padding_massively_larger_than_data_cases {
185+ let mut data: Vec < u8 > = vec ! [ 0x12 ; data_size] ;
186+ let error = pad_bytes ( & mut data, padding_size) ;
187+ assert ! (
188+ matches!( error, Err ( CryptoError :: InvalidPadding ) ) ,
189+ "Expected InvalidPadding error at size {} and padding {}, but got {:?}" ,
190+ data_size,
191+ padding_size,
192+ error
193+ ) ;
194+ }
195+ }
83196}
0 commit comments