Use zero to represent a nonexistent account

This also fixes a bug in the thin client where a nonexistent account
would have triggered a panic because we were using `balances[k]` instead
of `balances.get(key)`.

Fixes #534
This commit is contained in:
Greg Fitzgerald 2018-07-02 17:31:10 -06:00 committed by Greg Fitzgerald
parent d2bb4dc14a
commit 0dabdfd48e
4 changed files with 31 additions and 28 deletions

View File

@ -497,11 +497,11 @@ impl Bank {
self.process_transaction(&tx).map(|_| sig) self.process_transaction(&tx).map(|_| sig)
} }
pub fn get_balance(&self, pubkey: &PublicKey) -> Option<i64> { pub fn get_balance(&self, pubkey: &PublicKey) -> i64 {
let bals = self.balances let bals = self.balances
.read() .read()
.expect("'balances' read lock in get_balance"); .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 { pub fn transaction_count(&self) -> usize {
@ -541,11 +541,11 @@ mod tests {
bank.transfer(1_000, &mint.keypair(), pubkey, mint.last_id()) bank.transfer(1_000, &mint.keypair(), pubkey, mint.last_id())
.unwrap(); .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()) bank.transfer(500, &mint.keypair(), pubkey, mint.last_id())
.unwrap(); .unwrap();
assert_eq!(bank.get_balance(&pubkey).unwrap(), 1_500); assert_eq!(bank.get_balance(&pubkey), 1_500);
assert_eq!(bank.transaction_count(), 2); assert_eq!(bank.transaction_count(), 2);
} }
@ -588,8 +588,8 @@ mod tests {
assert_eq!(bank.transaction_count(), 1); assert_eq!(bank.transaction_count(), 1);
let mint_pubkey = mint.keypair().pubkey(); let mint_pubkey = mint.keypair().pubkey();
assert_eq!(bank.get_balance(&mint_pubkey).unwrap(), 10_000); assert_eq!(bank.get_balance(&mint_pubkey), 10_000);
assert_eq!(bank.get_balance(&pubkey).unwrap(), 1_000); assert_eq!(bank.get_balance(&pubkey), 1_000);
} }
#[test] #[test]
@ -599,7 +599,7 @@ mod tests {
let pubkey = KeyPair::new().pubkey(); let pubkey = KeyPair::new().pubkey();
bank.transfer(500, &mint.keypair(), pubkey, mint.last_id()) bank.transfer(500, &mint.keypair(), pubkey, mint.last_id())
.unwrap(); .unwrap();
assert_eq!(bank.get_balance(&pubkey).unwrap(), 500); assert_eq!(bank.get_balance(&pubkey), 500);
} }
#[test] #[test]
@ -612,26 +612,26 @@ mod tests {
.unwrap(); .unwrap();
// Mint's balance will be zero because all funds are locked up. // 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. // tx count is 1, because debits were applied.
assert_eq!(bank.transaction_count(), 1); assert_eq!(bank.transaction_count(), 1);
// pubkey's balance will be None because the funds have not been // pubkey's balance will be None because the funds have not been
// sent. // sent.
assert_eq!(bank.get_balance(&pubkey), None); assert_eq!(bank.get_balance(&pubkey), 0);
// Now, acknowledge the time in the condition occurred and // Now, acknowledge the time in the condition occurred and
// that pubkey's funds are now available. // that pubkey's funds are now available.
bank.apply_timestamp(mint.pubkey(), dt).unwrap(); 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 is still 1, because we chose not to count timestamp transactions
// tx count. // tx count.
assert_eq!(bank.transaction_count(), 1); assert_eq!(bank.transaction_count(), 1);
bank.apply_timestamp(mint.pubkey(), dt).unwrap(); // <-- Attack! Attempt to process completed transaction. 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] #[test]
@ -646,8 +646,8 @@ mod tests {
bank.transfer_on_date(1, &mint.keypair(), pubkey, dt, mint.last_id()) bank.transfer_on_date(1, &mint.keypair(), pubkey, dt, mint.last_id())
.unwrap(); .unwrap();
assert_eq!(bank.get_balance(&mint.pubkey()), None); assert_eq!(bank.get_balance(&mint.pubkey()), 0);
assert_eq!(bank.get_balance(&pubkey), Some(1)); assert_eq!(bank.get_balance(&pubkey), 1);
} }
#[test] #[test]
@ -663,22 +663,22 @@ mod tests {
assert_eq!(bank.transaction_count(), 1); assert_eq!(bank.transaction_count(), 1);
// Mint's balance will be zero because all funds are locked up. // 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 // pubkey's balance will be None because the funds have not been
// sent. // 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. // Now, cancel the trancaction. Mint gets her funds back, pubkey never sees them.
bank.apply_signature(mint.pubkey(), sig).unwrap(); bank.apply_signature(mint.pubkey(), sig).unwrap();
assert_eq!(bank.get_balance(&mint.pubkey()), Some(1)); assert_eq!(bank.get_balance(&mint.pubkey()), 1);
assert_eq!(bank.get_balance(&pubkey), None); assert_eq!(bank.get_balance(&pubkey), 0);
// Assert cancel doesn't cause count to go backward. // Assert cancel doesn't cause count to go backward.
assert_eq!(bank.transaction_count(), 1); assert_eq!(bank.transaction_count(), 1);
bank.apply_signature(mint.pubkey(), sig).unwrap(); // <-- Attack! Attempt to cancel completed transaction. 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] #[test]
@ -776,7 +776,7 @@ mod tests {
let genesis = mint.create_entries(); let genesis = mint.create_entries();
let bank = Bank::default(); let bank = Bank::default();
bank.process_ledger(genesis).unwrap(); 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<Item = Entry> { fn create_sample_block(mint: &Mint) -> impl Iterator<Item = Entry> {
@ -797,7 +797,7 @@ mod tests {
let (ledger, pubkey) = create_sample_ledger(); let (ledger, pubkey) = create_sample_ledger();
let bank = Bank::default(); let bank = Bank::default();
bank.process_ledger(ledger).unwrap(); 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. // 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(); let bank = Bank::default();
bank.process_ledger(ledger).unwrap(); bank.process_ledger(ledger).unwrap();
assert_eq!(bank.get_balance(&pubkey).unwrap(), 1); assert_eq!(bank.get_balance(&pubkey), 1);
} }
#[test] #[test]
@ -830,7 +830,7 @@ mod tests {
let bank = Bank::default(); let bank = Bank::default();
bank.process_ledger(genesis.chain(block)).unwrap(); bank.process_ledger(genesis.chain(block)).unwrap();
assert_eq!(bank.get_balance(&mint.pubkey()).unwrap(), 1); assert_eq!(bank.get_balance(&mint.pubkey()), 1);
} }
} }

