1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
//! PEM encoder.

use crate::{
    grammar, Base64Encoder, Error, LineEnding, Result, BASE64_WRAP_WIDTH,
    ENCAPSULATION_BOUNDARY_DELIMITER, POST_ENCAPSULATION_BOUNDARY, PRE_ENCAPSULATION_BOUNDARY,
};
use base64ct::{Base64, Encoding};
use core::str;

#[cfg(feature = "alloc")]
use alloc::string::String;

#[cfg(feature = "std")]
use std::io;

/// Compute the length of a PEM encoded document which encapsulates a
/// Base64-encoded body including line endings every 64 characters.
///
/// The `input_len` parameter specifies the length of the raw input
/// bytes prior to Base64 encoding.
///
/// Note that the current implementation of this function computes an upper
/// bound of the length and the actual encoded document may be slightly shorter
/// (typically 1-byte). Downstream consumers of this function should check the
/// actual encoded length and potentially truncate buffers allocated using this
/// function to estimate the encapsulated size.
///
/// Use [`encoded_len`] (when possible) to obtain a precise length.
///
/// ## Returns
/// - `Ok(len)` on success
/// - `Err(Error::Length)` on length overflow
pub fn encapsulated_len(label: &str, line_ending: LineEnding, input_len: usize) -> Result<usize> {
    encapsulated_len_wrapped(label, BASE64_WRAP_WIDTH, line_ending, input_len)
}

/// Compute the length of a PEM encoded document with the Base64 body
/// line wrapped at the specified `width`.
///
/// This is the same as [`encapsulated_len`], which defaults to a width of 64.
///
/// Note that per [RFC7468 § 2] encoding PEM with any other wrap width besides
/// 64 is technically non-compliant:
///
/// > Generators MUST wrap the base64-encoded lines so that each line
/// > consists of exactly 64 characters except for the final line, which
/// > will encode the remainder of the data (within the 64-character line
/// > boundary)
///
/// [RFC7468 § 2]: https://datatracker.ietf.org/doc/html/rfc7468#section-2
pub fn encapsulated_len_wrapped(
    label: &str,
    line_width: usize,
    line_ending: LineEnding,
    input_len: usize,
) -> Result<usize> {
    if line_width < 4 {
        return Err(Error::Length);
    }

    let base64_len = input_len
        .checked_mul(4)
        .and_then(|n| n.checked_div(3))
        .and_then(|n| n.checked_add(3))
        .ok_or(Error::Length)?
        & !3;

    let base64_len_wrapped = base64_len_wrapped(base64_len, line_width, line_ending)?;
    encapsulated_len_inner(label, line_ending, base64_len_wrapped)
}

/// Get the length of a PEM encoded document with the given bytes and label.
///
/// This function computes a precise length of the PEM encoding of the given
/// `input` data.
///
/// ## Returns
/// - `Ok(len)` on success
/// - `Err(Error::Length)` on length overflow
pub fn encoded_len(label: &str, line_ending: LineEnding, input: &[u8]) -> Result<usize> {
    let base64_len = Base64::encoded_len(input);
    let base64_len_wrapped = base64_len_wrapped(base64_len, BASE64_WRAP_WIDTH, line_ending)?;
    encapsulated_len_inner(label, line_ending, base64_len_wrapped)
}

/// Encode a PEM document according to RFC 7468's "Strict" grammar.
pub fn encode<'o>(
    type_label: &str,
    line_ending: LineEnding,
    input: &[u8],
    buf: &'o mut [u8],
) -> Result<&'o str> {
    let mut encoder = Encoder::new(type_label, line_ending, buf)?;
    encoder.encode(input)?;
    let encoded_len = encoder.finish()?;
    let output = &buf[..encoded_len];

    // Sanity check
    debug_assert!(str::from_utf8(output).is_ok());

    // Ensure `output` contains characters from the lower 7-bit ASCII set
    if output.iter().fold(0u8, |acc, &byte| acc | (byte & 0x80)) == 0 {
        // Use unchecked conversion to avoid applying UTF-8 checks to potentially
        // secret PEM documents (and therefore introducing a potential timing
        // sidechannel)
        //
        // SAFETY: contents of this buffer are controlled entirely by the encoder,
        // which ensures the contents are always a valid (ASCII) subset of UTF-8.
        // It's also additionally sanity checked by two assertions above to ensure
        // the validity (with the always-on runtime check implemented in a
        // constant time-ish manner.
        #[allow(unsafe_code)]
        Ok(unsafe { str::from_utf8_unchecked(output) })
    } else {
        Err(Error::CharacterEncoding)
    }
}

/// Encode a PEM document according to RFC 7468's "Strict" grammar, returning
/// the result as a [`String`].
#[cfg(feature = "alloc")]
pub fn encode_string(label: &str, line_ending: LineEnding, input: &[u8]) -> Result<String> {
    let expected_len = encoded_len(label, line_ending, input)?;
    let mut buf = vec![0u8; expected_len];
    let actual_len = encode(label, line_ending, input, &mut buf)?.len();
    debug_assert_eq!(expected_len, actual_len);
    String::from_utf8(buf).map_err(|_| Error::CharacterEncoding)
}

