openpgp: Improve use of armor writer in tests and examples.

3 files changed, 56 insertions, 63 deletions

diff --git a/net/src/lib.rs b/net/src/lib.rs
index 95216126..2619069b 100644
--- a/net/src/lib.rs
+++ b/net/src/lib.rs
@@ -217,19 +217,21 @@ impl KeyServer {
                 Ok(u) => u,
             };
 
-        let mut armored_blob = vec![];
-        {
-            let mut w = match Writer::new(&mut armored_blob,
-                                          Kind::PublicKey, &[]) {
-                Err(e) => return Box::new(future::err(e.into())),
-                Ok(w) => w,
-            };
+        let mut w = match Writer::new(Vec::new(),
+                                      Kind::PublicKey, &[]) {
+            Ok(v) => v,
+            Err(e) => return Box::new(future::err(e.into())),
+        };
 
-            if let Err(e) = key.serialize(&mut w) {
-                return Box::new(future::err(e));
-            }
+        if let Err(e) = key.serialize(&mut w) {
+            return Box::new(future::err(e));
         }
 
+        let armored_blob = match w.finalize() {
+            Ok(v) => v,
+            Err(e) => return Box::new(future::err(e.into())),
+        };
+
         // Prepare to send url-encoded data.
         let mut post_data = b"keytext=".to_vec();
         post_data.extend_from_slice(percent_encode(&armored_blob, KEYSERVER_ENCODE_SET)
diff --git a/openpgp/src/armor.rs b/openpgp/src/armor.rs
index d690bfc5..7e26f198 100644
--- a/openpgp/src/armor.rs
+++ b/openpgp/src/armor.rs
@@ -155,21 +155,17 @@ impl<W: Write> Writer<W> {
     /// # Example
     ///
     /// ```
-    /// # use std::io::Write;
-    /// # extern crate sequoia_openpgp as openpgp;
-    /// # use openpgp::armor::{Writer, Kind};
-    /// # use std::io::{self, Result};
+    /// use std::io::{Read, Write, Cursor};
+    /// use sequoia_openpgp as openpgp;
+    /// use openpgp::armor::{Writer, Kind};
     /// # fn main() { f().unwrap(); }
-    /// # fn f() -> Result<()> {
-    /// let mut buffer = io::Cursor::new(vec![]);
-    /// {
-    ///     let mut writer = Writer::new(&mut buffer, Kind::File,
-    ///         &[ ("Key", "Value") ][..])?;
-    ///     writer.write_all(b"Hello world!")?;
-    ///     // writer is drop()ed here.
-    /// }
+    /// # fn f() -> std::io::Result<()> {
+    /// let mut writer = Writer::new(Vec::new(), Kind::File,
+    ///     &[ ("Key", "Value") ][..])?;
+    /// writer.write_all(b"Hello world!")?;
+    /// let buffer = writer.finalize()?;
     /// assert_eq!(
-    ///     String::from_utf8_lossy(buffer.get_ref()),
+    ///     String::from_utf8_lossy(&buffer),
     ///     "-----BEGIN PGP ARMORED FILE-----
     /// Key: Value
     ///
@@ -1317,12 +1313,11 @@ mod test {
     #[test]
     fn enarmor() {
         for (bin, asc) in TEST_BIN.iter().zip(TEST_ASC.iter()) {
-            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();
-            }
+            let mut w =
+                Writer::new(Vec::new(), Kind::File, &[]).unwrap();
+            w.write(&[]).unwrap();  // Avoid zero-length optimization.
+            w.write_all(bin).unwrap();
+            let buf = w.finalize().unwrap();
             assert_eq!(String::from_utf8_lossy(&buf),
                        String::from_utf8_lossy(asc));
         }
@@ -1331,14 +1326,12 @@ mod test {
     #[test]
     fn enarmor_bytewise() {
         for (bin, asc) in TEST_BIN.iter().zip(TEST_ASC.iter()) {
-            let mut buf = Vec::new();
-            {
-                let mut w = Writer::new(&mut buf, Kind::File, &[]).unwrap();
-                w.write(&[]).unwrap();  // Avoid zero-length optimization.
-                for b in bin.iter() {
-                    w.write(&[*b]).unwrap();
-                }
+            let mut w = Writer::new(Vec::new(), Kind::File, &[]).unwrap();
+            w.write(&[]).unwrap();  // Avoid zero-length optimization.
+            for b in bin.iter() {
+                w.write(&[*b]).unwrap();
             }
+            let buf = w.finalize().unwrap();
             assert_eq!(String::from_utf8_lossy(&buf),
                        String::from_utf8_lossy(asc));
         }
@@ -1348,19 +1341,14 @@ mod 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());
+        assert!(Writer::new(Vec::new(), Kind::File, &[]).unwrap()
+                .finalize().unwrap().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();
-        }
+        let mut w = Writer::new(Vec::new(), Kind::File, &[]).unwrap();
+        w.write(&[]).unwrap();
+        let buf = w.finalize().unwrap();
         assert_eq!(
             &buf[..],
             &b"-----BEGIN PGP ARMORED FILE-----\n\
@@ -1599,10 +1587,9 @@ mod test {
                 return true;
             }
 
-            let mut encoded = Vec::new();
-            Writer::new(&mut encoded, kind, &[]).unwrap()
-                .write_all(&payload)
-                .unwrap();
+            let mut w = Writer::new(Vec::new(), kind, &[]).unwrap();
+            w.write_all(&payload).unwrap();
+            let encoded = w.finalize().unwrap();
 
             let mut recovered = Vec::new();
             Reader::new(Cursor::new(&encoded),
diff --git a/openpgp/src/autocrypt.rs b/openpgp/src/autocrypt.rs
index cb579211..c99ee4e1 100644
--- a/openpgp/src/autocrypt.rs
+++ b/openpgp/src/autocrypt.rs
@@ -467,23 +467,27 @@ impl AutocryptSetupMessage {
                 (&"Passphrase-Begin"[..], &begin[..]));
         }
 
-        let w = armor::Writer::new(w, armor::Kind::Message, &headers[..])?;
+        let mut armor_writer =
+            armor::Writer::new(w, armor::Kind::Message, &headers[..])?;
 
-        // Passphrase-Format header with value numeric9x4
-        let m = Message::new(w);
-        let w = Encryptor::with_password(m, self.passcode.clone().unwrap())
-            .build()?;
+        {
+            // Passphrase-Format header with value numeric9x4
+            let m = Message::new(&mut armor_writer);
+            let w = Encryptor::with_password(m, self.passcode.clone().unwrap())
+                .build()?;
 
-        let mut w = LiteralWriter::new(w).build()?;
+            let mut w = LiteralWriter::new(w).build()?;
 
-        // The inner message is an ASCII-armored encoded Cert.
-        let mut w = armor::Writer::new(
-            &mut w, armor::Kind::SecretKey,
-            &[ (&"Autocrypt-Prefer-Encrypt"[..],
-                self.prefer_encrypt().unwrap_or(&"nopreference"[..])) ])?;
+            // The inner message is an ASCII-armored encoded Cert.
+            let mut w = armor::Writer::new(
+                &mut w, armor::Kind::SecretKey,
+                &[ (&"Autocrypt-Prefer-Encrypt"[..],
+                    self.prefer_encrypt().unwrap_or(&"nopreference"[..])) ])?;
 
-        self.cert.as_tsk().serialize(&mut w)?;
-        w.finalize()?;
+            self.cert.as_tsk().serialize(&mut w)?;
+            w.finalize()?;
+        }
+        armor_writer.finalize()?;
         Ok(())
     }