From a037fec28383808ef7235ae5a19ec123cb694206 Mon Sep 17 00:00:00 2001
From: Paul Hauner <paul@paulhauner.com>
Date: Mon, 28 Jan 2019 11:22:25 +1100
Subject: [PATCH] Add `signable_message()` to `AttestationData`

---
 eth2/types/src/attestation/mod.rs                    |  8 +++++++-
 eth2/types/src/attestation/signing.rs                | 11 +++--------
 .../{attestation_data.rs => attestation_data/mod.rs} |  2 ++
 eth2/types/src/attestation_data/signing.rs           | 12 ++++++++++++
 .../types/src/beacon_state/attestation_validation.rs |  4 +++-
 5 files changed, 27 insertions(+), 10 deletions(-)
 rename eth2/types/src/{attestation_data.rs => attestation_data/mod.rs} (99%)
 create mode 100644 eth2/types/src/attestation_data/signing.rs

diff --git a/eth2/types/src/attestation/mod.rs b/eth2/types/src/attestation/mod.rs
index 3b278ef8ce..51c396a3b8 100644
--- a/eth2/types/src/attestation/mod.rs
+++ b/eth2/types/src/attestation/mod.rs
@@ -1,4 +1,4 @@
-use super::{AttestationData, Bitfield};
+use super::{AttestationData, Bitfield, Hash256};
 use crate::test_utils::TestRandom;
 use bls::AggregateSignature;
 use rand::RngCore;
@@ -15,6 +15,12 @@ pub struct Attestation {
     pub aggregate_signature: AggregateSignature,
 }
 
+impl Attestation {
+    pub fn canonical_root(&self) -> Hash256 {
+        Hash256::from(&self.hash_tree_root()[..])
+    }
+}
+
 impl Encodable for Attestation {
     fn ssz_append(&self, s: &mut SszStream) {
         s.append(&self.data);
diff --git a/eth2/types/src/attestation/signing.rs b/eth2/types/src/attestation/signing.rs
index a2ef13d2f8..1fefa068c5 100644
--- a/eth2/types/src/attestation/signing.rs
+++ b/eth2/types/src/attestation/signing.rs
@@ -1,12 +1,7 @@
-use crate::{Attestation, AttestationDataAndCustodyBit};
-use ssz::TreeHash;
+use crate::Attestation;
 
 impl Attestation {
-    pub fn signable_message(&self) -> Vec<u8> {
-        let attestation_data_and_custody_bit = AttestationDataAndCustodyBit {
-            data: self.data.clone(),
-            custody_bit: false,
-        };
-        attestation_data_and_custody_bit.hash_tree_root()
+    pub fn signable_message(&self, custody_bit: bool) -> Vec<u8> {
+        self.data.signable_message(custody_bit)
     }
 }
diff --git a/eth2/types/src/attestation_data.rs b/eth2/types/src/attestation_data/mod.rs
similarity index 99%
rename from eth2/types/src/attestation_data.rs
rename to eth2/types/src/attestation_data/mod.rs
index aab5f75632..f22b4007f1 100644
--- a/eth2/types/src/attestation_data.rs
+++ b/eth2/types/src/attestation_data/mod.rs
@@ -4,6 +4,8 @@ use rand::RngCore;
 use serde_derive::Serialize;
 use ssz::{hash, Decodable, DecodeError, Encodable, SszStream, TreeHash};
 
+mod signing;
+
 pub const SSZ_ATTESTION_DATA_LENGTH: usize = {
     8 +             // slot
     8 +             // shard
diff --git a/eth2/types/src/attestation_data/signing.rs b/eth2/types/src/attestation_data/signing.rs
new file mode 100644
index 0000000000..f14cc0ef48
--- /dev/null
+++ b/eth2/types/src/attestation_data/signing.rs
@@ -0,0 +1,12 @@
+use crate::{AttestationData, AttestationDataAndCustodyBit};
+use ssz::TreeHash;
+
+impl AttestationData {
+    pub fn signable_message(&self, custody_bit: bool) -> Vec<u8> {
+        let attestation_data_and_custody_bit = AttestationDataAndCustodyBit {
+            data: self.clone(),
+            custody_bit,
+        };
+        attestation_data_and_custody_bit.hash_tree_root()
+    }
+}
diff --git a/eth2/types/src/beacon_state/attestation_validation.rs b/eth2/types/src/beacon_state/attestation_validation.rs
index 160ff81282..26cc4e2b45 100644
--- a/eth2/types/src/beacon_state/attestation_validation.rs
+++ b/eth2/types/src/beacon_state/attestation_validation.rs
@@ -21,6 +21,8 @@ macro_rules! ensure {
     };
 }
 
+const PHASE_0_CUSTODY_BIT: bool = false;
+
 // TODO: define elsehwere.
 const DOMAIN_ATTESTATION: u64 = 1;
 
@@ -76,7 +78,7 @@ impl BeaconState {
         ensure!(
             bls_verify_aggregate(
                 &group_public_key,
-                &attestation.signable_message(),
+                &attestation.signable_message(PHASE_0_CUSTODY_BIT),
                 &attestation.aggregate_signature,
                 get_domain(&self.fork_data, attestation.data.slot, DOMAIN_ATTESTATION)
             ),