From 4283dee9dfb5666ce0a285654ef322fc69aa13c3 Mon Sep 17 00:00:00 2001
From: Jack Lloyd <jack.lloyd@fortanix.com>
Date: Thu, 16 Apr 2020 18:36:48 -0400
Subject: [PATCH] Allow reusing an Md context object

---
 mbedtls/src/hash/mod.rs | 37 ++++++++++++++++++++++++++++++++++++-
 1 file changed, 36 insertions(+), 1 deletion(-)

diff --git a/mbedtls/src/hash/mod.rs b/mbedtls/src/hash/mod.rs
index 7779607d0..e27134545 100644
--- a/mbedtls/src/hash/mod.rs
+++ b/mbedtls/src/hash/mod.rs
@@ -98,13 +98,14 @@ impl Md {
         Ok(())
     }
 
-    pub fn finish(mut self, out: &mut [u8]) -> Result<usize> {
+    pub fn finish(&mut self, out: &mut [u8]) -> Result<usize> {
         unsafe {
             let olen = (*self.inner.md_info).size as usize;
             if out.len() < olen {
                 return Err(Error::MdBadInputData);
             }
             md_finish(&mut self.inner, out.as_mut_ptr()).into_result()?;
+            md_starts(&mut self.inner).into_result()?;
             Ok(olen)
         }
     }
@@ -281,3 +282,37 @@ pub fn pbkdf_pkcs12(
         Ok(())
     }
 }
+
+
+#[test]
+fn test_digest_restart() {
+
+    use crate::hash::Type as MdType;
+
+    let mut digest1 = vec![0u8; 32];
+    Md::hash(MdType::Sha256, &[0,1,2,3], &mut digest1).unwrap();
+
+    let mut digest2 = vec![0u8; 32];
+    let mut md = Md::new(MdType::Sha256).unwrap();
+    md.update(&[0]).unwrap();
+    md.update(&[]).unwrap();
+    md.update(&[1,2]).unwrap();
+    md.update(&[3]).unwrap();
+    md.finish(&mut digest2).unwrap();
+
+    assert_eq!(digest1, digest2);
+
+    // hash some other data using same context
+    let mut digest3 = vec![0u8; 32];
+    md.update(&[9,8,7,6,5,4,3,2,2]).unwrap();
+    md.finish(&mut digest3).unwrap();
+    assert!(digest1 != digest3);
+
+    // finally hash same message as prior ensuring same output
+    let mut digest4 = vec![0u8; 32];
+    md.update(&[0,1]).unwrap();
+    md.update(&[2]).unwrap();
+    md.update(&[3]).unwrap();
+    md.finish(&mut digest4).unwrap();
+    assert_eq!(digest1, digest4);
+}