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}