1use cloud_sdk_sanitization::sanitize_bytes;
4
5pub struct SnapshotEncoder<'output, E> {
11 destination: Destination<'output>,
12 len: usize,
13 max_len: usize,
14 error: E,
15}
16
17enum Destination<'output> {
18 Measure,
19 Output(&'output mut [u8]),
20 Verify(&'output [u8]),
21}
22
23struct EncodeRollback<'output> {
24 target: &'output mut [u8],
25 armed: bool,
26}
27
28impl EncodeRollback<'_> {
29 fn new(target: &mut [u8]) -> EncodeRollback<'_> {
30 EncodeRollback {
31 target,
32 armed: true,
33 }
34 }
35
36 fn disarm(&mut self) {
37 self.armed = false;
38 }
39}
40
41impl Drop for EncodeRollback<'_> {
42 fn drop(&mut self) {
43 if self.armed {
44 sanitize_bytes(self.target);
45 }
46 }
47}
48
49impl<'output, E: Copy> SnapshotEncoder<'output, E> {
50 fn measuring(max_len: usize, error: E) -> Self {
51 Self {
52 destination: Destination::Measure,
53 len: 0,
54 max_len,
55 error,
56 }
57 }
58
59 fn writing(output: &'output mut [u8], max_len: usize, error: E) -> Self {
60 Self {
61 destination: Destination::Output(output),
62 len: 0,
63 max_len,
64 error,
65 }
66 }
67
68 fn verifying(output: &'output [u8], max_len: usize, error: E) -> Self {
69 Self {
70 destination: Destination::Verify(output),
71 len: 0,
72 max_len,
73 error,
74 }
75 }
76
77 #[must_use]
79 pub const fn len(&self) -> usize {
80 self.len
81 }
82
83 #[must_use]
85 pub const fn is_empty(&self) -> bool {
86 self.len == 0
87 }
88
89 pub fn byte(&mut self, value: u8) -> Result<(), E> {
91 self.bytes(core::slice::from_ref(&value))
92 }
93
94 pub fn bytes(&mut self, value: &[u8]) -> Result<(), E> {
96 let end = self.checked_end(value.len())?;
97 if let Destination::Output(output) = &mut self.destination {
98 let target = output.get_mut(self.len..end).ok_or(self.error)?;
99 target.copy_from_slice(value);
100 } else if let Destination::Verify(output) = &self.destination {
101 let actual = output.get(self.len..end).ok_or(self.error)?;
102 if actual != value {
103 return Err(self.error);
104 }
105 }
106 self.len = end;
107 Ok(())
108 }
109
110 pub fn string(&mut self, value: &str) -> Result<(), E> {
112 self.bytes(value.as_bytes())
113 }
114
115 pub fn u64(&mut self, mut value: u64) -> Result<(), E> {
117 if value == 0 {
118 return self.byte(b'0');
119 }
120
121 let mut digits = [0_u8; 20];
122 let mut cursor = digits.len();
123 while value != 0 {
124 cursor = cursor.checked_sub(1).ok_or(self.error)?;
125 let digit = u8::try_from(value % 10).map_err(|_| self.error)?;
126 let slot = digits.get_mut(cursor).ok_or(self.error)?;
127 *slot = b'0'.checked_add(digit).ok_or(self.error)?;
128 value /= 10;
129 }
130 self.bytes(digits.get(cursor..).ok_or(self.error)?)
131 }
132
133 pub fn percent_encoded(&mut self, value: &str) -> Result<(), E> {
135 for byte in value.bytes() {
136 if super::is_unreserved(byte) {
137 self.byte(byte)?;
138 } else {
139 self.byte(b'%')?;
140 self.byte(super::hex_digit(byte >> 4))?;
141 self.byte(super::hex_digit(byte & 0x0f))?;
142 }
143 }
144 Ok(())
145 }
146
147 pub fn form_component(&mut self, value: &str) -> Result<(), E> {
154 for byte in value.bytes() {
155 match byte {
156 b' ' => self.byte(b'+')?,
157 b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'*' | b'-' | b'.' | b'_' => {
158 self.byte(byte)?;
159 }
160 _ => {
161 self.byte(b'%')?;
162 self.byte(super::hex_digit(byte >> 4))?;
163 self.byte(super::hex_digit(byte & 0x0f))?;
164 }
165 }
166 }
167 Ok(())
168 }
169
170 pub fn json_string_escaped(&mut self, value: &str) -> Result<(), E> {
172 for byte in value.bytes() {
173 match byte {
174 b'"' => self.string("\\\"")?,
175 b'\\' => self.string("\\\\")?,
176 b'\n' => self.string("\\n")?,
177 b'\r' => self.string("\\r")?,
178 b'\t' => self.string("\\t")?,
179 0x00..=0x1f => {
180 self.string("\\u00")?;
181 self.byte(super::hex_digit(byte >> 4))?;
182 self.byte(super::hex_digit(byte & 0x0f))?;
183 }
184 _ => self.byte(byte)?,
185 }
186 }
187 Ok(())
188 }
189
190 pub fn json_string(&mut self, value: &str) -> Result<(), E> {
192 self.byte(b'"')?;
193 self.json_string_escaped(value)?;
194 self.byte(b'"')
195 }
196
197 pub fn query_pair(&mut self, first: &mut bool, key: &str, value: &str) -> Result<(), E> {
199 self.query_separator(first)?;
200 self.percent_encoded(key)?;
201 self.byte(b'=')?;
202 self.percent_encoded(value)
203 }
204
205 pub fn query_u64(&mut self, first: &mut bool, key: &str, value: u64) -> Result<(), E> {
207 self.query_separator(first)?;
208 self.percent_encoded(key)?;
209 self.byte(b'=')?;
210 self.u64(value)
211 }
212
213 pub fn query_separator(&mut self, first: &mut bool) -> Result<(), E> {
215 if *first {
216 *first = false;
217 Ok(())
218 } else {
219 self.byte(b'&')
220 }
221 }
222
223 fn checked_end(&self, additional: usize) -> Result<usize, E> {
224 let end = self.len.checked_add(additional).ok_or(self.error)?;
225 if end > self.max_len {
226 return Err(self.error);
227 }
228 Ok(end)
229 }
230}
231
232pub fn encode_snapshot<S: Copy, E: Copy>(
249 snapshot: S,
250 output: &mut [u8],
251 error: E,
252 encode: for<'encoder> fn(S, &mut SnapshotEncoder<'encoder, E>) -> Result<(), E>,
253) -> Result<usize, E> {
254 encode_snapshot_bounded(snapshot, output, usize::MAX, error, encode)
255}
256
257pub fn measure_snapshot<S: Copy, E: Copy>(
259 snapshot: S,
260 error: E,
261 encode: for<'encoder> fn(S, &mut SnapshotEncoder<'encoder, E>) -> Result<(), E>,
262) -> Result<usize, E> {
263 measure_snapshot_bounded(snapshot, usize::MAX, error, encode)
264}
265
266pub fn measure_snapshot_bounded<S: Copy, E: Copy>(
268 snapshot: S,
269 max_len: usize,
270 error: E,
271 encode: for<'encoder> fn(S, &mut SnapshotEncoder<'encoder, E>) -> Result<(), E>,
272) -> Result<usize, E> {
273 let mut measure = SnapshotEncoder::measuring(max_len, error);
274 encode(snapshot, &mut measure)?;
275 Ok(measure.len())
276}
277
278pub fn encode_snapshot_bounded<S: Copy, E: Copy>(
283 snapshot: S,
284 output: &mut [u8],
285 max_len: usize,
286 error: E,
287 encode: for<'encoder> fn(S, &mut SnapshotEncoder<'encoder, E>) -> Result<(), E>,
288) -> Result<usize, E> {
289 let required = measure_snapshot_bounded(snapshot, max_len, error, encode)?;
290 let target = output.get_mut(..required).ok_or(error)?;
291 let mut rollback = EncodeRollback::new(target);
292
293 let write_matches = {
294 let mut writer = SnapshotEncoder::writing(&mut *rollback.target, max_len, error);
295 encode(snapshot, &mut writer).is_ok() && writer.len() == required
296 };
297 if !write_matches {
298 return Err(error);
299 }
300
301 let verify_matches = {
302 let mut verifier = SnapshotEncoder::verifying(&*rollback.target, max_len, error);
303 encode(snapshot, &mut verifier).is_ok() && verifier.len() == required
304 };
305 if !verify_matches {
306 return Err(error);
307 }
308 rollback.disarm();
309 Ok(required)
310}
311
312#[cfg(test)]
313mod tests {
314 use core::sync::atomic::{AtomicUsize, Ordering};
315
316 use super::{SnapshotEncoder, encode_snapshot, encode_snapshot_bounded};
317
318 #[derive(Clone, Copy, Debug, Eq, PartialEq)]
319 enum TestError {
320 Rejected,
321 }
322
323 #[test]
324 fn exact_snapshot_encoding_preserves_the_tail() {
325 let mut output = [0xA5_u8; 16];
326 let result = encode_snapshot(
327 ("key", 42_u64),
328 &mut output,
329 TestError::Rejected,
330 |(key, value), encoder| {
331 encoder.string(key)?;
332 encoder.byte(b'=')?;
333 encoder.u64(value)
334 },
335 );
336
337 assert_eq!(result, Ok(6));
338 assert_eq!(output.get(..6), Some(b"key=42".as_slice()));
339 assert!(
340 output
341 .get(6..)
342 .is_some_and(|tail| tail.iter().all(|b| *b == 0xA5))
343 );
344 }
345
346 #[test]
347 fn every_undersized_capacity_is_unchanged() {
348 for capacity in 0..6 {
349 let mut output = [0xA5_u8; 6];
350 let result = encode_snapshot(
351 "secret",
352 output.get_mut(..capacity).unwrap_or_default(),
353 TestError::Rejected,
354 |value, encoder| encoder.string(value),
355 );
356 assert_eq!(result, Err(TestError::Rejected));
357 assert_eq!(output, [0xA5; 6]);
358 }
359 }
360
361 #[test]
362 fn aggregate_cap_is_checked_before_writing() {
363 let mut output = [0xA5_u8; 8];
364 assert_eq!(
365 encode_snapshot_bounded(
366 "1234",
367 &mut output,
368 3,
369 TestError::Rejected,
370 |value, encoder| encoder.string(value),
371 ),
372 Err(TestError::Rejected)
373 );
374 assert_eq!(output, [0xA5; 8]);
375 }
376
377 #[test]
378 fn form_component_uses_the_reviewed_html_form_grammar() {
379 let mut output = [0xA5_u8; 64];
380 let result = encode_snapshot(
381 "AZaz09 *-._~+&=\0\né",
382 &mut output,
383 TestError::Rejected,
384 |value, encoder| encoder.form_component(value),
385 );
386 let expected = b"AZaz09+*-._%7E%2B%26%3D%00%0A%C3%A9";
387
388 assert_eq!(result, Ok(expected.len()));
389 assert_eq!(output.get(..expected.len()), Some(expected.as_slice()));
390 }
391
392 #[test]
393 fn arithmetic_overflow_is_reported() {
394 let mut encoder = SnapshotEncoder::measuring(usize::MAX, TestError::Rejected);
395 encoder.len = usize::MAX;
396 assert_eq!(encoder.byte(b'x'), Err(TestError::Rejected));
397 }
398
399 #[test]
400 fn nondeterministic_same_length_output_is_rejected_and_cleared() {
401 static PASS: AtomicUsize = AtomicUsize::new(0);
402
403 fn changing(
404 _snapshot: (),
405 encoder: &mut SnapshotEncoder<'_, TestError>,
406 ) -> Result<(), TestError> {
407 let pass = PASS.fetch_add(1, Ordering::Relaxed);
408 encoder.byte(if pass == 2 { b'B' } else { b'A' })
409 }
410
411 PASS.store(0, Ordering::Relaxed);
412 let mut output = [0xA5_u8; 4];
413 assert_eq!(
414 encode_snapshot((), &mut output, TestError::Rejected, changing),
415 Err(TestError::Rejected)
416 );
417 assert_eq!(output, [0, 0xA5, 0xA5, 0xA5]);
418 }
419}