View File

@ -21,7 +21,7 @@ impl Request {
#[derive(Serialize, Deserialize, Debug)] #[derive(Serialize, Deserialize, Debug)]
pub enum Response { pub enum Response {
Balance { key: PublicKey, val: Option<i64> }, Balance { key: PublicKey, val: i64 },
LastId { id: Hash }, LastId { id: Hash },
TransactionCount { transaction_count: u64 }, TransactionCount { transaction_count: u64 },
SignatureStatus { signature_status: bool }, SignatureStatus { signature_status: bool },

View File

@ -20,7 +20,7 @@ pub struct ThinClient {
transactions_socket: UdpSocket, transactions_socket: UdpSocket,
last_id: Option<Hash>, last_id: Option<Hash>,
transaction_count: u64, transaction_count: u64,
balances: HashMap<PublicKey, Option<i64>>, balances: HashMap<PublicKey, i64>,
signature_status: bool, signature_status: bool,
} }
@ -124,7 +124,10 @@ impl ThinClient {
} }
self.process_response(resp); 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, /// Request the transaction count. If the response packet is dropped by the network,
@ -186,7 +189,7 @@ impl ThinClient {
let now = Instant::now(); let now = Instant::now();
loop { loop {
balance = self.get_balance(pubkey); 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; break;
} }
} }

View File

@ -257,10 +257,10 @@ pub mod tests {
trace!("msg: {:?}", msg); 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); 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); assert_eq!(bob_balance, starting_balance - alice_ref_balance);
exit.store(true, Ordering::Relaxed); exit.store(true, Ordering::Relaxed);