gaspersic/solana
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Cleanup Budget

Browse Source 
* BudgetProgram -> BudgetState
* Instruction -> BudgetInstruction
* Move BudgetState into its own module
* BudgetInstruction::NewBudget -> BudgetInstruction::InitializeAccount
* BudgetInstruction::new_budget -> BudgetInstruction::new_initialize_account
This commit is contained in:
Greg Fitzgerald
2019-03-03 14:17:51 -07:00
parent fc5fcd6cd4
commit baca3e6b6b
5 changed files with 222 additions and 197 deletions
Download Patch File Download Diff File Expand all files Collapse all files

278
programs/budget/src/budget_program.rs
View File

@ -1,95 +1,141 @@
//! budget program //! budget program
use bincode::{self, deserialize, serialize_into, serialized_size}; use crate::budget_state::{BudgetError, BudgetState};
use bincode::deserialize;
use chrono::prelude::{DateTime, Utc}; use chrono::prelude::{DateTime, Utc};
use log::*; use log::*;
use serde_derive::{Deserialize, Serialize}; use solana_budget_api::budget_instruction::BudgetInstruction;
use solana_budget_api::budget_expr::BudgetExpr;
use solana_budget_api::budget_instruction::Instruction;
use solana_budget_api::payment_plan::Witness; use solana_budget_api::payment_plan::Witness;
use solana_sdk::account::KeyedAccount; use solana_sdk::account::KeyedAccount;
use std::io;
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)] /// Process a Witness Signature. Any payment plans waiting on this signature
pub enum BudgetError { /// will progress one step.
InsufficientFunds, fn apply_signature(
ContractAlreadyExists, budget_state: &mut BudgetState,
ContractNotPending, keyed_accounts: &mut [KeyedAccount],
SourceIsPendingContract, ) -> Result<(), BudgetError> {
UninitializedContract, let mut final_payment = None;
NegativeTokens, if let Some(ref mut expr) = budget_state.pending_budget {
DestinationMissing, let key = match keyed_accounts[0].signer_key() {
FailedWitness, None => return Err(BudgetError::UnsignedKey),
UserdataTooSmall, Some(key) => key,
UserdataDeserializeFailure, };
UnsignedKey, expr.apply_witness(&Witness::Signature, key);
final_payment = expr.final_payment();
} }
#[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq)] if let Some(payment) = final_payment {
pub struct BudgetProgram { if let Some(key) = keyed_accounts[0].signer_key() {
pub initialized: bool, if &payment.to == key {
pub pending_budget: Option<BudgetExpr>, budget_state.pending_budget = None;
keyed_accounts[1].account.tokens -= payment.tokens;
keyed_accounts[0].account.tokens += payment.tokens;
return Ok(());
}
}
if &payment.to != keyed_accounts[2].unsigned_key() {
trace!("destination missing");
return Err(BudgetError::DestinationMissing);
}
budget_state.pending_budget = None;
keyed_accounts[1].account.tokens -= payment.tokens;
keyed_accounts[2].account.tokens += payment.tokens;
}
Ok(())
}
/// Process a Witness Timestamp. Any payment plans waiting on this timestamp
/// will progress one step.
fn apply_timestamp(
budget_state: &mut BudgetState,
keyed_accounts: &mut [KeyedAccount],
dt: DateTime<Utc>,
) -> Result<(), BudgetError> {
// Check to see if any timelocked transactions can be completed.
let mut final_payment = None;
if let Some(ref mut expr) = budget_state.pending_budget {
let key = match keyed_accounts[0].signer_key() {
None => return Err(BudgetError::UnsignedKey),
Some(key) => key,
};
expr.apply_witness(&Witness::Timestamp(dt), key);
final_payment = expr.final_payment();
}
if let Some(payment) = final_payment {
if &payment.to != keyed_accounts[2].unsigned_key() {
trace!("destination missing");
return Err(BudgetError::DestinationMissing);
}
budget_state.pending_budget = None;
keyed_accounts[1].account.tokens -= payment.tokens;
keyed_accounts[2].account.tokens += payment.tokens;
}
Ok(())
} }
fn apply_debits( fn apply_debits(
keyed_accounts: &mut [KeyedAccount], keyed_accounts: &mut [KeyedAccount],
instruction: &Instruction, instruction: &BudgetInstruction,
) -> Result<(), BudgetError> { ) -> Result<(), BudgetError> {
if !keyed_accounts[0].account.userdata.is_empty() { if !keyed_accounts[0].account.userdata.is_empty() {
trace!("source is pending"); trace!("source is pending");
return Err(BudgetError::SourceIsPendingContract); return Err(BudgetError::SourceIsPendingContract);
} }
match instruction { match instruction {
Instruction::NewBudget(expr) => { BudgetInstruction::InitializeAccount(expr) => {
let expr = expr.clone(); let expr = expr.clone();
if let Some(payment) = expr.final_payment() { if let Some(payment) = expr.final_payment() {
keyed_accounts[1].account.tokens += payment.tokens; keyed_accounts[1].account.tokens += payment.tokens;
Ok(()) Ok(())
} else { } else {
let existing = BudgetProgram::deserialize(&keyed_accounts[1].account.userdata).ok(); let existing = BudgetState::deserialize(&keyed_accounts[1].account.userdata).ok();
if Some(true) == existing.map(|x| x.initialized) { if Some(true) == existing.map(|x| x.initialized) {
trace!("contract already exists"); trace!("contract already exists");
Err(BudgetError::ContractAlreadyExists) Err(BudgetError::ContractAlreadyExists)
} else { } else {
let mut program = BudgetProgram::default(); let mut budget_state = BudgetState::default();
program.pending_budget = Some(expr); budget_state.pending_budget = Some(expr);
keyed_accounts[1].account.tokens += keyed_accounts[0].account.tokens; keyed_accounts[1].account.tokens += keyed_accounts[0].account.tokens;
keyed_accounts[0].account.tokens = 0; keyed_accounts[0].account.tokens = 0;
program.initialized = true; budget_state.initialized = true;
program.serialize(&mut keyed_accounts[1].account.userdata) budget_state.serialize(&mut keyed_accounts[1].account.userdata)
} }
} }
} }
Instruction::ApplyTimestamp(dt) => { BudgetInstruction::ApplyTimestamp(dt) => {
if let Ok(mut program) = BudgetProgram::deserialize(&keyed_accounts[1].account.userdata) if let Ok(mut budget_state) =
BudgetState::deserialize(&keyed_accounts[1].account.userdata)
{ {
if !program.is_pending() { if !budget_state.is_pending() {
Err(BudgetError::ContractNotPending) Err(BudgetError::ContractNotPending)
} else if !program.initialized { } else if !budget_state.initialized {
trace!("contract is uninitialized"); trace!("contract is uninitialized");
Err(BudgetError::UninitializedContract) Err(BudgetError::UninitializedContract)
} else { } else {
trace!("apply timestamp"); trace!("apply timestamp");
program.apply_timestamp(keyed_accounts, *dt)?; apply_timestamp(&mut budget_state, keyed_accounts, *dt)?;
trace!("apply timestamp committed"); trace!("apply timestamp committed");
program.serialize(&mut keyed_accounts[1].account.userdata) budget_state.serialize(&mut keyed_accounts[1].account.userdata)
} }
} else { } else {
Err(BudgetError::UninitializedContract) Err(BudgetError::UninitializedContract)
} }
} }
Instruction::ApplySignature => { BudgetInstruction::ApplySignature => {
if let Ok(mut program) = BudgetProgram::deserialize(&keyed_accounts[1].account.userdata) if let Ok(mut budget_state) =
BudgetState::deserialize(&keyed_accounts[1].account.userdata)
{ {
if !program.is_pending() { if !budget_state.is_pending() {
Err(BudgetError::ContractNotPending) Err(BudgetError::ContractNotPending)
} else if !program.initialized { } else if !budget_state.initialized {
trace!("contract is uninitialized"); trace!("contract is uninitialized");
Err(BudgetError::UninitializedContract) Err(BudgetError::UninitializedContract)
} else { } else {
trace!("apply signature"); trace!("apply signature");
program.apply_signature(keyed_accounts)?; apply_signature(&mut budget_state, keyed_accounts)?;
trace!("apply signature committed"); trace!("apply signature committed");
program.serialize(&mut keyed_accounts[1].account.userdata) budget_state.serialize(&mut keyed_accounts[1].account.userdata)
} }
} else { } else {
Err(BudgetError::UninitializedContract) Err(BudgetError::UninitializedContract)
@ -115,115 +161,9 @@ pub fn process_instruction(
} }
} }
impl BudgetProgram {
fn is_pending(&self) -> bool {
self.pending_budget != None
}
/// Process a Witness Signature. Any payment plans waiting on this signature
/// will progress one step.
fn apply_signature(&mut self, keyed_accounts: &mut [KeyedAccount]) -> Result<(), BudgetError> {
let mut final_payment = None;
if let Some(ref mut expr) = self.pending_budget {
let key = match keyed_accounts[0].signer_key() {
None => return Err(BudgetError::UnsignedKey),
Some(key) => key,
};
expr.apply_witness(&Witness::Signature, key);
final_payment = expr.final_payment();
}
if let Some(payment) = final_payment {
if let Some(key) = keyed_accounts[0].signer_key() {
if &payment.to == key {
self.pending_budget = None;
keyed_accounts[1].account.tokens -= payment.tokens;
keyed_accounts[0].account.tokens += payment.tokens;
return Ok(());
}
}
if &payment.to != keyed_accounts[2].unsigned_key() {
trace!("destination missing");
return Err(BudgetError::DestinationMissing);
}
self.pending_budget = None;
keyed_accounts[1].account.tokens -= payment.tokens;
keyed_accounts[2].account.tokens += payment.tokens;
}
Ok(())
}
/// Process a Witness Timestamp. Any payment plans waiting on this timestamp
/// will progress one step.
fn apply_timestamp(
&mut self,
keyed_accounts: &mut [KeyedAccount],
dt: DateTime<Utc>,
) -> Result<(), BudgetError> {
// Check to see if any timelocked transactions can be completed.
let mut final_payment = None;
if let Some(ref mut expr) = self.pending_budget {
let key = match keyed_accounts[0].signer_key() {
None => return Err(BudgetError::UnsignedKey),
Some(key) => key,
};
expr.apply_witness(&Witness::Timestamp(dt), key);
final_payment = expr.final_payment();
}
if let Some(payment) = final_payment {
if &payment.to != keyed_accounts[2].unsigned_key() {
trace!("destination missing");
return Err(BudgetError::DestinationMissing);
}
self.pending_budget = None;
keyed_accounts[1].account.tokens -= payment.tokens;
keyed_accounts[2].account.tokens += payment.tokens;
}
Ok(())
}
fn serialize(&self, output: &mut [u8]) -> Result<(), BudgetError> {
let len = serialized_size(self).unwrap() as u64;
if output.len() < len as usize {
warn!(
"{} bytes required to serialize, only have {} bytes",
len,
output.len()
);
return Err(BudgetError::UserdataTooSmall);
}
{
let writer = io::BufWriter::new(&mut output[..8]);
serialize_into(writer, &len).unwrap();
}
{
let writer = io::BufWriter::new(&mut output[8..8 + len as usize]);
serialize_into(writer, self).unwrap();
}
Ok(())
}
pub fn deserialize(input: &[u8]) -> bincode::Result<Self> {
if input.len() < 8 {
return Err(Box::new(bincode::ErrorKind::SizeLimit));
}
let len: u64 = deserialize(&input[..8]).unwrap();
if len < 2 {
return Err(Box::new(bincode::ErrorKind::SizeLimit));
}
if input.len() < 8 + len as usize {
return Err(Box::new(bincode::ErrorKind::SizeLimit));
}
deserialize(&input[8..8 + len as usize])
}
}
#[cfg(test)] #[cfg(test)]
mod test { mod test {
use super::*; use super::*;
use bincode::serialize;
use solana_budget_api::budget_transaction::BudgetTransaction; use solana_budget_api::budget_transaction::BudgetTransaction;
use solana_budget_api::id; use solana_budget_api::id;
use solana_sdk::account::Account; use solana_sdk::account::Account;
@ -256,26 +196,6 @@ mod test {
super::process_instruction(&mut keyed_accounts, &userdata) super::process_instruction(&mut keyed_accounts, &userdata)
} }
#[test] #[test]
fn test_serializer() {
let mut a = Account::new(0, 512, id());
let b = BudgetProgram::default();
b.serialize(&mut a.userdata).unwrap();
let buf = serialize(&b).unwrap();
assert_eq!(a.userdata[8..8 + buf.len()], buf[0..]);
let c = BudgetProgram::deserialize(&a.userdata).unwrap();
assert_eq!(b, c);
}
#[test]
fn test_serializer_userdata_too_small() {
let mut a = Account::new(0, 1, id());
let b = BudgetProgram::default();
assert_eq!(
b.serialize(&mut a.userdata),
Err(BudgetError::UserdataTooSmall)
);
}
#[test]
fn test_invalid_instruction() { fn test_invalid_instruction() {
let mut accounts = vec![Account::new(1, 0, id()), Account::new(0, 512, id())]; let mut accounts = vec![Account::new(1, 0, id()), Account::new(0, 512, id())];
let from = Keypair::new(); let from = Keypair::new();
@ -300,7 +220,7 @@ mod test {
Account::new(0, 0, id()), Account::new(0, 0, id()),
]; ];
// Initialize BudgetProgram // Initialize BudgetState
let from = Keypair::new(); let from = Keypair::new();
let contract = Keypair::new().pubkey(); let contract = Keypair::new().pubkey();
let to = Keypair::new().pubkey(); let to = Keypair::new().pubkey();
@ -339,7 +259,7 @@ mod test {
Account::new(0, 0, id()), Account::new(0, 0, id()),
]; ];
// Initialize BudgetProgram // Initialize BudgetState
let from = Keypair::new(); let from = Keypair::new();
let contract = Keypair::new().pubkey(); let contract = Keypair::new().pubkey();
let to = Keypair::new().pubkey(); let to = Keypair::new().pubkey();
@ -399,8 +319,8 @@ mod test {
process_transaction(&tx, &mut accounts).unwrap(); process_transaction(&tx, &mut accounts).unwrap();
assert_eq!(accounts[from_account].tokens, 0); assert_eq!(accounts[from_account].tokens, 0);
assert_eq!(accounts[contract_account].tokens, 1); assert_eq!(accounts[contract_account].tokens, 1);
let program = BudgetProgram::deserialize(&accounts[contract_account].userdata).unwrap(); let budget_state = BudgetState::deserialize(&accounts[contract_account].userdata).unwrap();
assert!(program.is_pending()); assert!(budget_state.is_pending());
// Attack! Try to payout to a rando key // Attack! Try to payout to a rando key
let tx = BudgetTransaction::new_timestamp( let tx = BudgetTransaction::new_timestamp(
@ -418,8 +338,8 @@ mod test {
assert_eq!(accounts[contract_account].tokens, 1); assert_eq!(accounts[contract_account].tokens, 1);
assert_eq!(accounts[to_account].tokens, 0); assert_eq!(accounts[to_account].tokens, 0);
let program = BudgetProgram::deserialize(&accounts[contract_account].userdata).unwrap(); let budget_state = BudgetState::deserialize(&accounts[contract_account].userdata).unwrap();
assert!(program.is_pending()); assert!(budget_state.is_pending());
// 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.
@ -435,8 +355,8 @@ mod test {
assert_eq!(accounts[contract_account].tokens, 0); assert_eq!(accounts[contract_account].tokens, 0);
assert_eq!(accounts[to_account].tokens, 1); assert_eq!(accounts[to_account].tokens, 1);
let program = BudgetProgram::deserialize(&accounts[contract_account].userdata).unwrap(); let budget_state = BudgetState::deserialize(&accounts[contract_account].userdata).unwrap();
assert!(!program.is_pending()); assert!(!budget_state.is_pending());
// try to replay the timestamp contract // try to replay the timestamp contract
assert_eq!( assert_eq!(
@ -474,8 +394,8 @@ mod test {
process_transaction(&tx, &mut accounts).unwrap(); process_transaction(&tx, &mut accounts).unwrap();
assert_eq!(accounts[from_account].tokens, 0); assert_eq!(accounts[from_account].tokens, 0);
assert_eq!(accounts[contract_account].tokens, 1); assert_eq!(accounts[contract_account].tokens, 1);
let program = BudgetProgram::deserialize(&accounts[contract_account].userdata).unwrap(); let budget_state = BudgetState::deserialize(&accounts[contract_account].userdata).unwrap();
assert!(program.is_pending()); assert!(budget_state.is_pending());
// Attack! try to put the tokens into the wrong account with cancel // Attack! try to put the tokens into the wrong account with cancel
let tx = let tx =
@ -539,7 +459,7 @@ mod test {
); );
assert!(process_transaction(&tx, &mut accounts).is_err()); assert!(process_transaction(&tx, &mut accounts).is_err());
assert!(BudgetProgram::deserialize(&accounts[1].userdata).is_err()); assert!(BudgetState::deserialize(&accounts[1].userdata).is_err());
let tx = BudgetTransaction::new_timestamp( let tx = BudgetTransaction::new_timestamp(
&from, &from,
@ -549,7 +469,7 @@ mod test {
Hash::default(), Hash::default(),
); );
assert!(process_transaction(&tx, &mut accounts).is_err()); assert!(process_transaction(&tx, &mut accounts).is_err());
assert!(BudgetProgram::deserialize(&accounts[1].userdata).is_err()); assert!(BudgetState::deserialize(&accounts[1].userdata).is_err());
// Success if there was no panic... // Success if there was no panic...
} }

98
programs/budget/src/budget_state.rs Normal file
View File

@ -0,0 +1,98 @@
//! budget state
use bincode::{self, deserialize, serialize_into, serialized_size};
use log::*;
use serde_derive::{Deserialize, Serialize};
use solana_budget_api::budget_expr::BudgetExpr;
use std::io;
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
pub enum BudgetError {
InsufficientFunds,
ContractAlreadyExists,
ContractNotPending,
SourceIsPendingContract,
UninitializedContract,
NegativeTokens,
DestinationMissing,
FailedWitness,
UserdataTooSmall,
UserdataDeserializeFailure,
UnsignedKey,
}
#[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq)]
pub struct BudgetState {
pub initialized: bool,
pub pending_budget: Option<BudgetExpr>,
}
impl BudgetState {
pub fn is_pending(&self) -> bool {
self.pending_budget.is_some()
}
pub fn serialize(&self, output: &mut [u8]) -> Result<(), BudgetError> {
let len = serialized_size(self).unwrap() as u64;
if output.len() < len as usize {
warn!(
"{} bytes required to serialize, only have {} bytes",
len,
output.len()
);
return Err(BudgetError::UserdataTooSmall);
}
{
let writer = io::BufWriter::new(&mut output[..8]);
serialize_into(writer, &len).unwrap();
}
{
let writer = io::BufWriter::new(&mut output[8..8 + len as usize]);
serialize_into(writer, self).unwrap();
}
Ok(())
}
pub fn deserialize(input: &[u8]) -> bincode::Result<Self> {
if input.len() < 8 {
return Err(Box::new(bincode::ErrorKind::SizeLimit));
}
let len: u64 = deserialize(&input[..8]).unwrap();
if len < 2 {
return Err(Box::new(bincode::ErrorKind::SizeLimit));
}
if input.len() < 8 + len as usize {
return Err(Box::new(bincode::ErrorKind::SizeLimit));
}
deserialize(&input[8..8 + len as usize])
}
}
#[cfg(test)]
mod test {
use super::*;
use bincode::serialize;
use solana_budget_api::id;
use solana_sdk::account::Account;
#[test]
fn test_serializer() {
let mut a = Account::new(0, 512, id());
let b = BudgetState::default();
b.serialize(&mut a.userdata).unwrap();
let buf = serialize(&b).unwrap();
assert_eq!(a.userdata[8..8 + buf.len()], buf[0..]);
let c = BudgetState::deserialize(&a.userdata).unwrap();
assert_eq!(b, c);
}
#[test]
fn test_serializer_userdata_too_small() {
let mut a = Account::new(0, 1, id());
let b = BudgetState::default();
assert_eq!(
b.serialize(&mut a.userdata),
Err(BudgetError::UserdataTooSmall)
);
}
}

1
programs/budget/src/lib.rs
View File

@ -1,4 +1,5 @@
mod budget_program; mod budget_program;
mod budget_state;
use crate::budget_program::process_instruction; use crate::budget_program::process_instruction;
use log::*; use log::*;

16
programs/budget_api/src/budget_instruction.rs
View File

@ -15,20 +15,24 @@ pub struct Contract {
/// An instruction to progress the smart contract. /// An instruction to progress the smart contract.
#[derive(Serialize, Deserialize, Debug, PartialEq, Eq, Clone)] #[derive(Serialize, Deserialize, Debug, PartialEq, Eq, Clone)]
pub enum Instruction { pub enum BudgetInstruction {
/// Declare and instantiate `BudgetExpr`. /// Declare and instantiate `BudgetExpr`.
NewBudget(BudgetExpr), InitializeAccount(BudgetExpr),
/// Tell a payment plan acknowledge the given `DateTime` has past. /// Tell a payment plan acknowledge the given `DateTime` has past.
ApplyTimestamp(DateTime<Utc>), ApplyTimestamp(DateTime<Utc>),
/// Tell the budget that the `NewBudget` with `Signature` has been /// Tell the budget that the `InitializeAccount` with `Signature` has been
/// signed by the containing transaction's `Pubkey`. /// signed by the containing transaction's `Pubkey`.
ApplySignature, ApplySignature,
} }
impl Instruction { impl BudgetInstruction {
pub fn new_budget(contract: Pubkey, expr: BudgetExpr) -> BuilderInstruction { pub fn new_initialize_account(contract: Pubkey, expr: BudgetExpr) -> BuilderInstruction {
BuilderInstruction::new(id(), &Instruction::NewBudget(expr), vec![(contract, false)]) BuilderInstruction::new(
id(),
&BudgetInstruction::InitializeAccount(expr),
vec![(contract, false)],
)
} }
} }

26
programs/budget_api/src/budget_transaction.rs
View File

@ -1,7 +1,7 @@
//! The `budget_transaction` module provides functionality for creating Budget transactions. //! The `budget_transaction` module provides functionality for creating Budget transactions.
use crate::budget_expr::{BudgetExpr, Condition}; use crate::budget_expr::{BudgetExpr, Condition};
use crate::budget_instruction::Instruction; use crate::budget_instruction::BudgetInstruction;
use crate::id; use crate::id;
use bincode::deserialize; use bincode::deserialize;
use chrono::prelude::*; use chrono::prelude::*;
@ -28,7 +28,7 @@ impl BudgetTransaction {
let payment = BudgetExpr::new_payment(tokens - fee, to); let payment = BudgetExpr::new_payment(tokens - fee, to);
TransactionBuilder::new(fee) TransactionBuilder::new(fee)
.push(SystemInstruction::new_move(from, contract, tokens)) .push(SystemInstruction::new_move(from, contract, tokens))
.push(Instruction::new_budget(contract, payment)) .push(BudgetInstruction::new_initialize_account(contract, payment))
.sign(&[from_keypair], recent_blockhash) .sign(&[from_keypair], recent_blockhash)
} }
@ -51,7 +51,7 @@ impl BudgetTransaction {
dt: DateTime<Utc>, dt: DateTime<Utc>,
recent_blockhash: Hash, recent_blockhash: Hash,
) -> Transaction { ) -> Transaction {
let instruction = Instruction::ApplyTimestamp(dt); let instruction = BudgetInstruction::ApplyTimestamp(dt);
Transaction::new( Transaction::new(
from_keypair, from_keypair,
&[contract, to], &[contract, to],
@ -69,7 +69,7 @@ impl BudgetTransaction {
to: Pubkey, to: Pubkey,
recent_blockhash: Hash, recent_blockhash: Hash,
) -> Transaction { ) -> Transaction {
let instruction = Instruction::ApplySignature; let instruction = BudgetInstruction::ApplySignature;
let mut keys = vec![contract]; let mut keys = vec![contract];
if from_keypair.pubkey() != to { if from_keypair.pubkey() != to {
keys.push(to); keys.push(to);
@ -105,7 +105,7 @@ impl BudgetTransaction {
Box::new(BudgetExpr::new_payment(tokens, to)), Box::new(BudgetExpr::new_payment(tokens, to)),
) )
}; };
let instruction = Instruction::NewBudget(expr); let instruction = BudgetInstruction::InitializeAccount(expr);
Transaction::new( Transaction::new(
from_keypair, from_keypair,
&[contract], &[contract],
@ -142,7 +142,7 @@ impl BudgetTransaction {
Box::new(BudgetExpr::new_payment(tokens, to)), Box::new(BudgetExpr::new_payment(tokens, to)),
) )
}; };
let instruction = Instruction::NewBudget(expr); let instruction = BudgetInstruction::InitializeAccount(expr);
Transaction::new( Transaction::new(
from_keypair, from_keypair,
&[contract], &[contract],
@ -157,14 +157,16 @@ impl BudgetTransaction {
deserialize(&tx.userdata(index)).ok() deserialize(&tx.userdata(index)).ok()
} }
pub fn instruction(tx: &Transaction, index: usize) -> Option<Instruction> { pub fn instruction(tx: &Transaction, index: usize) -> Option<BudgetInstruction> {
deserialize(&tx.userdata(index)).ok() deserialize(&tx.userdata(index)).ok()
} }
/// Verify only the payment plan. /// Verify only the payment plan.
pub fn verify_plan(tx: &Transaction) -> bool { pub fn verify_plan(tx: &Transaction) -> bool {
if let Some(SystemInstruction::Move { tokens }) = Self::system_instruction(tx, 0) { if let Some(SystemInstruction::Move { tokens }) = Self::system_instruction(tx, 0) {
if let Some(Instruction::NewBudget(expr)) = BudgetTransaction::instruction(&tx, 1) { if let Some(BudgetInstruction::InitializeAccount(expr)) =
BudgetTransaction::instruction(&tx, 1)
{
if !(tx.fee <= tokens && expr.verify(tokens - tx.fee)) { if !(tx.fee <= tokens && expr.verify(tokens - tx.fee)) {
return false; return false;
} }
@ -227,7 +229,7 @@ mod tests {
if let SystemInstruction::Move { ref mut tokens } = system_instruction { if let SystemInstruction::Move { ref mut tokens } = system_instruction {
*tokens = 1_000_000; // <-- attack, part 1! *tokens = 1_000_000; // <-- attack, part 1!
let mut instruction = BudgetTransaction::instruction(&tx, 1).unwrap(); let mut instruction = BudgetTransaction::instruction(&tx, 1).unwrap();
if let Instruction::NewBudget(ref mut expr) = instruction { if let BudgetInstruction::InitializeAccount(ref mut expr) = instruction {
if let BudgetExpr::Pay(ref mut payment) = expr { if let BudgetExpr::Pay(ref mut payment) = expr {
payment.tokens = *tokens; // <-- attack, part 2! payment.tokens = *tokens; // <-- attack, part 2!
} }
@ -248,7 +250,7 @@ mod tests {
let zero = Hash::default(); let zero = Hash::default();
let mut tx = BudgetTransaction::new(&keypair0, pubkey1, 42, zero); let mut tx = BudgetTransaction::new(&keypair0, pubkey1, 42, zero);
let mut instruction = BudgetTransaction::instruction(&tx, 1); let mut instruction = BudgetTransaction::instruction(&tx, 1);
if let Some(Instruction::NewBudget(ref mut expr)) = instruction { if let Some(BudgetInstruction::InitializeAccount(ref mut expr)) = instruction {
if let BudgetExpr::Pay(ref mut payment) = expr { if let BudgetExpr::Pay(ref mut payment) = expr {
payment.to = thief_keypair.pubkey(); // <-- attack! payment.to = thief_keypair.pubkey(); // <-- attack!
} }
@ -265,7 +267,7 @@ mod tests {
let zero = Hash::default(); let zero = Hash::default();
let mut tx = BudgetTransaction::new(&keypair0, keypair1.pubkey(), 1, zero); let mut tx = BudgetTransaction::new(&keypair0, keypair1.pubkey(), 1, zero);
let mut instruction = BudgetTransaction::instruction(&tx, 1).unwrap(); let mut instruction = BudgetTransaction::instruction(&tx, 1).unwrap();
if let Instruction::NewBudget(ref mut expr) = instruction { if let BudgetInstruction::InitializeAccount(ref mut expr) = instruction {
if let BudgetExpr::Pay(ref mut payment) = expr { if let BudgetExpr::Pay(ref mut payment) = expr {
payment.tokens = 2; // <-- attack! payment.tokens = 2; // <-- attack!
} }
@ -275,7 +277,7 @@ mod tests {
// Also, ensure all branchs of the plan spend all tokens // Also, ensure all branchs of the plan spend all tokens
let mut instruction = BudgetTransaction::instruction(&tx, 1).unwrap(); let mut instruction = BudgetTransaction::instruction(&tx, 1).unwrap();
if let Instruction::NewBudget(ref mut expr) = instruction { if let BudgetInstruction::InitializeAccount(ref mut expr) = instruction {
if let BudgetExpr::Pay(ref mut payment) = expr { if let BudgetExpr::Pay(ref mut payment) = expr {
payment.tokens = 0; // <-- whoops! payment.tokens = 0; // <-- whoops!
} }