From 7f01ec7c27dd0c5c7b7a95779782bea4c08574cc Mon Sep 17 00:00:00 2001
From: Paul Hauner <paul@paulhauner.com>
Date: Mon, 24 Sep 2018 13:16:39 +1000
Subject: [PATCH] Extend validator_store

---
 lighthouse/db/mod.rs                    |   2 +
 lighthouse/db/stores/mod.rs             |  15 +++
 lighthouse/db/stores/validator_store.rs | 116 ++++++++++++++++++++++--
 3 files changed, 125 insertions(+), 8 deletions(-)

diff --git a/lighthouse/db/mod.rs b/lighthouse/db/mod.rs
index e030fc009..5f8898861 100644
--- a/lighthouse/db/mod.rs
+++ b/lighthouse/db/mod.rs
@@ -6,6 +6,8 @@ mod memory_db;
 mod traits;
 pub mod stores;
 
+use super::bls;
+
 
 pub use self::disk_db::DiskDB;
 pub use self::memory_db::MemoryDB;
diff --git a/lighthouse/db/stores/mod.rs b/lighthouse/db/stores/mod.rs
index a7e5c60bb..58c0a52db 100644
--- a/lighthouse/db/stores/mod.rs
+++ b/lighthouse/db/stores/mod.rs
@@ -11,6 +11,21 @@ pub use self::block_store::BlockStore;
 pub use self::pow_chain_store::PoWChainStore;
 pub use self::validator_store::ValidatorStore;
 
+use super::bls;
+
 const BLOCKS_DB_COLUMN: &str = "blocks";
 const POW_CHAIN_DB_COLUMN: &str = "powchain";
 const VALIDATOR_DB_COLUMN: &str = "validator";
+
+#[derive(Debug, PartialEq)]
+pub enum StoreError {
+    DBError(String),
+    DecodeError,
+    EncodeError,
+}
+
+impl From<DBError> for StoreError {
+    fn from(error: DBError) -> Self {
+        StoreError::DBError(error.message)
+    }
+}
diff --git a/lighthouse/db/stores/validator_store.rs b/lighthouse/db/stores/validator_store.rs
index badea16e2..a20e7503a 100644
--- a/lighthouse/db/stores/validator_store.rs
+++ b/lighthouse/db/stores/validator_store.rs
@@ -1,9 +1,21 @@
+extern crate bytes;
+
+use self::bytes::{
+    BufMut,
+    BytesMut,
+};
 use std::sync::Arc;
 use super::{
     ClientDB,
-    DBError,
+    StoreError,
 };
 use super::VALIDATOR_DB_COLUMN as DB_COLUMN;
+use super::bls::PublicKey;
+
+#[derive(Debug, PartialEq)]
+enum KeyPrefixes {
+    PublicKey,
+}
 
 pub struct ValidatorStore<T>
     where T: ClientDB
@@ -18,17 +30,105 @@ impl<T: ClientDB> ValidatorStore<T> {
         }
     }
 
-    pub fn put_validator_record_by_index(&self, hash: &[u8], val: &[u8])
-        -> Result<(), DBError>
+    fn prefix_bytes(&self, key_prefix: KeyPrefixes)
+        -> Vec<u8>
     {
-        self.db.put(DB_COLUMN, hash, val)
+        match key_prefix {
+            KeyPrefixes::PublicKey => b"pubkey".to_vec(),
+        }
     }
 
-    pub fn get_validator_record_by_index(&self, hash: &[u8])
-        -> Result<bool, DBError>
+    fn get_db_key_for_index(&self, key_prefix: KeyPrefixes, index: usize)
+        -> Vec<u8>
     {
-        self.db.exists(DB_COLUMN, hash)
+        let mut buf = BytesMut::with_capacity(6 + 8);
+        buf.put(self.prefix_bytes(key_prefix));
+        buf.put_u64_be(index as u64);
+        buf.take().to_vec()
+    }
+
+    pub fn put_public_key_by_index(&self, index: usize, public_key: &PublicKey)
+        -> Result<(), StoreError>
+    {
+        let key = self.get_db_key_for_index(KeyPrefixes::PublicKey, index);
+        let val = public_key.as_bytes();
+        self.db.put(DB_COLUMN, &key[..], &val[..])
+                    .map_err(|e| StoreError::from(e))
+    }
+
+    pub fn get_public_key_by_index(&self, index: usize)
+        -> Result<Option<PublicKey>, StoreError>
+    {
+        let key = self.get_db_key_for_index(KeyPrefixes::PublicKey, index);
+        let val = self.db.get(DB_COLUMN, &key[..])?;
+        match val {
+            None => Ok(None),
+            Some(val) => {
+                match PublicKey::from_bytes(&val) {
+                    Ok(key) => Ok(Some(key)),
+                    Err(_) => Err(StoreError::DecodeError),
+                }
+            }
+        }
     }
 }
 
-// TODO: add tests
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use super::super::super::{
+        MemoryDB,
+        ClientDB,
+    };
+    use super::super::bls::Keypair;
+
+    fn open_client_db() -> MemoryDB {
+        let columns = vec![DB_COLUMN];
+        MemoryDB::open(Some(&columns))
+    }
+
+    #[test]
+    fn test_validator_store_put_get() {
+        let db = Arc::new(open_client_db());
+        let store = ValidatorStore::new(db);
+
+        let keys = vec![
+            Keypair::random(),
+            Keypair::random(),
+            Keypair::random(),
+            Keypair::random(),
+            Keypair::random(),
+        ];
+
+        for i in 0..keys.len() {
+            store.put_public_key_by_index(i, &keys[i].pk).unwrap();
+        }
+
+        /*
+         * Check all keys are retrieved correctly.
+         */
+        for i in 0..keys.len() {
+            let retrieved = store.get_public_key_by_index(i)
+                .unwrap().unwrap();
+            assert_eq!(retrieved, keys[i].pk);
+        }
+
+        /*
+         * Check that an index that wasn't stored returns None.
+         */
+        assert!(store.get_public_key_by_index(keys.len() + 1)
+                .unwrap().is_none());
+    }
+
+    #[test]
+    fn test_validator_store_bad_key() {
+        let db = Arc::new(open_client_db());
+        let store = ValidatorStore::new(db.clone());
+
+        let key = store.get_db_key_for_index(KeyPrefixes::PublicKey, 42);
+        db.put(DB_COLUMN, &key[..], "cats".as_bytes()).unwrap();
+
+        assert_eq!(store.get_public_key_by_index(42),
+            Err(StoreError::DecodeError));
+    }
+}