openpgp: Keep track of a writers position.

5 files changed, 62 insertions, 1 deletions

diff --git a/openpgp/src/serialize/partial_body.rs b/openpgp/src/serialize/partial_body.rs
index 398a553b..ab26e9d2 100644
--- a/openpgp/src/serialize/partial_body.rs
+++ b/openpgp/src/serialize/partial_body.rs
@@ -30,6 +30,9 @@ pub struct PartialBodyFilter<'a, C: 'a> {
     // The maximum size of a partial body chunk.  The standard allows
     // for chunks up to 1 GB in size.
     max_chunk_size: usize,
+
+    // The number of bytes written to this filter.
+    position: u64,
 }
 
 const PARTIAL_BODY_FILTER_MAX_CHUNK_SIZE : usize = 1 << 30;
@@ -74,6 +77,7 @@ impl<'a, C: 'a> PartialBodyFilter<'a, C> {
             buffer: Vec::with_capacity(buffer_threshold),
             buffer_threshold: buffer_threshold,
             max_chunk_size: max_chunk_size,
+            position: 0,
         })))
     }
 
@@ -159,6 +163,7 @@ impl<'a, C: 'a> io::Write for PartialBodyFilter<'a, C> {
         } else {
             self.buffer.append(buf.to_vec().as_mut());
         }
+        self.position += buf.len() as u64;
         Ok(buf.len())
     }
 
@@ -219,6 +224,9 @@ impl<'a, C: 'a> writer::Stackable<'a, C> for PartialBodyFilter<'a, C> {
     fn cookie_mut(&mut self) -> &mut C {
         &mut self.cookie
     }
+    fn position(&self) -> u64 {
+        self.position
+    }
 }
 
 #[cfg(test)]
diff --git a/openpgp/src/serialize/stream.rs b/openpgp/src/serialize/stream.rs
index 7afa925e..7d3eee84 100644
--- a/openpgp/src/serialize/stream.rs
+++ b/openpgp/src/serialize/stream.rs
@@ -177,6 +177,9 @@ impl<'a> writer::Stackable<'a, Cookie> for ArbitraryWriter<'a> {
     fn cookie_mut(&mut self) -> &mut Cookie {
         self.inner.cookie_mut()
     }
+    fn position(&self) -> u64 {
+        self.inner.position()
+    }
 }
 
 /// Signs a packet stream.
@@ -204,6 +207,7 @@ pub struct Signer<'a, R>
     detached: bool,
     hash: crypto::hash::Context,
     cookie: Cookie,
+    position: u64,
 }
 
 impl<'a, R> Signer<'a, R>
@@ -405,6 +409,7 @@ impl<'a, R> Signer<'a, R>
                 level: level,
                 private: Private::Signer,
             },
+            position: 0,
         })))
     }
 
@@ -478,6 +483,7 @@ impl<'a, R> Write for Signer<'a, R>
 
         if let Ok(amount) = written {
             self.hash.update(&buf[..amount]);
+            self.position += amount as u64;
         }
 
         written
@@ -530,6 +536,9 @@ impl<'a, R> writer::Stackable<'a, Cookie> for Signer<'a, R>
     fn cookie_mut(&mut self) -> &mut Cookie {
         &mut self.cookie
     }
+    fn position(&self) -> u64 {
+        self.position
+    }
 }
 
 
@@ -696,6 +705,9 @@ impl<'a> writer::Stackable<'a, Cookie> for LiteralWriter<'a> {
     fn cookie_mut(&mut self) -> &mut Cookie {
         self.inner.cookie_mut()
     }
+    fn position(&self) -> u64 {
+        self.inner.position()
+    }
 }
 
 /// Compresses a packet stream.
@@ -822,6 +834,9 @@ impl<'a> writer::Stackable<'a, Cookie> for Compressor<'a> {
     fn cookie_mut(&mut self) -> &mut Cookie {
         self.inner.cookie_mut()
     }
+    fn position(&self) -> u64 {
+        self.inner.position()
+    }
 }
 
 /// Encrypts a packet stream.
@@ -1184,6 +1199,9 @@ impl<'a> writer::Stackable<'a, Cookie> for Encryptor<'a> {
     fn cookie_mut(&mut self) -> &mut Cookie {
         &mut self.cookie
     }
+    fn position(&self) -> u64 {
+        self.inner.as_ref().map(|i| i.position()).unwrap_or(0)
+    }
 }
 
 #[cfg(test)]
diff --git a/openpgp/src/serialize/writer/mod.rs b/openpgp/src/serialize/writer/mod.rs
index a982d8ac..4545cdd2 100644
--- a/openpgp/src/serialize/writer/mod.rs
+++ b/openpgp/src/serialize/writer/mod.rs
@@ -111,6 +111,9 @@ pub(crate) trait Stackable<'a, C> : io::Write + fmt::Debug {
     /// Returns a mutable reference to the cookie.
     fn cookie_mut(&mut self) -> &mut C;
 
