Improved payment, fee, and reputation handling for contract sign

yanliu38 · yanliu38 · commit 792b7b2b5ae0 · 2025-01-02T16:14:54.000+01:00
diff --git a/common/src/account_transfers.rs b/common/src/account_transfers.rs
@@ -14,9 +14,10 @@ use sha2::{Digest, Sha256};
 use std::cell::RefCell;
 
 use crate::{
-    account_balance_add, account_balance_get, account_balance_sub, get_timestamp_ns,
-    ledger_add_reputation_change, slice_to_32_bytes_array, DccIdentity, TokenAmountE9s,
-    TransferError, LABEL_DC_TOKEN_TRANSFER, MINTING_ACCOUNT, MINTING_ACCOUNT_PRINCIPAL,
+    account_balance_add, account_balance_get, account_balance_sub, amount_as_string,
+    get_pubkey_from_principal, get_timestamp_ns, ledger_add_reputation_change,
+    slice_to_32_bytes_array, DccIdentity, TokenAmountE9s, TransferError, LABEL_DC_TOKEN_TRANSFER,
+    MINTING_ACCOUNT, MINTING_ACCOUNT_PRINCIPAL,
 };
 
 thread_local! {
@@ -63,34 +64,33 @@ pub fn ledger_funds_transfer(
 
 pub fn charge_fees_to_account_and_bump_reputation(
     ledger: &mut LedgerMap,
-    dcc_id_charge: &DccIdentity,
-    dcc_id_bump_reputation: &DccIdentity,
+    dcc_id: &DccIdentity,
     amount_e9s: TokenAmountE9s,
 ) -> Result<(), String> {
     if amount_e9s == 0 {
         return Ok(());
     }
-    let balance_from_after =
-        account_balance_get(&dcc_id_charge.as_icrc_compatible_account()) - amount_e9s;
+    let from_icrc1_account = dcc_id.as_icrc_compatible_account();
+    let balance_from_before = account_balance_get(&from_icrc1_account);
     match ledger_funds_transfer(
         ledger,
         // Burn 0 tokens, and transfer the entire amount_e9s to the fee accounts
         FundsTransfer::new(
-            dcc_id_charge.as_icrc_compatible_account(),
+            from_icrc1_account,
             MINTING_ACCOUNT,
             amount_e9s.into(),
             Some(fees_sink_accounts()),
             Some(get_timestamp_ns()),
             vec![],
             0,
-            balance_from_after,
+            balance_from_before.saturating_sub(amount_e9s),
             0,
         ),
     ) {
         Ok(_) => Ok(ledger_add_reputation_change(
             ledger,
-            dcc_id_bump_reputation,
-            amount_e9s.min(i64::MAX as TokenAmountE9s) as i64,
+            dcc_id,
+            amount_e9s as i64,
         )?),
         Err(e) => {
             info!("Failed to charge fees: {}", e);
@@ -101,26 +101,26 @@ pub fn charge_fees_to_account_and_bump_reputation(
 
 pub fn charge_fees_to_account_no_bump_reputation(
     ledger: &mut LedgerMap,
-    dcc_identity: &DccIdentity,
+    dcc_id_charge: &DccIdentity,
     amount_e9s: TokenAmountE9s,
 ) -> Result<(), String> {
     if amount_e9s == 0 {
         return Ok(());
     }
-    let balance_from_after =
-        account_balance_get(&dcc_identity.as_icrc_compatible_account()) - amount_e9s;
+    let from_icrc1_account = dcc_id_charge.as_icrc_compatible_account();
+    let balance_from_before = account_balance_get(&from_icrc1_account);
     match ledger_funds_transfer(
         ledger,
         // Burn 0 tokens, and transfer the entire amount_e9s to the fee accounts
         FundsTransfer::new(
-            dcc_identity.as_icrc_compatible_account(),
+            from_icrc1_account,
             MINTING_ACCOUNT,
             Some(amount_e9s),
             Some(fees_sink_accounts()),
             Some(get_timestamp_ns()),
             vec![],
             0,
-            balance_from_after,
+            balance_from_before.saturating_sub(amount_e9s),
             0,
         ),
     ) {
@@ -132,6 +132,78 @@ pub fn charge_fees_to_account_no_bump_reputation(
     }
 }
 
+pub enum IncreaseReputation {
+    None,
+    Sender,
+    Recipient,
+}
+
+pub fn do_funds_transfer(
+    ledger: &mut LedgerMap,
+    from_dcc_id: &DccIdentity,
+    to_icrc1_account: &IcrcCompatibleAccount,
+    transfer_amount_e9s: TokenAmountE9s,
+    fees_amount_e9s: TokenAmountE9s,
+    memo: &[u8],
+    increase_reputation: IncreaseReputation,
+) -> Result<String, String> {
+    let from_icrc1_account = from_dcc_id.as_icrc_compatible_account();
+
+    if transfer_amount_e9s == 0 {
+        return Ok("Nothing to transfer".to_string());
+    }
+    let balance_from_before = account_balance_get(&from_icrc1_account);
+    let balance_to_before = account_balance_get(to_icrc1_account);
+    if balance_from_before < transfer_amount_e9s + fees_amount_e9s {
+        return Err(format!(
+            "Not enough funds to transfer: {} < {} + {}",
+            balance_from_before, transfer_amount_e9s, fees_amount_e9s
+        ));
+    }
+    let balance_from_after = balance_from_before - transfer_amount_e9s - fees_amount_e9s;
+    match ledger_funds_transfer(
+        ledger,
+        FundsTransfer::new(
+            from_icrc1_account.clone(),
+            to_icrc1_account.clone(),
+            Some(fees_amount_e9s),
+            Some(fees_sink_accounts()),
+            Some(get_timestamp_ns()),
+            memo.to_vec(),
+            transfer_amount_e9s,
+            balance_from_after,
+            balance_to_before + transfer_amount_e9s,
+        ),
+    ) {
+        Ok(_) => {
+            let response = format!(
+                "Transferred {} tokens from {} \t to account {}, and charged fees {} tokens",
+                amount_as_string(transfer_amount_e9s),
+                from_icrc1_account,
+                to_icrc1_account,
+                amount_as_string(fees_amount_e9s)
+            );
+            match increase_reputation {
+                IncreaseReputation::None => (),
+                IncreaseReputation::Sender => {
+                    ledger_add_reputation_change(ledger, from_dcc_id, fees_amount_e9s as i64)?;
+                }
+                IncreaseReputation::Recipient => {
+                    let to_dcc_identity = to_icrc1_account
+                        .to_dcc_identity()
+                        .expect("Failed to get dcc identity from icrc1 account");
+                    ledger_add_reputation_change(ledger, &to_dcc_identity, fees_amount_e9s as i64)?;
+                }
+            }
+            Ok(response)
+        }
+        Err(e) => {
+            info!("Failed to charge fees: {}", e);
+            Err(e.to_string())
+        }
+    }
+}
+
 #[derive(CandidType, Deserialize, Clone, Debug, PartialEq, Eq, Hash)]
 pub struct IcrcCompatibleAccount {
     pub owner: Principal,
@@ -201,6 +273,14 @@ impl IcrcCompatibleAccount {
     pub fn from_bytes(bytes: &[u8]) -> Result<Self, std::io::Error> {
         borsh::from_slice(bytes)
     }
+
+    pub fn to_dcc_identity(&self) -> Option<DccIdentity> {
+        let owner_bytes = get_pubkey_from_principal(self.owner);
+        if owner_bytes.is_empty() {
+            return None;
+        }
+        DccIdentity::new_verifying_from_bytes(&owner_bytes).ok()
+    }
 }
 
 impl From<&str> for IcrcCompatibleAccount {
diff --git a/common/src/contract_sign_reply.rs b/common/src/contract_sign_reply.rs
@@ -6,8 +6,9 @@ use ledger_map::LedgerMap;
 use serde::{Deserialize, Serialize};
 
 use crate::{
-    amount_as_string, charge_fees_to_account_and_bump_reputation, contract_sign_fee_e9s,
-    contracts_cache_open_remove, fn_info, ContractSignRequestPayload, DccIdentity,
+    account_balance_get, amount_as_string, charge_fees_to_account_no_bump_reputation,
+    contract_sign_fee_e9s, contracts_cache_open_remove, do_funds_transfer, fn_info,
+    ledger_add_reputation_change, ContractSignRequestPayload, DccIdentity,
     LABEL_CONTRACT_SIGN_REPLY, LABEL_CONTRACT_SIGN_REQUEST,
 };
 
@@ -117,10 +118,11 @@ pub fn do_contract_sign_reply(
     reply_serialized: Vec<u8>,
     crypto_signature: Vec<u8>,
 ) -> Result<String, String> {
-    let dcc_id = DccIdentity::new_verifying_from_bytes(&pubkey_bytes).unwrap();
-    dcc_id.verify_bytes(&reply_serialized, &crypto_signature)?;
+    let provider_dcc_id = DccIdentity::new_verifying_from_bytes(&pubkey_bytes).unwrap();
+    provider_dcc_id.verify_bytes(&reply_serialized, &crypto_signature)?;
+    let provider_icrc1 = provider_dcc_id.as_icrc_compatible_account();
 
-    fn_info!("{}", dcc_id);
+    fn_info!("{}", provider_dcc_id);
 
     let cs_reply = ContractSignReply::try_from_slice(&reply_serialized).unwrap();
     let contract_id = cs_reply.contract_id();
@@ -134,14 +136,44 @@ pub fn do_contract_sign_reply(
     if pubkey_bytes != cs_req.provider_pubkey_bytes() {
         return Err(format!(
             "Contract signing reply signed and submitted by {} does not match the provider public key {} from contract req 0x{}",
-            dcc_id, DccIdentity::new_verifying_from_bytes(cs_req.provider_pubkey_bytes()).unwrap(), hex::encode(contract_id)
+            provider_dcc_id, DccIdentity::new_verifying_from_bytes(cs_req.provider_pubkey_bytes()).unwrap(), hex::encode(contract_id)
         ));
     }
+    let fees_e9s = contract_sign_fee_e9s(cs_req.payment_amount_e9s());
+
+    let requester_dcc_id =
+        DccIdentity::new_verifying_from_bytes(cs_req.requester_pubkey_bytes()).unwrap();
+    let requester_icrc1 = requester_dcc_id.as_icrc_compatible_account();
+    let requester_balance = account_balance_get(&requester_icrc1);
+
     let payload = ContractSignReplyPayload::new(reply_serialized, crypto_signature);
     let payload_serialized = borsh::to_vec(&payload).unwrap();
 
-    let fees = contract_sign_fee_e9s(cs_req.payment_amount_e9s());
-    charge_fees_to_account_and_bump_reputation(ledger, &dcc_id, &dcc_id, fees)?;
+    if cs_reply.sign_accepted() {
+        // Accepted: charge the requester the full amount
+        let charge_amount_e9s = cs_req.payment_amount_e9s() + fees_e9s;
+        if requester_balance < charge_amount_e9s {
+            return Err(format!(
+                "Requester {} does not have enough funds. At least {} tokens are required, and {} are available",
+                requester_icrc1,
+                amount_as_string(charge_amount_e9s),
+                amount_as_string(requester_balance)
+            ));
+        }
+        do_funds_transfer(
+            ledger,
+            &requester_dcc_id,
+            &provider_icrc1,
+            cs_req.payment_amount_e9s(),
+            fees_e9s,
+            cs_req.request_memo().as_bytes(),
+            crate::IncreaseReputation::Recipient,
+        )?;
+    } else {
+        // Else, charge the provider the response fee, and revert the requester's reputation
+        charge_fees_to_account_no_bump_reputation(ledger, &provider_dcc_id, fees_e9s)?;
+        ledger_add_reputation_change(ledger, &requester_dcc_id, -(fees_e9s as i64))?;
+    }
 
     ledger.upsert(
         LABEL_CONTRACT_SIGN_REPLY,
@@ -152,7 +184,7 @@ pub fn do_contract_sign_reply(
         contracts_cache_open_remove(contract_id);
         format!(
             "Contract signing reply submitted! Thank you. You have been charged {} tokens as a fee, and your reputation has been bumped accordingly",
-            amount_as_string(fees)
+            amount_as_string(fees_e9s)
         )
     })
     .map_err(|e| e.to_string())
diff --git a/common/src/contract_sign_request.rs b/common/src/contract_sign_request.rs
@@ -9,8 +9,8 @@ use serde::{Deserialize, Serialize, Serializer};
 use sha2::{Digest, Sha256};
 
 use crate::{
-    amount_as_string, charge_fees_to_account_and_bump_reputation, fn_info, AHashMap, DccIdentity,
-    TokenAmountE9s, LABEL_CONTRACT_SIGN_REQUEST,
+    account_balance_get, amount_as_string, charge_fees_to_account_and_bump_reputation, fn_info,
+    AHashMap, DccIdentity, TokenAmountE9s, LABEL_CONTRACT_SIGN_REQUEST,
 };
 
 pub type ContractId = Vec<u8>;
@@ -267,14 +267,29 @@ pub fn do_contract_sign_request(
 
     let contract_req = ContractSignRequest::try_from_slice(&request_serialized).unwrap();
 
-    let fees = contract_sign_fee_e9s(contract_req.payment_amount_e9s());
+    let fees_e9s = contract_sign_fee_e9s(contract_req.payment_amount_e9s());
+
+    let requester_dcc_id =
+        DccIdentity::new_verifying_from_bytes(contract_req.requester_pubkey_bytes()).unwrap();
+    let requester_icrc1 = requester_dcc_id.as_icrc_compatible_account();
+    let requester_balance = account_balance_get(&requester_icrc1);
+    let expected_min_balance = contract_req.payment_amount_e9s() + fees_e9s;
+
+    if requester_balance < expected_min_balance {
+        return Err(format!(
+            "Signing of this contract requires at least {} tokens. Requester {} has only {} tokens",
+            amount_as_string(expected_min_balance),
+            requester_icrc1,
+            amount_as_string(requester_balance)
+        ));
+    }
 
     let payload = ContractSignRequestPayload::new(&request_serialized, &crypto_signature).unwrap();
     let payload_bytes = borsh::to_vec(&payload).unwrap();
 
-    charge_fees_to_account_and_bump_reputation(ledger, &dcc_id, &dcc_id, fees)?;
-    let contract_id = payload.calc_contract_id();
+    charge_fees_to_account_and_bump_reputation(ledger, &dcc_id, fees_e9s)?;
 
+    let contract_id = payload.calc_contract_id();
     ledger.upsert(
         LABEL_CONTRACT_SIGN_REQUEST,
         contract_id,
@@ -284,7 +299,7 @@ pub fn do_contract_sign_request(
         format!(
             "Contract signing req 0x{} submitted! Thank you. You have been charged {} tokens as a fee, and your reputation has been bumped accordingly. Please check back for a response from the provider.",
             hex::encode(contract_id),
-            amount_as_string(fees)
+            amount_as_string(fees_e9s)
         )
     }).map_err(|e| e.to_string())
 }
diff --git a/common/src/lib.rs b/common/src/lib.rs
@@ -106,7 +106,7 @@ pub const MAX_PUBKEY_BYTES: usize = 32;
 pub const MEMO_BYTES_MAX: usize = 32;
 /// Reduction of reputations for all accounts, based on time (per block), in parts per million
 pub const REPUTATION_AGING_PPM: u64 = 1_000;
-pub const MAX_REPUTATION_INCREASE_PER_TX: i64 = DC_TOKEN_DECIMALS_DIV as i64 / 1000;
+pub const MAX_REPUTATION_INCREASE_PER_TX: i64 = DC_TOKEN_DECIMALS_DIV as i64 * 10; // Max 10 DC tokens per transaction
 pub const REWARD_HALVING_AFTER_BLOCKS: u64 = 210_000; // halve the rewards every 210000 reward distributions
 pub const DATA_PULL_BYTES_BEFORE_LEN: u16 = 16; // How many bytes before the pulled data should be compared as a quick sanity check
 
diff --git a/ic-canister/tests/test_canister.rs b/ic-canister/tests/test_canister.rs
