Skip to main content

ctaes_rs/
padding.rs

1// Copyright (c) 2023 Blockstream
2// Distributed under the MIT software license, see the accompanying
3// file COPYING or https://opensource.org/licenses/mit-license.php.
4
5use crate::Error;
6
7/// Trait defining interface for a Padding implementation
8pub trait Padding {
9    /// Given a length of data and block_size return the minimum size of a buffer required to hold
10    /// the padded data
11    fn padded_buffer_length(data_length: usize, block_size: usize) -> usize;
12    /// Take a buffer containing data of specified length and based on the blocksize apply the padding
13    /// to the buffer. The returned slice will reference the slice of padded data in the buffer
14    fn pad(buffer: &mut [u8], data_length: usize, block_size: usize) -> Result<&[u8], Error>;
15    /// Given a buffer of padded data, attempt to provide a slice that references
16    /// the unpadded data
17    fn unpad(buffer: &[u8]) -> Result<&[u8], Error>;
18}
19
20/// Implementation of basic Zero Padding. May not be reversible if the original data ends with one
21/// or more zero bytes. Does not add an extra block of padding if the data length is already a
22/// multiple of the block size
23pub 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
59/// Implementation of the PKCS7 padding scheme
60pub 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}