diff options
| -rw-r--r-- | openpgp/src/armor.rs | 67 |
1 files changed, 60 insertions, 7 deletions
diff --git a/openpgp/src/armor.rs b/openpgp/src/armor.rs index f7e2db60..7ba91108 100644 --- a/openpgp/src/armor.rs +++ b/openpgp/src/armor.rs @@ -29,7 +29,7 @@ use buffered_reader::{ BufferedReader, BufferedReaderGeneric, BufferedReaderMemory, }; use std::fs::File; -use std::io::{Read, Write}; +use std::io::{Cursor, Read, Write}; use std::io::{Result, Error, ErrorKind}; use std::path::Path; use std::cmp::min; @@ -140,6 +140,8 @@ pub struct Writer<W: Write> { stash: Vec<u8>, column: usize, crc: CRC, + epilogue: Vec<u8>, + dirty: bool, finalized: bool, } @@ -181,21 +183,37 @@ impl<W: Write> Writer<W> { stash: Vec::<u8>::with_capacity(2), column: 0, crc: CRC::new(), + epilogue: Vec::with_capacity(128), + dirty: false, finalized: false, }; - write!(w.sink, "{}{}", w.kind.begin(), LINE_ENDING)?; + { + let mut cur = Cursor::new(&mut w.epilogue); + write!(&mut cur, "{}{}", kind.begin(), LINE_ENDING)?; - for h in headers { - write!(w.sink, "{}: {}{}", h.0, h.1, LINE_ENDING)?; - } + for h in headers { + write!(&mut cur, "{}: {}{}", h.0, h.1, LINE_ENDING)?; + } - // A blank line separates the headers from the body. - write!(w.sink, "{}", LINE_ENDING)?; + // A blank line separates the headers from the body. + write!(&mut cur, "{}", LINE_ENDING)?; + } Ok(w) } + fn write_epilogue(&mut self) -> Result<()> { + if ! self.dirty { + self.dirty = true; + self.sink.write_all(&self.epilogue)?; + // Release memory. + self.epilogue.clear(); + self.epilogue.shrink_to_fit(); + } + Ok(()) + } + /// Writes the footer. /// /// No more data can be written after this call. If this is not @@ -206,6 +224,12 @@ impl<W: Write> Writer<W> { return Err(Error::new(ErrorKind::BrokenPipe, "Writer is finalized.")); } + if ! self.dirty { + // No data was written to us, don't emit anything. + return Ok(()); + } + self.write_epilogue()?; + // Write any stashed bytes and pad. if self.stash.len() > 0 { self.sink.write_all(base64::encode_config(&self.stash, @@ -250,6 +274,8 @@ impl<W: Write> Write for Writer<W> { return Err(Error::new(ErrorKind::BrokenPipe, "Writer is finalized.")); } + self.write_epilogue()?; + // Update CRC on the unencoded data. self.crc.update(buf); @@ -1031,6 +1057,7 @@ mod test { let mut buf = Vec::new(); { let mut w = Writer::new(&mut buf, Kind::File, &[][..]).unwrap(); + w.write(&[]).unwrap(); // Avoid zero-length optimization. w.write_all(&bin).unwrap(); } assert_eq!(String::from_utf8_lossy(&buf), @@ -1052,6 +1079,7 @@ mod test { let mut buf = Vec::new(); { let mut w = Writer::new(&mut buf, Kind::File, &[][..]).unwrap(); + w.write(&[]).unwrap(); // Avoid zero-length optimization. for (i, _) in bin.iter().enumerate() { w.write(&bin[i..i+1]).unwrap(); } @@ -1061,6 +1089,31 @@ mod test { } } + #[test] + fn drop_writer() { + // No ASCII frame shall be emitted if the writer is dropped + // unused. + let mut buf = Vec::new(); + { + drop(Writer::new(&mut buf, Kind::File, &[][..]).unwrap()); + } + assert!(buf.is_empty()); + + // However, if the user insists, we will encode a zero-byte + // string. + let mut buf = Vec::new(); + { + let mut w = Writer::new(&mut buf, Kind::File, &[][..]).unwrap(); + w.write(&[]).unwrap(); + } + assert_eq!( + &buf[..], + &b"-----BEGIN PGP ARMORED FILE-----\n\ + \n\ + =twTO\n\ + -----END PGP ARMORED FILE-----\n"[..]); + } + use super::Reader; #[test] |