From 07aa2e12602b684eb8f328af4b15239ce0549fb1 Mon Sep 17 00:00:00 2001
From: Greg Fitzgerald <garious@gmail.com>
Date: Mon, 2 Apr 2018 20:47:48 -0600
Subject: [PATCH] Add witness data to entry hash

Otherwise, witnesses can be dropped or reordered by a malicious
generator.
---
 src/entry.rs | 22 ++++++++++++++++++----
 src/event.rs | 25 +++++++++++++++++++------
 2 files changed, 37 insertions(+), 10 deletions(-)

diff --git a/src/entry.rs b/src/entry.rs
index 985bab83cc..e9d13157f2 100644
--- a/src/entry.rs
+++ b/src/entry.rs
@@ -43,6 +43,23 @@ impl Entry {
     }
 }
 
+fn add_event_data(hash_data: &mut Vec<u8>, event: &Event) {
+    match *event {
+        Event::Transaction(ref tr) => {
+            hash_data.push(0u8);
+            hash_data.extend_from_slice(&tr.sig);
+        }
+        Event::Signature { ref sig, .. } => {
+            hash_data.push(1u8);
+            hash_data.extend_from_slice(sig);
+        }
+        Event::Timestamp { ref sig, .. } => {
+            hash_data.push(2u8);
+            hash_data.extend_from_slice(sig);
+        }
+    }
+}
+
 /// Creates the hash `num_hashes` after `start_hash`. If the event contains
 /// signature, the final hash will be a hash of both the previous ID and
 /// the signature.
@@ -55,10 +72,7 @@ pub fn next_hash(start_hash: &Hash, num_hashes: u64, events: &[Event]) -> Hash {
     // Hash all the event data
     let mut hash_data = vec![];
     for event in events {
-        let sig = event.get_signature();
-        if let Some(sig) = sig {
-            hash_data.extend_from_slice(&sig);
-        }
+        add_event_data(&mut hash_data, event);
     }
 
     if !hash_data.is_empty() {
diff --git a/src/event.rs b/src/event.rs
index 9940959a3d..dbd2a93412 100644
--- a/src/event.rs
+++ b/src/event.rs
@@ -33,12 +33,13 @@ impl Event {
         }
     }
 
-    // TODO: Rename this to transaction_signature().
-    /// If the Event is a Transaction, return its Signature.
-    pub fn get_signature(&self) -> Option<Signature> {
-        match *self {
-            Event::Transaction(ref tr) => Some(tr.sig),
-            Event::Signature { .. } | Event::Timestamp { .. } => None,
+    /// Create and sign a new Witness Signature. Used for unit-testing.
+    pub fn new_signature(from: &KeyPair, tx_sig: Signature) -> Self {
+        let sig = Signature::clone_from_slice(from.sign(&tx_sig).as_ref());
+        Event::Signature {
+            from: from.pubkey(),
+            tx_sig,
+            sig,
         }
     }
 
@@ -52,3 +53,15 @@ impl Event {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use signature::{KeyPair, KeyPairUtil};
+
+    #[test]
+    fn test_event_verify() {
+        assert!(Event::new_timestamp(&KeyPair::new(), Utc::now()).verify());
+        assert!(Event::new_signature(&KeyPair::new(), Signature::default()).verify());
+    }
+}