+    /// Returns the number of bytes written to this filter.
+    fn position(&self) -> u64;
+
     /// Writes a byte.
     fn write_u8(&mut self, b: u8) -> io::Result<()> {
         let b : [u8; 1] = [b; 1];
@@ -159,6 +162,9 @@ impl <'a, C> Stackable<'a, C> for BoxStack<'a, C> {
     fn cookie_mut(&mut self) -> &mut C {
         self.as_mut().cookie_mut()
     }
+    fn position(&self) -> u64 {
+        self.as_ref().position()
+    }
 }
 
 /// Maps a function over the stack of writers.
@@ -272,12 +278,16 @@ impl<'a, C> Stackable<'a, C> for Identity<'a, C> {
     fn cookie_mut(&mut self) -> &mut C {
         &mut self.cookie
     }
+    fn position(&self) -> u64 {
+        self.inner.as_ref().map(|i| i.position()).unwrap_or(0)
+    }
 }
 
 /// Generic writer wrapping `io::Write`.
 pub struct Generic<W: io::Write, C> {
     inner: W,
     cookie: C,
+    position: u64,
 }
 
 impl<'a, W: 'a + io::Write, C: 'a> Generic<W, C> {
@@ -290,6 +300,7 @@ impl<'a, W: 'a + io::Write, C: 'a> Generic<W, C> {
         Generic {
             inner: inner,
             cookie: cookie,
+            position: 0,
         }
     }
 }
@@ -303,7 +314,13 @@ impl<W: io::Write, C> fmt::Debug for Generic<W, C> {
 
 impl<W: io::Write, C> io::Write for Generic<W, C> {
     fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
-        self.inner.write(bytes)
+        match self.inner.write(bytes) {
+            Ok(n) => {
+                self.position += n as u64;
+                Ok(n)
+            },
+            Err(e) => Err(e),
+        }
     }
 
     fn flush(&mut self) -> io::Result<()> {
@@ -338,6 +355,9 @@ impl<'a, W: io::Write, C> Stackable<'a, C> for Generic<W, C> {
     fn cookie_mut(&mut self) -> &mut C {
         &mut self.cookie
     }
+    fn position(&self) -> u64 {
+        self.position
+    }
 }
 
 
@@ -404,6 +424,9 @@ impl<'a, C: 'a> Stackable<'a, C> for Encryptor<'a, C> {
     fn cookie_mut(&mut self) -> &mut C {
         self.inner.cookie_mut()
     }
+    fn position(&self) -> u64 {
+        self.inner.position
+    }
 }
 
 
@@ -472,6 +495,9 @@ impl<'a, C: 'a> Stackable<'a, C> for AEADEncryptor<'a, C> {
     fn cookie_mut(&mut self) -> &mut C {
         self.inner.cookie_mut()
     }
+    fn position(&self) -> u64 {
+        self.inner.position
+    }
 }
 
 #[cfg(test)]
diff --git a/openpgp/src/serialize/writer/writer_bzip2.rs b/openpgp/src/serialize/writer/writer_bzip2.rs
index 8f51fe18..b951b596 100644
--- a/openpgp/src/serialize/writer/writer_bzip2.rs
+++ b/openpgp/src/serialize/writer/writer_bzip2.rs
@@ -66,4 +66,7 @@ impl<'a, C: 'a> Stackable<'a, C> for BZ<'a, C> {
     fn cookie_mut(&mut self) -> &mut C {
         self.inner.cookie_mut()
     }
+    fn position(&self) -> u64 {
+        self.inner.position
+    }
 }
diff --git a/openpgp/src/serialize/writer/writer_deflate.rs b/openpgp/src/serialize/writer/writer_deflate.rs
index a0e1ea60..664b5319 100644
--- a/openpgp/src/serialize/writer/writer_deflate.rs
+++ b/openpgp/src/serialize/writer/writer_deflate.rs
@@ -66,6 +66,9 @@ impl<'a, C: 'a> Stackable<'a, C> for ZIP<'a, C> {
     fn cookie_mut(&mut self) -> &mut C {
         self.inner.cookie_mut()
     }
+    fn position(&self) -> u64 {
+        self.inner.position
+    }
 }
 
 /// ZLIBing writer.
@@ -128,4 +131,7 @@ impl<'a, C: 'a> Stackable<'a, C> for ZLIB<'a, C> {
     fn cookie_mut(&mut self) -> &mut C {
         self.inner.cookie_mut()
     }
+    fn position(&self) -> u64 {
+        self.inner.position
+    }
 }