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}