From 9136636f4508f76d7a92169d05d29ce18caca24a Mon Sep 17 00:00:00 2001
From: joe <55120843+joebebel@users.noreply.github.com>
Date: Fri, 31 May 2024 05:21:01 -0700
Subject: [PATCH] digest: add HashWriter and HashReader structs (#1159)

---
 digest/src/hashreader.rs | 135 +++++++++++++++++++++++++++++++++++++++
 digest/src/hashwriter.rs | 120 ++++++++++++++++++++++++++++++++++
 digest/src/lib.rs        |   9 +++
 3 files changed, 264 insertions(+)
 create mode 100644 digest/src/hashreader.rs
 create mode 100644 digest/src/hashwriter.rs

diff --git a/digest/src/hashreader.rs b/digest/src/hashreader.rs
new file mode 100644
index 000000000..8c303af2a
--- /dev/null
+++ b/digest/src/hashreader.rs
@@ -0,0 +1,135 @@
+//! Adds hashing to any reader
+use super::{Digest, FixedOutputReset, Output, Reset};
+use std::io;
+
+/// Abstraction over a reader which hashes the data being read
+pub struct HashReader<D: Digest, R: io::Read> {
+    reader: R,
+    hasher: D,
+}
+
+impl<D: Digest, R: io::Read> HashReader<D, R> {
+    /// Construct a new `HashReader` given an existing `reader` by value.
+    pub fn new(reader: R) -> Self {
+        Self::new_from_parts(D::new(), reader)
+    }
+
+    /// Construct a new `HashReader` given an existing `hasher` and `reader` by value.
+    pub fn new_from_parts(hasher: D, reader: R) -> Self {
+        HashReader { reader, hasher }
+    }
+
+    /// Replace the reader with another reader
+    pub fn replace_reader(&mut self, reader: R) {
+        self.reader = reader;
+    }
+
+    /// Gets a reference to the underlying hasher
+    pub fn get_hasher(&self) -> &D {
+        &self.hasher
+    }
+
+    /// Gets a reference to the underlying reader
+    pub fn get_reader(&self) -> &R {
+        &self.reader
+    }
+
+    /// Gets a mutable reference to the underlying hasher
+    pub fn get_hasher_mut(&mut self) -> &mut D {
+        &mut self.hasher
+    }
+
+    /// Gets a mutable reference to the underlying reader
+    /// Direct reads from the underlying reader are not hashed
+    pub fn get_reader_mut(&mut self) -> &mut R {
+        &mut self.reader
+    }
+
+    /// Consume the HashReader and return its hasher
+    pub fn into_hasher(self) -> D {
+        self.hasher
+    }
+
+    /// Consume the HashReader and return its internal reader
+    pub fn into_inner_reader(self) -> R {
+        self.reader
+    }
+
+    /// Consume the HashReader and return its hasher and internal reader
+    pub fn into_parts(self) -> (D, R) {
+        (self.hasher, self.reader)
+    }
+
+    /// Retrieve result and consume HashReader instance.
+    pub fn finalize(self) -> Output<D> {
+        self.hasher.finalize()
+    }
+
+    /// Write result into provided array and consume the HashReader instance.
+    pub fn finalize_into(self, out: &mut Output<D>) {
+        self.hasher.finalize_into(out)
+    }
+
+    /// Get output size of the hasher
+    pub fn output_size() -> usize {
+        <D as Digest>::output_size()
+    }
+}
+
+impl<D: Digest + Clone, R: io::Read + Clone> Clone for HashReader<D, R> {
+    fn clone(&self) -> HashReader<D, R> {
+        HashReader {
+            reader: self.reader.clone(),
+            hasher: self.hasher.clone(),
+        }
+    }
+}
+
+impl<D: Digest, R: io::Read> io::Read for HashReader<D, R> {
+    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
+        let bytes = self.reader.read(buf)?;
+
+        if bytes > 0 {
+            self.hasher.update(&buf[0..bytes]);
+        }
+
+        Ok(bytes)
+    }
+}
+
+impl<D: Digest + FixedOutputReset, R: io::Read> HashReader<D, R> {
+    /// Retrieve result and reset hasher instance.
+    pub fn finalize_reset(&mut self) -> Output<D> {
+        Digest::finalize_reset(&mut self.hasher)
+    }
+
+    /// Rrite result into provided array and reset the hasher instance.
+    pub fn finalize_into_reset(&mut self, out: &mut Output<D>) {
+        Digest::finalize_into_reset(&mut self.hasher, out)
+    }
+}
+impl<D: Digest + Reset, R: io::Read> Reset for HashReader<D, R> {
+    fn reset(&mut self) {
+        Digest::reset(&mut self.hasher)
+    }
+}
+
+impl<D: Digest, R: io::BufRead> HashReader<D, R> {
+    /// Read and hash all bytes remaining in the reader, discarding the data
+    /// Based on implementation in b2sum crate, MIT License Copyright (c) 2017 John Downey
+    pub fn hash_to_end(&mut self) {
+        loop {
+            let count = {
+                let data = self.reader.fill_buf().unwrap();
+                if data.is_empty() {
+                    break;
+                }
+
+                self.hasher.update(data);
+                data.len()
+            };
+
+            self.reader.consume(count);
+        }
+    }
+}
diff --git a/digest/src/hashwriter.rs b/digest/src/hashwriter.rs
new file mode 100644
index 000000000..96ebe95a6
--- /dev/null
+++ b/digest/src/hashwriter.rs
@@ -0,0 +1,120 @@
+//! Adds hashing to any writer. Inspired by implemention in phase2 crate.
+use super::{Digest, FixedOutputReset, Output, Reset};
+use std::io;
+
+/// Abstraction over a writer which hashes the data being written.
+pub struct HashWriter<D: Digest, W: io::Write> {
+    writer: W,
+    hasher: D,
+}
+
+impl<D: Digest, W: io::Write> HashWriter<D, W> {
+    /// Construct a new `HashWriter` given an existing `writer` by value.
+    pub fn new(writer: W) -> Self {
+        Self::new_from_parts(D::new(), writer)
+    }
+
+    /// Construct a new `HashWriter` given an existing `hasher` and `writer` by value.
+    pub fn new_from_parts(hasher: D, writer: W) -> Self {
+        HashWriter { writer, hasher }
+    }
+
+    /// Replace the writer with another writer
+    pub fn replace_writer(&mut self, writer: W) {
+        self.writer = writer;
+    }
+
+    /// Gets a reference to the underlying hasher
+    pub fn get_hasher(&self) -> &D {
+        &self.hasher
+    }
+
+    /// Gets a reference to the underlying writer
+    pub fn get_writer(&self) -> &W {
+        &self.writer
+    }
+
+    /// Gets a mutable reference to the underlying hasher
+    /// Updates to the digest are not written to the underlying writer
+    pub fn get_hasher_mut(&mut self) -> &mut D {
+        &mut self.hasher
+    }
+
+    /// Gets a mutable reference to the underlying writer
+    /// Direct writes to the underlying writer are not hashed
+    pub fn get_writer_mut(&mut self) -> &mut W {
+        &mut self.writer
+    }
+
+    /// Consume the HashWriter and return its hasher
+    pub fn into_hasher(self) -> D {
+        self.hasher
+    }
+
+    /// Consume the HashWriter and return its internal writer
+    pub fn into_inner_writer(self) -> W {
+        self.writer
+    }
+
+    /// Consume the HashWriter and return its hasher and internal writer
+    pub fn into_parts(self) -> (D, W) {
+        (self.hasher, self.writer)
+    }
+
+    /// Retrieve result and consume HashWriter instance.
+    pub fn finalize(self) -> Output<D> {
+        self.hasher.finalize()
+    }
+
+    /// Write result into provided array and consume the HashWriter instance.
+    pub fn finalize_into(self, out: &mut Output<D>) {
+        self.hasher.finalize_into(out)
+    }
+
+    /// Get output size of the hasher
+    pub fn output_size() -> usize {
+        <D as Digest>::output_size()
+    }
+}
+
+impl<D: Digest + Clone, W: io::Write + Clone> Clone for HashWriter<D, W> {
+    fn clone(&self) -> HashWriter<D, W> {
+        HashWriter {
+            writer: self.writer.clone(),
+            hasher: self.hasher.clone(),
+        }
+    }
+}
+
+impl<D: Digest, W: io::Write> io::Write for HashWriter<D, W> {
+    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+        let bytes = self.writer.write(buf)?;
+
+        if bytes > 0 {
+            self.hasher.update(&buf[0..bytes]);
+        }
+
+        Ok(bytes)
+    }
+
+    fn flush(&mut self) -> io::Result<()> {
+        self.writer.flush()
+    }
+}
+
+impl<D: Digest + FixedOutputReset, W: io::Write> HashWriter<D, W> {
+    /// Retrieve result and reset hasher instance.
+    pub fn finalize_reset(&mut self) -> Output<D> {
+        Digest::finalize_reset(&mut self.hasher)
+    }
+
+    /// Write result into provided array and reset the hasher instance.
+    pub fn finalize_into_reset(&mut self, out: &mut Output<D>) {
+        Digest::finalize_into_reset(&mut self.hasher, out)
+    }
+}
+impl<D: Digest + Reset, W: io::Write> Reset for HashWriter<D, W> {
+    fn reset(&mut self) {
+        Digest::reset(&mut self.hasher)
+    }
+}
diff --git a/digest/src/lib.rs b/digest/src/lib.rs
index b535bc8aa..a9f5b7a4a 100644
--- a/digest/src/lib.rs
+++ b/digest/src/lib.rs
@@ -299,3 +299,12 @@ impl fmt::Display for InvalidBufferSize {
 
 #[cfg(feature = "std")]
 impl std::error::Error for InvalidBufferSize {}
+
+#[cfg(feature = "std")]
+mod hashwriter;
+#[cfg(feature = "std")]
+pub use hashwriter::HashWriter;
+#[cfg(feature = "std")]
+mod hashreader;
+#[cfg(feature = "std")]
+pub use hashreader::HashReader;