openpgp: Support setting ASCII-armor's headers

8 files changed, 49 insertions, 32 deletions

diff --git a/ffi/src/openpgp.rs b/ffi/src/openpgp.rs
index 330cb500..cd588ee2 100644
--- a/ffi/src/openpgp.rs
+++ b/ffi/src/openpgp.rs
@@ -235,13 +235,17 @@ pub extern "system" fn sq_armor_reader_new(inner: Option<&'static mut Box<Read>>
 ///
 /// A filter that applies ASCII Armor to the data written to it.
 #[no_mangle]
-pub extern "system" fn sq_armor_writer_new(inner: Option<&'static mut Box<Write>>,
+pub extern "system" fn sq_armor_writer_new(ctx: Option<&mut Context>,
+                                           inner: Option<&'static mut Box<Write>>,
                                            kind: c_int)
-                                           -> *mut Box<Write> {
+                                           -> *mut armor::Writer<
+                                                  &'static mut Box<Write>> {
+    let ctx = ctx.expect("Context is NULL");
     let inner = inner.expect("Inner is NULL");
     let kind = int_to_kind(kind);
 
-    box_raw!(Box::new(armor::Writer::new(inner, kind)))
+    // XXX: Expose header parameter.
+    fry_box!(ctx, armor::Writer::new(inner, kind, &[][..]).map_err(|e| e.into()))
 }
 
 
diff --git a/net/src/async.rs b/net/src/async.rs
index 65e6cfed..7af6f2ae 100644
--- a/net/src/async.rs
+++ b/net/src/async.rs
@@ -155,7 +155,12 @@ impl KeyServer {
 
         let mut armored_blob = vec![];
         {
-            let mut w = Writer::new(&mut armored_blob, Kind::PublicKey);
+            let mut w = match Writer::new(&mut armored_blob,
+                                          Kind::PublicKey, &[][..]) {
+                Err(e) => return Box::new(future::err(e.into())),
+                Ok(w) => w,
+            };
+
             if let Err(e) = key.serialize(&mut w) {
                 return Box::new(future::err(e));
             }
diff --git a/openpgp/examples/encrypt-for.rs b/openpgp/examples/encrypt-for.rs
index 9b55e142..f8a04c3b 100644
--- a/openpgp/examples/encrypt-for.rs
+++ b/openpgp/examples/encrypt-for.rs
@@ -40,7 +40,8 @@ fn main() {
     // Compose a writer stack corresponding to the output format and
     // packet structure we want.  First, we want the output to be as
     // armored.
-    let sink = armor::Writer::new(io::stdout(), armor::Kind::Message);
+    let sink = armor::Writer::new(io::stdout(), armor::Kind::Message, &[][..])
+        .expect("Failed to create an armored writer");
 
     // We want to encrypt a literal data packet.
     let encryptor = Encryptor::new(wrap(sink),
diff --git a/openpgp/examples/sign-detached.rs b/openpgp/examples/sign-detached.rs
index 7b49961b..bb9985ec 100644
--- a/openpgp/examples/sign-detached.rs
+++ b/openpgp/examples/sign-detached.rs
@@ -28,7 +28,8 @@ fn main() {
     // Compose a writer stack corresponding to the output format and
     // packet structure we want.  First, we want the output to be as
     // armored.
-    let sink = armor::Writer::new(io::stdout(), armor::Kind::Signature);
+    let sink = armor::Writer::new(io::stdout(), armor::Kind::Signature, &[][..])
+        .expect("Failed to create armored writer.");
 
     // Now, create a signer that emits a detached signature.
     let mut signer = Signer::detached(
diff --git a/openpgp/examples/sign.rs b/openpgp/examples/sign.rs
index 084f3e06..2caeefb0 100644
--- a/openpgp/examples/sign.rs
+++ b/openpgp/examples/sign.rs
@@ -29,7 +29,8 @@ fn main() {
     // Compose a writer stack corresponding to the output format and
     // packet structure we want.  First, we want the output to be as
     // armored.
-    let sink = armor::Writer::new(io::stdout(), armor::Kind::Message);
+    let sink = armor::Writer::new(io::stdout(), armor::Kind::Message, &[][..])
+        .expect("Failed to create an armored writer.");
 
     // Now, create a signer that emits a detached signature.
     let signer = Signer::new(
diff --git a/openpgp/src/armor.rs b/openpgp/src/armor.rs
index 30a8a344..2ef43a73 100644
--- a/openpgp/src/armor.rs
+++ b/openpgp/src/armor.rs
@@ -149,7 +149,6 @@ pub struct Writer<W: Write> {
     stash: Vec<u8>,
     column: usize,
     crc: CRC,
-    initialized: bool,
     finalized: bool,
 }
 
@@ -167,13 +166,15 @@ impl<W: Write> Writer<W> {
     /// # fn f() -> Result<()> {
     /// let mut buffer = io::Cursor::new(vec![]);
     /// {
-    ///     let mut writer = Writer::new(&mut buffer, Kind::File);
+    ///     let mut writer = Writer::new(&mut buffer, Kind::File,
+    ///         &[ ("Key", "Value") ][..])?;
     ///     writer.write_all(b"Hello world!")?;
     ///     // writer is drop()ed here.
     /// }
     /// assert_eq!(
     ///     String::from_utf8_lossy(buffer.get_ref()),
     ///     "-----BEGIN PGP ARMORED FILE-----
+    /// Key: Value
     ///
     /// SGVsbG8gd29ybGQh
     /// =s4Gu
@@ -182,28 +183,27 @@ impl<W: Write> Writer<W> {
     /// # Ok(())
     /// # }
     /// ```
-    pub fn new(inner: W, kind: Kind) -> Self {
+    pub fn new(inner: W, kind: Kind, headers: &[(&str, &str)]) -> Result<Self> {
         assert!(kind != Kind::Any);
-        Writer {
+        let mut w = Writer {
             sink: inner,
             kind: kind,
             stash: Vec::<u8>::with_capacity(2),
             column: 0,
             crc: CRC::new(),
-            initialized: false,
             finalized: false,
-        }
-    }
+        };
 
-    /// Writes the header if not already done.
-    fn initialize(&mut self) -> Result<()> {
-        if self.initialized { return Ok(()) }
+        write!(w.sink, "{}{}", w.kind.begin(), LINE_ENDING)?;
 
-        write!(self.sink, "{}{}{}", self.kind.begin(),
-               LINE_ENDING, LINE_ENDING)?;
+        for h in headers {
+            write!(w.sink, "{}: {}{}", h.0, h.1, LINE_ENDING)?;
+        }
 
-        self.initialized = true;
-        Ok(())
+        // A blank line separates the headers from the body.
+        write!(w.sink, "{}", LINE_ENDING)?;
+
+        Ok(w)
     }
 
     /// Writes the footer.
@@ -212,7 +212,6 @@ impl<W: Write> Writer<W> {
     /// called explicitly, the header is written once the writer is
     /// dropped.
     pub fn finalize(&mut self) -> Result<()> {
-        self.initialize()?;
         if self.finalized {
             return Err(Error::new(ErrorKind::BrokenPipe, "Writer is finalized."));
         }
@@ -257,7 +256,6 @@ impl<W: Write> Writer<W> {
 
 impl<W: Write> Write for Writer<W> {
     fn write(&mut self, buf: &[u8]) -> Result<usize> {
-        self.initialize()?;
         if self.finalized {
             return Err(Error::new(ErrorKind::BrokenPipe, "Writer is finalized."));
         }
@@ -848,7 +846,7 @@ mod test {
 
             let mut buf = Vec::new();
             {
-                let mut w = Writer::new(&mut buf, Kind::File);
+                let mut w = Writer::new(&mut buf, Kind::File, &[][..]).unwrap();
                 w.write_all(&bin).unwrap();
             }
             assert_eq!(String::from_utf8_lossy(&buf),
@@ -869,7 +867,7 @@ mod test {
 
             let mut buf = Vec::new();
             {
-                let mut w = Writer::new(&mut buf, Kind::File);
+                let mut w = Writer::new(&mut buf, Kind::File, &[][..]).unwrap();
                 for (i, _) in bin.iter().enumerate() {
                     w.write(&bin[i..i+1]).unwrap();
                 }
@@ -1042,7 +1040,7 @@ mod test {
             use std::io::Cursor;
 
             let mut encoded = Vec::new();
-            Writer::new(&mut encoded, kind)
+            Writer::new(&mut encoded, kind, &[][..]).unwrap()
                 .write_all(&payload)
                 .unwrap();
 
diff --git a/openpgp/src/tpk/mod.rs b/openpgp/src/tpk/mod.rs
index 45258d65..fe8803ab 100644
--- a/openpgp/src/tpk/mod.rs
+++ b/openpgp/src/tpk/mod.rs
@@ -2038,7 +2038,8 @@ mod test {
         let mut cur = io::Cursor::new(Vec::default());
 
         {
-            let mut a = armor::Writer::new(&mut cur, armor::Kind::SecretKey);
+            let mut a = armor::Writer::new(&mut cur, armor::Kind::SecretKey,
+                                           &[][..]).unwrap();
             t1.serialize(&mut a).unwrap();
         }
 
diff --git a/tool/src/sq.rs b/tool/src/sq.rs
index 2f70dc90..b8969065 100644
--- a/tool/src/sq.rs
+++ b/tool/src/sq.rs
@@ -75,7 +75,8 @@ fn real_main() -> Result<(), failure::Error> {
             let mut output = create_or_stdout(m.value_of("output"))?;
             let mut output = if ! m.is_present("binary") {
                 Box::new(armor::Writer::new(&mut output,
-                                            armor::Kind::Message))
+                                            armor::Kind::Message,
+                                            &[][..])?)
             } else {
                 output
             };
@@ -92,7 +93,8 @@ fn real_main() -> Result<(), failure::Error> {
         ("enarmor",  Some(m)) => {
             let mut input = open_or_stdin(m.value_of("input"))?;
             let mut output = create_or_stdout(m.value_of("output"))?;
-            let mut filter = armor::Writer::new(&mut output, armor::Kind::File);
+            let mut filter = armor::Writer::new(&mut output, armor::Kind::File,
+                                                &[][..])?;
             io::copy(&mut input, &mut filter)?;
         },
         ("dearmor",  Some(m)) => {
@@ -110,7 +112,7 @@ fn real_main() -> Result<(), failure::Error> {
                     for h in &ac.headers {
                         if let Some(ref tpk) = h.key {
                             let mut filter = armor::Writer::new(
-                                &mut output, armor::Kind::PublicKey);
+                                &mut output, armor::Kind::PublicKey, &[][..])?;
                             tpk.serialize(&mut filter)?;
                         }
                     }
@@ -157,7 +159,9 @@ fn real_main() -> Result<(), failure::Error> {
 
                     let mut output = create_or_stdout(m.value_of("output"))?;
                     let mut output = if ! m.is_present("binary") {
-                        Box::new(armor::Writer::new(&mut output, armor::Kind::PublicKey))
+                        Box::new(armor::Writer::new(&mut output,
+                                                    armor::Kind::PublicKey,
+                                                    &[][..])?)
                     } else {
                         output
                     };
@@ -208,7 +212,9 @@ fn real_main() -> Result<(), failure::Error> {
 
                     let mut output = create_or_stdout(m.value_of("output"))?;
                     let mut output = if ! m.is_present("binary") {
-                        Box::new(armor::Writer::new(&mut output, armor::Kind::PublicKey))
+                        Box::new(armor::Writer::new(&mut output,
+                                                    armor::Kind::PublicKey,
+                                                    &[][..])?)
                     } else {
                         output
                     };