Skip to main content

tzap_core/
padding.rs

1use crate::format::FormatError;
2
3pub fn suffix_pad_for_aead(
4    payload: &[u8],
5    tag_len: usize,
6    block_size: usize,
7) -> Result<Vec<u8>, FormatError> {
8    let total_before_padding = payload
9        .len()
10        .checked_add(tag_len)
11        .ok_or(FormatError::PaddingOverflow)?;
12    let mut envelope_total = round_up_to_block(total_before_padding, block_size)?;
13    let mut pad_len = envelope_total
14        .checked_sub(total_before_padding)
15        .ok_or(FormatError::PaddingOverflow)?;
16    if pad_len == 0 {
17        envelope_total = envelope_total
18            .checked_add(block_size)
19            .ok_or(FormatError::PaddingOverflow)?;
20        pad_len = block_size;
21    }
22
23    let plaintext_len = envelope_total
24        .checked_sub(tag_len)
25        .ok_or(FormatError::PaddingOverflow)?;
26    let mut plaintext = Vec::with_capacity(plaintext_len);
27    plaintext.extend_from_slice(payload);
28    append_suffix_padding(&mut plaintext, pad_len)?;
29    Ok(plaintext)
30}
31
32pub fn append_suffix_padding(plaintext: &mut Vec<u8>, pad_len: usize) -> Result<(), FormatError> {
33    if pad_len == 0 {
34        return Err(FormatError::InvalidSuffixPadding);
35    }
36    if pad_len <= 254 {
37        plaintext.resize(plaintext.len() + pad_len - 1, 0);
38        plaintext.push(pad_len as u8);
39        return Ok(());
40    }
41
42    if pad_len > u32::MAX as usize {
43        return Err(FormatError::PaddingOverflow);
44    }
45    if pad_len < 5 {
46        return Err(FormatError::InvalidSuffixPadding);
47    }
48    plaintext.resize(plaintext.len() + pad_len - 5, 0);
49    plaintext.extend_from_slice(&(pad_len as u32).to_le_bytes());
50    plaintext.push(0xff);
51    Ok(())
52}
53
54pub fn depad_suffix_padding(plaintext: &[u8]) -> Result<&[u8], FormatError> {
55    let n = plaintext.len();
56    if n == 0 {
57        return Err(FormatError::EmptyPaddedPlaintext);
58    }
59
60    let final_byte = plaintext[n - 1];
61    let (marker_size, pad_len) = if final_byte < 0xff {
62        (1usize, final_byte as usize)
63    } else {
64        if n < 5 {
65            return Err(FormatError::InvalidSuffixPadding);
66        }
67        let pad_len = u32::from_le_bytes(
68            plaintext[n - 5..n - 1]
69                .try_into()
70                .expect("slice length checked"),
71        ) as usize;
72        if pad_len < 255 {
73            return Err(FormatError::InvalidSuffixPadding);
74        }
75        (5usize, pad_len)
76    };
77
78    if pad_len < marker_size || pad_len > n {
79        return Err(FormatError::InvalidSuffixPadding);
80    }
81    let payload_len = n
82        .checked_sub(pad_len)
83        .ok_or(FormatError::InvalidSuffixPadding)?;
84    let zero_padding_end = n - marker_size;
85    if plaintext[payload_len..zero_padding_end]
86        .iter()
87        .any(|byte| *byte != 0)
88    {
89        return Err(FormatError::NonZeroPaddingBytes);
90    }
91    Ok(&plaintext[..payload_len])
92}
93
94fn round_up_to_block(value: usize, block_size: usize) -> Result<usize, FormatError> {
95    if block_size == 0 {
96        return Err(FormatError::PaddingOverflow);
97    }
98    let remainder = value % block_size;
99    if remainder == 0 {
100        Ok(value.max(block_size))
101    } else {
102        value
103            .checked_add(block_size - remainder)
104            .ok_or(FormatError::PaddingOverflow)
105    }
106}
107
108#[cfg(test)]
109mod tests {
110    use super::*;
111
112    #[test]
113    fn byte_form_padding_round_trips() {
114        let padded = suffix_pad_for_aead(b"hello", 16, 32).unwrap();
115        assert_eq!(padded.len() + 16, 32);
116        assert_eq!(depad_suffix_padding(&padded).unwrap(), b"hello");
117    }
118
119    #[test]
120    fn exact_fit_adds_full_extra_block_with_wide_form() {
121        let padded = suffix_pad_for_aead(&vec![7; 4080], 16, 4096).unwrap();
122        assert_eq!(padded.len() + 16, 8192);
123        assert_eq!(padded[padded.len() - 1], 0xff);
124        assert_eq!(depad_suffix_padding(&padded).unwrap(), vec![7; 4080]);
125    }
126
127    #[test]
128    fn wide_form_padding_round_trips() {
129        let padded = suffix_pad_for_aead(b"hello", 16, 4096).unwrap();
130        assert_eq!(padded[padded.len() - 1], 0xff);
131        assert_eq!(depad_suffix_padding(&padded).unwrap(), b"hello");
132    }
133
134    #[test]
135    fn rejects_noncanonical_wide_form_padding_below_255() {
136        assert_eq!(
137            depad_suffix_padding(&[0x05, 0x00, 0x00, 0x00, 0xff]).unwrap_err(),
138            FormatError::InvalidSuffixPadding
139        );
140    }
141
142    #[test]
143    fn payload_ending_with_ff_does_not_confuse_suffix_marker() {
144        let payload = [b'f', b'r', b'a', b'm', b'e', 0xff];
145        let padded = suffix_pad_for_aead(&payload, 16, 32).unwrap();
146
147        assert_ne!(padded[payload.len() - 1], padded[padded.len() - 1]);
148        assert_eq!(padded[payload.len() - 1], 0xff);
149        assert_eq!(depad_suffix_padding(&padded).unwrap(), payload);
150    }
151
152    #[test]
153    fn rejects_zero_padding_and_non_zero_padding_bytes() {
154        assert_eq!(
155            depad_suffix_padding(&[1, 2, 0]).unwrap_err(),
156            FormatError::InvalidSuffixPadding
157        );
158
159        let mut padded = suffix_pad_for_aead(b"hello", 16, 32).unwrap();
160        let payload_len = 5;
161        padded[payload_len] = 1;
162        assert_eq!(
163            depad_suffix_padding(&padded).unwrap_err(),
164            FormatError::NonZeroPaddingBytes
165        );
166    }
167}