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        (5usize, pad_len)
73    };
74
75    if pad_len < marker_size || pad_len > n {
76        return Err(FormatError::InvalidSuffixPadding);
77    }
78    let payload_len = n
79        .checked_sub(pad_len)
80        .ok_or(FormatError::InvalidSuffixPadding)?;
81    let zero_padding_end = n - marker_size;
82    if plaintext[payload_len..zero_padding_end]
83        .iter()
84        .any(|byte| *byte != 0)
85    {
86        return Err(FormatError::NonZeroPaddingBytes);
87    }
88    Ok(&plaintext[..payload_len])
89}
90
91fn round_up_to_block(value: usize, block_size: usize) -> Result<usize, FormatError> {
92    if block_size == 0 {
93        return Err(FormatError::PaddingOverflow);
94    }
95    let remainder = value % block_size;
96    if remainder == 0 {
97        Ok(value.max(block_size))
98    } else {
99        value
100            .checked_add(block_size - remainder)
101            .ok_or(FormatError::PaddingOverflow)
102    }
103}
104
105#[cfg(test)]
106mod tests {
107    use super::*;
108
109    #[test]
110    fn byte_form_padding_round_trips() {
111        let padded = suffix_pad_for_aead(b"hello", 16, 32).unwrap();
112        assert_eq!(padded.len() + 16, 32);
113        assert_eq!(depad_suffix_padding(&padded).unwrap(), b"hello");
114    }
115
116    #[test]
117    fn exact_fit_adds_full_extra_block_with_wide_form() {
118        let padded = suffix_pad_for_aead(&vec![7; 4080], 16, 4096).unwrap();
119        assert_eq!(padded.len() + 16, 8192);
120        assert_eq!(padded[padded.len() - 1], 0xff);
121        assert_eq!(depad_suffix_padding(&padded).unwrap(), vec![7; 4080]);
122    }
123
124    #[test]
125    fn wide_form_padding_round_trips() {
126        let padded = suffix_pad_for_aead(b"hello", 16, 4096).unwrap();
127        assert_eq!(padded[padded.len() - 1], 0xff);
128        assert_eq!(depad_suffix_padding(&padded).unwrap(), b"hello");
129    }
130
131    #[test]
132    fn rejects_zero_padding_and_non_zero_padding_bytes() {
133        assert_eq!(
134            depad_suffix_padding(&[1, 2, 0]).unwrap_err(),
135            FormatError::InvalidSuffixPadding
136        );
137
138        let mut padded = suffix_pad_for_aead(b"hello", 16, 32).unwrap();
139        let payload_len = 5;
140        padded[payload_len] = 1;
141        assert_eq!(
142            depad_suffix_padding(&padded).unwrap_err(),
143            FormatError::NonZeroPaddingBytes
144        );
145    }
146}