diff --git a/src/accountant.rs b/src/accountant.rs index 98e2d82dda..f244f788a7 100644 --- a/src/accountant.rs +++ b/src/accountant.rs @@ -6,7 +6,7 @@ use log::{Entry, Sha256Hash}; use event::{Event, PublicKey, Signature}; use historian::Historian; use ring::signature::Ed25519KeyPair; -use std::sync::mpsc::{RecvError, SendError}; +use std::sync::mpsc::SendError; use std::collections::HashMap; pub struct Accountant { @@ -25,33 +25,6 @@ impl Accountant { } } - pub fn process_event(self: &mut Self, event: &Event) { - match *event { - Event::Claim { key, data, .. } => { - if self.balances.contains_key(&key) { - if let Some(x) = self.balances.get_mut(&key) { - *x += data; - } - } else { - self.balances.insert(key, data); - } - } - Event::Transaction { from, to, data, .. } => { - if let Some(x) = self.balances.get_mut(&from) { - *x -= data; - } - if self.balances.contains_key(&to) { - if let Some(x) = self.balances.get_mut(&to) { - *x += data; - } - } else { - self.balances.insert(to, data); - } - } - _ => (), - } - } - pub fn sync(self: &mut Self) -> Vec> { let mut entries = vec![]; while let Ok(entry) = self.historian.receiver.try_recv() { @@ -62,25 +35,35 @@ impl Accountant { self.end_hash = last_entry.end_hash; } - for e in &entries { - self.process_event(&e.event); - } - entries } pub fn deposit_signed( - self: &Self, + self: &mut Self, key: PublicKey, data: u64, sig: Signature, ) -> Result<(), SendError>> { let event = Event::Claim { key, data, sig }; + if !self.historian.verify_event(&event) { + // TODO: Replace the SendError result with a custom one. + println!("Rejecting transaction: Invalid event"); + return Ok(()); + } + + if self.balances.contains_key(&key) { + if let Some(x) = self.balances.get_mut(&key) { + *x += data; + } + } else { + self.balances.insert(key, data); + } + self.historian.sender.send(event) } pub fn deposit( - self: &Self, + self: &mut Self, n: u64, keypair: &Ed25519KeyPair, ) -> Result>> { @@ -97,17 +80,36 @@ impl Accountant { data: u64, sig: Signature, ) -> Result<(), SendError>> { - if self.get_balance(&from).unwrap() < data { - // TODO: Replace the SendError result with a custom one. - println!("Error: Insufficient funds"); - return Ok(()); - } let event = Event::Transaction { from, to, data, sig, }; + if !self.historian.verify_event(&event) { + // TODO: Replace the SendError result with a custom one. + println!("Rejecting transaction: Invalid event"); + return Ok(()); + } + + if self.get_balance(&from).unwrap_or(0) < data { + // TODO: Replace the SendError result with a custom one. + println!("Rejecting transaction: Insufficient funds"); + return Ok(()); + } + + if let Some(x) = self.balances.get_mut(&from) { + *x -= data; + } + + if self.balances.contains_key(&to) { + if let Some(x) = self.balances.get_mut(&to) { + *x += data; + } + } else { + self.balances.insert(to, data); + } + self.historian.sender.send(event) } @@ -118,15 +120,13 @@ impl Accountant { to: PublicKey, ) -> Result>> { use event::{get_pubkey, sign_transaction_data}; - let from = get_pubkey(keypair); let sig = sign_transaction_data(&n, keypair, &to); self.transfer_signed(from, to, n, sig).map(|_| sig) } - pub fn get_balance(self: &mut Self, pubkey: &PublicKey) -> Result { - self.sync(); - Ok(*self.balances.get(pubkey).unwrap_or(&0)) + pub fn get_balance(self: &Self, pubkey: &PublicKey) -> Option { + self.balances.get(pubkey).map(|x| *x) } pub fn wait_on_signature(self: &mut Self, wait_sig: &Signature) { diff --git a/src/historian.rs b/src/historian.rs index be30f66867..c01947c1dc 100644 --- a/src/historian.rs +++ b/src/historian.rs @@ -6,7 +6,7 @@ //! The resulting stream of entries represents ordered events in time. use std::thread::JoinHandle; -use std::collections::HashMap; +use std::collections::HashSet; use std::sync::mpsc::{Receiver, SyncSender}; use std::time::{Duration, SystemTime}; use log::{hash, hash_event, Entry, Sha256Hash}; @@ -18,6 +18,7 @@ pub struct Historian { pub sender: SyncSender>, pub receiver: Receiver>, pub thread_hdl: JoinHandle<(Entry, ExitReason)>, + pub signatures: HashSet, } #[derive(Debug, PartialEq, Eq)] @@ -44,10 +45,25 @@ fn log_event( Ok(()) } +fn verify_event_and_reserve_signature( + signatures: &mut HashSet, + event: &Event, +) -> bool { + if !verify_event(&event) { + return false; + } + if let Some(sig) = get_signature(&event) { + if signatures.contains(&sig) { + return false; + } + signatures.insert(sig); + } + true +} + fn log_events( receiver: &Receiver>, sender: &SyncSender>, - signatures: &mut HashMap, num_hashes: &mut u64, end_hash: &mut Sha256Hash, epoch: SystemTime, @@ -65,15 +81,7 @@ fn log_events( } match receiver.try_recv() { Ok(event) => { - if verify_event(&event) { - if let Some(sig) = get_signature(&event) { - if signatures.contains_key(&sig) { - continue; - } - signatures.insert(sig, true); - } - log_event(sender, num_hashes, end_hash, event)?; - } + log_event(sender, num_hashes, end_hash, event)?; } Err(TryRecvError::Empty) => { return Ok(()); @@ -103,13 +111,11 @@ pub fn create_logger( let mut end_hash = start_hash; let mut num_hashes = 0; let mut num_ticks = 0; - let mut signatures = HashMap::new(); let epoch = SystemTime::now(); loop { if let Err(err) = log_events( &receiver, &sender, - &mut signatures, &mut num_hashes, &mut end_hash, epoch, @@ -130,12 +136,17 @@ impl Historian { let (sender, event_receiver) = sync_channel(1000); let (entry_sender, receiver) = sync_channel(1000); let thread_hdl = create_logger(*start_hash, ms_per_tick, event_receiver, entry_sender); + let signatures = HashSet::new(); Historian { sender, receiver, thread_hdl, + signatures, } } + pub fn verify_event(self: &mut Self, event: &Event) -> bool { + return verify_event_and_reserve_signature(&mut self.signatures, event); + } } #[cfg(test)] @@ -201,22 +212,28 @@ mod tests { } #[test] - fn test_bad_event_attack() { - let zero = Sha256Hash::default(); - let hist = Historian::new(&zero, None); + fn test_bad_event_signature() { let keypair = generate_keypair(); + let sig = sign_serialized(&hash(b"hello, world"), &keypair); let event0 = Event::Claim { key: get_pubkey(&keypair), data: hash(b"goodbye cruel world"), - sig: sign_serialized(&hash(b"hello, world"), &keypair), + sig, }; - hist.sender.send(event0).unwrap(); - drop(hist.sender); - assert_eq!( - hist.thread_hdl.join().unwrap().1, - ExitReason::RecvDisconnected - ); - let entries: Vec> = hist.receiver.iter().collect(); - assert_eq!(entries.len(), 0); + let mut sigs = HashSet::new(); + assert!(!verify_event_and_reserve_signature(&mut sigs, &event0)); + assert!(!sigs.contains(&sig)); + } + + #[test] + fn test_duplicate_event_signature() { + let keypair = generate_keypair(); + let key = get_pubkey(&keypair); + let data = &hash(b"hello, world"); + let sig = sign_serialized(data, &keypair); + let event0 = Event::Claim { key, data, sig }; + let mut sigs = HashSet::new(); + assert!(verify_event_and_reserve_signature(&mut sigs, &event0)); + assert!(!verify_event_and_reserve_signature(&mut sigs, &event0)); } }