Skip to main content

mail_builder/encoders/
base64.rs

1/*
2 * SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <hello@stalw.art>
3 *
4 * SPDX-License-Identifier: Apache-2.0 OR MIT
5 */
6
7use crate::writer::{IoWriter, Writer};
8use std::io::{self, Write};
9
10const CHARPAD: u8 = b'=';
11const CRLF: [u8; 2] = *b"\r\n";
12const LINE_GROUPS: usize = 19;
13const LINE_INPUT: usize = LINE_GROUPS * 3;
14const LINE_OUTPUT: usize = LINE_GROUPS * 4;
15const LINE_TOTAL: usize = LINE_OUTPUT + 2;
16const INLINE_BLOCK_INPUT: usize = 12288;
17const INLINE_BLOCK_OUTPUT: usize = INLINE_BLOCK_INPUT / 3 * 4;
18const WRAPPED_BLOCK_LINES: usize = 52;
19const WRAPPED_BLOCK_INPUT: usize = WRAPPED_BLOCK_LINES * LINE_INPUT;
20const WRAPPED_BLOCK_OUTPUT: usize = WRAPPED_BLOCK_LINES * LINE_TOTAL;
21
22const ALPHABET: &[u8; 64] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
23
24const PAIRS: [u16; 4096] = {
25    let mut table = [0u16; 4096];
26    let mut index = 0;
27    while index < table.len() {
28        table[index] = u16::from_ne_bytes([ALPHABET[index >> 6], ALPHABET[index & 0x3f]]);
29        index += 1;
30    }
31    table
32};
33
34/// Number of bytes `base64_encode_slice` writes for `input_len` input bytes.
35pub const fn base64_encoded_len(input_len: usize) -> usize {
36    input_len.div_ceil(3) * 4
37}
38
39#[inline(always)]
40fn encode_pair(index: usize) -> u32 {
41    let high = PAIRS[(index >> 12) & 0xfff] as u32;
42    let low = PAIRS[index & 0xfff] as u32;
43    if cfg!(target_endian = "little") {
44        high | (low << 16)
45    } else {
46        (high << 16) | low
47    }
48}
49
50#[inline(always)]
51fn encode_group(group: [u8; 3]) -> [u8; 4] {
52    let [b0, b1, b2] = group;
53    encode_pair(((b0 as usize) << 16) | ((b1 as usize) << 8) | b2 as usize).to_ne_bytes()
54}
55
56#[inline(always)]
57fn encode_block(block: &[u8; 12], slot: &mut [u8; 16]) {
58    let (Some(head), Some(foot)) = (block.first_chunk::<8>(), block.last_chunk::<8>()) else {
59        return;
60    };
61    let first = u64::from_be_bytes(*head);
62    let second = u64::from_be_bytes(*foot);
63    if let [a, b, c, d] = slot.as_chunks_mut::<4>().0 {
64        *a = encode_pair((first >> 40) as usize).to_ne_bytes();
65        *b = encode_pair((first >> 16) as usize & 0xff_ffff).to_ne_bytes();
66        *c = encode_pair((second >> 24) as usize & 0xff_ffff).to_ne_bytes();
67        *d = encode_pair(second as usize & 0xff_ffff).to_ne_bytes();
68    }
69}
70
71#[inline(always)]
72fn encode_nine(nine: &[u8; 9], slot: &mut [u8; 12]) {
73    let (Some(head), Some(&last)) = (nine.first_chunk::<8>(), nine.last()) else {
74        return;
75    };
76    let word = u64::from_be_bytes(*head);
77    if let [a, b, c] = slot.as_chunks_mut::<4>().0 {
78        *a = encode_pair((word >> 40) as usize).to_ne_bytes();
79        *b = encode_pair((word >> 16) as usize & 0xff_ffff).to_ne_bytes();
80        *c = encode_pair(((word as usize & 0xffff) << 8) | last as usize).to_ne_bytes();
81    }
82}
83
84#[inline(always)]
85fn encode_line(line: &[u8; LINE_INPUT], slot: &mut [u8; LINE_TOTAL]) {
86    let (blocks, rest) = line.as_chunks::<12>();
87    let (Some(nine), Some((head, tail))) =
88        (rest.first_chunk::<9>(), slot.split_first_chunk_mut::<64>())
89    else {
90        return;
91    };
92    for (block, slot) in blocks.iter().zip(head.as_chunks_mut::<16>().0.iter_mut()) {
93        encode_block(block, slot);
94    }
95    let Some((twelve, crlf)) = tail.split_first_chunk_mut::<12>() else {
96        return;
97    };
98    encode_nine(nine, twelve);
99    if let Some(crlf) = crlf.first_chunk_mut::<2>() {
100        *crlf = CRLF;
101    }
102}
103
104#[inline(always)]
105fn encode_tail(tail: &[u8]) -> [u8; 4] {
106    match *tail {
107        [b0] => {
108            let [c0, c1, _, _] = encode_group([b0, 0, 0]);
109            [c0, c1, CHARPAD, CHARPAD]
110        }
111        [b0, b1] => {
112            let [c0, c1, c2, _] = encode_group([b0, b1, 0]);
113            [c0, c1, c2, CHARPAD]
114        }
115        _ => [CHARPAD; 4],
116    }
117}
118
119#[inline(always)]
120fn encode_exact(input: &[u8], output: &mut [u8]) {
121    let (blocks, rest) = input.as_chunks::<12>();
122    let Some((slots, rest_slots)) = output.split_at_mut_checked(blocks.len() * 16) else {
123        return;
124    };
125    for (block, slot) in blocks.iter().zip(slots.as_chunks_mut::<16>().0.iter_mut()) {
126        encode_block(block, slot);
127    }
128
129    let (groups, tail) = rest.as_chunks::<3>();
130    let Some((group_slots, tail_slot)) = rest_slots.split_at_mut_checked(groups.len() * 4) else {
131        return;
132    };
133    for (group, slot) in groups
134        .iter()
135        .zip(group_slots.as_chunks_mut::<4>().0.iter_mut())
136    {
137        *slot = encode_group(*group);
138    }
139    if !tail.is_empty()
140        && let Some(slot) = tail_slot.first_chunk_mut::<4>()
141    {
142        *slot = encode_tail(tail);
143    }
144}
145
146#[inline(never)]
147fn encode_truncated(input: &[u8], output: &mut [u8]) -> usize {
148    let groups = (input.len() / 3).min(output.len() / 4);
149    let Some((body, tail)) = input.split_at_checked(groups * 3) else {
150        return 0;
151    };
152    let Some((slots, spare)) = output.split_at_mut_checked(groups * 4) else {
153        return 0;
154    };
155
156    encode_exact(body, slots);
157    let mut written = groups * 4;
158    if !tail.is_empty()
159        && let Some(slot) = spare.first_chunk_mut::<4>()
160    {
161        *slot = encode_tail(tail);
162        written += 4;
163    }
164
165    written
166}
167
168/// Encodes `input` as standard base64 (RFC 4648, `+` and `/`, `=` padding,
169/// no line breaks) into `output` and returns the number of bytes written.
170///
171/// `output` must hold at least [`base64_encoded_len`] bytes; otherwise only
172/// the complete 4-byte groups that fit are written.
173#[inline]
174pub fn base64_encode_slice(input: &[u8], output: &mut [u8]) -> usize {
175    let written = base64_encoded_len(input.len());
176    match output.get_mut(..written) {
177        Some(slots) => {
178            encode_exact(input, slots);
179            written
180        }
181        None => encode_truncated(input, output),
182    }
183}
184
185pub(crate) fn base64_encode_inline(input: &[u8], output: &mut impl Writer) -> usize {
186    let (blocks, rest) = input.as_chunks::<INLINE_BLOCK_INPUT>();
187    let mut written = 0;
188
189    for block in blocks {
190        output.write_with(INLINE_BLOCK_OUTPUT, |region| {
191            base64_encode_slice(block, region)
192        });
193        written += INLINE_BLOCK_OUTPUT;
194    }
195
196    if !rest.is_empty() {
197        let len = base64_encoded_len(rest.len());
198        output.write_with(len, |region| base64_encode_slice(rest, region));
199        written += len;
200    }
201
202    written
203}
204
205#[inline(always)]
206fn fill_lines(input: &[u8], region: &mut [u8]) -> usize {
207    let (lines, rest) = input.as_chunks::<LINE_INPUT>();
208    let Some((slots, spare)) = region.split_at_mut_checked(lines.len() * LINE_TOTAL) else {
209        return 0;
210    };
211    let mut written = 0;
212
213    for (line, slot) in lines
214        .iter()
215        .zip(slots.as_chunks_mut::<LINE_TOTAL>().0.iter_mut())
216    {
217        encode_line(line, slot);
218        written += LINE_TOTAL;
219    }
220
221    if !rest.is_empty() {
222        let encoded = base64_encode_slice(rest, spare);
223        written += encoded;
224        if let Some((_, after)) = spare.split_at_mut_checked(encoded)
225            && let Some(crlf) = after.first_chunk_mut::<2>()
226        {
227            *crlf = CRLF;
228            written += 2;
229        }
230    }
231
232    written
233}
234
235pub(crate) fn base64_encode_wrapped(input: &[u8], output: &mut impl Writer) -> usize {
236    let (blocks, rest) = input.as_chunks::<WRAPPED_BLOCK_INPUT>();
237    let mut written = 0;
238
239    for block in blocks {
240        output.write_with(WRAPPED_BLOCK_OUTPUT, |region| fill_lines(block, region));
241        written += WRAPPED_BLOCK_LINES * LINE_OUTPUT;
242    }
243
244    if !rest.is_empty() {
245        let encoded = base64_encoded_len(rest.len());
246        let len = encoded + rest.len().div_ceil(LINE_INPUT) * 2;
247        output.write_with(len, |region| fill_lines(rest, region));
248        written += encoded;
249    }
250
251    written
252}
253
254#[repr(transparent)]
255#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
256pub struct Base64Encoder {
257    wrap_lines: bool,
258}
259
260impl Base64Encoder {
261    #[inline(always)]
262    pub fn new() -> Self {
263        Self { wrap_lines: false }
264    }
265
266    #[inline(always)]
267    pub fn wrap_lines(mut self) -> Self {
268        self.wrap_lines = true;
269        self
270    }
271
272    #[inline(always)]
273    pub fn encode(&self, input: &[u8]) -> io::Result<Vec<u8>> {
274        let encoded = base64_encoded_len(input.len());
275        let capacity = if self.wrap_lines {
276            encoded + encoded.div_ceil(LINE_OUTPUT) * 2
277        } else {
278            encoded
279        };
280        let mut buf = Vec::with_capacity(capacity);
281        self.encode_into(input, &mut buf);
282        Ok(buf)
283    }
284
285    #[inline(always)]
286    pub fn encode_to_writer(&self, input: &[u8], output: &mut impl Write) -> io::Result<usize> {
287        let capacity = base64_encoded_len(input.len())
288            .saturating_add(input.len().div_ceil(LINE_INPUT).saturating_mul(2))
289            .clamp(64, 64 * 1024);
290        let mut writer = IoWriter::with_capacity(capacity, output);
291        let bytes_written = self.encode_into(input, &mut writer);
292        writer.into_result().map(|_| bytes_written)
293    }
294
295    #[inline(always)]
296    pub fn encode_into(&self, input: &[u8], output: &mut impl Writer) -> usize {
297        if self.wrap_lines {
298            base64_encode_wrapped(input, output)
299        } else {
300            base64_encode_inline(input, output)
301        }
302    }
303}
304
305#[cfg(test)]
306#[allow(clippy::items_after_test_module)]
307mod tests {
308    use super::*;
309
310    #[test]
311    fn encode_base64() {
312        for (input, expected_result) in [
313            ("Test".to_string(), "VGVzdA==\r\n"),
314            ("Ye".to_string(), "WWU=\r\n"),
315            ("A".to_string(), "QQ==\r\n"),
316            ("ro".to_string(), "cm8=\r\n"),
317            (
318                "Are you a Shimano or Campagnolo person?".to_string(),
319                "QXJlIHlvdSBhIFNoaW1hbm8gb3IgQ2FtcGFnbm9sbyBwZXJzb24/\r\n",
320            ),
321            (
322                "<!DOCTYPE html>\n<html>\n<body>\n</body>\n</html>\n".to_string(),
323                "PCFET0NUWVBFIGh0bWw+CjxodG1sPgo8Ym9keT4KPC9ib2R5Pgo8L2h0bWw+Cg==\r\n",
324            ),
325            ("áéíóú".to_string(), "w6HDqcOtw7PDug==\r\n"),
326            (
327                " ".repeat(100),
328                concat!(
329                    "ICAgICAgICAgICAgICAgICAgICAgICAgICAgICAg",
330                    "ICAgICAgICAgICAgICAgICAgICAgICAgICAg\r\n",
331                    "ICAgICAgICAgICAgICAgICAgICAgICAgICAgICAg",
332                    "ICAgICAgICAgICAgIA==\r\n",
333                ),
334            ),
335        ] {
336            let mut output = Vec::new();
337            base64_encode_wrapped(input.as_bytes(), &mut output);
338            assert_eq!(std::str::from_utf8(&output).unwrap(), expected_result);
339
340            let mut inline = Vec::new();
341            base64_encode_inline(input.as_bytes(), &mut inline);
342            assert_eq!(
343                std::str::from_utf8(&inline).unwrap(),
344                expected_result.replace("\r\n", "")
345            );
346
347            let mut slice = vec![0u8; base64_encoded_len(input.len())];
348            let written = base64_encode_slice(input.as_bytes(), &mut slice);
349            assert_eq!(written, slice.len());
350            assert_eq!(slice, inline);
351        }
352    }
353
354    fn reference_encode(input: &[u8]) -> Vec<u8> {
355        let alphabet = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
356        let mut out = Vec::new();
357        for group in input.chunks(3) {
358            let word = group.iter().enumerate().fold(0u32, |word, (pos, &byte)| {
359                word | (byte as u32) << (16 - pos * 8)
360            });
361            let chars = [
362                alphabet[(word >> 18) as usize & 63],
363                alphabet[(word >> 12) as usize & 63],
364                alphabet[(word >> 6) as usize & 63],
365                alphabet[word as usize & 63],
366            ];
367            out.extend_from_slice(&chars[..group.len() + 1]);
368            out.extend(std::iter::repeat_n(b'=', 3 - group.len()));
369        }
370        out
371    }
372
373    #[test]
374    fn every_length_matches_a_reference_encoder() {
375        let mut state = 0x9E37_79B9_7F4A_7C15u64;
376        let mut input = Vec::new();
377        for len in 0..=600usize {
378            input.clear();
379            while input.len() < len {
380                state ^= state << 13;
381                state ^= state >> 7;
382                state ^= state << 17;
383                input.extend_from_slice(&state.to_le_bytes());
384            }
385            input.truncate(len);
386            let expected = reference_encode(&input);
387
388            let mut slice = vec![0xAAu8; base64_encoded_len(len) + 8];
389            let written = base64_encode_slice(&input, &mut slice);
390            assert_eq!(written, expected.len(), "length {len}");
391            assert_eq!(&slice[..written], expected.as_slice(), "length {len}");
392            assert!(
393                slice[written..].iter().all(|&byte| byte == 0xAA),
394                "length {len}"
395            );
396
397            let mut inline = Vec::new();
398            assert_eq!(base64_encode_inline(&input, &mut inline), expected.len());
399            assert_eq!(inline, expected, "inline {len}");
400
401            let mut windowed = Vec::new();
402            for window in input.chunks(192) {
403                let mut buffer = [0u8; 256];
404                let written = base64_encode_slice(window, &mut buffer);
405                windowed.extend_from_slice(&buffer[..written]);
406            }
407            assert_eq!(windowed, expected, "windowed {len}");
408
409            let mut wrapped = Vec::new();
410            base64_encode_wrapped(&input, &mut wrapped);
411            let joined: Vec<u8> = wrapped
412                .iter()
413                .copied()
414                .filter(|&b| b != b'\r' && b != b'\n')
415                .collect();
416            assert_eq!(joined, expected, "wrapped {len}");
417            for line in wrapped.split(|&b| b == b'\n') {
418                assert!(line.len() <= 77, "wrapped line {len}");
419            }
420            if len > 0 {
421                assert!(wrapped.ends_with(b"\r\n"), "wrapped terminator {len}");
422            }
423
424            let encoded = Base64Encoder::new().wrap_lines().encode(&input).unwrap();
425            assert_eq!(encoded, wrapped);
426            assert_eq!(encoded.capacity(), encoded.len(), "exact capacity {len}");
427        }
428    }
429
430    #[test]
431    fn slice_encoding_truncates_when_output_is_short() {
432        let mut output = [0u8; 6];
433        assert_eq!(base64_encode_slice(b"abcdef", &mut output), 4);
434        assert_eq!(&output[..4], b"YWJj");
435        assert_eq!(base64_encode_slice(b"abcd", &mut output), 4);
436    }
437}
438
439/*
440 * Table adapted from Nick Galbreath's "High performance base64 encoder / decoder"
441 *
442 * Copyright 2005, 2006, 2007 Nick Galbreath -- nickg [at] modp [dot] com
443 * All rights reserved.
444 *
445 * http://code.google.com/p/stringencoders/
446 *
447 * Released under bsd license.
448 *
449 */
450
451pub static E0: &[u8] = b"AAAABBBBCCCCDDDDEEEEFFFFGGGGHHHHIIIIJJJJKKKKLLLLMMMMNNNNOOOOPPPPQQQQRRRRSSSSTTTTUUUUVVVVWWWWXXXXYYYYZZZZaaaabbbbccccddddeeeeffffgggghhhhiiiijjjjkkkkllllmmmmnnnnooooppppqqqqrrrrssssttttuuuuvvvvwwwwxxxxyyyyzzzz0000111122223333444455556666777788889999++++////";
452pub static E1: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
453pub static E2: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";