diff --git a/src/bank.rs b/src/bank.rs index 76e85f0d75..37c32acde0 100644 --- a/src/bank.rs +++ b/src/bank.rs @@ -497,11 +497,11 @@ impl Bank { self.process_transaction(&tx).map(|_| sig) } - pub fn get_balance(&self, pubkey: &PublicKey) -> Option { + pub fn get_balance(&self, pubkey: &PublicKey) -> i64 { let bals = self.balances .read() .expect("'balances' read lock in get_balance"); - bals.get(pubkey).map(|x| *x) + bals.get(pubkey).map(|x| *x).unwrap_or(0) } pub fn transaction_count(&self) -> usize { @@ -541,11 +541,11 @@ mod tests { bank.transfer(1_000, &mint.keypair(), pubkey, mint.last_id()) .unwrap(); - assert_eq!(bank.get_balance(&pubkey).unwrap(), 1_000); + assert_eq!(bank.get_balance(&pubkey), 1_000); bank.transfer(500, &mint.keypair(), pubkey, mint.last_id()) .unwrap(); - assert_eq!(bank.get_balance(&pubkey).unwrap(), 1_500); + assert_eq!(bank.get_balance(&pubkey), 1_500); assert_eq!(bank.transaction_count(), 2); } @@ -588,8 +588,8 @@ mod tests { assert_eq!(bank.transaction_count(), 1); let mint_pubkey = mint.keypair().pubkey(); - assert_eq!(bank.get_balance(&mint_pubkey).unwrap(), 10_000); - assert_eq!(bank.get_balance(&pubkey).unwrap(), 1_000); + assert_eq!(bank.get_balance(&mint_pubkey), 10_000); + assert_eq!(bank.get_balance(&pubkey), 1_000); } #[test] @@ -599,7 +599,7 @@ mod tests { let pubkey = KeyPair::new().pubkey(); bank.transfer(500, &mint.keypair(), pubkey, mint.last_id()) .unwrap(); - assert_eq!(bank.get_balance(&pubkey).unwrap(), 500); + assert_eq!(bank.get_balance(&pubkey), 500); } #[test] @@ -612,26 +612,26 @@ mod tests { .unwrap(); // Mint's balance will be zero because all funds are locked up. - assert_eq!(bank.get_balance(&mint.pubkey()), None); + assert_eq!(bank.get_balance(&mint.pubkey()), 0); // tx count is 1, because debits were applied. assert_eq!(bank.transaction_count(), 1); // pubkey's balance will be None because the funds have not been // sent. - assert_eq!(bank.get_balance(&pubkey), None); + assert_eq!(bank.get_balance(&pubkey), 0); // Now, acknowledge the time in the condition occurred and // that pubkey's funds are now available. bank.apply_timestamp(mint.pubkey(), dt).unwrap(); - assert_eq!(bank.get_balance(&pubkey), Some(1)); + assert_eq!(bank.get_balance(&pubkey), 1); // tx count is still 1, because we chose not to count timestamp transactions // tx count. assert_eq!(bank.transaction_count(), 1); bank.apply_timestamp(mint.pubkey(), dt).unwrap(); // <-- Attack! Attempt to process completed transaction. - assert_ne!(bank.get_balance(&pubkey), Some(2)); + assert_ne!(bank.get_balance(&pubkey), 2); } #[test] @@ -646,8 +646,8 @@ mod tests { bank.transfer_on_date(1, &mint.keypair(), pubkey, dt, mint.last_id()) .unwrap(); - assert_eq!(bank.get_balance(&mint.pubkey()), None); - assert_eq!(bank.get_balance(&pubkey), Some(1)); + assert_eq!(bank.get_balance(&mint.pubkey()), 0); + assert_eq!(bank.get_balance(&pubkey), 1); } #[test] @@ -663,22 +663,22 @@ mod tests { assert_eq!(bank.transaction_count(), 1); // Mint's balance will be zero because all funds are locked up. - assert_eq!(bank.get_balance(&mint.pubkey()), None); + assert_eq!(bank.get_balance(&mint.pubkey()), 0); // pubkey's balance will be None because the funds have not been // sent. - assert_eq!(bank.get_balance(&pubkey), None); + assert_eq!(bank.get_balance(&pubkey), 0); // Now, cancel the trancaction. Mint gets her funds back, pubkey never sees them. bank.apply_signature(mint.pubkey(), sig).unwrap(); - assert_eq!(bank.get_balance(&mint.pubkey()), Some(1)); - assert_eq!(bank.get_balance(&pubkey), None); + assert_eq!(bank.get_balance(&mint.pubkey()), 1); + assert_eq!(bank.get_balance(&pubkey), 0); // Assert cancel doesn't cause count to go backward. assert_eq!(bank.transaction_count(), 1); bank.apply_signature(mint.pubkey(), sig).unwrap(); // <-- Attack! Attempt to cancel completed transaction. - assert_ne!(bank.get_balance(&mint.pubkey()), Some(2)); + assert_ne!(bank.get_balance(&mint.pubkey()), 2); } #[test] @@ -776,7 +776,7 @@ mod tests { let genesis = mint.create_entries(); let bank = Bank::default(); bank.process_ledger(genesis).unwrap(); - assert_eq!(bank.get_balance(&mint.pubkey()).unwrap(), 1); + assert_eq!(bank.get_balance(&mint.pubkey()), 1); } fn create_sample_block(mint: &Mint) -> impl Iterator { @@ -797,7 +797,7 @@ mod tests { let (ledger, pubkey) = create_sample_ledger(); let bank = Bank::default(); bank.process_ledger(ledger).unwrap(); - assert_eq!(bank.get_balance(&pubkey).unwrap(), 1); + assert_eq!(bank.get_balance(&pubkey), 1); } // Write the given entries to a file and then return a file iterator to them. @@ -819,7 +819,7 @@ mod tests { let bank = Bank::default(); bank.process_ledger(ledger).unwrap(); - assert_eq!(bank.get_balance(&pubkey).unwrap(), 1); + assert_eq!(bank.get_balance(&pubkey), 1); } #[test] @@ -830,7 +830,7 @@ mod tests { let bank = Bank::default(); bank.process_ledger(genesis.chain(block)).unwrap(); - assert_eq!(bank.get_balance(&mint.pubkey()).unwrap(), 1); + assert_eq!(bank.get_balance(&mint.pubkey()), 1); } } diff --git a/src/request.rs b/src/request.rs index a27e71c73a..a1a006c38b 100644 --- a/src/request.rs +++ b/src/request.rs @@ -21,7 +21,7 @@ impl Request { #[derive(Serialize, Deserialize, Debug)] pub enum Response { - Balance { key: PublicKey, val: Option }, + Balance { key: PublicKey, val: i64 }, LastId { id: Hash }, TransactionCount { transaction_count: u64 }, SignatureStatus { signature_status: bool }, diff --git a/src/thin_client.rs b/src/thin_client.rs index 513fa41563..90cb1eab3e 100644 --- a/src/thin_client.rs +++ b/src/thin_client.rs @@ -20,7 +20,7 @@ pub struct ThinClient { transactions_socket: UdpSocket, last_id: Option, transaction_count: u64, - balances: HashMap>, + balances: HashMap, signature_status: bool, } @@ -124,7 +124,10 @@ impl ThinClient { } self.process_response(resp); } - self.balances[pubkey].ok_or(io::Error::new(io::ErrorKind::Other, "nokey")) + self.balances + .get(pubkey) + .map(|x| *x) + .ok_or(io::Error::new(io::ErrorKind::Other, "nokey")) } /// Request the transaction count. If the response packet is dropped by the network, @@ -186,7 +189,7 @@ impl ThinClient { let now = Instant::now(); loop { balance = self.get_balance(pubkey); - if balance.is_ok() || now.elapsed().as_secs() > 1 { + if balance.is_ok() && *balance.as_ref().unwrap() != 0 || now.elapsed().as_secs() > 1 { break; } } diff --git a/src/tvu.rs b/src/tvu.rs index a7aec358ba..d28f0d2c21 100644 --- a/src/tvu.rs +++ b/src/tvu.rs @@ -257,10 +257,10 @@ pub mod tests { trace!("msg: {:?}", msg); } - let alice_balance = bank.get_balance(&mint.keypair().pubkey()).unwrap(); + let alice_balance = bank.get_balance(&mint.keypair().pubkey()); assert_eq!(alice_balance, alice_ref_balance); - let bob_balance = bank.get_balance(&bob_keypair.pubkey()).unwrap(); + let bob_balance = bank.get_balance(&bob_keypair.pubkey()); assert_eq!(bob_balance, starting_balance - alice_ref_balance); exit.store(true, Ordering::Relaxed);