#[cfg(feature = "abi")]
use borsh::BorshSchema;
use std::cell::RefCell;
#[cfg(feature = "abi")]
use std::collections::BTreeMap;
use std::io::{Error, Write};
use std::num::NonZeroU128;
use std::rc::Rc;
use crate::env::migrate_to_allowance;
use crate::{AccountId, Gas, GasWeight, NearToken, PromiseIndex, PublicKey};
#[derive(Clone, Copy)]
pub enum Allowance {
Unlimited,
Limited(NonZeroU128),
}
impl Allowance {
pub fn unlimited() -> Allowance {
Allowance::Unlimited
}
pub fn limited(balance: NearToken) -> Option<Allowance> {
NonZeroU128::new(balance.as_yoctonear()).map(Allowance::Limited)
}
}
enum PromiseAction {
CreateAccount,
DeployContract {
code: Vec<u8>,
},
FunctionCall {
function_name: String,
arguments: Vec<u8>,
amount: NearToken,
gas: Gas,
},
FunctionCallWeight {
function_name: String,
arguments: Vec<u8>,
amount: NearToken,
gas: Gas,
weight: GasWeight,
},
Transfer {
amount: NearToken,
},
Stake {
amount: NearToken,
public_key: PublicKey,
},
AddFullAccessKey {
public_key: PublicKey,
nonce: u64,
},
AddAccessKey {
public_key: PublicKey,
allowance: Allowance,
receiver_id: AccountId,
function_names: String,
nonce: u64,
},
DeleteKey {
public_key: PublicKey,
},
DeleteAccount {
beneficiary_id: AccountId,
},
}
impl PromiseAction {
pub fn add(&self, promise_index: PromiseIndex) {
use PromiseAction::*;
match self {
CreateAccount => crate::env::promise_batch_action_create_account(promise_index),
DeployContract { code } => {
crate::env::promise_batch_action_deploy_contract(promise_index, code)
}
FunctionCall { function_name, arguments, amount, gas } => {
crate::env::promise_batch_action_function_call(
promise_index,
function_name,
arguments,
*amount,
*gas,
)
}
FunctionCallWeight { function_name, arguments, amount, gas, weight } => {
crate::env::promise_batch_action_function_call_weight(
promise_index,
function_name,
arguments,
*amount,
*gas,
GasWeight(weight.0),
)
}
Transfer { amount } => {
crate::env::promise_batch_action_transfer(promise_index, *amount)
}
Stake { amount, public_key } => {
crate::env::promise_batch_action_stake(promise_index, *amount, public_key)
}
AddFullAccessKey { public_key, nonce } => {
crate::env::promise_batch_action_add_key_with_full_access(
promise_index,
public_key,
*nonce,
)
}
AddAccessKey { public_key, allowance, receiver_id, function_names, nonce } => {
crate::env::promise_batch_action_add_key_allowance_with_function_call(
promise_index,
public_key,
*nonce,
*allowance,
receiver_id,
function_names,
)
}
DeleteKey { public_key } => {
crate::env::promise_batch_action_delete_key(promise_index, public_key)
}
DeleteAccount { beneficiary_id } => {
crate::env::promise_batch_action_delete_account(promise_index, beneficiary_id)
}
}
}
}
struct PromiseSingle {
pub account_id: AccountId,
pub actions: RefCell<Vec<PromiseAction>>,
pub after: RefCell<Option<Promise>>,
pub promise_index: RefCell<Option<PromiseIndex>>,
}
impl PromiseSingle {
pub fn construct_recursively(&self) -> PromiseIndex {
let mut promise_lock = self.promise_index.borrow_mut();
if let Some(res) = promise_lock.as_ref() {
return *res;
}
let promise_index = if let Some(after) = self.after.borrow().as_ref() {
crate::env::promise_batch_then(after.construct_recursively(), &self.account_id)
} else {
crate::env::promise_batch_create(&self.account_id)
};
let actions_lock = self.actions.borrow();
for action in actions_lock.iter() {
action.add(promise_index);
}
*promise_lock = Some(promise_index);
promise_index
}
}
pub struct PromiseJoint {
pub promise_a: Promise,
pub promise_b: Promise,
pub promise_index: RefCell<Option<PromiseIndex>>,
}
impl PromiseJoint {
pub fn construct_recursively(&self) -> PromiseIndex {
let mut promise_lock = self.promise_index.borrow_mut();
if let Some(res) = promise_lock.as_ref() {
return *res;
}
let res = crate::env::promise_and(&[
self.promise_a.construct_recursively(),
self.promise_b.construct_recursively(),
]);
*promise_lock = Some(res);
res
}
}
pub struct Promise {
subtype: PromiseSubtype,
should_return: RefCell<bool>,
}
#[cfg(feature = "abi")]
impl BorshSchema for Promise {
fn add_definitions_recursively(
definitions: &mut BTreeMap<borsh::schema::Declaration, borsh::schema::Definition>,
) {
<()>::add_definitions_recursively(definitions);
}
fn declaration() -> borsh::schema::Declaration {
<()>::declaration()
}
}
#[derive(Clone)]
enum PromiseSubtype {
Single(Rc<PromiseSingle>),
Joint(Rc<PromiseJoint>),
}
impl Promise {
pub fn new(account_id: AccountId) -> Self {
Self {
subtype: PromiseSubtype::Single(Rc::new(PromiseSingle {
account_id,
actions: RefCell::new(vec![]),
after: RefCell::new(None),
promise_index: RefCell::new(None),
})),
should_return: RefCell::new(false),
}
}
fn add_action(self, action: PromiseAction) -> Self {
match &self.subtype {
PromiseSubtype::Single(x) => x.actions.borrow_mut().push(action),
PromiseSubtype::Joint(_) => {
crate::env::panic_str("Cannot add action to a joint promise.")
}
}
self
}
pub fn create_account(self) -> Self {
self.add_action(PromiseAction::CreateAccount)
}
pub fn deploy_contract(self, code: Vec<u8>) -> Self {
self.add_action(PromiseAction::DeployContract { code })
}
pub fn function_call(
self,
function_name: String,
arguments: Vec<u8>,
amount: NearToken,
gas: Gas,
) -> Self {
self.add_action(PromiseAction::FunctionCall { function_name, arguments, amount, gas })
}
pub fn function_call_weight(
self,
function_name: String,
arguments: Vec<u8>,
amount: NearToken,
gas: Gas,
weight: GasWeight,
) -> Self {
self.add_action(PromiseAction::FunctionCallWeight {
function_name,
arguments,
amount,
gas,
weight,
})
}
pub fn transfer(self, amount: NearToken) -> Self {
self.add_action(PromiseAction::Transfer { amount })
}
pub fn stake(self, amount: NearToken, public_key: PublicKey) -> Self {
self.add_action(PromiseAction::Stake { amount, public_key })
}
pub fn add_full_access_key(self, public_key: PublicKey) -> Self {
self.add_full_access_key_with_nonce(public_key, 0)
}
pub fn add_full_access_key_with_nonce(self, public_key: PublicKey, nonce: u64) -> Self {
self.add_action(PromiseAction::AddFullAccessKey { public_key, nonce })
}
pub fn add_access_key_allowance(
self,
public_key: PublicKey,
allowance: Allowance,
receiver_id: AccountId,
function_names: String,
) -> Self {
self.add_access_key_allowance_with_nonce(
public_key,
allowance,
receiver_id,
function_names,
0,
)
}
#[deprecated(since = "5.0.0", note = "Use add_access_key_allowance instead")]
pub fn add_access_key(
self,
public_key: PublicKey,
allowance: NearToken,
receiver_id: AccountId,
function_names: String,
) -> Self {
let allowance = migrate_to_allowance(allowance);
self.add_access_key_allowance(public_key, allowance, receiver_id, function_names)
}
pub fn add_access_key_allowance_with_nonce(
self,
public_key: PublicKey,
allowance: Allowance,
receiver_id: AccountId,
function_names: String,
nonce: u64,
) -> Self {
self.add_action(PromiseAction::AddAccessKey {
public_key,
allowance,
receiver_id,
function_names,
nonce,
})
}
#[deprecated(since = "5.0.0", note = "Use add_access_key_allowance_with_nonce instead")]
pub fn add_access_key_with_nonce(
self,
public_key: PublicKey,
allowance: NearToken,
receiver_id: AccountId,
function_names: String,
nonce: u64,
) -> Self {
let allowance = migrate_to_allowance(allowance);
self.add_access_key_allowance_with_nonce(
public_key,
allowance,
receiver_id,
function_names,
nonce,
)
}
pub fn delete_key(self, public_key: PublicKey) -> Self {
self.add_action(PromiseAction::DeleteKey { public_key })
}
pub fn delete_account(self, beneficiary_id: AccountId) -> Self {
self.add_action(PromiseAction::DeleteAccount { beneficiary_id })
}
pub fn and(self, other: Promise) -> Promise {
Promise {
subtype: PromiseSubtype::Joint(Rc::new(PromiseJoint {
promise_a: self,
promise_b: other,
promise_index: RefCell::new(None),
})),
should_return: RefCell::new(false),
}
}
pub fn then(self, mut other: Promise) -> Promise {
match &mut other.subtype {
PromiseSubtype::Single(x) => {
let mut after = x.after.borrow_mut();
if after.is_some() {
crate::env::panic_str(
"Cannot callback promise which is already scheduled after another",
);
}
*after = Some(self)
}
PromiseSubtype::Joint(_) => crate::env::panic_str("Cannot callback joint promise."),
}
other
}
#[allow(clippy::wrong_self_convention)]
pub fn as_return(self) -> Self {
*self.should_return.borrow_mut() = true;
self
}
fn construct_recursively(&self) -> PromiseIndex {
let res = match &self.subtype {
PromiseSubtype::Single(x) => x.construct_recursively(),
PromiseSubtype::Joint(x) => x.construct_recursively(),
};
if *self.should_return.borrow() {
crate::env::promise_return(res);
}
res
}
}
impl Drop for Promise {
fn drop(&mut self) {
self.construct_recursively();
}
}
impl serde::Serialize for Promise {
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: serde::Serializer,
{
*self.should_return.borrow_mut() = true;
serializer.serialize_unit()
}
}
impl borsh::BorshSerialize for Promise {
fn serialize<W: Write>(&self, _writer: &mut W) -> Result<(), Error> {
*self.should_return.borrow_mut() = true;
Ok(())
}
}
#[cfg(feature = "abi")]
impl schemars::JsonSchema for Promise {
fn schema_name() -> String {
"Promise".to_string()
}
fn json_schema(_gen: &mut schemars::gen::SchemaGenerator) -> schemars::schema::Schema {
schemars::schema::Schema::Bool(true)
}
}
#[derive(serde::Serialize)]
#[serde(untagged)]
pub enum PromiseOrValue<T> {
Promise(Promise),
Value(T),
}
#[cfg(feature = "abi")]
impl<T> BorshSchema for PromiseOrValue<T>
where
T: BorshSchema,
{
fn add_definitions_recursively(
definitions: &mut BTreeMap<borsh::schema::Declaration, borsh::schema::Definition>,
) {
T::add_definitions_recursively(definitions);
}
fn declaration() -> borsh::schema::Declaration {
T::declaration()
}
}
impl<T> From<Promise> for PromiseOrValue<T> {
fn from(promise: Promise) -> Self {
PromiseOrValue::Promise(promise)
}
}
impl<T: borsh::BorshSerialize> borsh::BorshSerialize for PromiseOrValue<T> {
fn serialize<W: Write>(&self, writer: &mut W) -> Result<(), Error> {
match self {
PromiseOrValue::Value(x) => x.serialize(writer),
PromiseOrValue::Promise(p) => p.serialize(writer),
}
}
}
#[cfg(feature = "abi")]
impl<T: schemars::JsonSchema> schemars::JsonSchema for PromiseOrValue<T> {
fn schema_name() -> String {
format!("PromiseOrValue{}", T::schema_name())
}
fn json_schema(gen: &mut schemars::gen::SchemaGenerator) -> schemars::schema::Schema {
T::json_schema(gen)
}
}
#[cfg(not(target_arch = "wasm32"))]
#[cfg(test)]
mod tests {
use crate::mock::MockAction;
use crate::test_utils::get_created_receipts;
use crate::test_utils::test_env::{alice, bob};
use crate::{
test_utils::VMContextBuilder, testing_env, AccountId, Allowance, NearToken, Promise,
PublicKey,
};
fn pk() -> PublicKey {
"ed25519:6E8sCci9badyRkXb3JoRpBj5p8C6Tw41ELDZoiihKEtp".parse().unwrap()
}
fn get_actions() -> std::vec::IntoIter<MockAction> {
let receipts = get_created_receipts();
let first_receipt = receipts.into_iter().next().unwrap();
first_receipt.actions.into_iter()
}
fn has_add_key_with_full_access(public_key: PublicKey, nonce: Option<u64>) -> bool {
let public_key = near_crypto::PublicKey::try_from(public_key).unwrap();
get_actions().any(|el| {
matches!(
el,
MockAction::AddKeyWithFullAccess { public_key: p, nonce: n, receipt_index: _, }
if p == public_key
&& (nonce.is_none() || Some(n) == nonce)
)
})
}
fn has_add_key_with_function_call(
public_key: PublicKey,
allowance: u128,
receiver_id: AccountId,
function_names: String,
nonce: Option<u64>,
) -> bool {
let public_key = near_crypto::PublicKey::try_from(public_key).unwrap();
get_actions().any(|el| {
matches!(
el,
MockAction::AddKeyWithFunctionCall {
public_key: p,
allowance: a,
receiver_id: r,
method_names,
nonce: n,
receipt_index: _,
}
if p == public_key
&& a.unwrap() == NearToken::from_yoctonear(allowance)
&& r == receiver_id
&& method_names.clone() == function_names.split(',').collect::<Vec<_>>()
&& (nonce.is_none() || Some(n) == nonce)
)
})
}
#[test]
fn test_add_full_access_key() {
testing_env!(VMContextBuilder::new().signer_account_id(alice()).build());
let public_key: PublicKey = pk();
{
Promise::new(alice()).create_account().add_full_access_key(public_key.clone());
}
assert!(has_add_key_with_full_access(public_key, None));
}
#[test]
fn test_add_full_access_key_with_nonce() {
testing_env!(VMContextBuilder::new().signer_account_id(alice()).build());
let public_key: PublicKey = pk();
let nonce = 42;
{
Promise::new(alice())
.create_account()
.add_full_access_key_with_nonce(public_key.clone(), nonce);
}
assert!(has_add_key_with_full_access(public_key, Some(nonce)));
}
#[test]
fn test_add_access_key_allowance() {
testing_env!(VMContextBuilder::new().signer_account_id(alice()).build());
let public_key: PublicKey = pk();
let allowance = 100;
let receiver_id = bob();
let function_names = "method_a,method_b".to_string();
{
Promise::new(alice()).create_account().add_access_key_allowance(
public_key.clone(),
Allowance::Limited(allowance.try_into().unwrap()),
receiver_id.clone(),
function_names.clone(),
);
}
assert!(has_add_key_with_function_call(
public_key,
allowance,
receiver_id,
function_names,
None
));
}
#[test]
fn test_add_access_key() {
testing_env!(VMContextBuilder::new().signer_account_id(alice()).build());
let public_key: PublicKey = pk();
let allowance = NearToken::from_yoctonear(100);
let receiver_id = bob();
let function_names = "method_a,method_b".to_string();
{
#[allow(deprecated)]
Promise::new(alice()).create_account().add_access_key(
public_key.clone(),
allowance,
receiver_id.clone(),
function_names.clone(),
);
}
assert!(has_add_key_with_function_call(
public_key,
allowance.as_yoctonear(),
receiver_id,
function_names,
None
));
}
#[test]
fn test_add_access_key_allowance_with_nonce() {
testing_env!(VMContextBuilder::new().signer_account_id(alice()).build());
let public_key: PublicKey = pk();
let allowance = 100;
let receiver_id = bob();
let function_names = "method_a,method_b".to_string();
let nonce = 42;
{
Promise::new(alice()).create_account().add_access_key_allowance_with_nonce(
public_key.clone(),
Allowance::Limited(allowance.try_into().unwrap()),
receiver_id.clone(),
function_names.clone(),
nonce,
);
}
assert!(has_add_key_with_function_call(
public_key,
allowance,
receiver_id,
function_names,
Some(nonce)
));
}
#[test]
fn test_add_access_key_with_nonce() {
testing_env!(VMContextBuilder::new().signer_account_id(alice()).build());
let public_key: PublicKey = pk();
let allowance = NearToken::from_yoctonear(100);
let receiver_id = bob();
let function_names = "method_a,method_b".to_string();
let nonce = 42;
{
#[allow(deprecated)]
Promise::new(alice()).create_account().add_access_key_with_nonce(
public_key.clone(),
allowance,
receiver_id.clone(),
function_names.clone(),
nonce,
);
}
assert!(has_add_key_with_function_call(
public_key,
allowance.as_yoctonear(),
receiver_id,
function_names,
Some(nonce)
));
}
#[test]
fn test_delete_key() {
testing_env!(VMContextBuilder::new().signer_account_id(alice()).build());
let public_key: PublicKey = pk();
{
Promise::new(alice())
.create_account()
.add_full_access_key(public_key.clone())
.delete_key(public_key.clone());
}
let public_key = near_crypto::PublicKey::try_from(public_key).unwrap();
let has_action = get_actions().any(|el| {
matches!(
el,
MockAction::DeleteKey { public_key: p , receipt_index: _, } if p == public_key
)
});
assert!(has_action);
}
}