pgp/
line_writer.rs

1//! # Line writer module
2
3use std::io;
4
5use generic_array::{
6    typenum::{Sum, Unsigned, U2},
7    ArrayLength, GenericArray,
8};
9
10const CRLF: [u8; 2] = [b'\r', b'\n'];
11const CR: [u8; 1] = [b'\r'];
12const LF: [u8; 1] = [b'\n'];
13
14#[derive(Clone, Copy, PartialEq, Eq)]
15#[repr(u8)]
16pub enum LineBreak {
17    Crlf,
18    Lf,
19    Cr,
20}
21
22impl AsRef<[u8]> for LineBreak {
23    fn as_ref(&self) -> &[u8] {
24        match self {
25            LineBreak::Crlf => &CRLF[..],
26            LineBreak::Lf => &LF[..],
27            LineBreak::Cr => &CR[..],
28        }
29    }
30}
31
32/// A `Write` implementation that splits any written bytes into the given length lines.
33///
34///
35/// # Panics
36///
37/// Calling `write()` after `finish()` is invalid and will panic.
38pub struct LineWriter<'a, W, N>
39where
40    W: io::Write,
41    N: Unsigned + ArrayLength<u8>,
42    N: std::ops::Add<U2>,
43    Sum<N, U2>: ArrayLength<u8>,
44{
45    /// Which kind of line break to insert.
46    line_break: LineBreak,
47    /// Where encoded data is written to.
48    w: &'a mut W,
49    /// Holds a partial chunk, if any, after the last `write()`, so that we may then fill the chunk
50    /// with the next `write()`, write it, then proceed with the rest of the input normally.
51    extra: GenericArray<u8, N>,
52    /// How much of `extra` is occupied, in `[0, N]`.
53    extra_len: usize,
54    buffer: GenericArray<u8, Sum<N, U2>>,
55    /// True iff partial last chunk has been written.
56    finished: bool,
57    /// panic safety: don't write again in destructor if writer panicked while we were writing to it
58    panicked: bool,
59}
60
61impl<'a, W, N> LineWriter<'a, W, N>
62where
63    W: 'a + io::Write,
64    N: Unsigned + ArrayLength<u8>,
65    N: std::ops::Add<U2>,
66    Sum<N, U2>: ArrayLength<u8>,
67{
68    /// Creates a new encoder around an existing writer.
69    pub fn new(w: &'a mut W, line_break: LineBreak) -> Self {
70        LineWriter {
71            line_break,
72            w,
73            extra: Default::default(),
74            buffer: Default::default(),
75            extra_len: 0,
76            finished: false,
77            panicked: false,
78        }
79    }
80
81    /// Write all remaining buffered data.
82    ///
83    /// Once this succeeds, no further writes can be performed.
84    ///
85    /// # Errors
86    ///
87    /// Assuming the wrapped writer obeys the `Write` contract, if this returns `Err`, no data was
88    /// written, and `finish()` may be retried if appropriate for the type of error, etc.
89    pub fn finish(&mut self) -> io::Result<()> {
90        if self.finished {
91            return Ok(());
92        };
93
94        if self.extra_len > 0 {
95            self.panicked = true;
96            self.w.write_all(&self.extra[..self.extra_len])?;
97            self.w.write_all(self.line_break.as_ref())?;
98            self.panicked = false;
99            // write succeeded, do not write the encoding of extra again if finish() is retried
100            self.extra_len = 0;
101        }
102
103        self.finished = true;
104        Ok(())
105    }
106}
107
108impl<'a, W, N> io::Write for LineWriter<'a, W, N>
109where
110    W: 'a + io::Write,
111    N: Unsigned + ArrayLength<u8>,
112    N: std::ops::Add<U2>,
113    Sum<N, U2>: ArrayLength<u8>,
114{
115    fn write(&mut self, input: &[u8]) -> io::Result<usize> {
116        if self.finished {
117            panic!("Cannot write more after calling finish()");
118        }
119
120        if input.is_empty() {
121            return Ok(0);
122        }
123
124        // The contract of `Write::write` places some constraints on this implementation:
125        // - a call to `write()` represents at most one call to a wrapped `Write`, so we can't
126        // iterate over the input and encode multiple chunks.
127        // - Errors mean that "no bytes were written to this writer", so we need to reset the
128        // internal state to what it was before the error occurred
129
130        let sl = N::to_usize();
131        let line_break = self.line_break.as_ref();
132
133        let orig_extra_len = self.extra_len;
134
135        // process leftover stuff from last write
136        if self.extra_len + input.len() < sl {
137            // still not enough
138            self.extra_len += input.len();
139            self.extra[orig_extra_len..self.extra_len].copy_from_slice(input);
140            Ok(input.len())
141        } else {
142            let mut buffer_pos = 0;
143            let mut input_pos = 0;
144
145            if self.extra_len > 0 {
146                let copied = ::std::cmp::min(orig_extra_len, self.buffer.len());
147                self.buffer[buffer_pos..buffer_pos + copied].copy_from_slice(&self.extra[..copied]);
148                self.extra_len -= copied;
149                buffer_pos += copied;
150            }
151
152            if buffer_pos < sl {
153                let missing = ::std::cmp::min(sl - buffer_pos, input.len() - input_pos);
154
155                self.buffer[buffer_pos..buffer_pos + missing]
156                    .copy_from_slice(&input[input_pos..input_pos + missing]);
157
158                buffer_pos += missing;
159                input_pos += missing;
160            }
161
162            // still not enough
163            if buffer_pos < sl {
164                return Ok(input_pos);
165            }
166
167            // insert line break
168            self.buffer[buffer_pos..buffer_pos + line_break.len()].copy_from_slice(line_break);
169            buffer_pos += line_break.len();
170
171            self.panicked = true;
172            let r = self.w.write_all(&self.buffer[..buffer_pos]);
173            self.panicked = false;
174
175            match r {
176                Ok(_) => Ok(input_pos),
177                Err(err) => {
178                    // in case we filled and encoded `extra`, reset extra_len
179                    self.extra_len = orig_extra_len;
180                    Err(err)
181                }
182            }
183        }
184    }
185
186    /// Because this is usually treated as OK to call multiple times, it will *not* flush any
187    /// incomplete chunks of input or write padding.
188    fn flush(&mut self) -> io::Result<()> {
189        self.w.flush()
190    }
191}
192
193impl<'a, W, N> Drop for LineWriter<'a, W, N>
194where
195    W: 'a + io::Write,
196    N: Unsigned + ArrayLength<u8>,
197    N: std::ops::Add<U2>,
198    Sum<N, U2>: ArrayLength<u8>,
199{
200    fn drop(&mut self) {
201        if !self.panicked {
202            // like `BufWriter`, ignore errors during drop
203            let _ = self.finish();
204        }
205    }
206}
207
208#[cfg(test)]
209mod tests {
210    #![allow(clippy::unwrap_used)]
211
212    use std::io::Write;
213
214    use base64::engine::general_purpose;
215    use generic_array::typenum::{self, U10};
216
217    use super::*;
218
219    /// The same as the std lib, but doesn't choke on write 0. This is a hack, to be compatible with
220    /// rust-base64.
221    fn write_all(writer: &mut impl io::Write, mut buf: &[u8]) -> io::Result<()> {
222        while !buf.is_empty() {
223            match writer.write(buf) {
224                Ok(0) => {}
225                Ok(n) => buf = &buf[n..],
226                Err(ref e) if e.kind() == io::ErrorKind::Interrupted => {}
227                Err(e) => return Err(e),
228            }
229        }
230        Ok(())
231    }
232
233    #[test]
234    fn simple_writes() {
235        let mut buf = Vec::new();
236
237        {
238            let mut w = LineWriter::<_, U10>::new(&mut buf, LineBreak::Crlf);
239
240            // short write
241            assert_eq!(w.write(&[0, 1, 2, 3]).unwrap(), 4);
242            assert_eq!(w.write(&[4, 5, 6, 7]).unwrap(), 4);
243            assert_eq!(w.write(&[8, 9, 10, 11]).unwrap(), 2);
244            assert_eq!(w.write(&[10, 11]).unwrap(), 2);
245
246            // writer dropped, should flush now
247        }
248
249        assert_eq!(
250            &buf[..],
251            &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, b'\r', b'\n', 10, 11, b'\r', b'\n'][..]
252        );
253    }
254
255    macro_rules! test_len {
256        ( $name:ident, $len:ty ) => {
257            #[test]
258            fn $name() {
259                use rand::{Rng, SeedableRng};
260                use rand_xorshift::XorShiftRng;
261
262                let rng = &mut XorShiftRng::from_seed([
263                    0x3, 0x8, 0x3, 0xe, 0x3, 0x8, 0x3, 0xe, 0x3, 0x8, 0x3, 0xe, 0x3, 0x8, 0x3, 0xe,
264                ]);
265                let mut buf = Vec::new();
266
267                let mut list = Vec::new();
268                {
269                    let mut w = LineWriter::<_, $len>::new(&mut buf, LineBreak::Crlf);
270                    for i in 0..100 {
271                        let data = (0..i).map(|_| rng.gen()).collect::<Vec<_>>();
272                        w.write_all(&data).unwrap();
273                        list.extend(&data);
274                    }
275                }
276
277                let len = <$len as typenum::Unsigned>::to_usize();
278                let mut expected: Vec<u8> = list.chunks(len).fold(Vec::new(), |mut acc, line| {
279                    acc.extend(line);
280
281                    if line.len() == len {
282                        acc.push(b'\r');
283                        acc.push(b'\n');
284                    }
285                    acc
286                });
287
288                if expected.len() % len != 0 {
289                    expected.push(b'\r');
290                    expected.push(b'\n');
291                }
292
293                assert_eq!(&buf[..], &expected[..]);
294            }
295        };
296    }
297
298    test_len!(test_break_line_len_1, typenum::U1);
299    test_len!(test_break_line_len_2, typenum::U2);
300    test_len!(test_break_line_len_10, typenum::U10);
301    test_len!(test_break_line_len_74, typenum::U74);
302    test_len!(test_break_line_len_100, typenum::U100);
303    // test_len!(test_break_line_len_256, typenum::U256);
304
305    #[test]
306    fn test_key_roundtrip() {
307        let content = hex::decode("99010d04583c17af010800b9552ff9f4dae0e66c3ec0402793fbe2d02188f8ae6b1939b202bbb2fda892e2461f5098843eafc965809e350f464db24de4cf858afd5b870eb17847e5e05002fc0d14a37f5fbd448b247d95fbc953dc6c57b7c291631f0cbe8b6fa0a886d4346f827c11875aa26d7ebe86d25a84fc070d5894b85cf465e10f5a20b0ba830e10b9a24ecb845596d0b2fd3c07008ecc873733cb3add3b7030251ef8c061a9f46312eebf1adee68ae1865455c1f7a6a8d8ef6ed47b11edc523815429e89c062e02088a38f2d7aeccb3e7ca65eb6a03db73bb50b5480b4d622aa6014a2a186d581234bba00a6800a9870fe2c608c50f83977a6e1c3e3721e30c015624462fde41b70011010001b43554657374204b65792028646f206e6f742075736529203c61757468656e7469636174652d6f6e6c79406578616d706c652e6f72673e89014e041301080038162104927ef377fd1a1b6f795e40c02a87917d8ffba49f0502583c17af021b01050b09080702061508090a0b020416020301021e01021780000a09102a87917d8ffba49f7ecc07fd1556649c309608d638dbe448477e9fab69751acf0a7ccb17acd5ed7d83ed2ea8a83fc7f3d8b1342e8b9d4ac64c2d5ca0a273c6d190317485075dc15d52a3eb133b387c4c91e3169e392c6e8b643fbbafbd6e2ede8a5618cd53515a4bb2c764eb4506448cc5cc9ee25c5a9b466d15acfe2a5151904759f2e5dae74b97ed134482fb8678b4eb15421dbc04e6ccdc2e8cb3cef228c065400d716a786ee7b72eee44d64d003f9958c1f6274beca599544958bbcf55728330c8dd4e3648c26656a9d19880ac07740b23c36ce27c565cedb3f1f85e48572b3fe2c0718dd6e898272c2cefcaea20c675a67787f3af5881dc4d87732fecaf4720439c2dbac79046199396b9010d04583c17d9010800b4a71b058ac8aa1ddc453ab2663331c38f7645542815ac189a9af56d0e07a615469d3e08849650e03026d49259423cf00d089931cd700fd3a6e940bf83c81406e142a4b0a86f00738c7e1a9ff1b709f6bccc6cf900d0113a8e62e53d63be0a05105755b9efc6a4098c362c73fb422d40187d8e2382e88624d72caffceb13cec8fa0079c7d17883a46a1336471ab5be8cbb555c5d330d7fadb43318fa73b584edac312fa3302886bb5d04a05da3be2676c1fb94b3cf5c19d598659c3a7728ebab95f71721b662ac46aa9910726fe576d438f789c5ce2448f54546f254da814bcae1c35ee44b171e870ffa6403167a10e68573bdf155549274b431ff8e2418b6270011010001890136041801080020162104927ef377fd1a1b6f795e40c02a87917d8ffba49f0502583c17d9021b20000a09102a87917d8ffba49fac6b07ff7928f1c4082501da2517d94ad3bd2e566320ab81853ea27746a24f5058f010515260b5e48802e73065c34a639f3ea090e1cf8f5b0ba6161282cf83175430fcc7a9a2f59f87944c0831a0a7724ad24ee4393a0c0effabe4873e3639c80c6775909d67cd54cf236cd3fdbd7d6fae83de1cb15a3c6cbd28930cd19fe19422087b22bf6bdd335a54f950c7d1a2a35045b63f8a261d9a9bdfebd23d3e86c655ba3feda8594ee98f5b08e218eced3577aea39514680555c4c40160aa76c37c22976b07cd87d37d851233287ea14171e17973585ddb2f3b1a7d169eb8ad61e7e26ebc87229af539cd666d9f484ea62217a593a826fbccf6be43e19453e545f66e543c3f").unwrap();
308
309        // sanit check, ensure the base64 encoding works as expected
310        {
311            let mut buf = Vec::new();
312            {
313                let mut enc =
314                    base64::write::EncoderWriter::new(&mut buf, &general_purpose::STANDARD);
315                enc.write_all(&content).unwrap();
316            }
317
318            assert_eq!(
319                ::std::str::from_utf8(&buf).unwrap(),
320                "mQENBFg8F68BCAC5VS/59Nrg5mw+wEAnk/vi0CGI+K5rGTmyAruy/aiS4kYfUJiE\
321                 Pq/JZYCeNQ9GTbJN5M+Fiv1bhw6xeEfl4FAC/A0Uo39fvUSLJH2V+8lT3GxXt8KR\
322                 Yx8MvotvoKiG1DRvgnwRh1qibX6+htJahPwHDViUuFz0ZeEPWiCwuoMOELmiTsuE\
323                 VZbQsv08BwCOzIc3M8s63TtwMCUe+MBhqfRjEu6/Gt7miuGGVFXB96ao2O9u1HsR\
324                 7cUjgVQp6JwGLgIIijjy167Ms+fKZetqA9tzu1C1SAtNYiqmAUoqGG1YEjS7oApo\
325                 AKmHD+LGCMUPg5d6bhw+NyHjDAFWJEYv3kG3ABEBAAG0NVRlc3QgS2V5IChkbyBu\
326                 b3QgdXNlKSA8YXV0aGVudGljYXRlLW9ubHlAZXhhbXBsZS5vcmc+iQFOBBMBCAA4\
327                 FiEEkn7zd/0aG295XkDAKoeRfY/7pJ8FAlg8F68CGwEFCwkIBwIGFQgJCgsCBBYC\
328                 AwECHgECF4AACgkQKoeRfY/7pJ9+zAf9FVZknDCWCNY42+RIR36fq2l1Gs8KfMsX\
329                 rNXtfYPtLqioP8fz2LE0LoudSsZMLVygonPG0ZAxdIUHXcFdUqPrEzs4fEyR4xae\
330                 OSxui2Q/u6+9bi7eilYYzVNRWkuyx2TrRQZEjMXMnuJcWptGbRWs/ipRUZBHWfLl\
331                 2udLl+0TRIL7hni06xVCHbwE5szcLoyzzvIowGVADXFqeG7nty7uRNZNAD+ZWMH2\
332                 J0vspZlUSVi7z1VygzDI3U42SMJmVqnRmICsB3QLI8Ns4nxWXO2z8fheSFcrP+LA\
333                 cY3W6JgnLCzvyuogxnWmd4fzr1iB3E2Hcy/sr0cgQ5wtuseQRhmTlrkBDQRYPBfZ\
334                 AQgAtKcbBYrIqh3cRTqyZjMxw492RVQoFawYmpr1bQ4HphVGnT4IhJZQ4DAm1JJZ\
335                 QjzwDQiZMc1wD9Om6UC/g8gUBuFCpLCobwBzjH4an/G3Cfa8zGz5ANAROo5i5T1j\
336                 vgoFEFdVue/GpAmMNixz+0ItQBh9jiOC6IYk1yyv/OsTzsj6AHnH0XiDpGoTNkca\
337                 tb6Mu1VcXTMNf620Mxj6c7WE7awxL6MwKIa7XQSgXaO+JnbB+5Szz1wZ1ZhlnDp3\
338                 KOurlfcXIbZirEaqmRByb+V21Dj3icXOJEj1RUbyVNqBS8rhw17kSxcehw/6ZAMW\
339                 ehDmhXO98VVUknS0Mf+OJBi2JwARAQABiQE2BBgBCAAgFiEEkn7zd/0aG295XkDA\
340                 KoeRfY/7pJ8FAlg8F9kCGyAACgkQKoeRfY/7pJ+sawf/eSjxxAglAdolF9lK070u\
341                 VmMgq4GFPqJ3RqJPUFjwEFFSYLXkiALnMGXDSmOfPqCQ4c+PWwumFhKCz4MXVDD8\
342                 x6mi9Z+HlEwIMaCnckrSTuQ5OgwO/6vkhz42OcgMZ3WQnWfNVM8jbNP9vX1vroPe\
343                 HLFaPGy9KJMM0Z/hlCIIeyK/a90zWlT5UMfRoqNQRbY/iiYdmpvf69I9PobGVbo/\
344                 7ahZTumPWwjiGOztNXeuo5UUaAVVxMQBYKp2w3wil2sHzYfTfYUSMyh+oUFx4Xlz\
345                 WF3bLzsafRaeuK1h5+JuvIcimvU5zWZtn0hOpiIXpZOoJvvM9r5D4ZRT5UX2blQ8\
346                 Pw==",
347            );
348        }
349
350        let mut buf = Vec::new();
351        {
352            let mut line_wrapper = LineWriter::<_, typenum::U64>::new(&mut buf, LineBreak::Lf);
353            let mut enc =
354                base64::write::EncoderWriter::new(&mut line_wrapper, &general_purpose::STANDARD);
355            write_all(&mut enc, &content).unwrap();
356        }
357
358        assert_eq!(
359            ::std::str::from_utf8(&buf).unwrap(),
360            "mQENBFg8F68BCAC5VS/59Nrg5mw+wEAnk/vi0CGI+K5rGTmyAruy/aiS4kYfUJiE\n\
361             Pq/JZYCeNQ9GTbJN5M+Fiv1bhw6xeEfl4FAC/A0Uo39fvUSLJH2V+8lT3GxXt8KR\n\
362             Yx8MvotvoKiG1DRvgnwRh1qibX6+htJahPwHDViUuFz0ZeEPWiCwuoMOELmiTsuE\n\
363             VZbQsv08BwCOzIc3M8s63TtwMCUe+MBhqfRjEu6/Gt7miuGGVFXB96ao2O9u1HsR\n\
364             7cUjgVQp6JwGLgIIijjy167Ms+fKZetqA9tzu1C1SAtNYiqmAUoqGG1YEjS7oApo\n\
365             AKmHD+LGCMUPg5d6bhw+NyHjDAFWJEYv3kG3ABEBAAG0NVRlc3QgS2V5IChkbyBu\n\
366             b3QgdXNlKSA8YXV0aGVudGljYXRlLW9ubHlAZXhhbXBsZS5vcmc+iQFOBBMBCAA4\n\
367             FiEEkn7zd/0aG295XkDAKoeRfY/7pJ8FAlg8F68CGwEFCwkIBwIGFQgJCgsCBBYC\n\
368             AwECHgECF4AACgkQKoeRfY/7pJ9+zAf9FVZknDCWCNY42+RIR36fq2l1Gs8KfMsX\n\
369             rNXtfYPtLqioP8fz2LE0LoudSsZMLVygonPG0ZAxdIUHXcFdUqPrEzs4fEyR4xae\n\
370             OSxui2Q/u6+9bi7eilYYzVNRWkuyx2TrRQZEjMXMnuJcWptGbRWs/ipRUZBHWfLl\n\
371             2udLl+0TRIL7hni06xVCHbwE5szcLoyzzvIowGVADXFqeG7nty7uRNZNAD+ZWMH2\n\
372             J0vspZlUSVi7z1VygzDI3U42SMJmVqnRmICsB3QLI8Ns4nxWXO2z8fheSFcrP+LA\n\
373             cY3W6JgnLCzvyuogxnWmd4fzr1iB3E2Hcy/sr0cgQ5wtuseQRhmTlrkBDQRYPBfZ\n\
374             AQgAtKcbBYrIqh3cRTqyZjMxw492RVQoFawYmpr1bQ4HphVGnT4IhJZQ4DAm1JJZ\n\
375             QjzwDQiZMc1wD9Om6UC/g8gUBuFCpLCobwBzjH4an/G3Cfa8zGz5ANAROo5i5T1j\n\
376             vgoFEFdVue/GpAmMNixz+0ItQBh9jiOC6IYk1yyv/OsTzsj6AHnH0XiDpGoTNkca\n\
377             tb6Mu1VcXTMNf620Mxj6c7WE7awxL6MwKIa7XQSgXaO+JnbB+5Szz1wZ1ZhlnDp3\n\
378             KOurlfcXIbZirEaqmRByb+V21Dj3icXOJEj1RUbyVNqBS8rhw17kSxcehw/6ZAMW\n\
379             ehDmhXO98VVUknS0Mf+OJBi2JwARAQABiQE2BBgBCAAgFiEEkn7zd/0aG295XkDA\n\
380             KoeRfY/7pJ8FAlg8F9kCGyAACgkQKoeRfY/7pJ+sawf/eSjxxAglAdolF9lK070u\n\
381             VmMgq4GFPqJ3RqJPUFjwEFFSYLXkiALnMGXDSmOfPqCQ4c+PWwumFhKCz4MXVDD8\n\
382             x6mi9Z+HlEwIMaCnckrSTuQ5OgwO/6vkhz42OcgMZ3WQnWfNVM8jbNP9vX1vroPe\n\
383             HLFaPGy9KJMM0Z/hlCIIeyK/a90zWlT5UMfRoqNQRbY/iiYdmpvf69I9PobGVbo/\n\
384             7ahZTumPWwjiGOztNXeuo5UUaAVVxMQBYKp2w3wil2sHzYfTfYUSMyh+oUFx4Xlz\n\
385             WF3bLzsafRaeuK1h5+JuvIcimvU5zWZtn0hOpiIXpZOoJvvM9r5D4ZRT5UX2blQ8\n\
386             Pw==\n"
387        );
388    }
389}
pgp/line_writer.rs

pgp/
line_writer.rs