/// Compute the encapsulated length of Base64 data of the given length.
fn encapsulated_len_inner(
    label: &str,
    line_ending: LineEnding,
    base64_len: usize,
) -> Result<usize> {
    [
        PRE_ENCAPSULATION_BOUNDARY.len(),
        label.as_bytes().len(),
        ENCAPSULATION_BOUNDARY_DELIMITER.len(),
        line_ending.len(),
        base64_len,
        line_ending.len(),
        POST_ENCAPSULATION_BOUNDARY.len(),
        label.as_bytes().len(),
        ENCAPSULATION_BOUNDARY_DELIMITER.len(),
        line_ending.len(),
    ]
    .into_iter()
    .try_fold(0usize, |acc, len| acc.checked_add(len))
    .ok_or(Error::Length)
}

/// Compute Base64 length line-wrapped at the specified width with the given
/// line ending.
fn base64_len_wrapped(
    base64_len: usize,
    line_width: usize,
    line_ending: LineEnding,
) -> Result<usize> {
    base64_len
        .saturating_sub(1)
        .checked_div(line_width)
        .and_then(|lines| lines.checked_mul(line_ending.len()))
        .and_then(|len| len.checked_add(base64_len))
        .ok_or(Error::Length)
}

/// Buffered PEM encoder.
///
/// Stateful buffered encoder type which encodes an input PEM document according
/// to RFC 7468's "Strict" grammar.
pub struct Encoder<'l, 'o> {
    /// PEM type label.
    type_label: &'l str,

    /// Line ending used to wrap Base64.
    line_ending: LineEnding,

    /// Buffered Base64 encoder.
    base64: Base64Encoder<'o>,
}

impl<'l, 'o> Encoder<'l, 'o> {
    /// Create a new PEM [`Encoder`] with the default options which
    /// writes output into the provided buffer.
    ///
    /// Uses the default 64-character line wrapping.
    pub fn new(type_label: &'l str, line_ending: LineEnding, out: &'o mut [u8]) -> Result<Self> {
        Self::new_wrapped(type_label, BASE64_WRAP_WIDTH, line_ending, out)
    }

    /// Create a new PEM [`Encoder`] which wraps at the given line width.
    ///
    /// Note that per [RFC7468 § 2] encoding PEM with any other wrap width besides
    /// 64 is technically non-compliant:
    ///
    /// > Generators MUST wrap the base64-encoded lines so that each line
    /// > consists of exactly 64 characters except for the final line, which
    /// > will encode the remainder of the data (within the 64-character line
    /// > boundary)
    ///
    /// This method is provided with the intended purpose of implementing the
    /// OpenSSH private key format, which uses a non-standard wrap width of 70.
    ///
    /// [RFC7468 § 2]: https://datatracker.ietf.org/doc/html/rfc7468#section-2
    pub fn new_wrapped(
        type_label: &'l str,
        line_width: usize,
        line_ending: LineEnding,
        mut out: &'o mut [u8],
    ) -> Result<Self> {
        grammar::validate_label(type_label.as_bytes())?;

        for boundary_part in [
            PRE_ENCAPSULATION_BOUNDARY,
            type_label.as_bytes(),
            ENCAPSULATION_BOUNDARY_DELIMITER,
            line_ending.as_bytes(),
        ] {
            if out.len() < boundary_part.len() {
                return Err(Error::Length);
            }

            let (part, rest) = out.split_at_mut(boundary_part.len());
            out = rest;

            part.copy_from_slice(boundary_part);
        }

        let base64 = Base64Encoder::new_wrapped(out, line_width, line_ending)?;

        Ok(Self {
            type_label,
            line_ending,
            base64,
        })
    }

    /// Get the PEM type label used for this document.
    pub fn type_label(&self) -> &'l str {
        self.type_label
    }

    /// Encode the provided input data.
    ///
    /// This method can be called as many times as needed with any sized input
    /// to write data encoded data into the output buffer, so long as there is
    /// sufficient space in the buffer to handle the resulting Base64 encoded
    /// data.
    pub fn encode(&mut self, input: &[u8]) -> Result<()> {
        self.base64.encode(input)?;
        Ok(())
    }

    /// Borrow the inner [`Base64Encoder`].
    pub fn base64_encoder(&mut self) -> &mut Base64Encoder<'o> {
        &mut self.base64
    }

    /// Finish encoding PEM, writing the post-encapsulation boundary.
    ///
    /// On success, returns the total number of bytes written to the output
    /// buffer.
    pub fn finish(self) -> Result<usize> {
        let (base64, mut out) = self.base64.finish_with_remaining()?;

        for boundary_part in [
            self.line_ending.as_bytes(),
            POST_ENCAPSULATION_BOUNDARY,
            self.type_label.as_bytes(),
            ENCAPSULATION_BOUNDARY_DELIMITER,
            self.line_ending.as_bytes(),
        ] {
            if out.len() < boundary_part.len() {
                return Err(Error::Length);
            }

            let (part, rest) = out.split_at_mut(boundary_part.len());
            out = rest;

            part.copy_from_slice(boundary_part);
        }

        encapsulated_len_inner(self.type_label, self.line_ending, base64.len())
    }
}

#[cfg(feature = "std")]
impl<'l, 'o> io::Write for Encoder<'l, 'o> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.encode(buf)?;
        Ok(buf.len())
    }

    fn flush(&mut self) -> io::Result<()> {
        // TODO(tarcieri): return an error if there's still data remaining in the buffer?
        Ok(())
    }
}